ewise: caller-defined scalar compositions

ewise: caller-defined scalar compositions#

T.cuda.tile.ewise(function, dst, *inputs, scope="thread") applies a typed scalar TIRx function independently to register elements. This is the one CUDA tile operation whose atom is a caller-defined expression rather than one PTX instruction. It can represent exp, SiLU, scale/bias, casts, or another scalar combination. Use add / sub / mul, fma, or cvt when the native instruction and its qualifiers are the contract.

The scalar atom#

Tensor operands have coordinates (*group,*repeat). There is no trailing PTX register-group axis: every coordinate after the group prefix selects one scalar invocation. Inputs and destination have matching extents and belong to the executing thread. Read-only zero-stride input views express tensor broadcasts; a destination must not overlap itself.

Scalar arguments broadcast to every invocation. Their types must match the function parameters, and its result type must match the destination dtype. The function supports scalar bindings and returns, including an if/else with a scalar return on both paths. It does not express reductions, memory access algorithms, or inter-thread communication.

Example: SiLU, scale, bias#

This complete kernel uses four scalar register values per thread. Its memory instructions have a four-word suffix; ewise treats those same four coordinates as scalar repeats.

from tvm.script import tirx as T
from tvm.backend.cuda.axis import laneid, warpid
from tvm.tirx.layout import S, TileLayout

reg_layout = TileLayout(S[(2, 32, 4) : (1 @ warpid, 1 @ laneid, 1)])


@T.function(private=True)
def silu_affine(x: T.float32, scale: T.float32, bias: T.float32) -> T.float32:
    y = x / (T.float32(1) + T.exp(-x))
    return y * scale + bias


@T.function
def activation(A: T.Tensor((2, 32, 4), "float32"), Out: T.Tensor((2, 32, 4), "float32")):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
    R = T.alloc_local((2, 32, 4), "float32", layout=reg_layout)
    T.cuda.tile.apply(T.ptx["ld.global.v4.b32"], R, A, scope="cta")
    T.cuda.tile.ewise(silu_affine, R, R, T.float32(2), T.float32(1), scope="cta")
    T.cuda.tile.apply(T.ptx["st.global.v4.b32"], Out, R, scope="cta")

There is one vector load, four inlined scalar expressions, and one vector store per thread. The expression computes silu(x)*2+1; it does not mean silu(2*x+1). The ordinary scalar compiler lowers the expression; ewise does not promise one PTX instruction, choose a packed form, or fuse an unspecified activation. For an exact exponential instruction use ex2 / lg2.

Shared-memory load/compute/store#

Both the destination and tensor inputs of ewise must be local registers. The following complete kernel makes the shared staging explicit. The final read deliberately uses another thread’s row, so the second CTA barrier has a visible purpose.

@T.function(private=True)
def affine(x: T.float32) -> T.float32:
    return x * T.float32(2) + T.float32(1)


@T.function
def shared_affine(A: T.Tensor((64, 4), "float32"), Out: T.Tensor((64, 4), "float32")):
    T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=64))
    tid = T.cuda.thread_idx("x")
    S = T.alloc_shared((64, 4), "float32", layout="default", align=16)
    for j in range(4):
        S[tid, j] = A[tid, j]
    T.cuda.cta_sync()
    sv = S.view(2, 32, 4)
    R = T.alloc_local((2, 32, 4), "float32", layout=reg_layout)
    T.cuda.tile.apply(T.ptx["ld.shared.v4.b32"], R, sv, scope="cta")
    T.cuda.tile.ewise(affine, R, R, scope="cta")
    T.cuda.tile.apply(T.ptx["st.shared.v4.b32"], sv, R, scope="cta")
    T.cuda.cta_sync()
    for j in range(4):
        Out[tid, j] = S[(tid + 1) % 64, j]

The two views refer to the existing shared and local allocations. The tile operation adds no temporary storage or barrier. If the shared store used a different ownership map that overwrote another thread’s unread input, an additional barrier would be needed between all loads and all stores.

Compile and verify both examples:

import numpy as np
import tvm

x = np.linspace(-2, 2, 256, dtype="float32")
for kernel, shape, ref in [
    (activation, (2, 32, 4), (x / (1 + np.exp(-x)) * 2 + 1).reshape(2, 32, 4)),
    (shared_affine, (64, 4), np.roll((x * 2 + 1).reshape(64, 4), -1, axis=0)),
]:
    compiled = tvm.tirx.compile(kernel, target="cuda", pipeline="tirx")
    out = tvm.runtime.tensor(np.zeros(shape, dtype="float32"), tvm.cuda())
    compiled(tvm.runtime.tensor(x.reshape(shape), tvm.cuda()), out)
    np.testing.assert_allclose(out.numpy(), ref, rtol=2e-5, atol=2e-5)

Implementation: python/tvm/backend/cuda/tile/patterns/ewise.py.

Contents