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)[source]#

Mosaic GPU compiler parameters.

Parameters:

  • approx_math (bool)

  • dimension_semantics (Sequence[DimensionSemantics] | None)

  • max_concurrent_steps (int)

  • delay_release (int)

  • profile_space (int)

  • profile_dir (str)

  • lowering_semantics (mgpu.core.LoweringSemantics)

approx_math#

If True, the compiler is allowed to use approximate implementations of some math operations, e.g. exp. Defaults to False.

Type:

bool

dimension_semantics#

A list of dimension semantics for each grid dimension of the kernel. Either “parallel” for dimensions that can execute in any order, or “sequential” for dimensions that must be executed sequentially.

Type:

Sequence[DimensionSemantics] | None

max_concurrent_steps#

The maximum number of sequential stages that are active concurrently. Defaults to 1.

Type:

int

delay_release#

The number of steps to wait before reusing the input/output references. Defaults to 0, and must be strictly smaller than max_concurrent_steps. Generally, you’ll want to set it to 1 if you don’t await the WGMMA in the body.

Type:

int

profile_space#

The number of profiler events that can be collected in a single invocation. It is undefined behavior if a thread collects more events than this.

Type:

int

profile_dir#

The directory to which profiling traces will be written to.

Type:

str

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

  • approx_math (bool)

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

  • max_concurrent_steps (int)

  • delay_release (int)

  • profile_space (int)

  • profile_dir (str)

  • lowering_semantics (mgpu.core.LoweringSemantics)

Return type:

None

Methods

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

Attributes

BACKEND

approx_math

delay_release

dimension_semantics

lowering_semantics

max_concurrent_steps

profile_dir

profile_space