apache · spectrometerHBH · Jul 31, 2026 · Aug 1, 2026
diff --git a/.agents/skills/tir-test/SKILL.md b/.agents/skills/tir-test/SKILL.md
@@ -2,58 +2,226 @@ Run the full TIRX test suite.
 
 ## Steps
 
-1. Install the kernel package and select the least busy GPU:
+Run all commands below in the same Bash shell from the TVM repository root.
+
+1. Point Python at this workspace's repos and select every eligible idle Blackwell GPU.
+   `TIR_TEST_GPUS` may explicitly provide physical GPU IDs, and
+   `TIR_TEST_NUM_GPUS` may optionally cap the automatically selected count. Otherwise,
+   selection uses all eligible GPUs, ordered by least memory in use.
    ```bash
-   pip install -e /path/to/tirx-kernels-staging   # or sibling tirx-kernels checkout
-   export CUDA_VISIBLE_DEVICES=$(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits | sort -t',' -k2 -n | head -1 | cut -d',' -f1 | tr -d ' ')
-   export TVM_PATH=/path/to/tvm
-   export PYTHONPATH="${TVM_PATH}/python"
-   export TVM_LIBRARY_PATH="${TVM_PATH}/build/lib"
+   export WORKSPACE=/path/to/workspace
+   export TIR_TEST_NUM_GPUS="${TIR_TEST_NUM_GPUS:-}"
+   export TIR_TEST_MIN_FREE_GB="${TIR_TEST_MIN_FREE_GB:-32}"
+   export TIR_TEST_MAX_USED_MIB="${TIR_TEST_MAX_USED_MIB:-256}"
+   export TIR_TEST_XDIST_WORKERS="${TIR_TEST_XDIST_WORKERS:-16}"
+
+   [[ -z "$TIR_TEST_NUM_GPUS" || "$TIR_TEST_NUM_GPUS" =~ ^[1-9][0-9]*$ ]] || {
+     echo "TIR_TEST_NUM_GPUS must be empty or a positive integer" >&2
+     exit 2
+   }
+   [[ "$TIR_TEST_MIN_FREE_GB" =~ ^[0-9]+$ ]] || {
+     echo "TIR_TEST_MIN_FREE_GB must be a non-negative integer" >&2
+     exit 2
+   }
+   [[ "$TIR_TEST_MAX_USED_MIB" =~ ^[0-9]+$ ]] || {
+     echo "TIR_TEST_MAX_USED_MIB must be a non-negative integer" >&2
+     exit 2
+   }
+   [[ "$TIR_TEST_XDIST_WORKERS" =~ ^[1-9][0-9]*$ ]] || {
+     echo "TIR_TEST_XDIST_WORKERS must be a positive integer" >&2
+     exit 2
+   }
+
+   if [[ -n "${TIR_TEST_GPUS:-}" ]]; then
+     IFS=, read -r -a TIR_TEST_SELECTED_GPUS <<< "${TIR_TEST_GPUS// /}"
+   else
+     mapfile -t TIR_TEST_SELECTED_GPUS < <(
+       nvidia-smi \
+         --query-gpu=index,memory.total,memory.used,utilization.gpu,compute_cap \
+         --format=csv,noheader,nounits \
+       | awk -F, \
+           -v min_free_mib="$((TIR_TEST_MIN_FREE_GB * 1024))" \
+           -v max_used_mib="$TIR_TEST_MAX_USED_MIB" '
+           {
+             for (i = 1; i <= 5; ++i) {
+               gsub(/^[[:space:]]+|[[:space:]]+$/, "", $i)
+             }
+             free_mib = $2 - $3
+             if (($5 + 0) >= 10 && free_mib >= min_free_mib &&
+                 ($3 + 0) <= max_used_mib && ($4 + 0) == 0) {
+               print $1, $3
+             }
+           }' \
+       | sort -k2,2n -k1,1n \
+       | awk '{print $1}'
+     )
+     if [[ -n "$TIR_TEST_NUM_GPUS" ]] && \
+        (( ${#TIR_TEST_SELECTED_GPUS[@]} > TIR_TEST_NUM_GPUS )); then
+       TIR_TEST_SELECTED_GPUS=(
+         "${TIR_TEST_SELECTED_GPUS[@]:0:TIR_TEST_NUM_GPUS}"
+       )
+     fi
+   fi
+
+   (( ${#TIR_TEST_SELECTED_GPUS[@]} > 0 )) || {
+     echo "no eligible CUDA compute capability >= 10 GPU is available" >&2
+     exit 1
+   }
+   for gpu in "${TIR_TEST_SELECTED_GPUS[@]}"; do
+     [[ "$gpu" =~ ^[0-9]+$ ]] || {
+       echo "invalid physical GPU ID: $gpu" >&2
+       exit 2
+     }
+   done
+   if [[ -z "${TIR_TEST_GPUS:-}" && -n "$TIR_TEST_NUM_GPUS" ]] && \
+      (( ${#TIR_TEST_SELECTED_GPUS[@]} < TIR_TEST_NUM_GPUS )); then
+     echo "warning: requested $TIR_TEST_NUM_GPUS GPUs, using ${#TIR_TEST_SELECTED_GPUS[@]}"
+   fi
+
+   export CUDA_VISIBLE_DEVICES
+   CUDA_VISIBLE_DEVICES="$(IFS=,; printf '%s' "${TIR_TEST_SELECTED_GPUS[*]}")"
+   export TIRX_KERNEL_TEST_MIN_FREE_GB="$TIR_TEST_MIN_FREE_GB"
+   export PYTHONPATH="${WORKSPACE}/tirx-kernels:${WORKSPACE}/tvm/python"
+   export TVM_LIBRARY_PATH="${WORKSPACE}/tvm/build/lib"
+
+   PYTHON_BIN="${PYTHON_BIN:-$(command -v python || true)}"
+   [[ -n "$PYTHON_BIN" && -x "$PYTHON_BIN" ]] || {
+     echo "set PYTHON_BIN to the test environment's Python executable" >&2
+     exit 1
+   }
+   export PATH="$(dirname "$PYTHON_BIN"):$PATH"
+   "$PYTHON_BIN" -c 'import pytest, torch, tvm' >/dev/null || {
+     echo "PYTHON_BIN must provide pytest and torch and import this TVM worktree" >&2
+     exit 1
+   }
+   if ! command -v ninja >/dev/null; then
+     NINJA_BIN_DIR="$("$PYTHON_BIN" -c 'import ninja; print(ninja.BIN_DIR)' 2>/dev/null || true)"
+     if [[ -n "$NINJA_BIN_DIR" && -x "$NINJA_BIN_DIR/ninja" ]]; then
+       export PATH="$NINJA_BIN_DIR:$PATH"
+     fi
+   fi
+   command -v ninja >/dev/null || {
+     echo "ninja is required by kernel JIT setup but is not on PATH" >&2
+     exit 1
+   }
+
+   echo "selected physical GPUs: $CUDA_VISIBLE_DEVICES"
+   for logical_id in "${!TIR_TEST_SELECTED_GPUS[@]}"; do
+     echo "  cuda:$logical_id -> physical GPU ${TIR_TEST_SELECTED_GPUS[$logical_id]}"
+   done
    ```
 
