jax.experimental.pallas.mosaic_gpu.CompilerParams

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

Mosaic GPU 编译器参数。

参数:

• approx_math (bool)

• dimension_semantics (Sequence[DimensionSemantics] | None)

• max_concurrent_steps (int)

• delay_release (int)

• unsafe_no_auto_barriers (bool)

• profile_space (int)

• profile_dir (str)

• lowering_semantics (mgpu.core.LoweringSemantics)

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

unsafe_no_auto_barriers

如果为 True，则 Pallas 永远不会自动插入 barrier 指令，以确保加载和存储的同步语义。 目前，插入是以保守的方式完成的，可能会降低性能。 要安全使用此标志，必须满足（至少）两个条件。 首先，没有内存区域曾经被同一个线程读取和写入（异步复制由后台线程执行，不计入此规则）。 其次，没有线程曾经调用 commit_smem()，从提交的 SMEM 读取，然后发出覆盖该区域的异步复制（这是一个非常人为且极不可能的场景）。

类型:

bool

profile_space

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

类型:

int

profile_dir

将分析跟踪写入到的目录。

类型:

str

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

参数:

• approx_math (bool)

• dimension_semantics (Sequence[DimensionSemantics] | None)

• max_concurrent_steps (int)

• delay_release (int)

• unsafe_no_auto_barriers (bool)

• profile_space (int)

• profile_dir (str)

• lowering_semantics (mgpu.core.LoweringSemantics)

返回类型:

无

方法

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

属性

后端
approx_math
delay_release
dimension_semantics
lowering_semantics
max_concurrent_steps
profile_dir
profile_space
unsafe_no_auto_barriers