tvm_tirx.tirx.script.gpu_thread_allreduce

tvm_tirx.tirx.script.gpu_thread_allreduce#

tvm_tirx.tirx.script.gpu_thread_allreduce(combine, identity, values, predicate, destinations, thread_axes, *, ty=None, loc=UNKNOWN_LOC)#

Perform an all-reduce inside a thread block.

Parameters:

  • combine (tvm_tirx.ir.LambdaExpr) – Typed combining lambda with parameters ordered as all left-hand values followed by all right-hand values. Its body returns a scalar for one result or a Tuple of results.

  • identity (Expr or Sequence[Expr]) – Identity value for each reduction result.

  • values (Expr or Sequence[Expr]) – Values contributed by the current thread.

  • predicate (PrimExpr) – Boolean participation predicate. Inactive threads contribute identities.

  • destinations (Expr or Sequence[Expr]) – Tensor loads at index zero of one-element result temporaries, optionally cast for boolean storage. Each temporary must be accessed only at index zero.

  • thread_axes (Expr or Sequence[Expr]) – Thread variables participating in the reduction.

  • loc (Location)

Returns:

call (Expr) – The void call expression with six explicit operands.