jax.experimental.pallas.mosaic_gpu.GPUCompilerParams

class jax.experimental.pallas.mosaic_gpu.GPUCompilerParams(*, approx_math=False, dimension_semantics=None, max_concurrent_steps=1, delay_release=0, profile_space=0, profile_dir='', lowering_semantics=LoweringSemantics.Lane)

Mosaic GPU 编译器参数。

参数:

• approx_math (bool)

• dimension_semantics (DimensionSemantics 序列[DimensionSemantics] | None)

• max_concurrent_steps (int)

• delay_release (int)

• profile_space (int)

• profile_dir (str)

• lowering_semantics (mgpu.core.LoweringSemantics)

approx_math

如果为 True，则编译器允许使用某些数学运算的近似实现，例如 exp。默认为 False。

类型:

bool

dimension_semantics

内核的每个网格维度的维度语义列表。 “parallel” 用于可以以任何顺序执行的维度，或 “sequential” 用于必须按顺序执行的维度。

类型:

DimensionSemantics 序列[DimensionSemantics] | None

max_concurrent_steps

并发活动的最大顺序阶段数。默认为 1。

类型:

int

delay_release

重用输入/输出引用之前等待的步数。 默认为 0，并且必须严格小于 max_concurrent_steps。 通常，如果您不在主体中等待 WGMMA，则需要将其设置为 1。

类型:

int

profile_space

单次调用中可以收集的分析器事件数。 如果线程收集的事件数量超过此数量，则行为未定义。

类型:

int

profile_dir

将向其写入性能分析跟踪的目录。

类型:

str

__init__(*, approx_math=False, dimension_semantics=None, max_concurrent_steps=1, delay_release=0, profile_space=0, profile_dir='', lowering_semantics=LoweringSemantics.Lane)

参数:

• approx_math (bool)

• dimension_semantics (DimensionSemantics 序列[DimensionSemantics] | None | None)

• max_concurrent_steps (int)

• delay_release (int)

• profile_space (int)

• profile_dir (str)

• lowering_semantics (mgpu.core.LoweringSemantics)

返回类型:

None

方法

__init__(*[, approx_math, ...])

属性

PLATFORM
approx_math
delay_release
dimension_semantics
lowering_semantics
max_concurrent_steps
profile_dir
profile_space