Direct PTX Instructions

Direct PTX Instructions#

The CUDA backend makes a table-driven PTX instruction namespace available when authoring TIRx with TVMScript. This is an authoring namespace, not a separate IR dialect:

from tvm.script import tirx as Tx
Tx.ptx.add.rn.f32(dst, lhs, rhs)

Each Tx.ptx call must match a form registered in tvm.backend.cuda.ptx.table. The table validates the instruction family, modifier combination, operand count, and operand types while constructing one tirx.ptx.* IR call. A new instruction is added by extending the table and, when it needs a new representation, its code-generation support.

This layer is below the layout-aware Tx.tile.* primitives and does not run tile-primitive dispatch.

Supported instruction families#

The current instruction table has 318 entries across the following 94 families. These are the Python attribute names written after Tx.ptx; each family accepts only the modifier and operand combinations declared by its table entries.

  • abs

  • activemask

  • add

  • and_

  • applypriority

  • atom

  • bar

  • barrier

  • bfe

  • bfi

  • bfind

  • bmsk

  • brev

  • clusterlaunchcontrol

  • clz

  • cnot

  • copysign

  • cos

  • cp

  • createpolicy

  • cvt

  • cvt_pack

  • cvta

  • discard

  • div

  • dp2a

  • dp4a

  • elect_sync

  • ex2

  • fence

  • fma

  • fns

  • getctarank

  • griddepcontrol

  • isspacep

  • ld

  • ldmatrix

  • ldu

  • lg2

  • lop3

  • mad

  • mad24

  • mapa

  • match

  • max

  • mbarrier

  • min

  • mma

  • mov

  • movmatrix

  • mul

  • mul24

  • multimem_ld_reduce

  • multimem_red

  • multimem_st

  • neg

  • not_

  • or_

  • popc

  • prefetch

  • prefetchu

  • prmt

  • rcp

  • red

  • red_async

  • redux_sync

  • rem

  • rsqrt

  • sad

  • selp

  • set

  • setmaxnreg

  • setp

  • shf

  • shfl_sync

  • shl

  • shr

  • sin

  • slct

  • sqrt

  • st

  • st_async

  • st_bulk

  • stmatrix

  • sub

  • szext

  • tanh

  • tcgen05

  • tensormap_cp_fenceproxy

  • tensormap_replace

  • testp

  • vote_sync

  • wgmma

  • xor

The generated tvm/script/tirx.pyi stub provides editor completion for the valid modifier names and call signatures. The instruction table is the complete source of modifier domains, operand layouts, and combination constraints.

Instruction forms#

Attribute chains fill PTX modifier slots. Python keywords have a trailing underscore, and PTX’s :: separator is written as a double underscore:

Tx.ptx.ld.acquire.gpu.global_.b32(value, pointer)
Tx.ptx.mbarrier.arrive.shared__cluster.b64(barrier, count)

The indexed form accepts the exact PTX spelling of a registered form and uses the same table lookup:

Tx.ptx["st.weak.shared::cta.b32"](pointer, value)

Destination registers are leading operands, so PTX calls are statements. Predication is keyword-only:

Tx.ptx.ld.global_.b32(value, pointer, pred=predicate, preserve_dst=True)

Tx.ptx.pred(value) marks a register operand whose PTX type is .pred; Tx.ptx.addr(base, byte_offset) forms an immediate-offset address for instructions that accept one; and Tx.ptx.SINK represents PTX’s sink operand _ in table-marked positions. The generated tvm/script/tirx.pyi stub provides editor completion for addr, registered instruction families, and modifier chains. pred and SINK are available at runtime but are not declared in the current stub.

PTX versus CUDA helpers#

Use Tx.ptx for one table-described PTX instruction. Use Tx.cuda for backend helpers that require multiple statements, C/C++ expressions, descriptor packing, or other behavior that cannot be represented as one PTX instruction. Tx.ptx_legacy only preserves historical spellings needed by compatibility passes and is not the API for new kernels.

Namespace API#

class tvm.backend.cuda.ptx.PTXNamespace(table=None)#

T.ptx — table-driven PTX instruction namespace.

static pred(value)#

Tag an operand as a .pred register – see PredArg.

static addr(base, byte_offset)#

Form [base+byte_offset] for an eligible PTX address operand.

__getitem__(text)#

Exact-PTX-text form, e.g. T.ptx["st.weak.shared::cta.b32"].

Table-driven PTX dialect (T.ptx).

One table (table), one generic engine (engine), thin generators (gen_helpers`/gen_stubs`, :mod:`.gen_coverage`). Importing this package registers every table entry as a TVM Op with a generic codegen; the ``T.ptx namespace itself is installed by the CUDA backend’s register_backend() via script_namespaces().

tvm.backend.cuda.ptx.register_addr() → None#

Register the pure address-expression op consumed by PTX instructions.

tvm.backend.cuda.ptx.register_table(table: dict[str, InstructionEntry]) → None#

Register every table entry as a TVM Op + generic codegen.

Contents