tilelang.cuda.language.reduce_op¶

CUDA dialect of the reduction operators: the common ops plus CUDA knobs.

Functions¶

reduce_max(buffer, out[, dim, clear, batch, ...])

Perform reduce max, with the CUDA NaN-propagation knob.

reduce_min(buffer, out[, dim, clear, batch, ...])

Perform reduce min, with the CUDA NaN-propagation knob.

reduce_absmax(buffer, out[, dim, clear, batch, ...])

Perform reduce absolute max, with the CUDA NaN-propagation knob.

Module Contents¶

tilelang.cuda.language.reduce_op.reduce_max(buffer, out, dim=-1, clear=True, batch=1, nan_propagate=False, annotations=None)¶

Perform reduce max, with the CUDA NaN-propagation knob.

Same semantics as the common tilelang.language.reduce_op.reduce_max(). nan_propagate is meaningful for float16/bfloat16 only: when True the reduction lowers to __hmax_nan so NaNs propagate; when False (default) __hmax returns the non-NaN operand. Targets without these intrinsics reject the annotation at compile time.

Parameters:

  • buffer (tvm.tirx.Buffer)

  • out (tvm.tirx.Buffer)

  • dim (int)

  • clear (bool)

  • batch (int)

  • nan_propagate (bool)

  • annotations (dict | None)

Return type:

None

tilelang.cuda.language.reduce_op.reduce_min(buffer, out, dim=-1, clear=True, batch=1, nan_propagate=False, annotations=None)¶

Perform reduce min, with the CUDA NaN-propagation knob.

See reduce_max(); this lowers to __hmin_nan/__hmin.

Parameters:

  • buffer (tvm.tirx.Buffer)

  • out (tvm.tirx.Buffer)

  • dim (int)

  • clear (bool)

  • batch (int)

  • nan_propagate (bool)

  • annotations (dict | None)

Return type:

None

tilelang.cuda.language.reduce_op.reduce_absmax(buffer, out, dim=-1, clear=True, batch=1, nan_propagate=False, annotations=None)¶

Perform reduce absolute max, with the CUDA NaN-propagation knob.

See reduce_max().

Parameters:

  • buffer (tvm.tirx.Buffer)

  • out (tvm.tirx.Buffer)

  • dim (int)

  • clear (bool)

  • batch (int)

  • nan_propagate (bool)

  • annotations (dict | None)

Return type:

None