Convert LLType to C++ type string.

Get the code generators for a special struct type.

Unparse a node with custom attributes.

pass named attrs (e.g., tt.contiguity) from Triton to Triton

Check if the type is a tensor (ND Array)

Write a line with proper indentation.

pass named attrs (e.g., tt.contiguity) from Triton to Triton

NOTE: this function should only be called with atomic support. memory over PCI-e does not support atomic r/w. DON'T use this function on such pla

Promote types following C++ implicit conversion rules. Returns: (promoted_left, promoted_right, was_promoted): The promoted type

Resolve a name to its value. Priority: 1. If name is local (defined in kernel), return None (don't resolve)

Shuffle idx: read *value* from *lane* (broadcast if lane is constant).

Check CUDA error and raise exception if needed.

Initializes contexts for triton_dist AllGather-GEMM and GEMM-ReduceScatter operations.

Ensure the tensor is materialized before access.

Helper to define variables from AST nodes (targets of assignments, args, etc.) Handles: - ast.Name: Simple variable

Build a kernel from the given module and verify correctness.

Add a variable to the active scope

Check if this is the last PP stage

Applies RMS Normalization

Write line to struct buffer with newline.

Build a kernel from the given module and verify correctness.

Warp-level atomic add: lane 0 performs atomic_add, result is broadcast to all lanes. AMD equivalent of the NVIDIA atomic_add_per_warp. Uses wavef

Call shmem barrier on stream: mori_shmem when backend is mori_shmem, else rocshmem.

Get all special attributes from a node as a dictionary. Args: node: AST node to extract attributes from Returns:

Get type of a variable from scope/context.

Check if name is defined in any local scope (shadows global). This is the key method for determining symbol resolution prior

Format a value for display in dump.

Define a name in the current scope. This marks the name as locally defined, which means it shadows any global defini

Check if a node is an IR node (EnumDef or SpecialStructDef).

Applies RMS Normalization using flashinfer.

Pack two 32-bit values into one 64-bit value (val0: low 32, val1: high 32).

Raise a pass error with formatted message. Args: message: Error message node: AST node where the error occurred

Convenience function to register a pass.

Shuffle up: each lane reads from (laneid - delta), clamped to 0.

Initializes and shards MLP parameters for Tensor Parallelism. mlp: A standard nn.Module MLP (e.g., from HuggingFace Transformers).

Checks if 'npu-smi' is available on the system's PATH.

Recursively scan statements for variables assigned anywhere (including nested For/If/While). Returns dict of var_name -> (var_cpp_type, var_ll

allgather gemm C = all_gather(A) * B Args: A (torch.Tensor<float>): local matmul A matrix. shape: [M_per_rank, K] B (torch.Te