-2. Start the GPU monitor in the background so we can detect if anyone else lands on the same GPU mid-run:
+   Automatic selection uses every eligible idle device. Set
+   `TIR_TEST_NUM_GPUS=4` to cap the count or `TIR_TEST_GPUS=3,4,5,6` to override
+   selection with exact physical IDs.
+
+2. Start one GPU monitor per selected physical GPU. A baseline foreign process
+   is a selection race; stop and select a different GPU before running tests.
    ```bash
-   GPU_LOG="/tmp/tir_test_gpu_${CUDA_VISIBLE_DEVICES}.log"
-   bash .agents/scripts/monitor_gpu.sh --gpu "$CUDA_VISIBLE_DEVICES" --interval 5 --log "$GPU_LOG" &
-   MON_PID=$!
-   trap 'kill $MON_PID 2>/dev/null' EXIT
+   TIR_TEST_MONITOR_PIDS=()
+   TIR_TEST_GPU_LOGS=()
+
+   cleanup_tir_test_monitors() {
+     local pid
+     for pid in "${TIR_TEST_MONITOR_PIDS[@]}"; do
+       kill "$pid" 2>/dev/null || true
+     done
+     for pid in "${TIR_TEST_MONITOR_PIDS[@]}"; do
+       wait "$pid" 2>/dev/null || true
+     done
+   }
+   trap cleanup_tir_test_monitors EXIT
+
+   for gpu in "${TIR_TEST_SELECTED_GPUS[@]}"; do
+     log="/tmp/tir_test_gpu_${gpu}_$$.log"
+     bash .agents/scripts/monitor_gpu.sh \
+       --gpu "$gpu" --interval 5 --log "$log" &
+     TIR_TEST_MONITOR_PIDS+=("$!")
+     TIR_TEST_GPU_LOGS+=("$log")
+   done
+   sleep 1
+
+   baseline_conflict=0
+   for log in "${TIR_TEST_GPU_LOGS[@]}"; do
+     if grep -q '\[FOREIGN\]' "$log"; then
+       echo "foreign process present at baseline: $log" >&2
+       baseline_conflict=1
+     fi
+   done
+   if (( baseline_conflict )); then
+     cleanup_tir_test_monitors
+     trap - EXIT
+     exit 1
+   fi
    ```
 
