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='')

Mosaic GPU 编译器参数。

参数:

• approx_math (bool)

• dimension_semantics (Sequence[DimensionSemantics] | None)

• max_concurrent_steps (int)

• delay_release (int)

• profile_space (int)

• profile_dir (str)

approx_math

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

类型:

bool

dimension_semantics

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

类型:

Sequence[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='')

参数:

• approx_math (bool)

• dimension_semantics (Sequence[DimensionSemantics] | None | None)

• max_concurrent_steps (int)

• delay_release (int)

• profile_space (int)

• profile_dir (str)

返回类型:

None

方法

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

属性

PLATFORM
approx_math
delay_release
dimension_semantics
max_concurrent_steps
profile_dir
profile_space