-3. Import gate — bench workloads: fail fast if any kernel listed in `workloads.yaml` fails to import:
+3. Import gate: bench workloads. Fail fast if any kernel listed in
+   `workloads.yaml` fails to import.
    ```bash
-   python -m tirx_kernels.bench_suite --check-imports
+   "$PYTHON_BIN" -m tirx_kernels.bench_suite --check-imports
    ```
-   A non-zero exit means a pinned workload kernel failed to import — fix it before proceeding.
+   A non-zero exit means a pinned workload kernel failed to import. Fix it before
+   proceeding.
 
-4. Full kernel import gate (correctness test suite coverage):
+4. Full kernel import gate for correctness-suite coverage:
    ```bash
-   python -m tirx_kernels.registry --cc 10 --strict
+   "$PYTHON_BIN" -m tirx_kernels.registry --cc 10 --strict
    ```
 
 5. Run the full test suite with xdist parallelism:
    ```bash
-   pytest tests/python/tirx/ -n 16
+   "$PYTHON_BIN" -m pytest tests/python/tirx/ -n "$TIR_TEST_XDIST_WORKERS"
    ```
 
-6. Stop the monitor and check for foreign GPU usage during the run:
+   `tests/python/tirx/conftest.py` assigns each xdist worker's PyTorch current
+   device round-robin across the visible GPUs. Registry correctness tests also
+   take a per-device lock, so at most one large registered kernel runs on each
+   GPU concurrently. Tests that explicitly use `tvm.cuda(0)` still run on the
+   first selected physical GPU. MegaMoE remains skipped because it requires its
+   dedicated multi-process scheduler, not ordinary xdist device assignment.
+
+6. Stop every monitor and check every selected GPU for foreign usage:
    ```bash
-   kill $MON_PID 2>/dev/null; wait $MON_PID 2>/dev/null
-   grep -E 'FOREIGN USER|\[FOREIGN\]' "$GPU_LOG" || echo "no foreign GPU usage observed"
+   cleanup_tir_test_monitors
+   trap - EXIT
+
+   foreign_usage=0
+   for i in "${!TIR_TEST_GPU_LOGS[@]}"; do
+     gpu="${TIR_TEST_SELECTED_GPUS[$i]}"
+     log="${TIR_TEST_GPU_LOGS[$i]}"
+     if grep -qE 'FOREIGN USER|\[FOREIGN\]' "$log"; then
+       echo "GPU $gpu foreign-process events:"
+       grep -E 'FOREIGN USER|\[FOREIGN\]' "$log"
+       foreign_usage=1
+     else
+       echo "GPU $gpu: no foreign GPU usage observed"
+     fi
+   done
    ```
 
-7. Report results: total passed, failed, skipped, errors — and the import-gate results from steps 3–4. If any foreign-user events are present in step 6, mention them — flaky failures should be re-evaluated on a clean GPU before being attributed to code changes.
+7. Report results: selected physical GPUs, elapsed time, total passed, failed,
+   skipped, and errors, plus both import-gate results. If any foreign-user events
+   are present in step 6, mention the affected physical GPUs. Flaky failures
+   should be re-evaluated on clean GPUs before being attributed to code changes.
 
 ## Failure triage rules
 
-**CRITICAL: Never pipe test output to `tail` or `grep` when diagnosing failures. Always capture and read full logs.**
+**CRITICAL: Never pipe test output to `tail` or `grep` when diagnosing failures.
+Always capture and read full logs.**
 
 Classify every failure into one of these categories:
 
-- **A — Environment/import error**: Module not found, missing dependency, collection error. These are not caused by code changes.
-- **B — Real kernel correctness regression**: Assertion failures (cosine_sim, numerical diff), `CUDA: unspecified launch failure`, or wrong results. **These MUST be investigated and fixed if caused by current changes.**
-- **C — Secondary xdist crash**: `KeyError: <WorkerController gwXX>` after a worker abort. The KeyError itself is noise — find the underlying cause (usually category B in another worker).
+- **A - Environment/import error**: Module not found, missing dependency,
+  collection error. These are not caused by code changes.
+- **B - Real kernel correctness regression**: Assertion failures (`cosine_sim`,
+  numerical diff), `CUDA: unspecified launch failure`, or wrong results. These
+  MUST be investigated and fixed if caused by current changes.
+- **C - Secondary xdist crash**: `KeyError: <WorkerController gwXX>` after a
+  worker abort. The KeyError itself is noise; find the underlying cause (usually
+  category B in another worker).
+
+Never dismiss a failure as pre-existing without evidence. If a test fails:
 
-**Never dismiss a failure as "pre-existing" without evidence.** If a test fails:
 1. Check whether the test touches code you changed.
-2. If unclear, verify on the parent commit before claiming pre-existing.
-3. All failures caused by current changes MUST be fixed — not deferred.
+2. If unclear, verify on the parent commit before claiming it is pre-existing.
+3. Fix every failure caused by current changes; do not defer it.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -340,6 +340,7 @@ tvm_file_glob(GLOB CODEGEN_SRCS
   src/target/canonicalizer/llvm/*.cc
   src/backend/cuda/codegen/*.cc
   src/backend/cuda/op/*.cc
+  src/backend/cuda/transforms/*.cc
   src/backend/hexagon/codegen/*.cc
   src/backend/metal/codegen/*.cc
   src/backend/metal/op/*.cc

diff --git a/docs/tirx/api/backend.rst b/docs/tirx/api/backend.rst
@@ -42,7 +42,7 @@ tvm.backend.cuda.op
 .. automodule:: tvm.backend.cuda.op
    :members:
    :imported-members:
-   :exclude-members: PrimExpr, Op, Call
+   :exclude-members: PrimExpr, Op, Call, PrimType, PointerType
 
 tvm.backend.cuda.script
 ***********************

diff --git a/docs/tirx/layout.rst b/docs/tirx/layout.rst
@@ -324,14 +324,17 @@ scratchpad (partition ``P`` and free ``F`` axes), or NVIDIA Blackwell tensor
 memory with native 2D addressing (``TLane`` × ``TCol``). The demo includes
 presets for each.
 
-SwizzleLayout, ComposeLayout
-----------------------------
+ComposeLayout (swizzled tile)
+-----------------------------
 
 Some layouts also need a *swizzle*: a non-linear, XOR-based permutation of the
 linear memory address. It is not expressible as a strided ``TileLayout`` (which
-is affine), so TIRx represents it as a separate ``SwizzleLayout`` composed with
-the tile layout: ``ComposeLayout(swizzle, tile)``. The tile layout produces a
-linear memory address; the swizzle then permutes that address.
+is affine), so TIRx folds it into a ``ComposeLayout``: alongside a
+``tile_layout``, a ``ComposeLayout`` carries four swizzle parameters
+(``per_element``, ``swizzle_len``, ``atom_len``, ``swizzle_inner``). The tile
+layout produces a linear memory address; the swizzle then permutes that address.
+A *bare* swizzle (no meaningful tile) is a ``ComposeLayout`` over a trivial
+identity ``TileLayout`` covering one swizzle period.
 
 Why swizzle
 ~~~~~~~~~~~
@@ -351,10 +354,11 @@ conflict. Swizzle scatters those accesses across banks.
 The transform
 ~~~~~~~~~~~~~
 
-A ``SwizzleLayout`` has three integer parameters — ``per_element`` (M),
+The swizzle has three integer parameters — ``per_element`` (M),
 ``swizzle_len`` (B), and ``atom_len`` (S) — and maps a linear element address
 ``m`` as follows (keeping the low ``M`` bits untouched and XOR-ing a higher bit
-group down into a lower one):
+group down into a lower one; this is the ``swizzle_inner=True`` direction, and
+``swizzle_inner=False`` mirrors the XOR):
 
 .. math::
 
@@ -380,7 +384,7 @@ mode** (the 32B / 64B / 128B shared-memory swizzle widths):
    S = 3 .
 
 For example ``float16`` (16-bit) gives ``M = bitlen(8) - 1 = 3``; with 128B
-swizzle that is ``Swizzle(M=3, B=3, S=3)``. ``M`` keeps a 16-byte (128-bit)
+swizzle that is swizzle ``(M=3, B=3, S=3)``. ``M`` keeps a 16-byte (128-bit)
 contiguous run unswizzled, matching the minimum vector access.
 
 Bank and line of an element
@@ -400,7 +404,7 @@ swizzled element address ``a = addr(m)`` lands in
 Worked example: 128B swizzle, ``float16``, ``(8, 64)`` tile
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-With ``Swizzle(3,3,3)`` over ``m = 64i + j`` the address simplifies to
+With the swizzle ``(3,3,3)`` over ``m = 64i + j`` the address simplifies to
 
 .. math::