From f3e47d5571fadd0943c60b4621b571b493335be6 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Fri, 31 Jul 2026 01:41:19 -0400 Subject: [PATCH 1/2] feat(tirx): consolidated fork delta over apache main Squash of the fork's development stack. Major pieces: - op-dispatch: dense fp8/tf32 tcgen05 gemm_async paths, per-MMA SMEM descriptor hoist/recompute, TMA planning split/hardening, Layout F sub-slab selection in tcgen05 ld/st - layout: buffer dim-surgery views (unflatten/flatten/select/narrow/ sub/rearrange), SwizzleLayout folded into ComposeLayout, physical offset/rearrange axis-name preservation - lower-tirx: FlashMLA CUDA intrinsics + sparse decode lowering, PrimType dtype handling, unsigned swizzle iter patterns, IKET profiling, dynamic while loops kept unrolled-by-1 in CUDA codegen - tvmscript: PTX cvt instruction forms - infra: Triton-standard bench harness (proton timer, cooldown, timer alignment), TVM_CUDA_NVRTC_EXTRA_OPTS, nvcc arch-suffix and fast-math opt-out fixes, GemmComm distributed benchmarks - tests: MegaKernelMoE scheduler coverage, TMA host-init dtype coverage, remote mbarrier arrive API alignment --- .agents/skills/tir-test/SKILL.md | 224 +- CMakeLists.txt | 1 + docs/tirx/api/backend.rst | 2 +- docs/tirx/layout.rst | 22 +- .../tile_primitives/copy_async/tcgen05_cp.rst | 188 +- docs/tirx/tile_primitives/copy_async/tma.rst | 341 +- docs/tirx/tile_primitives/gemm_async.rst | 44 +- include/tvm/tirx/async_structs.h | 103 - include/tvm/tirx/function.h | 21 +- include/tvm/tirx/layout.h | 96 +- include/tvm/tirx/predicate.h | 66 - include/tvm/tirx/script/builder/ir.h | 2 +- include/tvm/tirx/stmt_functor.h | 2 +- .../tvm/tirx/{tirx_op.h => tile_primitive.h} | 122 +- include/tvm/tirx/tirx_stmt.h | 103 - python/tvm/arith/__init__.py | 10 +- python/tvm/arith/analyzer.py | 23 + python/tvm/backend/cuda/__init__.py | 4 +- python/tvm/backend/cuda/iket.py | 922 +++++ python/tvm/backend/cuda/lang/__init__.py | 2 - python/tvm/backend/cuda/lang/alloc_pool.py | 94 +- python/tvm/backend/cuda/lang/pipeline.py | 40 +- .../tvm/backend/cuda/lang/tile_scheduler.py | 7 +- python/tvm/backend/cuda/op.py | 925 ++++- .../cuda/operator/intrinsics/__init__.py | 3 +- .../cuda/operator/intrinsics/_schema.py | 1 + .../cuda/operator/intrinsics/cp_async.py | 294 +- .../backend/cuda/operator/intrinsics/cvt.py | 754 ++++ .../backend/cuda/operator/intrinsics/math.py | 174 +- .../cuda/operator/intrinsics/memory.py | 347 +- .../backend/cuda/operator/intrinsics/misc.py | 17 + .../backend/cuda/operator/intrinsics/sync.py | 429 ++- .../cuda/operator/intrinsics/tcgen05.py | 90 +- .../cuda/operator/tile_primitive/common.py | 2 +- .../operator/tile_primitive/copy/__init__.py | 7 +- .../operator/tile_primitive/copy/_common.py | 26 +- .../tile_primitive/copy/_swizzle_iter.py | 51 +- .../operator/tile_primitive/copy/fallback.py | 8 +- .../operator/tile_primitive/copy/gmem_smem.py | 4 +- .../tile_primitive/copy/ld_stmatrix.py | 10 +- .../cuda/operator/tile_primitive/copy/reg.py | 21 +- .../operator/tile_primitive/copy/utils.py | 34 +- .../operator/tile_primitive/copy/vec_auto.py | 53 + .../tile_primitive/copy/vec_auto_gmem_smem.py | 290 ++ .../tile_primitive/copy/vec_auto_reg.py | 588 +++ .../tile_primitive/copy/vec_forced.py | 221 ++ .../tile_primitive/copy_async/dsmem.py | 2 +- .../tile_primitive/copy_async/ldgsts.py | 73 +- .../tile_primitive/copy_async/tcgen05_cp.py | 489 ++- .../tile_primitive/copy_async/tcgen05_ldst.py | 232 +- .../operator/tile_primitive/copy_async/tma.py | 3184 +++++++++++------ .../tile_primitive/copy_async/utils.py | 10 +- .../tile_primitive/elementwise/_common.py | 4 +- .../tile_primitive/elementwise/ops/unary.py | 3 +- .../tile_primitive/elementwise/reg.py | 2 +- .../tile_primitive/elementwise/smem.py | 2 +- .../tile_primitive/exec_scope_utils.py | 2 +- .../tile_primitive/gemm/mma_m16n8k_.py | 2 +- .../tile_primitive/gemm_async/tcgen05.py | 659 +++- .../operator/tile_primitive/gemm_utils.py | 2 +- .../operator/tile_primitive/layout_utils.py | 36 +- .../permute_layout/warp_xor_swizzle.py | 2 +- .../tile_primitive/reduction/local.py | 2 +- .../tile_primitive/reduction/shared.py | 2 +- .../tile_primitive/reduction/sm100_packed.py | 2 +- .../tile_primitive/reduction/utils.py | 2 +- .../cuda/operator/tile_primitive/tma_utils.py | 68 +- python/tvm/backend/cuda/script.py | 107 +- .../tvm/backend/cuda/transforms/__init__.py | 41 + .../operator/tile_primitive/binary/default.py | 2 +- .../operator/tile_primitive/copy/default.py | 2 +- .../operator/tile_primitive/gemm/default.py | 2 +- .../operator/tile_primitive/private_alloc.py | 3 +- .../tile_primitive/reduction/utils.py | 2 +- .../operator/tile_primitive/unary/default.py | 2 +- .../tile_primitive/unary/with_bias_scale.py | 2 +- .../trn/transform/private_buffer_alloc.py | 3 +- python/tvm/script/parser/core/entry.py | 7 +- python/tvm/script/parser/core/evaluator.py | 65 +- python/tvm/script/parser/core/parser.py | 6 + python/tvm/support/nvcc.py | 26 +- python/tvm/tirx/__init__.py | 6 +- python/tvm/tirx/_buffer_view.py | 611 ++++ python/tvm/tirx/bench.py | 1191 +++--- python/tvm/tirx/buffer.py | 219 +- python/tvm/tirx/compilation_pipeline.py | 3 + python/tvm/tirx/layout.py | 323 +- python/tvm/tirx/op.py | 6 +- .../tvm/tirx/operator/intrinsics/_common.py | 6 + .../tirx/operator/tile_primitive/__init__.py | 2 +- .../tile_primitive/dispatch_context.py | 209 -- .../operator/tile_primitive/dispatcher.py | 4 +- .../tvm/tirx/operator/tile_primitive/ops.py | 11 +- .../tirx/operator/tile_primitive/registry.py | 3 +- python/tvm/tirx/predicate.py | 45 - python/tvm/tirx/script/__init__.py | 2 +- python/tvm/tirx/script/builder/__init__.py | 2 +- python/tvm/tirx/script/builder/ir.py | 5 +- python/tvm/tirx/script/builder/tirx.py | 152 +- python/tvm/tirx/script/builder/tmem_pool.py | 2 +- python/tvm/tirx/script/parser/__init__.py | 4 +- python/tvm/tirx/script/parser/entry.py | 96 +- python/tvm/tirx/script/parser/parser.py | 104 +- python/tvm/tirx/script/tile.py | 5 +- python/tvm/tirx/stmt.py | 188 +- python/tvm/tirx/tile_primitive.py | 421 +++ src/arith/analyzer.cc | 8 + src/arith/canonical_simplify.cc | 15 +- src/arith/const_fold.h | 27 + src/arith/rewrite_simplify.cc | 88 +- src/backend/cuda/codegen/codegen_cuda.cc | 76 +- src/backend/cuda/codegen/codegen_cuda.h | 1 + src/backend/cuda/codegen/ptx.cc | 17 - src/backend/cuda/codegen/ptx.h | 6 - src/backend/cuda/op/iket.cc | 48 + src/backend/cuda/op/target_builtin.cc | 14 +- src/backend/cuda/runtime/cuda_device_api.cc | 24 + src/backend/cuda/transforms/lower_iket.cc | 1581 ++++++++ .../trn/transform/lower_trainium_layout.cc | 2 +- src/relax/script/printer/call.cc | 6 +- .../extra/contrib/nvshmem/dist_gemm.cu | 52 +- src/runtime/extra/contrib/nvshmem/init.cc | 6 +- .../extra/contrib/nvshmem/memory_allocator.cc | 4 +- .../printer/doc_printer/python_doc_printer.cc | 2 +- src/target/source/codegen_c.cc | 7 +- src/tirx/analysis/verify_tirx_well_formed.cc | 2 +- src/tirx/ir/async_structs.cc | 87 - src/tirx/ir/{predicate.cc => lambdaexpr.cc} | 24 +- src/tirx/ir/layout/compose_layout.cc | 87 +- src/tirx/ir/layout/swizzle_layout.cc | 128 - src/tirx/ir/layout/tile_core.cc | 1 - src/tirx/ir/layout/tile_direct_sum_ops.cc | 82 +- src/tirx/ir/layout/tile_internal.h | 4 + src/tirx/ir/layout/tile_slice.cc | 6 + src/tirx/ir/layout/tile_tile_ops.cc | 214 +- src/tirx/ir/layout/utils.cc | 19 +- src/tirx/ir/layout/utils.h | 16 +- src/tirx/ir/stmt_functor.cc | 20 +- src/tirx/ir/tirx_stmt.cc | 2 +- src/tirx/op/tirx.cc | 4 +- src/tirx/script/builder/ir.cc | 2 +- src/tirx/script/printer/buffer.cc | 31 +- src/tirx/script/printer/expr.cc | 18 +- src/tirx/script/printer/stmt.cc | 81 +- src/tirx/script/printer/utils.h | 11 +- src/tirx/transform/common_subexpr_elim.cc | 36 +- src/tirx/transform/lower_tirx_cleanup.cc | 2 +- src/tirx/transform/split_host_device.cc | 81 +- src/tirx/transform/tile_primitive_dispatch.cc | 16 +- .../arith/test_arith_rewrite_simplify.py | 83 + tests/python/tirx-base/test_tir_op_types.py | 18 +- .../test_tir_transform_common_subexpr_elim.py | 25 + .../test_tir_transform_split_host_device.py | 32 + tests/python/tirx/codegen/conftest.py | 36 +- .../tirx/codegen/test_codegen_blackwell.py | 115 +- .../python/tirx/codegen/test_codegen_cuda.py | 476 ++- .../tirx/codegen/test_codegen_hopper.py | 66 +- tests/python/tirx/codegen/test_ptx_cvt.py | 151 + .../python/tirx/codegen/test_ptx_ld_st_ops.py | 307 +- .../python/tirx/iket/iket_profile_workload.py | 111 + .../oracle/generate_iket_official_oracle.py | 214 ++ .../iket_official_cutlass_4_6_1_oracle.json | 207 ++ .../tirx/iket/test_iket_orchestration.py | 502 +++ tests/python/tirx/iket/test_iket_profiler.py | 670 ++++ .../tirx/iket/verify_iket_official_patch.py | 219 ++ .../tile_primitive/cuda/copy/test_fallback.py | 8 +- .../cuda/copy/test_gmem_smem.py | 32 +- .../cuda/copy/test_ld_stmatrix.py | 14 +- .../tile_primitive/cuda/copy/test_reg.py | 315 +- .../cuda/copy/test_swizzle_iter.py | 49 +- .../cuda/copy/test_vec_forced_cache.py | 172 + .../cuda/copy_async/test_dsmem.py | 2 +- .../cuda/copy_async/test_ldgsts.py | 32 + .../cuda/copy_async/test_smem_tmem.py | 481 --- .../cuda/copy_async/test_tcgen05_cp.py | 1306 +++++++ ...est_tmem_16xnb.py => test_tcgen05_ldst.py} | 412 ++- .../cuda/copy_async/test_tma.py | 3150 ++++++++-------- .../cuda/copy_async/test_tmem.py | 344 -- .../cuda/elementwise/test_unary.py | 6 +- .../cuda/gemm_async/test_gemm_async.py | 1379 ++++++- .../permute_layout/test_permute_layout.py | 14 +- .../tile_primitive/test_dispatcher.py | 2 +- tests/python/tirx/test_alloc_pool.py | 2 +- tests/python/tirx/test_bench_utils.py | 419 ++- tests/python/tirx/test_hint.py | 2 +- tests/python/tirx/test_jit.py | 230 +- tests/python/tirx/test_layout.py | 272 +- tests/python/tirx/test_op.py | 2 +- .../python/tirx/test_op_namespace_cleanup.py | 2 +- tests/python/tirx/test_parser_printer.py | 539 ++- .../tirx/test_printer_tir_namespaces.py | 74 +- .../test_tirx_kernels_registry_correctness.py | 164 + tests/python/tirx/test_verifier.py | 8 +- .../tirx/transform/test_stmt_functor.py | 46 +- .../transform/test_transform_lower_tirx.py | 26 +- tests/scripts/setup-pytest-env.sh | 98 + 196 files changed, 24055 insertions(+), 7195 deletions(-) delete mode 100644 include/tvm/tirx/async_structs.h delete mode 100644 include/tvm/tirx/predicate.h rename include/tvm/tirx/{tirx_op.h => tile_primitive.h} (62%) delete mode 100644 include/tvm/tirx/tirx_stmt.h create mode 100644 python/tvm/backend/cuda/iket.py create mode 100644 python/tvm/backend/cuda/operator/intrinsics/cvt.py create mode 100644 python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py create mode 100644 python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py create mode 100644 python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py create mode 100644 python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py create mode 100644 python/tvm/backend/cuda/transforms/__init__.py create mode 100644 python/tvm/tirx/_buffer_view.py delete mode 100644 python/tvm/tirx/operator/tile_primitive/dispatch_context.py delete mode 100644 python/tvm/tirx/predicate.py create mode 100644 python/tvm/tirx/tile_primitive.py create mode 100644 src/backend/cuda/op/iket.cc create mode 100644 src/backend/cuda/transforms/lower_iket.cc delete mode 100644 src/tirx/ir/async_structs.cc rename src/tirx/ir/{predicate.cc => lambdaexpr.cc} (64%) delete mode 100644 src/tirx/ir/layout/swizzle_layout.cc create mode 100644 tests/python/tirx/codegen/test_ptx_cvt.py create mode 100644 tests/python/tirx/iket/iket_profile_workload.py create mode 100644 tests/python/tirx/iket/oracle/generate_iket_official_oracle.py create mode 100644 tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json create mode 100644 tests/python/tirx/iket/test_iket_orchestration.py create mode 100644 tests/python/tirx/iket/test_iket_profiler.py create mode 100644 tests/python/tirx/iket/verify_iket_official_patch.py create mode 100644 tests/python/tirx/operator/tile_primitive/cuda/copy/test_vec_forced_cache.py delete mode 100644 tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py create mode 100644 tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py rename tests/python/tirx/operator/tile_primitive/cuda/copy_async/{test_tmem_16xnb.py => test_tcgen05_ldst.py} (77%) delete mode 100644 tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem.py create mode 100644 tests/python/tirx/test_tirx_kernels_registry_correctness.py create mode 100755 tests/scripts/setup-pytest-env.sh diff --git a/.agents/skills/tir-test/SKILL.md b/.agents/skills/tir-test/SKILL.md index e4a5fe6ca2d9..fd159087fc50 100644 --- 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: ` 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: ` 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 index 567edc1dc698..9d1dfd9abb3a 100644 --- 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 index 3a3910390199..6fd4e037deef 100644 --- 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 index a71bb9f31923..8e1236006c22 100644 --- 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:: diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst index 8f15aaafdc9e..26d8e39e111d 100644 --- a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst +++ b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst @@ -19,35 +19,74 @@ copy_async → tcgen05_cp ======================= The ``tcgen05_cp`` variant lowers a ``copy_async`` from **shared memory to tensor -memory** (Blackwell ``tmem``). One elected thread issues -``tcgen05.cp.32x128b.warpx4``: a shared **matrix descriptor** names the source tile, -and the ``warpx4`` multicast routes 32 lanes × 128 bits into the tensor-memory lanes -owned by all four warps. The dispatch issues only the copy; the caller signals +memory** (Blackwell ``tmem``) through a generic planner covering every +``tcgen05.cp`` shape. A shared **matrix descriptor** names the source tile; all +descriptor fields (ldo/sdo/swizzle) and the cp issue sequence are derived from +the two buffer layouts. The dispatch issues only the copy; the caller signals completion with ``tcgen05.commit``. Source: ``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py``. +Shape selection +--------------- + +- ``shape=`` config: forces that PTX shape (with ``multicast=`` where the shape + has more than one legal qualifier). +- No ``shape`` config: the planner tries each candidate **widest atom first** + and takes the first whose plan validates against the layouts: + ``128x256b`` → ``4x256b`` → ``128x128b`` → ``64x128b.warpx2::02_13`` → + ``64x128b.warpx2::01_23`` → ``32x128b.warpx4``. All candidates but the + 256b/128b pair are mutually exclusive by the tmem (lane, replica) pattern, so + a bare warpx4 copy still resolves to ``32x128b.warpx4``. + +Each shape pins a tmem row→lane mapping and replica (multicast) pattern +(verified bit-exactly on B200 by ``test_tcgen05_cp.py``): + +.. list-table:: + :header-rows: 1 + :widths: 18 26 30 26 + + * - shape + - multicast + - t lane pattern + - t replica + * - ``128x256b`` + - (none) + - ``(128, 1@TLane)`` + - — + * - ``4x256b`` + - (none) + - ``(4, 32@TLane)`` + - — + * - ``128x128b`` + - (none) + - ``(128, 1@TLane)`` + - — + * - ``64x128b`` + - ``warpx2::02_13`` + - ``(64, 1@TLane)`` + - ``(2, 64@TLane)`` + * - ``64x128b`` + - ``warpx2::01_23`` + - ``(2,32):(64,1)@TLane`` + - ``(2, 32@TLane)`` + * - ``32x128b`` + - ``warpx4`` + - ``(32, 1@TLane)`` + - ``(4, 32@TLane)`` + What it accepts --------------- -Two predicates — a valid shared→tmem copy and a single-thread scope: +Two predicates — the memory-scope envelope and a single-thread exec scope; all +shape/layout validation happens in the planner with readable errors: .. code-block:: python # register_dispatch(..., variant="smem->tmem", priority=10, when=[ - predicate("validate_smem_tmem_copy", _is_valid_smem_tmem_copy), + predicate("validate_smem_tmem_copy", _validate_smem_tmem_copy), predicate("exec_scope", _single_thread_exec), # exec_scope == "thread" # ]) - def _is_valid_smem_tmem_copy(op, sctx): - if not (src.scope().startswith("shared") and dst.scope() == "tmem"): ... - if not (src.layout and dst.layout): ... - if dst.allocated_addr is None: ... - rep = dst.layout.replica # the warpx4 router - if not (len(rep) == 1 and int(rep[0].extent) == 4 - and int(rep[0].stride) == 32 and "TLane" in str(rep[0].axis)): - return False, f"requires R[4:32@TLane] on tmem, got {list(rep)}" - return True, None - .. list-table:: :header-rows: 1 :widths: 22 78 @@ -59,21 +98,22 @@ Two predicates — a valid shared→tmem copy and a single-thread scope: * - scope - **single thread** issues the copy * - memory pair - - source ``shared*`` → destination ``tmem`` (with ``allocated_addr`` set by a - prior ``tcgen05.alloc``) + - source ``shared*`` → destination ``tmem`` (with ``allocated_addr`` set by + a prior ``tcgen05.alloc``); both buffers carry layouts, dtypes match * - tmem layout - - the replica **must be exactly** ``R[4:32@TLane]`` — the warpx4 router that - fans the copy across all four warps' tensor-memory lanes - * - dtype - - sets ``elem_per_128b = 128 / dtype_bits`` (uint8 → 16) and the descriptor - swizzle mode + - must slice to one shape's (lane, replica) pattern from the table above + * - smem layout + - rows in 8-row descriptor core-matrix groups; the atom row width derives + the swizzle mode (K-byte ∈ {16, 32, 64, 128} → sw 0..3) and must match + the buffer's swizzle (if any) Demonstration program ---------------------- -A warpgroup allocates 16 tmem columns, fills a ``32×16`` ``uint8`` shared tile, and -copies it into tmem with ``tcgen05_cp`` (from ``test_smem_tmem.py``; the readback / -dealloc tail is elided): +A warpgroup allocates 16 tmem columns, fills a ``32×16`` ``uint8`` shared tile, +and copies it into tmem — no shape config, the planner infers +``32x128b.warpx4`` from the layouts (from ``test_tcgen05_cp.py``; readback / +dealloc tail elided): .. code-block:: python @@ -94,58 +134,40 @@ dealloc tail is elided): T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) # ... readback via tcgen05.ld, then tcgen05.dealloc ... -(The ``copy_async`` is auto-dispatched to the ``smem->tmem`` variant — the source is -shared and the destination is the ``R[4:32@TLane]`` tmem buffer.) - Algorithm --------- -**1. Verify the warpx4 router and re-order.** After slicing both layouts to the -region, the dispatch confirms the tmem replica is ``R[4:32@TLane]``, permutes to -TLane-first / TCol-stride-descending, isolates the broadcast, and groups the -remaining iters into ``(32, middle, elem_per_128b)`` — the 32×128-bit atom plus a -list of *middle* tiles to loop over. - -**2. Encode the matrix descriptor once.** A 64-bit shared descriptor (leading-dim -offset ``ldo``, stride-dim offset ``sdo``, swizzle mode) is encoded right after the -shared buffer is allocated, cached per ``(smem_buf, ldo, sdo, swizzle)``: +**1. Resolve the shape** — explicit ``shape=`` or layout inference (above). -.. code-block:: python +**2. Validate the plan.** Slice both layouts to the region, permute to +TLane-first / TCol-stride-descending, isolate broadcast iters, and split each +side into ``(lane, middle, col)`` segments against the atom: ``lane`` = one +instruction's rows (must match the shape's lane pattern), ``col`` = one +instruction's columns, ``middle`` = the loop dims — the outer columns and, for +lane-tiled atoms (e.g. 16 ``4x256b`` cps filling each warp's first 16 rows, the +M=64 Layout-F scatter), the extra row iters. The smem side derives the +descriptor fields: 8-row group strides → sdo + swizzle mode, 16B-unit stride → +ldo (256b atoms). - desc_buf = decl_buffer((1,), "uint64", scope="local") - T.ptx.tcgen05.encode_matrix_descriptor(desc_buf.data, s_buf.ptr_to([0, 0]), ldo, sdo, swizzle) +**3. Encode the matrix descriptor once.** A 64-bit shared descriptor is encoded +at the smem buffer base right after its allocation, cached per +``(smem_buf, ldo, sdo, swizzle)``; each cp patches only the 14-bit address +field. -**3. Issue the copy** — one ``tcgen05.cp`` for a single atom, or an unrolled loop -that bumps the tmem column offset and the descriptor's 16-byte shared offset per -middle tile: +**4. Issue the copies** — the middle dims flatten into one unrolled loop; each +step bumps the descriptor's 16-byte shared offset and the tmem address (column +bits, plus the lane half-word for lane-tiled atoms): .. code-block:: python - if total == 1: - T.ptx.tcgen05.cp(t_addr[0] + t_col0, - smem_desc_add_16B_offset(desc_buf[0], init_off_16B), - shape="32x128b", cta_group=cta_group, multicast="warpx4") - else: - for flat in T.unroll(total): - t_off, s_off = T.meta_var(compute_offsets(flat)) - T.ptx.tcgen05.cp(t_addr[0] + t_col0 + t_off, - smem_desc_add_16B_offset(desc_buf[0], init_off_16B + s_off), - shape="32x128b", cta_group=cta_group, multicast="warpx4") - -The dispatch emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits against -an mbarrier (as in the demo). - -Generated TIRx IR ------------------ + for flat in T.unroll(total): + t_off, s_off = T.meta_var(compute_offsets(flat)) + T.ptx.tcgen05.cp(t_addr[0] + t_addr_off + t_off, + smem_desc_add_16B_offset(desc_buf[0], init_off_16B + s_off), + shape=shape, cta_group=cta_group, multicast=multicast) -The ``32×16`` uint8 tile is a single atom (ldo=16, sdo=8, swizzle=0): - -.. code-block:: python - - T.ptx.tcgen05.encode_matrix_descriptor(cp_desc.data, T.address_of(A_smem[0]), 16, 8, 0) - T.ptx.tcgen05.cp(tmem_addr[0], - smem_desc_add_16B_offset(cp_desc[0], 0), - shape="32x128b", cta_group=1, multicast="warpx4") +The dispatch emits **no** ``tcgen05.commit`` / ``wait`` — the caller commits +against an mbarrier (as in the demo). Generated CUDA -------------- @@ -155,8 +177,8 @@ Generated CUDA // one warpx4 copy: shared (named by the matrix descriptor) -> tensor memory "tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" // [%0]=tmem addr, %1=descriptor -(Compiled for ``sm_100a``. End-to-end correctness — including the tmem readback — -is covered by ``test_smem_tmem.py``.) +(Compiled for ``sm_100a``. End-to-end correctness — including the tmem readback +— is covered by ``test_tcgen05_cp.py``.) How inputs change the algorithm ------------------------------- @@ -167,17 +189,19 @@ How inputs change the algorithm * - input - effect + * - tmem layout + - selects the shape under inference via its (lane, replica) pattern; a + ``R[4:32@TLane]`` replica resolves to ``32x128b.warpx4`` * - dtype - - sets ``elem_per_128b = 128 / dtype_bits`` (uint8 → 16, uint32 → 4) and the - descriptor swizzle mode (atom K-byte ∈ {16, 32, 64, 128} → swizzle 0/1/2/3) - * - number of tiles - - ``total`` middle tiles: ``1`` → a single ``tcgen05.cp``; ``> 1`` → an - unrolled loop, each step bumping the tmem column and the descriptor's 16-B - shared offset + - sets ``elem_per_atom = atom_bits / dtype_bits`` and, with the smem row + strides, the descriptor swizzle mode + * - region size + - middle dims: one atom → a single ``tcgen05.cp``; wider regions → an + unrolled loop over columns; more rows than one atom → lane-tiled cps + stepping the tmem lane half-word * - shared swizzle layout - - changes the encoded ``swizzle`` mode (must match the shared buffer's swizzle) - * - tmem layout (D vs F) / cta_group - - the permutation order sets per-tile column steps; ``cta_group`` selects the - multicast routing (``cta_group::1`` vs ``::2``) - * - atom shape - - fixed at ``32x128b`` ``warpx4`` — a different atom would need a new variant + - changes the encoded ``swizzle`` mode (must match the derived atom K-byte) + * - cta_group + - forwarded to the instruction; ``cta_group::2`` is pair-collective — one + even-CTA issue makes each CTA copy from its own smem into its own tmem + (B200-pinned by the cta_group=2 round-trip tests) diff --git a/docs/tirx/tile_primitives/copy_async/tma.rst b/docs/tirx/tile_primitives/copy_async/tma.rst index 5de6ed7f1b7a..596069f67db7 100644 --- a/docs/tirx/tile_primitives/copy_async/tma.rst +++ b/docs/tirx/tile_primitives/copy_async/tma.rst @@ -15,186 +15,213 @@ specific language governing permissions and limitations under the License. -copy_async → tma -================ +copy_async → tma_auto / tma_explicit +===================================== -The ``tma`` variant lowers ``copy_async`` between **global and shared** to the -hardware **Tensor Memory Accelerator**: a single elected thread issues a -descriptor-driven bulk copy (``cp.async.bulk.tensor``), and the hardware walks the -multi-dimensional tile described by a ``cuTensorMap``. The descriptor is built once -on the host (``cuTensorMapEncodeTiled``); the device only *issues* the copy — the -hardware signals the caller's mbarrier when the transfer completes (the dispatch -itself emits no completion op). Source: -``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py``. +The ``tma_auto`` and ``tma_explicit`` variants lower ``copy_async`` between +global and shared memory to CUDA Tensor Memory Accelerator instructions. Both +variants: -What it accepts ---------------- +* are issued by a single elected thread; +* construct and cache ``cuTensorMap`` descriptors on the host with + ``cuTensorMapEncodeTiled``; +* use the same hardware validator, descriptor cache, prefetch path, and PTX + emitter; and +* require the source and destination regions to contain the same total number + of bytes. Their ranks and per-dimension shapes need not match. -The dispatch registers two predicates — a valid copy and a **single-thread** scope: +The variants differ only in how the TensorMap and issue count are planned. -.. code-block:: python - - # register_dispatch(..., priority=10, when=[ - predicate("validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op")), - predicate("single_thread", lambda op, sctx: (single_thread(op, sctx), "expected single thread")), - # ]) +``tma_auto`` +------------ - def single_thread(op_call, sctx): - return sctx.is_thread # exactly one elected thread issues the TMA - -.. list-table:: - :header-rows: 1 - :widths: 22 78 - - * - Property - - Requirement - * - target / priority - - ``cuda``; priority ``10`` (the bulk path for ``copy_async`` global ↔ shared) - * - scope - - **single thread** (``sctx.is_thread``) — TMA is issued by one thread, not a - partitioned warp - * - direction - - ``global → shared`` (g2s) or ``shared → global`` (s2g), inferred from the - buffer scopes at lowering - * - dtype / shape - - ``validate_copy_op``: both sides have layouts, equal dtype, equal non-unit - extents - * - layout - - must form a legal descriptor: rank ≤ 5, innermost stride 1, innermost box - fits the shared swizzle atom (else the plan search shrinks / declines) - -Demonstration program ----------------------- - -One thread bulk-copies an ``8×256`` ``float16`` tile global → shared (with a -128-byte swizzled shared layout), signals an mbarrier, waits, then reads it back -(mirrors ``test_tma.py``'s G2S smoke test): +``tma_auto`` derives the largest legal TMA box from the shared-memory iteration +order. It is intended for ordinary full-tile copies whose layout relationship +can be proven statically: .. code-block:: python - from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout - - g_shape = s_shape = (8, 256); dtype = "float16" - shared_layout = mma_shared_layout(dtype, 3, (8, 256)) # 128-B swizzle - smem_bytes = 8 * 256 * 2 - - @T.prim_func - def copy_async(A_ptr: T.handle, B_ptr: T.handle): - A = T.match_buffer(A_ptr, g_shape, dtype, layout=TileLayout(S[8, 256])) - B = T.match_buffer(B_ptr, g_shape, dtype, layout=TileLayout(S[8, 256])) - T.device_entry(); T.cta_id([1]); tid = T.thread_id([8]) - dyn = T.alloc_buffer([smem_bytes + 8], "uint8", scope="shared.dyn") # arena - A_smem = T.decl_buffer(s_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - phase: T.int32 = 0 - if tid == 0: - T.ptx.mbarrier.init(mbarrier.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta"); T.cuda.cta_sync() - if tid == 0: - Tx.copy_async(A_smem[0:8, 0:256], A[0:8, 0:256], dispatch="tma", mbar=mbarrier.ptr_to([0])) - T.ptx.mbarrier.arrive.expect_tx(mbarrier.ptr_to([0]), smem_bytes) - T.ptx.mbarrier.try_wait(mbarrier.ptr_to([0]), phase) - T.ptx.fence.proxy_async("shared::cta"); T.cuda.cta_sync() - Tx.cta.copy(B[0:8, 0:256], A_smem[0:8, 0:256]) - -Algorithm ---------- - -**1. Infer direction from scopes.** ``global → shared`` is g2s, ``shared → global`` -is s2g (anything else is an error): + Tx.copy_async( + A_smem[:, :], + A[tile_m : tile_m + 64, tile_k : tile_k + 64], + dispatch="tma_auto", + mbar=mbar.ptr_to([0]), + ) + +The shared slice must contain only memory iterators and must form one complete +contiguous stride chain. The planner groups the global layout against that +chain, selects the maximum hardware-legal shared prefix for the TensorMap box, +and emits mixed-radix issue loops for the remaining dimensions. + +Every raw candidate runs the same address-preserving dimension +canonicalization: address-free unit dimensions are removed where legal, and +adjacent contiguous dimensions are merged when the inner dimension is copied +in full from coordinate zero. When an innermost contiguous, fully boxed, +coordinate-zero pair is blocked only because its merged box would exceed 256, +canonicalization may first apply one byte-preserving descriptor-unit promotion +and retry that same boundary, but only if that single promotion makes the merge +legal and a farther outer dimension exists. This avoids skipping outward and +changing the TensorMap tiling of the contiguous inner chain. Promotion is +otherwise a repair step used only when the canonical candidate fails a +repairable hardware rule. It preserves byte strides, payload size, +transaction size, global base, and shared pointer, and is not used for +reductions, TF32, packed dtypes, interleave, OOB fill, gather4, or issue-driven +innermost axes. + +``tma_auto`` does not accept ``gather4`` or ``src_selector`` and only accepts the +default no-OOB contract (``oob=None``). Symbolic facts that affect layout +mapping, prefix selection, or repair must be proven. A dynamic ``globalDim`` +whose range is otherwise unknown is the exception: the runtime TensorMap +encoder checks it is in ``(0, 2^32]`` immediately before the CUDA Driver call. +Use ``tma_explicit`` for explicit OOB behavior or when other descriptor facts +are only known at runtime. + +``tma_explicit`` +---------------- + +``tma_explicit`` maps the supplied global Buffer or view directly: + +* Buffer/view shape becomes ``globalDim``; +* layout strides become byte ``globalStrides``; +* region start becomes the instruction coordinates; +* region extent becomes ``boxDim``; and +* Buffer data plus ``elem_offset`` becomes the TensorMap base. + +It never regroups, compresses, promotes, shrinks, or splits a copy. One +``Tx.copy_async`` call emits exactly one TMA instruction, so a caller must +explicitly tile a wider transfer: .. code-block:: python - if src.scope() == "global" and dst.scope().startswith("shared"): - direction, s_buf, g_buf = "g2s", dst, src - elif src.scope().startswith("shared") and dst.scope() == "global": - direction, s_buf, g_buf = "s2g", src, dst - -**2. Plan the descriptor (L1 → L2 → L3).** The dispatch canonicalizes both -layouts (L1), then for each global iter finds the maximal contiguous stride-1 shard -chain and cuts the axis into descriptor **box** segments (L2), then stacks those -into a ``cuTensorMap`` and validates the hardware constraints — rank ≤ 5, innermost -stride 1, innermost box fits the shared swizzle atom — shrinking the chain prefix -and retrying if a constraint fails (L3). Adjacent fully-boxed contiguous dims are -merged, and an over-256 box may trigger element-type promotion. + for atom in T.unroll(8): + Tx.copy_async( + O[:, atom * 64 : (atom + 1) * 64], + O_smem[:, atom * 64 : (atom + 1) * 64], + dispatch="tma_explicit", + ) -**3. Emit the host descriptor once,** keyed by a cache so a repeated copy reuses it: +The sliced shared layout must canonicalize to a trivial box after its pointer +offset and swizzle are extracted. Global rank and memory-layout rank must +match. A statically illegal value is rejected; symbolic global shapes, +strides, and alignment are retained for validation by the runtime encoder. -.. code-block:: python +Gather4 +~~~~~~~ - T.call_packed("runtime.cuTensorMapEncodeTiled", tensormap, dtype_str, rank, - tensor_ptr, *reversed(shape), *reversed(strides[:-1]), - *reversed(box_dim), *element_strides, 0, swizzle_mode, 2, oob_fill) - -**4. Emit the device issue loop** — an unrolled loop over the issue axes, one -``cp.async.bulk.tensor`` per step, direction-specific: +Gather4 is an explicit global-to-shared operation on SM100 or newer. It +requires a two-dimensional TensorMap and exactly four absolute row +coordinates. Public axis zero is the four-row payload, and PTX receives +coordinates in ``{column, row0, row1, row2, row3}`` order: .. code-block:: python - if direction == "g2s": - T.ptx.cp_async.bulk.tensor.g2c(plan.rank, s_buf.ptr_to(s_st), mbar, - T.address_of(tensor_map), cta_mask, cta_group, - cache_hint, *tma_coords) - else: - T.ptx.cp_async.bulk.tensor.s2g(plan.rank, s_buf.ptr_to(s_st), - T.address_of(tensor_map), cache_hint, *tma_coords) + Tx.copy_async( + K_smem[0:4, :], + K[0:1, :], + dispatch="tma_explicit", + mbar=mbar.ptr_to([0]), + gather4=[row0, row1, row2, row3], + ) -Like all ``copy_async`` variants the dispatch emits no completion — the caller's -mbarrier ``arrive.expect_tx`` / ``try_wait`` (g2s) close the loop. +Longer gathers must be written as multiple four-row calls. Each destination +slice must be four-row box-linear and each source row must have the descriptor's +declared byte stride. -Generated TIRx IR ------------------ +Descriptor selection +~~~~~~~~~~~~~~~~~~~~ -The ``8×256`` swizzled tile produces a **rank-3** descriptor and a single issue: +``src_selector`` selects among alternate global Buffers or views while reusing +the main operand's region and gather coordinates: .. code-block:: python - # host (once): encode the tensor map (rank 3, reversed shape/box/strides, swizzle 3) - T.call_packed("runtime.cuTensorMapEncodeTiled", A_ptr_tensormap, "float16", 3, - A.data, 64, 8, 4, 512, 128, 64, 8, 4, 1, 1, 1, 0, 3, 2, 0) - # device: - for loop_vars in T.unroll(1): - T.ptx.cp_async.bulk.tensor.g2c(3, T.address_of(s_buf_w_offset[0]), - T.address_of(mbarrier[0]), - T.address_of(A_ptr_tensormap), 0, 1, ..., 0, 0, 0) - -Generated CUDA --------------- - -.. code-block:: c++ - - // one TMA instruction copies the whole rank-3 tile, async, into shared - "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes" - ".cta_group::1 [%0], [%1, {%3, %4, %5}], [%2];" - // call: ptx_cp_async_bulk_tensor_g2cluster_tile_3d(smem, mbar, tensormap, coords...) - -The three ``{%3, %4, %5}`` are the descriptor coordinates; ``[%1]`` is the -tensor-map address, ``[%2]`` the mbarrier. One thread launches the entire 8×256 -copy. (This was compiled for ``sm_100a`` — Blackwell — so the instruction carries -the ``.cta_group::1`` qualifier; on Hopper the qualifier is omitted.) - -How inputs change the algorithm -------------------------------- + Tx.copy_async( + K_smem[0:4, :], + K_main[0:1, :], + dispatch="tma_explicit", + mbar=mbar.ptr_to([0]), + gather4=[row0, row1, row2, row3], + src_selector=[ + (use_extra, K_extra.sub[base_row:, :]), + (use_backup, K_backup), + ], + ) + +Conditions use first-true priority and the main Buffer is the default. Every +candidate gets its own validated and encoded TensorMap. Candidates may have +different bases, global shapes, and strides, but must have the same descriptor +dtype, rank, box, swizzle, and transfer byte count. Lowering selects a +pointer-typed descriptor and emits one TMA instruction; it does not select +coordinates or generate instruction-level branches. + +``prefetch_tensormap=True`` deduplicates and prefetches only the main +descriptor. Selector candidates are not prefetched automatically. + +Common configuration +-------------------- .. list-table:: :header-rows: 1 - :widths: 26 74 - - * - input - - effect - * - direction - - ``g2s`` → ``cp.async.bulk.tensor.*.g2c``; ``s2g`` → ``…s2g``; with a reduce - op → ``…s2g_reduce`` (e.g. ``add``) - * - shared swizzle mode - - sets the ``swizzle_mode`` in the descriptor and the innermost-box constraint; - a 128-B swizzle on a 2-D tile yields a **rank-3** descriptor (the inner axis - splits into swizzle atoms), as in the demo - * - box shape / chain prefix - - more selected stride-1 shards → more box>1 descriptor dims; merge collapses - contiguous full-box dims; box > 256 triggers dtype promotion (1→2→4→8 B) - * - dtype - - sets element size and the descriptor's element strides / box byte width - -A copy whose layout cannot form a legal descriptor (rank > 5 after shrinking, or no -swizzle-atom-aligned innermost box) makes the plan search fail and the variant -declines. + :widths: 28 72 + + * - Option + - Meaning + * - ``mbar`` / ``mbarrier_addr`` + - Completion barrier for global-to-shared copies. ``mbarrier_addr`` + selects the PTX shared-address operand form; lowering converts the + supplied generic shared pointer once before the instruction. + * - ``cta_group`` / ``cta_mask`` + - CTA group and multicast mask. ``cta_group`` is one or two. + * - ``cache_hint`` + - Named cache hint or a runtime ``uint64`` cache-policy operand. + * - ``prefetch_tensormap`` + - Deduplicated device-side prefetch of the main descriptor. + * - ``tensormap_l2_promotion`` + - ``none``, ``L2::64B``, ``L2::128B``, or ``L2::256B``. + * - ``tma_dtype`` + - ``tf32`` or ``tfloat32`` for a float32 descriptor conversion. + * - ``use_tma_reduce`` + - Shared-to-global reduction operation. + * - ``oob`` + - Explicit global-to-shared only. ``zero`` is the default; ``nan`` is + limited to non-packed floating-point descriptors. + +Hardware validation +------------------- + +TensorMap dimensions use CUDA API order: dimension zero is innermost and +``globalStrides`` omits its unit stride. Layout order is reversed once when +the descriptor specification is constructed. + +The shared validator checks, among other rules: + +* rank 1 through 5 and target SM90 or newer; +* legal scalar descriptor dtype and conversion dtype; +* 16-byte global base alignment, raised where packed or interleaved modes + require 32 bytes, and 64-byte TensorMap object alignment; +* ``globalDim`` in ``(0, 2^32]``; +* each explicit byte stride non-negative, aligned, and less than ``2^40``; +* ``boxDim`` in ``[1, 256]`` and element stride in ``[1, 8]``; +* unit innermost stride and ordinary inner-box bytes divisible by 16; +* legal interleave, swizzle, L2 promotion, and OOB combinations; +* swizzled inner boxes no larger than their 32-, 64-, or 128-byte atom; and +* direction, reduction, mbarrier, CTA group, coordinate count, and load-mode + combinations. + +Packed sub-byte descriptors are validated in their actual TensorMap encoding +units rather than by truncating ``bits // 8``. + +For ``tma_auto``, a statically disproven rule always rejects the candidate. +Unknown planner-sensitive rules also reject it; only the dynamic +``globalDim`` range is deferred to the runtime encoder. ``tma_explicit`` +retains unknown descriptor values for runtime validation because it does not +choose among alternative boxes or issue loops. + +Completion +---------- + +The dispatch emits the TMA instruction but no completion operation. +Global-to-shared callers initialize the barrier, call +``arrive.expect_tx`` with the transferred byte count, and wait for its phase. +Shared-to-global callers use the bulk-group commit/wait operations required by +their surrounding algorithm. diff --git a/docs/tirx/tile_primitives/gemm_async.rst b/docs/tirx/tile_primitives/gemm_async.rst index c25cdfd57bcf..ef2250379f9a 100644 --- a/docs/tirx/tile_primitives/gemm_async.rst +++ b/docs/tirx/tile_primitives/gemm_async.rst @@ -62,6 +62,13 @@ A single predicate — single-thread or warp scope: 64 (fp4). With cta_group=2, the CTA pair covers twice the per-CTA M * - cta_group - ``1`` (one CTA) or ``2`` (two CTAs split the operand) + * - descriptor mode + - optional ``smem_desc`` controls shared matrix-descriptor construction: + ``"hoist"`` (default), ``"local_hoist"``, ``"encode"``, or + ``"recompute"`` + * - layout forms + - swizzled shared layouts, no-swizzle packed shared layouts, regular tmem + accumulators, and FlashMLA-style packed ``N/2`` tmem accumulator layouts Demonstration program ---------------------- @@ -94,9 +101,18 @@ into a tmem accumulator, after TMA-loading A/B into shared (from Algorithm --------- -**1. Encode shared matrix descriptors.** Each shared operand gets a 64-bit -descriptor (leading-dim offset ``ldo``, stride-dim offset ``sdo``, swizzle mode) -naming its tile to the tensor core: +**1. Encode or synthesize shared matrix descriptors.** Each shared operand is named +by a 64-bit descriptor (leading-dim offset ``ldo``, stride-dim offset ``sdo``, +swizzle mode). ``smem_desc`` selects where that descriptor comes from: + +* ``"hoist"`` (default): encode one uniform descriptor per operand after shared + allocation and add each per-MMA 16-byte offset. +* ``"local_hoist"``: encode at this ``gemm_async`` call site, under the caller's + control flow, then add offsets. This is for call sites where only the elected + issue thread should construct the descriptor. +* ``"encode"``: encode the exact shared pointer for each MMA issue. +* ``"recompute"``: synthesize the descriptor value inline per MMA without a local + descriptor cell. .. code-block:: python @@ -126,6 +142,12 @@ instruction descriptor is encoded at runtime. As with the other async ops, the dispatch emits **no** completion — the caller's ``tcgen05.commit`` + mbarrier wait close it. +For row-0 schedules, the lowering folds ``T.cuda.get_tmem_addr(base, 0, col)`` to +``base + col``. This keeps the generated issue loop close to hand-written +FlashMLA kernels while preserving the helper call for nonzero row offsets. When +``weight_stationary=True`` is passed, the flag is forwarded to the PTX wrapper so +the wrapper can select the matching tcgen05 MMA ABI. + Accumulator datapaths and readback ---------------------------------- @@ -206,4 +228,18 @@ How inputs change the algorithm - set the ``(mi, ni, ki)`` unrolled loop counts; K iterations accumulate into the same tmem accumulator * - shared swizzle - - sets the ``swizzle`` mode + ``ldo``/``sdo`` in the matrix descriptors + - sets the ``swizzle`` mode + ``ldo``/``sdo`` in the matrix descriptors; + no-swizzle packed layouts are accepted when the selected tile has a + hardware-compatible 16-byte packed stride + * - ``smem_desc`` + - selects hoisted, call-site-hoisted, per-MMA encoded, or recomputed shared + descriptor construction. The choice changes code shape only; the MMA + operands still describe the same selected shared tiles. + * - packed tmem accumulator + - layouts of the form ``TileLayout(S[(M, 2, N//2) : (1@TLane, 64@TLane, + 1@TCol)])`` are treated as packed ``N/2`` physical columns, matching + FlashMLA-style low/high accumulator placement. + * - ``weight_stationary`` + - forwarded to the low-level tcgen05 wrapper for kernels that require that + issue ABI; it is a lowering/configuration mode, not a separate PTX + instruction mnemonic. diff --git a/include/tvm/tirx/async_structs.h b/include/tvm/tirx/async_structs.h deleted file mode 100644 index 704ed089fba4..000000000000 --- a/include/tvm/tirx/async_structs.h +++ /dev/null @@ -1,103 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -/*! - * \file tvm/tirx/async_structs.h - * \brief Language structures for asynchronous execution in TIR+. - */ -#ifndef TVM_TIRX_ASYNC_STRUCTS_H_ -#define TVM_TIRX_ASYNC_STRUCTS_H_ - -#include -#include -#include -#include - -namespace tvm { -namespace tirx { - -// Pipeline -class PipelineNode : public ffi::Object { - public: - /*! \brief The thread scope of this pipeline */ - ExecScope thread_scope; - /*! \brief The pipeline depth */ - size_t depth; - /*! \brief Whether to separate producer and consumer threads */ - bool separate_pc; - /*! \brief The name hint of the pipeline. */ - ffi::String name_hint; - - /*! \brief The workspace of the pipeline. */ - ffi::Map workspace; - /*! \brief The schedule config of the pipeline. */ - ffi::Map schedule_config; - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("thread_scope", &PipelineNode::thread_scope) - .def_ro("name_hint", &PipelineNode::name_hint) - .def_ro("depth", &PipelineNode::depth) - .def_ro("separate_pc", &PipelineNode::separate_pc) - .def_ro("workspace", &PipelineNode::workspace) - .def_ro("schedule_config", &PipelineNode::schedule_config); - } - - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO("tirx.Pipeline", PipelineNode, ffi::Object); -}; - -class Pipeline : public ffi::ObjectRef { - public: - TVM_DLL explicit Pipeline(ExecScope thread_scope, size_t depth = 0, bool separate_pc = false, - ffi::String name_hint = "", - ffi::Map workspace = {}, - ffi::Map schedule_config = {}); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Pipeline, ffi::ObjectRef, PipelineNode); -}; - -// CopyPipeline -class CopyPipelineNode : public PipelineNode { - public: - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); - } - - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.CopyPipeline", CopyPipelineNode, PipelineNode); -}; - -class CopyPipeline : public Pipeline { - public: - TVM_DLL explicit CopyPipeline(ExecScope thread_scope, size_t depth = 0, bool separate_pc = false, - ffi::String name_hint = "", - ffi::Map workspace = {}, - ffi::Map schedule_config = {}); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CopyPipeline, Pipeline, CopyPipelineNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(CopyPipelineNode); -}; - -} // namespace tirx -} // namespace tvm - -#endif // TVM_TIRX_ASYNC_STRUCTS_H_ diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h index 3e413d5ca8b4..414e6c40b929 100644 --- a/include/tvm/tirx/function.h +++ b/include/tvm/tirx/function.h @@ -294,7 +294,8 @@ namespace attr { * [arg1, arg2, ..., arg_n, * work_size_1, work_size_2, ... work_size_m, dyn_shmem_size]) * - * Here n = len(arg), m = len(work_size) = len(launch_params)-1. + * Flag-only launch tags do not add packed operands. The dynamic shared-memory + * operand is present only when its value-bearing tag is listed. * * The list of kernel launch params indicates which additional * parameters will be provided to the ffi::Function by the calling @@ -320,13 +321,15 @@ namespace attr { * * - tvm::runtime::launch_param::kUseDynamicSharedMemoryTag * - * The size of the shared memory that may be allocated internally by - * the kernel. For example, exposed as the - * CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES attribute in - * CUDA. + * The dynamic shared-memory byte count passed for this launch. * * Defined as "tirx.use_dyn_shared_memory". * + * - tvm::runtime::launch_param::kUseProgramaticDependentLaunch + * - tvm::runtime::launch_param::kUseCooperativeLaunch + * + * Flag-only launch attributes. These tags add no packed operand. + * * \sa tvm::CallingConv::kDeviceKernelLaunch */ constexpr const char* kKernelLaunchParams = "tirx.kernel_launch_params"; @@ -338,6 +341,14 @@ constexpr const char* kKernelLaunchParams = "tirx.kernel_launch_params"; */ constexpr const char* kLaunchBoundsMinBlocksPerSM = "tirx.launch_bounds_min_blocks_per_sm"; +/*! + * \brief CUDA launch bound maximum CTAs per cluster. + * + * Type: IntImm + */ +constexpr const char* kLaunchBoundsMaxBlocksPerCluster = + "tirx.launch_bounds_max_blocks_per_cluster"; + /*! * \brief Whether to set noalias rule on the function arguments. * diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h index 87c6ca91da3c..ac60b7b4364b 100644 --- a/include/tvm/tirx/layout.h +++ b/include/tvm/tirx/layout.h @@ -408,99 +408,28 @@ class TileLayout : public Layout { TVM_DEFINE_OBJECT_REF_COW_METHOD(TileLayoutNode); }; -// SwizzleLayout -class SwizzleLayoutNode : public LayoutNode { +// ComposeLayout +class ComposeLayoutNode : public LayoutNode { public: int per_element; int swizzle_len; int atom_len; bool swizzle_inner; - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("per_element", &SwizzleLayoutNode::per_element) - .def_ro("swizzle_len", &SwizzleLayoutNode::swizzle_len) - .def_ro("atom_len", &SwizzleLayoutNode::atom_len) - .def_ro("swizzle_inner", &SwizzleLayoutNode::swizzle_inner) - .def_ro("inner_mask", &SwizzleLayoutNode::inner_mask) - .def_ro("outer_mask", &SwizzleLayoutNode::outer_mask); - } - - /*! \brief Check if the layout is compatible with the shape */ - bool CompatibleWithShape(const ffi::Array& shape) const final; - - /*! \brief Verify if the layout is well-formed */ - bool VerifyWellFormed() const final; - - /*! \brief Get the size of the layout */ - PrimExpr GetSize(ffi::Optional axis_name = std::nullopt) const final; - - /*! \brief Get the span of the layout */ - PrimExpr GetSpan(ffi::Optional axis_name = std::nullopt) const final; - - /*! \brief Apply the input coordinate and get the mapped output */ - ffi::Map Apply(ffi::Array coord) const final; - ffi::Map Apply(PrimExpr coord) const final; - - /*! \brief Turn the layout to canonical form */ - Layout Canonicalize() const final; - - /*! \brief Tile the layout with an outer layout */ - Layout Tile(const TileLayout& outer, const ffi::Array& outer_shape, - const ffi::Array& inner_shape) const final; - - Layout DirectSum(const TileLayout& left, const ffi::Array& left_shape, - const ffi::Array& right_shape) const final; - - /*! \brief Check if the layout is the inner layout of a tiled layout */ - ffi::Optional IsTileInner(const Layout& tile_layout, - const ffi::Array& tiled_shape, - const ffi::Array& inner_shape) const final; - - /*! \brief Check if the layout is the outer layout of a tiled layout */ - ffi::Optional IsTileOuter(const Layout& tile_layout, - const ffi::Array& tiled_shape, - const ffi::Array& outer_shape) const final; - - ffi::Optional IsDirectSumRight(const Layout& sum_layout, - const ffi::Array& interleaved_shape, - const ffi::Array& right_shape) const final; - - ffi::Optional IsDirectSumLeft(const Layout& sum_layout, - const ffi::Array& interleaved_shape, - const ffi::Array& left_shape) const final; - - /*! \brief Slice the layout with a given shape and region */ - ffi::Optional Slice(const ffi::Array& shape, const Region& region) const final; - - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SwizzleLayout", SwizzleLayoutNode, LayoutNode); - - private: - friend class SwizzleLayout; + TileLayout tile_layout; + // Cached swizzle masks: inner_mask = (1 << swizzle_len) - 1; + // outer_mask = inner_mask << atom_len. Set by the ComposeLayout ctor. int inner_mask; int outer_mask; -}; - -class SwizzleLayout : public Layout { - public: - TVM_DLL explicit SwizzleLayout(int per_element, int swizzle_len, int atom_len, - bool swizzle_inner); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SwizzleLayout, Layout, SwizzleLayoutNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(SwizzleLayoutNode); -}; - -// ComposeLayout -class ComposeLayoutNode : public LayoutNode { - public: - SwizzleLayout swizzle; - TileLayout tile_layout; static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef() - .def_ro("swizzle", &ComposeLayoutNode::swizzle) + .def_ro("per_element", &ComposeLayoutNode::per_element) + .def_ro("swizzle_len", &ComposeLayoutNode::swizzle_len) + .def_ro("atom_len", &ComposeLayoutNode::atom_len) + .def_ro("swizzle_inner", &ComposeLayoutNode::swizzle_inner) + .def_ro("inner_mask", &ComposeLayoutNode::inner_mask) + .def_ro("outer_mask", &ComposeLayoutNode::outer_mask) .def_ro("tile_layout", &ComposeLayoutNode::tile_layout); } @@ -556,7 +485,8 @@ class ComposeLayoutNode : public LayoutNode { class ComposeLayout : public Layout { public: - TVM_DLL explicit ComposeLayout(SwizzleLayout layout_A, TileLayout layout_B); + TVM_DLL explicit ComposeLayout(int per_element, int swizzle_len, int atom_len, + TileLayout tile_layout, bool swizzle_inner = true); TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ComposeLayout, Layout, ComposeLayoutNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(ComposeLayoutNode); diff --git a/include/tvm/tirx/predicate.h b/include/tvm/tirx/predicate.h deleted file mode 100644 index b7111f69f029..000000000000 --- a/include/tvm/tirx/predicate.h +++ /dev/null @@ -1,66 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - *//*! - * \file tvm/tir/predicate.h - * \brief Definition of predicate - */ - -#ifndef TVM_TIRX_PREDICATE_H_ -#define TVM_TIRX_PREDICATE_H_ - -#include -#include -#include -#include -#include -#include -namespace tvm { -namespace tirx { - -class PredicateNode : public ffi::Object { - public: - /*! \brief The variables in the predicate */ - Array vars; - /*! \brief The predicate */ - PrimExpr pred; - - /*! \brief Replace the variables in the predicate with the given indices */ - PrimExpr Apply(const Array& indices) const; - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("vars", &PredicateNode::vars, refl::AttachFieldFlag::SEqHashDefRecursive()) - .def_ro("pred", &PredicateNode::pred); - } - - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Predicate", PredicateNode, ffi::Object); -}; - -class Predicate : public ffi::ObjectRef { - public: - explicit Predicate(Array vars, PrimExpr pred); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Predicate, ffi::ObjectRef, PredicateNode); -}; - -} // namespace tirx -} // namespace tvm - -#endif // TVM_TIRX_PREDICATE_H_ diff --git a/include/tvm/tirx/script/builder/ir.h b/include/tvm/tirx/script/builder/ir.h index 40a92d194ab7..8346690102c5 100644 --- a/include/tvm/tirx/script/builder/ir.h +++ b/include/tvm/tirx/script/builder/ir.h @@ -27,7 +27,7 @@ #include #include #include -#include +#include namespace tvm { namespace script { diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index f1aca2b94712..5f8562161c32 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -31,7 +31,7 @@ #include #include #include -#include +#include #include #include diff --git a/include/tvm/tirx/tirx_op.h b/include/tvm/tirx/tile_primitive.h similarity index 62% rename from include/tvm/tirx/tirx_op.h rename to include/tvm/tirx/tile_primitive.h index ef485470468e..21f66eec1e41 100644 --- a/include/tvm/tirx/tirx_op.h +++ b/include/tvm/tirx/tile_primitive.h @@ -17,21 +17,63 @@ * under the License. */ /*! - * \file tvm/tirx/tirx_op.h - * \brief TIRX built-in operators. + * \file tvm/tirx/tile_primitive.h + * \brief TIRX tile primitive statements, operators, and reified lambda expressions. */ -#ifndef TVM_TIRX_TIRX_OP_H_ -#define TVM_TIRX_TIRX_OP_H_ +#ifndef TVM_TIRX_TILE_PRIMITIVE_H_ +#define TVM_TIRX_TILE_PRIMITIVE_H_ +#include #include #include #include #include -#include +#include namespace tvm { namespace tirx { +/*! + * \brief A reified Python lambda: a list of bound variables and a body over them. + * + * Used by tile primitive ops that take a per-element expression over the + * destination axes (e.g. ``tirx.tile.select``). ``vars`` are the abstract + * axis variables (lambda-bound); ``pred`` is the body referencing them. + * At lowering time the dispatch substitutes ``vars`` with the concrete + * instruction axes via ``Apply``. + */ +class LambdaExprNode : public ffi::Object { + public: + /*! \brief The bound variables of the lambda. */ + Array vars; + /*! \brief The lambda body over ``vars``. */ + PrimExpr pred; + + /*! \brief Replace the bound variables with the given indices, returning the substituted body. */ + PrimExpr Apply(const Array& indices) const; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("vars", &LambdaExprNode::vars, refl::AttachFieldFlag::SEqHashDefRecursive()) + .def_ro("pred", &LambdaExprNode::pred); + } + + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.LambdaExpr", LambdaExprNode, ffi::Object); +}; + +/*! + * \brief Managed reference to LambdaExprNode. + * \sa LambdaExprNode + */ +class LambdaExpr : public ffi::ObjectRef { + public: + explicit LambdaExpr(Array vars, PrimExpr pred); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LambdaExpr, ffi::ObjectRef, LambdaExprNode); +}; + /*! * \brief The type of the function that sanitizes the arguments of a TIRX operator. * \param op The operator. @@ -140,6 +182,74 @@ class DispatchContext : public ffi::ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DispatchContext, ffi::ObjectRef, DispatchContextNode); }; +/*! + * \brief TIRX TilePrimitiveCall stmt. + */ +class TilePrimitiveCallNode : public StmtNode { + public: + explicit TilePrimitiveCallNode(ffi::UnsafeInit tag) : op(tag) {} + + TilePrimitiveCallNode(tvm::Op op, ffi::Array args, + ffi::Map workspace, + ffi::Map config, ffi::Optional dispatch, + ExecScope scope) + : op(std::move(op)), + args(std::move(args)), + workspace(std::move(workspace)), + config(std::move(config)), + dispatch(std::move(dispatch)), + scope(std::move(scope)) {} + + // tvm::Op which corresponds to the TIRX operator. + tvm::Op op; + + // Arguments to the operator. + ffi::Array args; + + // Workspace (pre-allocated buffers) for the operator. + ffi::Map workspace; + + // Config for the operator/scheduler. + ffi::Map config; + + // Optional dispatch variant name registered via @register_dispatch. + ffi::Optional dispatch{std::nullopt}; + + // Cooperation scope of this call. Default thread (an unscoped call). + ExecScope scope = ExecScope(ScopeKind::kThread); + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef() + .def_ro("op", &TilePrimitiveCallNode::op) + .def_ro("args", &TilePrimitiveCallNode::args) + .def_ro("workspace", &TilePrimitiveCallNode::workspace) + .def_ro("config", &TilePrimitiveCallNode::config) + .def_ro("dispatch", &TilePrimitiveCallNode::dispatch) + .def_ro("scope", &TilePrimitiveCallNode::scope); + } + + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TilePrimitiveCall", TilePrimitiveCallNode, StmtNode); +}; + +/*! + * \brief Managed reference to TilePrimitiveCallNode + * \sa TilePrimitiveCallNode + */ +class TilePrimitiveCall : public Stmt { + public: + TVM_DLL TilePrimitiveCall(tvm::Op op, ffi::Array args, + ffi::Map workspace = {}, + ffi::Map config = {}, + ffi::Optional dispatch = std::nullopt, + ExecScope scope = ExecScope(ScopeKind::kThread)); + + static bool IsValidOpCallArgType(const ffi::Any& arg); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TilePrimitiveCall, Stmt, TilePrimitiveCallNode); + TVM_DEFINE_OBJECT_REF_COW_METHOD(TilePrimitiveCallNode); +}; + /*! * \brief See pesudo code below: * @@ -234,4 +344,4 @@ TVM_DLL const Op& permute_layout(); } // namespace tirx } // namespace tvm -#endif // TVM_TIRX_TIRX_OP_H_ +#endif // TVM_TIRX_TILE_PRIMITIVE_H_ diff --git a/include/tvm/tirx/tirx_stmt.h b/include/tvm/tirx/tirx_stmt.h deleted file mode 100644 index 207e44e0f7e1..000000000000 --- a/include/tvm/tirx/tirx_stmt.h +++ /dev/null @@ -1,103 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -/*! - * \file tvm/tirx/tirx_op.h - * \brief TIRX statements. - */ -#ifndef TVM_TIRX_TIRX_STMT_H_ -#define TVM_TIRX_TIRX_STMT_H_ - -#include -#include - -namespace tvm { -namespace tirx { - -/*! - * \brief TIRX TilePrimitiveCall stmt. - */ -class TilePrimitiveCallNode : public StmtNode { - public: - explicit TilePrimitiveCallNode(ffi::UnsafeInit tag) : op(tag) {} - - TilePrimitiveCallNode(tvm::Op op, ffi::Array args, - ffi::Map workspace, - ffi::Map config, ffi::Optional dispatch, - ExecScope scope) - : op(std::move(op)), - args(std::move(args)), - workspace(std::move(workspace)), - config(std::move(config)), - dispatch(std::move(dispatch)), - scope(std::move(scope)) {} - - // tvm::Op which corresponds to the TIRX operator. - tvm::Op op; - - // Arguments to the operator. - ffi::Array args; - - // Workspace (pre-allocated buffers) for the operator. - ffi::Map workspace; - - // Config for the operator/scheduler. - ffi::Map config; - - // Optional dispatch variant name registered via @register_dispatch. - ffi::Optional dispatch{std::nullopt}; - - // Cooperation scope of this call. Default thread (an unscoped call). - ExecScope scope = ExecScope(ScopeKind::kThread); - - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef() - .def_ro("op", &TilePrimitiveCallNode::op) - .def_ro("args", &TilePrimitiveCallNode::args) - .def_ro("workspace", &TilePrimitiveCallNode::workspace) - .def_ro("config", &TilePrimitiveCallNode::config) - .def_ro("dispatch", &TilePrimitiveCallNode::dispatch) - .def_ro("scope", &TilePrimitiveCallNode::scope); - } - - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TilePrimitiveCall", TilePrimitiveCallNode, StmtNode); -}; - -/*! - * \brief Managed reference to TilePrimitiveCallNode - * \sa TilePrimitiveCallNode - */ -class TilePrimitiveCall : public Stmt { - public: - TVM_DLL TilePrimitiveCall(tvm::Op op, ffi::Array args, - ffi::Map workspace = {}, - ffi::Map config = {}, - ffi::Optional dispatch = std::nullopt, - ExecScope scope = ExecScope(ScopeKind::kThread)); - - static bool IsValidOpCallArgType(const ffi::Any& arg); - - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TilePrimitiveCall, Stmt, TilePrimitiveCallNode); - TVM_DEFINE_OBJECT_REF_COW_METHOD(TilePrimitiveCallNode); -}; - -} // namespace tirx -} // namespace tvm - -#endif // TVM_TIRX_TIRX_STMT_H_ diff --git a/python/tvm/arith/__init__.py b/python/tvm/arith/__init__.py index b4a131cff7d0..c65d5d9805fe 100644 --- a/python/tvm/arith/__init__.py +++ b/python/tvm/arith/__init__.py @@ -25,7 +25,15 @@ estimate_region_strict_bound, estimate_region_upper_bound, ) -from .analyzer import ModularSet, ConstIntBound, Analyzer, ProofStrength, Extension, CompareResult +from .analyzer import ( + ModularSet, + ConstIntBound, + Analyzer, + ProofStrength, + Extension, + CompareResult, + allow_uint_as_index, +) from .bound import deduce_bound from .pattern import detect_linear_equation, detect_clip_bound from .int_solver import solve_linear_equations, solve_linear_inequalities diff --git a/python/tvm/arith/analyzer.py b/python/tvm/arith/analyzer.py index cd4cca83eca9..224244da8421 100644 --- a/python/tvm/arith/analyzer.py +++ b/python/tvm/arith/analyzer.py @@ -18,6 +18,7 @@ """Arithmetic data structure and utility""" import enum +from contextlib import contextmanager import tvm_ffi @@ -527,3 +528,25 @@ def enabled_extensions(self, flags: int | Extension): """ flags = Extension(flags).value _ffi_api.AnalyzerSetEnabledExtensions(self, flags) + + +@contextmanager +def allow_uint_as_index(): + """Opt-in no-overflow domain for unsigned index arithmetic. + + Within the scope, the analyzer treats uint32/uint64 scalars as index + types: it assumes their values never approach the type limit (no + wraparound), so the signed rewrite rule set (floordiv/floormod + decomposition, modular analysis) applies to them. Every unsigned + expression admitted under the assumption is logged once per process + (WARNING) so the assumption stays auditable. + + Off by default. Intended for compile-time layout proofs (offsets that + provably stay small), NOT for runtime codegen semantics where unsigned + wraparound is real behavior. + """ + _ffi_api.EnterAllowUintAsIndex() + try: + yield + finally: + _ffi_api.ExitAllowUintAsIndex() diff --git a/python/tvm/backend/cuda/__init__.py b/python/tvm/backend/cuda/__init__.py index a875f79e04f7..0db092201ca7 100644 --- a/python/tvm/backend/cuda/__init__.py +++ b/python/tvm/backend/cuda/__init__.py @@ -23,7 +23,7 @@ from tvm.base import _LOADED_LIBS -_LAZY_SUBMODULES = {"lang", "op", "operator", "script", "target_tags"} +_LAZY_SUBMODULES = {"iket", "lang", "op", "operator", "script", "target_tags", "transforms"} def _detect_target_from_device(dev): @@ -92,6 +92,7 @@ def __getattr__(name: str): __all__ = [ + "iket", "lang", "op", "operator", @@ -100,4 +101,5 @@ def __getattr__(name: str): "script_namespace", "script_namespaces", "target_tags", + "transforms", ] diff --git a/python/tvm/backend/cuda/iket.py b/python/tvm/backend/cuda/iket.py new file mode 100644 index 000000000000..1ccbfc11001b --- /dev/null +++ b/python/tvm/backend/cuda/iket.py @@ -0,0 +1,922 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""NVIDIA IKET annotations and ``run-iket`` orchestration. + +``profile`` runs an explicit replayable command. ``run`` is intended for a +script's ``__main__`` block: the parent process asks ``run-iket`` to replay the +original script or ``python -m`` invocation, while the injected tracker and +capture processes execute the supplied callable and then exit. +""" + +from __future__ import annotations + +import functools +import hashlib +import importlib.util +import inspect +import json +import math +import os +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import threading +import time +import uuid +import warnings +from collections import deque +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from importlib import metadata +from pathlib import Path +from types import ModuleType +from typing import Any + +import tvm +from tvm.script import tirx as T + +_PROFILE_ENV = "TVM_IKET_OFFICIAL_PROFILE" +_DEFAULT_PROFILE = "cutlass-4.6.1" +_POSTPROCESS_CHOICES = frozenset(("perfetto", "json", "html", "none", "all")) +_INJECTION_ENV_VARS = ("CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG") +_OUTPUT_TAIL_LINES = 100 +_TERMINATION_GRACE_SECONDS = 5.0 + +# Hashes are SHA-256 digests of files in the public 4.6.1 wheels. Only +# ABI-independent runtime/compiler binaries are pinned so this profile works +# with every Python version supported by that CUTLASS DSL release. +_OFFICIAL_PROFILES = { + "cutlass-4.6.1": { + "versions": { + "nvidia-cutlass-dsl": "4.6.1", + "nvidia-cutlass-dsl-libs-base": "4.6.1", + "nvidia-cutlass-dsl-libs-core": "4.6.1", + "nvidia-cutlass-dsl-libs-cu13": "4.6.1", + "nvidia-cuda-nvdisasm": "13.3.73", + "nvidia-cuda-nvrtc": "13.3.33", + }, + "files": { + "nvidia-cutlass-dsl-libs-base": { + "nvidia_cutlass_dsl/dsl_packages/iket/libiket_cubin_info.so": ( + "7ee839130c6bd129b04908a807c066118a459ebea644a59ecb6e41fbb323c103" + ), + "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so": ( + "83be54bd06e2cd82b2f6c17bbee6c925d049acae8d880242d4a5d5509a29e122" + ), + }, + "nvidia-cutlass-dsl-libs-cu13": { + "nvidia_cutlass_dsl/cu13/lib/libcute_dsl_runtime.so": ( + "2fa9809047485ae420ca99cab0678846de692e9608a179b0020834994311dd2f" + ), + }, + "nvidia-cuda-nvdisasm": { + "nvidia/cu13/bin/nvdisasm": ( + "5842e6adf9e232c9503a804915f158a576473e542577c070da3be49390474140" + ), + }, + "nvidia-cuda-nvrtc": { + "nvidia/cu13/lib/libnvrtc.so.13": ( + "e51d197b3b0d2d9d850d29977423e6ac60661d429a59c440fc04e52b6fc6750a" + ), + "nvidia/cu13/lib/libnvrtc-builtins.so.13.3": ( + "7394c640e5761d13d2bbcdbc4b4c5dbac7cb53cd5bc732d78f8a5cb38638e913" + ), + }, + }, + } +} + + +class IketProfileError(RuntimeError): + """An error while validating or running an official IKET profile.""" + + def __init__( + self, + message: str, + *, + returncode: int | None = None, + command: Sequence[str] = (), + output_tail: str = "", + timeout: float | None = None, + ) -> None: + self.returncode = returncode + self.command = tuple(command) + self.output_tail = output_tail + self.timeout = timeout + super().__init__(message) + + +def _requested_formats(postprocess: str) -> frozenset[str]: + if postprocess == "json": + return frozenset(("json",)) + if postprocess == "perfetto": + return frozenset(("perfetto",)) + if postprocess == "html": + return frozenset(("perfetto", "html")) + if postprocess == "all": + return frozenset(("json", "perfetto", "html")) + if postprocess == "none": + return frozenset() + raise ValueError( + f"postprocess must be one of {sorted(_POSTPROCESS_CHOICES)}, got {postprocess!r}" + ) + + +@dataclass(frozen=True) +class IketProfileResult: + """Published artifacts from a successful official IKET profile.""" + + output_dir: Path + postprocess: str + json_traces: tuple[Path, ...] + perfetto_traces: tuple[Path, ...] + html_reports: tuple[Path, ...] + command: tuple[str, ...] = () + + def _one(self, format_name: str, artifacts: tuple[Path, ...], plural_name: str) -> Path: + if format_name not in _requested_formats(self.postprocess): + raise IketProfileError( + f"The {format_name} artifact was not requested (postprocess={self.postprocess!r})" + ) + if not artifacts: + raise IketProfileError( + f"The requested {format_name} artifact is missing from {self.output_dir}" + ) + if len(artifacts) != 1: + raise IketProfileError( + f"The profile produced {len(artifacts)} {format_name} artifacts; " + f"use {plural_name} for multi-process results" + ) + return artifacts[0] + + @property + def trace_json(self) -> Path: + """Return the only JSON trace requested by this profile.""" + return self._one("json", self.json_traces, "json_traces") + + @property + def perfetto(self) -> Path: + """Return the only Perfetto trace requested by this profile.""" + return self._one("perfetto", self.perfetto_traces, "perfetto_traces") + + @property + def html(self) -> Path: + """Return the only HTML report requested by this profile.""" + return self._one("html", self.html_reports, "html_reports") + + @property + def trace(self) -> Any: + """Return the JSON trace without imposing a schema wrapper.""" + return json.loads(self.trace_json.read_text(encoding="utf-8")) + + @property + def launches(self) -> Any: + """Return ``trace[\"launches\"]`` without reshaping it.""" + return self.trace["launches"] + + +def _profile_error(message: str) -> IketProfileError: + return IketProfileError( + f"Official IKET profile validation failed: {message}. " + "Use tvm.tirx.cuda.iket.profile(command) or call " + "tvm.tirx.cuda.iket.run(main) from a replayable entry point." + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file_obj: + for chunk in iter(lambda: file_obj.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_run_iket_entrypoint() -> str: + executable = shutil.which("run-iket") + if executable is None or not os.access(executable, os.X_OK): + raise _profile_error("the run-iket executable is unavailable") + entry_points = metadata.distribution("nvidia-cutlass-dsl-libs-base").entry_points + if not any( + item.group == "console_scripts" + and item.name == "run-iket" + and item.value == "iket.cli.main:entrypoint" + for item in entry_points + ): + raise _profile_error("the run-iket entry point does not match CUTLASS DSL 4.6.1") + return executable + + +def _validate_nvrtc_13_3() -> None: + try: + from cuda.bindings import nvrtc + + error, major, minor = nvrtc.nvrtcVersion() + except (ImportError, OSError, RuntimeError) as err: + raise _profile_error("CUDA NVRTC 13.3 is unavailable") from err + if int(error) != 0 or (int(major), int(minor)) != (13, 3): + raise _profile_error(f"CUDA NVRTC 13.3 is required, got {int(major)}.{int(minor)}") + + +def _validate_official_installation(profile_name: str) -> str: + """Validate pinned host-side packages and return the official executable.""" + if profile_name not in _OFFICIAL_PROFILES: + raise _profile_error(f"unsupported profile {profile_name!r}; expected {_DEFAULT_PROFILE!r}") + profile_config = _OFFICIAL_PROFILES[profile_name] + distributions = {} + for distribution_name, expected_version in profile_config["versions"].items(): + try: + distribution = metadata.distribution(distribution_name) + except metadata.PackageNotFoundError as err: + raise _profile_error( + f"{distribution_name}=={expected_version} is not installed" + ) from err + if distribution.version != expected_version: + raise _profile_error( + f"{distribution_name} must be {expected_version}, got {distribution.version}" + ) + distributions[distribution_name] = distribution + + for distribution_name, expected_files in profile_config["files"].items(): + distribution = distributions[distribution_name] + for relative_path, expected_digest in expected_files.items(): + path = Path(distribution.locate_file(relative_path)) + if not path.is_file(): + raise _profile_error(f"profile binary is missing: {relative_path}") + actual_digest = _sha256(path) + if actual_digest != expected_digest: + raise _profile_error( + f"profile binary hash mismatch for {relative_path}: {actual_digest}" + ) + + executable = _validate_run_iket_entrypoint() + _validate_nvrtc_13_3() + return executable + + +def _validate_injection_environment(expected_injection_digest: str) -> None: + injection_value = os.environ.get("CUDA_INJECTION64_PATH") + injection_path = Path(injection_value) if injection_value else None + if injection_path is None or not injection_path.is_file(): + raise _profile_error("CUDA_INJECTION64_PATH was not supplied by run-iket") + if _sha256(injection_path) != expected_injection_digest: + raise _profile_error("CUDA_INJECTION64_PATH does not match the locked run-iket binary") + + config_value = os.environ.get("SMODEL_INJECTION_CONFIG") + config_path = Path(config_value) if config_value else None + if config_path is None or not config_path.is_file(): + raise _profile_error("SMODEL_INJECTION_CONFIG was not supplied by run-iket") + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, UnicodeDecodeError, json.JSONDecodeError) as err: + raise _profile_error("SMODEL_INJECTION_CONFIG is not valid JSON") from err + if not isinstance(config, dict): + raise _profile_error("SMODEL_INJECTION_CONFIG must contain a JSON object") + tool_name = config.get("toolName") + if tool_name not in ("tracker", "iket"): + raise _profile_error("SMODEL_INJECTION_CONFIG was not generated by run-iket profile") + if tool_name == "tracker": + return + tool_config = config.get("toolConfig") + if not isinstance(tool_config, dict): + raise _profile_error("run-iket capture config is missing toolConfig") + instrument_path = Path(tool_config.get("appInstrument", "")) + if not instrument_path.is_file(): + raise _profile_error("run-iket did not provide an application instrumentation manifest") + + +def _validate_official_environment() -> str: + """Validate the installation plus tracker/capture injection environment.""" + profile_name = os.environ.get(_PROFILE_ENV) + if profile_name not in _OFFICIAL_PROFILES: + raise _profile_error( + f"{_PROFILE_ENV} must be set to {_DEFAULT_PROFILE}, got {profile_name!r}" + ) + executable = _validate_official_installation(profile_name) + profile_config = _OFFICIAL_PROFILES[profile_name] + injection_digest = profile_config["files"]["nvidia-cutlass-dsl-libs-base"][ + "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so" + ] + _validate_injection_environment(injection_digest) + return executable + + +def validate_official_environment() -> None: + """Validate the exact official runtime before an instrumented CUBIN is loaded.""" + _validate_official_environment() + + +class _OfficialIketExecutable: + """Executable whose CUDA image is patched and decoded by NVIDIA run-iket.""" + + def __init__(self, executable) -> None: + self._executable = executable + self._jitted_module = None + self._jit_lock = threading.Lock() + + @property + def mod(self): + """Return the offline module without loading its CUDA image.""" + return self._executable.mod + + def jit(self, *args, **kwargs): + """Validate run-iket before the first load and every forced recompile.""" + force_recompile = bool(kwargs.get("force_recompile", False)) + with self._jit_lock: + if self._jitted_module is None or force_recompile: + _validate_official_environment() + self._jitted_module = self._executable.jit(*args, **kwargs) + return self._jitted_module + + def export_library(self, *_args, **_kwargs): + raise RuntimeError( + "Official IKET executables cannot be exported; run-iket must patch the " + "JIT-loaded CUDA image in the profiling process" + ) + + def __getitem__(self, name): + def invoke(*args, **kwargs): + return self.jit().get_function(name, query_imports=True)(*args, **kwargs) + + return invoke + + def __call__(self, *args, **kwargs): + return self.jit().main(*args, **kwargs) + + +@T.meta_class +class IketProfiler: + """TIRx annotations compiled for NVIDIA's official IKET runtime.""" + + def mark(self, name: str): + T.evaluate(T.cuda.iket.mark(name)) + + def range_start(self, name: str): + return T.cuda.iket.range_start(name) + + def range_end(self, token: tvm.tirx.Expr): + T.evaluate(T.cuda.iket.range_end(token)) + + def range_push(self, name: str): + T.evaluate(T.cuda.iket.range_push(name)) + + def range_pop(self): + T.evaluate(T.cuda.iket.range_pop()) + + def sentinel_token(self, name: str): + return T.cuda.iket.sentinel_token(name) + + def compile(self, mod, target=None, *, tir_pipeline="tirx"): + """Compile official IKET metadata and NativeDump placeholders.""" + if isinstance(mod, tvm.tirx.PrimFunc): + mod = tvm.IRModule.from_expr(mod) + if not isinstance(mod, tvm.IRModule): + raise TypeError("IketProfiler.compile expects a TIRx PrimFunc or IRModule") + enabled_mod = mod.with_attr("tirx.iket.enabled", True) + executable = tvm.compile(enabled_mod, target=target, tir_pipeline=tir_pipeline) + return _OfficialIketExecutable(executable) + + +def _normalize_command(command: Sequence[str | os.PathLike[str]]) -> tuple[str, ...]: + if isinstance(command, str | bytes | os.PathLike): + raise TypeError("command must be a sequence of arguments, not a shell command string") + try: + normalized = tuple(os.fspath(arg) for arg in command) + except TypeError as err: + raise TypeError("every command argument must be a string or path-like object") from err + if not normalized: + raise ValueError("command must contain at least one argument") + if any(not isinstance(arg, str) for arg in normalized): + raise TypeError("every command argument must resolve to a string") + if any("\x00" in arg for arg in normalized): + raise ValueError("command arguments cannot contain NUL bytes") + return normalized + + +def _validate_profile_options( + postprocess: str, timeout: float | None, max_ts_cnt_per_warp: int | None +) -> float | None: + _requested_formats(postprocess) + if timeout is not None: + if isinstance(timeout, bool): + raise TypeError("timeout must be a positive number or None") + timeout = float(timeout) + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("timeout must be positive or None") + if max_ts_cnt_per_warp is not None: + if isinstance(max_ts_cnt_per_warp, bool) or not isinstance(max_ts_cnt_per_warp, int): + raise TypeError("max_ts_cnt_per_warp must be an integer or None") + if max_ts_cnt_per_warp <= 0: + raise ValueError("max_ts_cnt_per_warp must be positive") + return timeout + + +def _child_environment( + profile_name: str, env: Mapping[str, str | os.PathLike[str]] | None +) -> dict[str, str]: + if not isinstance(profile_name, str): + raise TypeError("profile_name must be a string") + child_env = os.environ.copy() + if env is not None: + if not isinstance(env, Mapping): + raise TypeError("env must be a mapping or None") + for key, value in env.items(): + normalized_key = os.fspath(key) + normalized_value = os.fspath(value) + if not isinstance(normalized_key, str) or not isinstance(normalized_value, str): + raise TypeError("env keys and values must resolve to strings") + child_env[normalized_key] = normalized_value + inherited_profile = child_env.get(_PROFILE_ENV) + if inherited_profile is not None and inherited_profile != profile_name: + warnings.warn( + f"Ignoring inherited {_PROFILE_ENV}={inherited_profile!r}; " + f"profile_name={profile_name!r} takes precedence", + RuntimeWarning, + stacklevel=3, + ) + child_env[_PROFILE_ENV] = profile_name + return child_env + + +def _build_run_iket_argv( + executable: str, + command: Sequence[str], + *, + output_dir: Path, + postprocess: str, + max_ts_cnt_per_warp: int | None, + keep: bool, +) -> list[str]: + argv = [ + executable, + "--output-dir", + str(output_dir), + "--clobber", + "profile", + "--postprocess", + postprocess, + "--keep" if keep else "--no-keep", + ] + if max_ts_cnt_per_warp is not None: + argv.extend(("--max-ts-cnt-per-warp", str(max_ts_cnt_per_warp))) + argv.append("--") + argv.extend(command) + return argv + + +def _stream_output(pipe, output_tail: deque[str]) -> None: + try: + while True: + raw_line = pipe.readline() + if raw_line in ("", b""): + break + if isinstance(raw_line, bytes): + line = raw_line.decode("utf-8", errors="replace") + else: + line = raw_line + sys.stdout.write(line) + sys.stdout.flush() + split_lines = line.splitlines() + if not split_lines: + split_lines = [line.rstrip("\r\n")] + output_tail.extend(split_lines) + except (OSError, ValueError): + # The waiting thread closes the pipe to stop a reader after termination. + return + + +def _stop_output_thread(proc: subprocess.Popen, output_thread: threading.Thread | None) -> None: + if output_thread is None: + return + output_thread.join(timeout=1.0) + if output_thread.is_alive() and proc.stdout is not None: + try: + proc.stdout.close() + except OSError: + pass + output_thread.join(timeout=1.0) + + +def _process_group_exists(process_group: int) -> bool: + try: + os.killpg(process_group, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _terminate_process_group(proc: subprocess.Popen) -> None: + process_group = proc.pid + try: + os.killpg(process_group, signal.SIGTERM) + except ProcessLookupError: + pass + + deadline = time.monotonic() + _TERMINATION_GRACE_SECONDS + while _process_group_exists(process_group) and time.monotonic() < deadline: + remaining = deadline - time.monotonic() + try: + proc.wait(timeout=min(0.1, max(remaining, 0.0))) + except subprocess.TimeoutExpired: + continue + if _process_group_exists(process_group): + time.sleep(min(0.05, max(remaining, 0.0))) + + if _process_group_exists(process_group): + try: + os.killpg(process_group, signal.SIGKILL) + except ProcessLookupError: + pass + try: + proc.wait(timeout=_TERMINATION_GRACE_SECONDS) + except subprocess.TimeoutExpired: + # The process has been signalled and will be reaped by the interpreter + # if a broken test double or an unusual platform cannot report it here. + pass + + +def _format_process_error( + message: str, + *, + command: Sequence[str], + output_tail: deque[str], + returncode: int | None, + timeout: float | None, +) -> IketProfileError: + tail_text = "\n".join(output_tail) + detail = f"{message}\nCommand: {shlex.join(command)}" + if tail_text: + detail += f"\nLast {len(output_tail)} output line(s):\n{tail_text}" + return IketProfileError( + detail, + returncode=returncode, + command=command, + output_tail=tail_text, + timeout=timeout, + ) + + +def _run_process( + argv: Sequence[str], + *, + cwd: str | os.PathLike[str] | None, + env: Mapping[str, str], + timeout: float | None, +) -> None: + output_tail: deque[str] = deque(maxlen=_OUTPUT_TAIL_LINES) + try: + proc = subprocess.Popen( + argv, + cwd=cwd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + bufsize=1, + start_new_session=True, + shell=False, + ) + except OSError as err: + raise _format_process_error( + f"Failed to start official IKET profiler: {err}", + command=argv, + output_tail=output_tail, + returncode=None, + timeout=timeout, + ) from err + + output_thread = None + if proc.stdout is not None: + output_thread = threading.Thread( + target=_stream_output, + args=(proc.stdout, output_tail), + name="tvm-iket-output", + daemon=True, + ) + output_thread.start() + + try: + returncode = proc.wait(timeout=timeout) + except subprocess.TimeoutExpired as err: + _terminate_process_group(proc) + _stop_output_thread(proc, output_thread) + raise _format_process_error( + f"Official IKET profiling timed out after {timeout} seconds", + command=argv, + output_tail=output_tail, + returncode=None, + timeout=timeout, + ) from err + except KeyboardInterrupt: + _terminate_process_group(proc) + _stop_output_thread(proc, output_thread) + raise + + _stop_output_thread(proc, output_thread) + if returncode != 0: + raise _format_process_error( + f"Official IKET profiler exited with status {returncode}", + command=argv, + output_tail=output_tail, + returncode=returncode, + timeout=timeout, + ) + + +def _collect_artifacts( + output_dir: Path, +) -> tuple[tuple[Path, ...], tuple[Path, ...], tuple[Path, ...]]: + root_json = {path for path in output_dir.glob("*.json") if path.is_file()} + nested_trace_json = {path for path in output_dir.rglob("*.trace.json") if path.is_file()} + json_traces = tuple(sorted(root_json | nested_trace_json)) + perfetto_traces = tuple( + sorted(path for path in output_dir.rglob("*.pftrace") if path.is_file()) + ) + html_reports = tuple(sorted(path for path in output_dir.rglob("*.html") if path.is_file())) + return json_traces, perfetto_traces, html_reports + + +def _validate_requested_artifacts(output_dir: Path, postprocess: str) -> None: + json_traces, perfetto_traces, html_reports = _collect_artifacts(output_dir) + requested = _requested_formats(postprocess) + missing = [] + if "json" in requested and not json_traces: + missing.append("JSON trace") + if "perfetto" in requested and not perfetto_traces: + missing.append("Perfetto trace") + if "html" in requested and not html_reports: + missing.append("HTML report") + if missing: + raise IketProfileError( + f"Official IKET profiling completed without the requested " + f"{', '.join(missing)} in {output_dir}" + ) + + +def _path_exists(path: Path) -> bool: + return os.path.lexists(path) + + +def _remove_path(path: Path) -> None: + if path.is_symlink() or path.is_file(): + path.unlink(missing_ok=True) + elif path.exists(): + shutil.rmtree(path) + + +def _publish_staging(staging: Path, output_dir: Path, *, clobber: bool) -> None: + if not _path_exists(output_dir): + os.replace(staging, output_dir) + return + if not clobber: + raise FileExistsError(f"IKET output directory already exists: {output_dir}") + + backup = output_dir.with_name(f".{output_dir.name}.backup-{uuid.uuid4().hex}") + os.replace(output_dir, backup) + try: + os.replace(staging, output_dir) + except BaseException: + os.replace(backup, output_dir) + raise + try: + _remove_path(backup) + except OSError: + # The new result is already published and the prior result remains in + # a uniquely named backup rather than risking rollback of valid data. + pass + + +def profile( + command: Sequence[str | os.PathLike[str]], + *, + output_dir: str | os.PathLike[str], + profile_name: str = _DEFAULT_PROFILE, + postprocess: str = "all", + clobber: bool = False, + cwd: str | os.PathLike[str] | None = None, + env: Mapping[str, str | os.PathLike[str]] | None = None, + max_ts_cnt_per_warp: int | None = None, + keep: bool = False, + timeout: float | None = 600.0, +) -> IketProfileResult: + """Profile a replayable command with NVIDIA's official ``run-iket`` tool.""" + if _is_injected_process(): + raise IketProfileError( + "iket.profile() cannot start nested profiling from an injected IKET process" + ) + normalized_command = _normalize_command(command) + timeout = _validate_profile_options(postprocess, timeout, max_ts_cnt_per_warp) + target = Path(output_dir).expanduser().resolve() + if target == target.parent: + raise ValueError("output_dir cannot be a filesystem root") + if _path_exists(target) and not clobber: + raise FileExistsError(f"IKET output directory already exists: {target}") + + child_env = _child_environment(profile_name, env) + executable = _validate_official_installation(profile_name) + target.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=f".{target.name}.staging-", dir=target.parent)) + published = False + try: + argv = _build_run_iket_argv( + executable, + normalized_command, + output_dir=staging, + postprocess=postprocess, + max_ts_cnt_per_warp=max_ts_cnt_per_warp, + keep=keep, + ) + _run_process(argv, cwd=cwd, env=child_env, timeout=timeout) + _validate_requested_artifacts(staging, postprocess) + _publish_staging(staging, target, clobber=clobber) + published = True + json_traces, perfetto_traces, html_reports = _collect_artifacts(target) + return IketProfileResult( + output_dir=target, + postprocess=postprocess, + json_traces=json_traces, + perfetto_traces=perfetto_traces, + html_reports=html_reports, + command=normalized_command, + ) + finally: + if not published and _path_exists(staging): + _remove_path(staging) + + +def _is_injected_process() -> bool: + return any(name in os.environ for name in _INJECTION_ENV_VARS) + + +def _unwrap_partial(callable_obj): + while isinstance(callable_obj, functools.partial): + callable_obj = callable_obj.func + return callable_obj + + +def _replay_error(reason: str) -> ValueError: + return ValueError( + f"iket.run() cannot safely replay this entry point: {reason}. " + "Use iket.profile(command) with an explicit replayable command instead." + ) + + +def _parse_python_entry(orig_argv: Sequence[str]) -> tuple[str, str]: + if not orig_argv: + raise _replay_error("sys.orig_argv is unavailable") + options_with_value = frozenset(("-W", "-X", "--check-hash-based-pycs")) + index = 1 + while index < len(orig_argv): + arg = orig_argv[index] + if arg == "-m": + if index + 1 >= len(orig_argv): + raise _replay_error("python -m has no module name") + return "module", orig_argv[index + 1] + if arg.startswith("-m") and len(arg) > 2: + return "module", arg[2:] + if arg == "-c" or (arg.startswith("-c") and len(arg) > 2): + raise _replay_error("python -c has no replayable source file") + if arg == "-": + raise _replay_error("stdin has no replayable source file") + if arg == "--": + if index + 1 >= len(orig_argv) or orig_argv[index + 1] == "-": + raise _replay_error("stdin has no replayable source file") + return "script", orig_argv[index + 1] + if arg in options_with_value: + index += 2 + continue + if arg.startswith("-W") or arg.startswith("-X"): + index += 1 + continue + if arg.startswith("-"): + index += 1 + continue + return "script", arg + raise _replay_error("sys.orig_argv does not name a script or python -m module") + + +def _resolve_module_entry(module_name: str, main_module: ModuleType) -> Path: + main_spec = getattr(main_module, "__spec__", None) + valid_names = (module_name, f"{module_name}.__main__") + if main_spec is not None and main_spec.name in valid_names and main_spec.origin: + return Path(main_spec.origin).resolve() + try: + module_spec = importlib.util.find_spec(module_name) + if module_spec is not None and module_spec.submodule_search_locations is not None: + module_spec = importlib.util.find_spec(f"{module_name}.__main__") + except (ImportError, AttributeError, ValueError) as err: + raise _replay_error(f"cannot resolve python -m module {module_name!r}") from err + if module_spec is None or not module_spec.origin: + raise _replay_error(f"cannot resolve python -m module {module_name!r}") + return Path(module_spec.origin).resolve() + + +def _replay_command(main) -> tuple[str, ...]: + if not callable(main): + raise TypeError("main must be callable") + base_callable = _unwrap_partial(main) + if getattr(base_callable, "__module__", None) != "__main__": + raise _replay_error("the callable is not defined in __main__") + + main_module = sys.modules.get("__main__") + if main_module is None: + raise _replay_error("the __main__ module is unavailable") + if inspect.isfunction(base_callable) and base_callable.__globals__ is not vars(main_module): + raise _replay_error("the callable is not owned by the active __main__ module") + + try: + callable_source_name = inspect.getsourcefile(base_callable) or inspect.getfile( + base_callable + ) + except (OSError, TypeError) as err: + raise _replay_error("the callable source file cannot be located") from err + main_file_name = getattr(main_module, "__file__", None) + if not callable_source_name or not main_file_name: + raise _replay_error("the callable or __main__ source file cannot be located") + callable_source = Path(callable_source_name).resolve() + main_file = Path(main_file_name).resolve() + if not callable_source.is_file() or not main_file.is_file(): + raise _replay_error("the callable and __main__ must resolve to real files") + + orig_argv = tuple(getattr(sys, "orig_argv", ())) + entry_kind, entry_value = _parse_python_entry(orig_argv) + if entry_kind == "script": + replay_source = Path(entry_value).expanduser().resolve() + else: + replay_source = _resolve_module_entry(entry_value, main_module) + if not replay_source.is_file(): + raise _replay_error("the replay entry point does not resolve to a real file") + if not (callable_source == main_file == replay_source): + raise _replay_error( + "the callable source, __main__.__file__, and replay entry point are different files" + ) + return orig_argv + + +def run( + main, + *, + output_dir: str | os.PathLike[str], + profile_name: str = _DEFAULT_PROFILE, + postprocess: str = "all", + clobber: bool = False, + cwd: str | os.PathLike[str] | None = None, + env: Mapping[str, str | os.PathLike[str]] | None = None, + max_ts_cnt_per_warp: int | None = None, + keep: bool = False, + timeout: float | None = 600.0, +) -> IketProfileResult: + """Replay the active script under ``run-iket`` and execute ``main`` in its passes. + + Module-level code executes once in the parent, tracker, and capture + processes. ``main`` executes only in tracker and capture. Code following + ``run`` executes only in the parent after a successful profile. + """ + if not callable(main): + raise TypeError("main must be callable") + if _is_injected_process(): + _validate_official_environment() + main() + raise SystemExit(0) + + command = _replay_command(main) + return profile( + command, + output_dir=output_dir, + profile_name=profile_name, + postprocess=postprocess, + clobber=clobber, + cwd=cwd, + env=env, + max_ts_cnt_per_warp=max_ts_cnt_per_warp, + keep=keep, + timeout=timeout, + ) + + +__all__ = [ + "IketProfileError", + "IketProfileResult", + "IketProfiler", + "profile", + "run", +] diff --git a/python/tvm/backend/cuda/lang/__init__.py b/python/tvm/backend/cuda/lang/__init__.py index a70ce50196c6..eee5c6ba9851 100644 --- a/python/tvm/backend/cuda/lang/__init__.py +++ b/python/tvm/backend/cuda/lang/__init__.py @@ -34,7 +34,6 @@ "TCGen05Bar", "TMABar", "TMEMPool", - "TMEMStages", "WarpRole", "WarpgroupRole", ] @@ -54,7 +53,6 @@ "TCGen05Bar": ".pipeline", "TMABar": ".pipeline", "TMEMPool": ".alloc_pool", - "TMEMStages": ".alloc_pool", "WarpRole": ".warp_role", "WarpgroupRole": ".warp_role", "SmemDescriptor": ".smem_desc", diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py index 0dab5341d42c..75cf3034bbb2 100644 --- a/python/tvm/backend/cuda/lang/alloc_pool.py +++ b/python/tvm/backend/cuda/lang/alloc_pool.py @@ -125,7 +125,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): if len(shape) < 2: raise ValueError( - f"alloc_mma shape={tuple(shape)} has fewer than 2 dimensions; " + f"alloc_tcgen05_mma_AB shape={tuple(shape)} has fewer than 2 dimensions; " f"swizzled MMA layouts tile over the last two dims (rows, cols). " f"Use swizzle_mode='none' for 1-D allocations." ) @@ -148,7 +148,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): atom_bytes = _swizzle_atom_bytes(swizzle_mode) suggestion = _suggest_swizzle_for_row_bytes(col_bits // 8 if col_bits >= 8 else 0) raise ValueError( - f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces " + f"alloc_tcgen05_mma_AB shape={tuple(shape)} with dtype={dtype!r} produces " f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B " f"swizzle atom selected by {swizzle_mode.name}. " f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] " @@ -162,7 +162,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): suggestion = _suggest_swizzle_for_row_bytes(row_bytes) min_cols = atom_bytes // dtype_bytes raise ValueError( - f"alloc_mma shape={tuple(shape)} with dtype={dtype!r} produces " + f"alloc_tcgen05_mma_AB shape={tuple(shape)} with dtype={dtype!r} produces " f"{row_bytes}B rows, which is incompatible with the {atom_bytes}B " f"swizzle atom selected by {swizzle_mode.name}. " f"Use swizzle_mode=SwizzleMode.{suggestion}, or widen shape[-1] " @@ -173,65 +173,18 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode): atom_rows = 8 if rows < atom_rows or rows % atom_rows != 0: raise ValueError( - f"alloc_mma shape={tuple(shape)} has shape[-2]={rows}, but the " + f"alloc_tcgen05_mma_AB shape={tuple(shape)} has shape[-2]={rows}, but the " f"{swizzle_mode.name} atom requires shape[-2] to be a positive " f"multiple of {atom_rows}. Use swizzle_mode='none', or widen shape[-2] " f"to a multiple of {atom_rows}." ) -# --------------------------------------------------------------------------- -# TMEMStages -# --------------------------------------------------------------------------- - - def _meta_class(cls): """Apply @meta_class decorator from ir_builder.""" return _get_ir().meta_class(cls) -@_meta_class -class TMEMStages: - """Parse-time staged view over a TMEM buffer. - - Parameters - ---------- - buf : Buffer - The underlying TMEM buffer (e.g. f32 or f16 view). - col_start : int - First column of stage 0 in *buf*'s column space. - width : int - Number of columns per stage. - stages : int - Number of pipeline stages (default 1). - stride : int or None - Column distance between consecutive stages. When *None* (default), - equals *width* (stages are packed back-to-back). - """ - - def __init__(self, buf, col_start, width, stages=1, stride=None): - self.buf = buf - self.col_start = col_start - self.width = width - self.stages = stages - self.stride = width if stride is None else stride - - def _stage_base(self, stage): - return self.col_start + stage * self.stride - - def __getitem__(self, item): - if isinstance(item, tuple): - assert len(item) == 2, "TMEMStages expects region[stage] or region[stage, start:stop]" - stage, col_slice = item - assert isinstance(col_slice, slice), "TMEMStages tuple indexing requires a slice" - base = self._stage_base(stage) - start = 0 if col_slice.start is None else col_slice.start - stop = self.width if col_slice.stop is None else col_slice.stop - return self.buf[:, base + start : base + stop : col_slice.step] - base = self._stage_base(item) - return self.buf[:, base : base + self.width] - - # --------------------------------------------------------------------------- # TMEMPool # --------------------------------------------------------------------------- @@ -312,7 +265,7 @@ def _resolve_cols(self, shape, dtype, cols, layout=None): ) return total_bits // (32 * rows) - def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None): + def alloc(self, shape, dtype="float32", *, layout=None, cols=None): """Allocate a TMEM buffer. Parameters @@ -334,14 +287,6 @@ def alloc(self, shape, dtype="float32", *, layout=None, cols=None, datapath=None mapping — keep this for shape ``(128, X)`` buffers that hold an M=128 MMA accumulator. """ - from tvm.tirx.layout import tmem_datapath_layout - - if layout is not None and datapath is not None: - raise ValueError("TMEMPool.alloc: pass at most one of layout= and datapath=") - if datapath is not None: - assert len(shape) == 2, "TMEMPool.alloc: datapath= requires a 2-D shape" - layout = tmem_datapath_layout(datapath, shape[0], shape[1]) - ir = _get_ir() cols = self._resolve_cols(shape, dtype, cols, layout) col_start = self.offset @@ -382,6 +327,33 @@ def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1): ) return self.alloc(shape, dtype, layout=layout) + def alloc_tcgen05_mma_A( + self, shape, dtype="bfloat16", *, M, cta_group, ws=False, sparse=False, cols=None + ): + """Allocate a TMEM A operand (A-in-TMEM); layout resolved from MMA + instruction params. ``M`` is the PTX tcgen05.mma instruction M (256/128/ + 64), NOT per-CTA rows. See ``localdoc/claude_plan.txt`` S3a.""" + from tvm.tirx.layout import tmem_mma_operand_layout + + layout = tmem_mma_operand_layout( + "A", shape, dtype, M=M, cta_group=cta_group, ws=ws, sparse=sparse + ) + return self.alloc(shape, dtype, layout=layout, cols=cols) + + def alloc_tcgen05_mma_D( + self, shape, dtype="float32", *, M, cta_group, ws=False, sparse=False, group=None, cols=None + ): + """Allocate a TMEM D/C accumulator operand. ``M`` = PTX instruction M. + ``group=(s,2,n)`` gives a flat (m,N) buffer a grouped column layout.""" + from tvm.tirx.layout import tmem_mma_operand_layout + + layout = tmem_mma_operand_layout( + "D", shape, dtype, M=M, cta_group=cta_group, ws=ws, sparse=sparse, group=group + ) + return self.alloc(shape, dtype, layout=layout, cols=cols) + + alloc_tcgen05_mma_C = alloc_tcgen05_mma_D + def move_base_to(self, col): self.offset = col self.max_offset = max(self.max_offset, self.offset) @@ -479,7 +451,7 @@ def alloc( self.max_offset = max(self.max_offset, self.offset) return res - def alloc_mma(self, shape, dtype="float16", swizzle_mode="auto", align=1024): + def alloc_tcgen05_mma_AB(self, shape, dtype="float16", swizzle_mode="auto", align=1024): """Allocate MMA-compatible shared memory with an inferred swizzle layout.""" from tvm.backend.cuda.operator.tile_primitive.tma_utils import ( SwizzleMode, diff --git a/python/tvm/backend/cuda/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py index 040bdf056f4e..9e4246d5b1a0 100644 --- a/python/tvm/backend/cuda/lang/pipeline.py +++ b/python/tvm/backend/cuda/lang/pipeline.py @@ -121,22 +121,22 @@ def _wait(self, stage, phase): # so this returns only once the barrier has completed. T.ptx.mbarrier.try_wait(self.buf.ptr_to([stage]), phase ^ self.phase_offset) - def arrive(self, stage, cta_id=None, pred=None, count=None): + def arrive(self, stage, remote=None, pred=None, count=None): if self._remote_cta_id is not None: - if cta_id is not None: - raise ValueError("MBarrier.remote_view().arrive() cannot also specify cta_id") - cta_id = self._remote_cta_id + if remote is not None: + raise ValueError("MBarrier.remote_view().arrive() cannot also specify remote") + remote = self._remote_cta_id buf = self._local_buf else: buf = self.buf - self._arrive(buf.ptr_to([stage]), cta_id, pred, count) + self._arrive(buf.ptr_to([stage]), remote, pred, count) @T.inline - def _arrive(self, bar, cta_id=None, pred=None, count=None): + def _arrive(self, bar, remote=None, pred=None, count=None): # Default: local-CTA arrive — emits the simple # ``mbarrier.arrive.shared.b64`` form. To arrive on a remote # CTA's mbarrier in a cluster kernel, callers must pass - # ``cta_id=`` explicitly (e.g. ``bar.arrive(stage, cta_id=0)``) + # ``remote=`` explicitly (e.g. ``bar.arrive(stage, remote=0)``) # or use ``MBarrier.remote_view(rank).arrive(stage)``. Defaulting # the cross-CTA path was both surprising (``bar.arrive(stage)`` # silently ``mapa`` ed across the cluster) and a per-call cost @@ -146,11 +146,11 @@ def _arrive(self, bar, cta_id=None, pred=None, count=None): # operand, i.e. ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``. # When ``None`` the implicit count-of-1 form is emitted. Passing # ``count=1`` is semantically identical but spells the count explicitly. - if cta_id is None: + if remote is None: T.ptx.mbarrier.arrive(bar) else: actual_pred = True if pred is None else pred - T.ptx.mbarrier.arrive(bar, cta_id=cta_id, pred=actual_pred, count=count) + T.ptx.mbarrier.arrive(bar, remote=remote, pred=actual_pred, count=count) def ptr_to(self, idx): return self.buf.ptr_to(idx) @@ -190,36 +190,36 @@ class TMABar(MBarrier): (matching MBarrier.arrive defaults). """ - def arrive(self, stage, tx_count=None, cta_id=None, pred=None): + def arrive(self, stage, tx_count=None, remote=None, pred=None): if self._remote_cta_id is not None: - if cta_id is not None: - raise ValueError("TMABar.remote_view().arrive() cannot also specify cta_id") - cta_id = self._remote_cta_id + if remote is not None: + raise ValueError("TMABar.remote_view().arrive() cannot also specify remote") + remote = self._remote_cta_id buf = self._local_buf else: buf = self.buf - self._arrive_tma(buf.ptr_to([stage]), tx_count, cta_id, pred) + self._arrive_tma(buf.ptr_to([stage]), tx_count, remote, pred) @T.inline - def _arrive_tma(self, bar, tx_count=None, cta_id=None, pred=None): + def _arrive_tma(self, bar, tx_count=None, remote=None, pred=None): # NOTE: this arrive() kwarg set intentionally differs from # MBarrier.arrive (hardware necessity, LSP-incompatible by design). # ``tx_count``: TMA byte count for ``mbarrier.arrive.expect_tx``. - # ``cta_id`` / ``pred``: forwarded to the underlying + # ``remote`` / ``pred``: forwarded to the underlying # ``mbarrier.arrive`` (cluster path) when set; otherwise the # arrive is local-CTA only. See ``MBarrier.arrive`` for the # full default-local rationale. if tx_count is not None: - if cta_id is None: + if remote is None: T.ptx.mbarrier.arrive.expect_tx(bar, tx_count) else: actual_pred = True if pred is None else pred - T.ptx.mbarrier.arrive.expect_tx(bar, tx_count, cta_id=cta_id, pred=actual_pred) - elif cta_id is None: + T.ptx.mbarrier.arrive.expect_tx(bar, tx_count, remote=remote, pred=actual_pred) + elif remote is None: T.ptx.mbarrier.arrive(bar) else: actual_pred = True if pred is None else pred - T.ptx.mbarrier.arrive(bar, cta_id=cta_id, pred=actual_pred) + T.ptx.mbarrier.arrive(bar, remote=remote, pred=actual_pred) class TCGen05Bar(MBarrier): diff --git a/python/tvm/backend/cuda/lang/tile_scheduler.py b/python/tvm/backend/cuda/lang/tile_scheduler.py index 3ca8f8269ea7..199c4cfef35b 100644 --- a/python/tvm/backend/cuda/lang/tile_scheduler.py +++ b/python/tvm/backend/cuda/lang/tile_scheduler.py @@ -370,7 +370,6 @@ def _update_serpentine(self, tile_linear, set_tile_coords): BLOCK_SIZE = T.meta_var(self._BLOCK_SIZE) # S * S FULL_BLOCK_TILES = T.meta_var(self._FULL_BLOCK_TILES) M_TILE_ROWS = T.meta_var(self._M_TILE_ROWS) - T.meta_var(self._N_TILE_COLS) RESIDUAL_N = T.meta_var(self._RESIDUAL_N) RESIDUAL_M = T.meta_var(self._RESIDUAL_M) @@ -872,7 +871,7 @@ def consume(self): self._clc.sched_arr.full.wait(0, self._sa.phase) self._sa.advance() self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0])) - self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True) + self._clc.sched_fin.empty.arrive(0, remote=0, pred=True) @T.inline def consume_wg(self, wg_id, warp_id, lane_id): @@ -882,7 +881,7 @@ def consume_wg(self, wg_id, warp_id, lane_id): self._nxt = T.ptx.clc_query_cancel(T.address_of(self._clc.clc_handle[0])) T.cuda.warpgroup_sync(wg_id + 1) if (warp_id == 0) & (lane_id == 0): - self._clc.sched_fin.empty.arrive(0, cta_id=0, pred=True) + self._clc.sched_fin.empty.arrive(0, remote=0, pred=True) @T.inline def advance_coords(self): @@ -940,6 +939,6 @@ def run_scheduler(self, cbx): self.sched_arr.full.wait(0, sa.phase) sa.advance() self._s_nxt = T.ptx.clc_query_cancel(T.address_of(self.clc_handle[0])) - self.sched_fin.empty.arrive(0, cta_id=0, pred=True) + self.sched_fin.empty.arrive(0, remote=0, pred=True) if self._s_nxt == 0xFFFFFFFF: self._s_done = 1 diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py index 988e48f4d01d..7ab78c102bb0 100644 --- a/python/tvm/backend/cuda/op.py +++ b/python/tvm/backend/cuda/op.py @@ -21,8 +21,9 @@ from tvm import tirx from tvm.ir import Call, Op, is_prim_expr +from tvm.ir.type import PointerType, PrimType from tvm.runtime import const -from tvm.tirx.op import bitwise_and, call_intrin, tvm_access_ptr +from tvm.tirx.op import bitwise_and, call_intrin, reinterpret, tvm_access_ptr from tvm.tirx.operator.intrinsics._common import CLUSTER_BARRIER_SEM as _CLUSTER_BARRIER_SEM from tvm.tirx.operator.intrinsics._common import ( CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT, @@ -37,6 +38,18 @@ from tvm.tirx.operator.intrinsics._common import FENCE_SEM as _FENCE_SEM from tvm.tirx.operator.intrinsics._common import LDMATRIX_DTYPE as _LDMATRIX_DTYPE from tvm.tirx.operator.intrinsics._common import LDMATRIX_NUM as _LDMATRIX_NUM +from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SCOPE as _MBARRIER_ARRIVE_SCOPE +from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SEM as _MBARRIER_ARRIVE_SEM +from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SPACE as _MBARRIER_ARRIVE_SPACE +from tvm.tirx.operator.intrinsics._common import ( + MBARRIER_COMPLETE_TX_SCOPE as _MBARRIER_COMPLETE_TX_SCOPE, +) +from tvm.tirx.operator.intrinsics._common import ( + MBARRIER_COMPLETE_TX_SEM as _MBARRIER_COMPLETE_TX_SEM, +) +from tvm.tirx.operator.intrinsics._common import ( + MBARRIER_COMPLETE_TX_SPACE as _MBARRIER_COMPLETE_TX_SPACE, +) from tvm.tirx.operator.intrinsics._common import NVSHMEM_CMP as _NVSHMEM_CMP from tvm.tirx.operator.intrinsics._common import NVSHMEM_SIG_OP as _NVSHMEM_SIG_OP from tvm.tirx.operator.intrinsics._common import TCGEN05_CP_DECOMPRESS as _TCGEN05_CP_DECOMPRESS @@ -52,6 +65,41 @@ ######################################################## +def cuda_iket_mark(name): + """Create an NVIDIA IKET marker annotation.""" + return call_intrin("", "tirx.cuda.iket_mark", name) + + +def cuda_iket_range_start(name): + """Create an NVIDIA IKET token-range start annotation.""" + return call_intrin("uint32", "tirx.cuda.iket_range_start", name) + + +def cuda_iket_range_end(token): + """Create an NVIDIA IKET token-range end annotation.""" + return call_intrin("", "tirx.cuda.iket_range_end", token) + + +def cuda_iket_range_push(name): + """Create an NVIDIA IKET stack-range push annotation.""" + return call_intrin("", "tirx.cuda.iket_range_push", name) + + +def cuda_iket_range_pop(): + """Create an NVIDIA IKET stack-range pop annotation.""" + return call_intrin("", "tirx.cuda.iket_range_pop") + + +def cuda_iket_sentinel_token(name): + """Create a no-op NVIDIA IKET range token for warp-uniform control flow.""" + return call_intrin("uint32", "tirx.cuda.iket_sentinel_token", name) + + +def cuda_iket_official_event(event_id, source_code=""): + """Create an NVIDIA IKET official range-end event.""" + return call_intrin("uint32", "tirx.cuda.iket_official_event", event_id, source_code) + + def cuda_func_call(func_name, *args, source_code, return_type="void"): """TVM intrinsic to call a CUDA function. Source code is provided as a string. @@ -523,21 +571,41 @@ def ptx_cp_async_bulk_shared_to_cluster(dst_ptr, src_ptr, size, mbar): return call_intrin("", "tirx.ptx.cp_async_bulk_shared_to_cluster", dst_ptr, src_ptr, size, mbar) -def ptx_cp_async_mbarrier_arrive(barrier_id): +def ptx_cp_async_mbarrier_arrive(bar, noinc=False, space="shared"): """TVM intrinsic for ptx async copy barrier using cp.async.mbarrier.arrive https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-cp-async-mbarrier-arrive Parameters ---------- - barrier_id : int - The ID of the barrier shared memory pointer. + bar : Expr + Pointer to the shared-memory mbarrier object. + + noinc : bool + Whether to emit the ``.noinc`` modifier. + + space : str + Address-space modifier, either ``"shared"`` or ``"shared::cta"``. Returns ------- call : Expr The call expression. """ - return call_intrin("", "tirx.ptx.cp_async_mbarrier_arrive", barrier_id) + return call_intrin("", "tirx.ptx.cp_async_mbarrier_arrive", bar, noinc, space) + + +def ptx_cp_async_mbarrier_arrive_noinc(bar, space="shared::cta"): + """TVM intrinsic for ``cp.async.mbarrier.arrive.noinc.shared::cta.b64``. + + Parameters + ---------- + bar : Expr + Pointer to the shared-memory mbarrier object. + + space : str + Address-space modifier, either ``"shared"`` or ``"shared::cta"``. + """ + return ptx_cp_async_mbarrier_arrive(bar, True, space) def ptx_fence(sem: str, scope: str): @@ -601,82 +669,170 @@ def ptx_mbarrier_init(bar, thread_count): return call_intrin("", "tirx.ptx.mbarrier_init", bar, thread_count) -def ptx_mbarrier_arrive(bar, cta_id=None, pred=None, count=None): - """TVM intrinsic to call - mbarrier.arrive.shared::cta.b64 - or - @p mapa.shared::cluster.u32 - @p mbarrier.arrive.shared::cluster.b64 [, count] +def _validate_mbarrier_arrive_attrs(sem, scope, space, remote): + if (sem == "") != (scope == ""): + raise ValueError("mbarrier.arrive sem and scope must be specified together") + if sem not in _MBARRIER_ARRIVE_SEM: + raise ValueError(f"invalid sem={sem!r}; expected one of {_MBARRIER_ARRIVE_SEM}") + if scope not in _MBARRIER_ARRIVE_SCOPE: + raise ValueError(f"invalid scope={scope!r}; expected one of {_MBARRIER_ARRIVE_SCOPE}") + if space not in _MBARRIER_ARRIVE_SPACE: + raise ValueError(f"invalid space={space!r}; expected one of {_MBARRIER_ARRIVE_SPACE}") + if remote is not None and space != "shared::cluster": + raise ValueError("remote mbarrier.arrive requires space='shared::cluster'") - Parameters - ---------- - bar : Var - The pointer to barrier variable. - - cta_id : Optional[Expr] - The cta id. - pred : Optional[Expr] - The predicate to guard the operation. +def ptx_mbarrier_arrive( + bar, + *, + count=None, + sem="", + scope="", + space=None, + remote=None, + pred=None, + **kwargs, +): + """PTX ``mbarrier.arrive{.sem.scope}{.space}.b64 _, [bar]{, count}``. - count : Optional[Expr] - Explicit arrival count operand for the cross-CTA (cluster) form. When - ``None`` the implicit count-of-1 form is emitted; when given, emits - ``mbarrier.arrive.shared::cluster.b64 _, [addr], count``. + ``remote`` maps ``bar`` with ``mapa.shared::cluster.u32`` before the arrive + instruction. ``pred`` predicates the map and arrive instruction. """ - if cta_id is None and pred is None: - return call_intrin("", "tirx.ptx.mbarrier_arrive", bar) - assert cta_id is not None and pred is not None - if count is None: - return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, pred) - return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, pred, count) + if "cta_id" in kwargs: + raise ValueError("T.ptx.mbarrier.arrive uses remote= instead of cta_id=") + if kwargs: + raise TypeError(f"unexpected keyword argument(s): {', '.join(kwargs)}") + + if space is not None: + effective_space = space + elif remote is not None: + effective_space = "shared::cluster" + else: + effective_space = "shared" + _validate_mbarrier_arrive_attrs(sem, scope, effective_space, remote) + + args = [bar] + if count is not None: + args.append(count) + if remote is not None: + args.append(remote) + if pred is not None: + args.append(pred) + return call_intrin( + "", + "tirx.ptx.mbarrier_arrive", + *args, + sem, + scope, + effective_space, + int(count is not None), + int(remote is not None), + int(pred is not None), + ) -def ptx_mbarrier_arrive_cluster_count(bar, cta_id, count): - """Cross-CTA ``mbarrier.arrive`` on CTA ``cta_id`` with an explicit count. +def ptx_mbarrier_complete_tx( + bar, + tx_count, + *, + sem="relaxed", + scope="cluster", + space="shared::cluster", + remote=None, + pred=None, +): + """PTX ``mbarrier.complete_tx{.sem.scope}{.space}.b64 [bar], tx_count``. - Convenience for an already-elected thread: emits - ``@p mapa.shared::cluster.u32`` + ``@p mbarrier.arrive.shared::cluster.b64 _, - [addr], count`` with the guard defaulted to 1. + ``remote`` optionally maps ``bar`` with ``mapa.shared::cluster.u32`` before + the complete_tx instruction. ``pred`` optionally predicates the instruction. """ - return call_intrin("", "tirx.ptx.mbarrier_arrive", bar, cta_id, True, count) - + if sem not in _MBARRIER_COMPLETE_TX_SEM: + raise ValueError(f"invalid sem={sem!r}; expected one of {_MBARRIER_COMPLETE_TX_SEM}") + if scope not in _MBARRIER_COMPLETE_TX_SCOPE: + raise ValueError(f"invalid scope={scope!r}; expected one of {_MBARRIER_COMPLETE_TX_SCOPE}") + if space not in _MBARRIER_COMPLETE_TX_SPACE: + raise ValueError(f"invalid space={space!r}; expected one of {_MBARRIER_COMPLETE_TX_SPACE}") + if remote is not None and space != "shared::cluster": + raise ValueError("remote mbarrier.complete_tx requires space='shared::cluster'") + + args = [bar, tx_count] + if remote is not None: + args.append(remote) + if pred is not None: + args.append(pred) + return call_intrin( + "", + "tirx.ptx.mbarrier_complete_tx", + *args, + sem, + scope, + space, + int(remote is not None), + int(pred is not None), + ) -def ptx_mbarrier_arrive_expect_tx(bar, byte_count, cta_id=None, pred=None): - """TVM intrinsic to call - mbarrier.arrive_expect_tx.shared::cta.b64 - or - @p mapa.shared::cluster.u32 - @p mbarrier.arrive_expect_tx.shared::cluster.b64 - Parameters - ---------- - bar : Var - The pointer to barrier variable. +def ptx_mbarrier_arrive_expect_tx( + bar, + byte_count, + *, + sem="", + scope="", + space=None, + remote=None, + pred=None, + **kwargs, +): + """PTX ``mbarrier.arrive.expect_tx{.sem.scope}{.space}.b64 _, [bar], txCount``.""" + if "cta_id" in kwargs: + raise ValueError("T.ptx.mbarrier.arrive.expect_tx uses remote= instead of cta_id=") + if kwargs: + raise TypeError(f"unexpected keyword argument(s): {', '.join(kwargs)}") + + if space is not None: + effective_space = space + elif remote is not None: + effective_space = "shared::cluster" + else: + effective_space = "shared" + _validate_mbarrier_arrive_attrs(sem, scope, effective_space, remote) - byte_count : int - Increases the tx count of the mbarrier object to track completion of - addtional async transactions. + args = [bar, byte_count] + if remote is not None: + args.append(remote) + if pred is not None: + args.append(pred) + return call_intrin( + "", + "tirx.ptx.mbarrier_arrive_expect_tx", + *args, + sem, + scope, + effective_space, + int(remote is not None), + int(pred is not None), + ) - cta_id : Optional[Expr] - The cta id. - pred : Optional[Expr] - The predicate to guard the operation. +def ptx_mbarrier_arrive_no_complete(bar, count, *, space="shared", pred=None, **kwargs): + """PTX ``mbarrier.arrive.noComplete.release.cta.{space}.b64 _, [bar], count``.""" + if "remote" in kwargs: + raise ValueError("mbarrier.arrive.no_complete does not support remote=") + if kwargs: + raise TypeError(f"unexpected keyword argument(s): {', '.join(kwargs)}") + if space not in ("shared", "shared::cta"): + raise ValueError("mbarrier.arrive.no_complete space must be 'shared' or 'shared::cta'") - Returns - ------- - call : Expr - The call expression. - """ - if cta_id is None and pred is None: - return call_intrin("", "tirx.ptx.mbarrier_arrive_expect_tx", bar, byte_count) - assert cta_id is not None - # Cross-CTA expect_tx from an already-elected thread: default the guard to 1 - # (the caller has elected a single lane), so callers can pass cta_id alone. - if pred is None: - pred = True - return call_intrin("", "tirx.ptx.mbarrier_arrive_expect_tx", bar, byte_count, cta_id, pred) + args = [bar, count] + if pred is not None: + args.append(pred) + return call_intrin( + "", + "tirx.ptx.mbarrier_arrive_no_complete", + *args, + space, + int(pred is not None), + ) def ptx_mbarrier_try_wait(bar, phase): @@ -763,6 +919,29 @@ def ptx_bar_sync(name_bar_id, thread_count): return call_intrin("", "tirx.ptx.bar_sync", name_bar_id, thread_count) +def ptx_barrier_sync(name_bar_id, thread_count): + """TVM intrinsic to call unaligned ``barrier.sync a, b``. + + Unlike the aligned ``bar.sync`` alias, this spelling is valid when + participating lanes reach the named barrier through divergent control + flow or distinct static instruction sites. + + Parameters + ---------- + name_bar_id : int + The ID of the named barrier. + + thread_count : int + The number of threads expected to arrive at the barrier. + + Returns + ------- + call : Expr + The call expression. + """ + return call_intrin("", "tirx.ptx.barrier_sync", name_bar_id, thread_count) + + def ptx_cp_async( dst_ptr, src_ptr, @@ -889,10 +1068,31 @@ def ptx_cp_async_wait_group(num=0): return call_intrin("", "tirx.ptx.cp_async_wait_group", num) -def ptx_cp_async_bulk_tensor_global_to_cluster( - dim, dst_ptr, bar, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None +_CP_ASYNC_BULK_TENSOR_LOAD_MODE = ("tile", "tile_gather4") + + +def _is_static_unicast_cta_mask(cta_mask): + if isinstance(cta_mask, int): + return cta_mask == 0 or cta_mask & (cta_mask - 1) == 0 + if isinstance(cta_mask, tirx.IntImm): + value = int(cta_mask) + return value == 0 or value & (value - 1) == 0 + return False + + +def ptx_cp_async_bulk_tensor_g2s_cta( + dim, + dst_ptr, + mbar, + tensormap_addr, + cta_group, + cache_hint, + *coords, + cache_policy=None, + load_mode="tile", + mbar_is_shared_addr=False, ): - """TVM intrinsic to call cp.async.bulk.tensor.dim.shared::cluster.global.tile.mbarrier::complete_tx::bytes + """TVM intrinsic for cp.async.bulk.tensor global -> shared::cta. Parameters ---------- @@ -902,21 +1102,15 @@ def ptx_cp_async_bulk_tensor_global_to_cluster( dst_ptr : Expr The destination pointer to the shared memory. - bar : Expr - The pointer to mbarrier variable. + mbar : Expr + The mbarrier pointer, or a uint32 shared-memory address operand when + ``mbar_is_shared_addr`` is true. tensormap_addr : Expr The generic address of the tensor map object. - cta_mask : int - The mask of the cta for multicast. - cta_group : int Must be either 1 or 2. - If set to 1, mbarrier must be in the shared memory of the same CTA - as the shared memory destination. If set to 2, mbarrier can be in - shared memory of either the same CTA as the shared memory destination - or the shared memory of the peer CTA. cache_hint : str The cache hint. @@ -924,48 +1118,67 @@ def ptx_cp_async_bulk_tensor_global_to_cluster( coords : List[Expr] specifies the starting coordinates in the tensor data in the global memory + load_mode : str + Either ``"tile"`` or ``"tile_gather4"``. + + mbar_is_shared_addr : bool + Whether ``mbar`` is already the uint32 shared-memory address operand. + Returns ------- call : Expr The call expression. - """ # noqa: E501 + """ _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP) + _choice("load_mode", load_mode, _CP_ASYNC_BULK_TENSOR_LOAD_MODE) if is_prim_expr(cache_hint): - has_cache_policy, *coords = coords + has_cache_policy, load_mode, mbar_is_shared_addr, *coords = coords return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster", + "tirx.ptx.cp_async_bulk_tensor_g2s_cta", dim, dst_ptr, - bar, + mbar, tensormap_addr, - cta_mask, cta_group, cache_hint, has_cache_policy, + load_mode, + mbar_is_shared_addr, *coords, ) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster", + "tirx.ptx.cp_async_bulk_tensor_g2s_cta", dim, dst_ptr, - bar, + mbar, tensormap_addr, - cta_mask, cta_group, cache_policy, int(has_cache_policy), + load_mode, + int(bool(mbar_is_shared_addr)), *coords, ) -def ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster( - dim, dst_ptr, bar, tensormap_addr, cta_mask, cta_group, cache_hint, *coords, cache_policy=None +def ptx_cp_async_bulk_tensor_g2s_cluster( + dim, + dst_ptr, + mbar, + tensormap_addr, + cta_mask, + cta_group, + cache_hint, + *coords, + cache_policy=None, + load_mode="tile", + mbar_is_shared_addr=False, + multicast=None, ): - """TVM intrinsic to call - cp.async.bulk.tensor.dim.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes + """TVM intrinsic for cp.async.bulk.tensor global -> shared::cluster. Parameters ---------- @@ -975,8 +1188,9 @@ def ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster( dst_ptr : Expr The destination pointer to the shared memory. - bar : Expr - The pointer to mbarrier variable. + mbar : Expr + The mbarrier pointer, or a uint32 shared-memory address operand when + ``mbar_is_shared_addr`` is true. tensormap_addr : Expr The generic address of the tensor map object. @@ -991,7 +1205,17 @@ def ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster( The cache hint. coords : List[Expr] - The TMA coordinates followed by the 4 gather row indices. + Specifies the starting coordinates in the tensor data in global memory. + + load_mode : str + Either ``"tile"`` or ``"tile_gather4"``. + + mbar_is_shared_addr : bool + Whether ``mbar`` is already the uint32 shared-memory address operand. + + multicast : bool, optional + Whether to emit ``.multicast::cluster``. If omitted, infer it from a + static ``cta_mask``; dynamic masks are treated as multicast. Returns ------- @@ -999,33 +1223,42 @@ def ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster( The call expression. """ _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP) + _choice("load_mode", load_mode, _CP_ASYNC_BULK_TENSOR_LOAD_MODE) if is_prim_expr(cache_hint): - has_cache_policy, *coords = coords + has_cache_policy, load_mode, mbar_is_shared_addr, multicast, *coords = coords return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster", + "tirx.ptx.cp_async_bulk_tensor_g2s_cluster", dim, dst_ptr, - bar, + mbar, tensormap_addr, cta_mask, cta_group, cache_hint, has_cache_policy, + load_mode, + mbar_is_shared_addr, + multicast, *coords, ) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) + if multicast is None: + multicast = not _is_static_unicast_cta_mask(cta_mask) return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster", + "tirx.ptx.cp_async_bulk_tensor_g2s_cluster", dim, dst_ptr, - bar, + mbar, tensormap_addr, cta_mask, cta_group, cache_policy, int(has_cache_policy), + load_mode, + int(bool(mbar_is_shared_addr)), + int(bool(multicast)), *coords, ) @@ -1082,9 +1315,7 @@ def ptx_cp_async_bulk_tensor_shared_to_global( ) -def ptx_cp_async_bulk_tensor_global_to_cluster_prefetch( - dim, tensormap_addr, cache_hint, *coords, cache_policy=None -): +def ptx_cp_async_bulk_tensor_prefetch(dim, tensormap_addr, cache_hint, *coords, cache_policy=None): """TVM intrinsic to call cp.async.bulk.prefetch.tensor.dim.L2.global.tile Parameters @@ -1110,7 +1341,7 @@ def ptx_cp_async_bulk_tensor_global_to_cluster_prefetch( has_cache_policy, *coords = coords return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster_prefetch", + "tirx.ptx.cp_async_bulk_tensor_prefetch", dim, tensormap_addr, cache_hint, @@ -1120,7 +1351,7 @@ def ptx_cp_async_bulk_tensor_global_to_cluster_prefetch( cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) return call_intrin( "", - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster_prefetch", + "tirx.ptx.cp_async_bulk_tensor_prefetch", dim, tensormap_addr, cache_policy, @@ -1266,7 +1497,7 @@ def ptx_clc_try_cancel(handle, mbar): return call_intrin("", "tirx.ptx.clc_try_cancel", handle, mbar) -def ptx_clc_query_cancel(handle): +def ptx_clc_query_cancel(handle, *, use_ld_acquire=True): """TVM intrinsic to call clusterlaunchcontrol.query_cancel. Decodes the response handle written by :func:`ptx_clc_try_cancel`. Returns the @@ -1276,8 +1507,12 @@ def ptx_clc_query_cancel(handle): ---------- handle : Expr Pointer to the 16B (uint4) smem response handle. + + use_ld_acquire : bool + Whether to load the 16-byte response with ``ld.acquire.cta.shared.b128``. + """ - return call_intrin("uint32", "tirx.ptx.clc_query_cancel", handle) + return call_intrin("uint32", "tirx.ptx.clc_query_cancel", handle, int(bool(use_ld_acquire))) def ptx_elect_sync(): @@ -2283,6 +2518,7 @@ def ptx_tcgen05_mma( cta_group, enable_input_d=1, scale_input_d=0, + weight_stationary=False, pred=None, ): """TVM intrinsic to call tcgen05.mma.cta_group.kind without block scaling. @@ -2325,6 +2561,11 @@ def ptx_tcgen05_mma( The optional scaling factor to scale input matrix D. D = A*B+D * (2 ^ - scale-input-d) + weight_stationary : bool + Whether to emit the ``tcgen05.mma.ws`` weight-stationary form. This is + required by kernels whose TMEM layout and MMA issue sequence follow the + PTX ``.ws`` form, such as FlashMLA head64 phase1. + disable_output_lane : list The lanes that should not be updated in the resultant matrix D. @@ -2352,6 +2593,7 @@ def ptx_tcgen05_mma( cta_group, enable_input_d, scale_input_d, + int(weight_stationary), *disable_output_lane, ] if pred is not None: @@ -2649,15 +2891,28 @@ def _choice(name: str, value, options): validation; specialization later replaces them with concrete values that the C-side intrinsic body re-checks. """ - # Concrete int / IntImm value: validate. - try: - concrete = int(value) - except (TypeError, ValueError): - return # symbolic; defer check + if isinstance(value, str): + concrete = value + elif isinstance(value, tirx.StringImm): + concrete = value.value + else: + # Concrete int / IntImm value: validate. + try: + concrete = int(value) + except (TypeError, ValueError): + return # symbolic; defer check if concrete not in options: raise ValueError(f"invalid {name}={concrete!r}; expected one of {tuple(options)}") +def _static_str(value): + if isinstance(value, str): + return value + if isinstance(value, tirx.StringImm): + return value.value + return None + + # See top-of-file imports for `_FENCE_SEM` etc. (re-exported from _common). # Note: TCGEN05_LDST_SHAPES values must stay in sync with the shape branches # of codegen_ptx_tcgen05_ld/_st in intrinsics/cuda/tcgen05.py. @@ -2710,12 +2965,15 @@ def ptx_tcgen05_cp( _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP) _choice("multicast", multicast, _TCGEN05_CP_MULTICAST) _choice("decompress", decompress, _TCGEN05_CP_DECOMPRESS) - if shape == "32x128b" and multicast != "warpx4": - raise ValueError(f"shape=32x128b requires multicast='warpx4', got {multicast!r}") - if shape == "64x128b" and multicast not in ("warpx2::02_13", "warpx2::01_23"): - raise ValueError(f"shape=64x128b requires multicast in warpx2::*, got {multicast!r}") - if shape in ("128x128b", "128x256b", "4x256b") and multicast != "": - raise ValueError(f"shape={shape} requires multicast='', got {multicast!r}") + shape_s = _static_str(shape) + multicast_s = _static_str(multicast) + if shape_s is not None and multicast_s is not None: + if shape_s == "32x128b" and multicast_s != "warpx4": + raise ValueError(f"shape=32x128b requires multicast='warpx4', got {multicast_s!r}") + if shape_s == "64x128b" and multicast_s not in ("warpx2::02_13", "warpx2::01_23"): + raise ValueError(f"shape=64x128b requires multicast in warpx2::*, got {multicast_s!r}") + if shape_s in ("128x128b", "128x256b", "4x256b") and multicast_s != "": + raise ValueError(f"shape={shape_s} requires multicast='', got {multicast_s!r}") return call_intrin( "", @@ -3150,8 +3408,8 @@ def cuda_printf(fmt, *args): return call_intrin("", "tirx.cuda.printf", fmt, *args) -def cuda_ldg(addr, dtype): - """TVM intrinsic to call CUDA C++ __ldg() function +def cuda_ldg(addr, dtype, *, dst=None, vec=""): + """TVM intrinsic to call CUDA C++ ``__ldg()``. Parameters ---------- @@ -3161,9 +3419,32 @@ def cuda_ldg(addr, dtype): dtype : str The data type of the loaded value. + dst : Expr or tuple[Expr], optional + Destination pointers for vector loads. + + vec : str + CUDA vector width. Use ``"v2"`` or ``"v4"`` together with tuple/list + ``dst``. + Returns """ - return call_intrin(dtype, "tirx.cuda.ldg", addr, dtype) + if dst is None: + if vec: + raise ValueError("vector cuda.ldg requires dst") + return call_intrin(dtype, "tirx.cuda.ldg", addr, dtype) + if vec not in ("v2", "v4"): + raise ValueError(f"vector cuda.ldg expects vec in {{'v2', 'v4'}}, got {vec!r}") + if not isinstance(dst, list | tuple): + raise ValueError("vector cuda.ldg requires tuple/list dst") + vec_len = int(vec[1:]) + if len(dst) != vec_len: + raise ValueError(f"cuda.ldg dst length must match {vec}: got {len(dst)}") + return call_intrin("", "tirx.cuda.ldg", *dst, addr, dtype, vec, vec_len) + + +def cuda_fdividef(x, y): + """TVM intrinsic to call CUDA C++ ``__fdividef`` fast float division.""" + return call_intrin("float32", "tirx.cuda.fdividef", x, y) def cuda_get_tmem_addr(addr, row_offset, col_offset): @@ -3207,9 +3488,12 @@ def cuda_smem_addr_from_uint64(cluster_addr): return call_intrin("uint32", "tirx.cuda.smem_addr_from_uint64", cluster_addr) -def cuda_sm100_tma_2sm_mbarrier_addr(bar): - """Compute the SM100 2SM TMA mbarrier shared-address operand.""" - return bitwise_and(cuda_cvta_generic_to_shared(bar), const(0xFEFFFFFF, dtype="uint32")) +def cuda_sm100_2sm_leader_smem_addr(ptr): + """Return the SM100 2SM leader CTA shared-address operand. + + The input is a generic pointer to shared memory. + """ + return bitwise_and(cuda_cvta_generic_to_shared(ptr), const(0xFEFFFFFF, dtype="uint32")) def ptx_exp2(x): @@ -3313,6 +3597,63 @@ def _ptx_binary_arith(op_name, dtype, d, a, b, *, rounding="rn", ftz=False, sat= ) +_PTX_F32X2_VALUE_RETURN_DTYPES = ("uint64", "float32x2") + + +def _validate_f32x2_value_return_dtype(return_dtype: str) -> str: + if return_dtype not in _PTX_F32X2_VALUE_RETURN_DTYPES: + raise ValueError( + f"invalid return_dtype={return_dtype!r}; expected one of " + f"{_PTX_F32X2_VALUE_RETURN_DTYPES}" + ) + return return_dtype + + +def _as_f32x2_bits(value): + if is_prim_expr(value) and str(value.ty) == "float32x2": + return reinterpret("uint64", value) + return value + + +def _ptx_binary_f32x2(op_name, *args, rounding="rn", ftz=False, dps=True, return_dtype="uint64"): + """Shared helper for packed f32x2 binary forms. + + ``dps=True`` emits the destination-passing form and returns void. + ``dps=False`` returns either packed ``uint64`` bits or ``float32x2``. + """ + rounding_options = ("", *_F32X2_ROUND) if op_name in ("add", "mul") else _F32X2_ROUND + _choice("rounding", rounding, rounding_options) + if dps: + if len(args) != 3: + raise TypeError(f"ptx_{op_name}_f32x2 dps form expects (d_addr, a, b)") + d, a, b = args + return call_intrin( + "", + f"tirx.ptx.{op_name}_f32x2", + d, + _as_f32x2_bits(a), + _as_f32x2_bits(b), + rounding, + int(ftz), + int(True), + "void", + ) + if len(args) != 2: + raise TypeError(f"ptx_{op_name}_f32x2 value form expects (a, b)") + return_dtype = _validate_f32x2_value_return_dtype(return_dtype) + a, b = args + return call_intrin( + return_dtype, + f"tirx.ptx.{op_name}_f32x2", + _as_f32x2_bits(a), + _as_f32x2_bits(b), + rounding, + int(ftz), + int(False), + return_dtype, + ) + + def _ptx_fma(dtype, d, a, b, c, *, rounding="rn", ftz=False, sat=False): """Shared helper for fma over (f32 | f32x2 | f64), DPS form.""" _choice("rounding", rounding, _F32X2_ROUND) @@ -3333,17 +3674,56 @@ def _ptx_fma(dtype, d, a, b, c, *, rounding="rn", ftz=False, sat=False): ) +def _ptx_fma_f32x2(*args, rounding="rn", ftz=False, dps=True, return_dtype="uint64"): + """Shared helper for packed f32x2 fma forms.""" + _choice("rounding", rounding, _F32X2_ROUND) + if dps: + if len(args) != 4: + raise TypeError("ptx_fma_f32x2 dps form expects (d_addr, a, b, c)") + d, a, b, c = args + return call_intrin( + "", + "tirx.ptx.fma_f32x2", + d, + _as_f32x2_bits(a), + _as_f32x2_bits(b), + _as_f32x2_bits(c), + rounding, + int(ftz), + int(True), + "void", + ) + if len(args) != 3: + raise TypeError("ptx_fma_f32x2 value form expects (a, b, c)") + return_dtype = _validate_f32x2_value_return_dtype(return_dtype) + a, b, c = args + return call_intrin( + return_dtype, + "tirx.ptx.fma_f32x2", + _as_f32x2_bits(a), + _as_f32x2_bits(b), + _as_f32x2_bits(c), + rounding, + int(ftz), + int(False), + return_dtype, + ) + + def ptx_add_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False): """PTX ``add{.rnd}{.ftz}{.sat}.f32 [d_addr], a, b`` — DPS form.""" return _ptx_binary_arith("add", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat) -def ptx_add_f32x2(d_addr, a, b, *, rounding="rn", ftz=False): - """PTX ``add{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form. +def ptx_add_f32x2(*args, rounding="rn", ftz=False, dps=True, return_dtype="uint64"): + """PTX ``add{.rnd}{.ftz}.f32x2``. - a, b are packed-as-uint64 register operands (2 fp32 each). + DPS form: ``(d_addr, a, b, dps=True)`` returns void. + Value form: ``(a, b, dps=False)`` returns ``return_dtype``. """ - return _ptx_binary_arith("add", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz) + return _ptx_binary_f32x2( + "add", *args, rounding=rounding, ftz=ftz, dps=dps, return_dtype=return_dtype + ) def ptx_add_f64(d_addr, a, b, *, rounding="rn"): @@ -3356,9 +3736,11 @@ def ptx_sub_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False): return _ptx_binary_arith("sub", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat) -def ptx_sub_f32x2(d_addr, a, b, *, rounding="rn", ftz=False): - """PTX ``sub{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form.""" - return _ptx_binary_arith("sub", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz) +def ptx_sub_f32x2(*args, rounding="rn", ftz=False, dps=True, return_dtype="uint64"): + """PTX ``sub{.rnd}{.ftz}.f32x2``; see :func:`ptx_add_f32x2`.""" + return _ptx_binary_f32x2( + "sub", *args, rounding=rounding, ftz=ftz, dps=dps, return_dtype=return_dtype + ) def ptx_sub_f64(d_addr, a, b, *, rounding="rn"): @@ -3371,9 +3753,11 @@ def ptx_mul_f32(d_addr, a, b, *, rounding="rn", ftz=False, sat=False): return _ptx_binary_arith("mul", "f32", d_addr, a, b, rounding=rounding, ftz=ftz, sat=sat) -def ptx_mul_f32x2(d_addr, a, b, *, rounding="rn", ftz=False): - """PTX ``mul{.rnd}{.ftz}.f32x2 [d_addr], a, b`` — DPS form.""" - return _ptx_binary_arith("mul", "f32x2", d_addr, a, b, rounding=rounding, ftz=ftz) +def ptx_mul_f32x2(*args, rounding="", ftz=False, dps=True, return_dtype="uint64"): + """PTX ``mul{.rnd}{.ftz}.f32x2``; see :func:`ptx_add_f32x2`.""" + return _ptx_binary_f32x2( + "mul", *args, rounding=rounding, ftz=ftz, dps=dps, return_dtype=return_dtype + ) def ptx_mul_f64(d_addr, a, b, *, rounding="rn"): @@ -3386,12 +3770,13 @@ def ptx_fma_f32(d_addr, a, b, c, *, rounding="rn", ftz=False, sat=False): return _ptx_fma("f32", d_addr, a, b, c, rounding=rounding, ftz=ftz, sat=sat) -def ptx_fma_f32x2(d_addr, a, b, c, *, rounding="rn", ftz=False): - """PTX ``fma{.rnd}{.ftz}.f32x2 [d_addr], a, b, c`` — DPS form. +def ptx_fma_f32x2(*args, rounding="rn", ftz=False, dps=True, return_dtype="uint64"): + """PTX ``fma{.rnd}{.ftz}.f32x2``. - a, b, c are packed-as-uint64 register operands. + DPS form: ``(d_addr, a, b, c, dps=True)`` returns void. + Value form: ``(a, b, c, dps=False)`` returns ``return_dtype``. """ - return _ptx_fma("f32x2", d_addr, a, b, c, rounding=rounding, ftz=ftz) + return _ptx_fma_f32x2(*args, rounding=rounding, ftz=ftz, dps=dps, return_dtype=return_dtype) def ptx_fma_f64(d_addr, a, b, c, *, rounding="rn"): @@ -3418,6 +3803,125 @@ def ptx_max_f32(a, b, *, ftz=False, nan=False): return call_intrin("float32", "tirx.ptx.max_f32", a, b, int(ftz), int(nan)) +_PTX_CVT_TYPES = { + "u8", + "u16", + "u32", + "u64", + "s8", + "s16", + "s32", + "s64", + "bf16", + "f16", + "f32", + "f64", + "f16x2", + "bf16x2", + "tf32", + "e4m3x2", + "e5m2x2", + "e2m1x2", + "e2m3x2", + "e3m2x2", + "e4m3x4", + "e5m2x4", + "e2m1x4", + "e2m3x4", + "e3m2x4", + "ue8m0x2", + "s2f6x2", +} +_PTX_CVT_ROUNDING = {"", "rni", "rzi", "rmi", "rpi", "rn", "rz", "rm", "rp", "rna", "rs"} +_PTX_CVT_SCALED = {"", "n2::ue8m0"} +_PTX_CVT_RETURN_TYPE = { + "u8": "uint8", + "s8": "int8", + "u16": "uint16", + "s16": "int16", + "u32": "uint32", + "s32": "int32", + "u64": "uint64", + "s64": "int64", + "f32": "float32", + "f64": "float64", + "f16": "uint16", + "bf16": "uint16", + "e4m3x2": "uint16", + "e5m2x2": "uint16", + "e2m1x2": "uint8", + "e2m3x2": "uint16", + "e3m2x2": "uint16", + "ue8m0x2": "uint16", + "s2f6x2": "uint16", + "tf32": "uint32", + "f16x2": "uint32", + "bf16x2": "uint32", + "e2m1x4": "uint16", + "e4m3x4": "uint32", + "e5m2x4": "uint32", + "e2m3x4": "uint32", + "e3m2x4": "uint32", +} + + +def ptx_cvt( + dtype, + *values, + atype, + rounding="", + ftz=False, + sat=False, + relu=False, + satfinite=False, + scaled="", + rbits=None, + scale_factor=None, +): + """Build one PTX ``cvt`` instruction. + + The concrete ``dtype``/``atype`` pair and modifiers select one grammar + form from the PTX ISA ``cvt`` table. ``values``, ``rbits`` and + ``scale_factor`` are register operands; all other arguments are instruction + attrs. Invalid combinations are rejected when the form is classified. + + Alternate and packed floating-point values use raw PTX register carriers: + a ``.b8`` result is returned as ``uint8``, a ``.b16`` result as ``uint16``, + and a ``.b32`` result as ``uint32``. + """ + _choice("dtype", dtype, _PTX_CVT_TYPES) + _choice("atype", atype, _PTX_CVT_TYPES) + _choice("rounding", rounding, _PTX_CVT_ROUNDING) + _choice("scaled", scaled, _PTX_CVT_SCALED) + has_rbits = rbits is not None + has_scale = scale_factor is not None + if has_rbits != (rounding == "rs"): + raise ValueError("ptx.cvt requires rbits exactly when rounding='rs'") + if has_scale != bool(scaled): + raise ValueError("ptx.cvt scale_factor must be present exactly when scaled is set") + operands = [*values] + if has_rbits: + operands.append(rbits) + if has_scale: + operands.append(scale_factor) + return call_intrin( + _PTX_CVT_RETURN_TYPE[dtype], + "tirx.ptx.cvt", + *operands, + len(values), + int(has_rbits), + int(has_scale), + dtype, + atype, + rounding, + int(bool(ftz)), + int(bool(sat)), + int(bool(relu)), + int(bool(satfinite)), + scaled, + ) + + def ptx_griddepcontrol_wait(): """TVM intrinsic for PTX ``griddepcontrol.wait`` (sm_90+). @@ -3458,6 +3962,7 @@ def ptx_griddepcontrol_launch_dependents(): } _PTX_LD_TYPE = set(_PTX_SCALAR_TYPE) _PTX_LD_COP = {"", "ca", "cg", "cs", "lu", "cv"} +_PTX_LD_GLOBAL_NC_COP = {"", "ca", "cg", "cs"} _PTX_LD_VEC = {"", "v2", "v4", "v8"} _PTX_L1_EVICT = { "", @@ -3506,6 +4011,42 @@ def _resolve_cache_policy(cache_hint, cache_policy, choices=_CP_ASYNC_BULK_CACHE return const(0, dtype="uint64"), False +def _ptx_vec_len(vec): + return int(vec[1:]) if vec else 1 + + +def _normalize_ptx_ld_dst(dst, vec, op_name): + if dst is None: + if vec: + raise ValueError(f"vec {op_name} requires dst") + return [], 0 + if isinstance(dst, list | tuple): + if not vec: + raise ValueError(f"{op_name} scatter dst requires vec") + vec_len = _ptx_vec_len(vec) + if len(dst) != vec_len: + raise ValueError(f"{op_name} scatter dst length must match {vec}: got {len(dst)}") + return list(dst), vec_len + return [dst], 1 + + +def _validate_ptx_address(addr, space, op_name): + """Validate pointer and raw shared-memory address forms.""" + addr_ty = getattr(addr, "ty", None) + if isinstance(addr_ty, PointerType): + return + if isinstance(addr_ty, PrimType): + if addr_ty.dtype == "uint32": + if not str(space).startswith("shared"): + raise ValueError(f"{op_name} uint32 address requires shared state space") + return + if addr_ty.dtype.startswith(("int", "uint")): + raise ValueError( + f"{op_name} integer address must be uint32 in shared state space, " + f"got {addr_ty.dtype}" + ) + + def ptx_ld_acquire( addr, return_type, @@ -3555,14 +4096,12 @@ def ptx_ld_acquire( _choice("ptx_type", ptx_type, _PTX_LD_TYPE) _choice("vec", vec, _PTX_LD_VEC) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) - to_dst = int(dst is not None) - if vec and dst is None: - raise ValueError("vec ld.acquire requires dst") + dst_args, to_dst = _normalize_ptx_ld_dst(dst, vec, "ld.acquire") if to_dst: return call_intrin( "", "tirx.ptx.ld_acquire", - dst, + *dst_args, addr, cache_policy, return_type, @@ -3611,18 +4150,23 @@ def ptx_ld( ): """TVM intrinsic for PTX ``ld{.weak}{.ss}{.cop}...`` loads.""" _choice("space", space, _PTX_LD_SPACE | {"const", "param::entry", "param::func"}) + if isinstance(cop, str) and cop not in _PTX_LD_COP: + if cop == "nc": + raise ValueError( + 'T.ptx.ld(..., cop="nc") is no longer supported; use T.ptx.ld_global_nc' + ) + raise ValueError(f"invalid cop={cop!r}; expected one of {tuple(_PTX_LD_COP)}") _choice("cop", cop, _PTX_LD_COP) _choice("ptx_type", ptx_type, _PTX_LD_TYPE) _choice("vec", vec, _PTX_LD_VEC) + _validate_ptx_address(addr, space, "ptx_ld") cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) - to_dst = int(dst is not None) - if vec and dst is None: - raise ValueError("vec ld requires dst") + dst_args, to_dst = _normalize_ptx_ld_dst(dst, vec, "ld") if to_dst: return call_intrin( "", "tirx.ptx.ld", - dst, + *dst_args, addr, cache_policy, return_type, @@ -3656,6 +4200,66 @@ def ptx_ld( ) +def ptx_ld_global_nc( + addr, + return_type, + ptx_type, + *, + dst=None, + cop="", + vec="", + cache_hint="", + cache_policy=None, + l1_evict="", + l2_evict="", + prefetch_size="", +): + """TVM intrinsic for PTX ``ld.global{.cop}.nc...`` loads.""" + if isinstance(cop, str) and cop not in _PTX_LD_GLOBAL_NC_COP: + raise ValueError( + f"invalid ld.global.nc cop={cop!r}; expected one of {tuple(_PTX_LD_GLOBAL_NC_COP)}" + ) + if isinstance(cop, str) and cop and (l1_evict or l2_evict): + raise ValueError("ld.global.nc with cop cannot use l1_evict or l2_evict") + _choice("cop", cop, _PTX_LD_GLOBAL_NC_COP) + _choice("ptx_type", ptx_type, _PTX_LD_TYPE) + _choice("vec", vec, _PTX_LD_VEC) + cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) + dst_args, to_dst = _normalize_ptx_ld_dst(dst, vec, "ld.global.nc") + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld_global_nc", + *dst_args, + addr, + cache_policy, + return_type, + cop, + vec, + ptx_type, + int(has_cache_policy), + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) + return call_intrin( + return_type, + "tirx.ptx.ld_global_nc", + addr, + cache_policy, + return_type, + cop, + "", + ptx_type, + int(has_cache_policy), + 0, + l1_evict, + l2_evict, + prefetch_size, + ) + + def ptx_ld_relaxed( addr, return_type, @@ -3676,14 +4280,12 @@ def ptx_ld_relaxed( _choice("ptx_type", ptx_type, _PTX_LD_TYPE) _choice("vec", vec, _PTX_LD_VEC) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) - to_dst = int(dst is not None) - if vec and dst is None: - raise ValueError("vec ld.relaxed requires dst") + dst_args, to_dst = _normalize_ptx_ld_dst(dst, vec, "ld.relaxed") if to_dst: return call_intrin( "", "tirx.ptx.ld_relaxed", - dst, + *dst_args, addr, cache_policy, return_type, @@ -3729,14 +4331,12 @@ def ptx_ld_volatile( _choice("space", space, _PTX_LD_VOLATILE_SPACE) _choice("ptx_type", ptx_type, _PTX_LD_TYPE) _choice("vec", vec, _PTX_LD_VEC) - to_dst = int(dst is not None) - if vec and dst is None: - raise ValueError("vec ld.volatile requires dst") + dst_args, to_dst = _normalize_ptx_ld_dst(dst, vec, "ld.volatile") if to_dst: return call_intrin( "", "tirx.ptx.ld_volatile", - dst, + *dst_args, addr, return_type, space, @@ -3895,7 +4495,24 @@ def ptx_st( _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"}) _choice("cop", cop, _PTX_ST_COP) _choice("vec", vec, _PTX_ST_VEC) - _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE) + _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE | {"b128"}) + _validate_ptx_address(address, space, "ptx_st") + if ptx_type == "b128" and not ( + src is not None + and not values + and weak + and space == "shared::cta" + and not cop + and not vec + and not cache_hint + and cache_policy is None + and not l1_evict + and not l2_evict + ): + raise ValueError( + "ptx_st b128 requires src=, weak=True, space='shared::cta', " + "and no cache/vector modifiers" + ) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) attrs = ( cache_policy, diff --git a/python/tvm/backend/cuda/operator/intrinsics/__init__.py b/python/tvm/backend/cuda/operator/intrinsics/__init__.py index 58c097149e2f..e30734a355ea 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/__init__.py +++ b/python/tvm/backend/cuda/operator/intrinsics/__init__.py @@ -21,6 +21,7 @@ - ``cp_async`` — cp.async + cp.async.bulk + cp.async.bulk.tensor (TMA), incl. TMA address helpers. - ``sync`` — barriers, fences, mbarrier, cluster.barrier, warp vote, elect, sync helpers. - ``math`` — packed-f32x2 arithmetic, exp2/rcp/reduce3, warp/CTA reductions. +- ``cvt`` — PTX data movement and conversion instruction forms. - ``memory`` — typed copies, ldg, ld.global.acquire, atomics, type conversions, address casts. - ``nvshmem`` — NVSHMEM RMA / signal / collective. - ``misc`` — register-allocation control, profiler timer, debug helpers (printf / trap). @@ -34,7 +35,7 @@ """ # Import op modules to register their codegen functions. -from . import cp_async, math, memory, misc, mma, nvshmem, sync, tcgen05, wgmma +from . import cp_async, cvt, math, memory, misc, mma, nvshmem, sync, tcgen05, wgmma from .header import TAGS, header_generator from .registry import CODEGEN_REGISTRY, get_codegen, register_codegen from .types import PTXDataType diff --git a/python/tvm/backend/cuda/operator/intrinsics/_schema.py b/python/tvm/backend/cuda/operator/intrinsics/_schema.py index 67c8b8cebd31..e9da9c3e4fdf 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/_schema.py +++ b/python/tvm/backend/cuda/operator/intrinsics/_schema.py @@ -61,6 +61,7 @@ "unsigned long long": "uint64", "long long": "int64", "unsigned short": "uint16", + "float2": "float32x2", "bool": "bool", "unsigned int": "uint32", "int": "int32", diff --git a/python/tvm/backend/cuda/operator/intrinsics/cp_async.py b/python/tvm/backend/cuda/operator/intrinsics/cp_async.py index 05bab2f31aff..f12817ad18fd 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/cp_async.py +++ b/python/tvm/backend/cuda/operator/intrinsics/cp_async.py @@ -18,9 +18,9 @@ """PTX cp.async / cp.async.bulk / cp.async.bulk.tensor intrinsics. Each PTX form table entry is registered as one ``device_intrinsic``. -User-facing wrappers in ``tvm.tirx.op`` keep their v1 signatures; -``register_codegen`` dispatchers below decode the (cp_size, fill_mode, -predicate) / (dim, cta_mask, tile_mode) arguments to pick the right form. +User-facing wrappers in ``tvm.tirx.op`` expose separate PTX instruction +families; ``register_codegen`` dispatchers below decode attrs such as +``load_mode`` / ``cta_group`` / ``multicast`` to pick the right form variant. Bodies are hand-written ``asm volatile(...)`` strings. The file is grouped as cp.async, cp.async.bulk.tensor, cp.async.bulk non-TMA, and CUDA compatibility helpers. @@ -59,6 +59,35 @@ def _safe(s): body=lambda n: f' asm volatile("cp.async.wait_group {int(n)};");', ) +_CP_ASYNC_MBARRIER_ARRIVE_SPACES = ("shared", "shared::cta") + + +def _cp_async_mbarrier_arrive_parts(*args): + noinc = _bool_attr(args[-2]) + space = parse_str(args[-1]) + assert space in _CP_ASYNC_MBARRIER_ARRIVE_SPACES, ( + f"invalid cp.async.mbarrier.arrive space {space!r}, " + f"expected one of {_CP_ASYNC_MBARRIER_ARRIVE_SPACES}" + ) + noinc_suffix = ".noinc" if noinc else "" + noinc_name = "_noinc" if noinc else "" + space_name = "_" + _safe(space) + return ( + f"tvm_builtin_ptx_cp_async_mbarrier_arrive{noinc_name}{space_name}", + " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" + f' asm volatile("cp.async.mbarrier.arrive{noinc_suffix}.{space}.b64 [%0];"\n' + ' :: "r"(barrier_addr) : "memory");', + ) + + +device_intrinsic( + "ptx_cp_async_mbarrier_arrive", + n_attrs=2, + helper_name=lambda *a: _cp_async_mbarrier_arrive_parts(*a)[0], + c_signature="(void* barrier)", + body=lambda *a: _cp_async_mbarrier_arrive_parts(*a)[1], +) + # cp.async non-bulk copy forms: # Form 1: cp.async.ca.shared.global ... [dst], [src], cp-size{, src-size}{, cache-policy} @@ -421,41 +450,158 @@ def _coord_sig(n): return ", ".join(f"int coord{i}" for i in range(n)) +def _ptx_shared_u32_addr_decl(var_name, ptr_name, ptx_reg_name): + return ( + f" unsigned int {var_name};\n" + " asm volatile(\n" + ' "{\\n\\t"\n' + f' ".reg .u64 {ptx_reg_name};\\n\\t"\n' + f' "cvta.to.shared.u64 {ptx_reg_name}, %1;\\n\\t"\n' + f' "cvt.u32.u64 %0, {ptx_reg_name};\\n\\t"\n' + ' "}\\n"\n' + f' : "=r"({var_name}) : "l"({ptr_name}));\n' + ) + + +def _g2s_cta_mbar_addr_decl(mbar_is_shared_addr): + if mbar_is_shared_addr: + # The caller passed the final uint32 mbarrier operand and owns any + # architecture-specific address masking before this helper is called. + return "" + return _ptx_shared_u32_addr_decl("mbar_addr", "mbar", "mbar_addr64") + + +def _g2s_cluster_mbar_addr_decl(mbar_is_shared_addr): + if mbar_is_shared_addr: + # The caller passed the final uint32 mbarrier operand and owns any + # architecture-specific address masking before this helper is called. + return "" + return " unsigned int mbar_addr = __cvta_generic_to_shared(mbar);\n" + + +def _tensor_tile_modifier(tile_mode): + assert tile_mode in _TILE_MODE_CHOICES, ( + f"invalid cp.async.bulk.tensor load_mode {tile_mode!r}, " + f"expected one of {_TILE_MODE_CHOICES}" + ) + return ".tile::gather4" if tile_mode == "tile_gather4" else "" + + +def _g2s_cta_body(instr, coord_count, has_cache, mbar_is_shared_addr): + coord_start = 4 if has_cache else 3 + coord_tpl = _coord_template(coord_count, coord_start) + cache_slot = ", %3" if has_cache else "" + cache_arg = ',\n "l"(cache_policy)' if has_cache else "" + return ( + f"{_ptx_shared_u32_addr_decl('dst_addr', 'dst', 'smem_addr64')}" + f"{_g2s_cta_mbar_addr_decl(mbar_is_shared_addr)}" + " asm volatile(\n" + f' "{instr} [%0], [%1, {coord_tpl}], [%2]{cache_slot};"\n' + " :\n" + f' : "r"(dst_addr), "l"(tensormap_addr), "r"(mbar_addr){cache_arg},\n' + f" {_coord_constraints(coord_count)}\n" + ' : "memory"\n' + " );" + ) + + +def _g2s_cluster_body( + instr, + coord_count, + coord_tpl, + mbar_is_shared_addr, + mask_arg, + mask_slot, + cache_arg, + cache_slot, +): + return ( + " unsigned int dst_addr = __cvta_generic_to_shared(dst);\n" + f"{_g2s_cluster_mbar_addr_decl(mbar_is_shared_addr)}" + " asm volatile(\n" + f' "{instr} [%0], [%1, {coord_tpl}], [%2]{mask_slot}{cache_slot};"\n' + " :\n" + f' : "r"(dst_addr), "l"(tensormap_addr), "r"(mbar_addr){mask_arg}{cache_arg},\n' + f" {_coord_constraints(coord_count)}\n" + ' : "memory"\n' + " );" + ) + + +# PTX cp.async.bulk.tensor global -> shared::cta; this registration supports +# tile/tile::gather4 load modes. +def _g2s_cta_parts(*args): + attrs = args[-5:] + dim = int(attrs[0]) + cta_group = int(attrs[1]) + has_cache = _bool_attr(attrs[2]) + tile_mode = parse_str(attrs[3]) + mbar_is_shared_addr = _bool_attr(attrs[4]) + coord_count = 5 if tile_mode == "tile_gather4" else dim + mbar_type = "unsigned int mbar_addr" if mbar_is_shared_addr else "void* mbar" + sig = ( + f"(void* dst, {mbar_type}, unsigned long long tensormap_addr, " + "unsigned long long cache_policy" + + (", " + _coord_sig(coord_count) if coord_count else "") + + ")" + ) + name = ( + f"ptx_cp_async_bulk_tensor_g2s_cta_{tile_mode}_{dim}d" + f"{'_cache_hint' if has_cache else ''}{'_mbar_addr' if mbar_is_shared_addr else ''}" + ) + tile_modifier = _tensor_tile_modifier(tile_mode) + cta_group_str = _resolve_cta_group_str(cta_group) + cache_inst = ".L2::cache_hint" if has_cache else "" + instr = ( + f"cp.async.bulk.tensor.{dim}d.shared::cta.global{tile_modifier}" + f".mbarrier::complete_tx::bytes{cta_group_str}{cache_inst}" + ) + return name, sig, _g2s_cta_body(instr, coord_count, has_cache, mbar_is_shared_addr) + + +device_intrinsic( + "_ptx_cp_async_bulk_tensor_g2s_cta", + n_attrs=5, + helper_name=lambda *a: _g2s_cta_parts(*a)[0], + c_signature=lambda *a: _g2s_cta_parts(*a)[1], + body=lambda *a: _g2s_cta_parts(*a)[2], +) + + # PTX cp.async.bulk.tensor global -> shared::cluster form: # cp.async.bulk.tensor.dim.dst.src{.load_mode}.completion_mechanism # {.multicast}{.cta_group}{.level::cache_hint} -# [dstMem], [tensorMap, tensorCoords], [mbar]{, im2colInfo} -# {, ctaMask} {, cache-policy} +# [dstMem], [tensorMap, tensorCoords], [mbar]{, ctaMask} {, cache-policy} # .dst = {.shared::cluster}; .src = {.global} # .completion_mechanism = {.mbarrier::complete_tx::bytes} # .multicast = {.multicast::cluster} # .cta_group = {.cta_group::1, .cta_group::2} # .load_mode = {.tile, .tile::gather4, .im2col, .im2col::w, .im2col::w::128} # .level::cache_hint = {.L2::cache_hint} -# This registration supports tile/tile::gather4 modes; ctaMask is only used +# This registration supports tile/tile::gather4 modes. ctaMask is only used # when the optional ``.multicast::cluster`` modifier is enabled. -def _g2cluster_parts(*args): +def _g2s_cluster_parts(*args): attrs = args[-6:] dim = int(attrs[0]) cta_group = int(attrs[1]) has_cache = _bool_attr(attrs[2]) tile_mode = parse_str(attrs[3]) - bar_is_addr = _bool_attr(attrs[4]) + mbar_is_shared_addr = _bool_attr(attrs[4]) multicast = _bool_attr(attrs[5]) coord_count = 5 if tile_mode == "tile_gather4" else dim - bar_type = "unsigned int bar_addr" if bar_is_addr else "void* bar" + mbar_type = "unsigned int mbar_addr" if mbar_is_shared_addr else "void* mbar" sig = ( - f"(void* dst, {bar_type}, unsigned long long tensormap_addr, " + f"(void* dst, {mbar_type}, unsigned long long tensormap_addr, " "uint16_t cta_mask, unsigned long long cache_policy" + (", " + _coord_sig(coord_count) if coord_count else "") + ")" ) name = ( - f"ptx_cp_async_bulk_tensor_g2cluster_{tile_mode}_{dim}d" + f"ptx_cp_async_bulk_tensor_g2s_cluster_{tile_mode}_{dim}d" f"{'_multicast' if multicast else ''}" - f"{'_cache_hint' if has_cache else ''}{'_bar_addr' if bar_is_addr else ''}" + f"{'_cache_hint' if has_cache else ''}{'_mbar_addr' if mbar_is_shared_addr else ''}" ) - tile_modifier = ".tile::gather4" if tile_mode == "tile_gather4" else "" + tile_modifier = _tensor_tile_modifier(tile_mode) cta_group_str = _resolve_cta_group_str(cta_group) multicast_inst = ".multicast::cluster" if multicast else "" cache_inst = ".L2::cache_hint" if has_cache else "" @@ -470,29 +616,28 @@ def _g2cluster_parts(*args): f".mbarrier::complete_tx::bytes{multicast_inst}" f"{cta_group_str}{cache_inst}" ) - bar_addr_decl = ( - "" if bar_is_addr else " unsigned int bar_addr = __cvta_generic_to_shared(bar);\n" - ) - body = ( - " unsigned int dst_addr = __cvta_generic_to_shared(dst);\n" - f"{bar_addr_decl}" - " asm volatile(\n" - f' "{instr} [%0], [%1, {coord_tpl}], [%2]{mask_slot}{cache_slot};"\n' - " :\n" - f' : "r"(dst_addr), "l"(tensormap_addr), "r"(bar_addr){mask_arg}{cache_arg},\n' - f" {_coord_constraints(coord_count)}\n" - ' : "memory"\n' - " );" + return ( + name, + sig, + _g2s_cluster_body( + instr, + coord_count, + coord_tpl, + mbar_is_shared_addr, + mask_arg, + mask_slot, + cache_arg, + cache_slot, + ), ) - return name, sig, body device_intrinsic( - "ptx_cp_async_bulk_tensor_g2cluster", + "_ptx_cp_async_bulk_tensor_g2s_cluster", n_attrs=6, - helper_name=lambda *a: _g2cluster_parts(*a)[0], - c_signature=lambda *a: _g2cluster_parts(*a)[1], - body=lambda *a: _g2cluster_parts(*a)[2], + helper_name=lambda *a: _g2s_cluster_parts(*a)[0], + c_signature=lambda *a: _g2s_cluster_parts(*a)[1], + body=lambda *a: _g2s_cluster_parts(*a)[2], ) @@ -559,10 +704,7 @@ def _prefetch_parts(*args): + (", " + _coord_sig(dim) if dim else "") + ")" ) - name = ( - f"ptx_cp_async_bulk_tensor_global_to_cluster_prefetch_{dim}d" - f"{'_cache_hint' if has_cache else ''}" - ) + name = f"ptx_cp_async_bulk_tensor_prefetch_{dim}d{'_cache_hint' if has_cache else ''}" cache_inst = ".L2::cache_hint" if has_cache else "" cache_arg = ', "l"(cache_policy)' if has_cache else "" cache_slot = ", %1" if has_cache else "" @@ -644,50 +786,57 @@ def _reduce_parts(*args): ) -# User-facing dispatchers for tensor global -> shared::cluster. The same -# backend root handles the optional ``.multicast::cluster`` modifier. +@register_codegen("ptx_cp_async_bulk_tensor_g2s_cta") +def codegen_g2s_cta(dim, dst_ptr, mbar, tensormap, *args): + cta_group, cache_policy, has_cache, tile_mode, mbar_is_shared_addr, *coords = args + result = CODEGEN_REGISTRY["tirx._ptx_cp_async_bulk_tensor_g2s_cta"]( + [ + dst_ptr, + mbar, + tensormap, + cache_policy, + *coords, + int(dim), + int(cta_group), + has_cache, + tile_mode, + mbar_is_shared_addr, + ] + ) + return result[0] if isinstance(result, tuple) else result -def _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode): - cta_mask, cta_group, cache_policy, has_cache, *rest = args - coord_count = 5 if tile_mode == "tile_gather4" else int(dim) - if len(rest) == coord_count + 1: - bar_is_addr = _bool_attr(rest[0]) - coords = rest[1:] - else: - bar_is_addr = False - coords = rest - is_unicast = isinstance(cta_mask, tvm.tirx.IntImm) and bin(int(cta_mask)).count("1") <= 1 - cg = int(cta_group) - op = "tirx.ptx_cp_async_bulk_tensor_g2cluster" - call_args = [ - dst_ptr, - bar, - tensormap, +@register_codegen("ptx_cp_async_bulk_tensor_g2s_cluster") +def codegen_g2s_cluster(dim, dst_ptr, mbar, tensormap, *args): + ( cta_mask, + cta_group, cache_policy, - *coords, - int(dim), - cg, has_cache, tile_mode, - bar_is_addr, - int(not is_unicast), - ] - result = CODEGEN_REGISTRY[op](call_args) + mbar_is_shared_addr, + multicast, + *coords, + ) = args + result = CODEGEN_REGISTRY["tirx._ptx_cp_async_bulk_tensor_g2s_cluster"]( + [ + dst_ptr, + mbar, + tensormap, + cta_mask, + cache_policy, + *coords, + int(dim), + int(cta_group), + has_cache, + tile_mode, + mbar_is_shared_addr, + multicast, + ] + ) return result[0] if isinstance(result, tuple) else result -@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster") -def codegen_g2c(dim, dst_ptr, bar, tensormap, *args): - return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile") - - -@register_codegen("ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster") -def codegen_g2c_gather4(dim, dst_ptr, bar, tensormap, *args): - return _g2c_dispatch(dim, dst_ptr, bar, tensormap, *args, tile_mode="tile_gather4") - - @register_codegen("ptx_cp_async_bulk_tensor_shared_to_global") def codegen_s2g(dim, src_ptr, tensormap, *args): cache_policy, has_cache, *coords = args @@ -697,7 +846,7 @@ def codegen_s2g(dim, src_ptr, tensormap, *args): return result[0] if isinstance(result, tuple) else result -@register_codegen("ptx_cp_async_bulk_tensor_global_to_cluster_prefetch") +@register_codegen("ptx_cp_async_bulk_tensor_prefetch") def codegen_prefetch(dim, tensormap, *args): cache_policy, has_cache, *coords = args result = CODEGEN_REGISTRY["tirx.ptx_cp_async_bulk_tensor_prefetch"]( @@ -731,7 +880,8 @@ def _ptx_cp_async_bulk_wait_group_parts(n, read): read_b = bool(int(read)) if hasattr(read, "value") else bool(read) return ( f"ptx_cp_async_bulk_wait_group{'_read' if read_b else ''}_{n}", - f' asm volatile("cp.async.bulk.wait_group{".read" if read_b else ""} {n};");', + f' asm volatile("cp.async.bulk.wait_group{".read" if read_b else ""} {n};" ' + '::: "memory");', ) diff --git a/python/tvm/backend/cuda/operator/intrinsics/cvt.py b/python/tvm/backend/cuda/operator/intrinsics/cvt.py new file mode 100644 index 000000000000..b0f3f6c7e638 --- /dev/null +++ b/python/tvm/backend/cuda/operator/intrinsics/cvt.py @@ -0,0 +1,754 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name, too-many-arguments, too-many-branches, too-many-locals +"""PTX ``cvt`` instruction forms. + +The PTX ISA specifies ``cvt`` as a family of grammar forms. The public +``tirx.ptx.cvt`` op carries concrete PTX types and optional modifiers as attrs. +Its codegen classifies the invocation into one of the forms below. Each form +then receives only its register operands followed by the attrs admitted by that +specific grammar line. + +Alternate and packed floating-point formats use their PTX register carriers: +``uint8`` for ``.b8`` values, ``uint16`` for ``.b16`` values, and ``uint32`` +for ``.b32`` values. +""" + +from dataclasses import dataclass + +from ._schema import device_intrinsic +from .registry import CODEGEN_REGISTRY, register_codegen +from .utils import parse_str + +_BASIC_TYPES = { + "u8", + "u16", + "u32", + "u64", + "s8", + "s16", + "s32", + "s64", + "bf16", + "f16", + "f32", + "f64", +} +_INTEGER_ROUNDING = {"rni", "rzi", "rmi", "rpi"} +_FLOAT_ROUNDING = {"rn", "rz", "rm", "rp"} +_F8X2_TYPES = {"e4m3x2", "e5m2x2"} +_F4X2_TYPES = {"e2m1x2"} +_F6X2_TYPES = {"e2m3x2", "e3m2x2"} +_F8X4_TYPES = {"e4m3x4", "e5m2x4"} +_F4X4_TYPES = {"e2m1x4"} +_F6X4_TYPES = {"e2m3x4", "e3m2x4"} +_FP16X2_TYPES = {"f16x2", "bf16x2"} + + +def _as_bool(value) -> bool: + return bool(int(value)) if hasattr(value, "value") else bool(value) + + +def _safe(value: str) -> str: + return value.replace("::", "_").replace(".", "_") + + +@dataclass(frozen=True) +class _CvtForm: + """One syntax line from the PTX ``cvt`` grammar.""" + + attrs: tuple[str, ...] + primary_count: int + dtype: str + atype: str + modifiers: tuple[str, ...] + rbits: bool = False + grouped_primary: bool = False + + +# Modifier descriptors: +# literal -> mandatory modifier, e.g. ``rn`` +# ``$rounding`` -> string-valued modifier attr +# ``?relu`` -> boolean modifier attr +# ``~scaled`` -> ``.scaled::`` and an extra scale-factor operand +_CVT_FORMS = { + # Generic scalar forms. + "scalar": _CvtForm(("dtype", "atype", "ftz", "sat"), 1, "$dtype", "$atype", ("?ftz", "?sat")), + "irnd": _CvtForm( + ("dtype", "atype", "rounding", "ftz", "sat"), + 1, + "$dtype", + "$atype", + ("$rounding", "?ftz", "?sat"), + ), + "frnd": _CvtForm( + ("dtype", "atype", "rounding", "ftz", "sat"), + 1, + "$dtype", + "$atype", + ("$rounding", "?ftz", "?sat"), + ), + # f16 / bf16 / tf32 destination forms. + "frnd2_f16_f32": _CvtForm( + ("rounding", "relu", "satfinite"), + 1, + "f16", + "f32", + ("$rounding", "?relu", "?satfinite"), + ), + "frnd2_f16x2_f32": _CvtForm( + ("rounding", "relu", "satfinite"), + 2, + "f16x2", + "f32", + ("$rounding", "?relu", "?satfinite"), + ), + "rs_f16x2_f32": _CvtForm( + ("relu", "satfinite"), + 2, + "f16x2", + "f32", + ("rs", "?relu", "?satfinite"), + rbits=True, + ), + "frnd2_bf16_f32": _CvtForm( + ("rounding", "relu", "satfinite"), + 1, + "bf16", + "f32", + ("$rounding", "?relu", "?satfinite"), + ), + "frnd2_bf16x2_f32": _CvtForm( + ("rounding", "relu", "satfinite"), + 2, + "bf16x2", + "f32", + ("$rounding", "?relu", "?satfinite"), + ), + "rs_bf16x2_f32": _CvtForm( + ("relu", "satfinite"), + 2, + "bf16x2", + "f32", + ("rs", "?relu", "?satfinite"), + rbits=True, + ), + "rna_tf32_f32": _CvtForm(("satfinite",), 1, "tf32", "f32", ("rna", "?satfinite")), + "frnd2_tf32_f32": _CvtForm( + ("rounding", "satfinite", "relu"), + 1, + "tf32", + "f32", + ("$rounding", "?satfinite", "?relu"), + ), + # f8 forms. + "rn_satfinite_f8x2_f32": _CvtForm( + ("f8x2type", "relu"), + 2, + "$f8x2type", + "f32", + ("rn", "satfinite", "?relu"), + ), + "rn_satfinite_f8x2_fp16x2": _CvtForm( + ("f8x2type", "fp16x2type", "relu"), + 1, + "$f8x2type", + "$fp16x2type", + ("rn", "satfinite", "?relu"), + ), + "rn_f16x2_f8x2": _CvtForm(("f8x2type", "relu"), 1, "f16x2", "$f8x2type", ("rn", "?relu")), + "rn_bf16x2_f8x2": _CvtForm( + ("f8x2type", "relu", "satfinite", "scaled"), + 1, + "bf16x2", + "$f8x2type", + ("rn", "?relu", "?satfinite", "~scaled"), + ), + "rs_satfinite_f8x4_f32": _CvtForm( + ("f8x4type", "relu"), + 4, + "$f8x4type", + "f32", + ("rs", "?relu", "satfinite"), + rbits=True, + grouped_primary=True, + ), + # f4 forms. + "rn_satfinite_f4x2_f32": _CvtForm( + ("f4x2type", "relu"), + 2, + "$f4x2type", + "f32", + ("rn", "satfinite", "?relu"), + ), + "rn_satfinite_f4x2_fp16x2": _CvtForm( + ("f4x2type", "fp16x2type", "relu"), + 1, + "$f4x2type", + "$fp16x2type", + ("rn", "satfinite", "?relu"), + ), + "rn_f16x2_f4x2": _CvtForm(("f4x2type", "relu"), 1, "f16x2", "$f4x2type", ("rn", "?relu")), + "rn_bf16x2_f4x2": _CvtForm( + ("f4x2type", "relu", "satfinite", "scaled"), + 1, + "bf16x2", + "$f4x2type", + ("rn", "?relu", "?satfinite", "~scaled"), + ), + "rs_satfinite_f4x4_f32": _CvtForm( + ("f4x4type", "relu"), + 4, + "$f4x4type", + "f32", + ("rs", "?relu", "satfinite"), + rbits=True, + grouped_primary=True, + ), + # f6 forms. + "rn_satfinite_f6x2_f32": _CvtForm( + ("f6x2type", "relu"), + 2, + "$f6x2type", + "f32", + ("rn", "satfinite", "?relu"), + ), + "rn_satfinite_f6x2_fp16x2": _CvtForm( + ("f6x2type", "fp16x2type", "relu"), + 1, + "$f6x2type", + "$fp16x2type", + ("rn", "satfinite", "?relu"), + ), + "rn_f16x2_f6x2": _CvtForm(("f6x2type", "relu"), 1, "f16x2", "$f6x2type", ("rn", "?relu")), + "rn_bf16x2_f6x2": _CvtForm( + ("f6x2type", "relu", "satfinite", "scaled"), + 1, + "bf16x2", + "$f6x2type", + ("rn", "?relu", "?satfinite", "~scaled"), + ), + "rs_satfinite_f6x4_f32": _CvtForm( + ("f6x4type", "relu"), + 4, + "$f6x4type", + "f32", + ("rs", "?relu", "satfinite"), + rbits=True, + grouped_primary=True, + ), + # ue8m0 forms. + "frnd3_ue8m0x2_f32": _CvtForm( + ("rounding", "satfinite"), + 2, + "ue8m0x2", + "f32", + ("$rounding", "?satfinite"), + ), + "frnd3_ue8m0x2_bf16x2": _CvtForm( + ("rounding", "satfinite"), + 1, + "ue8m0x2", + "bf16x2", + ("$rounding", "?satfinite"), + ), + "rn_bf16x2_ue8m0x2": _CvtForm((), 1, "bf16x2", "ue8m0x2", ("rn",)), + # s2f6 forms. + "rn_satfinite_s2f6x2_f32": _CvtForm( + ("relu", "scaled"), + 2, + "s2f6x2", + "f32", + ("rn", "satfinite", "?relu", "~scaled"), + ), + "rn_satfinite_s2f6x2_bf16x2": _CvtForm( + ("relu", "scaled"), + 1, + "s2f6x2", + "bf16x2", + ("rn", "satfinite", "?relu", "~scaled"), + ), + "rn_bf16x2_s2f6x2": _CvtForm( + ("satfinite", "relu", "scaled"), + 1, + "bf16x2", + "s2f6x2", + ("rn", "?satfinite", "?relu", "~scaled"), + ), +} + + +_STRING_ATTRS = { + "dtype", + "atype", + "rounding", + "f8x2type", + "f8x4type", + "f4x2type", + "f4x4type", + "f6x2type", + "f6x4type", + "fp16x2type", + "scaled", +} + + +def _parse_form_attrs(form: _CvtForm, args): + raw_attrs = args[-len(form.attrs) :] if form.attrs else () + attrs = {} + for name, value in zip(form.attrs, raw_attrs): + attrs[name] = parse_str(value) if name in _STRING_ATTRS else _as_bool(value) + operands = args[: -len(form.attrs)] if form.attrs else args + return operands, attrs + + +def _resolve(value: str, attrs: dict) -> str: + return attrs[value[1:]] if value.startswith("$") else value + + +def _modifier_string(form: _CvtForm, attrs: dict) -> str: + out = "" + for modifier in form.modifiers: + if modifier.startswith("$"): + value = attrs[modifier[1:]] + if value: + out += f".{value}" + elif modifier.startswith("?"): + name = modifier[1:] + if attrs[name]: + out += f".{name}" + elif modifier.startswith("~"): + value = attrs[modifier[1:]] + if value: + out += f".scaled::{value}" + else: + out += f".{modifier}" + return out + + +def _result_info(dtype: str): + """Return (C return type, result temporary type, constraint, TVM dtype).""" + if dtype == "u8": + return "uint8_t", "uint32_t", "r", "uint8" + if dtype == "s8": + return "int8_t", "int32_t", "r", "int8" + if dtype == "u16": + return "uint16_t", "uint16_t", "h", "uint16" + if dtype == "s16": + return "int16_t", "int16_t", "h", "int16" + if dtype == "u32": + return "uint32_t", "uint32_t", "r", "uint32" + if dtype == "s32": + return "int32_t", "int32_t", "r", "int32" + if dtype == "u64": + return "uint64_t", "uint64_t", "l", "uint64" + if dtype == "s64": + return "int64_t", "int64_t", "l", "int64" + if dtype == "f32": + return "float", "float", "f", "float32" + if dtype == "f64": + return "double", "double", "d", "float64" + if dtype == "e2m1x2": + # CUDA extended asm has no 8-bit register constraint. The helper + # stages the PTX .b8 result through a .b16 output register. + return "uint8_t", "uint16_t", "h", "uint8" + if dtype in {"f16", "bf16", "e4m3x2", "e5m2x2", "e2m1x4"}: + return "uint16_t", "uint16_t", "h", "uint16" + if dtype in {"e2m3x2", "e3m2x2", "ue8m0x2", "s2f6x2"}: + return "uint16_t", "uint16_t", "h", "uint16" + if dtype in { + "tf32", + "f16x2", + "bf16x2", + "e4m3x4", + "e5m2x4", + "e2m3x4", + "e3m2x4", + }: + return "uint32_t", "uint32_t", "r", "uint32" + raise ValueError(f"Unsupported PTX cvt destination type {dtype!r}") + + +def _input_info(atype: str): + """Return (C parameter type, asm constraint) for one source operand.""" + if atype == "u8": + return "uint32_t", "r" + if atype == "s8": + return "int32_t", "r" + if atype == "u16": + return "uint16_t", "h" + if atype == "s16": + return "int16_t", "h" + if atype == "u32": + return "uint32_t", "r" + if atype == "s32": + return "int32_t", "r" + if atype == "u64": + return "uint64_t", "l" + if atype == "s64": + return "int64_t", "l" + if atype == "f32": + return "float", "f" + if atype == "f64": + return "double", "d" + if atype in { + "f16", + "bf16", + "e4m3x2", + "e5m2x2", + "e2m3x2", + "e3m2x2", + "ue8m0x2", + "s2f6x2", + }: + return "uint16_t", "h" + if atype == "e2m1x2": + # CUDA extended asm has no 8-bit register constraint. PTX permits a + # .b8 value to be staged through a wider register. + return "uint16_t", "h" + if atype in { + "tf32", + "f16x2", + "bf16x2", + "e4m3x4", + "e5m2x4", + "e2m1x4", + "e2m3x4", + "e3m2x4", + }: + return "uint32_t", "r" + raise ValueError(f"Unsupported PTX cvt source type {atype!r}") + + +def _cvt_form_parts(form_name: str, *args): + form = _CVT_FORMS[form_name] + operands, attrs = _parse_form_attrs(form, args) + dtype = _resolve(form.dtype, attrs) + atype = _resolve(form.atype, attrs) + has_scale = bool(attrs.get("scaled", "")) + expected_operands = form.primary_count + int(form.rbits) + int(has_scale) + if len(operands) != expected_operands: + raise ValueError( + f"PTX cvt form {form_name!r} expects {expected_operands} operands, got {len(operands)}" + ) + + source_types = [atype] * form.primary_count + if form.rbits: + source_types.append("u32") + if has_scale: + source_types.append("u16") + source_info = [_input_info(source_type) for source_type in source_types] + source_names = ["a", "b", "e", "f"][: form.primary_count] + if form.rbits: + source_names.append("rbits") + if has_scale: + source_names.append("scale_factor") + + return_type, result_type, out_constraint, tvm_return_type = _result_info(dtype) + signature = ( + "(" + + ", ".join( + f"{c_type} {name}" for name, (c_type, _constraint) in zip(source_names, source_info) + ) + + ")" + ) + + instruction = f"cvt{_modifier_string(form, attrs)}.{dtype}.{atype}" + name = f"tvm_builtin_ptx_{_safe(instruction)}" + placeholders = [f"%{index}" for index in range(1, len(source_names) + 1)] + if form.grouped_primary: + primary = "{" + ", ".join(placeholders[: form.primary_count]) + "}" + asm_operands = [primary, *placeholders[form.primary_count :]] + else: + asm_operands = placeholders + asm_prefix = [] + asm_suffix = [] + result_operand = "%0" + if atype == "e2m1x2": + asm_prefix.extend((".reg .b8 raw_a;", "cvt.u8.u16 raw_a, %1;")) + asm_operands[0] = "raw_a" + if dtype == "e2m1x2": + asm_prefix.append(".reg .b8 raw_result;") + result_operand = "raw_result" + asm_suffix.append("cvt.u16.u8 %0, raw_result;") + operand_text = ", ".join([result_operand, *asm_operands]) + inputs = ", ".join( + f'"{constraint}"({source_name})' + for source_name, (_c_type, constraint) in zip(source_names, source_info) + ) + asm_instructions = [*asm_prefix, f"{instruction} {operand_text};", *asm_suffix] + if asm_prefix or asm_suffix: + asm_text = "\\n\\t".join(asm_instructions) + body = ( + f" {result_type} result;\n" + f' asm volatile("{{\\n\\t{asm_text}\\n\\t}}"\n' + f' : "={out_constraint}"(result)\n' + f" : {inputs});\n" + " return result;" + ) + else: + body = ( + f" {result_type} result;\n" + f' asm volatile("{instruction} {operand_text};"\n' + f' : "={out_constraint}"(result)\n' + f" : {inputs});\n" + " return result;" + ) + return name, signature, return_type, tvm_return_type, body + + +def _register_cvt_form(form_name: str) -> None: + form = _CVT_FORMS[form_name] + device_intrinsic( + f"_ptx_cvt_{form_name}", + n_attrs=len(form.attrs), + helper_name=lambda *a, _form=form_name: _cvt_form_parts(_form, *a)[0], + c_signature=lambda *a, _form=form_name: _cvt_form_parts(_form, *a)[1], + return_type=lambda *a, _form=form_name: _cvt_form_parts(_form, *a)[2], + tvm_return_type=lambda *a, _form=form_name: _cvt_form_parts(_form, *a)[3], + body=lambda *a, _form=form_name: _cvt_form_parts(_form, *a)[4], + ) + + +for _form_name in _CVT_FORMS: + _register_cvt_form(_form_name) +del _form_name + + +def _reject_modifiers(form_name, *, ftz, sat, relu, satfinite, scaled): + if ftz or sat: + raise ValueError(f"PTX cvt form {form_name} does not accept .ftz or .sat") + if relu and "relu" not in _CVT_FORMS[form_name].attrs: + raise ValueError(f"PTX cvt form {form_name} does not accept .relu") + if satfinite and "satfinite" not in _CVT_FORMS[form_name].attrs: + raise ValueError(f"PTX cvt form {form_name} does not accept .satfinite") + if scaled and "scaled" not in _CVT_FORMS[form_name].attrs: + raise ValueError(f"PTX cvt form {form_name} does not accept .scaled") + + +def _classify_cvt_form( + dtype, + atype, + rounding, + ftz, + sat, + relu, + satfinite, + scaled, + n_values, + has_rbits, +): + """Return ``(form_name, form_attrs)`` for one public ``ptx.cvt`` call.""" + + def require_values(expected): + if n_values != expected: + raise ValueError( + f"PTX cvt {dtype}.{atype} expects {expected} primary operands, got {n_values}" + ) + + # f16 / bf16 destination forms from f32. + if atype == "f32" and dtype in {"f16", "f16x2", "bf16", "bf16x2"}: + _reject_modifiers( + "frnd2_f16_f32", + ftz=ftz, + sat=sat, + relu=relu, + satfinite=satfinite, + scaled=scaled, + ) + if rounding == "rs": + if dtype not in {"f16x2", "bf16x2"}: + raise ValueError("PTX cvt.rs from f32 is available only for packed x2 outputs") + require_values(2) + if not has_rbits: + raise ValueError("PTX cvt.rs requires the rbits operand") + form_name = f"rs_{dtype}_f32" + return form_name, [relu, satfinite] + if rounding not in {"rn", "rz"}: + raise ValueError( + "PTX f16/bf16 destination forms require rounding in {'rn', 'rz', 'rs'}" + ) + if has_rbits: + raise ValueError("rbits is valid only with rounding='rs'") + require_values(2 if dtype.endswith("x2") else 1) + form_name = f"frnd2_{dtype}_f32" + return form_name, [rounding, relu, satfinite] + + # tf32 destination forms. + if dtype == "tf32" and atype == "f32": + require_values(1) + if has_rbits or scaled or ftz or sat: + raise ValueError("PTX tf32 destination forms do not accept rbits, scaled, ftz, or sat") + if rounding == "rna": + if relu: + raise ValueError("cvt.rna.tf32.f32 does not accept .relu") + return "rna_tf32_f32", [satfinite] + if rounding not in {"rn", "rz"}: + raise ValueError("PTX tf32 destination forms require rounding in {'rna', 'rn', 'rz'}") + return "frnd2_tf32_f32", [rounding, satfinite, relu] + + # f8/f4/f6 families. + families = ( + ("f8", _F8X2_TYPES, _F8X4_TYPES), + ("f4", _F4X2_TYPES, _F4X4_TYPES), + ("f6", _F6X2_TYPES, _F6X4_TYPES), + ) + for family, x2_types, x4_types in families: + if dtype in x2_types and atype == "f32": + require_values(2) + if rounding != "rn" or not satfinite: + raise ValueError(f"PTX {family}x2.f32 requires rounding='rn' and satfinite=True") + if has_rbits or scaled or ftz or sat: + raise ValueError(f"PTX {family}x2.f32 does not accept rbits, scaled, ftz, or sat") + return f"rn_satfinite_{family}x2_f32", [dtype, relu] + if dtype in x2_types and atype in _FP16X2_TYPES: + require_values(1) + if rounding != "rn" or not satfinite: + raise ValueError( + f"PTX {family}x2.{atype} requires rounding='rn' and satfinite=True" + ) + if has_rbits or scaled or ftz or sat: + raise ValueError( + f"PTX {family}x2.{atype} does not accept rbits, scaled, ftz, or sat" + ) + return f"rn_satfinite_{family}x2_fp16x2", [dtype, atype, relu] + if dtype == "f16x2" and atype in x2_types: + require_values(1) + if rounding != "rn": + raise ValueError(f"PTX f16x2.{family}x2 requires rounding='rn'") + if has_rbits or scaled or ftz or sat or satfinite: + raise ValueError( + f"PTX f16x2.{family}x2 does not accept rbits, scaled, ftz, sat, or satfinite" + ) + return f"rn_f16x2_{family}x2", [atype, relu] + if dtype == "bf16x2" and atype in x2_types: + require_values(1) + if rounding != "rn": + raise ValueError(f"PTX bf16x2.{family}x2 requires rounding='rn'") + if has_rbits or ftz or sat: + raise ValueError(f"PTX bf16x2.{family}x2 does not accept rbits, ftz, or sat") + return f"rn_bf16x2_{family}x2", [atype, relu, satfinite, scaled] + if dtype in x4_types and atype == "f32": + require_values(4) + if rounding != "rs" or not satfinite or not has_rbits: + raise ValueError( + f"PTX {family}x4.f32 requires rounding='rs', satfinite=True, and rbits" + ) + if scaled or ftz or sat: + raise ValueError(f"PTX {family}x4.f32 does not accept scaled, ftz, or sat") + return f"rs_satfinite_{family}x4_f32", [dtype, relu] + + # ue8m0 forms. + if dtype == "ue8m0x2" and atype in {"f32", "bf16x2"}: + require_values(2 if atype == "f32" else 1) + if rounding not in {"rz", "rp"}: + raise ValueError("PTX ue8m0x2 destination forms require rounding in {'rz', 'rp'}") + if has_rbits or scaled or ftz or sat or relu: + raise ValueError("PTX ue8m0x2 destination forms accept only optional .satfinite") + return f"frnd3_ue8m0x2_{atype}", [rounding, satfinite] + if dtype == "bf16x2" and atype == "ue8m0x2": + require_values(1) + if rounding != "rn" or has_rbits or scaled or ftz or sat or relu or satfinite: + raise ValueError("PTX bf16x2.ue8m0x2 accepts only rounding='rn'") + return "rn_bf16x2_ue8m0x2", [] + + # s2f6 forms. + if dtype == "s2f6x2" and atype in {"f32", "bf16x2"}: + require_values(2 if atype == "f32" else 1) + if rounding != "rn" or not satfinite: + raise ValueError(f"PTX s2f6x2.{atype} requires rounding='rn' and satfinite=True") + if has_rbits or ftz or sat: + raise ValueError(f"PTX s2f6x2.{atype} does not accept rbits, ftz, or sat") + return f"rn_satfinite_s2f6x2_{atype}", [relu, scaled] + if dtype == "bf16x2" and atype == "s2f6x2": + require_values(1) + if rounding != "rn" or has_rbits or ftz or sat: + raise ValueError("PTX bf16x2.s2f6x2 requires rounding='rn'") + return "rn_bf16x2_s2f6x2", [satfinite, relu, scaled] + + # Generic scalar grammar forms. + if dtype in _BASIC_TYPES and atype in _BASIC_TYPES: + require_values(1) + if has_rbits or relu or satfinite or scaled: + raise ValueError( + "generic scalar PTX cvt does not accept rbits, relu, satfinite, or scaled" + ) + if rounding in _INTEGER_ROUNDING: + return "irnd", [dtype, atype, rounding, ftz, sat] + if rounding in _FLOAT_ROUNDING: + return "frnd", [dtype, atype, rounding, ftz, sat] + if not rounding: + return "scalar", [dtype, atype, ftz, sat] + raise ValueError(f"Unsupported generic scalar PTX cvt rounding {rounding!r}") + + raise ValueError( + f"No PTX cvt instruction form for dtype={dtype!r}, atype={atype!r}, rounding={rounding!r}" + ) + + +@register_codegen("ptx_cvt") +def codegen_ptx_cvt(*args): + """Classify one public ``ptx.cvt`` invocation into an ISA grammar form.""" + if len(args) < 11: + raise ValueError("ptx.cvt is missing its trailing instruction attrs") + ( + n_values, + has_rbits, + has_scale, + dtype, + atype, + rounding, + ftz, + sat, + relu, + satfinite, + scaled, + ) = args[-11:] + operands = list(args[:-11]) + n_values = int(n_values) + has_rbits = _as_bool(has_rbits) + has_scale = _as_bool(has_scale) + dtype = parse_str(dtype) + atype = parse_str(atype) + rounding = parse_str(rounding) + ftz = _as_bool(ftz) + sat = _as_bool(sat) + relu = _as_bool(relu) + satfinite = _as_bool(satfinite) + scaled = parse_str(scaled) + + if len(operands) != n_values + int(has_rbits) + int(has_scale): + raise ValueError("ptx.cvt operand metadata does not match the provided operands") + if has_rbits != (rounding == "rs"): + raise ValueError("ptx.cvt rbits operand must be present exactly when rounding='rs'") + if has_scale != bool(scaled): + raise ValueError("ptx.cvt scale-factor operand must match the scaled attr") + + form_name, form_attrs = _classify_cvt_form( + dtype, + atype, + rounding, + ftz, + sat, + relu, + satfinite, + scaled, + n_values, + has_rbits, + ) + result = CODEGEN_REGISTRY[f"tirx._ptx_cvt_{form_name}"](operands + form_attrs) + return result[0] if isinstance(result, tuple) else result diff --git a/python/tvm/backend/cuda/operator/intrinsics/math.py b/python/tvm/backend/cuda/operator/intrinsics/math.py index 2786abbad1df..d1559cdc4194 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/math.py +++ b/python/tvm/backend/cuda/operator/intrinsics/math.py @@ -64,16 +64,23 @@ } -def _ptx_arith_modifier_string(dtype, rounding, ftz, sat): +def _ptx_arith_modifier_string(dtype, rounding, ftz, sat, op=None): """Build the `.rnd.ftz.sat` modifier substring + name suffix.""" rnd = parse_str(rounding) - assert rnd in _PACKED_ROUNDING, f"invalid rounding {rnd!r}, expected one of {_PACKED_ROUNDING}" ftz_b = bool(int(ftz)) if hasattr(ftz, "value") else bool(ftz) sat_b = bool(int(sat)) if hasattr(sat, "value") else bool(sat) if dtype == "f64" and (ftz_b or sat_b): raise ValueError("PTX .f64 does not accept .ftz or .sat") if dtype == "f32x2" and sat_b: raise ValueError("PTX .f32x2 does not accept .sat") + if dtype == "f32x2" and op in ("add", "mul") and rnd in ("", "none"): + mod = "" + name_suffix = "" + if ftz_b: + mod += ".ftz" + name_suffix += "_ftz" + return mod, name_suffix + assert rnd in _PACKED_ROUNDING, f"invalid rounding {rnd!r}, expected one of {_PACKED_ROUNDING}" mod = f".{rnd}" if ftz_b: mod += ".ftz" @@ -95,7 +102,7 @@ def _ptx_binary_arith_parts(op, dtype): sig = f"(void* d, {info['c_in']} a, {info['c_in']} b)" def _name(d, a, b, rounding, ftz, sat): - _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat) + _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat, op=op) return f"tvm_builtin_ptx_{op}_{dtype}{suf}" out_c = info["out_cstr"] @@ -103,7 +110,7 @@ def _name(d, a, b, rounding, ftz, sat): out_cast = info["out_cast"] def _body(d, a, b, rounding, ftz, sat): - mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat) + mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat, op=op) return ( f' asm volatile("{op}{mod}.{dtype} %0, %1, %2;"\n' f' : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n' @@ -113,13 +120,82 @@ def _body(d, a, b, rounding, ftz, sat): return _name, sig, _body +_F32X2_VALUE_RETURN_TYPES = { + "uint64": ("unsigned long long", "ret_u64"), + "float32x2": ("float2", "ret_f32x2"), +} + + +def _as_bool_attr(value): + return bool(int(value)) + + +def _f32x2_value_return(return_dtype): + ret = parse_str(return_dtype) + if ret not in _F32X2_VALUE_RETURN_TYPES: + raise ValueError( + f"invalid f32x2 return dtype {ret!r}, expected one of " + f"{tuple(_F32X2_VALUE_RETURN_TYPES)}" + ) + return ret + + +def _ptx_binary_f32x2_parts(op): + """Return helper pieces for ptx_{op}_f32x2 DPS and value forms.""" + + def _name(*args): + rounding, ftz, dps, return_dtype = args[-4:] + _, suf = _ptx_arith_modifier_string("f32x2", rounding, ftz, False, op=op) + if _as_bool_attr(dps): + return f"tvm_builtin_ptx_{op}_f32x2{suf}" + ret = _f32x2_value_return(return_dtype) + _, ret_suffix = _F32X2_VALUE_RETURN_TYPES[ret] + return f"tvm_builtin_ptx_{op}_f32x2_{ret_suffix}{suf}" + + def _sig(*args): + if _as_bool_attr(args[-2]): + return "(void* d, unsigned long long a, unsigned long long b)" + return "(unsigned long long a, unsigned long long b)" + + def _return_type(*args): + if _as_bool_attr(args[-2]): + return "void" + ret = _f32x2_value_return(args[-1]) + return _F32X2_VALUE_RETURN_TYPES[ret][0] + + def _body(*args): + rounding, ftz, dps, return_dtype = args[-4:] + mod, _ = _ptx_arith_modifier_string("f32x2", rounding, ftz, False, op=op) + if _as_bool_attr(dps): + return ( + f' asm volatile("{op}{mod}.f32x2 %0, %1, %2;"\n' + ' : "=l"(*reinterpret_cast(d))\n' + ' : "l"(a), "l"(b));' + ) + ret = _f32x2_value_return(return_dtype) + ret_stmt = ( + " return result;" + if ret == "uint64" + else " return *reinterpret_cast(&result);" + ) + return ( + " unsigned long long result;\n" + f' asm volatile("{op}{mod}.f32x2 %0, %1, %2;"\n' + ' : "=l"(result)\n' + ' : "l"(a), "l"(b));\n' + f"{ret_stmt}" + ) + + return _name, _sig, _return_type, _body + + def _ptx_fma_parts(dtype): """Return (name_fn, sig, body_fn) for ptx_fma_{dtype}.""" info = _DTYPE_INFO[dtype] sig = f"(void* d, {info['c_in']} a, {info['c_in']} b, {info['c_in']} c)" def _name(d, a, b, c, rounding, ftz, sat): - _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat) + _, suf = _ptx_arith_modifier_string(dtype, rounding, ftz, sat, op="fma") return f"tvm_builtin_ptx_fma_{dtype}{suf}" out_c = info["out_cstr"] @@ -127,7 +203,7 @@ def _name(d, a, b, c, rounding, ftz, sat): out_cast = info["out_cast"] def _body(d, a, b, c, rounding, ftz, sat): - mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat) + mod, _ = _ptx_arith_modifier_string(dtype, rounding, ftz, sat, op="fma") return ( f' asm volatile("fma{mod}.{dtype} %0, %1, %2, %3;"\n' f' : "={out_c}"(*reinterpret_cast<{out_cast}>(d))\n' @@ -137,8 +213,8 @@ def _body(d, a, b, c, rounding, ftz, sat): return _name, sig, _body -# Register 12 ops: {add, sub, mul, fma} x {f32, f32x2, f64}. -for _dtype in ("f32", "f32x2", "f64"): +# Register scalar DPS ops: {add, sub, mul, fma} x {f32, f64}. +for _dtype in ("f32", "f64"): for _op in ("add", "sub", "mul"): _name_fn, _sig, _body_fn = _ptx_binary_arith_parts(_op, _dtype) device_intrinsic( @@ -159,6 +235,79 @@ def _body(d, a, b, c, rounding, ftz, sat): del _dtype, _op, _name_fn, _sig, _body_fn +for _op in ("add", "sub", "mul"): + _name_fn, _sig_fn, _ret_fn, _body_fn = _ptx_binary_f32x2_parts(_op) + device_intrinsic( + f"ptx_{_op}_f32x2", + n_attrs=4, # rounding, ftz, dps, return_dtype + helper_name=_name_fn, + c_signature=_sig_fn, + body=_body_fn, + return_type=_ret_fn, + ) + + +def _ptx_fma_f32x2_parts(): + """Return helper pieces for ptx_fma_f32x2 DPS and value forms.""" + + def _name(*args): + rounding, ftz, dps, return_dtype = args[-4:] + _, suf = _ptx_arith_modifier_string("f32x2", rounding, ftz, False, op="fma") + if _as_bool_attr(dps): + return f"tvm_builtin_ptx_fma_f32x2{suf}" + ret = _f32x2_value_return(return_dtype) + _, ret_suffix = _F32X2_VALUE_RETURN_TYPES[ret] + return f"tvm_builtin_ptx_fma_f32x2_{ret_suffix}{suf}" + + def _sig(*args): + if _as_bool_attr(args[-2]): + return "(void* d, unsigned long long a, unsigned long long b, unsigned long long c)" + return "(unsigned long long a, unsigned long long b, unsigned long long c)" + + def _return_type(*args): + if _as_bool_attr(args[-2]): + return "void" + ret = _f32x2_value_return(args[-1]) + return _F32X2_VALUE_RETURN_TYPES[ret][0] + + def _body(*args): + rounding, ftz, dps, return_dtype = args[-4:] + mod, _ = _ptx_arith_modifier_string("f32x2", rounding, ftz, False, op="fma") + if _as_bool_attr(dps): + return ( + f' asm volatile("fma{mod}.f32x2 %0, %1, %2, %3;"\n' + ' : "=l"(*reinterpret_cast(d))\n' + ' : "l"(a), "l"(b), "l"(c));' + ) + ret = _f32x2_value_return(return_dtype) + ret_stmt = ( + " return result;" + if ret == "uint64" + else " return *reinterpret_cast(&result);" + ) + return ( + " unsigned long long result;\n" + f' asm volatile("fma{mod}.f32x2 %0, %1, %2, %3;"\n' + ' : "=l"(result)\n' + ' : "l"(a), "l"(b), "l"(c));\n' + f"{ret_stmt}" + ) + + return _name, _sig, _return_type, _body + + +_name_fn, _sig_fn, _ret_fn, _body_fn = _ptx_fma_f32x2_parts() +device_intrinsic( + "ptx_fma_f32x2", + n_attrs=4, # rounding, ftz, dps, return_dtype + helper_name=_name_fn, + c_signature=_sig_fn, + body=_body_fn, + return_type=_ret_fn, +) +del _op, _name_fn, _sig_fn, _ret_fn, _body_fn + + # ============================================================================= # ex2.approx.ftz.f32 / rcp.approx.ftz.f32 — 1 form each. # ============================================================================= @@ -183,6 +332,15 @@ def _body(d, a, b, c, rounding, ftz, sat): ), ) +device_intrinsic( + "cuda_fdividef", + helper_name="tvm_builtin_cuda_fdividef", + c_signature="(float x, float y)", + return_type="float", + tvm_return_type="float32", + body=" return __fdividef(x, y);", +) + # ============================================================================= # 3-operand max.f32 / min.f32 — the f32, 3-operand form-table entry of the diff --git a/python/tvm/backend/cuda/operator/intrinsics/memory.py b/python/tvm/backend/cuda/operator/intrinsics/memory.py index 4b66ed26a25f..9a39c5fe3d05 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/memory.py +++ b/python/tvm/backend/cuda/operator/intrinsics/memory.py @@ -38,23 +38,81 @@ from .registry import CODEGEN_REGISTRY, register_codegen from .utils import parse_str +# __ldg — typed read-only cached load. Source is ``void*`` so callers may pass +# a typed pointer or handle_add_byte_offset result; the helper casts per ``dtype``. +_CUDA_LDG_CTYPES = { + "int8": "signed char", + "uint8": "unsigned char", + "int16": "short", + "uint16": "unsigned short", + "int32": "int", + "uint32": "unsigned int", + "int64": "long long", + "uint64": "unsigned long long", + "float16": "half", + "bfloat16": "nv_bfloat16", + "float32": "float", + "float64": "double", +} + +_CUDA_LDG_VECTOR_CTYPES = { + "int32": "int", + "uint32": "unsigned int", + "float32": "float", +} +_CUDA_LDG_VECTOR_BASES = {"int32": "int", "uint32": "uint", "float32": "float"} + + +def _cuda_ldg_suffix(dtype: str) -> str: + return dtype.replace("float", "f").replace("uint", "u").replace("int", "i") + -# ============================================================================= -# __ldg — templated read-only cached load; ``T`` resolved at call time from -# the ``dtype`` argument. Hand-written because the helper signature uses a -# template parameter for both arg and return. -# ============================================================================= @register_codegen("cuda_ldg") -def codegen_cuda_ldg(addr, dtype): - dtype = DataType(parse_str(dtype)) - func_name = "tvm_builtin_cuda_ldg" +def codegen_cuda_ldg(*args): + if len(args) == 2: + addr, dtype = args + dtype = str(DataType(parse_str(dtype))) + if dtype not in _CUDA_LDG_CTYPES: + raise ValueError(f"Unsupported CUDA __ldg dtype {dtype!r}") + c_type = _CUDA_LDG_CTYPES[dtype] + func_name = f"tvm_builtin_cuda_ldg_{_cuda_ldg_suffix(dtype)}" + source_code = f""" +__forceinline__ __device__ {c_type} {func_name}(void* src) {{ + return __ldg(reinterpret_cast(src)); +}} +""" + return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype) + + if len(args) < 5: + raise ValueError(f"cuda_ldg expects 2 args or vector form, got {len(args)}") + *dsts, addr, dtype, vec, dst_count = args + dtype = str(DataType(parse_str(dtype))) + vec = parse_str(vec) + dst_count = _int_attr(dst_count) + vec_len = int(vec[1:]) if vec else 1 + if dtype not in _CUDA_LDG_VECTOR_CTYPES: + raise ValueError(f"Unsupported vector CUDA __ldg dtype {dtype!r}") + if vec not in ("v2", "v4") or dst_count != vec_len or len(dsts) != vec_len: + raise ValueError( + f"vector CUDA __ldg expects dst_count=len(dsts)=vec_len for v2/v4, " + f"got vec={vec!r}, dst_count={dst_count}, len(dsts)={len(dsts)}" + ) + c_type = _CUDA_LDG_VECTOR_CTYPES[dtype] + vec_type = f"{_CUDA_LDG_VECTOR_BASES[dtype]}{vec_len}" + members = ("x", "y", "z", "w")[:vec_len] + func_name = f"tvm_builtin_cuda_ldg_{_cuda_ldg_suffix(dtype)}_{vec}_to_dst{dst_count}" + params = ", ".join(f"void* dst{i}" for i in range(vec_len)) + stores = "\n".join( + f" *reinterpret_cast<{c_type}*>(dst{i}) = v.{member};" + for i, member in enumerate(members) + ) source_code = f""" -template -__forceinline__ __device__ T {func_name}(T* src) {{ - return __ldg(src); +__forceinline__ __device__ void {func_name}({params}, void* src) {{ + {vec_type} v = __ldg(reinterpret_cast(src)); +{stores} }} """ - return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype) + return cuda_func_call(func_name, *dsts, addr, source_code=source_code, return_type="void") # Shared PTX scalar type metadata (ld/st/red/atom). @@ -85,6 +143,35 @@ def codegen_cuda_ldg(addr, dtype): 2: "unsigned short", 1: "unsigned char", } +_PTX_ST_SRC_CTYPES = { + "b8": "unsigned char", + "u8": "unsigned char", + "s8": "signed char", + "b16": "unsigned short", + "u16": "unsigned short", + "s16": "short", + "b32": "unsigned int", + "u32": "unsigned int", + "s32": "int", + "b64": "unsigned long long", + "u64": "unsigned long long", + "s64": "long long", + "f32": "float", + "f64": "double", +} +_PTX_ST_SRC_VECTOR_CTYPES = { + "b8": "uchar", + "u8": "uchar", + "s8": "char", + "b16": "ushort", + "u16": "ushort", + "s16": "short", + "b32": "uint", + "u32": "uint", + "s32": "int", + "f32": "float", + "f64": "double", +} def _safe_attr(value): @@ -113,6 +200,8 @@ def _type_info(ptx_type): # PTX ld forms (ISA table entries registered via ``ptx_ld`` and siblings): # ld{.weak}{.ss}{.cop}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; # ld{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; +# ld.global{.cop}.nc{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; +# ld.global.nc{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; # ld.volatile{.ss}{.level::prefetch_size}{.vec}.type d, [a]; # ld.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; # ld.acquire.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.level::prefetch_size}{.vec}.type d, [a]{, cache-policy}; @@ -123,6 +212,7 @@ def _type_info(ptx_type): _PTX_LD_VOLATILE_SPACES = _PTX_LD_SPACES | {"const"} _PTX_LD_WEAK_SPACES = _PTX_LD_SPACES | {"const", "param::entry", "param::func"} _PTX_LD_COPS = {"", "ca", "cg", "cs", "lu", "cv"} +_PTX_LD_GLOBAL_NC_COPS = {"", "ca", "cg", "cs"} _PTX_VEC = {"", "v2", "v4", "v8"} _PTX_L1_EVICT = { "", @@ -140,6 +230,10 @@ def _bool_attr(value): return bool(int(value)) if hasattr(value, "value") else bool(value) +def _int_attr(value): + return int(value.value) if hasattr(value, "value") else int(value) + + def _parse_ld_attrs(return_dtype, ptx_type, scope=None, space="global"): return_dtype = parse_str(return_dtype) ptx_type = parse_str(ptx_type) @@ -189,6 +283,21 @@ def _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size): return suffix +def _ptx_global_nc_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size): + suffix = "" + l1_evict = parse_str(l1_evict) + l2_evict = parse_str(l2_evict) + prefetch_size = parse_str(prefetch_size) + if l1_evict: + suffix += f".{l1_evict}" + if l2_evict: + suffix += f".{l2_evict}" + suffix += _cache_suffix("cache" if has_cache else "") + if prefetch_size: + suffix += f".{prefetch_size}" + return suffix + + def _ptx_shared_addr(space, ptr_name="address"): if parse_str(space).startswith("shared"): return ( @@ -198,21 +307,41 @@ def _ptx_shared_addr(space, ptr_name="address"): return "", f'"l"({ptr_name})' -def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type): +def _ptx_addr_operand(space, ptr_name, raw_address): + if raw_address: + return "", f'"r"({ptr_name})' + return _ptx_shared_addr(space, ptr_name) + + +def _is_raw_u32_address(address): + return str(address.ty) == "uint32" + + +def _ptx_ld_signature(dst_count, vec_len, raw_address): + address_type = "unsigned int" if raw_address else "void*" + if dst_count == 1: + return f"(void* dst_ptr, {address_type} src_ptr, unsigned long long cache_policy)" + if dst_count == vec_len: + dst_params = ", ".join(f"void* dst{i}" for i in range(vec_len)) + return f"({dst_params}, {address_type} src_ptr, unsigned long long cache_policy)" + raise ValueError(f"PTX ld dst count must be 1 or vec_len={vec_len}, got {dst_count}") + + +def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type, c_type, dst_count): if ptx_type == "u8" and vec_len == 1: return " *reinterpret_cast(dst_ptr) = static_cast(r0);" - store_type = _PTX_VEC_STORE_TYPE[num_bytes] - if vec_len > 1: - return ( - f" *reinterpret_cast<{store_type}*>(dst_ptr) = " - + "{" - + ", ".join(f"r{i}" for i in range(vec_len)) - + "};" + if dst_count == vec_len and vec_len > 1: + return "\n".join( + f" *reinterpret_cast<{c_type}*>(dst{i}) = r{i};" for i in range(vec_len) ) + if vec_len > 1: + stores = "".join(f" dst[{i}] = r{i};\n" for i in range(vec_len)) + return f" {c_type}* dst = reinterpret_cast<{c_type}*>(dst_ptr);\n{stores}".rstrip() + store_type = _PTX_VEC_STORE_TYPE[num_bytes] return f" *reinterpret_cast<{store_type}*>(dst_ptr) = r0;" -def _ptx_ld_form_parts(form, attr_args): +def _ptx_ld_form_parts(form, attr_args, raw_address): if form == "weak": ( return_dtype, @@ -264,6 +393,19 @@ def _ptx_ld_form_parts(form, attr_args): return_dtype, sem, scope, space, ptx_type, to_dst = attr_args weak, cop, vec = False, "", "" has_cache_hint, l1_evict, l2_evict, prefetch_size = False, "", "", "" + elif form == "global_nc": + ( + return_dtype, + cop, + vec, + ptx_type, + has_cache_hint, + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) = attr_args + sem, scope, weak, space = "", "", False, "global" else: raise ValueError(f"unknown ld form {form!r}") @@ -280,8 +422,11 @@ def _ptx_ld_form_parts(form, attr_args): prefetch_size = parse_str(prefetch_size) weak = _bool_attr(weak) has_cache = _bool_attr(has_cache_hint) - to_dst = _bool_attr(to_dst) - if cop and cop not in _PTX_LD_COPS: + to_dst = _int_attr(to_dst) + if form == "global_nc": + if cop and cop not in _PTX_LD_GLOBAL_NC_COPS: + raise ValueError(f"Unsupported PTX ld.global.nc cache operation {cop!r}") + elif cop and cop not in _PTX_LD_COPS: raise ValueError(f"Unsupported PTX ld cache operation {cop!r}") if vec and vec not in _PTX_VEC: raise ValueError(f"Unsupported PTX ld vector modifier {vec!r}") @@ -295,6 +440,10 @@ def _ptx_ld_form_parts(form, attr_args): if sem not in ("acquire", "relaxed") or scope != "sys" or space != "global": raise ValueError("ld.mmio requires sem in {acquire, relaxed}, scope=sys, space=global") prefix = f"ld.mmio.{sem}.{scope}" + elif form == "global_nc": + if cop and (l1_evict or l2_evict): + raise ValueError("ld.global.nc with cop cannot use l1_evict or l2_evict") + prefix = f"ld.global{_dot(cop)}.nc" elif form == "relaxed": if not scope: raise ValueError("ld.relaxed requires scope") @@ -311,7 +460,11 @@ def _ptx_ld_form_parts(form, attr_args): else: _validate_ld_space(space, _PTX_LD_WEAK_SPACES) prefix = f"ld{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}" - level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size) + level = ( + _ptx_global_nc_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size) + if form == "global_nc" + else _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size) + ) vec_len = int(vec[1:]) if vec else 1 if vec and not to_dst: raise ValueError("vector ld requires to_dst") @@ -325,23 +478,32 @@ def _ptx_ld_form_parts(form, attr_args): else 4 ) num_bytes = vec_len * elem_bytes if vec else elem_bytes - name_parts = [ - "tvm_builtin_ptx_ld", - form if form != "weak" else ("weak" if weak else "plain"), - ] - if sem: - name_parts.append(_safe_attr(sem)) - if scope: - name_parts.append(_safe_attr(scope)) - name_parts.extend( - [ - _safe_attr(space), + if form == "global_nc": + name_parts = [ + "tvm_builtin_ptx_ld_global_nc", _safe_attr(cop) if cop else "", _safe_attr(vec) if vec else "", ptx_type, - return_dtype if not to_dst else "to_dst", + return_dtype if not to_dst else ("to_dst" if to_dst == 1 else f"to_dst{to_dst}"), ] - ) + else: + name_parts = [ + "tvm_builtin_ptx_ld", + form if form != "weak" else ("weak" if weak else "plain"), + ] + if sem: + name_parts.append(_safe_attr(sem)) + if scope: + name_parts.append(_safe_attr(scope)) + name_parts.extend( + [ + _safe_attr(space), + _safe_attr(cop) if cop else "", + _safe_attr(vec) if vec else "", + ptx_type, + return_dtype if not to_dst else ("to_dst" if to_dst == 1 else f"to_dst{to_dst}"), + ] + ) if has_cache: name_parts.append("cache_hint") if l1_evict: @@ -350,10 +512,16 @@ def _ptx_ld_form_parts(form, attr_args): name_parts.append(_safe_attr(l2_evict)) if prefetch_size: name_parts.append(_safe_attr(prefetch_size)) + if raw_address: + name_parts.append("raw_u32") name = "_".join(p for p in name_parts if p) cache_operand = ', "l"(cache_policy)' if has_cache else "" - addr_decl, addr_operand = _ptx_shared_addr(space, "src_ptr" if to_dst else "address") + addr_decl, addr_operand = _ptx_addr_operand( + space, "src_ptr" if to_dst else "address", raw_address + ) if to_dst: + if to_dst not in (1, vec_len): + raise ValueError(f"PTX ld dst count must be 1 or vec_len={vec_len}, got {to_dst}") reg_decls = "".join(f" {c_type} r{i};\n" for i in range(vec_len)) if vec_len > 1: out_slot = "{" + ", ".join(f"%{i}" for i in range(vec_len)) + "}" @@ -370,11 +538,11 @@ def _ptx_ld_form_parts(form, attr_args): f' asm volatile("{instr} {out_slot}, [%{addr_idx}]{cache_slot};"\n' f" : {out_constraints}\n" f" : {addr_operand}{cache_operand});\n" - f"{_ptx_ld_vec_store(num_bytes, vec_len, ptx_type)}" + f"{_ptx_ld_vec_store(num_bytes, vec_len, ptx_type, c_type, to_dst)}" ) return ( name, - "(void* dst_ptr, void* src_ptr, unsigned long long cache_policy)", + _ptx_ld_signature(to_dst, vec_len, raw_address), "void", "", body, @@ -389,15 +557,22 @@ def _ptx_ld_form_parts(form, attr_args): f" : {addr_operand}{cache_operand});\n" " return ret;" ) + address_type = "unsigned int" if raw_address else "void*" sig = ( - "(void* address, unsigned long long cache_policy)" if form == "weak" else "(void* address)" + f"({address_type} address, unsigned long long cache_policy)" + if form in ("weak", "global_nc") + else f"({address_type} address)" ) return name, sig, c_type, return_dtype, body def _register_ptx_ld(op_name, form, n_attrs): def _parts(*args): - return _ptx_ld_form_parts(form, args[-n_attrs:]) + attrs = args[-n_attrs:] + forward = args[:-n_attrs] + # The weak ld form always forwards (..., address, cache_policy). + raw_address = form == "weak" and _is_raw_u32_address(forward[-2]) + return _ptx_ld_form_parts(form, attrs, raw_address) device_intrinsic( op_name, @@ -413,6 +588,7 @@ def _parts(*args): _register_ptx_ld("ptx_ld", "weak", 11) +_register_ptx_ld("ptx_ld_global_nc", "global_nc", 9) _register_ptx_ld("ptx_ld_relaxed", "relaxed", 10) _register_ptx_ld("ptx_ld_acquire", "acquire", 10) _register_ptx_ld("ptx_ld_volatile", "volatile", 6) @@ -740,18 +916,28 @@ def _prefetch_tensormap_parts(_tensor_map, tensormap_space): _PTX_ST_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local", "param::func"} -def _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type): - if ptx_type == "u8" and vec_len == 1: - return " unsigned int r0 = *reinterpret_cast(src_ptr);\n" - store_type = _PTX_VEC_STORE_TYPE[num_bytes] +def _ptx_st_load_src(_num_bytes, vec_len, ptx_type, c_type): + src_c_type = _PTX_ST_SRC_CTYPES[ptx_type] + + def _assign(dst, src): + if src_c_type == c_type: + return f" {c_type} {dst} = {src};\n" + return f" {c_type} {dst} = static_cast<{c_type}>({src});\n" + if vec_len > 1: - return f" {store_type} src_ = *reinterpret_cast<{store_type}*>(src_ptr);\n" + "".join( - f" {c_type} r{i} = src_.{c};\n" for i, c in enumerate("xyzw"[:vec_len]) - ) - return f" {c_type} r0 = *reinterpret_cast<{c_type}*>(src_ptr);\n" + vec_base = _PTX_ST_SRC_VECTOR_CTYPES.get(ptx_type) + if vec_base is not None and vec_len in (2, 3, 4): + vec_type = f"{vec_base}{vec_len}" + loads = "".join(_assign(f"r{i}", f"src_.{c}") for i, c in enumerate("xyzw"[:vec_len])) + return f" {vec_type} src_ = *reinterpret_cast<{vec_type}*>(src_ptr);\n{loads}" + + loads = "".join(_assign(f"r{i}", f"src[{i}]") for i in range(vec_len)) + return f" {src_c_type}* src = reinterpret_cast<{src_c_type}*>(src_ptr);\n{loads}" + return _assign("r0", f"*reinterpret_cast<{src_c_type}*>(src_ptr)") -def _ptx_st_form_parts(form, attr_args, from_src): + +def _ptx_st_form_parts(form, attr_args, from_src, raw_address): if form == "weak": weak, space, cop, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args sem, scope = "", "" @@ -781,6 +967,45 @@ def _ptx_st_form_parts(form, attr_args, from_src): l2_evict = parse_str(l2_evict) weak = _bool_attr(weak) has_cache = _bool_attr(has_cache_hint) + ptx_type = parse_str(ptx_type) + if ptx_type == "b128": + # The b128 form mirrors the FlashMLA dequantization store exactly. + # Its source pointer materializes one scalar 128-bit register for the + # PTX ``q`` constraint. Unlike generic scalar/vector stores, this is + # weak shared::cta only, accepts no cache policy, and deliberately + # omits a memory clobber to preserve the source dependency shape. + if not ( + form == "weak" + and weak + and space == "shared::cta" + and not cop + and not vec + and from_src + and not has_cache + and not l1_evict + and not l2_evict + ): + raise ValueError( + "PTX st.b128 requires src=, weak=True, space='shared::cta', " + "and no cache/vector modifiers" + ) + name = "tvm_builtin_ptx_st_weak_shared_cta_b128_from_src" + if raw_address: + name += "_raw_u32" + address_type = "unsigned int" if raw_address else "void*" + addr_decl, addr_operand = _ptx_addr_operand(space, "address", raw_address) + body = ( + f"{addr_decl}" + " unsigned __int128 value = *reinterpret_cast(src_ptr);\n" + ' asm volatile("st.weak.shared::cta.b128 [%0], %1;"\n' + " :\n" + f' : {addr_operand}, "q"(value));' + ) + return ( + name, + (f"({address_type} address, void* src_ptr, unsigned long long cache_policy)"), + body, + ) ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type) if cop and cop not in _PTX_ST_COPS: raise ValueError(f"Unsupported PTX st cache operation {cop!r}") @@ -834,12 +1059,15 @@ def _ptx_st_form_parts(form, attr_args, from_src): ) if has_cache: name_parts.append("cache_hint") + if raw_address: + name_parts.append("raw_u32") name = "_".join(p for p in name_parts if p) values = f"{{{', '.join(f'%{i + 1}' for i in range(vec_len))}}}" if vec_len > 1 else "%1" value_constraints = "".join(f', "{constraint}"(value{i})' for i in range(vec_len)) cache_slot = f", %{vec_len + 1}" if has_cache else "" cache_operand = ', "l"(cache_policy)' if has_cache else "" - addr_decl, addr_operand = _ptx_shared_addr(space, "address") + addr_decl, addr_operand = _ptx_addr_operand(space, "address", raw_address) + address_type = "unsigned int" if raw_address else "void*" if from_src: load_regs = _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type) if vec_len > 1: @@ -855,9 +1083,9 @@ def _ptx_st_form_parts(form, attr_args, from_src): ' : "memory");' ) if use_cache_policy: - sig = "(void* address, void* src_ptr, unsigned long long cache_policy)" + sig = f"({address_type} address, void* src_ptr, unsigned long long cache_policy)" else: - sig = "(void* address, void* src_ptr)" + sig = f"({address_type} address, void* src_ptr)" else: body = ( f"{addr_decl}" @@ -867,22 +1095,23 @@ def _ptx_st_form_parts(form, attr_args, from_src): ' : "memory");' ) if form == "mmio": - sig = f"(void* address, {c_type} value0)" + sig = f"({address_type} address, {c_type} value0)" elif use_cache_policy: value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len)) - sig = f"(void* address, {value_params}, unsigned long long cache_policy)" + sig = f"({address_type} address, {value_params}, unsigned long long cache_policy)" else: value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len)) - sig = f"(void* address, {value_params})" + sig = f"({address_type} address, {value_params})" return name, sig, body -def _register_ptx_st(op_name, form, n_attrs, *, with_cache_policy=True): +def _register_ptx_st(op_name, form, n_attrs): def codegen(*args): from_src = _bool_attr(args[-1]) st_attrs = args[-n_attrs:-1] - parts = _ptx_st_form_parts(form, st_attrs, from_src) forward = args[:-(n_attrs)] + raw_address = form == "weak" and _is_raw_u32_address(forward[0]) + parts = _ptx_st_form_parts(form, st_attrs, from_src, raw_address) name, sig, body_str = parts source_code = f"\n__forceinline__ __device__ void {name}{sig} {{\n{body_str}\n}}\n" return cuda_func_call(name, *forward, source_code=source_code) diff --git a/python/tvm/backend/cuda/operator/intrinsics/misc.py b/python/tvm/backend/cuda/operator/intrinsics/misc.py index 9a664a815893..b6bbb9780854 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/misc.py +++ b/python/tvm/backend/cuda/operator/intrinsics/misc.py @@ -193,6 +193,23 @@ def _write_event(event_bits: str) -> str: ) +# ============================================================================= +# Official IKET NativeDump placeholder. +# ============================================================================= +@register_codegen("cuda_iket_official_event") +def codegen_cuda_iket_official_event(event_id, source_code): + if isinstance(source_code, tvm.tirx.StringImm): + source_code = source_code.value + else: + source_code = parse_str(source_code) + return cuda_func_call( + "tvm_builtin_iket_official_event", + event_id, + source_code=source_code, + return_type="uint32", + ) + + # ============================================================================= # Debug helpers — ``printf`` (variadic templated) and ``trap`` on assert. # ============================================================================= diff --git a/python/tvm/backend/cuda/operator/intrinsics/sync.py b/python/tvm/backend/cuda/operator/intrinsics/sync.py index 791d9cc981fc..fd5f6e742dc1 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/sync.py +++ b/python/tvm/backend/cuda/operator/intrinsics/sync.py @@ -18,7 +18,8 @@ """Synchronization primitives. PTX side: -* ``bar.arrive`` / ``bar.sync`` — named-barrier alias of ``barrier.arrive/sync`` +* ``bar.arrive`` / ``bar.sync`` — aligned named-barrier aliases +* ``barrier.sync`` — unaligned named barrier for divergent control flow * ``fence{.sem}.scope`` / ``fence.proxy.async`` / ``fence.mbarrier_init`` * ``barrier.cluster.arrive`` / ``barrier.cluster.wait`` * ``mbarrier.init`` / ``mbarrier.arrive[.expect_tx]`` (local + remote) / ``mbarrier.try_wait`` @@ -37,16 +38,32 @@ FENCE_PROXY_ASYNC_SPACE, FENCE_SCOPE, FENCE_SEM, + MBARRIER_ARRIVE_SCOPE, + MBARRIER_ARRIVE_SEM, + MBARRIER_ARRIVE_SPACE, + MBARRIER_COMPLETE_TX_SCOPE, + MBARRIER_COMPLETE_TX_SEM, + MBARRIER_COMPLETE_TX_SPACE, ) from ._schema import device_intrinsic -from .registry import CODEGEN_REGISTRY, register_codegen from .utils import parse_str + +def _safe_attr(value): + return parse_str(value).replace("::", "_").replace(".", "_") + + +def _as_bool(value) -> bool: + return bool(int(value)) if hasattr(value, "value") else bool(value) + + # ============================================================================= -# bar.arrive / bar.sync — alias of barrier.arrive/sync. 1 form each. +# bar.arrive / bar.sync — aligned named-barrier aliases. 1 form each. # bar.sync a, b ; # bar.arrive a, b ; +# barrier.sync — unaligned named barrier. 1 form. +# barrier.sync a, b ; # ============================================================================= device_intrinsic( "ptx_bar_arrive", @@ -62,6 +79,14 @@ ' asm volatile("bar.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory");' ), ) +device_intrinsic( + "ptx_barrier_sync", + c_signature="(int name_bar_id, int thread_count)", + body=( + ' asm volatile("barrier.sync %0, %1;" : : ' + '"r"(name_bar_id), "r"(thread_count) : "memory");' + ), +) # ============================================================================= @@ -191,19 +216,20 @@ def _ptx_barrier_cluster_wait(acquire, aligned): ) -device_intrinsic( - "ptx_clc_query_cancel", - c_signature="(void* handle)", - return_type="uint32_t", - tvm_return_type="uint32", - body=( +def _ptx_clc_query_cancel_parts(use_ld_acquire): + use_ld_acquire = ( + bool(int(use_ld_acquire)) if hasattr(use_ld_acquire, "value") else bool(use_ld_acquire) + ) + name = f"tvm_builtin_ptx_clc_query_cancel{'_ld_acquire' if use_ld_acquire else ''}" + load_instr = "ld.acquire.cta.shared.b128" if use_ld_acquire else "ld.shared.b128" + body = ( " unsigned int addr = (unsigned int)__cvta_generic_to_shared(handle);\n" " unsigned int first_ctaid_x;\n" " asm volatile(\n" ' "{\\n"\n' ' ".reg .pred canceled;\\n"\n' ' ".reg .b128 response;\\n"\n' - ' "ld.shared.b128 response, [%1];\\n"\n' + f' "{load_instr} response, [%1];\\n"\n' ' "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 canceled, response;\\n"\n' ' "mov.u32 %0, 0xffffffff;\\n"\n' ' "@canceled clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128"\n' @@ -212,7 +238,18 @@ def _ptx_barrier_cluster_wait(acquire, aligned): ' : "=r"(first_ctaid_x) : "r"(addr) : "memory");\n' ' asm volatile("fence.proxy.async.shared::cta;\\n" ::: "memory");\n' " return first_ctaid_x;" - ), + ) + return name, body + + +device_intrinsic( + "ptx_clc_query_cancel", + n_attrs=1, + helper_name=lambda *args: _ptx_clc_query_cancel_parts(args[-1])[0], + c_signature="(void* handle)", + return_type="uint32_t", + tvm_return_type="uint32", + body=lambda *args: _ptx_clc_query_cancel_parts(args[-1])[1], ) @@ -231,122 +268,294 @@ def _ptx_barrier_cluster_wait(acquire, aligned): # ============================================================================= -# mbarrier.arrive — local + remote (cluster-mapped) forms. 2 PTX forms. -# Form local: mbarrier.arrive.shared.b64 _, [bar]; -# Form remote: { setp+@p mapa.shared::cluster.u32 + @p mbarrier.arrive.shared::cluster.b64 } -# Dispatcher picks by arg count (1 vs 3). +# mbarrier.arrive{.sem.scope}{.space}.b64 _, [addr]{, count}; # ============================================================================= -device_intrinsic( - "_ptx_mbarrier_arrive_local", - helper_name="tvm_builtin_ptx_mbarrier_arrive", - c_signature="(void* barrier)", - body=( - " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" - ' asm volatile("mbarrier.arrive.shared.b64 _, [%0];"\n' - ' :: "r"(barrier_addr) : "memory");' - ), -) -device_intrinsic( - "_ptx_mbarrier_arrive_remote", - helper_name="tvm_builtin_ptx_mbarrier_arrive_remote", - c_signature="(void* barrier, int cta_id, int pred)", - body=( - " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" - " asm volatile(\n" - ' "{\\n"\n' - ' ".reg .pred p;\\n"\n' - ' ".reg .b32 remAddr32;\\n"\n' - ' "setp.ne.s32 p, %2, 0;\\n"\n' - ' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n' - ' "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];\\n"\n' - ' "}\\n"\n' - ' :: "r"(barrier_addr), "r"(cta_id), "r"(pred) : "memory");' - ), -) +def _check_mbarrier_arrive_attrs(sem, scope, space): + sem = parse_str(sem) + scope = parse_str(scope) + space = parse_str(space) + if (sem == "") != (scope == ""): + raise ValueError("mbarrier.arrive sem and scope must be specified together") + if sem not in MBARRIER_ARRIVE_SEM: + raise ValueError(f"invalid mbarrier.arrive sem {sem!r}") + if scope not in MBARRIER_ARRIVE_SCOPE: + raise ValueError(f"invalid mbarrier.arrive scope {scope!r}") + if space not in MBARRIER_ARRIVE_SPACE: + raise ValueError(f"invalid mbarrier.arrive space {space!r}") + return sem, scope, space + + +def _ptx_mbarrier_arrive_parts(*args): + sem, scope, space, has_count, has_remote, has_pred = args[-6:] + sem, scope, space = _check_mbarrier_arrive_attrs(sem, scope, space) + has_count = _as_bool(has_count) + has_remote = _as_bool(has_remote) + has_pred = _as_bool(has_pred) + if has_remote and space != "shared::cluster": + raise ValueError("remote mbarrier.arrive requires space='shared::cluster'") + + name = "tvm_builtin_ptx_mbarrier_arrive" + if sem: + name += f"_{_safe_attr(sem)}_{_safe_attr(scope)}" + name += f"_{_safe_attr(space)}" + if has_count: + name += "_count" + if has_remote: + name += "_remote" + if has_pred: + name += "_pred" + + params = ["void* barrier"] + arg_idx = 1 + if has_count: + params.append("int count") + count_idx = arg_idx + arg_idx += 1 + if has_remote: + params.append("int remote") + remote_idx = arg_idx + arg_idx += 1 + if has_pred: + params.append("int pred") + pred_idx = arg_idx + + instr_suffix = f".{sem}.{scope}" if sem else "" + instr = f"mbarrier.arrive{instr_suffix}.{space}.b64" + body = " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" + + if has_remote or has_pred: + body += " asm volatile(\n" + body += ' "{\\n"\n' + if has_pred: + body += ' ".reg .pred p;\\n"\n' + if has_remote: + body += ' ".reg .b32 remAddr32;\\n"\n' + if has_pred: + body += f' "setp.ne.s32 p, %{pred_idx}, 0;\\n"\n' + if has_remote: + prefix = "@p " if has_pred else "" + body += ( + f' "{prefix}mapa.shared::cluster.u32 remAddr32, %0, %{remote_idx};\\n"\n' + ) + addr = "remAddr32" + else: + addr = "%0" + pred_prefix = "@p " if has_pred else "" + count_suffix = f", %{count_idx}" if has_count else "" + body += f' "{pred_prefix}{instr} _, [{addr}]{count_suffix};\\n"\n' + body += ' "}\\n"\n' + constraints = ['"r"(barrier_addr)'] + if has_count: + constraints.append('"r"(count)') + if has_remote: + constraints.append('"r"(remote)') + if has_pred: + constraints.append('"r"(pred)') + body += f' :: {", ".join(constraints)} : "memory");' + else: + count_suffix = ", %1" if has_count else "" + constraints = '"r"(barrier_addr)' + if has_count: + constraints += ', "r"(count)' + body += f' asm volatile("{instr} _, [%0]{count_suffix};"\n' + body += f' :: {constraints} : "memory");' + + return name, f"({', '.join(params)})", body -# Same cross-CTA arrive, but with an explicit arrival-count operand -# (``..., [remAddr32], count``). Matches the ``tma::cluster::arrive`` spelling. device_intrinsic( - "_ptx_mbarrier_arrive_remote_count", - helper_name="tvm_builtin_ptx_mbarrier_arrive_remote_count", - c_signature="(void* barrier, int cta_id, int pred, int count)", - body=( - " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" - " asm volatile(\n" - ' "{\\n"\n' - ' ".reg .pred p;\\n"\n' - ' ".reg .b32 remAddr32;\\n"\n' - ' "setp.ne.s32 p, %2, 0;\\n"\n' - ' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n' - ' "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %3;\\n"\n' - ' "}\\n"\n' - ' :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(count) : "memory");' - ), + "ptx_mbarrier_arrive", + n_attrs=6, + helper_name=lambda *a: _ptx_mbarrier_arrive_parts(*a)[0], + c_signature=lambda *a: _ptx_mbarrier_arrive_parts(*a)[1], + body=lambda *a: _ptx_mbarrier_arrive_parts(*a)[2], ) -@register_codegen("ptx_mbarrier_arrive") -def _codegen_mbarrier_arrive(*args): - """Dispatch by arg count: 1 -> local, 3 -> remote, 4 -> remote+count.""" - if len(args) == 1: - result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_local"](list(args)) - elif len(args) == 3: - result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote"](list(args)) - elif len(args) == 4: - result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_remote_count"](list(args)) +# mbarrier.complete_tx{.sem.scope}{.space}.b64 [addr], txCount; +# sem={.relaxed} scope={.cta,.cluster} space={.shared{::cta},.shared::cluster} +def _ptx_mbarrier_complete_tx_parts(*args): + sem, scope, space, has_remote, has_pred = args[-5:] + sem = parse_str(sem) + scope = parse_str(scope) + space = parse_str(space) + has_remote = _as_bool(has_remote) + has_pred = _as_bool(has_pred) + if sem not in MBARRIER_COMPLETE_TX_SEM: + raise ValueError(f"invalid mbarrier.complete_tx sem {sem!r}") + if scope not in MBARRIER_COMPLETE_TX_SCOPE: + raise ValueError(f"invalid mbarrier.complete_tx scope {scope!r}") + if space not in MBARRIER_COMPLETE_TX_SPACE: + raise ValueError(f"invalid mbarrier.complete_tx space {space!r}") + if has_remote and space != "shared::cluster": + raise ValueError("remote mbarrier.complete_tx requires space='shared::cluster'") + + name = ( + "tvm_builtin_ptx_mbarrier_complete_tx" + f"_{_safe_attr(sem)}_{_safe_attr(scope)}_{_safe_attr(space)}" + ) + if has_remote: + name += "_remote" + if has_pred: + name += "_pred" + + params = ["void* barrier", "int tx_count"] + if has_remote: + params.append("int remote") + if has_pred: + params.append("int pred") + + instr = f"mbarrier.complete_tx.{sem}.{scope}.{space}.b64" + body = " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" + addr = "barrier_addr" + if has_remote: + body += ( + " unsigned int remote_addr;\n" + ' asm volatile("mapa.shared::cluster.u32 %0, %1, %2;"\n' + ' : "=r"(remote_addr) : "r"(barrier_addr), "r"(remote));\n' + ) + addr = "remote_addr" + if has_pred: + body += ( + " asm volatile(\n" + ' "{\\n"\n' + ' ".reg .pred p;\\n"\n' + ' "setp.ne.s32 p, %2, 0;\\n"\n' + f' "@p {instr} [%0], %1;\\n"\n' + ' "}\\n"\n' + f' :: "r"({addr}), "r"(tx_count), "r"(pred) : "memory");' + ) else: - raise ValueError(f"ptx_mbarrier_arrive expects 1, 3, or 4 args, got {len(args)}") - return result[0] if isinstance(result, tuple) else result + body += ( + f' asm volatile("{instr} [%0], %1;"\n' + f' :: "r"({addr}), "r"(tx_count) : "memory");' + ) + + return name, f"({', '.join(params)})", body -# ============================================================================= -# mbarrier.arrive.expect_tx — local + remote (cluster-mapped) forms. -# ============================================================================= device_intrinsic( - "_ptx_mbarrier_arrive_expect_tx_local", - helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx", - c_signature="(void* barrier, int byte_count)", - body=( - " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" - ' asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;"\n' - ' :: "r"(barrier_addr), "r"(byte_count) : "memory");' - ), -) + "ptx_mbarrier_complete_tx", + n_attrs=5, + helper_name=lambda *a: _ptx_mbarrier_complete_tx_parts(*a)[0], + c_signature=lambda *a: _ptx_mbarrier_complete_tx_parts(*a)[1], + body=lambda *a: _ptx_mbarrier_complete_tx_parts(*a)[2], +) + + +# mbarrier.arrive.expect_tx{.sem.scope}{.space}.b64 _, [addr], txCount; +def _ptx_mbarrier_arrive_expect_tx_parts(*args): + sem, scope, space, has_remote, has_pred = args[-5:] + sem, scope, space = _check_mbarrier_arrive_attrs(sem, scope, space) + has_remote = _as_bool(has_remote) + has_pred = _as_bool(has_pred) + if has_remote and space != "shared::cluster": + raise ValueError("remote mbarrier.arrive.expect_tx requires space='shared::cluster'") + + name = "tvm_builtin_ptx_mbarrier_arrive_expect_tx" + if sem: + name += f"_{_safe_attr(sem)}_{_safe_attr(scope)}" + name += f"_{_safe_attr(space)}" + if has_remote: + name += "_remote" + if has_pred: + name += "_pred" + + params = ["void* barrier", "int byte_count"] + arg_idx = 2 + if has_remote: + params.append("int remote") + remote_idx = arg_idx + arg_idx += 1 + if has_pred: + params.append("int pred") + pred_idx = arg_idx + + instr_suffix = f".{sem}.{scope}" if sem else "" + instr = f"mbarrier.arrive.expect_tx{instr_suffix}.{space}.b64" + body = " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" + + if has_remote or has_pred: + body += " asm volatile(\n" + body += ' "{\\n"\n' + if has_pred: + body += ' ".reg .pred p;\\n"\n' + if has_remote: + body += ' ".reg .b32 remAddr32;\\n"\n' + if has_pred: + body += f' "setp.ne.s32 p, %{pred_idx}, 0;\\n"\n' + if has_remote: + prefix = "@p " if has_pred else "" + body += ( + f' "{prefix}mapa.shared::cluster.u32 remAddr32, %0, %{remote_idx};\\n"\n' + ) + addr = "remAddr32" + else: + addr = "%0" + pred_prefix = "@p " if has_pred else "" + body += f' "{pred_prefix}{instr} _, [{addr}], %1;\\n"\n' + body += ' "}\\n"\n' + constraints = ['"r"(barrier_addr)', '"r"(byte_count)'] + if has_remote: + constraints.append('"r"(remote)') + if has_pred: + constraints.append('"r"(pred)') + body += f' :: {", ".join(constraints)} : "memory");' + else: + body += f' asm volatile("{instr} _, [%0], %1;"\n' + body += ' :: "r"(barrier_addr), "r"(byte_count) : "memory");' + + return name, f"({', '.join(params)})", body + + device_intrinsic( - "_ptx_mbarrier_arrive_expect_tx_remote", - helper_name="tvm_builtin_ptx_mbarrier_arrive_expect_tx_remote", - c_signature="(void* barrier, int cta_id, int pred, int byte_count)", - body=( - " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" - " asm volatile(\n" - ' "{\\n"\n' - ' ".reg .pred p;\\n"\n' - ' ".reg .b32 remAddr32;\\n"\n' - ' "setp.ne.s32 p, %2, 0;\\n"\n' - ' "@p mapa.shared::cluster.u32 remAddr32, %0, %1;\\n"\n' - ' "@p mbarrier.arrive.expect_tx.shared::cluster.b64 _, [remAddr32], %3;\\n"\n' - ' "}\\n"\n' - ' :: "r"(barrier_addr), "r"(cta_id), "r"(pred), "r"(byte_count) : "memory");' - ), + "ptx_mbarrier_arrive_expect_tx", + n_attrs=5, + helper_name=lambda *a: _ptx_mbarrier_arrive_expect_tx_parts(*a)[0], + c_signature=lambda *a: _ptx_mbarrier_arrive_expect_tx_parts(*a)[1], + body=lambda *a: _ptx_mbarrier_arrive_expect_tx_parts(*a)[2], ) -@register_codegen("ptx_mbarrier_arrive_expect_tx") -def _codegen_mbarrier_arrive_expect_tx(*args): - """Dispatch by arg count: 2 -> local, 4 -> remote. Remote arg order from - the user is (bar, byte_count, cta_id, pred); reorder to match the helper - signature (bar, cta_id, pred, byte_count).""" - if len(args) == 2: - result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_local"](list(args)) - elif len(args) == 4: - bar, byte_count, cta_id, pred = args - result = CODEGEN_REGISTRY["tirx._ptx_mbarrier_arrive_expect_tx_remote"]( - [bar, cta_id, pred, byte_count] +def _ptx_mbarrier_arrive_no_complete_parts(*args): + space, has_pred = args[-2:] + space = parse_str(space) + has_pred = _as_bool(has_pred) + if space not in ("shared", "shared::cta"): + raise ValueError("mbarrier.arrive.noComplete space must be 'shared' or 'shared::cta'") + + name = f"tvm_builtin_ptx_mbarrier_arrive_no_complete_{_safe_attr(space)}" + if has_pred: + name += "_pred" + + params = ["void* barrier", "int count"] + if has_pred: + params.append("int pred") + instr = f"mbarrier.arrive.noComplete.release.cta.{space}.b64" + body = " unsigned int barrier_addr = __cvta_generic_to_shared(barrier);\n" + if has_pred: + body += ( + " asm volatile(\n" + ' "{\\n"\n' + ' ".reg .pred p;\\n"\n' + ' "setp.ne.s32 p, %2, 0;\\n"\n' + f' "@p {instr} _, [%0], %1;\\n"\n' + ' "}\\n"\n' + ' :: "r"(barrier_addr), "r"(count), "r"(pred) : "memory");' ) else: - raise ValueError(f"ptx_mbarrier_arrive_expect_tx expects 2 or 4 args, got {len(args)}") - return result[0] if isinstance(result, tuple) else result + body += f' asm volatile("{instr} _, [%0], %1;"\n' + body += ' :: "r"(barrier_addr), "r"(count) : "memory");' + return name, f"({', '.join(params)})", body + + +device_intrinsic( + "ptx_mbarrier_arrive_no_complete", + n_attrs=2, + helper_name=lambda *a: _ptx_mbarrier_arrive_no_complete_parts(*a)[0], + c_signature=lambda *a: _ptx_mbarrier_arrive_no_complete_parts(*a)[1], + body=lambda *a: _ptx_mbarrier_arrive_no_complete_parts(*a)[2], +) # ============================================================================= diff --git a/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py index 029a6bf87ce8..16f93140d2c0 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py +++ b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py @@ -442,6 +442,21 @@ def _get_tcgen05_mma_kind(d_dtype, a_dtype, b_dtype, sfa_dtype="", sfb_dtype="") "mxf4nvf4": (64, 128), } +# Operand dtypes admitting a transposed (MN-major) SMEM read in dense +# tcgen05.mma; shared with the gemm_async instruction-descriptor fold. +_TCGEN05_MMA_TRANS_DTYPES = frozenset( + { + PTXDataType.FLOAT8_E4M3FN, + PTXDataType.FLOAT8_E4M3FNUZ, + PTXDataType.FLOAT8_E5M2, + PTXDataType.INT8, + PTXDataType.UINT8, + PTXDataType.FLOAT16, + PTXDataType.BFLOAT16, + PTXDataType.TENSOR_FLOAT32, + } +) + def _check_tcgen05_mma_matrix_shape(kind, cta_group, m, n, k, is_sparse): err = ( @@ -573,19 +588,9 @@ def codegen_ptx_tcgen05_encode_instr_descriptor( a_format = format_map[atype] b_format = format_map[btype] - valid_dtypes_for_trans = { - PTXDataType.FLOAT8_E4M3FN, - PTXDataType.FLOAT8_E4M3FNUZ, - PTXDataType.FLOAT8_E5M2, - PTXDataType.INT8, - PTXDataType.UINT8, - PTXDataType.FLOAT16, - PTXDataType.BFLOAT16, - PTXDataType.TENSOR_FLOAT32, - } - if trans_a and atype not in valid_dtypes_for_trans: + if trans_a and atype not in _TCGEN05_MMA_TRANS_DTYPES: raise ValueError(f"Invalid a_dtype for transpose: {a_dtype}") - if trans_b and btype not in valid_dtypes_for_trans: + if trans_b and btype not in _TCGEN05_MMA_TRANS_DTYPES: raise ValueError(f"Invalid b_dtype for transpose: {b_dtype}") if (neg_a or neg_b) and kind not in ["f16", "tf32", "f8f6f4"]: raise ValueError(f"Invalid kind for negate: {kind}") @@ -732,7 +737,8 @@ def _mma_dense_parts(*args): Args layout: (d_tmem_addr, a_operand, b_desc[, sp_tmem_addr], i_desc, enable_input_d, mask0..maskN-1[, pred], - kind, sparse, use_a_tmem, cta_group, scale_input_d, has_pred) + kind, sparse, use_a_tmem, cta_group, scale_input_d, + weight_stationary) """ attrs = args[-6:] kind = parse_str(attrs[0]) @@ -744,7 +750,7 @@ def _mma_dense_parts(*args): ) cta_group = int(attrs[3]) scale_input_d = int(attrs[4]) - has_pred = bool(int(attrs[5])) + weight_stationary = bool(int(attrs[5])) if not 0 <= scale_input_d <= 15: raise ValueError( @@ -754,8 +760,26 @@ def _mma_dense_parts(*args): raise ValueError(f"scale_input_d is only valid for kind 'f16' or 'tf32', not '{kind!r}'") if scale_input_d > 0 and kind == "i8": raise ValueError("Int form: scale_input_d not supported (only valid for f16/tf32)") + if weight_stationary and sparse: + raise ValueError("tcgen05.mma.ws sparse form is not supported by this intrinsic") + if weight_stationary and cta_group != 1: + raise ValueError("tcgen05.mma.ws is currently supported only for cta_group=1") + if weight_stationary and scale_input_d > 0: + raise ValueError("tcgen05.mma.ws does not support scale_input_d in this intrinsic") num_masks = 8 if cta_group == 2 else 4 + operand_count = len(args) - len(attrs) + expected_operand_count = (6 if sparse else 5) + num_masks + if operand_count == expected_operand_count: + has_pred = False + elif operand_count == expected_operand_count + 1: + has_pred = True + else: + raise ValueError( + "The number of operands for ptx_tcgen05_mma is incorrect, expected " + f"{expected_operand_count} or {expected_operand_count + 1}, got {operand_count}." + ) + a_type = "uint32_t" if use_a_tmem else "uint64_t" a_constraint = "r" if use_a_tmem else "l" @@ -773,6 +797,7 @@ def _mma_dense_parts(*args): name = ( f"ptx_tcgen05_mma_cta_{cta_group}_kind_{kind}" f"{'_sp' if sparse else ''}{'_TS' if use_a_tmem else '_SS'}" + f"{'_ws' if weight_stationary else ''}" f"{('_' + str(scale_input_d)) if scale_input_d > 0 else ''}" f"{'_pred' if has_pred else ''}" ) @@ -805,13 +830,31 @@ def _mma_dense_parts(*args): asm_inputs.append('"r"(pred)') inputs_str = ", ".join(asm_inputs) + pred_prefix = "@p_issue " if has_pred else "" + pred_reg = ", p_issue" if has_pred else "" + pred_setp = f' "setp.ne.b32 p_issue, %{pred_idx}, 0;\\n"\n' if has_pred else "" + if weight_stationary: + # Weight-stationary ABI takes the input-D predicate and a literal ``0`` + # instead of the regular dense disable-output-lane mask vector. + instr = f"tcgen05.mma.ws.cta_group::{cta_group}.kind::{kind} [%0], {a_str}, %2, %3" + body = ( + " asm volatile(\n" + ' "{\\n"\n' + f' ".reg .pred p{pred_reg};\\n"\n' + f' "setp.ne.b32 p, %{p_idx}, 0;\\n"\n' + f"{pred_setp}" + f' "{pred_prefix}{instr}, p, 0;\\n"\n' + ' "}\\n"\n' + " :\n" + f" : {inputs_str}\n" + " );" + ) + return name, sig, body + instr = ( f"tcgen05.mma{sparse_suffix}.cta_group::{cta_group}.kind::{kind}" f" [%0], {a_str}, %2, {sp_str}" ) - pred_prefix = "@p_issue " if has_pred else "" - pred_reg = ", p_issue" if has_pred else "" - pred_setp = f' "setp.ne.b32 p_issue, %{pred_idx}, 0;\\n"\n' if has_pred else "" body = ( " asm volatile(\n" ' "{\\n"\n' @@ -851,6 +894,7 @@ def _dispatch_tcgen05_mma( cta_group, enable_input_d, scale_input_d, + weight_stationary, *disable_output_lane, pred=None, sparse=False, @@ -889,7 +933,14 @@ def _dispatch_tcgen05_mma( if has_pred: operand_args.append(pred) - attr_args = [kind, sparse, use_a_tmem_b, cta_group_i, scale_input_d_i, int(has_pred)] + attr_args = [ + kind, + sparse, + use_a_tmem_b, + cta_group_i, + scale_input_d_i, + int(weight_stationary), + ] return CODEGEN_REGISTRY[f"tirx.{op}"](operand_args + attr_args) @@ -906,6 +957,7 @@ def codegen_ptx_tcgen05_mma( cta_group, enable_input_d, scale_input_d, + weight_stationary, *rest, ): # `rest` = disable_output_lane (4 or 8) + optional pred (1 extra). @@ -929,6 +981,7 @@ def codegen_ptx_tcgen05_mma( cta_group, enable_input_d, scale_input_d, + weight_stationary, *disable_output_lane, pred=pred, sparse=False, @@ -964,6 +1017,7 @@ def codegen_ptx_tcgen05_mma_sp( cta_group, enable_input_d, scale_input_d, + False, *disable_output_lane, sparse=True, sp_tmem_addr=sp_tmem_addr, diff --git a/python/tvm/backend/cuda/operator/tile_primitive/common.py b/python/tvm/backend/cuda/operator/tile_primitive/common.py index 08c56deaecdd..3f6d792b9051 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/common.py @@ -27,7 +27,7 @@ from tvm.script import tirx as T from tvm.tirx import Buffer, BufferRegion, PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def next_power_of_2(x: int) -> int: diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py index 3f9f0f1bc35f..6ba001eaf3a9 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py @@ -16,12 +16,13 @@ # under the License. from .fallback import * -from .gmem_smem import * from .ld_stmatrix import * -from .reg import * from .utils import ( _is_valid_copy, - _is_valid_smem_tmem_copy, _scope_allowed, _single_thread_exec, ) +from .vec_auto import * +from .vec_auto_gmem_smem import * +from .vec_auto_reg import * +from .vec_forced import * diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py index 5ff8100c397d..883b2009d25a 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py @@ -26,9 +26,11 @@ """ from tvm import arith -from tvm.tirx.layout import ComposeLayout, Iter, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import ComposeLayout, Iter, TileLayout from tvm.tirx.operator.tile_primitive.registry import DispatchContext +from ..layout_utils import strip_swizzle_to_tile + def _alignment_ok(vec_len: int, terms) -> bool: """Every term must be a multiple of ``vec_len``. Constants checked @@ -317,16 +319,13 @@ def _vec_len_candidates(elem_bits: int, allowed_bits: tuple | None = None) -> li def _extract_tile(layout, region): """Strip swizzle so we can perm/group as a TileLayout.""" - if isinstance(layout, ComposeLayout): - return layout.tile_layout - if isinstance(layout, SwizzleLayout): - # Region bounds may be constant-valued but remain as unfolded - # expressions after substitution. Simplify before converting to a - # Python integer; genuinely symbolic tile extents still raise. - analyzer = arith.Analyzer() - extents = [int(analyzer.simplify(end - start)) for (start, end) in region] - return TileLayout(S[tuple(extents)]) - return layout + # Region bounds may be constant-valued but remain as unfolded expressions + # after substitution. Simplify before converting to a Python integer; + # genuinely symbolic extents fall back inside ``strip_swizzle_to_tile``. + analyzer = arith.Analyzer() + return strip_swizzle_to_tile( + layout, lambda: [int(analyzer.simplify(end - start)) for (start, end) in region] + ) def _sort_by_stride_desc(layout: TileLayout) -> TileLayout: @@ -414,14 +413,12 @@ def align_layouts_gs( """ g = g_layout.slice(list(g_shape), g_region) s = s_layout.slice(list(s_shape), s_region) - # Detect a SwizzleLayout on the S side BEFORE _extract_tile strips it. + # Detect a swizzle on the S side BEFORE _extract_tile strips it. # vec_len must fit inside one swizzle chunk (C = 2^per_element elements); # otherwise the vec ld/st crosses a swizzle XOR boundary and hits the # wrong physical bytes mid-vec. s_swizzle_chunk_elems = None if isinstance(s_layout, ComposeLayout): - s_swizzle_chunk_elems = 1 << int(s_layout.swizzle.per_element) - elif isinstance(s_layout, SwizzleLayout): s_swizzle_chunk_elems = 1 << int(s_layout.per_element) g = _extract_tile(g, g_region) s = _extract_tile(s, s_region) @@ -547,6 +544,7 @@ def _outer_offsets(outer_iters_s, outer_iters_g, flat_idx): def copy_ptx_form(num_bytes: int) -> tuple[str, str]: """Map copy width (bytes) to PTX ``(vec, ptx_type)`` for ``T.ptx.ld`` / ``T.ptx.st``.""" return { + 32: ("v8", "u32"), 16: ("v4", "u32"), 8: ("v2", "u32"), 4: ("", "u32"), diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py index 7dfa3a8b9bc8..b470e45c1dd4 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py @@ -18,7 +18,7 @@ """Generic swizzle-aware iter pattern for CUDA copy dispatches. When the per-thread outer-iter loop satisfies (C1)+(C2) below for a -``SwizzleLayout(per_element=p, swizzle_len=sw, atom_len=at, +``ComposeLayout(per_element=p, swizzle_len=sw, atom_len=at, swizzle_inner=True)`` on the SMEM side, the swizzled physical address at unrolled iter ``k`` reduces to @@ -71,7 +71,7 @@ from tvm import arith from tvm.script import tirx as T from tvm.tirx.expr import IntImm as _IntImm -from tvm.tirx.layout import ComposeLayout, SwizzleLayout +from tvm.tirx.layout import ComposeLayout, S, TileLayout @dataclass @@ -114,7 +114,7 @@ class SwizzlePattern: degenerate case (no outer iter, just base_off). """ - swizzle: SwizzleLayout + swizzle: ComposeLayout bit_positions: list[int] iter_strides_elems: list[int] outer_iters: "list[_BitIter | _LinearIter]" @@ -124,16 +124,20 @@ def n_binary_iters(self) -> int: return len(self.bit_positions) -def get_swizzle(layout) -> SwizzleLayout | None: - """Return the SwizzleLayout from ``layout`` if present, else ``None``. +def get_swizzle(layout) -> ComposeLayout | None: + """Return a bare-swizzle view of ``layout`` if it is swizzled, else ``None``. - Accepts ``ComposeLayout(SwizzleLayout, TileLayout)`` (the common case - when a TileLayout is wrapped by a swizzle), or a bare ``SwizzleLayout``. + The result carries the swizzle params and an identity tile over the swizzle + period, so ``.apply()`` reproduces the bare swizzle XOR and + ``.per_element`` / ``.swizzle_len`` / ``.atom_len`` / ``.swizzle_inner`` + read the params directly. """ if isinstance(layout, ComposeLayout): - return layout.swizzle - if isinstance(layout, SwizzleLayout): - return layout + p = int(layout.per_element) + sw = int(layout.swizzle_len) + at = int(layout.atom_len) + period = 1 << (p + sw + at) + return ComposeLayout(p, sw, at, TileLayout(S[(period,)]), bool(layout.swizzle_inner)) return None @@ -142,7 +146,7 @@ def _is_pow2(n: int) -> bool: def try_recognize( - swizzle: SwizzleLayout, + swizzle: ComposeLayout, iter_extents: list[int], iter_strides: list[int], s_off_template, @@ -152,7 +156,7 @@ def try_recognize( ``iter_extents`` / ``iter_strides``: the outer-iter list on the S side (excluding T iter and vec iter), in outermost-first order matching - ``s_p.shard[:-2]`` (or the atom-derived analog in ``reg.py``). + ``s_p.shard[:-2]`` (or the atom-derived analog in ``vec_auto_reg.py``). Strides are in element units. Each outer iter with ``extent=2^k`` and ``stride=s`` is conceptually @@ -248,18 +252,23 @@ def try_recognize( # free lane / warp placeholders in s_off_template — ``can_prove_equal`` # returns False if the analyzer can't discharge the equality # universally, conservatively forcing a fallback. + # + # These are compile-time SMEM-offset proofs (values stay far below + # 2**32), so unsigned index terms are analyzed in the no-overflow + # domain; every uint expression admitted is logged once by the analyzer. analyzer = arith.Analyzer() if var_bounds: for var, rng in var_bounds.items(): analyzer.bind(var, rng) - for bj in bj_set: - divisor = C * (1 << bj) - check = tvm.tirx.floormod( - tvm.tirx.floordiv(s_off_template, _IntImm("int32", divisor)), - _IntImm("int32", 2), - ) - if not analyzer.can_prove_equal(check, _IntImm("int32", 0)): - return None + with arith.allow_uint_as_index(): + for bj in bj_set: + divisor = C * (1 << bj) + check = tvm.tirx.floormod( + tvm.tirx.floordiv(s_off_template, _IntImm(s_off_template.expr_ty().dtype, divisor)), + _IntImm(s_off_template.expr_ty().dtype, 2), + ) + if not analyzer.can_prove_equal(check, _IntImm(s_off_template.expr_ty().dtype, 0)): + return None return SwizzlePattern( swizzle=swizzle, @@ -394,7 +403,7 @@ def emit_iter_offset(pattern: SwizzlePattern, signed_strides, base_off, k): return off -def emit_fallback_offset(swizzle: SwizzleLayout, s_off_resolved, ds_k): +def emit_fallback_offset(swizzle: ComposeLayout, s_off_resolved, ds_k): """Slow but always-correct path: full ``swizzle.apply(s_off + ds_k)`` per iter. Use when ``try_recognize`` returns ``None``. diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py index 622a8e7a158d..1a78fce1d9cd 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py @@ -27,11 +27,11 @@ register_dispatch, ) from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ._common import _TID_AXIS_FOR_SCOPE -from .reg import _axis_decl from .utils import _is_valid_copy +from .vec_auto_reg import _axis_decl def _region_st_extent(buffer_region): @@ -75,6 +75,10 @@ def _src_coord(lvs): coord[src_indices[k]] += lv return coord + if not copy_extents: + T.buffer_store(dst_buf, src_buf[tuple(src_st)], dst_st) + return + with T.grid(*copy_extents) as lvs: T.buffer_store(dst_buf, src_buf[tuple(_src_coord(lvs))], _dst_coord(lvs)) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py index 123852d8c815..db2ba99946b6 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py @@ -36,7 +36,7 @@ register_dispatch, ) from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ._common import ( _TID_AXIS_FOR_SCOPE, @@ -184,7 +184,7 @@ def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu else: outer_iters_s = list(s_p.shard[:-1]) - # SwizzleLayout on s_buf: try the closed-form signed-strides pattern + # Swizzle on s_buf: try the closed-form signed-strides pattern # (precomputed once per thread, then per-iter is a sum of register # adds); fall back to per-iter ``swizzle.apply`` (one full XOR + # decompose per iter). Closure picked at parse time so the TIRx parser diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py index 1116b4407e3d..5405c5280910 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py @@ -32,15 +32,15 @@ from tvm.tirx.layout import S, TileLayout from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ._common import ( # noqa: F401 (_carve_tail reserved for future variants) _carve_tail, _extract_tile, ) from ._swizzle_iter import emit_init, emit_iter_offset, get_swizzle, try_recognize -from .reg import _all_threads_active, _ptr_off from .utils import _is_valid_copy, _scope_allowed +from .vec_auto_reg import _all_threads_active, _ptr_off _REG_SMEM_PAIRS = [ ("local", "shared*"), @@ -107,10 +107,8 @@ def _emit(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: # TileLayout. The ComposeLayout doesn't have ``.replica`` / ``.shard``, # so we must peel *before* the structural checks below. Capture the # swizzle separately for use at emit time. - # NB: read swizzle from the *buffer* layout, not the post-canon ``s``. - # When the underlying tile is trivial, ``ComposeLayoutNode::Canonicalize`` - # returns a bare ``SwizzleLayout``; isinstance(s, ComposeLayout) is then - # False and we'd miss the swizzle here. + # NB: read swizzle from the *buffer* layout, not the post-canon ``s`` -- + # the buffer layout is the reliable source for the swizzle params. s_swizzle = get_swizzle(s_buf.layout) if s_swizzle is not None and s_swizzle.per_element < 3: # ldmatrix/stmatrix .b16 reads/writes 8 fp16 = 128b per lane in one diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py index ab6a57c86bc1..be4583d1d551 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py @@ -34,11 +34,12 @@ from tvm.tirx import Buffer, PrimFunc from tvm.tirx import Var as _TirVar from tvm.tirx.expr import IntImm as _IntImm -from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import TileLayout from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall +from ..layout_utils import strip_swizzle_to_tile from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_return_type from ._swizzle_iter import ( emit_fallback_offset, @@ -51,20 +52,8 @@ def _extract_tile(layout, region): - """Strip swizzle off ``layout`` so we can perm/group it as a TileLayout. - - ``region`` is the per-axis ``(start, end)`` pair list — we only consume - its extents when ``layout`` is a bare ``SwizzleLayout`` (rebuilding a - trivial TileLayout for it). Plain ``TileLayout`` / ``ComposeLayout`` - don't need the extent, so symbolic regions are fine for them. - """ - if isinstance(layout, ComposeLayout): - return layout.tile_layout - if isinstance(layout, SwizzleLayout): - # TODO: keep swizzle info around for later (addressing in emit). - extents = [int(end - start) for (start, end) in region] - return TileLayout(S[tuple(extents)]) - return layout + """Strip swizzle off ``layout`` so it can be perm/grouped as a TileLayout.""" + return strip_swizzle_to_tile(layout, lambda: [int(end - start) for (start, end) in region]) _REG_PAIRS = [ diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py index 5a4ce4381296..cccf8b2b99b2 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py @@ -18,44 +18,12 @@ from collections.abc import Iterable -from tvm.tirx import Buffer from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import match_scope, validate_copy_op -def _is_valid_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext): - """Validate smem->tmem copy operation. - - The new tcgen05.cp.32x128b.warpx4 dispatch requires the destination tmem - buffer to declare warpx4 broadcast as ``R[4 : 32@TLane]``. The legacy - 128-row dispatch path (no replica) goes through a separate code path and - is not handled here. - """ - dst_region, src_region = op_call.args[:2] - src: Buffer = src_region.buffer - dst: Buffer = dst_region.buffer - if not (src.scope().startswith("shared") and dst.scope() == "tmem"): - return (False, f"expected shared->tmem, got {src.scope()}->{dst.scope()}") - if not (src.layout and dst.layout): - return (False, "both buffers must have layouts") - if dst.allocated_addr is None: - return (False, "tmem buffer must have allocated_addr") - # Require warpx4 router on TMEM side so this dispatch only handles the - # 32x128b.warpx4 case; other shapes (128x256b/128x128b etc.) fall back - # to the legacy dispatch. - rep = dst.layout.replica - if not ( - len(rep) == 1 - and int(rep[0].extent) == 4 - and int(rep[0].stride) == 32 - and "TLane" in str(rep[0].axis) - ): - return (False, f"requires R[4:32@TLane] on tmem, got replica={list(rep)}") - return (True, None) - - def _single_thread_exec(op_call: TilePrimitiveCall, sctx: DispatchContext): """Predicate: exec scope must be single-thread.""" exec_scope = sctx.scope_kind diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py new file mode 100644 index 000000000000..b54781f2bd16 --- /dev/null +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py @@ -0,0 +1,53 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Auto-vectorizing synchronous copy dispatch.""" + +from tvm.tirx import PrimFunc +from tvm.tirx.operator.tile_primitive.dispatcher import fail, predicate, register_dispatch +from tvm.tirx.operator.tile_primitive.registry import DispatchContext +from tvm.tirx.tile_primitive import TilePrimitiveCall + +from .vec_auto_gmem_smem import _emit_gmem_smem, _is_gmem_smem +from .vec_auto_reg import _emit_reg, _is_reg_copy + + +def _is_vec_auto_copy(op_call: TilePrimitiveCall, sctx: DispatchContext): + g_ok, g_reason = _is_gmem_smem(op_call, sctx) + if g_ok: + return True, None + r_ok, r_reason = _is_reg_copy(op_call, sctx) + if r_ok: + return True, None + return False, f"gmem_smem: {g_reason}; reg: {r_reason}" + + +@register_dispatch( + "copy", + "cuda", + variant="vec_auto", + priority=10, + when=[predicate("vec_auto_applicable", _is_vec_auto_copy)], +) +def copy_schedule_vec_auto(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + g_ok, g_reason = _is_gmem_smem(op_call, sctx) + if g_ok: + return _emit_gmem_smem(op_call, sctx) + r_ok, r_reason = _is_reg_copy(op_call, sctx) + if r_ok: + return _emit_reg(op_call, sctx) + fail(f"gmem_smem: {g_reason}; reg: {r_reason}") diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py new file mode 100644 index 000000000000..30a80a8f617a --- /dev/null +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py @@ -0,0 +1,290 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Copy dispatch for ``global ↔ shared`` (no register side). + +There's no per-thread register side to inherit a partition from — both sides +are cross-thread storage. The partition is synthesized from the surrounding +scope context (warp / warpgroup / cta / thread): ``thread_cnt`` is derived +from ``sctx.intra`` and each thread takes ``n_elements / thread_cnt`` +consecutive fused-index slots. Layout / partition algorithm lives in +``_common.py`` and is shared with ``ldgsts.py``. +""" + +import tvm +from tvm.runtime import DataType +from tvm.script import tirx as T +from tvm.tirx import Buffer, PrimFunc +from tvm.tirx import Var as _TirVar +from tvm.tirx.expr import IntImm as _IntImm +from tvm.tirx.operator.tile_primitive.registry import DispatchContext +from tvm.tirx.tile_primitive import TilePrimitiveCall + +from ._common import ( + _TID_AXIS_FOR_SCOPE, + _thread_cnt, + align_layouts_gs, + copy_ptx_form, + copy_ptx_ld_return_type, +) +from ._swizzle_iter import ( + emit_init, + emit_iter_offset, + get_swizzle, + try_recognize, +) +from .utils import _is_valid_copy, _scope_allowed +from .vec_auto_reg import _all_threads_active, _axis_decl, _ptr_off + +_GMEM_SMEM_PAIRS = [ + ("global", "shared*"), + ("shared*", "global"), +] + + +def _divides_thread_cnt( + op_call: TilePrimitiveCall, sctx: DispatchContext +) -> tuple[bool, str | None]: + """Reject copies whose region element count does not divide ``thread_cnt``. + + Without this guard the emit's ``[outer, T, vec]`` partition has no + integer solution: either every thread gets fractional work, or + ``thread_cnt=0`` (degenerate scope) hits a modulo-by-zero. Both cases + indicate a poorly-shaped copy (e.g. 1024-thread CTA writing a 64-elem + tail) that this dispatch refuses to paper over with a slow scalar emit. + """ + op_call = TilePrimitiveCall.downcast(op_call) + thread_cnt = _thread_cnt(sctx) + if thread_cnt <= 0: + return False, f"degenerate thread_cnt={thread_cnt} (scope has empty intra)" + g_br = op_call.src if op_call.src.buffer.scope() == "global" else op_call.dst + n_elements = 1 + for r in g_br.region: + ext = r.extent + try: + n_elements *= int(ext) + except (TypeError, ValueError): + return False, f"non-constant region extent {ext}" + if n_elements % thread_cnt != 0: + return False, (f"region size {n_elements} not divisible by thread_cnt={thread_cnt}") + return True, None + + +def _is_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]: + if not sctx.is_target("cuda"): + return False, "non-cuda target" + if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"): + return False, f"unsupported exec_scope {sctx.scope_kind}" + for check in ( + lambda: _all_threads_active(sctx), + lambda: _is_valid_copy(op_call, sctx), + lambda: _scope_allowed(op_call, sctx, allowed_pairs=_GMEM_SMEM_PAIRS), + lambda: _divides_thread_cnt(op_call, sctx), + ): + ok, msg = check() + if not ok: + return False, msg + return True, None + + +def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + op_call = TilePrimitiveCall.downcast(op_call) + src: Buffer = op_call.src.buffer + dst: Buffer = op_call.dst.buffer + if src.scope() == "global": + g_buf, g_br, s_buf, s_br = src, op_call.src, dst, op_call.dst + g_is_src = True + else: + g_buf, g_br, s_buf, s_br = dst, op_call.dst, src, op_call.src + g_is_src = False + + g_region = [(r.min, r.min + r.extent) for r in g_br.region] + s_region = [(r.min, r.min + r.extent) for r in s_br.region] + + elem_bits = DataType(src.dtype).bits + thread_cnt = _thread_cnt(sctx) + + with sctx.target: + g_p, s_p, vec_len = align_layouts_gs( + g_buf.layout, + g_buf.shape, + g_region, + s_buf.layout, + s_buf.shape, + s_region, + elem_bits, + thread_cnt, + ) + + # vec_len=1 is the scalar fallback — uses the same unified + # [outer x thread x vec] coord scheme below. + + vec_bits = vec_len * elem_bits + num_bytes = vec_bits // 8 + vec, ptx_type = copy_ptx_form(num_bytes) + + # Express the per-thread per-round address as a 3D coord ``(f, tid, 0)`` vs + # ``[total_outer, thread_cnt, vec_len]``; ``layout.apply`` flattens the rest. + n_elements = 1 + for it in s_p.shard: + n_elements *= int(it.extent) + assert n_elements % (thread_cnt * vec_len) == 0, ( + f"partition produced {n_elements} elements but thread_cnt({thread_cnt}) * " + f"vec_len({vec_len}) = {thread_cnt * vec_len} doesn't divide it" + ) + total_outer = n_elements // (thread_cnt * vec_len) + apply_shape = [ + _IntImm("int32", total_outer), + _IntImm("int32", thread_cnt), + _IntImm("int32", vec_len), + ] + + s_zero = [0] * len(s_buf.shape) + g_zero = [0] * len(g_buf.shape) + + tid_axis_name = _TID_AXIS_FOR_SCOPE[sctx.scope_kind] if thread_cnt > 1 else None + + # Walk shard back from the vec iter to the prefix covering T exactly + # (∏ext == thread_cnt); the leading prefix is the outer iter list. + if thread_cnt > 1: + acc, _i = 1, len(s_p.shard) - 2 + while _i >= 0 and acc < thread_cnt: + _ext = int(s_p.shard[_i].extent) + if acc * _ext > thread_cnt: + break + acc *= _ext + _i -= 1 + outer_iters_s = list(s_p.shard[: _i + 1]) if acc == thread_cnt else [] + else: + outer_iters_s = list(s_p.shard[:-1]) + + # Swizzle on s_buf: try the closed-form signed-strides pattern, else + # fall back to per-iter ``swizzle.apply``; closure picked at parse time. + swizzle = get_swizzle(s_buf.layout) + swizzle_pattern = None + if swizzle is not None and outer_iters_s: + if tid_axis_name is not None: + _tid_placeholder = _TirVar(tid_axis_name, "int32") + else: + _tid_placeholder = _IntImm("int32", 0) + s_off_template = s_p.apply( + _IntImm("int32", 0), + _tid_placeholder, + _IntImm("int32", 0), + shape=apply_shape, + )["m"] + # Bind the tid range so the (C1) analyzer can discharge + # ``bit_bj(s_off // C) == 0``; without bounds it rejects. + var_bounds = {} + if tid_axis_name is not None: + var_bounds[_tid_placeholder] = tvm.ir.Range.from_min_extent(0, thread_cnt) + swizzle_pattern = try_recognize( + swizzle, + [int(it.extent) for it in outer_iters_s], + [int(it.stride) for it in outer_iters_s], + s_off_template, + var_bounds=var_bounds or None, + ) + + class _SwizzleState: + def __init__(self): + self.signed_strides = None + self.base_off = None + + state = _SwizzleState() + + def _decl_tid(): + if tid_axis_name is not None: + return _axis_decl(tid_axis_name, sctx) + return _IntImm("int32", 0) + + def _setup_swizzle(tid): + if swizzle_pattern is None: + return + s_off_resolved = s_p.apply( + _IntImm("int32", 0), + tid, + _IntImm("int32", 0), + shape=apply_shape, + )["m"] + state.signed_strides, state.base_off = emit_init( + swizzle_pattern, + s_off_resolved, + ) + + if swizzle_pattern is not None: + + def _s_off(f, s_lin): + return emit_iter_offset( + swizzle_pattern, + state.signed_strides, + state.base_off, + f, + ) + elif swizzle is not None: + _sw = swizzle + + def _s_off(f, s_lin): + return _sw.apply(s_lin)["m"] + else: + + def _s_off(f, s_lin): + return s_lin + + v0 = _IntImm("int32", 0) + + # fmt: off + @T.prim_func(check_well_formed=False) + def impl(): + tid = _decl_tid() + _setup_swizzle(tid) + tmp = T.alloc_local((vec_len,), src.dtype) + tmp_ptr = tmp.ptr_to([0]) + # Pass typed ptr_to(...) directly to _ptr_off (caching → byte math, + # misaligned vec ops); keep a serial loop, T.unroll floods the kernel. + for f in range(total_outer): + s_lin = s_p.apply(f, tid, v0, shape=apply_shape)["m"] + g_lin = g_p.apply(f, tid, v0, shape=apply_shape)["m"] + s_off = _s_off(f, s_lin) + s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off) + g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin) + if g_is_src: + T.ptx.ld( + g_ptr, + copy_ptx_ld_return_type(ptx_type), + ptx_type, + dst=tmp_ptr, + space="global", + vec=vec, + ) + T.ptx.st( + s_ptr, src=tmp_ptr, space="shared", vec=vec, ptx_type=ptx_type + ) + else: + T.ptx.ld( + s_ptr, + copy_ptx_ld_return_type(ptx_type), + ptx_type, + dst=tmp_ptr, + space="shared", + vec=vec, + ) + T.ptx.st( + g_ptr, src=tmp_ptr, space="global", vec=vec, ptx_type=ptx_type + ) + # fmt: on + return impl diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py new file mode 100644 index 000000000000..9e5677fd3dfe --- /dev/null +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py @@ -0,0 +1,588 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Non-ldmatrix copy dispatch for register ↔ memory. + +This file owns every copy where one side is per-thread local (``R`` = +register). That R side carries the partition: its ``TileLayout`` ``shard`` +has thread-axis iters telling us which thread owns which logical coordinate. +The other side (``S``) can be ``shared*`` or ``global`` — the algorithm is +identical either way. + +Slice/canonicalize both sides, align via perm+group, then emit a per-thread +vectorized copy loop. Direction-symmetric: covers R2S / S2R / R2G / G2R. +""" + +import tvm +from tvm.arith import Analyzer +from tvm.runtime import DataType +from tvm.script import tirx as T +from tvm.tirx import Buffer, PrimFunc +from tvm.tirx import Var as _TirVar +from tvm.tirx.expr import IntImm as _IntImm +from tvm.tirx.layout import TileLayout +from tvm.tirx.operator.tile_primitive.registry import DispatchContext +from tvm.tirx.tile_primitive import TilePrimitiveCall + +from ..layout_utils import strip_swizzle_to_tile +from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_return_type +from ._swizzle_iter import ( + emit_fallback_offset, + emit_init, + emit_iter_offset, + get_swizzle, + try_recognize, +) +from .utils import _is_valid_copy, _scope_allowed +from .vec_forced import _ld_cache_config + + +def _extract_tile(layout, region): + """Strip swizzle off ``layout`` so it can be perm/grouped as a TileLayout.""" + return strip_swizzle_to_tile(layout, lambda: [int(end - start) for (start, end) in region]) + + +_REG_PAIRS = [ + ("local", "shared*"), + ("shared*", "local"), + ("local", "global"), + ("global", "local"), +] +_SCOPE_RANK = {"thread": 0, "warp": 1, "warpgroup": 2, "cta": 3} +_VALID_R_SUBSCOPES = {"thread", "warp", "warpgroup"} + + +def _all_threads_active(sctx: DispatchContext) -> tuple[bool, str | None]: + if sctx.scope_kind == "thread": + return True, None + required: dict[str, int] = {} + if sctx.scope_kind in ("warp", "warpgroup", "cta"): + required["laneid"] = 32 + if sctx.scope_kind == "warpgroup": + required["wid_in_wg"] = 4 + if sctx.scope_kind == "cta": + tx_iv = sctx.launch_params.get("threadIdx.x") + if tx_iv is None: + return False, "cta scope missing threadIdx.x launch_params" + try: + required["warpid"] = int(tx_iv.dom.extent) // 32 + except (TypeError, ValueError): + return False, f"non-static threadIdx.x extent: {tx_iv.dom.extent}" + for axis_name, expected in required.items(): + if axis_name not in sctx.intra: + return False, f"sctx.intra missing {axis_name!r}" + ext_raw, off_raw = sctx.intra[axis_name] + try: + ext, off = int(ext_raw), int(off_raw) + except (TypeError, ValueError): + return False, f"non-static range for {axis_name}: ({ext_raw}, {off_raw})" + if ext != expected or off != 0: + return False, f"{axis_name} narrowed to [{off}, {off + ext}) vs full [0, {expected})" + return True, None + + +def _r_side_layout_valid( + op_call: TilePrimitiveCall, sctx: DispatchContext +) -> tuple[bool, str | None]: + op_call = TilePrimitiveCall.downcast(op_call) + src: Buffer = op_call.src.buffer + dst: Buffer = op_call.dst.buffer + r_buf = src if src.scope() == "local" else dst + layout = r_buf.layout + if layout is None: + return False, "R has no layout" + if layout.is_swizzle(): + return False, "R layout is swizzle" + if not isinstance(layout, TileLayout): + return False, f"R layout is {type(layout).__name__}, not TileLayout" + + scope_rank = _SCOPE_RANK[sctx.scope_kind] + for it in layout.shard: + ax = it.axis + if not ax.is_thread(): + continue + ax_scope = ax.get_scope() + ax_sub = ax.get_subscope() + if ax_scope is None or ax_sub is None: + return False, f"R thread axis {ax.name!r} missing scope/subscope" + if ax_sub.name not in _VALID_R_SUBSCOPES: + return False, f"R thread axis {ax.name!r} subscope={ax_sub.name!r} (not register-level)" + if ax_scope.name not in _SCOPE_RANK or _SCOPE_RANK[ax_scope.name] > scope_rank: + return ( + False, + f"R thread axis {ax.name!r} scope={ax_scope.name!r} > exec {sctx.scope_kind!r}", + ) + r_br = op_call.src if src.scope() == "local" else op_call.dst + region = [(r.min, r.min + r.extent) for r in r_br.region] + sliced = layout.slice(list(r_buf.shape), region) + if sliced is None: + return False, "R layout slice failed" + analyzer = Analyzer() + for axis, off in sliced.offset.items(): + if axis.is_thread() and not analyzer.can_prove_equal(off, 0): + return False, f"R sliced offset on thread axis {axis.name!r} = {off}" + return True, None + + +def _s_side_slice_ok(op_call: TilePrimitiveCall) -> tuple[bool, str | None]: + """S is the non-local side (shared* or global). Slice must succeed.""" + op_call = TilePrimitiveCall.downcast(op_call) + src_br = op_call.src + dst_br = op_call.dst + s_br = dst_br if src_br.buffer.scope() == "local" else src_br + s_buf: Buffer = s_br.buffer + layout = s_buf.layout + if layout is None: + return False, "S has no layout" + region = [(r.min, r.min + r.extent) for r in s_br.region] + if layout.slice(list(s_buf.shape), region) is None: + return False, "S layout slice failed" + return True, None + + +def _is_reg_copy(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str | None]: + if not sctx.is_target("cuda"): + return False, "non-cuda target" + if sctx.scope_kind not in ("thread", "warp", "warpgroup", "cta"): + return False, f"unsupported exec_scope {sctx.scope_kind}" + for check in ( + lambda: _all_threads_active(sctx), + lambda: _is_valid_copy(op_call, sctx), + lambda: _scope_allowed(op_call, sctx, allowed_pairs=_REG_PAIRS), + lambda: _r_side_layout_valid(op_call, sctx), + lambda: _s_side_slice_ok(op_call), + ): + ok, msg = check() + if not ok: + return False, msg + return True, None + + +def _compute_perm_r(r): + # thread axes first, then by stride descending + def key(p): + it = p[1] + return (0 if it.axis.is_thread() else 1, -int(it.stride)) + + return [i for i, _ in sorted(enumerate(r.shard), key=key)] + + +def align_layouts_raw(r_sliced, s_sliced, s_region): + """Returns (r_p, s_p, s_seps, r_perm). + + ``r_p`` for ``_s_thread_offset``; ``r_perm`` for ``_split_thread_loop`` pairing. + """ + r = r_sliced.canonicalize() + s = s_sliced.canonicalize() + s = _extract_tile(s, s_region) + perm = _compute_perm_r(r) + r_shape_for_group = [int(it.extent) for it in r.shard] + s_grp, seps = s.group(r_shape_for_group) + s_p = s_grp.permute_by_groups(list(seps), perm) + r_perm = r.permute_dims(perm) + r_p = r_perm.canonicalize() + sizes = [seps[i + 1] - seps[i] for i in range(len(seps) - 1)] + s_seps = [0] + for p in perm: + s_seps.append(s_seps[-1] + sizes[p]) + return r_p, s_p, s_seps, r_perm + + +def _split_thread_loop(r_perm, s_p, s_seps): + """Drop R's thread-axis positions and return per-R-position bundles: + (r_iters, s_groups) — same length lists; s_groups[k] is the list of S + iters belonging to the k-th kept R position. + + ``r_perm`` is the permuted-but-uncanonicalized R layout: its iter list lines + up one-to-one with ``s_seps`` (which is built from the same ``perm``), so a + multi-group memory axis keeps each of its groups paired with the matching S + group instead of collapsing into one (see ``align_layouts_raw``).""" + r_iters = [] + s_groups = [] + for k, r_it in enumerate(r_perm.shard): + if r_it.axis.is_thread(): + continue + r_iters.append(r_it) + s_groups.append(list(s_p.shard[s_seps[k] : s_seps[k + 1]])) + return r_iters, s_groups + + +def _build_atoms(r_iters, s_groups): + """One atom per (R position, intra-group S iter): (extent, s_stride, r_mul). + r_mul = R_stride_at_position * (product of S group extents to the right + of this intra-position index) — i.e. how much R address advances per unit + of this iter's loop input.""" + atoms = [] + for r_it, s_group in zip(r_iters, s_groups, strict=True): + rs = int(r_it.stride) + extents = [int(it.extent) for it in s_group] + for j, s_it in enumerate(s_group): + inner_prod = 1 + for e in extents[j + 1 :]: + inner_prod *= e + atoms.append((int(s_it.extent), int(s_it.stride), rs * inner_prod)) + return atoms + + +def _atoms_contiguous_tail_extent(atoms) -> int: + """Like _contiguous_tail_extent but on atoms (uses s_stride for chaining).""" + if not atoms or atoms[-1][1] != 1: + return 0 + acc = atoms[-1][0] + for k in range(len(atoms) - 2, -1, -1): + if atoms[k][1] == acc: + acc *= atoms[k][0] + else: + break + return acc + + +def _split_atoms_for_vec(atoms, vec_len): + """Returns outer atoms (the inner vec_len-element tail is consumed by one + vec ld/st and dropped). Splits the boundary atom if needed.""" + outer = list(atoms) + acc = 1 + while outer: + ext, ss, rm = outer[-1] + new_acc = acc * ext + if new_acc == vec_len: + outer.pop() + return outer + if new_acc > vec_len: + inner_factor = vec_len // acc + outer[-1] = (ext // inner_factor, ss * inner_factor, rm * inner_factor) + return outer + acc = new_acc + outer.pop() + raise ValueError(f"tail too short for vec_len {vec_len}") + + +def _align_layouts(op_call: TilePrimitiveCall, sctx: DispatchContext): + op_call = TilePrimitiveCall.downcast(op_call) + src_br = op_call.src + dst_br = op_call.dst + if src_br.buffer.scope() == "local": + r_br, s_br = src_br, dst_br + else: + r_br, s_br = dst_br, src_br + r_buf = r_br.buffer + s_buf = s_br.buffer + r_region = [(r.min, r.min + r.extent) for r in r_br.region] + s_region = [(r.min, r.min + r.extent) for r in s_br.region] + with sctx.target: + r_sliced = r_buf.layout.slice(list(r_buf.shape), r_region) + s_sliced = s_buf.layout.slice(list(s_buf.shape), s_region) + return align_layouts_raw(r_sliced, s_sliced, s_region) + + +def _make_thread_placeholders(r_p) -> dict[str, _TirVar]: + placeholders: dict[str, _TirVar] = {} + for it in r_p.shard: + name = it.axis.name + if it.axis.is_thread() and name not in placeholders: + placeholders[name] = _TirVar(name, "int32") + return placeholders + + +def _s_thread_offset(r_p, s_p, placeholders: dict[str, _TirVar]): + """Per-thread S base offset. Coord per R position is placeholder (thread + axis) or 0 (memory axis); apply_to_shape decomposes across s_p iters. + Includes layout-level offsets (e.g. from slicing a non-zero S region).""" + coord = [ + placeholders[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0) + for it in r_p.shard + ] + input_shape = [int(it.extent) for it in r_p.shard] + per_iter = s_p.apply_to_shape(coord, input_shape) + off = _IntImm("int32", 0) + for c, it in zip(per_iter, s_p.shard, strict=True): + off = off + c * it.stride + for _axis, val in s_p.offset.items(): + off = off + val + return off + + +_VEC_BITS_CANDIDATES = (128, 64, 32, 16, 8) + + +def _vec_len_candidates(elem_bits: int) -> list[int]: + """Widest-first element counts to try; always ends with scalar (1).""" + out: list[int] = [] + for vb in _VEC_BITS_CANDIDATES: + if vb < elem_bits or vb % elem_bits != 0: + continue + n = vb // elem_bits + if n not in out: + out.append(n) + if 1 not in out: + out.append(1) + return out + + +def _choose_vec_len(elem_bits: int, atoms, r_p, s_p) -> int: + """Widest candidate that: + 1. divides the atom contiguous-tail extent (so vec_len consecutive + R-side regs map to vec_len contiguous S-side elements), AND + 2. keeps every per-thread / per-round address-offset term a + multiple of vec_len, so the resulting vec ld/st pointer is + naturally aligned to vec_bits/8 bytes. + + Only **mem-axis** strides contribute to physical address. Thread-axis + iter strides live in partition-coord space (which thread owns which + logical position), not in the buffer's storage space — they're + redistributed through ``apply_to_shape`` into the mem iters and don't + appear directly in the per-thread address. So neither r-side nor + s-side thread-axis strides belong in the alignment check. + + The contig-tail atoms (whose extents the vec ld/st consumes) have + stride 1 by definition; they live entirely inside the vec and + contribute nothing to the per-round address delta. Only the + **post-vec-split** outer atom strides matter for the per-round delta. + """ + t = _atoms_contiguous_tail_extent(atoms) + # Region-base offsets are real address contributions; thread-iter + # strides are partition-virtual, not storage-physical — exclude them. + shared_terms = list(s_p.offset.values()) + list(r_p.offset.values()) + for n in _vec_len_candidates(elem_bits): + if n == 1: + return n + if t % n != 0 or t < n: + continue + # Post-vec-split outer atoms: these are the strides that contribute + # to per-round address deltas after the vec consumes the inner tail. + outer = _split_atoms_for_vec(atoms, n) + outer_atom_terms = [a[1] for a in outer] + [a[2] for a in outer] + if not _alignment_ok(n, outer_atom_terms + shared_terms): + continue + return n + return 1 + + +def _axis_decl(axis_name: str, sctx: DispatchContext): + """Declare the runtime Var for one thread axis (called inside impl body). + + Each scope_id declarator emits a ``ScopeIdDef`` stmt at the current + builder frame. ``TilePrimitiveDispatch`` re-gathers + resolves all + ScopeIdDefs after dispatch (see ``ResolveAllScopeBinds`` in + ``tile_primitive_dispatch.cc``), so dispatch-introduced vars are bound + alongside kernel-declared ones. + + Extents are deferred: the kernel header is expected to declare the full + scope-id chain (``cta_id`` / ``warpgroup_id`` / ``warp_id_in_wg`` / + ``lane_id`` / ``thread_id`` / ``thread_id_in_wg``) — the verifier then + fills our deferred defs from those siblings. + """ + if axis_name == "tx": + return sctx.launch_params["threadIdx.x"].var + if axis_name == "laneid": + return T.lane_id() + if axis_name == "wid_in_wg": + return T.warp_id_in_wg() + if axis_name == "tid_in_wg": + return T.thread_id_in_wg() + if axis_name == "warpid": + return T.warp_id() + if axis_name == "wgid": + return T.warpgroup_id() + raise ValueError(f"unsupported thread axis {axis_name}") + + +def _s_thread_offset_with_vars(r_p, s_p, axis_var_map: dict): + coord = [ + axis_var_map[it.axis.name] if it.axis.is_thread() else _IntImm("int32", 0) + for it in r_p.shard + ] + input_shape = [int(it.extent) for it in r_p.shard] + per_iter = s_p.apply_to_shape(coord, input_shape) + off = _IntImm("int32", 0) + for c, it in zip(per_iter, s_p.shard, strict=True): + off = off + c * it.stride + for _ax, val in s_p.offset.items(): + off = off + val + return off + + +def _substitute_axes(s_off_template, placeholders: dict[str, _TirVar], sctx: DispatchContext): + """Inside an impl body: declare real scope_ids and rewrite the + placeholder-built ``s_off_template`` to use them.""" + vmap = {placeholders[name]: _axis_decl(name, sctx) for name in placeholders} + return tvm.tirx.stmt_functor.substitute(s_off_template, vmap) + + +def _flat_coords(outer_atoms, flat_idx: int) -> list[int]: + coords = [] + rem = flat_idx + for a in reversed(outer_atoms): + coords.append(rem % a[0]) + rem //= a[0] + coords.reverse() + return coords + + +_POINTER_OFFSET_SRC = ( + "\ntemplate \n" + "__forceinline__ __device__ T* tvm_builtin_pointer_offset(T* ptr, int offset) {\n" + " return ptr + offset;\n" + "}\n" +) + + +def _ptr_off(base_ptr, off): + return T.cuda.func_call( + "tvm_builtin_pointer_offset", + base_ptr, + off, + source_code=_POINTER_OFFSET_SRC, + return_type="handle", + ) + + +def _outer_const_offsets(outer_atoms, flat_idx: int) -> tuple[int, int]: + """Returns (s_offset_const, r_offset_const) for one outer-loop flat index.""" + coords = _flat_coords(outer_atoms, flat_idx) + ds = sum(c * a[1] for c, a in zip(coords, outer_atoms)) + dr = sum(c * a[2] for c, a in zip(coords, outer_atoms)) + return ds, dr + + +def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + op_call = TilePrimitiveCall.downcast(op_call) + src: Buffer = op_call.src.buffer + dst: Buffer = op_call.dst.buffer + if src.scope() == "local": + r_buf, s_buf, r_is_src = src, dst, True + else: + r_buf, s_buf, r_is_src = dst, src, False + + with sctx.target: + r_p, s_p, s_seps, r_perm = _align_layouts(op_call, sctx) + r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps) + atoms = _build_atoms(r_iters, s_groups) + elem_bits = DataType(src.dtype).bits + vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p) + vec_bits = vec_len * elem_bits + outer = _split_atoms_for_vec(atoms, vec_len) + per_thread_r_total = 1 + for it in r_iters: + per_thread_r_total *= int(it.extent) + per_thread_r_shape = [per_thread_r_total or 1] + + # Build the per-thread S offset outside the impl with placeholder Vars; + # inside the impl, real scope_ids are declared and substituted in. + placeholders = _make_thread_placeholders(r_p) + s_off_template = _s_thread_offset(r_p, s_p, placeholders) + + # R-side base offset from slicing (e.g. ``R[i*8:i*8+8]`` ⇒ ``i*8``), from + # ``r_p.offset``; sum across all axes on R's local stride-1 storage. + r_off_base = _IntImm("int32", 0) + for _ax, val in r_p.offset.items(): + r_off_base = r_off_base + val + + s_is_shared = s_buf.scope().startswith("shared") + num_bytes = vec_bits // 8 + vec, ptx_type = copy_ptx_form(num_bytes) + space = "shared" if s_is_shared else "global" + + # Honor the load cache hint (cache="nc") only on global -> register loads; + # ``nc`` is a global qualifier, so stores and shared copies stay plain. + _cache, _cache_hints = _ld_cache_config(op_call) + _use_nc = _cache == "nc" and space == "global" and not r_is_src + _l1_evict = _cache_hints.get("l1_evict", "") + _l2_evict = _cache_hints.get("l2_evict", "") + _prefetch = _cache_hints.get("prefetch_size", "") + + total_outer = 1 + for a in outer: + total_outer *= a[0] + + # Recognize the S-side swizzle iter-pattern from atom extents/strides + # (atom = (extent, s_stride, r_mul); a[1] is the S stride per outer round). + swizzle = get_swizzle(s_buf.layout) + swizzle_pattern = None + if swizzle is not None: + swizzle_pattern = try_recognize( + swizzle, + [a[0] for a in outer], + [a[1] for a in outer], + s_off_template, + ) + + class _SwizzleState: + def __init__(self): + self.signed_strides = None + self.base_off = None + + state = _SwizzleState() + + def _setup_swizzle(s_off): + if swizzle_pattern is None: + return + state.signed_strides, state.base_off = emit_init(swizzle_pattern, s_off) + + def _s_iter_off(f, ds, s_off): + if swizzle_pattern is not None: + return emit_iter_offset(swizzle_pattern, state.signed_strides, state.base_off, f) + if swizzle is not None: + return emit_fallback_offset(swizzle, s_off, ds) + return s_off + ds + + # fmt: off + s_zero_indices = [0] * len(s_buf.shape) + + @T.prim_func(check_well_formed=False) + def impl(): + s_off = _substitute_axes(s_off_template, placeholders, sctx) + _setup_swizzle(s_off) + r_local = r_buf.local(*per_thread_r_shape) + # Keep a serial TIR loop and let ptxas unroll; explicit ``T.unroll`` + # replicates per-iter scratch arrays and pressures registers. + for f in range(total_outer): + ds, dr = _outer_const_offsets(outer, f) + s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), _s_iter_off(f, ds, s_off)) + r_ptr = _ptr_off(r_local.ptr_to([0]), r_off_base + dr) + if r_is_src: + T.ptx.st(s_ptr, src=r_ptr, space=space, vec=vec, ptx_type=ptx_type) + elif _use_nc: + T.ptx.ld_global_nc( + s_ptr, + copy_ptx_ld_return_type(ptx_type), + ptx_type, + dst=r_ptr, + vec=vec, + l1_evict=_l1_evict, + l2_evict=_l2_evict, + prefetch_size=_prefetch, + ) + else: + T.ptx.ld( + s_ptr, + copy_ptx_ld_return_type(ptx_type), + ptx_type, + dst=r_ptr, + space=space, + vec=vec, + l1_evict=_l1_evict, + l2_evict=_l2_evict, + prefetch_size=_prefetch, + ) + # fmt: on + import os + + if os.environ.get("R2S_DUMP"): + print("=== emitted impl ===") + print(impl.script()) + return impl diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py new file mode 100644 index 000000000000..34459ea2ba2b --- /dev/null +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py @@ -0,0 +1,221 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Explicit fixed-width vector copy dispatches. + +Cache-semantics config (all variants, read from ``op_call.config``): + +- ``cache``: ``None`` (default, plain ``T.ptx.ld``) or ``"nc"`` (load via + ``T.ptx.ld_global_nc``, i.e. PTX ``ld.global.nc``). Requires a global src. +- ``l1_evict`` / ``l2_evict`` / ``prefetch_size``: L1/L2 cache hint kwargs + forwarded verbatim to the emitted load (same values ``T.ptx.ld`` / + ``T.ptx.ld_global_nc`` accept, e.g. ``"L1::no_allocate"``, + ``"L2::evict_first"``, ``"L2::256B"``). Requires a global src. + +Example:: + + Tx.copy(dst_local[0:8], src_global[base : base + 8], dispatch="vec_256b", + cache="nc", l1_evict="L1::no_allocate", prefetch_size="L2::256B") +""" + +from tvm.arith.analyzer import Analyzer +from tvm.runtime import DataType +from tvm.script import tirx as T +from tvm.tirx import Buffer, PrimFunc +from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch +from tvm.tirx.operator.tile_primitive.registry import DispatchContext +from tvm.tirx.stmt import BufferRegion +from tvm.tirx.tile_primitive import TilePrimitiveCall + +from ._common import copy_ptx_form, copy_ptx_ld_return_type +from .utils import _scope_allowed + + +def _region_start(buffer_region: BufferRegion): + return [r.min for r in buffer_region.region] + + +def _region_elements(buffer_region: BufferRegion): + product = 1 + for r in buffer_region.region: + product *= r.extent + return product + + +def _can_prove_equal(lhs, rhs) -> bool: + if isinstance(lhs, int) and isinstance(rhs, int): + return lhs == rhs + return Analyzer().can_prove_equal(lhs, rhs) + + +def _ptx_space(scope: str) -> str: + if scope.startswith("shared"): + return "shared" + return scope + + +_LD_CACHE_HINT_KEYS = ("l1_evict", "l2_evict", "prefetch_size") + + +def _ld_cache_config(op_call: TilePrimitiveCall) -> tuple[str | None, dict[str, str]]: + """Read the cache-semantics config from ``op_call.config``. + + Returns ``(cache, hints)``: ``cache`` is ``None`` or ``"nc"``, and + ``hints`` maps the ``T.ptx.ld``/``T.ptx.ld_global_nc`` L1/L2 hint kwargs + (``l1_evict``, ``l2_evict``, ``prefetch_size``) to their string values. + """ + cache = op_call.config.get("cache", None) + if cache is not None: + cache = str(cache) + hints: dict[str, str] = {} + for key in _LD_CACHE_HINT_KEYS: + value = op_call.config.get(key, None) + if value is not None and str(value): + hints[key] = str(value) + return cache, hints + + +def _is_forced_vec_copy( + op_call: TilePrimitiveCall, + sctx: DispatchContext, + *, + variant: str, + num_bytes: int, +): + if getattr(op_call, "dispatch", None) != variant: + return False, f"requires explicit dispatch={variant!r}" + if sctx.scope_kind != "thread": + return False, f"expected thread exec_scope, got {sctx.scope_kind}" + + scope_ok, scope_reason = _scope_allowed(op_call, sctx) + if not scope_ok: + return False, scope_reason + + op_call = TilePrimitiveCall.downcast(op_call) + src: Buffer = op_call.src.buffer + dst: Buffer = op_call.dst.buffer + if src.dtype != dst.dtype: + return False, f"dtype mismatch: src={src.dtype}, dst={dst.dtype}" + + elem_bits = DataType(src.dtype).bits + width_bits = num_bytes * 8 + if width_bits % elem_bits != 0: + return False, f"{variant} is not an integral number of {src.dtype} elements" + expected_elements = width_bits // elem_bits + + src_elements = _region_elements(op_call.src) + dst_elements = _region_elements(op_call.dst) + if not _can_prove_equal(src_elements, expected_elements): + return False, f"src region does not contain exactly {expected_elements} elements" + if not _can_prove_equal(dst_elements, expected_elements): + return False, f"dst region does not contain exactly {expected_elements} elements" + + cache, hints = _ld_cache_config(op_call) + if cache not in (None, "nc"): + return False, f"unsupported cache config {cache!r}; expected None or 'nc'" + if (cache is not None or hints) and src.scope() != "global": + return False, ( + "cache/L1/L2 hint config applies to global loads; " + f"requires src scope 'global', got {src.scope()!r}" + ) + return True, None + + +def _emit_forced_vec_copy(op_call: TilePrimitiveCall, _sctx: DispatchContext, num_bytes: int): + op_call = TilePrimitiveCall.downcast(op_call) + src: Buffer = op_call.src.buffer + dst: Buffer = op_call.dst.buffer + src_scope = src.scope() + dst_scope = dst.scope() + src_is_local = src_scope == "local" + dst_is_local = dst_scope == "local" + src_space = _ptx_space(src_scope) + dst_space = _ptx_space(dst_scope) + src_ptr = src.ptr_to(_region_start(op_call.src)) + dst_ptr = dst.ptr_to(_region_start(op_call.dst)) + elem_bits = DataType(src.dtype).bits + n_elements = num_bytes * 8 // elem_bits + vec, ptx_type = copy_ptx_form(num_bytes) + return_type = copy_ptx_ld_return_type(ptx_type) + cache, hints = _ld_cache_config(op_call) + use_nc = cache == "nc" + l1_evict = hints.get("l1_evict", "") + l2_evict = hints.get("l2_evict", "") + prefetch_size = hints.get("prefetch_size", "") + + # fmt: off + @T.prim_func(check_well_formed=False) + def impl(): + if src_is_local: + T.ptx.st(dst_ptr, src=src_ptr, space=dst_space, vec=vec, ptx_type=ptx_type) + elif dst_is_local: + if use_nc: + T.ptx.ld_global_nc( + src_ptr, return_type, ptx_type, dst=dst_ptr, vec=vec, + l1_evict=l1_evict, l2_evict=l2_evict, prefetch_size=prefetch_size, + ) + else: + T.ptx.ld( + src_ptr, return_type, ptx_type, dst=dst_ptr, space=src_space, vec=vec, + l1_evict=l1_evict, l2_evict=l2_evict, prefetch_size=prefetch_size, + ) + else: + tmp = T.alloc_local((n_elements,), src.dtype) + tmp_ptr = tmp.ptr_to([0]) + if use_nc: + T.ptx.ld_global_nc( + src_ptr, return_type, ptx_type, dst=tmp_ptr, vec=vec, + l1_evict=l1_evict, l2_evict=l2_evict, prefetch_size=prefetch_size, + ) + else: + T.ptx.ld( + src_ptr, return_type, ptx_type, dst=tmp_ptr, space=src_space, vec=vec, + l1_evict=l1_evict, l2_evict=l2_evict, prefetch_size=prefetch_size, + ) + T.ptx.st(dst_ptr, src=tmp_ptr, space=dst_space, vec=vec, ptx_type=ptx_type) + # fmt: on + return impl + + +def _register_forced_vec_copy(variant: str, num_bytes: int) -> None: + @register_dispatch( + "copy", + "cuda", + variant=variant, + priority=20, + when=[ + predicate( + f"{variant}_applicable", + _is_forced_vec_copy, + variant=variant, + num_bytes=num_bytes, + ) + ], + ) + def _copy_schedule_forced_vec( + op_call: TilePrimitiveCall, + sctx: DispatchContext, + _num_bytes=num_bytes, + ) -> PrimFunc: + return _emit_forced_vec_copy(op_call, sctx, _num_bytes) + + +_register_forced_vec_copy("vec_256b", 32) +_register_forced_vec_copy("vec_128b", 16) +_register_forced_vec_copy("vec_64b", 8) +_register_forced_vec_copy("vec_32b", 4) +_register_forced_vec_copy("vec_16b", 2) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py index c0e3e7cfdcc8..1485938e1f26 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py @@ -29,7 +29,7 @@ predicate, register_dispatch, ) -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import validate_copy_op from ..exec_scope_utils import single_thread diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py index e57228db6b2e..78179005d524 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py @@ -19,7 +19,7 @@ (SASS: ``LDGSTS``). Shares the partition / layout-alignment algorithm with -``cuda/copy/gmem_smem.py`` (sync ``T.copy`` global ↔ shared); differs at +``cuda/copy/vec_auto_gmem_smem.py`` (sync ``T.copy`` global ↔ shared); differs at emit time only: * direction: ``cp.async`` is global → shared only (hardware restriction). @@ -27,6 +27,13 @@ candidate set is restricted to ``{32, 64, 128}`` bits. * emit: ``T.evaluate(T.ptx.cp_async(dst, src, cp_size))`` instead of the synchronous ``T.cuda.copy_{vec_bits}b(dst, src)``. +* config: ``prefetch_size``, ``predicate`` and ``fill_mode`` are forwarded to + ``T.ptx.cp_async``. This covers predicated zero-fill loads such as FlashMLA + RoPE KV staging while preserving the same layout partitioner. +* config: ``direct=True`` is an explicit thread-scope fast path for callers + that already selected a physically contiguous 4/8/16-byte slice. It emits + one ``cp.async`` from the region starts and bypasses the synthesized + partitioner, matching hand-written per-thread copies. Note: ``cp.async`` does **not** sync at emit time — caller is responsible for ``commit_group`` / ``wait_group`` / ``cta_sync`` plumbing around the @@ -43,7 +50,7 @@ register_dispatch, ) from tvm.tirx.operator.tile_primitive.registry import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..copy._common import ( _TID_AXIS_FOR_SCOPE, @@ -56,8 +63,8 @@ get_swizzle, try_recognize, ) -from ..copy.reg import _all_threads_active, _axis_decl, _ptr_off from ..copy.utils import _is_valid_copy, _scope_allowed +from ..copy.vec_auto_reg import _all_threads_active, _axis_decl, _ptr_off # cp.async is unidirectional: global → shared. _LDGSTS_PAIRS = [("global", "shared*")] @@ -65,6 +72,14 @@ _LDGSTS_VEC_BITS = (128, 64, 32) +def _config_bool(value) -> bool: + if isinstance(value, bool): + return value + if isinstance(value, _IntImm): + return bool(int(value.value)) + return bool(value) + + def _divides_thread_cnt_ldgsts( op_call: TilePrimitiveCall, sctx: DispatchContext ) -> tuple[bool, str | None]: @@ -113,10 +128,49 @@ def _emit_ldgsts(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: g_buf, g_br = src, op_call.src s_buf, s_br = dst, op_call.dst + elem_bits = DataType(src.dtype).bits + prefetch_size = op_call.config.get("prefetch_size", -1) + predicate_expr = op_call.config.get("predicate", -1) + fill_mode = op_call.config.get("fill_mode", "") + + if _config_bool(op_call.config.get("direct", False)): + if sctx.scope_kind != "thread": + raise ValueError("ldgsts direct=True is only valid in thread scope") + n_elements = 1 + for r in s_br.region: + try: + n_elements *= int(r.extent) + except (TypeError, ValueError) as err: + raise ValueError( + f"ldgsts direct=True requires constant extent, got {r.extent}" + ) from err + cp_size = n_elements * elem_bits // 8 + if n_elements * elem_bits % 8 != 0 or cp_size not in (4, 8, 16): + raise ValueError( + f"ldgsts direct=True requires a 4/8/16-byte region, got {n_elements} " + f"elements of {elem_bits} bits" + ) + s_start = [r.min for r in s_br.region] + g_start = [r.min for r in g_br.region] + + @T.prim_func(check_well_formed=False) + def impl(): + T.evaluate( + T.ptx.cp_async( + s_buf.ptr_to(s_start), + g_buf.ptr_to(g_start), + cp_size, + prefetch_size=prefetch_size, + predicate=predicate_expr, + fill_mode=fill_mode, + ) + ) + + return impl + g_region = [(r.min, r.min + r.extent) for r in g_br.region] s_region = [(r.min, r.min + r.extent) for r in s_br.region] - elem_bits = DataType(src.dtype).bits thread_cnt = _thread_cnt(sctx) with sctx.target: @@ -257,7 +311,16 @@ def impl(): s_off = _s_off(f, s_lin) s_ptr = _ptr_off(s_buf.ptr_to(s_zero), s_off) g_ptr = _ptr_off(g_buf.ptr_to(g_zero), g_lin) - T.evaluate(T.ptx.cp_async(s_ptr, g_ptr, cp_size)) + T.evaluate( + T.ptx.cp_async( + s_ptr, + g_ptr, + cp_size, + prefetch_size=prefetch_size, + predicate=predicate_expr, + fill_mode=fill_mode, + ) + ) # cp.async is caller-synced — no cta_sync here (commit_group / # wait_group / cta_sync are the caller's responsibility). # fmt: on diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py index 1acf76bc52aa..6af48623770c 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py @@ -15,35 +15,86 @@ # specific language governing permissions and limitations # under the License. -"""smem->tmem dispatch via tcgen05.cp.32x128b.warpx4. +"""smem->tmem dispatch via ``tcgen05.cp`` — generic planner for all PTX shapes. ``tcgen05.cp`` is inherently async; this dispatch emits the cp loop only and leaves completion signaling (``tcgen05.commit`` against a barrier) to the caller. Callers who want sync semantics should issue ``tcgen05.commit`` themselves after the copy. +Routing +------- +- ``shape=`` config: forces that shape (see ``_CP_SHAPE_MULTICASTS`` below; + atom rows/bits are parsed from the shape name itself). +- No ``shape`` config: the shape is inferred from the buffer layouts — each + candidate (widest atom first) is planned until one validates. Bare warpx4 + copies resolve to ``32x128b.warpx4`` exactly as the historical dispatch. + +Descriptor fields are always derived from the layouts. + +Shape table (PTX ISA 8.8 §9.7.16.9.2) +------------------------------------- +Each ``tcgen05.cp`` shape is ``rows x bits-per-row``; some shapes REQUIRE a +multicast qualifier. Multicast replicates the copied data across warp lane +slabs of the destination TMEM: + +=========== ========================== ====================== ============== +shape multicast t lane pattern t replica +=========== ========================== ====================== ============== +128x256b (none) (128, 1@TLane) — +4x256b (none) (4, 32@TLane) — +128x128b (none) (128, 1@TLane) — +64x128b warpx2::02_13 (required) (64, 1@TLane) (2, 64@TLane) +64x128b warpx2::01_23 (required) (2,32):(64,1)@TLane (2, 32@TLane) +32x128b warpx4 (required) (32, 1@TLane) (4, 32@TLane) +=========== ========================== ====================== ============== + +The ``warpx2`` row→lane mappings follow PTX ISA 8.8 p675 ("each warp in the +warp pair receives half of the data", pairs 02_13 = {0,2},{1,3} and 01_23 = +{0,1},{2,3}): the two warps of a pair mirror each other and the pairs split +the 64 rows in listed order. For 02_13 this is exactly the Layout E / B +datapath organization (Figures 213/207, §9.7.16.10.5: M 0-31 in warp-ranks 0 +and 2, M 32-63 in warp-ranks 1 and 3). ``4x256b`` writes one row per +warp-quadrant datapath (lanes 0/32/64/96). All row→lane mappings and replica +structures are verified bit-exactly on B200 hardware by the round-trip tests +in ``test_tcgen05_cp.py``. + Algorithm --------- -Given ``Tx.copy_async(t_region, s_region)`` where t is in tmem (with -R[4:32@TLane] indicating warpx4 broadcast), and s is in shared memory: +Given ``Tx.copy_async(t_region, s_region)`` where t is in tmem (declared with +the replica pattern of the requested (shape, multicast)), and s is in shared +memory: A. Slice + canonicalize both layouts at the given regions. -B. Verify ``t.replica == [4:32@TLane]`` (warpx4 router). +B. Verify ``t.replica`` matches the shape table entry (empty for + non-multicast shapes). C. Compute permutation that puts TLane first, then TCol stride-descending; apply to t.permute_dims and to s via group + permute_by_groups. D. Canonicalize again. E. Isolate broadcast: split-by-stride-zero on both t and s; their split sequences must match (same distinct prefix prods + broadcast extents). Drop stride-0 iters → ``t_iso`` and ``s_iso``. -F. Group both into ``(32, middle, elem_per_128b)``. Validate: - - t_lane = (32, 1@TLane) - - t_col = (elem_per_128b, 1@TCol) - - s_col = (elem_per_128b, 1) - - s_lane refines into (4, 8) on m axis with strides (SDO_stride, atom_K_stride) - - atom_K_byte ∈ {16, 32, 64, 128} → swizzle_mode 0..3 - - swizzle_mode matches s_buf.layout's SwizzleLayout (if any) +F. Split each side's iters into three segments by grouping the flattened + element space as ``atom_rows x n_mid x elem_per_atom`` + (``elem_per_atom = atom_bits / dtype_bits``). ``*_lane`` covers exactly one + instruction's rows (product == atom_rows); ``*_col`` = one instruction's + own columns; ``*_middle`` = everything between — the outer columns (which + instruction) and, when the region has more rows than one atom (e.g. 4x256b + tiling each warp's first 16 rows: the M=64 Layout-F scatter), the extra row + iters, stepped through the taddr lane half-word. ``t_*`` = tmem side, + ``s_*`` = smem side. + Validate: + - t_lane (row → TMEM lane) matches the shape table lane pattern (above) + - t_col: one row's elements land in contiguous TMEM columns + - s_col: one row = contiguous 16B units in smem; for 256b atoms the two + units may sit at a non-contiguous stride when unswizzled (→ LDO field) + - s_lane: rows sit in smem as (atom_rows/8, 8) groups with strides + (SDO_stride, atom_K_stride); 4x256b is a single 4-row group (SDO=0) + - atom_K_byte ∈ {16, 32, 64, 128} → swizzle_mode 0..3, which must match + s_buf.layout's swizzle (if any) G. Alignment checks: - - t_iso TCol offset ≡ 0 (mod 32-bit) + - t_iso TCol offset ≡ 0 (mod 128-bit) + - t_iso TLane offset folds into the taddr lane half-word - s_iso m offset ≡ 0 (mod 16B for sw=0; mod atom_size for sw>0) - middle iter strides 16B-aligned H. middle 1-1 correspondence (simple-mode): t_middle and s_middle have same @@ -51,7 +102,7 @@ I. Emit: - SmemDescriptor encoded once at SMEM base (hoisted via post_buffer_def_stmt). - Loop over middle iters; each cp uses ``desc.add_16B_offset(init + loop)`` - and writes to ``tmem_addr + t_col0 + Σ i_j * t_step_j``. + and writes to ``tmem_addr + t_addr_off + Σ i_j * t_step_j``. """ import functools @@ -62,12 +113,111 @@ from tvm.runtime import DataType from tvm.script import tirx as T from tvm.tirx import Buffer, PrimFunc -from tvm.tirx.layout import ComposeLayout, SwizzleLayout, TCol, TileLayout, TLane +from tvm.tirx.layout import ComposeLayout, TCol, TileLayout, TLane from tvm.tirx.layout import m as m_axis from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch -from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt, TilePrimitiveCall +from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt +from tvm.tirx.tile_primitive import TilePrimitiveCall + +from ..copy import _single_thread_exec + +# Allowed .multicast qualifiers per shape (PTX ISA 8.8 §9.7.16.9.2): "" = none, +# 64x128b needs an explicit warpx2 pick. +_CP_SHAPE_MULTICASTS = { + "128x256b": ("",), + "4x256b": ("",), + "128x128b": ("",), + "64x128b": ("warpx2::02_13", "warpx2::01_23"), + "32x128b": ("warpx4",), +} + + +def _shape_dims(shape): + """(atom_rows, atom_bits) parsed from the PTX shape name "xb".""" + rows, bits = shape.split("x") + return int(rows), int(bits[:-1]) + + +# Inference order for bare copies: widest atom first (a 256b source also fits +# 128b at 2x instructions); other candidates are mutually exclusive. +_CP_SHAPE_CANDIDATES = ( + ("128x256b", ""), + ("4x256b", ""), + ("128x128b", ""), + ("64x128b", "warpx2::02_13"), + ("64x128b", "warpx2::01_23"), + ("32x128b", "warpx4"), +) + + +def _cp_lane_replica_pattern(shape: str, multicast: str): + """Expected (t_lane pattern, t.replica pattern) for a (shape, multicast). -from ..copy import _is_valid_smem_tmem_copy, _single_thread_exec + Both are lists of ``(extent, stride)`` on the TLane axis. The replica + strides are the warp-slab offsets that receive multicast copies; the lane + pattern is the row → lane mapping of one copy. + + Multicast semantics (PTX ISA 8.8 §9.7.16.9.2, p675): for warpx2 the data + is "multicasted into warp pairs and each warp in the warp pair receives + half of the data" — the two warps OF a pair MIRROR each other (hold the + same half), and the pairs split the 64 rows in listed order: + + - ``warpx2::02_13`` (pairs {0,2}, {1,3}): rows 0-31 to warps 0/2, rows + 32-63 to warps 1/3. One copy occupies lanes 0-63 in row order (warp0 + + warp1), replicated at lane offset +64 (warp2 + warp3) → lane + (64, 1@TLane), replica (2, 64@TLane). This is exactly the Layout E / + Layout B datapath organization (PTX ISA 8.8 Figures 213 and 207: + M 0-31 in warp-ranks 0 and 2, M 32-63 in warp-ranks 1 and 3). + - ``warpx2::01_23`` (pairs {0,1}, {2,3}): rows 0-31 to warps 0/1, rows + 32-63 to warps 2/3 → one copy at lanes 0-31 (rows 0-31) + lanes 64-95 + (rows 32-63), replicated at +32 → lane (2, 32):(64, 1)@TLane, replica + (2, 32@TLane). + + ``4x256b`` writes one row per warp-quadrant datapath: rows 0..3 land at + lanes 0, 32, 64, 96 → lane (4, 32@TLane), no replica. + + All mappings are verified bit-exactly on B200 hardware by the round-trip + tests in ``test_tcgen05_cp.py``. + """ + rows, _ = _shape_dims(shape) + if multicast == "": + if rows == 4: + return [(4, 32)], [] + return [(rows, 1)], [] + if multicast == "warpx4": + return [(32, 1)], [(4, 32)] + if multicast == "warpx2::02_13": + return [(64, 1)], [(2, 64)] + if multicast == "warpx2::01_23": + return [(2, 64), (32, 1)], [(2, 32)] + raise ValueError(f"unknown tcgen05.cp multicast {multicast!r}") + + +def _resolve_cp_shape(op_call: TilePrimitiveCall): + """Resolve (shape, multicast) from an explicit ``shape=`` config.""" + shape = str(op_call.config["shape"]) + multicast = op_call.config.get("multicast") + allowed = _CP_SHAPE_MULTICASTS.get(shape) + if allowed is None: + raise ValueError( + f"unknown tcgen05.cp shape {shape!r}; expected one of {sorted(_CP_SHAPE_MULTICASTS)}" + ) + if multicast is None: + if len(allowed) == 1: + multicast = allowed[0] + else: + raise ValueError( + f"tcgen05.cp shape {shape!r} requires an explicit multicast config; " + f"choose one of {list(allowed)}" + ) + else: + multicast = str(multicast) + if multicast not in allowed: + raise ValueError( + f"illegal multicast {multicast!r} for tcgen05.cp shape {shape!r}; " + f"allowed: {list(allowed)}" + ) + return shape, multicast # ----------------------------------------------------------------------------- @@ -162,53 +312,80 @@ def subgroup(iters): # ----------------------------------------------------------------------------- # Plan (state object) # ----------------------------------------------------------------------------- -def _build_plan(op_call: TilePrimitiveCall, sctx: DispatchContext): +def _build_plan(op_call: TilePrimitiveCall): """Run A..H and return a dispatch plan. Plan fields: - s_buf, t_buf - dtype, dtype_bits - - elem_per_128b, elem_per_32b + - shape, multicast (PTX qualifier strings) + - elem_per_atom, elem_per_32b - SmemSwizzleMode (int) - - SDO_field, atom_K_byte + - LDO_field, SDO_field, atom_K_byte - middle_iters: list of (extent, s_step_16B, t_step_32bcol) - init_off_16B (Expr) - - t_col0 (Expr, TMEM 32-bit col offset for cp's first call) + - t_addr_off (Expr, taddr offset of the first cp: 32-bit col + offset plus the region row offset in the lane half-word) """ op_call = TilePrimitiveCall.downcast(op_call) + if op_call.config.get("shape") is not None: + shape, multicast = _resolve_cp_shape(op_call) + return _plan_for_shape(op_call, shape, multicast) + # No shape config: infer from the buffer layouts. + multicast_cfg = op_call.config.get("multicast") + errors = [] + for shape, multicast in _CP_SHAPE_CANDIDATES: + if multicast_cfg is not None and str(multicast_cfg) != multicast: + continue + try: + return _plan_for_shape(op_call, shape, multicast) + except ValueError as err: + errors.append(f"{shape}{('.' + multicast) if multicast else ''}: {err}") + raise ValueError( + "no tcgen05.cp shape fits the given layouts; candidates rejected:\n " + "\n ".join(errors) + ) + + +def _plan_for_shape(op_call: TilePrimitiveCall, shape: str, multicast: str): + """Run A..I for one (shape, multicast); raises ValueError on any mismatch.""" dst_region, src_region = op_call.args[:2] s_buf: Buffer = src_region.buffer t_buf: Buffer = dst_region.buffer dtype = s_buf.dtype dtype_bits = DataType(dtype).bits - elem_per_128b = 128 // dtype_bits + elem_per_128b = 128 // dtype_bits # elements per 16B descriptor unit elem_per_32b = 32 // dtype_bits + atom_rows, atom_bits = _shape_dims(shape) + elem_per_atom = atom_bits // dtype_bits # per-lane elements per cp # C: slice + canonicalize. s_region = [(r.min, r.min + r.extent) for r in src_region.region] t_region = [(r.min, r.min + r.extent) for r in dst_region.region] - s = s_buf.layout.slice(list(s_buf.shape), s_region).canonicalize() + s_sliced = s_buf.layout.slice(list(s_buf.shape), s_region) + s = s_sliced.canonicalize() t = t_buf.layout.slice(list(t_buf.shape), t_region).canonicalize() - # If s is ComposeLayout (SwizzleLayout∘TileLayout), peel off the swizzle - # for stride analysis; record swizzle_len for cross-check. + # If s is ComposeLayout (swizzle over a tile), peel off the swizzle for + # stride analysis; record the swizzle object for cross-checks. s_swizzle_mode_from_layout = 0 + s_swizzle_obj = None if isinstance(s, ComposeLayout): - s_swizzle_mode_from_layout = int(s.swizzle.swizzle_len) + s_swizzle_obj = s + s_swizzle_mode_from_layout = int(s.swizzle_len) s = s.tile_layout - elif isinstance(s, SwizzleLayout): - raise ValueError("s slice produced bare SwizzleLayout (unexpected)") - # B: warpx4 router check. + # B: replica router check per (shape, multicast). + lane_pattern, replica_pattern = _cp_lane_replica_pattern(shape, multicast) rep = t.replica - if not ( - len(rep) == 1 - and int(rep[0].extent) == 4 - and int(rep[0].stride) == 32 - and rep[0].axis == TLane - ): + rep_ok = len(rep) == len(replica_pattern) and all( + int(r.extent) == e and int(r.stride) == st and r.axis == TLane + for r, (e, st) in zip(rep, replica_pattern) + ) + if not rep_ok: raise ValueError( - f"warpx4 router fail: t.replica = " + f"replica mismatch for tcgen05.cp {shape}" + f"{('.' + multicast) if multicast else ''}: expected " + f"{[f'({e}, {st}@TLane)' for e, st in replica_pattern]}, got t.replica = " f"{[(int(r.extent), int(r.stride), str(r.axis)) for r in rep]}" ) @@ -227,13 +404,20 @@ def _build_plan(op_call: TilePrimitiveCall, sctx: DispatchContext): s_iso = TileLayout.from_iters(keep_s, list(s_p.replica), dict(s_p.offset)) t_iso = TileLayout.from_iters(keep_t, list(t_p.replica), dict(t_p.offset)) - # F: group into (32, middle, elem_per_128b). + # F: group into (atom_rows, middle, elem_per_atom). def shard_prod(lay): return functools.reduce(operator.mul, [int(it.extent) for it in lay.shard], 1) - n_lane, n_col = 32, elem_per_128b - n_mid_t = shard_prod(t_iso) // (n_lane * n_col) - n_mid_s = shard_prod(s_iso) // (n_lane * n_col) + n_lane, n_col = atom_rows, elem_per_atom + atom_elems = n_lane * n_col + for side, total in (("t", shard_prod(t_iso)), ("s", shard_prod(s_iso))): + if total < atom_elems or total % atom_elems: + raise ValueError( + f"{side} region ({total} elems) is not a multiple of one " + f"{shape} atom ({n_lane}x{n_col} elems)" + ) + n_mid_t = shard_prod(t_iso) // atom_elems + n_mid_s = shard_prod(s_iso) // atom_elems t_grp, t_seps = t_iso.group([n_lane, n_mid_t, n_col]) s_grp2, s_seps = s_iso.group([n_lane, n_mid_s, n_col]) t_seps = list(t_seps) @@ -255,34 +439,49 @@ def _canon_segment(iters): t_middle, s_middle = _align_middles(t_middle, s_middle) # F.1: lane / col validation. - if len(t_lane) != 1: - raise ValueError(f"t_lane must canonicalize to single iter, got {t_lane}") + if len(t_lane) != len(lane_pattern) or not all( + int(li.extent) == e and int(li.stride) == st and li.axis == TLane + for li, (e, st) in zip(t_lane, lane_pattern) + ): + raise ValueError( + f"t_lane for tcgen05.cp {shape}{('.' + multicast) if multicast else ''} " + f"must canonicalize to {[f'({e}, {st}@TLane)' for e, st in lane_pattern]}, " + f"got {t_lane}" + ) if len(t_col) != 1: raise ValueError(f"t_col must canonicalize to single iter, got {t_col}") - if len(s_col) != 1: - raise ValueError(f"s_col must canonicalize to single iter, got {s_col}") - li = t_lane[0] - if not (int(li.extent) == 32 and int(li.stride) == 1 and li.axis == TLane): - raise ValueError(f"t_lane must be (32, 1@TLane), got {li}") ci = t_col[0] - if not (int(ci.extent) == elem_per_128b and int(ci.stride) == 1 and ci.axis == TCol): - raise ValueError(f"t_col must be ({elem_per_128b}, 1@TCol), got {ci}") - sci = s_col[0] - if not (int(sci.extent) == elem_per_128b and int(sci.stride) == 1): - raise ValueError(f"s_col must be ({elem_per_128b}, 1, m), got {sci}") + if not (int(ci.extent) == elem_per_atom and int(ci.stride) == 1 and ci.axis == TCol): + raise ValueError(f"t_col must be ({elem_per_atom}, 1@TCol), got {ci}") - # F.2: s_lane → group (4, 8) → (SDO_stride, atom_K_stride) + # F.2: s_lane → (atom_rows/8, 8) groups: atom_K = row stride within an 8-row + # core matrix, SDO = stride between 8-row groups (single group: SDO 0). s_lane_layout = TileLayout.from_iters(s_lane, [], {}) - s_lane_grp, s_lane_seps = s_lane_layout.group([4, 8]) - s_lane_seps = list(s_lane_seps) - blk_4 = list(s_lane_grp.shard[s_lane_seps[0] : s_lane_seps[1]]) - blk_8 = list(s_lane_grp.shard[s_lane_seps[1] : s_lane_seps[2]]) - if len(blk_4) != 1 or len(blk_8) != 1: - raise ValueError( - f"s_lane must group into single iter per block: blk_4={blk_4}, blk_8={blk_8}" - ) - SDO_byte = int(blk_4[0].stride) * dtype_bits // 8 - atom_K_byte = int(blk_8[0].stride) * dtype_bits // 8 + rows_per_group = min(atom_rows, 8) + n_row_groups = atom_rows // rows_per_group + if n_row_groups == 1: + blk_rows = list(_canon_segment(s_lane)) + if len(blk_rows) != 1: + raise ValueError(f"s_lane must canonicalize to a single row iter, got {blk_rows}") + SDO_byte = 0 + atom_K_byte = int(blk_rows[0].stride) * dtype_bits // 8 + else: + s_lane_grp, s_lane_seps = s_lane_layout.group([n_row_groups, rows_per_group]) + s_lane_seps = list(s_lane_seps) + blk_grp = list(s_lane_grp.shard[s_lane_seps[0] : s_lane_seps[1]]) + blk_8 = list(s_lane_grp.shard[s_lane_seps[1] : s_lane_seps[2]]) + if len(blk_grp) != 1 or len(blk_8) != 1: + raise ValueError( + f"s_lane must group into single iter per ({n_row_groups}, {rows_per_group}) " + f"block: blk_{n_row_groups}={blk_grp}, blk_8={blk_8}" + ) + SDO_byte = int(blk_grp[0].stride) * dtype_bits // 8 + atom_K_byte = int(blk_8[0].stride) * dtype_bits // 8 + if SDO_byte % 16 != 0: + raise ValueError( + f"s_lane row-group stride {SDO_byte}B not 16B-aligned " + "(descriptor SBO is encoded in 16B units)" + ) sw_candidates = {16: 0, 32: 1, 64: 2, 128: 3} if atom_K_byte not in sw_candidates: raise ValueError(f"atom_K_byte {atom_K_byte} not in {{16,32,64,128}}") @@ -293,18 +492,100 @@ def _canon_segment(iters): f"{s_swizzle_mode_from_layout} but atom_K_byte={atom_K_byte} " f"implies sw={derived_sw}" ) + if s_swizzle_obj is not None: + # The descriptor's address permutation assumes the canonical mma atom + # family (128-bit atoms of 8 rows); other per_element/atom_len mis-plans. + expected_per_element = (128 // dtype_bits).bit_length() - 1 + if ( + int(s_swizzle_obj.per_element) != expected_per_element + or int(s_swizzle_obj.atom_len) != 3 + ): + raise ValueError( + "swizzle family mismatch: expected canonical mma atom " + f"ComposeLayout(per_element={expected_per_element}, " + f"swizzle_len={derived_sw}, atom_len=3); got " + f"per_element={int(s_swizzle_obj.per_element)}, " + f"atom_len={int(s_swizzle_obj.atom_len)}" + ) + # The hardware walk is the swizzle_inner=True permutation; inner=False + # would be mis-copied. sw==0: both directions identity, don't-care. + if int(s_swizzle_obj.swizzle_len) > 0 and not bool(s_swizzle_obj.swizzle_inner): + raise ValueError( + "swizzle direction mismatch: the tcgen05.cp smem descriptor " + "implements the swizzle_inner=True address permutation " + "(x ^ ((x & outer_mask) >> atom_len)); got a ComposeLayout " + f"with swizzle_inner=False (per_element=" + f"{int(s_swizzle_obj.per_element)}, swizzle_len=" + f"{int(s_swizzle_obj.swizzle_len)}, atom_len=" + f"{int(s_swizzle_obj.atom_len)}), whose mirrored permutation " + "would be silently mis-copied" + ) + + # F.3: s_col = one cp's per-lane columns. 128b: one 16B unit (LDO=0). 256b: + # two 16B units — adjacent when swizzled (LBO=1), else stride in LDO field. + if atom_bits == 128: + if len(s_col) != 1: + raise ValueError(f"s_col must canonicalize to single iter, got {s_col}") + sci = s_col[0] + if not (int(sci.extent) == elem_per_128b and int(sci.stride) == 1): + raise ValueError(f"s_col must be ({elem_per_128b}, 1, m), got {sci}") + LDO_field = 0 + else: + s_col_layout = TileLayout.from_iters(s_col, [], {}) + s_col_grp, s_col_seps = s_col_layout.group([2, elem_per_128b]) + s_col_seps = list(s_col_seps) + blk_u = list(s_col_grp.shard[s_col_seps[0] : s_col_seps[1]]) + blk_e = list(s_col_grp.shard[s_col_seps[1] : s_col_seps[2]]) + if len(blk_u) != 1 or len(blk_e) != 1: + raise ValueError( + f"s_col must group into (2, {elem_per_128b}) 16B units: " + f"units={blk_u}, elems={blk_e}" + ) + if not (int(blk_e[0].extent) == elem_per_128b and int(blk_e[0].stride) == 1): + raise ValueError(f"s_col 16B unit must be ({elem_per_128b}, 1, m), got {blk_e[0]}") + ldo_byte = int(blk_u[0].stride) * dtype_bits // 8 + if ldo_byte % 16 != 0: + raise ValueError(f"s_col unit stride {ldo_byte}B not 16B-aligned") + if derived_sw == 0: + LDO_field = ldo_byte // 16 + else: + if ldo_byte != 16: + raise ValueError( + f"swizzled {atom_bits}b atom requires the two 16B units " + f"adjacent in the linear layout, got stride {ldo_byte}B" + ) + LDO_field = 1 # ignored by hardware for swizzled layouts (assumed 1) analyzer = Analyzer() # G: alignments. - # G.1: t_iso TCol offset ≡ 0 (mod 32-bit element count). + # G.1: tcgen05.cp taddr is 128-bit aligned. The layout offset is in + # element units, so require four complete 32-bit TMEM cells. t_col_offset_expr = 0 for ax, val in t_iso.offset.items(): if ax == TCol: t_col_offset_expr = val break - if not analyzer.can_prove_equal(t_col_offset_expr % elem_per_32b, 0): - raise ValueError(f"t TCol offset {t_col_offset_expr} not provably 32b-aligned") + elem_per_128b_tmem = 4 * elem_per_32b + if not analyzer.can_prove_equal(t_col_offset_expr % elem_per_128b_tmem, 0): + raise ValueError(f"t TCol offset {t_col_offset_expr} not provably 128b-aligned") + + # G.1b: a region row offset folds into the taddr lane half-word ([31:16]). + # The full lane footprint must still fit the 128-lane space. + t_lane_offset_expr = 0 + for ax, val in t_iso.offset.items(): + if ax == TLane: + t_lane_offset_expr = val + break + lane_span = 1 + for it in list(t_p.shard) + list(t_p.replica): + if it.axis == TLane: + lane_span += (int(it.extent) - 1) * int(it.stride) + if not analyzer.can_prove(t_lane_offset_expr + lane_span <= 128): + raise ValueError( + f"t TLane offset {t_lane_offset_expr} overflows the 128-lane space " + f"(footprint spans {lane_span} lanes)" + ) # G.2: s_iso m offset alignment. s_m_offset_expr = 0 @@ -339,30 +620,46 @@ def _canon_segment(iters): n = int(ti.extent) if n == 1: continue - if ti.axis != TCol: - raise ValueError(f"middle[{i}] t axis must be TCol, got {ti.axis}") + if ti.axis == TCol: + if int(ti.stride) % elem_per_32b != 0: + raise ValueError( + f"t_middle[{i}] stride {int(ti.stride)} elems not 32b-aligned " + f"(need multiple of {elem_per_32b})" + ) + t_step = int(ti.stride) // elem_per_32b + elif ti.axis == TLane: + # Lane-tiled atoms (e.g. 4x256b filling each warp's first 16 rows + # — the M=64 Layout-F scatter): step the taddr lane half-word. + t_step = int(ti.stride) << 16 + else: + raise ValueError(f"middle[{i}] t axis must be TCol or TLane, got {ti.axis}") s_stride_byte = int(si.stride) * dtype_bits // 8 if s_stride_byte % 16 != 0: raise ValueError(f"s_middle[{i}] stride {s_stride_byte}B not 16B-aligned") - middle_iters.append((n, s_stride_byte // 16, int(ti.stride) // elem_per_32b)) + middle_iters.append((n, s_stride_byte // 16, t_step)) SDO_field = SDO_byte // 16 init_off_16B = s_m_offset_expr * dtype_bits // 8 // 16 - t_col0 = t_col_offset_expr // elem_per_32b + t_addr_off = t_col_offset_expr // elem_per_32b + if not (isinstance(t_lane_offset_expr, int) and t_lane_offset_expr == 0): + t_addr_off = t_addr_off + t_lane_offset_expr * 65536 return { "s_buf": s_buf, "t_buf": t_buf, "dtype": dtype, "dtype_bits": dtype_bits, - "elem_per_128b": elem_per_128b, + "shape": shape, + "multicast": multicast, + "elem_per_atom": elem_per_atom, "elem_per_32b": elem_per_32b, "swizzle_mode": derived_sw, + "LDO_field": LDO_field, "SDO_field": SDO_field, "atom_K_byte": atom_K_byte, "middle_iters": middle_iters, "init_off_16B": init_off_16B, - "t_col0": t_col0, + "t_addr_off": t_addr_off, } @@ -401,23 +698,50 @@ def _desc_set_addr(desc_val, addr_ptr): return T.bitwise_or(T.bitwise_and(desc_val, T.bitwise_not(T.uint64(0x3FFF))), start_addr) +def _validate_smem_tmem_copy(op_call: TilePrimitiveCall, sctx: DispatchContext): + """Memory-scope envelope only; shape resolution/inference and the detailed + layout validation raise readable ValueErrors in ``_build_plan``.""" + dst_region, src_region = op_call.args[:2] + src: Buffer = src_region.buffer + dst: Buffer = dst_region.buffer + return ( + src.scope().startswith("shared") + and dst.scope() == "tmem" + and src.layout is not None + and dst.layout is not None + # Byte reinterpret is legal (nvfp4 stages sf as uint8, reads fp8); a + # true width change needs decompress, rejected in copy_smem_tmem_impl. + and DataType(src.dtype).bits == DataType(dst.dtype).bits + and dst.allocated_addr is not None + ) + + # ----------------------------------------------------------------------------- # Core impl: emits the cp loop given a plan + cp config. Async only — caller # is responsible for issuing ``tcgen05.commit`` against a barrier if they # need synchronization. # ----------------------------------------------------------------------------- def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: - plan = _build_plan(op_call, sctx) + if op_call.config.get("decompress"): + # fp4/fp6->fp8 in-flight decompression needs dtype-pair plan derivation + # the planner can't lower; reject loudly rather than copy undecompressed. + raise ValueError("tcgen05.cp planner does not support decompress") + # NOTE: descriptor templates use base_offset=0, so the smem buffer base must + # align to the swizzle period (8 * atom_K bytes); align=1024 discharges this. + plan = _build_plan(op_call) s_buf = plan["s_buf"] t_buf = plan["t_buf"] + shape = plan["shape"] + multicast = plan["multicast"] SDO_field = plan["SDO_field"] sw = plan["swizzle_mode"] middle_iters = plan["middle_iters"] init_off_16B = plan["init_off_16B"] - t_col0 = plan["t_col0"] + t_addr_off = plan["t_addr_off"] - # cp ignores LDO for data; non-zero LDO bloats the addr-patch codegen. - LDO_field = 0 + # 128b atoms never read the descriptor LDO (single 16B unit per lane): + # keep the legacy LDO=0 encoding. 256b atoms carry the derived field. + LDO_field = plan["LDO_field"] cta_group = op_call.config.get("cta_group", 1) @@ -429,9 +753,8 @@ def _cp_desc(off_16B): addr = T.ptr_byte_offset(s_buf.ptr_to([0] * s_rank), off_16B * 16, s_buf.dtype) return _desc_set_addr(desc_buf[0], addr) - # Flatten the N-D middle iteration into a single T.unroll. Each iteration's - # per-dim index is (flat // stride) % extent, summed into the t/s offsets. - # Works uniformly for n_mid ∈ {0, 1, 2, ...}; total == 1 (no middle dims) is + # Flatten the N-D middle loop into one T.unroll: per-dim index is + # (flat // div) % extent, summed into the t/s offsets. total == 1 is # special-cased to avoid a degenerate T.unroll(1). total = functools.reduce(operator.mul, [n for n, _, _ in middle_iters], 1) @@ -440,9 +763,9 @@ def _cp_desc(off_16B): @T.prim_func(check_well_formed=False) def impl(): T.ptx.tcgen05.cp( - t_addr[0] + t_col0, + t_addr[0] + t_addr_off, _cp_desc(init_off_16B), - shape="32x128b", cta_group=cta_group, multicast="warpx4", + shape=shape, cta_group=cta_group, multicast=multicast, ) else: def compute_offsets(flat): @@ -461,9 +784,9 @@ def impl(): for flat in T.unroll(total): t_off, s_off = T.meta_var(compute_offsets(flat)) T.ptx.tcgen05.cp( - t_addr[0] + t_col0 + t_off, + t_addr[0] + t_addr_off + t_off, _cp_desc(init_off_16B + s_off), - shape="32x128b", cta_group=cta_group, multicast="warpx4", + shape=shape, cta_group=cta_group, multicast=multicast, ) # fmt: on @@ -477,7 +800,7 @@ def impl(): variant="smem->tmem", priority=10, when=[ - predicate("validate_smem_tmem_copy", _is_valid_smem_tmem_copy), + predicate("validate_smem_tmem_copy", _validate_smem_tmem_copy), predicate("exec_scope", _single_thread_exec), ], ) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py index 27144a8e227e..4d03a5f8bdc5 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py @@ -33,11 +33,10 @@ TileLayout, TLane, tcgen05_atom_layout, - tid_in_wg, tmem_datapath_layout, ) from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import get_st_extent from ..copy import _is_valid_copy, _scope_allowed @@ -81,6 +80,43 @@ def _match_tcgen05_atom_layout(buf): return None +# Compatibility matrix between the TMEM buffer's datapath layout and the +# tcgen05 ld/st atom requested by ``T.copy_async``: +# +# datapath x atom | accepted? | rationale +# ---------------------------- | --------- | -------------------------------- +# D (M=128 full) x .32x32b | yes | full 128 lanes, all 32 per warp +# D (M=128 full) x .16x*b M=64| yes | reads first half-slab (lanes +# | | 0..15 of each warp partition) +# | | — the rest of acc is wasted +# | | for this atom but valid data +# D (M=128 full) x .16x*b M=128| yes | reads all 128 lanes via row=0 +# | | and row=16 PTX issues +# F (M=64 scatter)x .16x*b M=64| yes | canonical pairing - F's row +# | | indexing matches the atom's +# | | scatter access +# F (M=64 scatter)x .16x*b M=128| no | F only writes the low slab; the +# | | high slab (row=16) is garbage +# F (M=64 scatter)x .32x32b | no | F only utilizes 16 of each +# | | warp's 32 lanes +# B (M=64 2x2) x bare atom | no | B splits N into two N/2 +# | | lane-halves; use the dedicated +# | | logical (64, N) Layout B image +# | | handled by +# | | _emit_datapath_b_path +_TMEM_ATOM_COMPAT = { + ("D", "32x32b", 128): True, + ("D", "16x*b", 64): True, + ("D", "16x*b", 128): True, + ("F", "32x32b", 128): False, + ("F", "16x*b", 64): True, + ("F", "16x*b", 128): False, + ("B", "32x32b", 128): False, + ("B", "16x*b", 64): False, + ("B", "16x*b", 128): False, +} + + def _classify_tmem_datapath(tmem_buf): """Return ``(datapath, sub_slab)`` if ``tmem_buf.layout`` matches a known tcgen05 datapath (PTX ISA §9.7.16.10.5), else ``None``. @@ -134,69 +170,56 @@ def _classify_tmem_datapath(tmem_buf): return None -# Compatibility matrix between the TMEM buffer's datapath layout and the -# tcgen05 ld/st atom requested by ``T.copy_async``: -# -# datapath x atom | accepted? | rationale -# ---------------------------- | --------- | -------------------------------- -# D (M=128 full) x .32x32b | yes | full 128 lanes, all 32 per warp -# D (M=128 full) x .16x*b M=64| yes | reads first half-slab (lanes -# | | 0..15 of each warp partition) -# | | — the rest of acc is wasted -# | | for this atom but valid data -# D (M=128 full) x .16x*b M=128| yes | reads all 128 lanes via row=0 -# | | and row=16 PTX issues -# F (M=64 scatter)x .16x*b M=64| yes | canonical pairing - F's row -# | | indexing matches the atom's -# | | scatter access -# F (M=64 scatter)x .16x*b M=128| no | F only writes the low slab; the -# | | high slab (row=16) is garbage -# F (M=64 scatter)x .32x32b | no | F only utilizes 16 of each -# | | warp's 32 lanes -# B (M=64 2x2) x bare atom | no | B splits N into two N/2 -# | | lane-halves; use the dedicated -# | | logical (64, N) Layout B image -# | | handled by -# | | _emit_datapath_b_path -_TMEM_ATOM_COMPAT = { - ("D", "32x32b", 128): True, - ("D", "16x*b", 64): True, - ("D", "16x*b", 128): True, - ("F", "32x32b", 128): False, - ("F", "16x*b", 64): True, - ("F", "16x*b", 128): False, - ("B", "32x32b", 128): False, - ("B", "16x*b", 64): False, - ("B", "16x*b", 128): False, -} +def _tmem_window(tmem_buf, tmem_region, atom_kind, frag_rows, analyzer): + """Resolve a tcgen05 ld/st TMEM operand region to its + ``(width, col_off, sub_slab)`` column window (element units). + The region is ``(frag_rows, width)`` optionally prefixed by point-indexed + dims (a staged view's stage axis), which fold into the TCol offset when the + layout is sliced. ``frag_rows`` (128 or 64) selects datapath D or F; the + datapath x atom pairing is gated by ``_TMEM_ATOM_COMPAT`` (PTX ISA + §9.7.16.10.5) and the sliced window is checked against the datapath layout. -def _check_tmem_layout_for_atom(tmem_buf, atom_kind, frag_rows): - """Raise ``ValueError`` if the TMEM buffer's datapath layout is - incompatible with the requested ``tcgen05`` atom. - - ``atom_kind`` is ``"32x32b"`` or ``"16x*b"``; ``frag_rows`` is the - register-side fragment row count (128 for ``.32x32b`` and ``.16x*b`` - M=128 variants, 64 for ``.16x*b`` M=64). If the buffer's layout is - unrecognized (i.e. it isn't Layout D or Layout F), the dispatch falls - back to the structural assertions below. - - Returns ``(datapath, sub_slab)`` for a recognized compatible layout. + ``sub_slab`` is 0 for Layout D. A Layout F view may occupy either 16-lane + half of each warp's 32-lane TMEM partition; the layout carries that choice + as a ``+16`` TLane offset, and it is reported here so the caller can bias + the PTX row immediate rather than silently dropping the half-slab. """ - classified = _classify_tmem_datapath(tmem_buf) - if classified is None: - return None - datapath, sub_slab = classified - allowed = _TMEM_ATOM_COMPAT.get((datapath, atom_kind, frag_rows), False) - if not allowed: + _, extent = get_st_extent(tmem_region) + for d in range(len(extent) - 2): + assert analyzer.can_prove_equal(extent[d], 1), ( + f"tcgen05 ld/st: leading tmem region dims must be points, got extent {extent[d]}" + ) + assert analyzer.can_prove_equal(extent[-2], frag_rows), ( + f"tcgen05 ld/st: tmem row extent must be {frag_rows}, got {extent[-2]}" + ) + width = extent[-1] + datapath = "D" if int(tmem_buf.shape[-2]) == 128 else "F" + if not _TMEM_ATOM_COMPAT.get((datapath, atom_kind, frag_rows), False): raise ValueError( f"tcgen05 dispatch: TMEM buffer with datapath={datapath!r} is " f"incompatible with atom={atom_kind!r} (frag_rows={frag_rows}). " - f"See PTX ISA §9.7.16.10.5 for datapath/atom pairings; the " - f"buffer was allocated via tmem_pool.alloc(..., " - f"datapath={datapath!r})." + f"See PTX ISA §9.7.16.10.5 for datapath/atom pairings." + ) + window = tmem_buf.layout.slice(tmem_buf.shape, tmem_region.region).canonicalize() + if datapath == "D": + base = TileLayout(S[(frag_rows, width) : (1 @ TLane, 1 @ TCol)]) + else: + base = tmem_datapath_layout("F", 64, width) + expected = TileLayout.from_iters(base.shard, base.replica, window.offset).canonicalize() + tvm.ir.assert_structural_equal(window, expected) + lane_off = int(window.offset.get(TLane, 0)) + if lane_off not in (0, 16): + raise ValueError( + f"tcgen05 dispatch: TMEM window has TLane offset {lane_off}; only 0 or 16 " + "(the Layout F sub-slab selector) are representable in the PTX row immediate." + ) + if lane_off and datapath != "F": + raise ValueError( + f"tcgen05 dispatch: datapath={datapath!r} spans both sub-slabs and cannot " + f"carry a TLane offset (got {lane_off})." ) - return (datapath, sub_slab) + return width, window.offset.get(TCol, 0), lane_off // 16 def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc | None: @@ -224,7 +247,7 @@ def copy_tmem_local_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> P analyzer = Analyzer() elem_size = DataType(local_buf.dtype).bits elem_per_32b = 32 // elem_size - assert len(local_buf.shape) == len(tmem_buf.shape) == 2 + assert len(local_buf.shape) == 2 # Datapath B is identified from the TMEM side before probing ordinary # register atoms. A logical (64, N) Layout B tile is physically a @@ -282,9 +305,7 @@ def _emit_32x32b_path( # local: 128xWIDTH <-> tmem: 128xSHAPE[1] # ``.32x32b`` accesses 32 lanes per warp — the full warp partition — so # the TMEM buffer must be Layout D (M=128 full datapath). Reject Layout F. - _check_tmem_layout_for_atom(tmem_buf, "32x32b", 128) assert analyzer.can_prove_equal(local_buf.shape[0], 128) - assert analyzer.can_prove_equal(tmem_buf.shape[0], 128) # Check width is valid for 32x32b, and determine num width = local_region.region[1].extent @@ -301,40 +322,39 @@ def _emit_32x32b_path( else: raise ValueError(f"Width {width} is not valid for tcgen05.ld/st with shape 32x32b") - tmem_st, tmem_extent = get_st_extent(tmem_region) local_st, local_extent = get_st_extent(local_region) - # tmem layout (128, WIDTH):(1@TLane, 1@TCol) - tmem_layout = TileLayout(S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)]).canonicalize() - # local layout - TileLayout(S[(128, width) : (1 @ tid_in_wg, 1)]).canonicalize() - - tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout) - # local: [0:128, 0:WIDTH] <-> tmem: [0:128, st:st+WIDTH] - assert analyzer.can_prove_equal(tmem_st[0], 0) - assert analyzer.can_prove_equal(tmem_extent[0], 128) + # local: [0:128, 0:WIDTH] <-> tmem window: [0:128, off:off+WIDTH] + tmem_width, offset, _ = _tmem_window(tmem_buf, tmem_region, "32x32b", 128, analyzer) assert analyzer.can_prove_equal(local_st[0], 0) assert analyzer.can_prove_equal(local_extent[0], 128) - offset = tmem_st[1] assert analyzer.can_prove_equal(tvm.tirx.floormod(offset, elem_per_32b), 0) offset_32b = tvm.tirx.floordiv(offset, elem_per_32b) - assert analyzer.can_prove_equal(tmem_extent[1], width), ( - f"tmem_extent[1]: {tmem_extent[1]}, width: {width}" - ) + assert analyzer.can_prove_equal(tmem_width, width), f"tmem width: {tmem_width}, width: {width}" # assert analyzer.can_prove_equal(local_st[1], 0) assert analyzer.can_prove_equal(local_extent[1], width) op = T.ptx.tcgen05.ld if direction == "tmem2local" else T.ptx.tcgen05.st - # fmt: off - @T.prim_func(check_well_formed=False) - def impl(): - local_storage = local_buf.view(local_buf.shape[1] * elem_per_32b, layout=TileLayout(S[num * elem_per_32b])) # noqa: E501 - local_32b = local_storage.view("uint32") - op(tmem_buf.allocated_addr[0], *[local_32b[local_st[1] // elem_per_32b+i] for i in range(num)], shape="32x32b", num=num, row=0, col=offset_32b) # noqa: E501 - # fmt: on + if elem_per_32b == 1: + # Keep 32-bit fragments in source dtype; b32 helper makes a uint32 view change codegen. + # fmt: off + @T.prim_func(check_well_formed=False) + def impl(): + local_storage = local_buf.view(local_buf.shape[1], layout=TileLayout(S[num])) + op(T.uint32(tmem_buf.allocated_addr[0]), *[local_storage[local_st[1]+i] for i in range(num)], shape="32x32b", num=num, row=0, col=offset_32b) # noqa: E501 + # fmt: on + else: + # 16-bit fragments are packed two elements per b32 register operand. + # fmt: off + @T.prim_func(check_well_formed=False) + def impl(): + local_storage = local_buf.view(local_buf.shape[1] * elem_per_32b, layout=TileLayout(S[num * elem_per_32b])) # noqa: E501 + local_32b = local_storage.view("uint32") + op(T.uint32(tmem_buf.allocated_addr[0]), *[local_32b[local_st[1] // elem_per_32b+i] for i in range(num)], shape="32x32b", num=num, row=0, col=offset_32b) # noqa: E501 + # fmt: on return impl @@ -385,42 +405,18 @@ def _emit_16xnb_path( f".16x*b path expects local_buf cols={width_elems}, got {local_buf.shape[1]}" ) - # TMEM-side: structurally classify the buffer's datapath (D or F) and - # reject incompatible pairings. The PTX is identical in either case (the - # warp partition rule and the atom's lane access pattern are baked into - # the hardware); the layout classification just keeps the buffer's - # logical row indexing in sync with the physical TMEM occupation. - classified = _check_tmem_layout_for_atom(tmem_buf, "16x*b", frag_rows) - datapath, sub_slab = classified if classified is not None else (None, 0) - # A .16x*b issue covers one 16-lane half-slab. A 64-row fragment issues - # once; a 128-row fragment issues for both halves. ``sub_slab`` comes from - # the TMEM layout and shifts the first issue to the upper half. + # TMEM-side: resolve the region to its column window; datapath (D or F) + # classification and atom gating happen inside _tmem_window. + tmem_width, col_off, sub_slab = _tmem_window( + tmem_buf, tmem_region, "16x*b", frag_rows, analyzer + ) + # A .16x*b issue covers one 16-lane half-slab. A 64-row fragment issues once; + # a 128-row fragment issues for both halves. ``sub_slab`` comes from the TMEM + # layout and shifts the first issue to the upper half. assert sub_slab + n_slabs <= 2, ( f".16x*b sub_slab={sub_slab} with frag_rows={frag_rows} exceeds the 2 " "sub-slabs of each warp's 32-lane TMEM partition" ) - - if datapath == "F": - # Layout F: buffer shape (64, W), scattered row→lane. - assert analyzer.can_prove_equal(tmem_buf.shape[0], 64), ( - f".16x*b Layout F expects tmem_buf rows=64, got {tmem_buf.shape[0]}" - ) - tmem_rows = 64 - else: - # Layout D (or untagged legacy buffers): shape (128, W), identity. - # The legacy structural check below still fires for untagged buffers - # so we don't silently accept arbitrary layouts. - assert analyzer.can_prove_equal(tmem_buf.shape[0], 128), ( - f".16x*b path expects tmem_buf rows=128, got {tmem_buf.shape[0]}" - ) - if datapath is None: - tmem_layout = TileLayout( - S[(128, tmem_buf.shape[1]) : (1 @ TLane, 1 @ TCol)] - ).canonicalize() - tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), tmem_layout) - tmem_rows = 128 - - tmem_st, tmem_extent = get_st_extent(tmem_region) local_st, local_extent = get_st_extent(local_region) # Rows must span the full frag. The COLUMN extent may be a sub-multiple of @@ -433,17 +429,13 @@ def _emit_16xnb_path( # atom (the common case), num_eff == num and reg offset == 0 (no change). assert analyzer.can_prove_equal(local_st[0], 0) assert analyzer.can_prove_equal(local_extent[0], frag_rows) - assert analyzer.can_prove_equal(tmem_st[0], 0) - assert analyzer.can_prove_equal(tmem_extent[0], frag_rows) # local and tmem column slices must match and divide the atom's full width. - assert analyzer.can_prove_equal(local_extent[1], tmem_extent[1]) + assert analyzer.can_prove_equal(local_extent[1], tmem_width) slice_w = int(local_extent[1]) assert width_elems % slice_w == 0, f"slice width {slice_w} must divide atom width {width_elems}" num_eff = num * slice_w // width_elems regs_eff = regs_per_thread_per_slab * slice_w // width_elems - del tmem_rows # only used for the structural check above - col_off = tmem_st[1] assert analyzer.can_prove_equal(tvm.tirx.floormod(col_off, elem_per_32b), 0) col_off_32b = tvm.tirx.floordiv(col_off, elem_per_32b) local_col_off = local_st[1] @@ -477,7 +469,7 @@ def impl(): for slab in range(n_slabs): reg_base = slab * regs_per_thread_per_slab op( - tmem_buf.allocated_addr[0], + T.uint32(tmem_buf.allocated_addr[0]), *[local_32b[local_reg_base + reg_base + i] for i in range(regs_eff)], shape=shape, num=num_eff, row=(sub_slab + slab) * 16, col=col_off_32b, ) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py index e83cd827c2c1..f5d678567c4d 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py @@ -15,1288 +15,2276 @@ # specific language governing permissions and limitations # under the License. -"""copy_async dispatch variant: tma (unified algorithm). - -One algorithm handles all global↔shared TMA copies, respecting the user's -logical OOB spec through alignment conditions on the reshape. No more -aggressive vs exact family split; ``oob`` only selects the hardware fill -kind (0 = zero, 1 = NaN) in the cuTensorMap. - -Pipeline: - -L1 Canonicalize smem; group gmem by buffer shape and canonicalize within each group; split any - multi-iter gmem group into t separate iters (requires g_st, copy_ext - divisible by the inner-product u); slice smem by copy region; regroup - smem by the "copy shape with ext=1 dropped". -L2 For each ext>1 gmem iter (paired with one smem shard sequence), choose - a contiguous chain prefix of selected smem shards (j from max to 0). - Cut the gmem axis into segments at each selected position; each segment - reduces to Case 1 (has selected → box>1 desc dim) or Case 2 (no - selected → box=1 desc dim). Segment 0 absorbs the G-vs-copy_ext slack - via a non-full copy_range; alignment requires g_st, G divisible by - u_{p_0}. Every unselected shard becomes an issue axis. -L3 Stack desc dims across all gmem iters; nest issue axes as an unrolled - loop; validate hardware constraints (rank≤5, swizzle atom, unit inner - stride). Shrink j and retry on failure; bail out when j=0 fails. -Emit Single unrolled loop over the flat mixed-radix decomposition; each - iter computes (smem offset, per-desc-dim tma coord) and emits one - cp_async_bulk_tensor. Host init emits one cuTensorMapEncodeTiled - (deduped by cache key). +"""CUDA Tensor Memory Accelerator dispatches. + +``tma_auto`` derives the largest hardware-legal TMA box from the shared +memory iteration order. ``tma_explicit`` maps the user's global tensor, +layout, and region directly to one TensorMap and one TMA instruction. + +Both planners construct :class:`TensorMapSpec` in CUDA Driver API order: +dimension zero is innermost and ``global_strides`` omits dimension zero. +Validation, descriptor caching/encoding, prefetch, and PTX emission are +shared after that boundary. """ -from dataclasses import dataclass +import re +from dataclasses import dataclass, replace +from enum import Enum +from itertools import pairwise import tvm from tvm.arith import Analyzer from tvm.script import tirx as T -from tvm.tirx import Buffer, PrimFunc -from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout +from tvm.tirx import Buffer, IntImm, PrimFunc, is_buffer_var +from tvm.tirx.layout import Layout, TileLayout from tvm.tirx.operator.tile_primitive import ( DispatchContext, fail, predicate, register_dispatch, ) -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall -from ..common import validate_copy_op from ..exec_scope_utils import single_thread -from ..tma_utils import SwizzleMode, get_swizzle_mode_from_layout, tma_atom_shape - -# ============================================================================== -# Data types -# ============================================================================== - - -@dataclass(frozen=True) -class GmemIter: - """One gmem logical dim after multi-iter group splitting. +from ..layout_utils import strip_swizzle_to_tile +from ..tma_utils import SwizzleMode, get_swizzle_mode_from_layout - ``shape`` and ``stride`` come from the grouped gmem layout for this - dim. ``copy_start`` / ``copy_ext`` carve out the user-requested - sub-range. ``copy_ext == 1`` collapses the iter into a trivial - coord-only descriptor dim (no smem shards, no issue axes). - """ - shape: object - stride: object - copy_start: object - copy_ext: object +class ProofStatus(Enum): + """Static proof result used by the shared TensorMap validator.""" - @property - def is_ext1(self) -> bool: - return Analyzer().can_prove_equal(self.copy_ext, 1) + PROVEN = "proven" + DISPROVEN = "disproven" + UNKNOWN = "unknown" @dataclass(frozen=True) -class SmemShard: - """One canonicalized smem shard inside a group (after slice+regroup).""" - - extent: object - smem_stride: object +class ValidationFinding: + """One named TensorMap validation rule.""" + rule: str + status: ProofStatus + message: str + repairable: bool = False -@dataclass -class SmemGroup: - """Smem shards paired with a single ext>1 gmem iter, outer→inner. - After L1, each ext>1 gmem iter has a matching smem group whose shards' - extents multiply to the iter's ``copy_ext``. - """ - - shards: list # list[SmemShard], outer→inner - bound_gmem_iter_idx: int - - -@dataclass -class Segment: - """One reshape segment produced by the chain-prefix cut. +@dataclass(frozen=True) +class TensorMapSpec: + """A complete TensorMap plus instruction contract in CUDA API order.""" + + descriptor_dtype: str + descriptor_bits: int + effective_bytes: int + packed_kind: str | None + force_cu_dtype: int + base: object + base_key: str + descriptor_name: str + base_byte_offset: object + global_dims: tuple + global_strides: tuple + inner_stride: object + box_dims: tuple + element_strides: tuple + interleave: int + swizzle: int + l2_promotion: int + oob_fill: int + direction: str + load_mode: str + target_arch: str + coordinates: tuple + gather4: tuple + smem_buffer: Buffer + smem_start: tuple + smem_base_offset: object + mbar: object | None + mbar_is_shared_addr: bool + cta_group: int + cta_mask: object + cache_hint: object + cache_policy: object + use_tma_reduce: object | None + payload_bits: object + transaction_bits: object - ``local_shape * local_stride`` is the axis's gmem span; the - ``local_copy_range`` is where the user-requested slice lives on this - axis. A segment is "selected" when it ends with a chosen smem shard - (→ Case 1: box = selected extent); "trailing" otherwise (→ Case 2: - box = 1). - """ + @property + def rank(self) -> int: + return len(self.global_dims) - local_shape: object - local_stride: object - local_copy_start: object # lo endpoint of local_copy_range - local_copy_extent: object # width of local_copy_range - # ``selected_shard_extent`` is the extent of the selected smem shard at - # this segment's inner end (only meaningful when ``is_selected``). - is_selected: bool - selected_shard_extent: object - # Unselected shards within this segment become issue axes contributing - # to this segment's descriptor dim. Each entry is (extent, u_k) where - # u_k is the shard's gmem-units-per-step value divided by the - # segment's selected u (so coord_advance = u_k directly); see - # ``_segment_issue_contribs``. - unselected_contribs: list # list[(extent, coord_advance, smem_stride)] + @property + def is_packed(self) -> bool: + return self.packed_kind is not None @dataclass(frozen=True) -class DescDim: - """One cuTensorMap descriptor dim.""" +class IssueCoord: + """One mixed-radix contribution from an auto issue axis.""" - shape: object - stride: object # gmem stride (elements, not bytes) - box: object - coord_base: object + dim_idx: int + divisor: object + modulus: object @dataclass(frozen=True) -class IssueAxis: - """One issue axis = one unselected smem shard becoming a loop iter. - - Each iteration advances one desc dim's coord by ``coord_advance`` and - one smem region by ``smem_stride``. ``dim_idx`` is the index of the - owning desc dim in the final ``TmaPlan.dims`` list. - """ +class AutoIssueAxis: + """One unboxed shared-memory iterator lowered as an issue loop.""" extent: object - dim_idx: int - coord_advance: object smem_stride: object + coords: tuple[IssueCoord, ...] @dataclass(frozen=True) -class TmaPlan: - """Final descriptor + loop plan.""" - - swizzle_mode: SwizzleMode - dims: list # list[DescDim], in cuTensorMap outer→inner order - issue_axes: list # list[IssueAxis], outer→inner nesting order - tensor_ptr: object - # Element size used by the cuTensorMap descriptor. Defaults to the - # underlying buffer's dtype size; merge can promote this (e.g. uint8 → - # uint16) when adjacent contiguous dims would exceed boxDim≤256 in the - # native dtype. Strides/extents/boxes in ``dims`` are in this unit. - elem_bytes: int = 1 - elem_dtype: str = "uint8" - - @property - def rank(self) -> int: - return len(self.dims) - - @property - def shape(self) -> list: - return [d.shape for d in self.dims] +class TMAPlan: + """A validated spec plus optional auto issue loops.""" - @property - def box_dim(self) -> list: - return [d.box for d in self.dims] - - @property - def g_strides(self) -> list: - return [d.stride for d in self.dims] + spec: TensorMapSpec + issue_axes: tuple[AutoIssueAxis, ...] = () + shared_layout: TileLayout | None = None - def flatten_total_extent(self) -> object: - total: object = 1 + def issue_extent(self): + total = 1 for axis in self.issue_axes: total = total * axis.extent return total def offsets_and_coords(self, loop_var): - """Decompose ``loop_var`` into (smem offset, per-dim coord vector). + """Return shared offset and CUDA-order coordinates for one issue.""" - Axes are stored outer→inner. The innermost axis has cum=1; each - outer axis's cum is the product of inner axes' extents. - """ - total = 1 - cum_per_axis: list = [None] * len(self.issue_axes) + cumulative = 1 + divisors = [None] * len(self.issue_axes) for idx in range(len(self.issue_axes) - 1, -1, -1): - cum_per_axis[idx] = total - total = total * self.issue_axes[idx].extent - - s_offset: object = 0 - coords: list = [d.coord_base for d in self.dims] - for axis, cum in zip(self.issue_axes, cum_per_axis): - iter_val = tvm.tirx.floormod(tvm.tirx.floordiv(loop_var, cum), axis.extent) - s_offset = s_offset + iter_val * axis.smem_stride - coords[axis.dim_idx] = coords[axis.dim_idx] + iter_val * axis.coord_advance - return s_offset, coords + divisors[idx] = cumulative + cumulative = cumulative * self.issue_axes[idx].extent + + smem_offset = 0 + coordinates = list(self.spec.coordinates) + for axis, flat_divisor in zip(self.issue_axes, divisors): + value = tvm.tirx.floormod(tvm.tirx.floordiv(loop_var, flat_divisor), axis.extent) + smem_offset = smem_offset + value * axis.smem_stride + for contribution in axis.coords: + digit = tvm.tirx.floormod( + tvm.tirx.floordiv(value, contribution.divisor), + contribution.modulus, + ) + coordinates[contribution.dim_idx] = coordinates[contribution.dim_idx] + digit + return smem_offset, coordinates -# ============================================================================== -# Common helpers -# ============================================================================== +@dataclass(frozen=True) +class _Gt: + """One global-layout iterator after the two grouping steps.""" + extent: object + stride: object + smem_idx: int | None + copy_dim: int + global_dim: object + coordinate: object + + +_COMMON_CONFIG = { + "cache_hint", + "cta_group", + "cta_mask", + "mbar", + "mbarrier_addr", + "oob", + "prefetch_tensormap", + "tensormap_l2_promotion", + "tma_dtype", + "use_tma_reduce", +} +_EXPLICIT_CONFIG = _COMMON_CONFIG | {"gather4", "src_selector"} +_FLOAT_DTYPES = {"float16", "float32", "float64", "bfloat16"} +_PROMOTE_DTYPE = {1: ("uint16", 16), 2: ("uint32", 32), 4: ("uint64", 64)} +_REPAIRABLE_RULES = { + "rank", + "global_stride_alignment", + "box_dim", + "inner_box_bytes", +} +# These bounds depend only on runtime tensor arguments and are checked by +# runtime.cuTensorMapEncodeTiled before the CUDA driver encodes the descriptor. +_RUNTIME_VALIDATED_RULES = {"global_dim"} + + +def _proof(predicate, analyzer: Analyzer) -> ProofStatus: + if isinstance(predicate, bool): + return ProofStatus.PROVEN if predicate else ProofStatus.DISPROVEN + if analyzer.can_prove(predicate): + return ProofStatus.PROVEN + if analyzer.can_prove(predicate == 0): + return ProofStatus.DISPROVEN + return ProofStatus.UNKNOWN + + +def _proof_all(analyzer: Analyzer, *conditions) -> ProofStatus: + statuses = [_proof(condition, analyzer) for condition in conditions] + if ProofStatus.DISPROVEN in statuses: + return ProofStatus.DISPROVEN + if ProofStatus.UNKNOWN in statuses: + return ProofStatus.UNKNOWN + return ProofStatus.PROVEN + + +def _proof_equal(lhs, rhs, analyzer: Analyzer) -> ProofStatus: + if analyzer.can_prove_equal(lhs, rhs): + return ProofStatus.PROVEN + return _proof(lhs == rhs, analyzer) + + +def _require_proven(predicate, analyzer: Analyzer, stage: str, detail: str) -> None: + status = _proof(predicate, analyzer) + if status != ProofStatus.PROVEN: + _auto_fail(stage, f"{detail} ({status.value})") + + +def _auto_fail(stage: str, detail: str): + fail( + f'tma_auto stage={stage}: {detail}; use dispatch="tma_explicit" ' + "when the mapping or hardware legality is only known at runtime" + ) -def _to_tile_layout(layout: Layout, shape: list) -> TileLayout: - """Normalize the shared layout so pointer arithmetic always sees a TileLayout.""" - if isinstance(layout, ComposeLayout): - return layout.tile_layout - if isinstance(layout, SwizzleLayout): - return TileLayout(S[tuple(shape)]) - return layout +def _to_tile_layout(layout: Layout, shape) -> TileLayout: + tile = strip_swizzle_to_tile(layout, lambda: list(shape)) + if not isinstance(tile, TileLayout): + raise ValueError(f"expected TileLayout after removing swizzle, got {type(tile).__name__}") + return tile -def _assert_memory_only(layout: TileLayout, label: str) -> None: - for shard in layout.shard: - if not shard.axis.is_memory(): +def _assert_plain_memory_layout(layout: TileLayout, label: str) -> None: + if layout.replica: + raise ValueError(f"{label} layout contains replica iterators") + for iterator in layout.shard: + if not iterator.axis.is_memory(): raise ValueError( - f"TMA {label} layout must be pure memory; saw non-memory axis " - f"{shard.axis} in {layout}" + f"{label} layout must contain only memory iterators; got axis " + f"{iterator.axis.name!r}" ) + for axis, offset in layout.offset.items(): + if not axis.is_memory() and not Analyzer().can_prove_equal(offset, 0): + raise ValueError(f"{label} layout has non-memory offset {axis.name}={offset}") -def _normalize_oob_mode(dtype: str, oob_mode): - """Validate the user-visible ``oob`` contract flag. +def _layout_offset(layout: TileLayout): + value = 0 + for axis, offset in layout.offset.items(): + if axis.is_memory(): + value = value + offset + return Analyzer().simplify(value) - ``None`` / ``"zero"`` → hardware fill kind 0. - ``"nan"`` → hardware fill kind 1 (floating-point only). - """ - if oob_mode is None: - return None - if oob_mode not in ("zero", "nan"): - fail(f"Unsupported TMA oob mode: {oob_mode!r}. Expected None, 'zero', or 'nan'.") - if oob_mode == "nan" and dtype not in ("float16", "float32", "float64", "bfloat16"): - fail("TMA oob='nan' requires a floating-point dtype") - return oob_mode +def _slice_layout(buffer: Buffer, starts, extents, label: str) -> tuple[TileLayout, TileLayout]: + tile = _to_tile_layout(buffer.layout, buffer.shape) + _assert_plain_memory_layout(tile, label) + region = [(start, start + extent) for start, extent in zip(starts, extents)] + sliced = tile.slice(list(buffer.shape), region) + if sliced is None: + raise ValueError( + f"{label} layout cannot be sliced at start={list(starts)}, extent={list(extents)}" + ) + if not isinstance(sliced, TileLayout): + raise ValueError(f"{label} sliced layout is not a TileLayout") + _assert_plain_memory_layout(sliced, f"sliced {label}") + return tile, sliced -def _oob_fill_kind(oob_mode) -> int: - if oob_mode is None or oob_mode == "zero": + +def _slice_global_layout(buffer: Buffer, starts, extents) -> tuple[TileLayout, TileLayout]: + """Slice each logical global dimension without fusing across its boundary. + + ``TileLayout.slice`` canonicalizes the complete layout before grouping it + by ``buffer.shape``. For an unsigned dynamic shape, that can turn + ``(n * C, K)`` into a single wrapping ``uint32`` product, after which the + analyzer correctly refuses to prove the original dimension boundary. + Global TensorMap planning needs those semantic buffer boundaries, so group + the original memory layout first and slice each proven group separately. + """ + + analyzer = Analyzer() + tile = _to_tile_layout(buffer.layout, buffer.shape) + _assert_plain_memory_layout(tile, "global") + try: + grouped, separators = tile.group(list(buffer.shape)) + except (TypeError, ValueError, tvm.error.InternalError) as error: + _auto_fail("global-slice", f"cannot group global layout by buffer shape: {error}") + + sliced_shard = [] + sliced_offset = dict(grouped.offset.items()) + for dim, (start, extent) in enumerate(zip(starts, extents)): + group = TileLayout.from_iters( + grouped.shard[separators[dim] : separators[dim + 1]], + ) + sliced = group.slice([buffer.shape[dim]], [(start, start + extent)]) + if sliced is None or not isinstance(sliced, TileLayout): + _auto_fail( + "global-slice", + f"global dimension {dim} cannot be sliced at start={start}, extent={extent}", + ) + sliced_shard.extend(sliced.shard) + for axis, value in sliced.offset.items(): + sliced_offset[axis] = analyzer.simplify(sliced_offset.get(axis, 0) + value) + + result = TileLayout.from_iters(sliced_shard, grouped.replica, sliced_offset) + _assert_plain_memory_layout(result, "sliced global") + return tile, result + + +def _target_sm(arch: str) -> int: + match = re.search(r"sm_(\d+)", arch or "") + return int(match.group(1)) if match else 0 + + +def _normalize_l2_promotion(value) -> int: + if value is None: + return 2 + if isinstance(value, IntImm): + value = int(value) + if isinstance(value, int): + if 0 <= value <= 3: + return value + fail("TensorMap L2 promotion integer must be in [0, 3]") + names = { + "none": 0, + "L2::none": 0, + "L2::64B": 1, + "L2::128B": 2, + "L2::256B": 3, + } + if value in names: + return names[value] + fail("TensorMap L2 promotion must be None, 0..3, 'none', 'L2::64B', 'L2::128B', or 'L2::256B'") + + +def _normalize_oob(value) -> int: + if value is None or value == "zero": return 0 - if oob_mode == "nan": + if value == "nan": return 1 - raise ValueError(f"Unexpected oob mode: {oob_mode}") + fail(f"unsupported TensorMap oob={value!r}; expected None, 'zero', or 'nan'") + + +def _normalize_cache_hint(cache_hint): + if cache_hint is None: + return "", None + if isinstance(cache_hint, str): + return cache_hint, None + if isinstance(cache_hint, tvm.tirx.Expr): + return "", cache_hint + fail(f"cache_hint must be a string or TIR expression, got {type(cache_hint).__name__}") + + +def _dtype_contract(dtype, tma_dtype=None): + data_type = tvm.DataType(dtype) + name = str(data_type) + if data_type.lanes != 1: + fail(f"TensorMap descriptor dtype must have lanes=1, got {data_type}") + + valid = { + "int8", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + "float16", + "float32", + "float64", + "bfloat16", + "float8_e4m3fn", + "float8_e5m2", + "float4_e2m1fn", + } + if name not in valid: + fail(f"unsupported TensorMap descriptor dtype {name}") + + force_cu_dtype = -1 + if tma_dtype is not None: + if tma_dtype not in ("tf32", "tfloat32"): + fail("tma_dtype must be 'tf32' or 'tfloat32'") + if name != "float32": + fail(f"tma_dtype={tma_dtype!r} requires a float32 global buffer, got {name}") + force_cu_dtype = 11 + + bits = data_type.bits + packed_kind = "16u4_align16" if name == "float4_e2m1fn" else None + return name, bits, (bits + 7) // 8, packed_kind, force_cu_dtype + + +def _elements_to_bytes( + elements, bits: int, analyzer: Analyzer, *, auto: bool, stage: str, label: str +): + total_bits = analyzer.simplify(elements * bits) + divisible = _proof_equal(tvm.tirx.floormod(total_bits, 8), 0, analyzer) + if divisible == ProofStatus.DISPROVEN or (auto and divisible == ProofStatus.UNKNOWN): + message = f"{label}={elements} elements at {bits} bits is not provably byte aligned" + if auto: + _auto_fail(stage, message) + fail(f"tma_explicit stage={stage}: {message}") + return analyzer.simplify(tvm.tirx.floordiv(total_bits, 8)) -def _swizzle_inner_box_fits(dtype: str, swizzle_mode: SwizzleMode, inner_box) -> bool: - """Hardware check: innermost ``boxDim[0] * elementSize`` fits swizzle atom.""" - if swizzle_mode == SwizzleMode.SWIZZLE_NONE: - return True - atom = tma_atom_shape(dtype, swizzle_mode) - return bool(Analyzer().can_prove(inner_box <= atom[-1])) +def _buffer_base(buffer: Buffer, *, auto: bool, stage: str): + analyzer = Analyzer() + layout = _to_tile_layout(buffer.layout, buffer.shape) + layout_offset = _layout_offset(layout) + element_offset = analyzer.simplify(buffer.elem_offset + layout_offset) + byte_offset = _elements_to_bytes( + element_offset, + tvm.DataType(buffer.dtype).bits, + analyzer, + auto=auto, + stage=stage, + label="view base offset", + ) + base = tvm.tirx.handle_add_byte_offset(buffer.data, byte_offset) + return base, f"{hash(buffer.data)}:{base}", byte_offset -def _divides(a, b, analyzer: Analyzer) -> bool: - """Return True when ``a`` divides ``b`` (``b % a == 0``).""" - return analyzer.can_prove_equal(tvm.tirx.floormod(b, a), 0) +def _finding( + findings: list[ValidationFinding], + rule: str, + status: ProofStatus, + message: str, + repairable: bool = False, +): + findings.append(ValidationFinding(rule, status, message, repairable)) -def _simplify_with_var_ranges(exprs, var_ranges, sctx: DispatchContext): - """Simplify expressions under dispatch-context and loop-variable ranges.""" - local_analyzer = Analyzer() - for var, value_range in sctx.var_range_map.items(): - local_analyzer.bind(var, value_range) - for var, extent in var_ranges: - if tvm.ir.is_prim_var(var): - local_analyzer.bind(var, tvm.ir.Range.from_min_extent(0, extent)) - return [local_analyzer.simplify(expr) for expr in exprs] +def validate_tensor_map_spec(spec: TensorMapSpec) -> tuple[ValidationFinding, ...]: + """Validate one TensorMap using the static half of the runtime rule matrix.""" + analyzer = Analyzer() + findings: list[ValidationFinding] = [] + dtype_bits = { + "int8": 8, + "int32": 32, + "int64": 64, + "uint8": 8, + "uint16": 16, + "uint32": 32, + "uint64": 64, + "float16": 16, + "float32": 32, + "float64": 64, + "bfloat16": 16, + "float8_e4m3fn": 8, + "float8_e5m2": 8, + "float4_e2m1fn": 4, + } + + rank_ok = _proof(1 <= spec.rank <= 5, analyzer) + _finding( + findings, + "rank", + rank_ok, + f"descriptor rank must be in [1, 5], got {spec.rank}", + repairable=True, + ) -# ============================================================================== -# L1: layout prerequisite analysis -# ============================================================================== + if spec.descriptor_dtype not in dtype_bits: + _finding( + findings, + "dtype", + ProofStatus.DISPROVEN, + f"unsupported descriptor dtype {spec.descriptor_dtype}", + ) + else: + _finding(findings, "dtype", ProofStatus.PROVEN, "descriptor dtype is supported") + _finding( + findings, + "dtype_bits", + _proof(spec.descriptor_bits == dtype_bits[spec.descriptor_dtype], analyzer), + f"descriptor dtype {spec.descriptor_dtype} requires " + f"{dtype_bits[spec.descriptor_dtype]} bits, got {spec.descriptor_bits}", + ) + _finding( + findings, + "effective_bytes", + _proof(spec.effective_bytes == (spec.descriptor_bits + 7) // 8, analyzer), + f"effective bytes {spec.effective_bytes} do not match " + f"{spec.descriptor_bits}-bit descriptor units", + ) + expected_packed = "16u4_align16" if spec.descriptor_dtype == "float4_e2m1fn" else None + _finding( + findings, + "packed_kind", + _proof(spec.packed_kind == expected_packed, analyzer), + f"descriptor dtype {spec.descriptor_dtype} requires packed_kind={expected_packed!r}, " + f"got {spec.packed_kind!r}", + ) + force_dtype_ok = spec.force_cu_dtype == -1 or ( + spec.force_cu_dtype == 11 + and spec.descriptor_dtype == "float32" + and spec.descriptor_bits == 32 + ) + _finding( + findings, + "forced_cuda_dtype", + _proof(force_dtype_ok, analyzer), + "forced CUDA dtype must be -1, or TFLOAT32 (11) for a float32 descriptor", + ) -@dataclass -class L1Result: - """Output of L1: all gmem iters (ext=1 and ext>1), paired smem groups.""" + array_lengths_ok = ( + len(spec.global_strides) == max(spec.rank - 1, 0) + and len(spec.box_dims) == spec.rank + and len(spec.element_strides) == spec.rank + and len(spec.coordinates) == spec.rank + ) + _finding( + findings, + "array_lengths", + _proof(array_lengths_ok, analyzer), + "TensorMap arrays must contain rank global dimensions, boxes, element strides, " + "and coordinates, plus rank-1 byte strides", + ) - swizzle_mode: SwizzleMode - # All gmem iters in positional order (outer→inner across the splitted - # logical dims). Mix of ext=1 and ext>1. - gmem_iters: list # list[GmemIter] - # One entry per ext>1 gmem iter, in the same order they appear in - # ``gmem_iters`` (but excluding ext=1 iters). - smem_groups: list # list[SmemGroup] + required_alignment = 32 if spec.interleave == 2 or spec.packed_kind else 16 + base_alignment = _proof_equal( + tvm.tirx.floormod(spec.base_byte_offset, required_alignment), 0, analyzer + ) + _finding( + findings, + "global_base_alignment", + base_alignment, + f"view base byte offset {spec.base_byte_offset} must preserve " + f"{required_alignment}B global alignment", + ) + for idx, dim in enumerate(spec.global_dims): + status = _proof_all(analyzer, dim > 0, dim <= (1 << 32)) + _finding( + findings, + "global_dim", + status, + f"globalDim[{idx}]={dim} must be in (0, 2^32]", + ) -def _gmem_layout(g_buf: Buffer) -> TileLayout: - layout = g_buf.layout - if not isinstance(layout, TileLayout): - # cuTensorMap requires a plain memory layout on gmem side. - raise ValueError(f"TMA gmem layout must be a TileLayout; got {type(layout).__name__}") - return layout + for idx, stride in enumerate(spec.global_strides): + nonnegative = _proof(stride >= 0, analyzer) + _finding( + findings, + "global_stride_range", + nonnegative, + f"globalStrides[{idx}]={stride} must be non-negative", + ) + bounded = _proof(stride < (1 << 40), analyzer) + _finding( + findings, + "global_stride_range", + bounded, + f"globalStrides[{idx}]={stride} must be less than 2^40", + ) + alignment = 32 if spec.interleave == 2 or spec.packed_kind else 16 + aligned = _proof_equal(tvm.tirx.floormod(stride, alignment), 0, analyzer) + _finding( + findings, + "global_stride_alignment", + aligned, + f"globalStrides[{idx}]={stride} must be a multiple of {alignment}B", + repairable=True, + ) + inner_stride = _proof_equal(spec.inner_stride, 1, analyzer) + _finding( + findings, + "inner_stride", + inner_stride, + f"innermost global element stride must be 1, got {spec.inner_stride}", + ) -def _canonicalize_smem(s_buf: Buffer) -> TileLayout: - return _to_tile_layout(s_buf.layout, s_buf.shape).canonicalize() + for idx, box in enumerate(spec.box_dims): + status = _proof_all(analyzer, box >= 1, box <= 256) + _finding( + findings, + "box_dim", + status, + f"boxDim[{idx}]={box} must be in [1, 256]", + repairable=True, + ) + for idx, stride in enumerate(spec.element_strides): + status = _proof_all(analyzer, stride >= 1, stride <= 8) + _finding( + findings, + "element_stride", + status, + f"elementStride[{idx}]={stride} must be in [1, 8]", + ) + if spec.element_strides: + _finding( + findings, + "inner_stride", + _proof_equal(spec.element_strides[0], 1, analyzer), + f"innermost elementStride must be 1, got {spec.element_strides[0]}", + ) + _finding( + findings, + "interleave", + _proof(spec.interleave in (0, 1, 2), analyzer), + f"interleave enum {spec.interleave} is invalid", + ) + _finding( + findings, + "swizzle", + _proof(spec.swizzle in (0, 1, 2, 3), analyzer), + f"swizzle enum {spec.swizzle} is invalid", + ) + _finding( + findings, + "l2_promotion", + _proof(spec.l2_promotion in (0, 1, 2, 3), analyzer), + f"L2 promotion enum {spec.l2_promotion} is invalid", + ) + _finding( + findings, + "oob", + _proof(spec.oob_fill in (0, 1), analyzer), + f"OOB enum {spec.oob_fill} is invalid", + ) -def _group_gmem_by_buffer_shape(gmem_raw: TileLayout, buffer_shape: list): - """Group raw gmem first; each group is canonicalized locally before splitting.""" - try: - return gmem_raw.group(list(buffer_shape)) - except Exception as err: - raise ValueError(f"Cannot group gmem layout by buffer shape: {err}") from err + if spec.interleave != 0: + _finding( + findings, + "interleave_rank", + _proof(spec.rank >= 3, analyzer), + "interleaved TensorMaps require rank >= 3", + ) + if spec.interleave == 2: + _finding( + findings, + "interleave_swizzle", + _proof(spec.swizzle == 1, analyzer), + "32B interleave requires the 32B swizzle mode", + ) + if spec.box_dims and spec.interleave == 0 and not spec.is_packed: + inner_bytes = analyzer.simplify(spec.box_dims[0] * spec.effective_bytes) + _finding( + findings, + "inner_box_bytes", + _proof_equal(tvm.tirx.floormod(inner_bytes, 16), 0, analyzer), + f"innermost box is {inner_bytes}B; non-interleaved TensorMaps require a 16B multiple", + repairable=True, + ) + atom_bytes = {1: 32, 2: 64, 3: 128}.get(spec.swizzle) + if atom_bytes is not None: + _finding( + findings, + "swizzle_inner_box", + _proof(inner_bytes <= atom_bytes, analyzer), + f"innermost box {inner_bytes}B exceeds the {atom_bytes}B swizzle atom", + ) -def _canonicalize_gmem_group_shards(shards: list, analyzer: Analyzer) -> list: - canon = TileLayout.from_iters(shards).canonicalize() - return [sh for sh in canon.shard if not analyzer.can_prove_equal(sh.extent, 1)] + if spec.packed_kind == "16u4_align16" and spec.rank and spec.box_dims: + _finding( + findings, + "packed_shape", + _proof_equal(tvm.tirx.floormod(spec.global_dims[0], 128), 0, analyzer), + "packed 16U4 align16 requires globalDim[0] to be a multiple of 128", + ) + _finding( + findings, + "packed_box", + _proof_equal(spec.box_dims[0], 128, analyzer), + "packed 16U4 align16 requires boxDim[0] == 128", + ) + _finding( + findings, + "packed_swizzle", + _proof(spec.swizzle in (0, 3), analyzer), + "packed 16U4 align16 supports only NONE or 128B swizzle in this layout model", + ) + _finding( + findings, + "packed_direction", + _proof(spec.direction == "g2s", analyzer), + "packed 16U4 align16 TensorMaps are load-only", + ) + if spec.oob_fill == 1: + floating = spec.descriptor_dtype in _FLOAT_DTYPES or spec.force_cu_dtype == 11 + _finding( + findings, + "nan_oob_dtype", + _proof(floating and not spec.is_packed, analyzer), + "NaN OOB fill requires a non-packed floating-point descriptor", + ) -def _split_multi_iter_group( - grouped: TileLayout, separators: list, group_idx: int, copy_start, copy_ext, analyzer: Analyzer -): - """Handle a gmem group containing t ≥ 1 iters. - - Returns a list of ``GmemIter`` for this group (outer→inner within the - group). For t=1 → one iter (direct passthrough). For t≥2 → requires - ``copy_start % u == 0`` and ``copy_ext % u == 0`` where - ``u = prod(x_1, ..., x_{t-1})`` (everything except the outermost iter - of this group); splits into t iters where the outermost carries the - partial copy range and the inner t-1 carry full ranges. - """ - start = separators[group_idx] - end = separators[group_idx + 1] - raw_shards = _canonicalize_gmem_group_shards(grouped.shard[start:end], analyzer) - if not raw_shards: - # Degenerate extent-1 group (e.g. batch dim with size 1); emit a - # placeholder iter that's flagged ext=1 by copy_ext==1. - return [GmemIter(shape=1, stride=0, copy_start=copy_start, copy_ext=copy_ext)] - - # Canonicalize ordering: outer→inner is the same order as in ``grouped`` - # (TileLayout.group gives outer-first shards per group by construction). - # t = len(raw_shards). - if len(raw_shards) == 1: - sh = raw_shards[0] - return [ - GmemIter(shape=sh.extent, stride=sh.stride, copy_start=copy_start, copy_ext=copy_ext) - ] - - # Multi-iter group: require alignment. - u: object = 1 - for sh in raw_shards[1:]: - u = u * sh.extent - - if not _divides(u, copy_start, analyzer): - fail( - f"TMA multi-iter gmem group requires copy_start % {u} == 0; got copy_start={copy_start}" + if spec.direction not in ("g2s", "s2g"): + _finding( + findings, + "direction", + ProofStatus.DISPROVEN, + f"invalid TMA direction {spec.direction!r}", ) - if not _divides(u, copy_ext, analyzer): - fail(f"TMA multi-iter gmem group requires copy_ext % {u} == 0; got copy_ext={copy_ext}") - - outer = raw_shards[0] - outer_start = analyzer.simplify(tvm.tirx.floordiv(copy_start, u)) - outer_ext = analyzer.simplify(tvm.tirx.floordiv(copy_ext, u)) - iters = [ - GmemIter( - shape=outer.extent, stride=outer.stride, copy_start=outer_start, copy_ext=outer_ext + if spec.load_mode == "tile_gather4": + _finding( + findings, + "gather4_rank", + _proof(spec.rank == 2, analyzer), + f"gather4 requires a rank-2 TensorMap, got rank {spec.rank}", + ) + _finding( + findings, + "gather4_rows", + _proof(len(spec.gather4) == 4, analyzer), + f"gather4 requires exactly four row coordinates, got {len(spec.gather4)}", + ) + _finding( + findings, + "gather4_box", + _proof_equal(spec.box_dims[1], 1, analyzer) + if spec.rank == 2 and len(spec.box_dims) == 2 + else ProofStatus.DISPROVEN, + "gather4 public axis 0 must map to a hardware row box of one", + ) + _finding( + findings, + "gather4_interleave", + _proof(spec.interleave == 0, analyzer), + "gather4 does not support interleaved TensorMaps", + ) + _finding( + findings, + "target", + _proof(_target_sm(spec.target_arch) >= 100, analyzer), + f"gather4 requires SM100+, got {spec.target_arch!r}", + ) + else: + _finding( + findings, + "coordinate_count", + _proof(len(spec.coordinates) == spec.rank, analyzer), + f"coordinate count {len(spec.coordinates)} must equal rank {spec.rank}", + ) + _finding( + findings, + "target", + _proof(_target_sm(spec.target_arch) >= 90, analyzer), + f"TMA requires SM90+, got {spec.target_arch!r}", ) - ] - for sh in raw_shards[1:]: - iters.append(GmemIter(shape=sh.extent, stride=sh.stride, copy_start=0, copy_ext=sh.extent)) - return iters + _finding( + findings, + "cta_group", + _proof(spec.cta_group in (1, 2), analyzer), + f"cta_group must be 1 or 2, got {spec.cta_group}", + ) + if spec.cta_group == 2: + _finding( + findings, + "target", + _proof(_target_sm(spec.target_arch) >= 100, analyzer), + f"cta_group=2 requires SM100+, got {spec.target_arch!r}", + ) + _finding( + findings, + "mbarrier_address", + _proof(spec.mbar_is_shared_addr, analyzer), + "cta_group=2 requires a precomputed shared mbarrier address", + ) + if spec.direction == "g2s": + _finding( + findings, + "mbar", + _proof(spec.mbar is not None, analyzer), + "global-to-shared TMA requires mbar", + ) + _finding( + findings, + "reduce_direction", + _proof(spec.use_tma_reduce is None, analyzer), + "TMA reduction is only valid for shared-to-global", + ) + else: + _finding( + findings, + "gather_direction", + _proof(spec.load_mode == "tile", analyzer), + "gather4 is only valid for global-to-shared", + ) + _finding( + findings, + "cta_mask_direction", + _proof( + (isinstance(spec.cta_mask, int) and spec.cta_mask == 0) + or (isinstance(spec.cta_mask, IntImm) and int(spec.cta_mask) == 0), + analyzer, + ), + "cta_mask is only valid for global-to-shared", + ) -def _slice_and_canonicalize_smem( - smem_canon: TileLayout, buffer_shape: list, s_st: list, s_ext: list -) -> TileLayout: - region = [(st, st + ext) for st, ext in zip(s_st, s_ext)] - sliced = smem_canon.slice(list(buffer_shape), region) - if sliced is None: - raise ValueError("Cannot slice smem layout for TMA copy") - return sliced.canonicalize() + return tuple(findings) -def _regroup_smem_by_extgt1_shape(sliced_smem: TileLayout, extgt1_shape: list) -> tuple: - """Group the sliced smem layout by the ext>1 copy shape.""" - try: - return sliced_smem.group(list(extgt1_shape)) - except Exception: - return None +def _validation_failures(spec: TensorMapSpec, *, auto: bool): + findings = validate_tensor_map_spec(spec) + return [ + finding + for finding in findings + if finding.status == ProofStatus.DISPROVEN + or ( + auto + and finding.status == ProofStatus.UNKNOWN + and finding.rule not in _RUNTIME_VALIDATED_RULES + ) + ] -def _build_l1_result( - s_buf: Buffer, g_buf: Buffer, g_st: list, g_ext: list, s_st: list, s_ext: list -) -> L1Result: - """Run the L1 pipeline. Raises ``ValueError`` or ``DispatchFail`` on - prerequisite violations; the caller treats these as bail-outs.""" +def _raise_validation(stage: str, failures, *, auto: bool): + details = "; ".join( + f"{finding.rule}: {finding.message} [{finding.status.value}]" for finding in failures + ) + if auto: + _auto_fail(stage, details) + fail(f"tma_explicit stage={stage}: {details}") - analyzer = Analyzer() - swizzle_mode = get_swizzle_mode_from_layout(s_buf.layout) - if swizzle_mode is None: - raise ValueError(f"Cannot determine swizzle mode from layout: {s_buf.layout}") - - smem_canon = _canonicalize_smem(s_buf) - _assert_memory_only(smem_canon, "shared") - gmem_raw = _gmem_layout(g_buf) - _assert_memory_only(gmem_raw, "global") - - # --- gmem: group by buffer shape, then split each group --- - grouped_g, sep_g = _group_gmem_by_buffer_shape(gmem_raw, g_buf.shape) - - gmem_iters: list = [] - # Track which gmem_iters correspond to each original buffer dim to - # later align with the copy region's extent!=1 dims. - per_group_iter_slices: list = [] # list of (start_idx, end_idx) in gmem_iters - for d in range(len(g_buf.shape)): - before = len(gmem_iters) - gmem_iters.extend(_split_multi_iter_group(grouped_g, sep_g, d, g_st[d], g_ext[d], analyzer)) - per_group_iter_slices.append((before, len(gmem_iters))) - - # --- smem: slice then regroup by "copy shape with ext=1 dropped" --- - sliced_smem = _slice_and_canonicalize_smem(smem_canon, s_buf.shape, s_st, s_ext) - - # The post-split "copy shape" (per iter): for ext=1 iters, skip; for - # ext>1 iters, use copy_ext. - extgt1_iter_indices = [i for i, it in enumerate(gmem_iters) if not it.is_ext1] - extgt1_shape = [gmem_iters[i].copy_ext for i in extgt1_iter_indices] - - if not extgt1_shape: - # Entire copy is ext=1 everywhere: single element. Emit one - # trivial DescDim per ext=1 iter at assembly time; no smem groups. - return L1Result(swizzle_mode=swizzle_mode, gmem_iters=gmem_iters, smem_groups=[]) - - regrouped = _regroup_smem_by_extgt1_shape(sliced_smem, extgt1_shape) - if regrouped is None: - raise ValueError(f"Cannot regroup smem layout by ext>1 copy shape {extgt1_shape}") - grouped_s, sep_s = regrouped - - smem_groups: list = [] - for logical_idx, iter_idx in enumerate(extgt1_iter_indices): - start = sep_s[logical_idx] - end = sep_s[logical_idx + 1] - shards = [ - SmemShard(extent=sh.extent, smem_stride=sh.stride) - for sh in grouped_s.shard[start:end] - if not analyzer.can_prove_equal(sh.extent, 1) - ] - smem_groups.append(SmemGroup(shards=shards, bound_gmem_iter_idx=iter_idx)) - - return L1Result(swizzle_mode=swizzle_mode, gmem_iters=gmem_iters, smem_groups=smem_groups) - - -# ============================================================================== -# L2: segment algorithm -# ============================================================================== - - -def _find_contiguous_chain_prefix(smem_groups: list) -> list: - """Return the indices (flat, across groups) of the maximal stride-1 - contiguous chain within the innermost smem group(s). - - Returns a list of (group_idx, shard_idx_within_group) tuples, ordered - from inner to outer. Length of this list = max candidate j. - """ - analyzer = Analyzer() - # Concatenate all shards across groups, innermost→outermost. The chain - # must start with stride 1 and each successive stride equals the product - # of prior extents. - flat = [] - for gi, group in enumerate(smem_groups): - for si, sh in enumerate(group.shards): - flat.append((gi, si, sh)) +def _validate_explicit(spec: TensorMapSpec, stage: str) -> None: + failures = _validation_failures(spec, auto=False) + if failures: + _raise_validation(stage, failures, auto=False) - if not flat: - return [] - chain: list = [] - consumed: set = set() - expected_stride: object = 1 +def _simplify_with_ranges(exprs, var_ranges, sctx: DispatchContext): + analyzer = Analyzer() + for var, extent in var_ranges: + analyzer.bind(var, tvm.ir.Range.from_min_extent(0, extent)) + for var, value in sctx.var_range_map.items(): + if var not in [item[0] for item in var_ranges]: + analyzer.bind(var, value) + return [analyzer.simplify(expr) for expr in exprs] - while True: - for key, (gi, si, sh) in enumerate(flat): - if key in consumed: - continue - if analyzer.can_prove_equal(sh.smem_stride, expected_stride): - consumed.add(key) - chain.append((gi, si)) - expected_stride = analyzer.simplify(expected_stride * sh.extent) - break - else: - break - return chain +def _smem_iter_order(sliced: TileLayout): + analyzer = Analyzer() + active = [ + idx + for idx, iterator in enumerate(sliced.shard) + if not analyzer.can_prove_equal(iterator.extent, 1) + ] + constant_strides = [] + for idx in active: + stride = analyzer.simplify(sliced.shard[idx].stride) + if not isinstance(stride, IntImm): + _auto_fail( + "shared-chain", + f"shared iter {idx} stride={stride} is symbolic, so physical order is unknown", + ) + constant_strides.append((int(stride), idx)) + constant_strides.sort() + order = [idx for _, idx in constant_strides] + if not order: + _auto_fail("shared-chain", "copy region has no non-unit shared memory iterator") + + previous = sliced.shard[order[0]] + _require_proven( + previous.stride == 1, + analyzer, + "shared-chain", + f"smallest shared stride must be 1, got {previous.stride}", + ) + for position, idx in enumerate(order[1:], start=1): + current = sliced.shard[idx] + expected = analyzer.simplify(previous.stride * previous.extent) + _require_proven( + current.stride == expected, + analyzer, + "shared-chain", + f"shared iter {idx} at sorted position {position} has stride={current.stride}, " + f"expected {expected} after extent={previous.extent}", + ) + previous = current + return order -def _distribute_selection(chain: list, smem_groups: list) -> dict: - """From a chain prefix (inner→outer), return a per-group mapping - ``group_idx -> sorted list of selected shard indices (outer→inner)``. +def _copy_region_parts(region): + return [item.min for item in region], [item.extent for item in region] - Only the first ``prefix_len`` chain entries are used; caller slices - ``chain[:prefix_len]`` before passing in. - """ - per_group: dict = {} - for gi, si in chain: - per_group.setdefault(gi, []).append(si) - for gi in per_group: - per_group[gi].sort() - # Each selected position in the chain must be a contiguous prefix of - # the selected positions within that group (no gaps by construction of - # the chain walk). Caller relies on this for u_{p_0} arithmetic. - return per_group - - -def _check_alignment( - gmem_iter: GmemIter, selected_positions: list, shards: list, analyzer: Analyzer -) -> bool: - """Alignment: when j ≥ 1, ``u_{p_0} | G`` and ``u_{p_0} | copy_start``. - - ``p_0`` is the outermost selected position; ``u_{p_0}`` is the product - of shard extents strictly inside ``p_0`` in the group's outer→inner - order. - """ - if not selected_positions: - return True # j=0: trivially ok - p0 = selected_positions[0] - u_p0: object = 1 - for si in range(p0 + 1, len(shards)): - u_p0 = u_p0 * shards[si].extent - u_p0 = analyzer.simplify(u_p0) +def _build_auto_gt( + g_buf: Buffer, + g_starts, + g_extents, + sliced_smem: TileLayout, + sctx: DispatchContext, +): + analyzer = Analyzer() + for var, value in sctx.var_range_map.items(): + analyzer.bind(var, value) + _, sliced_gmem = _slice_global_layout(g_buf, g_starts, g_extents) - if not _divides(u_p0, gmem_iter.shape, analyzer): - return False - if not _divides(u_p0, gmem_iter.copy_start, analyzer): - return False - return True - - -def _build_segments( - gmem_iter: GmemIter, selected_positions: list, shards: list, analyzer: Analyzer -) -> list: - """Cut the gmem axis into segments per the chain-prefix-selection rule. - - Segments (outer→inner): - * Segment 0 (if j≥1): positions [0, p_0], extent G/u_{p_0}, - stride s·u_{p_0}, copy_range [g_st/u_{p_0}, g_st/u_{p_0}+E_0). - * Segment i (i=1..j-1): positions [p_{i-1}+1, p_i], extent E_i, - stride s·u_{p_i}, copy_range [0, E_i). - * Trailing (if p_{j-1} < q-1): positions [p_{j-1}+1, q-1], - extent E_j, stride s·1, copy_range [0, E_j). - * j=0: single "trailing"-style segment covering the whole axis: - extent G, stride s, copy_range [copy_start, copy_start+copy_ext). - """ - G = gmem_iter.shape - s = gmem_iter.stride - copy_start = gmem_iter.copy_start - copy_ext = gmem_iter.copy_ext - q = len(shards) - - def _u_at(k: int) -> object: - """u_k = prod(shards[m].extent for m > k).""" - out: object = 1 - for m in range(k + 1, q): - out = out * shards[m].extent - return analyzer.simplify(out) - - # Helper: for a segment spanning positions [lo, hi] (inclusive), the - # unselected shards inside contribute issue axes on the segment's desc - # dim. Each contribution is (extent, coord_advance, smem_stride) where - # coord_advance (in the segment's desc coord units) = u_k / u_{hi}. - def _unselected_contribs(lo: int, hi: int) -> list: - u_hi = _u_at(hi) - out: list = [] - for m in range(lo, hi + 1): - if m in selected_positions: - continue - u_m = _u_at(m) - coord_advance = ( - analyzer.simplify(tvm.tirx.floordiv(u_m, u_hi)) - if not analyzer.can_prove_equal(u_hi, 1) - else u_m - ) - out.append((shards[m].extent, coord_advance, shards[m].smem_stride)) - return out - - segments: list = [] - - if not selected_positions: - # Case 2 applied to entire axis. The "selected position" at the - # inner end is effectively q-1 with u=1, so unselected contribs - # keep their full u_m as coord_advance. - trailing_contribs = [] - for m in range(q): - trailing_contribs.append((shards[m].extent, _u_at(m), shards[m].smem_stride)) - segments.append( - Segment( - local_shape=G, - local_stride=s, - local_copy_start=copy_start, - local_copy_extent=copy_ext, - is_selected=False, - selected_shard_extent=1, - unselected_contribs=trailing_contribs, + smem_shape = tuple(iterator.extent for iterator in sliced_smem.shard) + if not smem_shape: + _auto_fail("group-global", "sliced shared layout has no iterators") + try: + grouped, (smem_sep, copy_sep) = sliced_gmem.group_many((smem_shape, tuple(g_extents))) + except (TypeError, ValueError, tvm.error.TVMError) as error: + _auto_fail( + "group-global", + f"cannot jointly group global layout by shared shape={smem_shape} " + f"and copy shape={tuple(g_extents)}: {error}", + ) + _assert_plain_memory_layout(grouped, "grouped global") + + shard_to_smem = {} + for smem_idx in range(len(smem_shape)): + for shard_idx in range(smem_sep[smem_idx], smem_sep[smem_idx + 1]): + shard_to_smem[shard_idx] = smem_idx + + records = [None] * len(grouped.shard) + for copy_dim, (start_idx, end_idx) in enumerate(pairwise(copy_sep)): + iterators = list(grouped.shard[start_idx:end_idx]) + product = 1 + for iterator in iterators: + product = analyzer.simplify(product * iterator.extent) + _require_proven( + product == g_extents[copy_dim], + analyzer, + "copy-group", + f"group={copy_dim} iterator product={product} does not equal " + f"copy extent={g_extents[copy_dim]}", + ) + if not iterators: + _auto_fail( + "copy-group", + f"group={copy_dim} start={g_starts[copy_dim]} " + f"extent={g_extents[copy_dim]} has no global iterator", ) + + suffix_products = [None] * len(iterators) + suffix_product = 1 + for local_idx in range(len(iterators) - 1, -1, -1): + suffix_products[local_idx] = suffix_product + suffix_product = analyzer.simplify(suffix_product * iterators[local_idx].extent) + inner_product = suffix_products[0] + _require_proven( + tvm.tirx.floormod(g_starts[copy_dim], inner_product) == 0, + analyzer, + "coordinate", + f"group={copy_dim} start={g_starts[copy_dim]} is not divisible by " + f"inner_product={inner_product}", ) - return segments - - j = len(selected_positions) - p_first = selected_positions[0] - p_last = selected_positions[-1] - - # Segment 0 (outermost selected segment: positions [0, p_0]) - u_p0 = _u_at(p_first) - E0: object = 1 - for m in range(0, p_first + 1): - E0 = E0 * shards[m].extent - E0 = analyzer.simplify(E0) - - seg0_shape = analyzer.simplify(tvm.tirx.floordiv(G, u_p0)) - seg0_stride = analyzer.simplify(s * u_p0) - seg0_copy_start = analyzer.simplify(tvm.tirx.floordiv(copy_start, u_p0)) - segments.append( - Segment( - local_shape=seg0_shape, - local_stride=seg0_stride, - local_copy_start=seg0_copy_start, - local_copy_extent=E0, - is_selected=True, - selected_shard_extent=shards[p_first].extent, - unselected_contribs=_unselected_contribs(0, p_first), + _require_proven( + tvm.tirx.floormod(g_buf.shape[copy_dim], inner_product) == 0, + analyzer, + "global-shape", + f"group={copy_dim} tensor extent={g_buf.shape[copy_dim]} is not divisible " + f"by inner_product={inner_product}", ) - ) - # Inner selected segments (i=1..j-1): positions [p_{i-1}+1, p_i] - for i in range(1, j): - lo = selected_positions[i - 1] + 1 - hi = selected_positions[i] - Ei: object = 1 - for m in range(lo, hi + 1): - Ei = Ei * shards[m].extent - Ei = analyzer.simplify(Ei) - u_pi = _u_at(hi) - segments.append( - Segment( - local_shape=Ei, - local_stride=analyzer.simplify(s * u_pi), - local_copy_start=0, - local_copy_extent=Ei, - is_selected=True, - selected_shard_extent=shards[hi].extent, - unselected_contribs=_unselected_contribs(lo, hi), + for local_idx, iterator in enumerate(iterators): + shard_idx = start_idx + local_idx + smem_idx = shard_to_smem.get(shard_idx) + if analyzer.can_prove_equal(iterator.extent, 1): + # Unit copy dimensions remain real TensorMap dimensions, but + # they do not contribute shared-memory payload. + smem_idx = None + elif smem_idx is None: + _require_proven( + iterator.extent == 1, + analyzer, + "coordinate-only-dim", + f"global iterator={shard_idx} copy group={copy_dim} extent={iterator.extent} " + "has no shared-memory iterator", + ) + if local_idx == 0: + global_dim = analyzer.simplify( + tvm.tirx.floordiv(g_buf.shape[copy_dim], inner_product) + ) + coordinate = analyzer.simplify(tvm.tirx.floordiv(g_starts[copy_dim], inner_product)) + else: + global_dim = iterator.extent + coordinate = 0 + records[shard_idx] = _Gt( + extent=iterator.extent, + stride=iterator.stride, + smem_idx=smem_idx, + copy_dim=copy_dim, + global_dim=global_dim, + coordinate=coordinate, ) + + if any(record is None for record in records): + _auto_fail("copy-group", "global iterator was not assigned to a copy-region dimension") + return tuple(records) + + +def _auto_spec_for_prefix( + *, + op_call, + sctx, + direction, + s_buf, + g_buf, + s_starts, + smem_layout_offset, + sliced_smem, + smem_order, + gt_records, + prefix, + swizzle, + descriptor, + runtime, +): + analyzer = Analyzer() + selected_smem = set(smem_order[:prefix]) + gt_by_smem = {idx: [] for idx in range(len(sliced_smem.shard))} + detached_gt = [] + for gt_idx, gt in enumerate(gt_records): + if gt.smem_idx is None: + detached_gt.append(gt_idx) + else: + gt_by_smem[gt.smem_idx].append((gt_idx, gt)) + + cuda_gt_indices = [] + for smem_idx in smem_order: + cuda_gt_indices.extend(idx for idx, _ in reversed(gt_by_smem[smem_idx])) + unit_smem = [ + idx + for idx, iterator in enumerate(sliced_smem.shard) + if analyzer.can_prove_equal(iterator.extent, 1) + ] + for smem_idx in unit_smem: + cuda_gt_indices.extend(idx for idx, _ in reversed(gt_by_smem[smem_idx])) + cuda_gt_indices.extend(detached_gt) + if len(cuda_gt_indices) != len(gt_records): + _auto_fail( + "descriptor-order", + "not every global iterator maps to shared data or a coordinate-only dimension", ) - # Trailing (if p_{j-1} < q-1): positions [p_{j-1}+1, q-1] - if p_last < q - 1: - Ej: object = 1 - for m in range(p_last + 1, q): - Ej = Ej * shards[m].extent - Ej = analyzer.simplify(Ej) - # For trailing, every position is unselected; "selected u" at the - # inner end is u_{q-1} = 1, so coord_advance = u_m. - trailing_contribs = [] - for m in range(p_last + 1, q): - trailing_contribs.append((shards[m].extent, _u_at(m), shards[m].smem_stride)) - segments.append( - Segment( - local_shape=Ej, - local_stride=s, - local_copy_start=0, - local_copy_extent=Ej, - is_selected=False, - selected_shard_extent=1, - unselected_contribs=trailing_contribs, + gt_to_dim = {gt_idx: dim_idx for dim_idx, gt_idx in enumerate(cuda_gt_indices)} + ordered = [gt_records[idx] for idx in cuda_gt_indices] + global_dims = tuple(gt.global_dim for gt in ordered) + coordinates = tuple(gt.coordinate for gt in ordered) + box_dims = tuple(gt.extent if gt.smem_idx in selected_smem else 1 for gt in ordered) + + descriptor_dtype, descriptor_bits, effective_bytes, packed_kind, force_cu_dtype = descriptor + full_byte_strides = [ + _elements_to_bytes( + gt.stride, + descriptor_bits, + analyzer, + auto=True, + stage="descriptor-stride", + label=f"global stride for copy group {gt.copy_dim}", + ) + for gt in ordered + ] + inner_stride = ordered[0].stride + global_strides = tuple(full_byte_strides[1:]) + + issue_axes = [] + for smem_idx in smem_order[prefix:]: + iterator = sliced_smem.shard[smem_idx] + group = gt_by_smem[smem_idx] + group_product = 1 + for _, gt in group: + group_product = analyzer.simplify(group_product * gt.extent) + _require_proven( + group_product == iterator.extent, + analyzer, + "issue-axis", + f"shared group={smem_idx} extent={iterator.extent} maps to " + f"global product={group_product}", + ) + contributions = [] + inner_product = 1 + local = [] + for gt_idx, gt in reversed(group): + local.append( + IssueCoord( + dim_idx=gt_to_dim[gt_idx], + divisor=inner_product, + modulus=gt.extent, + ) + ) + inner_product = analyzer.simplify(inner_product * gt.extent) + contributions.extend(reversed(local)) + issue_axes.append( + AutoIssueAxis( + extent=iterator.extent, + smem_stride=iterator.stride, + coords=tuple(contributions), ) ) - return segments + issue_count = 1 + for axis in issue_axes: + issue_count = analyzer.simplify(issue_count * axis.extent) + box_elements = 1 + for box in box_dims: + box_elements = analyzer.simplify(box_elements * box) + transaction_bits = analyzer.simplify(box_elements * descriptor_bits) + payload_bits = analyzer.simplify(transaction_bits * issue_count) + + s_elements = 1 + for iterator in sliced_smem.shard: + s_elements = analyzer.simplify(s_elements * iterator.extent) + expected_bits = analyzer.simplify(s_elements * tvm.DataType(s_buf.dtype).bits) + _require_proven( + payload_bits == expected_bits, + analyzer, + "byte-equivalence", + f"prefix={prefix} payload={payload_bits} bits, expected={expected_bits} bits", + ) + base, base_key, base_byte_offset = _buffer_base(g_buf, auto=True, stage="global-base") + spec = TensorMapSpec( + descriptor_dtype=descriptor_dtype, + descriptor_bits=descriptor_bits, + effective_bytes=effective_bytes, + packed_kind=packed_kind, + force_cu_dtype=force_cu_dtype, + base=base, + base_key=base_key, + descriptor_name=g_buf.name, + base_byte_offset=base_byte_offset, + global_dims=global_dims, + global_strides=global_strides, + inner_stride=inner_stride, + box_dims=box_dims, + element_strides=(1,) * len(global_dims), + interleave=0, + swizzle=swizzle.value, + l2_promotion=runtime["l2_promotion"], + oob_fill=runtime["oob_fill"], + direction=direction, + load_mode="tile", + target_arch=sctx.target.arch, + coordinates=coordinates, + gather4=(), + smem_buffer=s_buf, + smem_start=tuple(s_starts), + smem_base_offset=analyzer.simplify(s_buf.elem_offset + smem_layout_offset), + mbar=runtime["mbar"], + mbar_is_shared_addr=runtime["mbar_is_shared_addr"], + cta_group=runtime["cta_group"], + cta_mask=runtime["cta_mask"], + cache_hint=runtime["cache_hint"], + cache_policy=runtime["cache_policy"], + use_tma_reduce=runtime["use_tma_reduce"], + payload_bits=payload_bits, + transaction_bits=transaction_bits, + ) + return TMAPlan( + spec=spec, + issue_axes=tuple(issue_axes), + shared_layout=_to_tile_layout(s_buf.layout, s_buf.shape), + ) -# ============================================================================== -# L3: assembly + hardware constraint validation + shrink -# ============================================================================== +def _validate_auto_shared_mapping( + plan: TMAPlan, + sliced_smem: TileLayout, + swizzle: SwizzleMode, + sctx: DispatchContext, +) -> None: + """Prove that issue loops partition the sliced shared region exactly once.""" -def _assemble_plan( - l1: L1Result, per_iter_selected: dict, chain: list, g_buf: Buffer, analyzer: Analyzer -) -> TmaPlan: - """Build the final ``TmaPlan`` by stacking desc dims from all gmem iters. + analyzer = Analyzer() + for var, value in sctx.var_range_map.items(): + analyzer.bind(var, value) + spec = plan.spec + shared_bits = tvm.DataType(spec.smem_buffer.dtype).bits + start_elements = spec.smem_base_offset + start_bytes = _elements_to_bytes( + start_elements, + shared_bits, + analyzer, + auto=True, + stage="shared-pointer", + label="shared slice pointer", + ) + atom_alignment = {0: 16, 1: 32, 2: 64, 3: 128}[swizzle.value] + _require_proven( + tvm.tirx.floormod(start_bytes, atom_alignment) == 0, + analyzer, + "shared-pointer", + f"shared slice offset={start_bytes}B must preserve {atom_alignment}B alignment", + ) - Emission (natural) order: - * ext=1 gmem iters (in positional order) → one desc dim each (box=1). - * ext>1 gmem iters (in positional order): for each, segments in - outer→inner order produce desc dims; selected segments contribute - box>1 dims, trailing contributes a box=1 dim. + transaction_elements = analyzer.simplify(tvm.tirx.floordiv(spec.transaction_bits, shared_bits)) + _require_proven( + tvm.tirx.floormod(spec.transaction_bits, shared_bits) == 0, + analyzer, + "issue-coverage", + f"transaction={spec.transaction_bits} bits is not an integral number of " + f"{shared_bits}-bit shared elements", + ) + covered = transaction_elements + for axis_idx, axis in enumerate(plan.issue_axes): + _require_proven( + axis.smem_stride == covered, + analyzer, + "issue-coverage", + f"issue axis={axis_idx} stride={axis.smem_stride} does not follow " + f"the covered prefix={covered}", + ) + covered = analyzer.simplify(covered * axis.extent) + + sliced_elements = 1 + for iterator in sliced_smem.shard: + sliced_elements = analyzer.simplify(sliced_elements * iterator.extent) + _require_proven( + covered == sliced_elements, + analyzer, + "issue-coverage", + f"issue loops plus one box cover {covered} shared elements, " + f"but the sliced region contains {sliced_elements}", + ) - Then we **reorder** the desc dims so: - * All box=1 dims (ext=1 iters and trailing segments) come first, in - natural order. - * All box>1 dims (selected segments) come last, in the reverse of - the chain order — i.e. the outermost selected shard in the chain - walk becomes the outermost box>1 desc dim, and the innermost - selected shard (chain[0]) becomes the innermost desc dim. This - matches how the TMA hardware writes the tile into swizzled smem: - the innermost box dim (stride = 1 in gmem, ideally stride = 1 in - smem too) must align with the innermost smem atom axis. - Issue axes' ``dim_idx`` are remapped to the new positions. - """ +def _remap_issue_axes_after_remove(issue_axes, removed_idx): + remapped = [] + for axis in issue_axes: + coords = [] + for coord in axis.coords: + if coord.dim_idx == removed_idx: + raise ValueError("attempted to remove an issue-driven descriptor dimension") + coords.append( + replace( + coord, + dim_idx=coord.dim_idx - 1 if coord.dim_idx > removed_idx else coord.dim_idx, + ) + ) + remapped.append(replace(axis, coords=tuple(coords))) + return tuple(remapped) - dims_natural: list = [] - origins: list = [] # parallel to dims_natural: 'ext1' | 'trailing' | ('selected', chain_idx) - issue_axes_natural: list = [] - # --- First pass: ext=1 iters --- - for _, it in enumerate(l1.gmem_iters): - if not it.is_ext1: - continue - dims_natural.append( - DescDim(shape=it.shape, stride=it.stride, box=1, coord_base=it.copy_start) +def _auto_inner_dimension_is_legal(spec, global_dim, box_dim, analyzer: Analyzer) -> bool: + """Check rules that can change when an auto dimension becomes innermost.""" + + if spec.interleave != 0: + return True + if spec.is_packed: + checks = ( + _proof_equal(tvm.tirx.floormod(global_dim, 128), 0, analyzer), + _proof_equal(box_dim, 128, analyzer), ) - origins.append("ext1") - - # --- Second pass: ext>1 iters --- - for gi, group in enumerate(l1.smem_groups): - iter_idx = group.bound_gmem_iter_idx - gmem_iter = l1.gmem_iters[iter_idx] - shards = group.shards - selected_positions = per_iter_selected.get(gi, []) - segments = _build_segments(gmem_iter, selected_positions, shards, analyzer) - - # For each selected position in this group, pre-compute its chain index. - selected_chain_idx: dict = {} - for p in selected_positions: - for ci, (cgi, csi) in enumerate(chain): - if cgi == gi and csi == p: - selected_chain_idx[p] = ci - break - - for i_seg, seg in enumerate(segments): - dim_idx = len(dims_natural) - box = seg.selected_shard_extent if seg.is_selected else 1 - dims_natural.append( - DescDim( - shape=seg.local_shape, - stride=seg.local_stride, - box=box, - coord_base=seg.local_copy_start, + else: + inner_bytes = analyzer.simplify(box_dim * spec.effective_bytes) + checks = [_proof_equal(tvm.tirx.floormod(inner_bytes, 16), 0, analyzer)] + atom_bytes = {1: 32, 2: 64, 3: 128}.get(spec.swizzle) + if atom_bytes is not None: + checks.append(_proof(inner_bytes <= atom_bytes, analyzer)) + return all(status == ProofStatus.PROVEN for status in checks) + + +def _canonicalize_auto_plan(plan: TMAPlan) -> TMAPlan: + """Canonicalize auto-only TensorMap dimensions without changing byte addresses.""" + + analyzer = Analyzer() + current = plan + while True: + spec = current.spec + full_strides = [spec.effective_bytes * spec.inner_stride, *spec.global_strides] + issue_dims = {coord.dim_idx for axis in current.issue_axes for coord in axis.coords} + + # A unit global dimension carries no address information for an + # in-bounds tma_auto copy. Non-innermost units can always disappear; + # an innermost unit can disappear only when the next dimension already + # has the implicit innermost byte stride. + removed = False + if spec.rank > 1: + for dim_idx in range(spec.rank): + if dim_idx in issue_dims: + continue + checks = ( + _proof_equal(spec.global_dims[dim_idx], 1, analyzer), + _proof_equal(spec.box_dims[dim_idx], 1, analyzer), + _proof_equal(spec.element_strides[dim_idx], 1, analyzer), ) - ) - if seg.is_selected: - # Selected segments are emitted in the same order as - # selected_positions (Segment 0 anchors p_0, etc.), so - # i_seg directly indexes selected_positions for selected - # segments. Trailing segments don't anchor any selection. - p_anchor = selected_positions[i_seg] - origins.append(("selected", selected_chain_idx[p_anchor])) - else: - origins.append("trailing") - # Segment's unselected shards become issue axes on this dim. - for extent, coord_advance, smem_stride in seg.unselected_contribs: - issue_axes_natural.append( - IssueAxis( - extent=extent, - dim_idx=dim_idx, - coord_advance=coord_advance, - smem_stride=smem_stride, - ) + if any(status != ProofStatus.PROVEN for status in checks): + continue + if dim_idx == 0: + if _proof_equal( + full_strides[1], full_strides[0], analyzer + ) != ProofStatus.PROVEN or not _auto_inner_dimension_is_legal( + spec, + spec.global_dims[1], + spec.box_dims[1], + analyzer, + ): + continue + + global_dims = list(spec.global_dims) + box_dims = list(spec.box_dims) + coordinates = list(spec.coordinates) + element_strides = list(spec.element_strides) + for values in (global_dims, box_dims, coordinates, element_strides): + values.pop(dim_idx) + full_strides.pop(dim_idx) + current = replace( + current, + spec=replace( + spec, + global_dims=tuple(global_dims), + global_strides=tuple(full_strides[1:]), + box_dims=tuple(box_dims), + coordinates=tuple(coordinates), + element_strides=tuple(element_strides), + ), + issue_axes=_remap_issue_axes_after_remove(current.issue_axes, dim_idx), ) + removed = True + break + if removed: + continue - # --- Permute: box=1 first (natural order), box>1 last (chain DESC) --- - non_sel_indices = [ - idx for idx, o in enumerate(origins) if not (isinstance(o, tuple) and o[0] == "selected") - ] - sel_entries = [ - (idx, o[1]) for idx, o in enumerate(origins) if isinstance(o, tuple) and o[0] == "selected" - ] - sel_entries.sort(key=lambda x: -x[1]) # chain index descending = outer selected first - new_order = non_sel_indices + [idx for idx, _ in sel_entries] - old_to_new = {old: new for new, old in enumerate(new_order)} - - dims = [dims_natural[old] for old in new_order] - issue_axes = [ - IssueAxis( - extent=ax.extent, - dim_idx=old_to_new[ax.dim_idx], - coord_advance=ax.coord_advance, - smem_stride=ax.smem_stride, - ) - for ax in issue_axes_natural - ] + # Flatten an adjacent pair when the inner dimension is copied in full + # from coordinate zero and the outer byte stride follows it + # contiguously. The outer dimension may have a partial box and a + # non-zero coordinate; both are scaled into the flattened dimension. + merged = False + promoted_for_merge = False + for inner_idx in range(spec.rank - 1): + outer_idx = inner_idx + 1 + if inner_idx in issue_dims or outer_idx in issue_dims: + continue + merged_global_dim = analyzer.simplify( + spec.global_dims[inner_idx] * spec.global_dims[outer_idx] + ) + merged_box_dim = analyzer.simplify(spec.box_dims[inner_idx] * spec.box_dims[outer_idx]) + checks = ( + _proof_equal(spec.box_dims[inner_idx], spec.global_dims[inner_idx], analyzer), + _proof_equal(spec.coordinates[inner_idx], 0, analyzer), + _proof_equal( + full_strides[outer_idx], + spec.global_dims[inner_idx] * full_strides[inner_idx], + analyzer, + ), + _proof_equal(spec.element_strides[inner_idx], 1, analyzer), + _proof_equal(spec.element_strides[outer_idx], 1, analyzer), + _proof_all(analyzer, merged_global_dim > 0, merged_global_dim <= (1 << 32)), + ) + if inner_idx == 0 and not _auto_inner_dimension_is_legal( + spec, merged_global_dim, merged_box_dim, analyzer + ): + continue + if any(status != ProofStatus.PROVEN for status in checks): + continue + box_status = _proof_all(analyzer, merged_box_dim >= 1, merged_box_dim <= 256) + if box_status != ProofStatus.PROVEN: + # Preserve the innermost contiguous-chain boundary. Skipping a + # box-size-blocked inner merge and merging an outer pair instead + # changes the TensorMap tiling even though one byte-preserving + # descriptor-unit promotion may make this merge legal. + if ( + inner_idx == 0 + and spec.rank > 2 + and _proof(merged_box_dim > 256, analyzer) == ProofStatus.PROVEN + and _proof_equal( + spec.box_dims[outer_idx], spec.global_dims[outer_idx], analyzer + ) + == ProofStatus.PROVEN + and _proof_equal(spec.coordinates[outer_idx], 0, analyzer) == ProofStatus.PROVEN + ): + promoted = _promote_auto_once(current) + if promoted is not None: + promoted_spec = promoted.spec + promoted_global_dim = analyzer.simplify( + promoted_spec.global_dims[0] * promoted_spec.global_dims[1] + ) + promoted_box_dim = analyzer.simplify( + promoted_spec.box_dims[0] * promoted_spec.box_dims[1] + ) + if _proof_all( + analyzer, + promoted_box_dim >= 1, + promoted_box_dim <= 256, + promoted_global_dim > 0, + promoted_global_dim <= (1 << 32), + ) == ProofStatus.PROVEN and _auto_inner_dimension_is_legal( + promoted_spec, + promoted_global_dim, + promoted_box_dim, + analyzer, + ): + current = promoted + promoted_for_merge = True + break + continue - elem_bytes = tvm.DataType(g_buf.dtype).bits // 8 - plan = TmaPlan( - swizzle_mode=l1.swizzle_mode, - dims=dims, - issue_axes=issue_axes, - tensor_ptr=g_buf.data, - elem_bytes=elem_bytes, - elem_dtype=str(g_buf.dtype), - ) - return _merge_contig_full_box_dims(plan, analyzer) + global_dims = list(spec.global_dims) + box_dims = list(spec.box_dims) + coordinates = list(spec.coordinates) + element_strides = list(spec.element_strides) + global_dims[inner_idx] = merged_global_dim + box_dims[inner_idx] = merged_box_dim + coordinates[inner_idx] = analyzer.simplify( + spec.coordinates[outer_idx] * spec.global_dims[inner_idx] + ) + for values in (global_dims, box_dims, coordinates, element_strides): + values.pop(outer_idx) + full_strides.pop(outer_idx) + current = replace( + current, + spec=replace( + spec, + global_dims=tuple(global_dims), + global_strides=tuple(full_strides[1:]), + box_dims=tuple(box_dims), + coordinates=tuple(coordinates), + element_strides=tuple(element_strides), + ), + issue_axes=_remap_issue_axes_after_remove(current.issue_axes, outer_idx), + ) + merged = True + break + if promoted_for_merge: + continue + if not merged: + return current -def _plan_needs_alignment_fix(dims, elem_bytes, analyzer: Analyzer) -> bool: - """``True`` iff some non-innermost dim has a byte-stride that isn't a - multiple of 16. cuTensorMap rejects such descriptors; merge+promote is - the way out. If the plan already satisfies the constraint, leave it - alone — the natural shape is what kernels expect and what existing - codegen tests pin. - """ - if len(dims) <= 1: +def _promotion_allowed(plan: TMAPlan) -> bool: + spec = plan.spec + analyzer = Analyzer() + if spec.effective_bytes not in _PROMOTE_DTYPE: return False - for d in dims[:-1]: - byte_stride = analyzer.simplify(d.stride * elem_bytes) - if not analyzer.can_prove_equal(tvm.tirx.floormod(byte_stride, 16), 0): - return True - return False - - -def _merge_contig_full_box_dims(plan: TmaPlan, analyzer: Analyzer) -> TmaPlan: - """Collapse adjacent fully-boxed dims that are physically contiguous. - - Two adjacent dims ``outer`` (at i) and ``inner`` (at i+1) merge when ALL of: - - 1. Physically contiguous: ``outer.stride == inner.shape * inner.stride``. - Walking inner.shape elements at inner.stride lands exactly on the - next outer element, so the two dims jointly cover one stride-1 run. - 2. Both fully boxed (``box == shape``). A partial box is a strided - slice; flattening it would change which elements the descriptor - touches. - 3. Runtime coord on each dim is provably 0. The descriptor coord for - dim d at iteration t equals - d.coord_base + Σ(iter_val · ax.coord_advance for ax in issue_axes - if ax.dim_idx == d) - For the merged dim's coord to be a constant 0 (matching the implicit - coord of the collapsed pair), both halves must satisfy: - * static term: ``coord_base == 0``, - * dynamic term: no ``IssueAxis`` binds this dim_idx. - 4. Merged ``box <= 256`` (TMA hardware limit on boxDim). - - Scan inner→outer (greedy from rank-2 down to 0) so the innermost stride - boundary is fixed first. - - When a candidate pair is blocked solely by ``merged_box > 256`` and the - layout admits an element-type promotion (current ``elem_bytes < 8``, - innermost extent even, all non-innermost element-strides even, no - issue_axis on innermost), promote ``elem_bytes`` one step (x2), halve - the innermost extent/box and the non-innermost strides, and retry the - merge. Promotion preserves byte-level semantics: byte-stride is - ``stride * elem_bytes`` and stays unchanged across promotion. - - Repeats until no merges and no promotions are possible. ``issue_axes`` - dim indices are shifted to track removed dims; the innermost - ``coord_advance`` is also halved on each promotion (it's in element - units). - """ - dims = list(plan.dims) - issue_axes = list(plan.issue_axes) - elem_bytes = plan.elem_bytes - elem_dtype = plan.elem_dtype - - # Only attempt the merge+promote rewrite when the original plan - # already violates cuTensorMap's 16-byte non-innermost-stride rule. - # An aligned plan is left intact: descriptor shape matches the - # natural buffer layout, which is what users (and goldens) expect. - if not _plan_needs_alignment_fix(dims, elem_bytes, analyzer): - return plan - - def has_issue_axis(idx): - return any(ax.dim_idx == idx for ax in issue_axes) - - def shift_issue_axes_after_remove(axes, removed_i): - return [ - IssueAxis( - extent=ax.extent, - dim_idx=ax.dim_idx if ax.dim_idx <= removed_i else ax.dim_idx - 1, - coord_advance=ax.coord_advance, - smem_stride=ax.smem_stride, - ) - for ax in axes - ] - - def try_merge_at(i, dims_, axes_): - outer, inner = dims_[i], dims_[i + 1] - if any(ax.dim_idx in (i, i + 1) for ax in axes_): - return None, None - if not analyzer.can_prove_equal(outer.coord_base, 0): - return None, None - if not analyzer.can_prove_equal(inner.coord_base, 0): - return None, None - if not analyzer.can_prove_equal(outer.box, outer.shape): - return None, None - if not analyzer.can_prove_equal(inner.box, inner.shape): - return None, None - if not analyzer.can_prove_equal(outer.stride, inner.shape * inner.stride): - return None, None - merged_box = analyzer.simplify(outer.box * inner.box) - if not analyzer.can_prove(merged_box <= 256): - # signal "blocked only by box>256" so caller can try promotion - return "blocked_box", merged_box - merged = DescDim( - shape=analyzer.simplify(outer.shape * inner.shape), - stride=inner.stride, - box=merged_box, - coord_base=0, - ) - new_dims = [*dims_[:i], merged, *dims_[i + 2 :]] - new_axes = shift_issue_axes_after_remove(axes_, i) - return new_dims, new_axes + if spec.is_packed or spec.interleave != 0 or spec.oob_fill != 0: + return False + if spec.load_mode != "tile" or spec.use_tma_reduce is not None: + return False + if spec.force_cu_dtype >= 0: + return False + if _proof_equal(spec.inner_stride, 1, analyzer) != ProofStatus.PROVEN: + return False + if any( + _proof_equal(stride, 1, analyzer) != ProofStatus.PROVEN for stride in spec.element_strides + ): + return False + return not any(coord.dim_idx == 0 for axis in plan.issue_axes for coord in axis.coords) - _PROMOTE_CHAIN = {1: ("uint16", 2), 2: ("uint32", 4), 4: ("uint64", 8)} - def try_promote(dims_, axes_, eb, edt): - if eb not in _PROMOTE_CHAIN: - return None - if not dims_: - return None - innermost_idx = len(dims_) - 1 - if any(ax.dim_idx == innermost_idx for ax in axes_): - return None - inner = dims_[innermost_idx] - if not analyzer.can_prove_equal(inner.stride, 1): +def _promote_auto_once(plan: TMAPlan) -> TMAPlan | None: + """Promote descriptor units while preserving every byte address.""" + + if not _promotion_allowed(plan): + return None + analyzer = Analyzer() + spec = plan.spec + new_dtype, new_bits = _PROMOTE_DTYPE[spec.effective_bytes] + for value, label in ( + (spec.global_dims[0], "innermost global shape"), + (spec.box_dims[0], "innermost box"), + (spec.coordinates[0], "innermost coordinate"), + ): + if _proof_equal(tvm.tirx.floormod(value, 2), 0, analyzer) != ProofStatus.PROVEN: return None - if not analyzer.can_prove_equal(tvm.tirx.floormod(inner.shape, 2), 0): + for stride in spec.global_strides: + if ( + _proof_equal(tvm.tirx.floormod(stride, spec.effective_bytes * 2), 0, analyzer) + != ProofStatus.PROVEN + ): return None - for d in dims_[:-1]: - if not analyzer.can_prove_equal(tvm.tirx.floormod(d.stride, 2), 0): - return None - new_dtype, new_eb = _PROMOTE_CHAIN[eb] - new_dims = [] - for j, d in enumerate(dims_): - if j == innermost_idx: - new_dims.append( - DescDim( - shape=analyzer.simplify(tvm.tirx.floordiv(d.shape, 2)), - stride=d.stride, - box=analyzer.simplify(tvm.tirx.floordiv(d.box, 2)), - coord_base=analyzer.simplify(tvm.tirx.floordiv(d.coord_base, 2)), - ) - ) - else: - new_dims.append( - DescDim( - shape=d.shape, - stride=analyzer.simplify(tvm.tirx.floordiv(d.stride, 2)), - box=d.box, - coord_base=d.coord_base, - ) - ) - new_axes = [ - IssueAxis( - extent=ax.extent, - dim_idx=ax.dim_idx, - coord_advance=( - analyzer.simplify(tvm.tirx.floordiv(ax.coord_advance, 2)) - if ax.dim_idx == innermost_idx - else ax.coord_advance - ), - smem_stride=ax.smem_stride, - ) - for ax in axes_ - ] - return new_dims, new_axes, new_eb, new_dtype + global_dims = list(spec.global_dims) + box_dims = list(spec.box_dims) + coordinates = list(spec.coordinates) + global_dims[0] = analyzer.simplify(tvm.tirx.floordiv(global_dims[0], 2)) + box_dims[0] = analyzer.simplify(tvm.tirx.floordiv(box_dims[0], 2)) + coordinates[0] = analyzer.simplify(tvm.tirx.floordiv(coordinates[0], 2)) + new_spec = replace( + spec, + descriptor_dtype=new_dtype, + descriptor_bits=new_bits, + effective_bytes=new_bits // 8, + global_dims=tuple(global_dims), + box_dims=tuple(box_dims), + coordinates=tuple(coordinates), + ) + transaction_bits = new_bits + for box in box_dims: + transaction_bits = analyzer.simplify(transaction_bits * box) + issue_count = 1 + for axis in plan.issue_axes: + issue_count = analyzer.simplify(issue_count * axis.extent) + new_spec = replace( + new_spec, + transaction_bits=transaction_bits, + payload_bits=analyzer.simplify(transaction_bits * issue_count), + ) + if _proof_equal(new_spec.payload_bits, spec.payload_bits, analyzer) != ProofStatus.PROVEN: + return None + return replace(plan, spec=new_spec) + + +def _repair_auto_candidate(plan: TMAPlan): + candidate = plan while True: - # Greedy inner→outer merge sweep. - merged_any = False - blocked_by_box = False - for i in range(len(dims) - 2, -1, -1): - res, _info = try_merge_at(i, dims, issue_axes) - if res == "blocked_box": - blocked_by_box = True - continue - if res is not None: - dims, issue_axes = res, _info - merged_any = True + candidate = _canonicalize_auto_plan(candidate) + failures = _validation_failures(candidate.spec, auto=True) + if not failures: + return candidate, () + if any(finding.rule not in _REPAIRABLE_RULES for finding in failures): + return None, failures + + promoted = _promote_auto_once(candidate) + if promoted is None: + return None, failures + candidate = promoted + + +def _runtime_config(op_call, sctx, direction: str, *, explicit: bool): + allowed = _EXPLICIT_CONFIG if explicit else _COMMON_CONFIG + unknown = sorted(set(op_call.config) - allowed) + if unknown: + fail( + f"dispatch={'tma_explicit' if explicit else 'tma_auto'} does not support " + f"config key(s) {unknown}" + ) + + if not explicit and op_call.config.get("oob") is not None: + fail('tma_auto does not support non-default oob; use dispatch="tma_explicit"') + if direction != "g2s" and op_call.config.get("oob") is not None: + fail("TensorMap oob is only valid for explicit global-to-shared copies") + + cta_group = op_call.config.get("cta_group", 1) + if isinstance(cta_group, IntImm): + cta_group = int(cta_group) + if cta_group not in (1, 2): + fail(f"cta_group must be 1 or 2, got {cta_group}") + + cta_mask = op_call.config.get("cta_mask", 0) + if isinstance(cta_mask, IntImm): + cta_mask_value = int(cta_mask) + if not 0 <= cta_mask_value <= 0xFFFF: + fail(f"cta_mask must fit uint16, got {cta_mask_value}") + elif isinstance(cta_mask, int): + if not 0 <= cta_mask <= 0xFFFF: + fail(f"cta_mask must fit uint16, got {cta_mask}") + elif not isinstance(cta_mask, tvm.tirx.Expr): + fail("cta_mask must be an integer or TIR expression") + mbar = op_call.config.get("mbar") + mbarrier_addr = op_call.config.get("mbarrier_addr", False) + if isinstance(mbarrier_addr, IntImm): + mbarrier_addr = bool(int(mbarrier_addr)) + if not isinstance(mbarrier_addr, bool | tvm.tirx.Expr): + fail("mbarrier_addr must be bool or Expr") + if direction == "g2s": + if mbar is None: + fail("global-to-shared TMA requires mbar") + else: + if mbar is not None: + fail("mbar is only valid for global-to-shared TMA") + if mbarrier_addr not in (False, None): + fail("mbarrier_addr is only valid for global-to-shared TMA") + if not ( + (isinstance(cta_mask, int) and cta_mask == 0) + or (isinstance(cta_mask, IntImm) and int(cta_mask) == 0) + ): + fail("cta_mask is only valid for global-to-shared TMA") + + use_tma_reduce = op_call.config.get("use_tma_reduce") + if use_tma_reduce is not None: + if direction != "s2g": + fail("use_tma_reduce is only valid for shared-to-global TMA") + if use_tma_reduce not in ("add", "min", "max", "inc", "dec", "and", "or", "xor"): + fail(f"unsupported TMA reduce operation {use_tma_reduce!r}") + + cache_hint, cache_policy = _normalize_cache_hint(op_call.config.get("cache_hint", "")) + return { + "cta_group": cta_group, + "cta_mask": cta_mask, + "mbar": mbar, + "mbarrier_addr": mbarrier_addr, + # A dynamic form selects between two equivalent address + # representations. Normalize it once to the shared-address form so + # there remains exactly one TMA instruction. + "mbar_is_shared_addr": ( + direction == "g2s" + and ( + cta_group == 2 or mbarrier_addr is True or isinstance(mbarrier_addr, tvm.tirx.Expr) + ) + ), + "use_tma_reduce": use_tma_reduce, + "cache_hint": cache_hint, + "cache_policy": cache_policy, + "l2_promotion": _normalize_l2_promotion(op_call.config.get("tensormap_l2_promotion")), + "oob_fill": _normalize_oob(op_call.config.get("oob")), + "prefetch": bool(op_call.config.get("prefetch_tensormap", False)), + "tma_dtype": op_call.config.get("tma_dtype"), + "target_arch": sctx.target.arch, + } + + +def _copy_direction(op_call): + op_call = TilePrimitiveCall.downcast(op_call) + dst_region, src_region = op_call.dst, op_call.src + src_scope = src_region.buffer.scope() + dst_scope = dst_region.buffer.scope() + if src_scope == "global" and dst_scope.startswith("shared"): + return "g2s", dst_region, src_region + if src_scope.startswith("shared") and dst_scope == "global": + return "s2g", src_region, dst_region + fail(f"TMA requires global<->shared operands, got src={src_scope}, dst={dst_scope}") + + +def _build_auto_plan(op_call: TilePrimitiveCall, sctx: DispatchContext) -> TMAPlan: + direction, shared_region, global_region = _copy_direction(op_call) + s_buf = shared_region.buffer + g_buf = global_region.buffer + if str(s_buf.dtype) != str(g_buf.dtype): + _auto_fail( + "dtype", + f"shared dtype={s_buf.dtype} and global dtype={g_buf.dtype} differ", + ) + runtime = _runtime_config(op_call, sctx, direction, explicit=False) + if "gather4" in op_call.config or "src_selector" in op_call.config: + fail('gather4 and src_selector are only supported by dispatch="tma_explicit"') + + s_starts, s_extents = _copy_region_parts(shared_region.region) + g_starts, g_extents = _copy_region_parts(global_region.region) + try: + swizzle = get_swizzle_mode_from_layout(s_buf.layout) + except ValueError as error: + _auto_fail("shared-swizzle", str(error)) + if swizzle is None: + _auto_fail("shared-layout", f"cannot recognize shared swizzle in {s_buf.layout}") + try: + _, sliced_smem_with_offset = _slice_layout(s_buf, s_starts, s_extents, "shared") + except ValueError as error: + _auto_fail("shared-slice", str(error)) + smem_layout_offset = _layout_offset(sliced_smem_with_offset) + sliced_smem = TileLayout.from_iters( + sliced_smem_with_offset.shard, + sliced_smem_with_offset.replica, + {}, + ).canonicalize() + if not isinstance(sliced_smem, TileLayout): + _auto_fail( + "shared-canonicalize", + f"canonical sliced shared layout is not a TileLayout: {sliced_smem}", + ) + _assert_plain_memory_layout(sliced_smem, "canonical sliced shared") + canonical_smem_shape = tuple(iterator.extent for iterator in sliced_smem.shard) + try: + sliced_smem, _ = sliced_smem.group_many((canonical_smem_shape, tuple(g_extents))) + except (TypeError, ValueError, tvm.error.TVMError) as error: + _auto_fail( + "group-shared", + f"cannot refine canonical shared shape={canonical_smem_shape} " + f"with copy shape={tuple(g_extents)}: {error}", + ) + smem_order = _smem_iter_order(sliced_smem) + gt_records = _build_auto_gt(g_buf, g_starts, g_extents, sliced_smem, sctx) + descriptor = _dtype_contract(g_buf.dtype, runtime["tma_dtype"]) + + best = None + last_failures = () + for prefix in range(1, len(smem_order) + 1): + raw = _auto_spec_for_prefix( + op_call=op_call, + sctx=sctx, + direction=direction, + s_buf=s_buf, + g_buf=g_buf, + s_starts=s_starts, + smem_layout_offset=smem_layout_offset, + sliced_smem=sliced_smem, + smem_order=smem_order, + gt_records=gt_records, + prefix=prefix, + swizzle=swizzle, + descriptor=descriptor, + runtime=runtime, + ) + _validate_auto_shared_mapping(raw, sliced_smem, swizzle, sctx) + repaired, failures = _repair_auto_candidate(raw) + if repaired is None: + last_failures = failures + if best is not None: break - if merged_any: continue - # Nothing merged this pass; try promotion if any pair was box-blocked. - if not blocked_by_box: - break - promoted = try_promote(dims, issue_axes, elem_bytes, elem_dtype) - if promoted is None: - break - dims, issue_axes, elem_bytes, elem_dtype = promoted - - return TmaPlan( - swizzle_mode=plan.swizzle_mode, - dims=dims, - issue_axes=issue_axes, - tensor_ptr=plan.tensor_ptr, - elem_bytes=elem_bytes, - elem_dtype=elem_dtype, + best = repaired + + if best is None: + if last_failures: + _raise_validation("prefix-search", last_failures, auto=True) + _auto_fail("prefix-search", "no legal shared-memory prefix") + return best + + +def _explicit_smem_layout(s_buf: Buffer, starts, extents, swizzle: SwizzleMode): + _, sliced = _slice_layout(s_buf, starts, extents, "shared") + # LayoutSlice preserves the layout's physical base and adds the selected + # region offset, so the sliced offset is already the complete layout-side + # contribution to the shared pointer. + offset = _layout_offset(sliced) + normalized = TileLayout.from_iters(sliced.shard, sliced.replica, {}) + canonical = normalized.canonicalize() + if not canonical.is_trivial(): + fail( + "tma_explicit stage=shared-layout: sliced shared layout must canonicalize " + f"to trivial after extracting its pointer offset; got {canonical}" + ) + analyzer = Analyzer() + byte_offset = _elements_to_bytes( + s_buf.elem_offset + offset, + tvm.DataType(s_buf.dtype).bits, + analyzer, + auto=False, + stage="shared-layout", + label="shared pointer offset", + ) + atom_alignment = {0: 16, 1: 32, 2: 64, 3: 128}[swizzle.value] + status = _proof_equal(tvm.tirx.floormod(byte_offset, atom_alignment), 0, analyzer) + if status == ProofStatus.DISPROVEN: + fail( + "tma_explicit stage=shared-layout: shared slice pointer offset " + f"{byte_offset}B violates {atom_alignment}B swizzle/base alignment" + ) + return ( + _to_tile_layout(s_buf.layout, s_buf.shape), + analyzer.simplify(s_buf.elem_offset + offset), ) -def _validate_hw_constraints(plan: TmaPlan, dtype: str) -> tuple: - """Return ``(ok, reason)``. ``reason`` is the error string when ``ok`` is False.""" +def _direct_global_layout(g_buf: Buffer): + layout = g_buf.layout + if not isinstance(layout, TileLayout): + fail( + "tma_explicit stage=global-layout: global Buffer/view layout must be " + f"TileLayout, got {type(layout).__name__}" + ) + _assert_plain_memory_layout(layout, "explicit global") + if len(g_buf.shape) != len(layout.shard): + fail( + "tma_explicit stage=global-layout: tensor rank " + f"{len(g_buf.shape)} != memory-layout rank {len(layout.shard)}" + ) analyzer = Analyzer() + for dim, (shape, iterator) in enumerate(zip(g_buf.shape, layout.shard)): + status = _proof_equal(shape, iterator.extent, analyzer) + if status != ProofStatus.PROVEN: + fail( + "tma_explicit stage=global-layout: " + f"dim={dim} shape={shape} must provably equal layout extent=" + f"{iterator.extent}; got {status.value}" + ) + return layout - if plan.rank == 0: - return False, "TMA descriptor rank must be ≥ 1" - if plan.rank > 5: - return False, f"TMA descriptor rank {plan.rank} exceeds hardware limit of 5" - # Innermost dim stride must be 1 (unit stride). - inner = plan.dims[-1] - if not analyzer.can_prove_equal(inner.stride, 1): - return False, f"TMA innermost dim must have unit stride; got {inner.stride}" +def _explicit_spec_for_gmem( + *, + g_buf, + s_buf, + s_starts, + g_starts, + g_extents, + swizzle, + runtime, + sctx, + gather4, + smem_base_offset, +): + analyzer = Analyzer() + layout = _direct_global_layout(g_buf) + descriptor = _dtype_contract(g_buf.dtype, runtime["tma_dtype"]) + descriptor_dtype, descriptor_bits, effective_bytes, packed_kind, force_cu_dtype = descriptor + strides_outer = [ + _elements_to_bytes( + iterator.stride, + descriptor_bits, + analyzer, + auto=False, + stage="global-layout", + label=f"global layout stride dim={idx}", + ) + for idx, iterator in enumerate(layout.shard) + ] + base, base_key, base_byte_offset = _buffer_base(g_buf, auto=False, stage="global-base") + global_dims = tuple(reversed(g_buf.shape)) + box_dims = tuple(reversed(g_extents)) + coordinates = tuple(reversed(g_starts)) + full_strides = tuple(reversed(strides_outer)) + inner_stride = layout.shard[-1].stride + if _proof_equal(inner_stride, 1, analyzer) != ProofStatus.PROVEN: + fail( + "tma_explicit stage=global-layout: the innermost memory stride must be " + f"provably one because CUDA omits it from globalStrides; got {inner_stride}" + ) + transaction_bits = descriptor_bits + for box in box_dims: + transaction_bits = analyzer.simplify(transaction_bits * box) + if gather4: + transaction_bits = analyzer.simplify(transaction_bits * 4) + return TensorMapSpec( + descriptor_dtype=descriptor_dtype, + descriptor_bits=descriptor_bits, + effective_bytes=effective_bytes, + packed_kind=packed_kind, + force_cu_dtype=force_cu_dtype, + base=base, + base_key=base_key, + descriptor_name=g_buf.name, + base_byte_offset=base_byte_offset, + global_dims=global_dims, + global_strides=full_strides[1:], + inner_stride=inner_stride, + box_dims=box_dims, + element_strides=(1,) * len(global_dims), + interleave=0, + swizzle=swizzle.value, + l2_promotion=runtime["l2_promotion"], + oob_fill=runtime["oob_fill"], + direction="g2s" if runtime["mbar"] is not None else "s2g", + load_mode="tile_gather4" if gather4 else "tile", + target_arch=sctx.target.arch, + coordinates=coordinates, + gather4=tuple(gather4), + smem_buffer=s_buf, + smem_start=tuple(s_starts), + smem_base_offset=smem_base_offset, + mbar=runtime["mbar"], + mbar_is_shared_addr=runtime["mbar_is_shared_addr"], + cta_group=runtime["cta_group"], + cta_mask=runtime["cta_mask"], + cache_hint=runtime["cache_hint"], + cache_policy=runtime["cache_policy"], + use_tma_reduce=runtime["use_tma_reduce"], + payload_bits=transaction_bits, + transaction_bits=transaction_bits, + ) - # Innermost box times element size must fit the swizzle atom. - if not _swizzle_inner_box_fits(dtype, plan.swizzle_mode, inner.box): - return False, "TMA innermost box exceeds the swizzle atom size" - return True, "" +def _normalize_gather4(value): + if value is None: + return () + if not isinstance(value, list | tuple | tvm.ir.Array) or len(value) != 4: + fail("tma_explicit gather4 must contain exactly four row coordinates") + return tuple(value) -def _build_plan_with_shrink(l1: L1Result, g_buf: Buffer, s_buf: Buffer) -> TmaPlan: - """Enumerate chain prefix length j from max down to 0, validate - alignment per gmem iter, build and validate the plan. Return the first - plan that passes everything. Raise when j=0 still fails. - """ +def _validate_gather4_dst(s_buf: Buffer, s_starts, s_extents, spec: TensorMapSpec) -> None: analyzer = Analyzer() - chain = _find_contiguous_chain_prefix(l1.smem_groups) - max_j = len(chain) - - # Empty-smem_groups case (all ext=1): the assembly still yields a - # valid plan (trivial desc dims). - if not l1.smem_groups: - plan = _assemble_plan(l1, {}, [], g_buf, analyzer) - ok, reason = _validate_hw_constraints(plan, s_buf.dtype) - if ok: - return plan - fail(f"TMA plan (no smem groups) failed hardware check: {reason}") - - last_reason = "no valid plan" - for j in range(max_j, -1, -1): - per_iter_selected: dict = _distribute_selection(chain[:j], l1.smem_groups) - - # Check alignment for each ext>1 iter. - aligned = True - for gi, group in enumerate(l1.smem_groups): - iter_idx = group.bound_gmem_iter_idx - sel = per_iter_selected.get(gi, []) - if not _check_alignment(l1.gmem_iters[iter_idx], sel, group.shards, analyzer): - aligned = False - last_reason = f"alignment fails for gmem iter {iter_idx} at j={j}" - break - if not aligned: - continue - - plan = _assemble_plan(l1, per_iter_selected, chain[:j], g_buf, analyzer) - ok, reason = _validate_hw_constraints(plan, s_buf.dtype) - if ok: - return plan - last_reason = reason + if len(s_extents) < 1: + fail("tma_explicit gather4 requires a non-scalar shared destination") + if _proof_equal(s_extents[0], 4, analyzer) != ProofStatus.PROVEN: + fail("tma_explicit gather4 requires public shared axis 0 to provably contain four rows") + _, sliced = _slice_layout(s_buf, s_starts, s_extents, "gather4 shared") + normalized = TileLayout.from_iters(sliced.shard, sliced.replica, {}) + canonical = normalized.canonicalize() + if not canonical.is_trivial(): + fail(f"tma_explicit gather4 destination is not four-row box-linear: {canonical}") + row_bits = spec.transaction_bits + if spec.gather4: + row_bits = analyzer.simplify(tvm.tirx.floordiv(row_bits, 4)) + shared_bits = tvm.DataType(s_buf.dtype).bits + row_elements = analyzer.simplify(tvm.tirx.floordiv(row_bits, shared_bits)) + grouped, sep = sliced.group(list(s_extents)) + row_iters = grouped.shard[sep[0] : sep[1]] + if len(row_iters) != 1: + fail("tma_explicit gather4 destination row axis must map to one memory iterator") + row = row_iters[0] + if ( + _proof_equal(row.extent, 4, analyzer) != ProofStatus.PROVEN + or _proof_equal(row.stride, row_elements, analyzer) != ProofStatus.PROVEN + ): + fail( + "tma_explicit gather4 destination rows are not box-linear at " + f"payload width {row_elements}; got extent={row.extent}, stride={row.stride}" + ) - fail(f"TMA plan: all chain prefix lengths rejected; last reason: {last_reason}") +def _normalize_src_selector(value): + if value is None: + return () + if not isinstance(value, list | tuple | tvm.ir.Array): + fail("tma_explicit src_selector must be a list of (condition, global Buffer/view)") + result = [] + for idx, item in enumerate(value): + if not isinstance(item, list | tuple | tvm.ir.Array) or len(item) != 2: + fail(f"tma_explicit src_selector[{idx}] must be a (condition, Buffer/view) pair") + condition, buffer = item + if not isinstance(condition, tvm.tirx.Expr): + fail(f"tma_explicit src_selector[{idx}] condition must be a TIR expression") + if not is_buffer_var(buffer): + fail( + f"tma_explicit src_selector[{idx}] candidate must be a global " + "Buffer/view, not a region" + ) + if buffer.scope() != "global": + fail( + f"tma_explicit src_selector[{idx}] candidate scope must be global, " + f"got {buffer.scope()}" + ) + result.append((condition, buffer)) + return tuple(result) -# ============================================================================== -# Emit layer + entry point -# ============================================================================== +def _selector_compatibility(main: TensorMapSpec, candidate: TensorMapSpec, index: int): + analyzer = Analyzer() + fields = ( + ("descriptor dtype", main.descriptor_dtype, candidate.descriptor_dtype), + ("descriptor bits", main.descriptor_bits, candidate.descriptor_bits), + ("forced CUDA dtype", main.force_cu_dtype, candidate.force_cu_dtype), + ("rank", main.rank, candidate.rank), + ("box", main.box_dims, candidate.box_dims), + ("interleave", main.interleave, candidate.interleave), + ("swizzle", main.swizzle, candidate.swizzle), + ("load mode", main.load_mode, candidate.load_mode), + ("transaction bits", main.transaction_bits, candidate.transaction_bits), + ) + for label, lhs, rhs in fields: + if isinstance(lhs, tuple): + equal = len(lhs) == len(rhs) and all( + _proof_equal(a, b, analyzer) == ProofStatus.PROVEN for a, b in zip(lhs, rhs) + ) + elif isinstance(lhs, int | str): + equal = lhs == rhs + else: + equal = _proof_equal(lhs, rhs, analyzer) == ProofStatus.PROVEN + if not equal: + fail( + f"tma_explicit src_selector[{index}] has incompatible {label}: " + f"main={lhs}, candidate={rhs}" + ) -def copy_tma_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: - """Lower global<->shared copy_async to TMA using the unified algorithm. - Emits a device-side unrolled loop over the flat issue-axis extent and - a host-side ``cuTensorMapEncodeTiled`` (deduped via cache key). - """ - op_call = TilePrimitiveCall.downcast(op_call) - dst_buffer_region, src_buffer_region = op_call.dst, op_call.src - src: Buffer = src_buffer_region.buffer - dst: Buffer = dst_buffer_region.buffer +def _build_explicit_plan(op_call: TilePrimitiveCall, sctx: DispatchContext): + direction, shared_region, global_region = _copy_direction(op_call) + s_buf = shared_region.buffer + g_buf = global_region.buffer + runtime = _runtime_config(op_call, sctx, direction, explicit=True) + gather4 = _normalize_gather4(op_call.config.get("gather4")) + selectors = _normalize_src_selector(op_call.config.get("src_selector")) + if (gather4 or selectors) and direction != "g2s": + fail("tma_explicit gather4 and src_selector are only valid for global-to-shared") + if gather4 and len(g_buf.shape) != 2: + fail("tma_explicit gather4 requires a rank-2 global tensor") + + s_starts, s_extents = _copy_region_parts(shared_region.region) + g_starts, g_extents = _copy_region_parts(global_region.region) + if gather4: + analyzer = Analyzer() + if _proof_equal(g_extents[0], 1, analyzer) == ProofStatus.DISPROVEN: + fail("tma_explicit gather4 global public axis 0 must describe one row") + if _proof_equal(s_extents[0], 4, analyzer) == ProofStatus.DISPROVEN: + fail("tma_explicit gather4 shared public axis 0 must describe four rows") - src_scope, dst_scope = src.scope(), dst.scope() - if src_scope == "global" and dst_scope.startswith("shared"): - direction = "g2s" - s_buf, g_buf = dst, src - shared_region, global_region = dst_buffer_region, src_buffer_region - elif src_scope.startswith("shared") and dst_scope == "global": - direction = "s2g" - s_buf, g_buf = src, dst - shared_region, global_region = src_buffer_region, dst_buffer_region - else: - raise ValueError( - f"Unsupported combination of src and dst scopes: src={src_scope} dst={dst_scope}" + try: + swizzle = get_swizzle_mode_from_layout(s_buf.layout) + except ValueError as error: + fail(f"tma_explicit stage=shared-swizzle: {error}") + if swizzle is None: + fail(f"tma_explicit cannot recognize shared swizzle in {s_buf.layout}") + shared_layout, smem_base_offset = _explicit_smem_layout(s_buf, s_starts, s_extents, swizzle) + main_spec = _explicit_spec_for_gmem( + g_buf=g_buf, + s_buf=s_buf, + s_starts=s_starts, + g_starts=g_starts, + g_extents=g_extents, + swizzle=swizzle, + runtime=runtime, + sctx=sctx, + gather4=gather4, + smem_base_offset=smem_base_offset, + ) + _validate_explicit(main_spec, "main-descriptor") + if gather4: + _validate_gather4_dst(s_buf, s_starts, s_extents, main_spec) + + candidate_specs = [] + for idx, (condition, candidate_buffer) in enumerate(selectors): + candidate = _explicit_spec_for_gmem( + g_buf=candidate_buffer, + s_buf=s_buf, + s_starts=s_starts, + g_starts=g_starts, + g_extents=g_extents, + swizzle=swizzle, + runtime=runtime, + sctx=sctx, + gather4=gather4, + smem_base_offset=smem_base_offset, ) + _validate_explicit(candidate, f"src-selector[{idx}]") + _selector_compatibility(main_spec, candidate, idx) + candidate_specs.append((condition, candidate)) + return ( + TMAPlan(spec=main_spec, issue_axes=(), shared_layout=shared_layout), + tuple(candidate_specs), + ) - g_st = [region.min for region in global_region.region] - g_ext = [region.extent for region in global_region.region] - s_st = [region.min for region in shared_region.region] - s_ext = [region.extent for region in shared_region.region] - - oob_mode = _normalize_oob_mode(s_buf.dtype, op_call.config.get("oob", None)) - oob_fill_kind = _oob_fill_kind(oob_mode) - - # L1 → L2 → L3 - l1 = _build_l1_result(s_buf, g_buf, g_st, g_ext, s_st, s_ext) - plan = _build_plan_with_shrink(l1, g_buf, s_buf) - - # Optional descriptor dtype override. An fp32 buffer feeding a tf32 MMA should - # be loaded via a TFLOAT32 (== 11) descriptor so the TMA hardware RN-truncates - # fp32 -> tf32 ON LOAD (matching the MMA's operand precision and a torch - # allow_tf32 / DeepGEMM reference); leaving it as FLOAT32 loads full fp32 and - # the tf32 MMA then RZ-truncates, diverging by ~1 tf32 ULP. -1 = derive from - # the buffer dtype (the C++ ``runtime.cuTensorMapEncodeTiled`` default). - _TMA_DTYPE_TO_CU = {"tf32": 11, "tfloat32": 11} - _tma_dtype = op_call.config.get("tma_dtype", None) - if _tma_dtype is not None and _tma_dtype not in _TMA_DTYPE_TO_CU: - fail(f"Unsupported tma_dtype={_tma_dtype!r}; expected one of {sorted(_TMA_DTYPE_TO_CU)}") - if _tma_dtype is not None and plan.elem_dtype not in ("float32", "tfloat32"): - fail(f"tma_dtype={_tma_dtype!r} requires a float32 descriptor; got {plan.elem_dtype}") - force_cu_dtype = _TMA_DTYPE_TO_CU.get(_tma_dtype, -1) - - # Direction / runtime-config bits that don't affect the plan itself. - cta_group = op_call.config.get("cta_group", None) - if cta_group is None: - cta_group = 1 if sctx.target.arch == "sm_100a" else -1 - - cta_mask = op_call.config.get("cta_mask", None) - if cta_mask is not None: - assert direction == "g2s", "cta_mask is only supported for global to shared copy" - else: - cta_mask = 0 - if direction == "g2s": - mbar = op_call.config.get("mbar", None) - if mbar is None: - raise ValueError("mbar is not set in config") - use_tma_reduce = op_call.config.get("use_tma_reduce", None) +def _descriptor_cache_key(spec: TensorMapSpec) -> str: + fields = ( + spec.base_key, + spec.descriptor_dtype, + spec.descriptor_bits, + spec.force_cu_dtype, + spec.global_dims, + spec.global_strides, + spec.box_dims, + spec.element_strides, + spec.interleave, + spec.swizzle, + spec.l2_promotion, + spec.oob_fill, + ) + return "tensormap:" + ":".join(str(field) for field in fields) - dtype_bytes = plan.elem_bytes - tma_global_strides = [stride * dtype_bytes for stride in plan.g_strides] - # cuTensorMap omits the last dim's stride (implicit element size). - tma_g_strides_for_map = tma_global_strides[:-1] if plan.rank > 1 else [] - element_strides = [1] * plan.rank - flat_total_extent = plan.flatten_total_extent() +def _get_or_encode_descriptor(spec: TensorMapSpec, sctx: DispatchContext): + key = _descriptor_cache_key(spec) + cached = sctx.cache_get(key) + if cached is not None: + return cached, key - def compute_offsets_and_tma_coords(loop_var): - s_offset, coords = plan.offsets_and_coords(loop_var) - simplified = _simplify_with_var_ranges( - [s_offset, *coords], [(loop_var, flat_total_extent)], sctx + tensor_map = T.Var(f"{spec.descriptor_name}_tensormap", ty=T.handle("tensormap").ty) + + # fmt: off + @T.prim_func(check_well_formed=False) + def create_tensor_map(): + T.Bind(T.tvm_stack_alloca("tensormap", 1), var=tensor_map) + T.call_packed( + "runtime.cuTensorMapEncodeTiled", + tensor_map, + spec.descriptor_dtype, + spec.rank, + spec.base, + *spec.global_dims, + *spec.global_strides, + *spec.box_dims, + *spec.element_strides, + spec.interleave, + spec.swizzle, + spec.l2_promotion, + spec.oob_fill, + *([spec.force_cu_dtype] if spec.force_cu_dtype >= 0 else []), ) - return simplified[0], reversed(simplified[1:]) + T.tvm_kernel_replace_point() + # fmt: on - def val_key(value) -> str: - return str(value) + sctx.add_init_stmt(create_tensor_map.body, host=True) + sctx.cache_set(key, tensor_map) + return tensor_map, key - tensormap_cache_key = ( - f"tensormap:{hash(plan.tensor_ptr)}:{g_buf.dtype}:{val_key(plan.rank)}" - f":{tuple(val_key(v) for v in plan.shape)}" - f":{tuple(val_key(v) for v in tma_g_strides_for_map)}" - f":{tuple(val_key(v) for v in plan.box_dim)}" - f":{val_key(plan.swizzle_mode.value)}:{oob_fill_kind}:{force_cu_dtype}" - ) - cached_tensormap = sctx.cache_get(tensormap_cache_key) - if cached_tensormap is not None: - tensor_map = cached_tensormap - tensormap_is_cached = True - else: - tensor_map = T.Var(g_buf.data.name + "_tensormap", ty=T.handle("tensormap").ty) - tensormap_is_cached = False +def _prefetch_main_descriptor(tensor_map, key: str, sctx: DispatchContext) -> None: + cache_key = f"prefetch_tensormap:{key}" + if sctx.cache_get(cache_key) is not None: + return + if "warp_id_in_cta" not in sctx.launch_params: + fail("prefetch_tensormap requires warp_id_in_cta launch param") + warp_id = sctx.launch_params["warp_id_in_cta"].var # fmt: off @T.prim_func(check_well_formed=False) - def impl(): - for loop_vars in T.unroll(flat_total_extent): - s_offset, tma_coords = T.meta_var(compute_offsets_and_tma_coords(loop_vars)) - s_buf_w_offset = T.decl_buffer( - s_buf.shape, - s_buf.dtype, - s_buf.data, - elem_offset=s_buf.elem_offset + s_offset, - scope=s_buf.scope(), - layout=_to_tile_layout(s_buf.layout, s_buf.shape), - ) + def prefetch_tensor_map(): + if warp_id == 0: + if T.ptx.elect_sync() != T.uint32(0): + T.ptx.prefetch_tensormap(T.address_of(tensor_map)) + T.tvm_kernel_replace_point() + # fmt: on + + sctx.add_init_stmt(prefetch_tensor_map.body) + sctx.cache_set(cache_key, tensor_map) - if direction == "g2s": - T.ptx.cp_async.bulk.tensor.g2c( - plan.rank, - s_buf_w_offset.ptr_to(s_st), - mbar, - T.address_of(tensor_map), - cta_mask, - cta_group, - op_call.config.get("cache_hint", ""), - *tma_coords, + +def _selected_descriptor(main_map, candidate_maps): + if not candidate_maps: + return main_map, None, False + selected_expr = T.address_of(main_map) + for condition, candidate_map in reversed(candidate_maps): + selected_expr = tvm.tirx.Select( + condition, + T.address_of(candidate_map), + selected_expr, + ) + selected = T.Var("selected_tensormap", "uint64") + return selected, tvm.tirx.Bind(selected, selected_expr), True + + +def _emit_plan( + plan: TMAPlan, + tensor_map, + selector_bind, + tensor_map_is_address: bool, + sctx: DispatchContext, +) -> PrimFunc: + spec = plan.spec + tensor_map_address = tensor_map if tensor_map_is_address else T.address_of(tensor_map) + + def tma_coordinates(coordinates): + if spec.load_mode == "tile_gather4": + return [coordinates[0], *spec.gather4] + return list(coordinates) + + def emit_at(shared_ptr, coordinates): + coords = tma_coordinates(coordinates) + if spec.direction == "g2s": + mbar_operand = ( + T.cuda.cvta_generic_to_shared(spec.mbar) if spec.mbar_is_shared_addr else spec.mbar + ) + T.evaluate( + T.ptx.cp_async.bulk.tensor.g2s_cluster( + spec.rank, + shared_ptr, + mbar_operand, + tensor_map_address, + spec.cta_mask, + spec.cta_group, + spec.cache_hint, + *coords, + cache_policy=spec.cache_policy, + load_mode=spec.load_mode, + mbar_is_shared_addr=spec.mbar_is_shared_addr, ) - else: - if use_tma_reduce is None: - T.ptx.cp_async.bulk.tensor.s2g( - plan.rank, - s_buf_w_offset.ptr_to(s_st), - T.address_of(tensor_map), - op_call.config.get("cache_hint", ""), - *tma_coords, - ) - else: - T.ptx.cp_async.bulk.tensor.s2g_reduce( - plan.rank, - s_buf_w_offset.ptr_to(s_st), - T.address_of(tensor_map), - op_call.config.get("cache_hint", ""), - use_tma_reduce, - *tma_coords, - ) - # fmt: on + ) + elif spec.use_tma_reduce is None: + T.evaluate( + T.ptx.cp_async.bulk.tensor.s2g( + spec.rank, + shared_ptr, + tensor_map_address, + spec.cache_hint, + *coords, + cache_policy=spec.cache_policy, + ) + ) + else: + T.evaluate( + T.ptx.cp_async.bulk.tensor.s2g_reduce( + spec.rank, + shared_ptr, + tensor_map_address, + spec.cache_hint, + spec.use_tma_reduce, + *coords, + cache_policy=spec.cache_policy, + ) + ) + + def shared_ptr(element_offset=0): + # Keep the sliced offset in the pointer index instead of the Buffer's + # elem_offset. The latter is part of a flat DeclBuffer definition; + # after loop unrolling, CSE may otherwise lift an offset containing a + # locally bound coordinate above that coordinate's Bind statement. + smem_view = T.decl_buffer( + (1,), + spec.smem_buffer.dtype, + spec.smem_buffer.data, + elem_offset=0, + scope=spec.smem_buffer.scope(), + ) + return smem_view.ptr_to([spec.smem_base_offset + element_offset]) - if not tensormap_is_cached: + if not plan.issue_axes: # fmt: off @T.prim_func(check_well_formed=False) - def create_tensor_map(): - T.Bind(T.tvm_stack_alloca("tensormap", 1), var=tensor_map) - T.call_packed( - "runtime.cuTensorMapEncodeTiled", - tensor_map, - plan.elem_dtype, - plan.rank, - plan.tensor_ptr, - *reversed(plan.shape), - *reversed(tma_g_strides_for_map) if plan.rank > 1 else [], - *reversed(plan.box_dim), - *element_strides, - 0, # CU_TENSOR_MAP_INTERLEAVE_NONE - plan.swizzle_mode.value, - 2, # CU_TENSOR_MAP_L2_PROMOTION_L2_128B - oob_fill_kind, - # Append the dtype override ONLY when requested, so the default - # (derive-from-dtype) path emits a byte-identical encode call (the - # C++ wrapper treats the trailing arg as optional). - *([force_cu_dtype] if force_cu_dtype >= 0 else []), - ) - T.tvm_kernel_replace_point() + def impl(): + emit_at(shared_ptr(), spec.coordinates) # fmt: on - sctx.add_init_stmt(create_tensor_map.body, host=True) - sctx.cache_set(tensormap_cache_key, tensor_map) + else: + flat_extent = plan.issue_extent() + + # fmt: off + @T.prim_func(check_well_formed=False) + def impl(): + for issue in T.unroll(flat_extent): + smem_offset, coordinates = T.meta_var(plan.offsets_and_coords(issue)) + simplified = T.meta_var( + _simplify_with_ranges( + [smem_offset, *coordinates], [(issue, flat_extent)], sctx + ) + ) + emit_at( + shared_ptr(simplified[0]), + simplified[1:], + ) + # fmt: on + + if selector_bind is not None: + body = tvm.tirx.SeqStmt([selector_bind, impl.body]) + impl = PrimFunc([], body, ret_type=None, buffer_map={}).with_attr("global_symbol", "impl") + return impl + + +def copy_tma_auto_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + """Lower one ``tma_auto`` call.""" + plan = _build_auto_plan(op_call, sctx) + tensor_map, key = _get_or_encode_descriptor(plan.spec, sctx) + impl = _emit_plan(plan, tensor_map, None, False, sctx) if bool(op_call.config.get("prefetch_tensormap", False)): - if "warp_id_in_cta" not in sctx.launch_params: - fail("tma prefetch_tensormap requires warp_id_in_cta launch param") - prefetch_cache_key = f"prefetch_tensormap:{tensormap_cache_key}" - if sctx.cache_get(prefetch_cache_key) is None: - warp_id_in_cta = sctx.launch_params["warp_id_in_cta"].var - - # fmt: off - @T.prim_func(check_well_formed=False) - def prefetch_tensor_map(): - if warp_id_in_cta == 0: - T.ptx.prefetch_tensormap(T.address_of(tensor_map)) - T.tvm_kernel_replace_point() - # fmt: on - - sctx.add_init_stmt(prefetch_tensor_map.body) - sctx.cache_set(prefetch_cache_key, tensor_map) + _prefetch_main_descriptor(tensor_map, key, sctx) + return impl + + +def copy_tma_explicit_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + """Lower one direct TensorMap and exactly one TMA instruction.""" + plan, candidates = _build_explicit_plan(op_call, sctx) + main_map, main_key = _get_or_encode_descriptor(plan.spec, sctx) + candidate_maps = [] + for condition, candidate_spec in candidates: + candidate_map, _ = _get_or_encode_descriptor(candidate_spec, sctx) + candidate_maps.append((condition, candidate_map)) + selected_map, selector_bind, selected_is_address = _selected_descriptor( + main_map, candidate_maps + ) + impl = _emit_plan(plan, selected_map, selector_bind, selected_is_address, sctx) + if bool(op_call.config.get("prefetch_tensormap", False)): + _prefetch_main_descriptor(main_map, main_key, sctx) return impl -# Variant: copy_async/tma (priority=10). Applies at single-thread exec scope -# on Hopper+ (SM90+) for global↔shared copies; DispatchFail otherwise. +def _validate_tma_copy_op(op_call: TilePrimitiveCall, _sctx: DispatchContext) -> bool: + dst_region, src_region = op_call.args[:2] + src = src_region.buffer + dst = dst_region.buffer + if src.layout is None or dst.layout is None: + return False + src_scope, dst_scope = src.scope(), dst.scope() + return (src_scope == "global" and dst_scope.startswith("shared")) or ( + src_scope.startswith("shared") and dst_scope == "global" + ) + + +_COMMON_PREDICATES = [ + predicate( + "validate_tma_copy_op", + lambda op, sctx: (_validate_tma_copy_op(op, sctx), "not a global<->shared TMA copy"), + ), + predicate( + "single_thread", + lambda op, sctx: ( + single_thread(op, sctx), + f"unsupported exec_scope {sctx.exec_scope}, expected single thread", + ), + ), +] + + @register_dispatch( "copy_async", "cuda", - variant="tma", + variant="tma_auto", priority=10, - when=[ - predicate( - "validate_copy_op", lambda op, sctx: (validate_copy_op(op, sctx), "not a valid copy op") - ), - predicate( - "single_thread", - lambda op, sctx: ( - single_thread(op, sctx), - f"unsupported exec_scope {sctx.exec_scope}, expected single thread", - ), - ), - ], + when=_COMMON_PREDICATES, +) +def copy_async_dispatch_tma_auto(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + return copy_tma_auto_impl(op, sctx) + + +@register_dispatch( + "copy_async", + "cuda", + variant="tma_explicit", + priority=10, + when=_COMMON_PREDICATES, ) -def copy_async_dispatch_tma(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: - return copy_tma_impl(op, sctx) +def copy_async_dispatch_tma_explicit(op: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: + return copy_tma_explicit_impl(op, sctx) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py index 2603e7ac0345..d6d4cf754d78 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py @@ -22,7 +22,9 @@ """ from tvm.arith import Analyzer -from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import Layout, TileLayout + +from ..layout_utils import strip_swizzle_to_tile def find_contiguous_region(layout: TileLayout) -> tuple: @@ -71,8 +73,4 @@ def find_contiguous_region(layout: TileLayout) -> tuple: def to_tile_layout(layout: Layout, shape: list[int]) -> TileLayout: """Normalize any layout kind to a TileLayout for pointer arithmetic.""" - if isinstance(layout, ComposeLayout): - return layout.tile_layout - if isinstance(layout, SwizzleLayout): - return TileLayout(S[tuple(shape)]) - return layout + return strip_swizzle_to_tile(layout, lambda: shape) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py index 4cb242bf6038..16e400728b11 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py @@ -17,7 +17,7 @@ """Shared layout / vec-selection / emit helpers for ``reg.py`` and ``smem.py``. -Borrows directly from ``cuda/copy/reg.py`` (induced partition) and +Borrows directly from ``cuda/copy/vec_auto_reg.py`` (induced partition) and ``cuda/copy/_common.py`` (synthesized partition), extended to N operands. The dispatch split mirrors copy: @@ -40,7 +40,7 @@ # Re-use copy's primitives (PR-640) — same algorithm, same scope_id machinery. from ..copy._common import _TID_AXIS_FOR_SCOPE, _extract_tile, _thread_cnt -from ..copy.reg import _all_threads_active, _axis_decl, _compute_perm_r +from ..copy.vec_auto_reg import _all_threads_active, _axis_decl, _compute_perm_r # ----------------------------------------------------------------------------- diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py index b782a7b7be4e..be4de45257fe 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -"""Unary elementwise ops: zero / fill / reciprocal / sqrt / exp / exp2 / silu. +"""Unary elementwise ops: zero / fill / reciprocal / sqrt / exp / exp2 / log2 / silu. All carry the same ``T.(dst, src[, bias, scale])`` shape (bias / scale optional; ``silu`` ignores bias/scale to preserve legacy behavior). @@ -117,5 +117,6 @@ def _compute_silu(src_vals, extras, dt): "sqrt": OpSpec("sqrt", _parse_unary, _with_bias_scale(T.sqrt), _check_unary_extras), "exp": OpSpec("exp", _parse_unary, _with_bias_scale(T.exp), _check_unary_extras), "exp2": OpSpec("exp2", _parse_unary, _with_bias_scale(T.exp2), _check_unary_extras), + "log2": OpSpec("log2", _parse_unary, _with_bias_scale(T.log2), _check_unary_extras), "silu": OpSpec("silu", _parse_unary, _compute_silu, _check_unary_extras), } diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py index 2c276d71eecb..8288954b9950 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py @@ -17,7 +17,7 @@ """Elementwise dispatch when all operands live in ``local`` (registers). -Mirrors ``cuda/copy/reg.py``: the partition is *induced* by the layout that +Mirrors ``cuda/copy/vec_auto_reg.py``: the partition is *induced* by the layout that carries thread-axis info (the "anchor" operand). The region slice is absorbed into the sliced layout up front via ``align_operands_to_anchor`` — emit operates on a flat 1D per-thread view and indexes it with a scalar offset, so diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py index 404d7d0c1415..5abee7ab10db 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py @@ -17,7 +17,7 @@ """Elementwise dispatch when all operands live in ``shared*``. -Mirrors ``cuda/copy/gmem_smem.py``: no operand carries a per-thread partition, +Mirrors ``cuda/copy/vec_auto_gmem_smem.py``: no operand carries a per-thread partition, so partition is *synthesized* from ``sctx.intra`` (``thread_cnt = ∏ intra``). Each thread takes ``ceildiv(total, vec_chunk * thread_cnt)`` strided chunks. diff --git a/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py index 54deb5108c68..907ccc3b7b59 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py @@ -21,7 +21,7 @@ from tvm.script import tirx as T from tvm.tirx import PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def macro_or_prim_func(macro: Callable, need_macro: bool = False) -> Callable: diff --git a/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py index b069556843e0..9ada98240f29 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py @@ -29,7 +29,7 @@ predicate, register_dispatch, ) -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall @dataclass(frozen=True) diff --git a/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py index 637ca1378796..bfdde9b3eeb4 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py @@ -40,12 +40,21 @@ TileLayout, TLane, tmem_datapath_layout, + tmem_mma_operand_layout, ) from tvm.tirx.operator.tile_primitive import DispatchContext, predicate, register_dispatch -from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt, TilePrimitiveCall +from tvm.tirx.stmt import AllocBuffer, Evaluate, SeqStmt +from tvm.tirx.tile_primitive import TilePrimitiveCall +from ...intrinsics.tcgen05 import ( + _TCGEN05_MMA_TRANS_DTYPES, + _check_tcgen05_mma_matrix_shape, + _get_tcgen05_mma_kind, +) +from ...intrinsics.types import PTXDataType from ..common import get_st_extent, smem_desc_add_16B_offset from ..exec_scope_utils import single_thread +from ..layout_utils import strip_swizzle_to_tile from ..tma_utils import ( SwizzleMode, get_swizzle_mode_from_layout, @@ -85,11 +94,13 @@ def _dtype_name(dtype) -> str: def _encode_instr_descriptor_dense_uint32( M, N, + K, d_dtype, a_dtype, b_dtype, trans_a, trans_b, + cta_group=1, neg_a=False, neg_b=False, sat_d=False, @@ -102,10 +113,34 @@ def _encode_instr_descriptor_dense_uint32( ``T.ptx.tcgen05.mma`` instead of allocating + encoding a per-dispatch local descriptor on every gemm_async call (which forces an inline ``asm`` block that ptxas cannot hoist out of the i_kv loop body). + + Mirrors the runtime encoder's validation + (``codegen_ptx_tcgen05_encode_instr_descriptor``): the compile-time fold + must not accept kind/shape/trans combinations the runtime encoder would + reject — e.g. cta_group=1 M=128 requires N % 16 == 0, which the tile + chooser alone does not guarantee (it only enforces N % 8). """ d_dtype = _dtype_name(d_dtype) a_dtype = _dtype_name(a_dtype) b_dtype = _dtype_name(b_dtype) + + # PTXDataType spells tensor-float32 as "tf32". + _PTX_NAME = {"tensor_float32": "tf32"} + d_ptx_name = _PTX_NAME.get(d_dtype, d_dtype) + a_ptx_name = _PTX_NAME.get(a_dtype, a_dtype) + b_ptx_name = _PTX_NAME.get(b_dtype, b_dtype) + kind = _get_tcgen05_mma_kind(d_ptx_name, a_ptx_name, b_ptx_name) + if kind not in ("f16", "tf32", "f8f6f4", "i8"): + raise ValueError( + f"Check failed for Data Type Kind. d_dtype: {d_dtype}, " + f"a_dtype: {a_dtype}, b_dtype: {b_dtype}" + ) + _check_tcgen05_mma_matrix_shape(kind, cta_group, int(M), int(N), int(K), is_sparse) + if trans_a and PTXDataType.from_string(a_ptx_name) not in _TCGEN05_MMA_TRANS_DTYPES: + raise ValueError(f"Invalid a_dtype for transpose: {a_dtype}") + if trans_b and PTXDataType.from_string(b_ptx_name) not in _TCGEN05_MMA_TRANS_DTYPES: + raise ValueError(f"Invalid b_dtype for transpose: {b_dtype}") + d_format = _INSTR_DESC_FORMAT_MAP[d_dtype] a_format = _INSTR_DESC_FORMAT_MAP[a_dtype] b_format = _INSTR_DESC_FORMAT_MAP[b_dtype] @@ -316,10 +351,83 @@ def _choose_mma_tile(M, N, cta_group, MMA_N_MIN): return M_mma, N_mma +def _get_explicit_mma_tile(config): + """Parse an optional physical tcgen05 instruction tile from op config. + + ``mma_m`` is the descriptor M (cluster-wide for ``cta_group=2``), while + ``mma_n`` is the descriptor N. The pair is all-or-nothing so a caller + cannot accidentally combine one explicit dimension with the heuristic + choice for the other. + """ + has_m = "mma_m" in config + has_n = "mma_n" in config + if has_m != has_n: + missing = "mma_n" if has_m else "mma_m" + raise ValueError( + "gemm_async[tcgen05]: config fields mma_m and mma_n must be provided " + f"together; missing {missing}" + ) + if not has_m: + return None + + values = [] + for name in ("mma_m", "mma_n"): + value = config[name] + if isinstance(value, bool): + raise ValueError( + f"gemm_async[tcgen05]: {name} must be a positive integer, got {value!r}" + ) + try: + value = operator.index(value) + except TypeError as err: + raise ValueError( + f"gemm_async[tcgen05]: {name} must be a positive integer, got {value!r}" + ) from err + if value <= 0: + raise ValueError(f"gemm_async[tcgen05]: {name} must be a positive integer, got {value}") + values.append(value) + return tuple(values) + + +def _validate_explicit_mma_tile( + explicit_mma_tile, + *, + M, + N, + cta_group, + MMA_K, + mma_kind, +): + """Validate and convert a physical instruction tile to per-CTA loop extents.""" + M_desc, N_mma = explicit_mma_tile + M_total = M * cta_group + if M_desc != M_total: + raise ValueError( + "gemm_async[tcgen05]: explicit mma_m must match the physical " + f"instruction M derived from the operand/layout tile: got mma_m={M_desc}, " + f"but M={M} and cta_group={cta_group} require {M_total}. A smaller " + "mma_m would select a different TMEM datapath layout." + ) + if N % N_mma != 0: + raise ValueError( + "gemm_async[tcgen05]: explicit mma_n must divide the derived output " + f"N exactly, got mma_n={N_mma}, N={N}" + ) + _check_tcgen05_mma_matrix_shape( + mma_kind, + cta_group, + M_desc, + N_mma, + MMA_K, + False, + ) + return M_desc // cta_group, N_mma + + def _layout_matches_datapath_f(tmem_buf) -> bool: """Return True if ``tmem_buf.layout`` structurally equals Layout F (M=64 scattered) over the buffer's full (64, X) shape — i.e. the buffer was - allocated via ``tmem_pool.alloc((64, X), datapath="F")``. + allocated with ``layout=tmem_datapath_layout("F", 64, X)``. Used by the C-operand layout check to accept M=64 MMA writes into Layout F C buffers (the canonical pairing for M=64 outputs that are read back @@ -328,7 +436,10 @@ def _layout_matches_datapath_f(tmem_buf) -> bool: if tmem_buf.layout is None or int(tmem_buf.shape[0]) != 64: return False try: - expected = tmem_datapath_layout("F", 64, tmem_buf.shape[1]).canonicalize() + base = tmem_datapath_layout("F", 64, tmem_buf.shape[1]) + expected = TileLayout.from_iters( + base.shard, base.replica, tmem_buf.layout.offset + ).canonicalize() tvm.ir.assert_structural_equal(tmem_buf.layout.canonicalize(), expected) return True except (AssertionError, ValueError): @@ -355,7 +466,12 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - - args[5:8]: transA, transB, accum flags Config: - config["cta_group"]: CTA group in tcgen05 instructions (default 1) + - config["mma_m"], config["mma_n"]: Optional explicit physical + instruction tile. Both fields are required together. ``mma_m`` + is cluster-wide for cta_group=2. - config["descI"]: Optional pre-encoded instruction descriptor + (block-scaled only; the dense path always self-encodes and + raises if descI is passed) sctx: Schedule context (single-thread or warp execution scope) Returns: @@ -403,6 +519,9 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - # fp32/bf16 storage may still use tf32 MMA semantics via is_AB_tf32. is_AB_tf32 = op_call.config.get("is_AB_tf32", False) + # Emit the PTX ``tcgen05.mma.ws`` weight-stationary form only for kernels + # that explicitly require that tcgen05 ABI. + weight_stationary = bool(op_call.config.get("weight_stationary", False)) A_sem = "tf32" if is_AB_tf32 else A_type B_sem = "tf32" if is_AB_tf32 else B_type @@ -477,8 +596,13 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - cta_group = op_call.config.get("cta_group", 1) assert cta_group in [1, 2], f"tcgen05 schedule expected cta_group=1 or 2, got {cta_group}" - # descI: pre-encoded instruction descriptor (uint32), if None we encode it locally + # descI (pre-encoded uint32 instruction descriptor): rejected on the dense + # path (dispatcher encodes it); block-scaled callers may still pass it in. descI = op_call.config.get("descI", None) + if descI is not None and not is_block_scaled: + raise ValueError( + "descI was removed: the dispatcher encodes the instruction descriptor itself" + ) C_elem_size = DataType(C_type).bits C_elem_per_32b = 32 // C_elem_size @@ -488,19 +612,21 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - A_slice_layout = A_buffer.layout.slice(A_buffer.shape, A_buffer_region.region) B_slice_layout = B_buffer.layout.slice(B_buffer.shape, B_buffer_region.region) C_slice_layout = C_buffer.layout.slice(C_buffer.shape, C_buffer_region.region) - # Extract pre-swizzle tile layout for descriptor offset computation + # Extract pre-swizzle tile layout for descriptor offset computation. + # strip_swizzle_to_tile rebuilds an identity tile over the slice extents for + # a bare swizzle (ComposeLayout with a trivial tile), so .apply() maps over + # the full operand extent instead of wrapping inside one swizzle period. if not a_is_tmem: - A_slice_tile = ( - A_slice_layout.tile_layout - if isinstance(A_slice_layout, ComposeLayout) - else A_slice_layout - ) - B_slice_tile = ( - B_slice_layout.tile_layout if isinstance(B_slice_layout, ComposeLayout) else B_slice_layout - ) - - assert len(C_extent) == 2 and len(A_extent) >= 2 and len(B_extent) >= 2, ( - "Only 2D C, A, B are supported for gemm" + A_slice_tile = strip_swizzle_to_tile(A_slice_layout, lambda: A_extent) + B_slice_tile = strip_swizzle_to_tile(B_slice_layout, lambda: B_extent) + + # C may be batched C[2, M, N]: the M=64 .ws Layout-E lane fold (partials at + # lanes m, m+64), normalized to packed C_slice_layout. Squeeze leading unit dims. + while len(C_extent) > 2 and int(C_extent[0]) == 1: + C_extent = C_extent[1:] + C_batched = len(C_extent) == 3 and int(C_extent[0]) == 2 + assert (len(C_extent) == 2 or C_batched) and len(A_extent) >= 2 and len(B_extent) >= 2, ( + "Only 2D C (or the batched .ws form C[2, M, N]), A, B are supported for gemm" ) def _mat_dim_vals(extent, name): @@ -511,10 +637,41 @@ def _mat_dim_vals(extent, name): ) return vals[0], vals[1] - M = int(C_extent[-2]) - N = int(C_extent[-1]) + if C_batched: + M = int(C_extent[1]) + N = int(C_extent[2]) * 2 + N_half = N // 2 + batched_base = TileLayout(S[(2, M, N_half) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]) + expected_batched = TileLayout.from_iters( + batched_base.shard, batched_base.replica, C_slice_layout.offset + ).canonicalize() + tvm.ir.assert_structural_equal(C_slice_layout.canonicalize(), expected_batched) + packed_norm = TileLayout(S[(M, 2, N_half) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + C_slice_layout = TileLayout.from_iters( + packed_norm.shard, packed_norm.replica, C_slice_layout.offset + ).canonicalize() + # Batched C[2, M, N] is unambiguously the M=64 .ws fold, so .ws is + # inferred — the caller need not pass weight_stationary=True. + if cta_group == 1: + weight_stationary = True + else: + M = int(C_extent[-2]) + N = int(C_extent[-1]) is_2x2 = M == 64 and cta_group == 2 + # A may be batched A[2, M, K] for the M=64 .ws A-in-TMEM datapath: leading dim + # = the two 64-lane halves; normalized to 2D [M, K]. TMEM A only; SMEM A stays 2D. + A_batched = a_is_tmem and len(A_extent) == 3 and int(A_extent[0]) == 2 + if A_batched: + A_bM = int(A_extent[1]) + A_bK = int(A_extent[2]) + a_batched_base = TileLayout(S[(2, A_bM, A_bK) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]) + expected_a_batched = TileLayout.from_iters( + a_batched_base.shard, a_batched_base.replica, A_slice_layout.offset + ).canonicalize() + tvm.ir.assert_structural_equal(A_slice_layout.canonicalize(), expected_a_batched) + A_extent = A_extent[1:] + # Majorness (a_mn_major / b_mn_major) is determined later by # compute_canonical_params via dual-atom matching on the physical # SMEM layout. Extract dim extents here for cross-validation. @@ -564,6 +721,11 @@ def _try_atom(atom, atom_shape): return None tiler_shape = [s // a for s, a in zip(shape_2d, atom_shape)] tiler_grouped, seps = tiler.canonicalize().group(tiler_shape) + # Each tiler dim must group into exactly one iter: ldo/sdo read only + # the innermost, so multi-iter groups walk wrong addresses — reject. + seps = list(seps) + if seps != list(range(len(tiler_shape) + 1)): + return None elem_per_128b = 128 // tvm.DataType(dtype).bits # extent==1 leading dim -> unused LBO/SBO offset. @@ -585,11 +747,17 @@ def _atom_off(dim): base_shape = mma_atom_shape(dtype, mode) # [8, T*s] swapped_shape = [base_shape[1], base_shape[0]] # [T*s, 8] - # MN-major atom: compose SwizzleLayout with stride-reversed TileLayout - # so the first dim (T*s) is contiguous instead of the second. - # Needed when the penultimate dim is physically contiguous. + # MN-major atom: carry the swizzle params over a stride-reversed + # TileLayout so the first dim (T*s) is contiguous instead of the + # second. Needed when the penultimate dim is physically contiguous. mn_tile = TileLayout(S[tuple(swapped_shape) : (1, swapped_shape[0])]) - mn_atom = ComposeLayout(swizzle_atom, mn_tile) + mn_atom = ComposeLayout( + swizzle_atom.per_element, + swizzle_atom.swizzle_len, + swizzle_atom.atom_len, + mn_tile, + swizzle_atom.swizzle_inner, + ) # Determine K-major vs MN-major based on which dim is contiguous. # K-major: K dim contiguous (last dim for [MN,K], first dim for [K,MN]) @@ -678,6 +846,61 @@ def _atom_off(dim): assert len(shape_2d) == 2, ( f"Expected exactly 2 non-unit dims in region {[int(r.extent) for r in desc_region]}" ) + if phys_mode == SwizzleMode.SWIZZLE_NONE: + if not isinstance(slice_layout, TileLayout): + raise ValueError( + "gemm_async: no-swizzle SMEM descriptors require a plain TileLayout" + ) + grouped, _ = slice_layout.canonicalize().group(shape_2d) + shard = list(grouped.shard) + elem_per_16B = 128 // DataType(dtype).bits + if len(shard) >= 3: + dim0, dim1, dim2 = shard[-3], shard[-2], shard[-1] + is_packed_16B = ( + int(dim0.extent) == shape_2d[0] + and int(dim0.stride) == elem_per_16B + and int(dim1.extent) * int(dim2.extent) == shape_2d[1] + and int(dim1.stride) == shape_2d[0] * elem_per_16B + and int(dim2.extent) == elem_per_16B + and int(dim2.stride) == 1 + ) + if is_packed_16B: + # SBO = literal 8 (8-row group pitch = 128B) regardless of dtype + # (PTX §9.7.16.3.2); only validated for 16-bit — reject others. + if elem_per_16B != 8: + raise ValueError( + f"gemm_async: no-swizzle packed-16B SMEM descriptors are " + f"only supported for 16-bit dtypes (8 elements per 16B " + f"line); got dtype {dtype} with {elem_per_16B} elements " + f"per 16B line" + ) + return SwizzleMode.SWIZZLE_NONE, shape_2d[0], 8, is_transposed + + dim0, dim1 = shard[-2], shard[-1] + is_col_major = int(dim0.stride) == 1 and int(dim1.stride) == shape_2d[0] + if not is_col_major: + raise ValueError( + "gemm_async: no-swizzle SMEM descriptors currently require the " + "penultimate matrix dimension to be contiguous or 16B-packed" + ) + if shape_2d[0] != 8 * elem_per_16B: + # SBO=8 (128B row-group pitch) only agrees with the packed ABI when the + # contiguous dim spans exactly 128B (8*elem_per_16B, 64 for bf16) — reject else. + raise ValueError( + f"gemm_async: no-swizzle column-major-view SMEM descriptors " + f"require the contiguous matrix dimension to span exactly 128B " + f"(= {8 * elem_per_16B} elements for {dtype}); got " + f"{shape_2d[0]}. Use a swizzled or packed-16B layout for " + f"larger tiles." + ) + # No-swizzle branch for FlashMLA's column-major S-tile views: ldo=K-column + # stride, sdo advances one 16B group. Encoding is K-major, so return is_transposed. + ldo = int(dim1.stride) + # SBO = literal 8 (8-row group pitch = 128B in 16B units, + # PTX §9.7.16.3.2), valid on the domain enforced above. + sdo = 8 + return SwizzleMode.SWIZZLE_NONE, ldo, sdo, is_transposed + result = _match(slice_layout, shape_2d) if result is not None: return result @@ -715,6 +938,21 @@ def _atom_off(dim): B_buffer, B_buffer_region, B_type, transB ) + # tf32 is only PTX-legal MN-major with the 128B-32B-atomicity swizzle (PTX Table + # 52), which the matcher never produces — reject tf32 MN-major; keep tf32 K-major. + if A_sem in ("tf32", "tensor_float32") and not a_is_tmem and a_mn_major: + raise ValueError( + "gemm_async: tf32 A operand matched an MN-major SMEM layout, which is " + "PTX-illegal without 128B-32B-atomicity swizzle support; lay A out " + "K-major (K contiguous) instead" + ) + if B_sem in ("tf32", "tensor_float32") and b_mn_major: + raise ValueError( + "gemm_async: tf32 B operand matched an MN-major SMEM layout, which is " + "PTX-illegal without 128B-32B-atomicity swizzle support; lay B out " + "K-major (K contiguous) instead" + ) + # Extract K from A dims using transA (shape order). # transA tells us which dim is K; a_mn_major tells us the layout orientation. # transA=False [M, K]: K = dim[-1]; transA=True [K, M]: K = dim[-2] @@ -731,7 +969,22 @@ def _atom_off(dim): MMA_K = 16 MMA_N_MIN = 8 if cta_group == 1 else 16 # Minimum N dimension - M_mma, N_mma = _choose_mma_tile(M, N, cta_group, MMA_N_MIN) + explicit_mma_tile = _get_explicit_mma_tile(op_call.config) + if explicit_mma_tile is None: + M_mma, N_mma = _choose_mma_tile(M, N, cta_group, MMA_N_MIN) + else: + if is_block_scaled: + mma_kind = _get_tcgen05_mma_kind(C_type, A_type, B_type, SFA_type, SFB_type) + else: + mma_kind = _get_tcgen05_mma_kind("float32", A_sem, B_sem) + M_mma, N_mma = _validate_explicit_mma_tile( + explicit_mma_tile, + M=M, + N=N, + cta_group=cta_group, + MMA_K=MMA_K, + mma_kind=mma_kind, + ) M_tiles = M // M_mma N_tiles = N // N_mma K_iters = K // MMA_K @@ -772,14 +1025,78 @@ def _atom_off(dim): # 2x2 layout: (M, 2, N//2):(1@TLane, 64@TLane, 1@TCol) # Layout F (M=64 scatter): the full TMEM buffer is shape (64, X) with the # scattered row→lane mapping from tmem_datapath_layout("F", 64, X). When - # the user allocates with ``tmem_pool.alloc(..., datapath="F")`` and slices + # the user allocates with that layout and slices # the full row range, the slice layout structurally matches Layout F over # (M=64, N) — assert against that base instead of the Layout D identity. - if is_2x2: - # Layout B (per-CTA M=64, .cta_group::2 "2x2"). Single-source the - # write layout with allocation and readback so the two sides cannot - # drift. - base = tmem_datapath_layout("B", M, N) + packed_n2 = False + if M == 64 and N % 2 == 0: + N_half = N // 2 + packed_base = TileLayout(S[(M, 2, N_half) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + expected_packed_layout = TileLayout.from_iters( + packed_base.shard, packed_base.replica, C_slice_layout.offset + ).canonicalize() + try: + tvm.ir.assert_structural_equal(C_slice_layout.canonicalize(), expected_packed_layout) + packed_n2 = True + except (AssertionError, ValueError): + packed_n2 = False + + # Packed Layout-E C (M, 2, N//2) is uniquely the cta_group::1 M=64 .ws datapath + # (PTX §9.7.16.10.5), so .ws is inferred; weight_stationary=False on it is rejected. + if packed_n2 and not is_2x2: + if op_call.config.get("weight_stationary") is False: + raise ValueError( + "gemm_async[tcgen05]: C uses the packed (M, 2, N//2):(1@TLane, " + "64@TLane, 1@TCol) Layout-E TMEM layout, which is the M=64 " + "weight-stationary datapath — but weight_stationary=False was " + "passed. Drop the flag (it is inferred from the layout), or " + "declare C with the Layout F / identity datapath layout for a " + "non-ws M=64 MMA." + ) + weight_stationary = True + + # Converse check: .ws (M=64, cta_group::1) writes the Layout-E fold, so reject a + # C declared as identity/Layout-F (the two banks would land at the wrong lanes). + if weight_stationary and cta_group == 1 and M == 64 and not packed_n2 and not is_2x2: + raise ValueError( + "gemm_async[tcgen05]: weight_stationary .ws (M=64, cta_group::1) writes " + "the Layout-E fold (M, 2, N//2):(1@TLane, 64@TLane, 1@TCol), but C is " + "declared with the identity / Layout-F layout, so the two banks would be " + "read from the wrong lanes. Declare C in the Layout-E (packed) or the " + "batched C[2, M, N//2] form for a .ws output, or drop weight_stationary " + "for a non-ws (Layout-F) M=64 MMA." + ) + + # A-side check: an M=64 .ws A-in-TMEM reads both 64-lane halves, which flat 2D + # A[M, K] (lanes 0-63 only) can't express — require the batched A[2, M, K] fold. + if a_is_tmem and weight_stationary and cta_group == 1 and M == 64 and not A_batched: + raise ValueError( + "gemm_async[tcgen05]: an M=64 cta_group::1 .ws GEMM with A in TMEM reads " + "A from both 64-lane halves (lanes 0-63 and 64-127), but A was given as a " + "flat 2D [M, K] tile that only describes lanes 0-63. Declare A in the " + "batched A[2, M, K] fold (layout (2, M, K):(64@TLane, 1@TLane, 1@TCol)), " + "the A-side dual of the batched C[2, M, N//2] form, so the two-half " + "occupancy is explicit and structurally checked." + ) + + # Converse guard: batched A[2, M, K] is valid only in the two per-CTA M=64 folds + # (Layout E: cta_group::1 .ws; Layout B: cta_group::2 dense) — reject any other. + _batched_ok = (weight_stationary and cta_group == 1 and M == 64) or ( + not weight_stationary and cta_group == 2 and M == 64 + ) + if A_batched and not _batched_ok: + raise ValueError( + "gemm_async[tcgen05]: batched A[2, M, K] is only valid for Layout E " + "(cta_group::1 .ws, per-CTA M=64) or Layout B (cta_group::2 dense, " + f"per-CTA M=64), but got weight_stationary={weight_stationary}, " + f"cta_group={cta_group}, M={M}. Declare A as a flat 2D [M, K] identity " + "for any other datapath." + ) + + if is_2x2 or packed_n2: + # Some cta_group::1 M=64 tiles use the 2x2 logical-N/physical-column packing: + # N is split across two TLane banks, so N=128 occupies 64 physical TMEM columns. + base = TileLayout(S[(M, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) elif ( M == 64 and int(C_buffer.shape[0]) == 64 @@ -795,18 +1112,39 @@ def _atom_off(dim): tvm.ir.assert_structural_equal(C_slice_layout.canonicalize(), expected_c_layout) assert C_buffer.allocated_addr is not None tmem_addr = C_buffer.allocated_addr[0] + tmem_lane_offset = C_slice_layout.offset.get(TLane, 0) tmem_offset_32b = C_slice_layout.offset.get(TCol, 0) - # Validate TMEM A layout: (A_dim2, A_dim1):(1@TLane, 1@TCol) + # Validate TMEM A via the same resolver as tmem_pool.alloc_tcgen05_mma_A. + # `M` here is per-CTA rows; the resolver takes the PTX instruction M. if a_is_tmem: - A_tmem_base = TileLayout(S[(A_dim2, A_dim1) : (1 @ TLane, 1 @ TCol)]) - expected_a_layout = TileLayout.from_iters( - A_tmem_base.shard, A_tmem_base.replica, A_slice_layout.offset - ).canonicalize() - tvm.ir.assert_structural_equal(A_slice_layout.canonicalize(), expected_a_layout) + if not A_batched: + instruction_M = M * cta_group + try: + A_tmem_base = tmem_mma_operand_layout( + "A", + (A_dim2, A_dim1), + A_type, + M=instruction_M, + cta_group=cta_group, + ws=weight_stationary, + ) + except ValueError as err: + raise ValueError(f"gemm_async[tcgen05]: invalid TMEM A layout: {err}") from err + expected_a_layout = TileLayout.from_iters( + A_tmem_base.shard, A_tmem_base.replica, A_slice_layout.offset + ).canonicalize() + try: + tvm.ir.assert_structural_equal(A_slice_layout.canonicalize(), expected_a_layout) + except (AssertionError, ValueError) as err: + raise ValueError( + "gemm_async[tcgen05]: TMEM A layout does not match " + "tmem_pool.alloc_tcgen05_mma_A semantic layout" + ) from err assert A_buffer.allocated_addr is not None, "TMEM A buffer must have allocated_addr" A_tmem_addr = A_buffer.allocated_addr[0] A_elem_per_32b = 32 // DataType(A_type).bits + A_tmem_lane_offset = A_slice_layout.offset.get(TLane, 0) # TCol offset is in element units (not 32-bit columns) for sub-32-bit dtypes. # Convert to 32-bit column units for get_tmem_addr. A_tmem_offset_32b = A_slice_layout.offset.get(TCol, 0) // A_elem_per_32b @@ -844,6 +1182,10 @@ def _make_lo_uniform(desc): ) def _make_desc(smem_buf, ldo, sdo, swizzle_val, name): + # Warp-scoped calls encode in all active lanes before electing the MMA + # issuer, so make the descriptor low word uniform there. A + # single-thread caller is already elected: a full-mask shuffle in that + # scope is invalid, and the same thread consumes the descriptor anyway. desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name=name, scope="local") encode_call = tvm.tirx.call_intrin( "", @@ -854,14 +1196,11 @@ def _make_desc(smem_buf, ldo, sdo, swizzle_val, name): sdo, swizzle_val, ) - wrap = SeqStmt( - [ - AllocBuffer(desc_buf), - Evaluate(encode_call), - Evaluate(_make_lo_uniform(desc_buf[0])), - _krp, - ] - ) + wrap_stmts = [AllocBuffer(desc_buf), Evaluate(encode_call)] + if warp_scope: + wrap_stmts.append(Evaluate(_make_lo_uniform(desc_buf[0]))) + wrap_stmts.append(_krp) + wrap = SeqStmt(wrap_stmts) sctx.add_post_buffer_def_stmt(smem_buf, wrap) return desc_buf @@ -878,6 +1217,23 @@ def _uniform_desc(smem_buf, off16, ldo, sdo, swizzle): lo = T.bitwise_or(T.uint32(lo_const), sa) if lo_const else sa return T.bitwise_or(T.shift_left(T.uint64(const_hi), T.uint64(32)), T.cast(lo, "uint64")) + def _desc_ptr(smem_buf, off16): + base_ptr = smem_buf.ptr_to([0] * len(smem_buf.shape)) + return T.handle_add_byte_offset(base_ptr, off16 * 16) + + def _encoded_desc_val(smem_buf, off16, ldo, sdo, swizzle): + desc = T.alloc_local((1,), "uint64") + T.evaluate( + T.ptx.tcgen05.encode_matrix_descriptor( + T.address_of(desc[0]), + _desc_ptr(smem_buf, off16), + ldo=ldo, + sdo=sdo, + swizzle=swizzle, + ) + ) + return desc[0] + def _b_offset(ni, ki): B_linear = ( ki * MMA_K * B_extent[-1] + ni * N_mma_per_cta @@ -894,14 +1250,24 @@ def _a_offset(mi, ki): ) return tvm.tirx.floordiv(A_slice_tile.apply(A_linear)["m"], A_elem_per_16B) - # smem_desc: "hoist" (default) or "recompute" per-MMA descriptor update. - use_add = op_call.config.get("smem_desc", "hoist") != "recompute" + # smem_desc modes: "hoist" (default, encode once after alloc), "local_hoist" + # (encode at call site, reuse via add_16B_offset), "encode"/"recompute" (per MMA). + smem_desc_mode = op_call.config.get("smem_desc", "hoist") + local_hoist = smem_desc_mode == "local_hoist" + encode_per_mma = smem_desc_mode == "encode" + use_add = smem_desc_mode not in ("recompute", "encode") + if encode_per_mma and is_block_scaled: + raise ValueError("tcgen05 gemm_async smem_desc='encode' is only implemented for dense MMA") descB_buf = ( - _make_desc(B_buffer, B_ldo, B_sdo, B_swizzle_mode.value, "descB") if use_add else None + _make_desc(B_buffer, B_ldo, B_sdo, B_swizzle_mode.value, "descB") + if use_add and not local_hoist + else None ) A_use_add = use_add and not a_is_tmem descA_buf = ( - _make_desc(A_buffer, A_ldo, A_sdo, A_swizzle_mode.value, "descA") if A_use_add else None + _make_desc(A_buffer, A_ldo, A_sdo, A_swizzle_mode.value, "descA") + if A_use_add and not local_hoist + else None ) B_use_add = use_add @@ -909,19 +1275,32 @@ def _a_offset(mi, ki): def _b_desc_val(descB_in, ni, ki): B_offset = _b_offset(ni, ki) if B_use_add: - return smem_desc_add_16B_offset(descB_buf[0], B_offset) + descB_base = descB_in[0] if local_hoist else descB_buf[0] + return smem_desc_add_16B_offset(descB_base, B_offset) return _uniform_desc(B_buffer, B_offset, B_ldo, B_sdo, B_swizzle_mode.value) + def _get_tmem_addr_fast(base, row_offset, col_offset): + row_offset = analyzer.simplify(row_offset) + col_offset = analyzer.simplify(col_offset) + # get_tmem_addr(base, 0, col) folds to base + col for all layouts here + # (small constant col, no 16-bit overflow), saving a device helper call. + if analyzer.can_prove_equal(row_offset, 0): + if analyzer.can_prove_equal(col_offset, 0): + return base + return analyzer.simplify(base + col_offset) + return T.cuda.get_tmem_addr(base, row_offset, col_offset) + # Helper: compute A operand (TMEM address or SMEM descriptor) for a given (mi, ki) tile def _a_operand(mi, ki, descA_in=None): if a_is_tmem: # A is [M, K] non-transposed: M→TLane (rows), K→TCol (cols) - a_row = mi * M_mma + a_row = A_tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma) a_col = A_tmem_offset_32b + ki * (MMA_K // A_elem_per_32b) - return T.cuda.get_tmem_addr(A_tmem_addr, a_row, a_col) + return _get_tmem_addr_fast(A_tmem_addr, a_row, a_col) A_offset = _a_offset(mi, ki) if A_use_add: - return smem_desc_add_16B_offset(descA_buf[0], A_offset) + descA_base = descA_in[0] if local_hoist else descA_buf[0] + return smem_desc_add_16B_offset(descA_base, A_offset) return _uniform_desc(A_buffer, A_offset, A_ldo, A_sdo, A_swizzle_mode.value) if is_block_scaled: @@ -932,15 +1311,21 @@ def _a_operand(mi, ki, descA_in=None): sfa_base = SFA_buffer.allocated_addr[0] sfb_base = SFB_buffer.allocated_addr[0] - # Compute initial SFA/SFB addresses (for ki=0) - # apply(0)["TCol"] at row 0 gives physical TCol offset - sfa_tcol_0 = SFA_slice_layout.apply(0).get("TCol", 0) - sfb_tcol_0 = SFB_slice_layout.apply(0).get("TCol", 0) - SFA_init_addr = analyzer.simplify( - sfa_base + tvm.tirx.floordiv(sfa_tcol_0, SFA_elem_per_col) + # Compute initial SFA/SFB addresses (for ki=0). Both physical axes are + # part of a TMEM address: TLane occupies the high half-word and TCol + # the low half-word. Dropping TLane silently redirects an explicitly + # banded scale view to row zero. + sfa_coord_0 = SFA_slice_layout.apply(0) + sfb_coord_0 = SFB_slice_layout.apply(0) + sfa_tlane_0 = sfa_coord_0.get("TLane", 0) + sfb_tlane_0 = sfb_coord_0.get("TLane", 0) + sfa_tcol_0 = sfa_coord_0.get("TCol", 0) + sfb_tcol_0 = sfb_coord_0.get("TCol", 0) + SFA_init_addr = _get_tmem_addr_fast( + sfa_base, sfa_tlane_0, tvm.tirx.floordiv(sfa_tcol_0, SFA_elem_per_col) ) - SFB_init_addr = analyzer.simplify( - sfb_base + tvm.tirx.floordiv(sfb_tcol_0, SFB_elem_per_col) + SFB_init_addr = _get_tmem_addr_fast( + sfb_base, sfb_tlane_0, tvm.tirx.floordiv(sfb_tcol_0, SFB_elem_per_col) ) # Rotate sf_id per ki when multiple ki share one SF column. @@ -952,7 +1337,7 @@ def _a_operand(mi, ki, descA_in=None): # Physical TMEM columns per MMA N tile. # 2x2 layout (Layout B): each MMA tile spans N_mma/2 physical columns # and uses rows 64-127 for the other half. - N_mma_phys_cols = N_mma // 2 if is_2x2 else N_mma + N_mma_phys_cols = N_mma // 2 if is_2x2 or packed_n2 else N_mma # Build main_impl: descA_in is None when A is in TMEM (ignored by _a_operand). # fmt: off @@ -968,14 +1353,26 @@ def main_impl(descA_in, descB_in, descI_in): should_accum = T.meta_var(tvm.tirx.any(ki != 0, accum_expr)) sfa_linear = mi * M_mma * SFA_K_total + ki * sfa_elems_per_ki sfb_linear = ni * N_mma_per_cta * SFB_K_total + ki * sfb_elems_per_ki - # Fold tcol to constants shared by SF addr and sf_id. - sfa_tcol = T.meta_var(analyzer.simplify(SFA_slice_layout.apply(sfa_linear).get("TCol", 0))) # noqa: E501 - sfb_tcol = T.meta_var(analyzer.simplify(SFB_slice_layout.apply(sfb_linear).get("TCol", 0))) # noqa: E501 + # Fold the physical coordinates shared by SF addr and sf_id. + sfa_coord = T.meta_var(SFA_slice_layout.apply(sfa_linear)) + sfb_coord = T.meta_var(SFB_slice_layout.apply(sfb_linear)) + sfa_tlane = T.meta_var(analyzer.simplify(sfa_coord.get("TLane", 0))) + sfb_tlane = T.meta_var(analyzer.simplify(sfb_coord.get("TLane", 0))) + sfa_tcol = T.meta_var(analyzer.simplify(sfa_coord.get("TCol", 0))) + sfb_tcol = T.meta_var(analyzer.simplify(sfb_coord.get("TCol", 0))) sfa_addr = T.meta_var( - analyzer.simplify(sfa_base + tvm.tirx.floordiv(sfa_tcol, SFA_elem_per_col)) + _get_tmem_addr_fast( + sfa_base, + sfa_tlane, + tvm.tirx.floordiv(sfa_tcol, SFA_elem_per_col), + ) ) sfb_addr = T.meta_var( - analyzer.simplify(sfb_base + tvm.tirx.floordiv(sfb_tcol, SFB_elem_per_col)) + _get_tmem_addr_fast( + sfb_base, + sfb_tlane, + tvm.tirx.floordiv(sfb_tcol, SFB_elem_per_col), + ) ) if needs_sf_id: sf_id = T.meta_var(analyzer.simplify(tvm.tirx.floormod(sfa_tcol, SFA_elem_per_col))) # noqa: E501 @@ -983,9 +1380,12 @@ def main_impl(descA_in, descB_in, descI_in): tmem_col = T.meta_var( tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) ) + tmem_row = T.meta_var( + tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma) + ) if elect_pred: T.ptx.tcgen05.mma.block_scale( - T.cuda.get_tmem_addr(tmem_addr, mi * M_mma, tmem_col), + _get_tmem_addr_fast(tmem_addr, tmem_row, tmem_col), a_val, descB_val, sfa_addr, sfb_addr, descI_in, @@ -1003,39 +1403,120 @@ def main_impl(descA_in, descB_in, descI_in): # ``should_accum_ptr``, ``tmem_col_ptr``) which ptxas cannot fold # back into the operand and the resulting LMEM round-trips show up # on the fa4 hot path. - @T.inline - def main_impl(descA_in, descB_in, descI_in): - for mi in T.unroll(M_tiles): - for ni in T.unroll(N_tiles): - for ki in T.unroll(K_iters): - a_val = T.meta_var(_a_operand(mi, ki, descA_in)) - descB_val = T.meta_var(_b_desc_val(descB_in, ni, ki)) - should_accum = T.meta_var(tvm.tirx.any(ki != 0, accum_expr)) - tmem_col = T.meta_var( - tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) - ) - if elect_pred: - T.ptx.tcgen05.mma( - T.cuda.get_tmem_addr(tmem_addr, mi * M_mma, tmem_col), - a_val, descB_val, descI_in, - d_dtype="float32", a_dtype=A_sem, b_dtype=B_sem, - use_a_tmem=a_is_tmem, cta_group=cta_group, - enable_input_d=should_accum, + if encode_per_mma: + @T.inline + def main_impl(descA_in, descB_in, descI_in): + for mi in T.unroll(M_tiles): + for ni in T.unroll(N_tiles): + for ki in T.unroll(K_iters): + should_accum = T.meta_var(tvm.tirx.any(ki != 0, accum_expr)) + tmem_col = T.meta_var( + tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) + ) + tmem_row = T.meta_var( + tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma) + ) + if elect_pred: + descB_mma = T.meta_var( + _encoded_desc_val( + B_buffer, _b_offset(ni, ki), B_ldo, B_sdo, B_swizzle_mode.value + ) + ) + if a_is_tmem: + a_val = T.meta_var(_a_operand(mi, ki, descA_in)) + T.ptx.tcgen05.mma( + _get_tmem_addr_fast(tmem_addr, tmem_row, tmem_col), + a_val, descB_mma, descI_in, + d_dtype="float32", a_dtype=A_sem, b_dtype=B_sem, + use_a_tmem=a_is_tmem, cta_group=cta_group, + enable_input_d=should_accum, + weight_stationary=weight_stationary, + ) + else: + descA_mma = T.meta_var( + _encoded_desc_val( + A_buffer, + _a_offset(mi, ki), + A_ldo, + A_sdo, + A_swizzle_mode.value, + ) + ) + T.ptx.tcgen05.mma( + _get_tmem_addr_fast(tmem_addr, tmem_row, tmem_col), + descA_mma, descB_mma, descI_in, + d_dtype="float32", a_dtype=A_sem, b_dtype=B_sem, + use_a_tmem=a_is_tmem, cta_group=cta_group, + enable_input_d=should_accum, + weight_stationary=weight_stationary, + ) + else: + @T.inline + def main_impl(descA_in, descB_in, descI_in): + for mi in T.unroll(M_tiles): + for ni in T.unroll(N_tiles): + for ki in T.unroll(K_iters): + a_val = T.meta_var(_a_operand(mi, ki, descA_in)) + descB_val = T.meta_var(_b_desc_val(descB_in, ni, ki)) + should_accum = T.meta_var(tvm.tirx.any(ki != 0, accum_expr)) + tmem_col = T.meta_var( + tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) ) + tmem_row = T.meta_var( + tmem_lane_offset + (0 if M_tiles == 1 else mi * M_mma) + ) + if elect_pred: + T.ptx.tcgen05.mma( + _get_tmem_addr_fast(tmem_addr, tmem_row, tmem_col), + a_val, descB_val, descI_in, + d_dtype="float32", a_dtype=A_sem, b_dtype=B_sem, + use_a_tmem=a_is_tmem, cta_group=cta_group, + enable_input_d=should_accum, + weight_stationary=weight_stationary, + ) descA_val = None # descriptors built per-MMA from SMEM addr via _uniform_desc + @T.inline + def call_main(descI_arg): + if local_hoist: + descB_local = T.alloc_local((1,), "uint64") + T.ptx.tcgen05.encode_matrix_descriptor( + T.address_of(descB_local[0]), + B_buffer.ptr_to([0] * len(B_buffer.shape)), + ldo=B_ldo, + sdo=B_sdo, + swizzle=B_swizzle_mode.value, + ) + # local_hoist runs at the call site under elected-thread control, so the + # descriptor is consumed by the same thread and needs no warp shuffle. + if A_use_add: + descA_local = T.alloc_local((1,), "uint64") + T.ptx.tcgen05.encode_matrix_descriptor( + T.address_of(descA_local[0]), + A_buffer.ptr_to([0] * len(A_buffer.shape)), + ldo=A_ldo, + sdo=A_sdo, + swizzle=A_swizzle_mode.value, + ) + main_impl(descA_local, descB_local, descI_arg) + else: + main_impl(None, descB_local, descI_arg) + else: + main_impl(descA_val, descB_buf, descI_arg) + + # descI is not None only for block-scaled calls (dense descI raises above). if descI is not None and not needs_sf_id: @T.prim_func(check_well_formed=False) def impl(): - main_impl(descA_val, None, descI) + call_main(descI) elif descI is not None: # Local copy: main_impl rotates descI in-place per ki. @T.prim_func(check_well_formed=False) def impl(): descI_local: T.uint32 descI_local = descI - main_impl(descA_val, None, descI_local) + call_main(descI_local) elif is_block_scaled: @T.prim_func(check_well_formed=False) def impl(): @@ -1043,7 +1524,7 @@ def impl(): T.ptx.tcgen05.encode_instr_descriptor_block_scaled(T.address_of(descI_local), d_dtype=C_type, a_dtype=A_type, b_dtype=B_type, sfa_dtype=SFA_type, sfb_dtype=SFB_type, # noqa: E501, F821 sfa_tmem_addr=SFA_init_addr, sfb_tmem_addr=SFB_init_addr, # noqa: E501 M=M_mma * cta_group, N=N_mma, K=MMA_K, trans_a=a_mn_major, trans_b=b_mn_major, n_cta_groups=cta_group) # noqa: E501 - main_impl(descA_val, None, descI_local) # noqa: F821 + call_main(descI_local) # noqa: F821 else: # Pre-compute the dense instruction descriptor at dispatcher time so # the MMA's 4th operand is a literal ``uint32`` instead of a per-call @@ -1053,17 +1534,19 @@ def impl(): descI_value = _encode_instr_descriptor_dense_uint32( M=M_mma * cta_group, N=N_mma, + K=MMA_K, d_dtype="float32", a_dtype=A_sem, b_dtype=B_sem, trans_a=a_mn_major, trans_b=b_mn_major, + cta_group=cta_group, ) descI_const = tvm.tirx.const(descI_value, "uint32") @T.prim_func(check_well_formed=False) def impl(): - main_impl(descA_val, None, descI_const) + call_main(descI_const) # fmt: on return impl diff --git a/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py index 7531ffce838e..207999cdc2f8 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py @@ -20,7 +20,7 @@ from tvm.arith.analyzer import Analyzer from tvm.tirx import Buffer from tvm.tirx.operator.tile_primitive import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def validate_gemm_op(op_call: TilePrimitiveCall, sctx: DispatchContext) -> bool: diff --git a/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py index 2a46d33d9945..457d836e2e66 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py @@ -23,11 +23,45 @@ """ import functools +import math import operator from collections import defaultdict from tvm.arith import Analyzer -from tvm.tirx.layout import TileLayout +from tvm.tirx.layout import ComposeLayout, S, TileLayout + + +def strip_swizzle_to_tile(layout, get_extents): + """Strip swizzle off ``layout`` to a TileLayout for perm/group/slice. + + For a plain TileLayout, returns it unchanged. For a ComposeLayout, returns + the inner tile. For a *bare* swizzle (ComposeLayout with a trivial tile), + the tile carries only the swizzle period, which may not match the buffer or + region extents: rebuild an identity tile over the extents so downstream + group/slice sees the real per-axis sizes. When the trivial tile already + spans the extents (wrapped identity, or a bare swizzle whose period equals + the product), keep it as-is to preserve its dimension structure. + + ``get_extents`` is a no-arg callable returning the per-axis extent list + (buffer shape or region extents). It is only invoked for a bare swizzle, so + callers with symbolic regions pay no int()-coercion for plain TileLayouts. + If the extents or tile size can't be coerced to int (symbolic/dynamic + region), fall back to the period tile -- the pre-refactor behavior -- so + dynamic swizzled copies keep compiling instead of raising. + """ + if isinstance(layout, ComposeLayout): + tile = layout.tile_layout + if tile.is_trivial(): + try: + extents = get_extents() + tile_size = int(tile.size()) + buf_size = math.prod(int(s) for s in extents) + except (TypeError, ValueError): + return tile + if tile_size != buf_size: + return TileLayout(S[tuple(extents)]) + return tile + return layout def get_sublayout_from_region(layout, buffer_shape, region_st, region_extent): diff --git a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py index cf2d9b15561f..fa0398cc72f7 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py @@ -78,7 +78,7 @@ from tvm.tirx import BufferRegion, IntImm, PrimFunc, is_buffer_var from tvm.tirx.layout import TileLayout, _flatten_coord from tvm.tirx.operator.tile_primitive import DispatchContext, fail, register_dispatch -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import get_indices, get_st_extent diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py index 3316aaf2a1fb..446b027ff90c 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py @@ -68,7 +68,7 @@ from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import ReduceOpType from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import get_indices, get_st_extent from ..layout_utils import get_local_region, get_sublayout_from_region diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py index 38986df86bb0..fced38e81a3b 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py @@ -63,7 +63,7 @@ from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import ReduceOpType from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import get_indices, get_st_extent, next_power_of_2 from .utils import ( diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py index fd506fa14bf9..3725366f9b6f 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py @@ -50,7 +50,7 @@ from tvm.tirx.operator.tile_primitive import DispatchContext from tvm.tirx.operator.tile_primitive.common import ReduceOpType from tvm.tirx.operator.tile_primitive.dispatcher import predicate, register_dispatch -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import sm_version_ok from ..exec_scope_utils import exec_scope_ok diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py b/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py index 6248cddea117..ffb9b5cf1a1c 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py @@ -26,7 +26,7 @@ from tvm.tirx import BufferRegion from tvm.tirx.operator.tile_primitive import DispatchContext from tvm.tirx.operator.tile_primitive.common import ReduceOpType -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import match_scope diff --git a/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py b/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py index 625b8ff5caca..99886ff363ca 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py @@ -22,7 +22,7 @@ import tvm from tvm.arith.analyzer import Analyzer -from tvm.tirx.layout import ComposeLayout, Layout, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import ComposeLayout, Layout, S, TileLayout class SwizzleMode(Enum): @@ -34,14 +34,16 @@ class SwizzleMode(Enum): SWIZZLE_128B_ATOM = 3 -def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> SwizzleLayout: +def mma_atom_layout(dtype: str, swizzle_mode: SwizzleMode | int) -> ComposeLayout: """Generate the MMA-compatible shared-memory atom layout.""" bits = tvm.DataType(dtype).bits if isinstance(swizzle_mode, int): swizzle_mode = SwizzleMode(swizzle_mode) - return SwizzleLayout( - per_element=(128 // bits).bit_length() - 1, swizzle_len=swizzle_mode.value, atom_len=3 - ) + per_element = (128 // bits).bit_length() - 1 + swizzle_len = swizzle_mode.value + atom_len = 3 + period = 1 << (per_element + swizzle_len + atom_len) + return ComposeLayout(per_element, swizzle_len, atom_len, TileLayout(S[(period,)])) def mma_atom_shape(dtype: str, swizzle_mode: SwizzleMode | int, shape: list[int] | None = None): @@ -69,6 +71,21 @@ def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Lay if isinstance(swizzle_mode, int): swizzle_mode = SwizzleMode(swizzle_mode) if swizzle_mode == SwizzleMode.SWIZZLE_NONE: + # No-swizzle MMA smem is the packed-16B atom: offset e16*m + + # M*e16*(k//e16) + k%e16 (e16=128/bits); plain tile if K % e16 != 0. + bits = tvm.DataType(dtype).bits + e16 = 128 // bits + m, k = int(shape[-2]), int(shape[-1]) + if k % e16 == 0: + lead = [int(s) for s in shape[:-2]] + extents = [*lead, m, k // e16, e16] + strides = [] + stride = m * k + for e in reversed(lead): + strides.insert(0, stride) + stride *= e + strides += [e16, m * e16, 1] + return TileLayout(S[tuple(extents) : tuple(strides)]).canonicalize() return TileLayout(S[tuple(shape)]).canonicalize() atom_shape = mma_atom_shape(dtype, swizzle_mode, shape) layout = mma_atom_layout(dtype, swizzle_mode) @@ -77,12 +94,6 @@ def mma_shared_layout(dtype: str, swizzle_mode: SwizzleMode | int, shape) -> Lay return layout.tile_to(tile_to_shape, atom_shape).tile_to(shape, tile_to_shape).canonicalize() -# Backward-compatible aliases kept during the alloc_mma migration. -tma_atom_layout = mma_atom_layout -tma_atom_shape = mma_atom_shape -tma_shared_layout = mma_shared_layout - - def tma_atom_compatible(dst_shape, dst_st, dst_extent, atom_shape): """Check if the copy region in dst is compatible with the TMA atom shape.""" analyzer = Analyzer() @@ -96,18 +107,41 @@ def tma_atom_compatible(dst_shape, dst_st, dst_extent, atom_shape): def get_swizzle_mode_from_layout(layout: Layout) -> SwizzleMode | None: - """Extract swizzle mode from a shared memory layout.""" + """Extract swizzle mode from a shared memory layout. + + The hardware swizzle modes this maps to (TMA descriptor swizzles and the + tcgen05 smem matrix descriptor walk) implement the ``swizzle_inner=True`` + address permutation ``x ^ ((x & outer_mask) >> atom_len)`` — the atom-row + bits XORed into the 16B-unit bits — pinned bit-exactly on hardware by the + TMA / tcgen05.cp GPU round-trip tests. ``swizzle_inner=False`` is the + mirrored permutation ``x ^ ((x & inner_mask) << atom_len)``; for + ``swizzle_len=sw >= 1`` (``atom_len=3``) the two coincide only on the + ``2^(3-sw)`` of the ``2^(3+sw)`` chunks per atom period whose inner + ``[0, sw)`` and outer ``[3, 3+sw)`` chunk bits are all zero (exhaustive + enumeration over the full domain), so planning a flipped layout by its + ``swizzle_len`` alone would silently misplace data — reject it instead. + For ``swizzle_len == 0`` both masks are 0 and either value is the + identity, so ``swizzle_inner`` is a don't-care there. + """ if isinstance(layout, ComposeLayout): - swizzle = layout.swizzle # SwizzleLayout is named 'swizzle' in ComposeLayout - swizzle_len = swizzle.swizzle_len - elif isinstance(layout, SwizzleLayout): - swizzle_len = layout.swizzle_len + swizzle = layout # ComposeLayout carries the swizzle params directly elif isinstance(layout, TileLayout): - # TileLayout without SwizzleLayout means no swizzle (mode 0) + # TileLayout alone means no swizzle (mode 0) return SwizzleMode.SWIZZLE_NONE else: return None + swizzle_len = swizzle.swizzle_len + if int(swizzle_len) > 0 and not bool(swizzle.swizzle_inner): + raise ValueError( + "swizzle direction mismatch: the hardware swizzle modes implement " + "the swizzle_inner=True address permutation " + "(x ^ ((x & outer_mask) >> atom_len)); got a ComposeLayout with " + f"swizzle_inner=False (per_element={int(swizzle.per_element)}, " + f"swizzle_len={int(swizzle_len)}, atom_len={int(swizzle.atom_len)}), " + "whose mirrored permutation would be silently misplaced" + ) + # Map swizzle_len to SwizzleMode return { 0: SwizzleMode.SWIZZLE_NONE, diff --git a/python/tvm/backend/cuda/script.py b/python/tvm/backend/cuda/script.py index 062f27d38085..b49a956dd741 100644 --- a/python/tvm/backend/cuda/script.py +++ b/python/tvm/backend/cuda/script.py @@ -57,6 +57,7 @@ def __init__(self): self.clc_query_cancel = _op_wrapper(_cuda_op.ptx_clc_query_cancel) self.fetch_register: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_fetch_register) self.ld = _op_wrapper(_cuda_op.ptx_ld) + self.ld_global_nc = _op_wrapper(_cuda_op.ptx_ld_global_nc) self.ld_acquire = _op_wrapper(_cuda_op.ptx_ld_acquire) self.ld_relaxed = _op_wrapper(_cuda_op.ptx_ld_relaxed) self.ld_volatile = _op_wrapper(_cuda_op.ptx_ld_volatile) @@ -86,6 +87,7 @@ def __init__(self): self.rcp = _op_wrapper(_cuda_op.ptx_rcp) self.reduce3_min_f32 = _op_wrapper(_cuda_op.ptx_reduce3_min_f32) self.reduce3_max_f32 = _op_wrapper(_cuda_op.ptx_reduce3_max_f32) + self.cvt = _dtype_forward(_cuda_op.ptx_cvt) # add/sub/mul/fma DPS form: (d_addr, a, b[, c], *, rounding, ftz[, sat]) self.add_f32 = _op_wrapper(_cuda_op.ptx_add_f32) self.add_f32x2 = _op_wrapper(_cuda_op.ptx_add_f32x2) @@ -181,80 +183,32 @@ class CpAsyncBulkTensorNamespace: """The CpAsyncBulkTensor instruction submodule.""" def __init__(self): - self.g2c = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster) - self.g2c_tile_gather4 = _op_wrapper( - _cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster - ) + self.g2s_cta = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_g2s_cta) + self.g2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_g2s_cluster) self.s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global) self.s2g_reduce = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_shared_to_global_reduce) - self.g2c_prefetch = _op_wrapper( - _cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch - ) - - @staticmethod - def g2c_bar_addr( - dim, - dst_ptr, - bar_addr, - tensormap_addr, - cta_mask, - cta_group, - cache_hint, - *coords, - cache_policy=None, - ): - _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP) - cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy) - return _tir_op.call_intrin( - "", - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster", - dim, - dst_ptr, - bar_addr, - tensormap_addr, - cta_mask, - cta_group, - cache_policy, - int(has_cache_policy), - 1, - *coords, - ) - - @staticmethod - def g2c_tile_gather4_bar_addr( - dim, - dst_ptr, - bar_addr, - tensormap_addr, - cta_mask, - cta_group, - cache_hint, - *coords, - cache_policy=None, - ): - _cuda_op._choice("cta_group", cta_group, _cuda_op._TCGEN05_CTA_GROUP) - cache_policy, has_cache_policy = _cuda_op._resolve_cache_policy(cache_hint, cache_policy) - return _tir_op.call_intrin( - "", - "tirx.ptx.cp_async_bulk_tensor_tile_gather4_global_to_cluster", - dim, - dst_ptr, - bar_addr, - tensormap_addr, - cta_mask, - cta_group, - cache_policy, - int(has_cache_policy), - 1, - *coords, - ) + self.prefetch = _op_wrapper(_cuda_op.ptx_cp_async_bulk_tensor_prefetch) class CpAsyncMbarrierNamespace: """The CpAsyncMbarrier instruction submodule.""" def __init__(self): - self.arrive = _op_wrapper(_cuda_op.ptx_cp_async_mbarrier_arrive) + self.arrive = CpAsyncMbarrierArriveNamespace() + + +class CpAsyncMbarrierArriveNamespace: + """The CpAsyncMbarrier Arrive instruction submodule.""" + + @staticmethod + def noinc(*args, **kwds): + return _cuda_op.ptx_cp_async_mbarrier_arrive_noinc(*args, **kwds) + + def __call__(self, *args, **kwds): + return _op_wrapper(_cuda_op.ptx_cp_async_mbarrier_arrive)(*args, **kwds) + + # __call__ corresponds to ptx_cp_async_mbarrier_arrive + __tir_call_op_name__ = "ptx_cp_async_mbarrier_arrive" class WgmmaNamespace: @@ -282,6 +236,7 @@ class MbarrierNamespace: def __init__(self): self.init = _op_wrapper(_cuda_op.ptx_mbarrier_init) + self.complete_tx = _op_wrapper(_cuda_op.ptx_mbarrier_complete_tx) self.try_wait = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait) self.try_wait_once = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_once) self.try_wait_acquire_cluster = _op_wrapper(_cuda_op.ptx_mbarrier_try_wait_acquire_cluster) @@ -293,7 +248,7 @@ class MbarrierArriveNamespace: def __init__(self): self.expect_tx = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_expect_tx) - self.cluster_count = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_cluster_count) + self.no_complete = _op_wrapper(_cuda_op.ptx_mbarrier_arrive_no_complete) def __call__(self, *args, **kwds): return _op_wrapper(_cuda_op.ptx_mbarrier_arrive)(*args, **kwds) @@ -379,6 +334,7 @@ class BarrierNamespace: """The Barrier instruction submodule.""" def __init__(self): + self.sync = _op_wrapper(_cuda_op.ptx_barrier_sync) self.cluster = BarrierClusterNamespace() @@ -411,10 +367,24 @@ def __init__(self): self.launch_dependents = _op_wrapper(_cuda_op.ptx_griddepcontrol_launch_dependents) +class IketNamespace: + """Frontend-only NVIDIA IKET annotations.""" + + def __init__(self): + self.mark = _op_wrapper(_cuda_op.cuda_iket_mark) + self.range_start = _op_wrapper(_cuda_op.cuda_iket_range_start) + self.range_end = _op_wrapper(_cuda_op.cuda_iket_range_end) + self.range_push = _op_wrapper(_cuda_op.cuda_iket_range_push) + self.range_pop = _op_wrapper(_cuda_op.cuda_iket_range_pop) + self.sentinel_token = _op_wrapper(_cuda_op.cuda_iket_sentinel_token) + self.official_event = _op_wrapper(_cuda_op.cuda_iket_official_event) + + class CUDANamespace: """The CUDA intrinsics submodule.""" def __init__(self): + self.iket = IketNamespace() self.atomic_add = _op_wrapper(_cuda_op.cuda_atomic_add) self.thread_fence = _op_wrapper(_cuda_op.cuda_thread_fence) self.warpgroup_sync = _op_wrapper(_cuda_op.cuda_warpgroup_sync) @@ -445,10 +415,11 @@ def __init__(self): self.func_call = _op_wrapper(_cuda_op.cuda_func_call) self.printf = _op_wrapper(_cuda_op.cuda_printf) self.ldg = _op_wrapper(_cuda_op.cuda_ldg) + self.fdividef = _op_wrapper(_cuda_op.cuda_fdividef) self.get_tmem_addr = _op_wrapper(_cuda_op.cuda_get_tmem_addr) self.cvta_generic_to_shared = _op_wrapper(_cuda_op.cuda_cvta_generic_to_shared) self.smem_addr_from_uint64 = _op_wrapper(_cuda_op.cuda_smem_addr_from_uint64) - self.sm100_tma_2sm_mbarrier_addr = _op_wrapper(_cuda_op.cuda_sm100_tma_2sm_mbarrier_addr) + self.sm100_2sm_leader_smem_addr = _op_wrapper(_cuda_op.cuda_sm100_2sm_leader_smem_addr) self.uint_as_float = _op_wrapper(_cuda_op.cuda_uint_as_float) self.float_as_uint = _op_wrapper(_cuda_op.cuda_float_as_uint) self.ballot_sync = _op_wrapper(_cuda_op.cuda_ballot_sync) diff --git a/python/tvm/backend/cuda/transforms/__init__.py b/python/tvm/backend/cuda/transforms/__init__.py new file mode 100644 index 000000000000..b43403d337e3 --- /dev/null +++ b/python/tvm/backend/cuda/transforms/__init__.py @@ -0,0 +1,41 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""CUDA-specific TIRx transformations.""" + +# pylint: disable=invalid-name + +from tvm_ffi import get_global_func + + +def LowerIket(): + """Lower frontend-only NVIDIA IKET annotations. + + This pass must run after ``tvm.tirx.transform.SplitHostDevice`` and before + ``tvm.tirx.transform.MakePackedAPI``. It strips annotations for regular + compilation and emits NVIDIA IKET metadata and NativeDump placeholders + when the IRModule is explicitly IKET-enabled. Trace collection and + postprocessing are owned by the external ``run-iket`` process. + + Returns + ------- + fpass : tvm.transform.Pass + The result pass. + """ + return get_global_func("tirx.backend.cuda.transforms.LowerIket")() + + +__all__ = ["LowerIket"] diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/default.py b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py index 6b1baee43163..998532d93938 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/binary/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/binary/default.py @@ -21,7 +21,7 @@ from tvm.tirx import FloatImm, PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import MapOpType -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import init_analyzer, nki_dim from ..instruction_generator import InstructionGenerator diff --git a/python/tvm/backend/trn/operator/tile_primitive/copy/default.py b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py index 112b78ab6488..1dd18ed77627 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/copy/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/copy/default.py @@ -26,7 +26,7 @@ predicate, register_dispatch, ) -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import init_analyzer, nki_dim from ..dim_utils import get_ewise_dim_map diff --git a/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py index d0ed02ba2bee..1a52792fda46 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py @@ -31,7 +31,7 @@ predicate, register_dispatch, ) -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import init_analyzer from ..dim_utils import normalize_and_group diff --git a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py index 91f6f3b619c5..1ce89ef4f0a4 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py +++ b/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py @@ -22,7 +22,6 @@ from tvm.backend.trn.operator.tile_primitive.instruction_generator import InstructionGenerator from tvm.script import tirx as T from tvm.tirx import Buffer, FloatImm, Stmt -from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext from tvm.tirx.operator.tile_primitive.ops import ( BinaryReduce, Copy, @@ -32,7 +31,7 @@ UnaryReduce, ) from tvm.tirx.operator.tile_primitive.registry import f_op_dispatcher -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import DispatchContext, TilePrimitiveCall def _scalar_dtype(scalar) -> str: diff --git a/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py index 7fa26b11a1f6..e600ebda6de8 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py +++ b/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py @@ -22,7 +22,7 @@ from tvm.tirx import PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import ReduceOpType -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import init_analyzer, nki_dim from ..dim_utils import get_reduction_dim_map diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/default.py b/python/tvm/backend/trn/operator/tile_primitive/unary/default.py index 69facea7f97b..b6516c1d0517 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/unary/default.py +++ b/python/tvm/backend/trn/operator/tile_primitive/unary/default.py @@ -20,7 +20,7 @@ from tvm.tirx import FloatImm, PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import MapOpType -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..common import init_analyzer from ..instruction_generator import InstructionGenerator diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py b/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py index 26fb5670270a..c75f18544f47 100644 --- a/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py +++ b/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py @@ -20,7 +20,7 @@ from tvm.tirx import BufferRegion, PrimFunc from tvm.tirx.operator.tile_primitive import DispatchContext, fail from tvm.tirx.operator.tile_primitive.common import MapOpType -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall from ..binary import try_find_inst_nary from ..common import init_analyzer diff --git a/python/tvm/backend/trn/transform/private_buffer_alloc.py b/python/tvm/backend/trn/transform/private_buffer_alloc.py index 77908210bafd..129483b3dde4 100644 --- a/python/tvm/backend/trn/transform/private_buffer_alloc.py +++ b/python/tvm/backend/trn/transform/private_buffer_alloc.py @@ -19,16 +19,15 @@ from tvm.ir import Range from tvm.target import Target from tvm.tirx.buffer import Buffer -from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext from tvm.tirx.stmt import ( AllocBuffer, AttrStmt, For, SeqStmt, Stmt, - TilePrimitiveCall, ) from tvm.tirx.stmt_functor import StmtMutator, StmtVisitor +from tvm.tirx.tile_primitive import DispatchContext, TilePrimitiveCall from tvm.tirx.transform.common import seek_kernel_replace_point from tvm.tirx.transform.function_pass import prim_func_pass diff --git a/python/tvm/script/parser/core/entry.py b/python/tvm/script/parser/core/entry.py index 7db6335ad67e..418639db299f 100644 --- a/python/tvm/script/parser/core/entry.py +++ b/python/tvm/script/parser/core/entry.py @@ -78,6 +78,7 @@ def parse( extra_vars: dict[str, Any] | None = None, check_well_formed: bool = True, s_tir: bool = False, + absent_params: dict[str, None] | None = None, ) -> Any: """Register a method for a operand type, AST operator node and operand index. @@ -92,6 +93,10 @@ def parse( check_well_formed : bool Whether to check well-formedness after parsing. + absent_params : Optional[Dict[str, None]] + Function parameters removed by a compile-time specialization. The + dialect-specific function parser decides how to bind these names. + Returns ------- func : Any @@ -111,7 +116,7 @@ def parse( all_pyfuncs[name] = func source = Source(program) - parser = Parser(source, ann) + parser = Parser(source, ann, absent_params=absent_params) with IRBuilder() as builder: try: parser.parse(extra_vars=extra_vars) diff --git a/python/tvm/script/parser/core/evaluator.py b/python/tvm/script/parser/core/evaluator.py index ccb582bef0c6..71a1b3bc8fff 100644 --- a/python/tvm/script/parser/core/evaluator.py +++ b/python/tvm/script/parser/core/evaluator.py @@ -185,24 +185,36 @@ def _visit(self, node: doc.AST) -> Any: def _visit_node(self, node: doc.AST) -> Any: """Evaluate one AST node while its source span is active.""" + # Boolean operators and conditional expressions require AST-level + # control over operand evaluation. Walking all fields first would + # eagerly evaluate dead compile-time branches. + if isinstance(node, doc.BoolOp): + try: + value = self._eval_bool_op(node) + except Exception as err: # pylint: disable=broad-except + self.parser.report_error(node, err) + return self._add_intermediate_result(value) + if isinstance(node, doc.IfExp): + try: + value = self._eval_if_exp(node) + except Exception as err: # pylint: disable=broad-except + self.parser.report_error(node, err) + return self._add_intermediate_result(value) + args = [] if ( isinstance(node, doc.Call) and hasattr(node.func, "attr") and node.func.attr not in ["reads", "writes", "match_buffer"] - ) or isinstance(node, doc.BinOp | doc.UnaryOp | doc.Compare | doc.BoolOp | doc.IfExp): + ) or isinstance(node, doc.BinOp | doc.UnaryOp | doc.Compare): if isinstance(node, doc.BinOp): args = [node.left, node.right] elif isinstance(node, doc.UnaryOp): args = [node.operand] elif isinstance(node, doc.Compare): args = [node.left, *node.comparators] - elif isinstance(node, doc.IfExp): - args = [node.test, node.body, node.orelse] elif isinstance(node, doc.Call): args = node.args - elif isinstance(node, doc.BoolOp): - args = node.values for arg in args: if isinstance(arg, doc.Subscript) and isinstance(arg.slice, doc.Slice | doc.Tuple): if isinstance(arg.slice, doc.Slice): @@ -260,16 +272,12 @@ def _visit_node(self, node: doc.AST) -> Any: else: fields[field] = attr try: - if isinstance(node, doc.BoolOp): - value = self._eval_bool_op(fields) - elif isinstance(node, doc.Compare): + if isinstance(node, doc.Compare): value = self._eval_compare(fields) elif isinstance(node, doc.UnaryOp): value = self._eval_unary_op(fields) elif isinstance(node, doc.BinOp): value = self._eval_bin_op(fields) - elif isinstance(node, doc.IfExp): - value = self._eval_if_exp(fields) elif isinstance(node, doc.Slice): value = self._eval_slice(fields) else: @@ -297,26 +305,31 @@ def _eval_lambda(self, node: doc.Lambda) -> Any: self.parser.report_error(node, err) return self._add_intermediate_result(value) - def _eval_bool_op(self, fields: dict[str, Any]) -> Any: + def _eval_bool_op(self, node: doc.BoolOp) -> Any: """The doc AST boolean operator node evaluating method. Parameters ---------- - fields : Dict[str, Any] - The dictionary of boolean operation information, - e.g., operator types, operand values. + node : doc.BoolOp + The boolean operation whose operands are evaluated in order. Returns ------- res : Any The evaluation result. """ - op = fields["op"] + op = node.op if not isinstance(op, doc.And | doc.Or): raise TypeError(f"Unexpected operator: {op}") - value = self._eval_expr(fields["values"][0]) - for rhs in fields["values"][1:]: - value = _eval_op(op, values=[value, self._eval_expr(rhs)]) + value = self._eval_expr(self._visit(node.values[0])) + for rhs_node in node.values[1:]: + if isinstance(value, bool): + if isinstance(op, doc.And) and not value: + return value + if isinstance(op, doc.Or) and value: + return value + rhs = self._eval_expr(self._visit(rhs_node)) + value = _eval_op(op, values=[value, rhs]) return value def _eval_compare(self, fields: dict[str, Any]) -> Any: @@ -386,26 +399,26 @@ def _eval_bin_op(self, fields: dict[str, Any]) -> Any: ], ) - def _eval_if_exp(self, fields: dict[str, Any]) -> Any: + def _eval_if_exp(self, node: doc.IfExp) -> Any: """The doc AST if-else expression node evaluating method. Parameters ---------- - fields : Dict[str, Any] - The dictionary of if-else expression information, - e.g., test, body, orelse. + node : doc.IfExp + The conditional expression to evaluate. Returns ------- res : Any The evaluation result. """ - test = self._eval_expr(fields["test"]) - body = self._eval_expr(fields["body"]) - orelse = self._eval_expr(fields["orelse"]) + test = self._eval_expr(self._visit(node.test)) if isinstance(test, bool): - return body if test else orelse + selected = node.body if test else node.orelse + return self._eval_expr(self._visit(selected)) elif tvm.ir.is_prim_expr(test) and test.ty.matches_code(tvm.DataTypeCode.BOOL): + body = self._eval_expr(self._visit(node.body)) + orelse = self._eval_expr(self._visit(node.orelse)) return tvm.tirx.op.if_then_else(test, body, orelse) else: raise TypeError(f"Expected Python bool or TIR bool, but got {type(test)}") diff --git a/python/tvm/script/parser/core/parser.py b/python/tvm/script/parser/core/parser.py index ff09ee117738..36349b26a4c9 100644 --- a/python/tvm/script/parser/core/parser.py +++ b/python/tvm/script/parser/core/parser.py @@ -373,11 +373,15 @@ class Parser(doc.NodeVisitor): var_table : VarTable The variable table for parsing. + + absent_params : Dict[str, None] + Function parameters removed by a compile-time specialization. """ diag: Diagnostics dispatch_tokens: list[str] function_annotations: dict[str, dict[str, Any]] | None + absent_params: dict[str, None] var_table: VarTable inside_function: bool # whether we are within a function current_class: str | None = None # current class being parsed @@ -387,10 +391,12 @@ def __init__( self, source: Source, function_annotations: dict[str, dict[str, Any]], + absent_params: dict[str, None] | None = None, ) -> None: self.diag = Diagnostics(source) self.dispatch_tokens = ["default"] self.function_annotations = function_annotations + self.absent_params = absent_params or {} self.var_table = VarTable() self.inside_function = False diff --git a/python/tvm/support/nvcc.py b/python/tvm/support/nvcc.py index 84976dfb15b2..f9b0ebcc895c 100644 --- a/python/tvm/support/nvcc.py +++ b/python/tvm/support/nvcc.py @@ -149,10 +149,13 @@ def _compile_cuda_nvcc( # "-gencode", "arch=compute_52,code=sm_52", # "-gencode", "arch=compute_70,code=sm_70" # ] - compute_version = "".join( - get_target_compute_version(Target.current(allow_none=True)).split(".") - ) - arch = ["-gencode", f"arch=compute_{compute_version},code=sm_{compute_version}"] + target = Target.current(allow_none=True) + target_arch = getattr(target, "arch", None) if target is not None else None + if isinstance(target_arch, str) and target_arch.startswith("sm_"): + suffix = target_arch[3:] + else: + suffix = "".join(get_target_compute_version(target).split(".")) + arch = ["-gencode", f"arch=compute_{suffix},code=sm_{suffix}"] temp = utils.tempdir() file_name = "tvm_kernels" @@ -213,7 +216,7 @@ def _compile_cuda_nvcc( "-U__CUDA_NO_BFLOAT162_CONVERSIONS__", "--expt-relaxed-constexpr", "--expt-extended-lambda", - "--use_fast_math", + *([] if os.environ.get("TVM_CUDA_NVCC_NO_FAST_MATH") else ["--use_fast_math"]), f"--ptxas-options={','.join(_ptxas_option_flags())}", ] @@ -255,6 +258,12 @@ def _compile_cuda_nvcc( # Second stage for NVSHMEM if use_nvshmem: + target = Target.current(allow_none=True) + target_arch = getattr(target, "arch", None) if target is not None else None + if isinstance(target_arch, str) and target_arch.startswith("sm_"): + compute_version = target_arch[3:] + else: + compute_version = "".join(get_target_compute_version(target).split(".")) cmd = ["nvlink"] cmd += [f"-arch=sm_{compute_version}"] cmd += ["-L", nvshmem_lib_path] @@ -550,6 +559,13 @@ def _compile_cuda_nvrtc( for flag in _ptxas_option_flags(): compile_opts.append(f"--ptxas-options={flag}".encode()) + # Extra NVRTC frontend flags (shell-tokenized), appended after all built-in + # defaults so they can override them (e.g. TVM_CUDA_NVRTC_EXTRA_OPTS="--ftz=false" + # to undo the -ftz=true implied by --use_fast_math). + nvrtc_extra = os.environ.get("TVM_CUDA_NVRTC_EXTRA_OPTS", "").strip() + if nvrtc_extra: + compile_opts.extend(t.encode() for t in shlex.split(nvrtc_extra)) + # Add user-provided options, filtering out nvcc-specific flags that nvrtc doesn't support if options: nvcc_only_prefixes = ( diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py index e1c84297a311..a71105044bd6 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -53,7 +53,8 @@ from .stmt import SeqStmt from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list from .stmt import BufferRegion, MatchBufferRegion, SBlock, SBlockRealize -from .stmt import TilePrimitiveCall, ScopeIdDefStmt +from .stmt import ScopeIdDefStmt +from .tile_primitive import DispatchContext, LambdaExpr, TilePrimitiveCall from .function import PrimFunc, TensorIntrin, IndexMap @@ -94,8 +95,7 @@ # TIRX-specific imports (must come before subpackage imports to avoid circular imports) from .exec_scope import ExecScope, ScopeIdDef -from .layout import TileLayout, Layout, SwizzleLayout, ComposeLayout -from .predicate import Predicate +from .layout import TileLayout, Layout, ComposeLayout from .expr_functor import ExprFunctor from . import transform diff --git a/python/tvm/tirx/_buffer_view.py b/python/tvm/tirx/_buffer_view.py new file mode 100644 index 000000000000..4448a8a72d71 --- /dev/null +++ b/python/tvm/tirx/_buffer_view.py @@ -0,0 +1,611 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Private implementation of derived :class:`Buffer` views.""" + +from __future__ import annotations + +import functools +import re +from numbers import Integral +from typing import TYPE_CHECKING + +import tvm + +if TYPE_CHECKING: + from .buffer import Buffer + + +def _redecl(buf: Buffer, shape, layout, *, dtype=None, elem_offset=None, addr_offset=None): + """Re-declare a derived view over ``buf`` storage. + + Routes the buffer identity by storage kind: a ``tmem`` buffer aliases by + column address, while every other scope carries + ``data``/``strides``/``elem_offset`` through. + """ + if buf.scope() == "tmem" and buf.allocated_addr is not None and len(buf.allocated_addr) > 0: + addr = buf.allocated_addr[0] + if addr_offset is not None: + addr = addr + addr_offset + return tvm.tirx.script.builder.decl_buffer( + shape, + buf.dtype if dtype is None else dtype, + None, + None, + None, + None, + buf.scope(), + buf.data_alignment, + 0, + layout, + allocated_addr=addr, + ) + return tvm.tirx.script.builder.decl_buffer( + shape, + buf.dtype if dtype is None else dtype, + buf.data, + buf.strides, + buf.elem_offset if elem_offset is None else elem_offset, + None, + buf.scope(), + buf.data_alignment, + buf.offset_factor, + layout, + ) + + +def view(buf: Buffer, *args, **kwargs) -> Buffer: + """Implement :meth:`Buffer.view`.""" + + def _infer_shape(shape): + shape = list(shape) + if -1 in shape and shape.count(-1) == 1: + size = functools.reduce(lambda x, y: x * y, buf.shape) + n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1) + shape[shape.index(-1)] = size // n_size + else: + # A PrimExpr comparison returns an EQ node, not a Python bool. + if all(isinstance(s, int) for s in shape) and all( + isinstance(s, int) for s in buf.shape + ): + assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce( + lambda x, y: x * y, buf.shape + ), ( + "The shape of the buffer " + + str(buf.shape) + + " and the new shape " + + str(shape) + + " are not compatible" + ) + return shape + + if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs: + cast_dtype = tvm.DataType(args[0]) + cur_dtype = tvm.DataType(buf.dtype) + if cast_dtype.bits > cur_dtype.bits: + assert cast_dtype.bits % cur_dtype.bits == 0 + ratio = cast_dtype.bits // cur_dtype.bits + layout = buf.layout.pack(ratio) + shape = [s for s in buf.shape[:-1]] + [buf.shape[-1] // ratio] + new_elem_offset = buf.elem_offset // ratio + else: + assert cur_dtype.bits % cast_dtype.bits == 0 + ratio = cur_dtype.bits // cast_dtype.bits + layout = buf.layout.unpack(ratio) + shape = [s for s in buf.shape[:-1]] + [buf.shape[-1] * ratio] + new_elem_offset = buf.elem_offset * ratio + return _redecl(buf, shape, layout, dtype=cast_dtype, elem_offset=new_elem_offset) + + shape = args + assert all( + isinstance(arg, int) or (tvm.ir.is_prim_expr(arg) and arg.ty.dtype in ["int32", "int64"]) + for arg in shape + ), "shape must be a list of integers or Exprs with dtype int32 or int64" + layout = kwargs.get("layout", None) + assert set(kwargs.keys()).issubset({"layout"}), ( + f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}" + ) + if layout is None: + shape = _infer_shape(shape) + return _redecl(buf, shape, buf.layout if layout is None else layout) + + +def local(buf: Buffer, *shape, layout=None) -> Buffer: + """Implement :meth:`Buffer.local`.""" + if not shape: + local_layout = buf.layout.storage() + total = functools.reduce(lambda x, y: x * y, [it.extent for it in local_layout.shard], 1) + shape = (total,) + return _redecl(buf, shape, buf.layout.storage() if layout is None else layout) + + +def permute(buf: Buffer, *dims) -> Buffer: + """Implement :meth:`Buffer.permute`.""" + new_shape = [buf.shape[d] for d in dims] + layout, swizzle = _surgery_parts(buf) + grouped, seps = layout.group(list(buf.shape)) + new_layout = grouped.permute_by_groups(seps, list(dims)) + if swizzle is not None: + new_layout = _rewrap_swizzle(new_layout, swizzle) + return _redecl(buf, new_shape, new_layout) + + +def rearrange(buf: Buffer, pattern: str, /, **sizes) -> Buffer: + """Implement :meth:`Buffer.rearrange`.""" + + def _groups(side): + out = [] + for group, bare in re.findall(r"\(([^)]*)\)|(\S+)", side.strip()): + out.append(group.split() if group else [bare]) + return out + + if "->" not in pattern: + raise ValueError(f"rearrange: pattern must contain '->', got {pattern!r}") + lhs_s, rhs_s = pattern.split("->") + lhs_groups, rhs_groups = _groups(lhs_s), _groups(rhs_s) + if len(lhs_groups) != len(buf.shape): + raise ValueError( + f"rearrange: lhs has {len(lhs_groups)} groups but buffer has " + f"{len(buf.shape)} dims (pattern {pattern!r}, shape {list(buf.shape)})" + ) + axis_size: dict = {} + for group, dim in zip(lhs_groups, buf.shape): + dim_c = int(dim) + known = {axis: int(sizes[axis]) for axis in group if axis in sizes} + unknown = [axis for axis in group if axis not in sizes] + product = 1 + for value in known.values(): + product *= value + if len(unknown) == 0: + if product != dim_c: + raise ValueError(f"rearrange: group {group} product {product} != dim {dim_c}") + elif len(unknown) == 1: + if dim_c % product != 0: + raise ValueError( + f"rearrange: dim {dim_c} not divisible by known product {product} " + f"in group {group}" + ) + known[unknown[0]] = dim_c // product + else: + raise ValueError( + f"rearrange: group {group} has >1 unknown axis; give sizes= for all but one" + ) + axis_size.update(known) + flat_lhs = [axis for group in lhs_groups for axis in group] + flat_rhs = [axis for group in rhs_groups for axis in group] + if sorted(flat_lhs) != sorted(flat_rhs): + raise ValueError( + f"rearrange: lhs axes {flat_lhs} != rhs axes {flat_rhs} (pattern {pattern!r})" + ) + out = buf.view(*[axis_size[axis] for axis in flat_lhs]) + out = out.permute(*[flat_lhs.index(axis) for axis in flat_rhs]) + merged_shape = [] + for group in rhs_groups: + size = 1 + for axis in group: + size *= axis_size[axis] + merged_shape.append(size) + return out.view(*merged_shape) + + +def sub(buf: Buffer) -> SubIndexer: + """Return the indexer used by :attr:`Buffer.sub`.""" + return SubIndexer(buf) + + +def tile(buf: Buffer, *specs) -> TileIndexer: + """Implement :meth:`Buffer.tile`.""" + if specs and isinstance(specs[0], int | Integral): + if len(specs) != 2: + raise ValueError("tile(dim, factors) takes a dim and a factors tuple") + specs = (specs,) + normalized = [] + seen = set() + for spec in specs: + if not isinstance(spec, tuple | list) or len(spec) != 2: + raise ValueError(f"tile: each spec must be (dim, factors), got {spec!r}") + dim = _normalized_dim(buf, spec[0], "tile") + if dim in seen: + raise ValueError(f"tile: dim {dim} tiled more than once") + seen.add(dim) + factors = spec[1] + if not isinstance(factors, tuple | list) or len(factors) < 1: + raise ValueError(f"tile: factors for dim {dim} must be a non-empty tuple") + normalized.append((dim, tuple(factors))) + return TileIndexer(buf, normalized) + + +def chunk(buf: Buffer, spec) -> ChunkIndexer: + """Implement :meth:`Buffer.chunk`.""" + if not isinstance(spec, tuple | list): + raise ValueError(f"chunk: spec must be a per-dim tuple, got {spec!r}") + if len(spec) != len(buf.shape): + raise ValueError(f"chunk: spec length {len(spec)} != buffer rank {len(buf.shape)}") + for dim, count in enumerate(spec): + if count is not None and (not isinstance(count, Integral) or int(count) < 1): + raise ValueError(f"chunk: spec[{dim}] must be None or a positive int, got {count!r}") + return ChunkIndexer(buf, tuple(spec)) + + +def _normalized_dim(buf: Buffer, dim, name): + ndim = len(buf.shape) + if dim < 0: + dim += ndim + if not 0 <= dim < ndim: + raise ValueError(f"{name}: dim {dim} out of range for buffer of rank {ndim}") + return dim + + +def _concrete_int(value): + """Return ``value`` as a Python int when it is statically known.""" + if isinstance(value, bool): + return None + if isinstance(value, Integral): + return int(value) + if isinstance(value, tvm.tirx.IntImm): + return int(value) + return None + + +def _surgery_parts(buf: Buffer): + """Split a layout into its inner tile layout and optional swizzle params.""" + layout = buf.layout + swizzle = None + if isinstance(layout, tvm.tirx.layout.ComposeLayout): + swizzle = ( + layout.per_element, + layout.swizzle_len, + layout.atom_len, + layout.swizzle_inner, + ) + layout = layout.tile_layout + return layout, swizzle + + +def _rewrap_swizzle(layout, swizzle): + if swizzle is not None: + per_element, swizzle_len, atom_len, swizzle_inner = swizzle + return tvm.tirx.layout.ComposeLayout( + per_element, swizzle_len, atom_len, layout, swizzle_inner + ) + return layout + + +def _dim_group_offset(group, index): + """Return the offset of ``index`` along a grouped logical dimension.""" + if len(group) == 1: + return index * group[0].stride + extents = [] + for iterator in group[1:]: + extent = iterator.extent + if not isinstance(extent, tvm.tirx.IntImm): + raise ValueError( + "select/narrow: dim with multiple layout iters requires concrete iter extents" + ) + extents.append(int(extent)) + offset = None + remaining = index + for pos, iterator in enumerate(group): + inner = functools.reduce(lambda a, b: a * b, extents[pos:], 1) + if inner == 1: + coord = remaining + elif isinstance(remaining, Integral): + coord = remaining // inner + remaining = remaining % inner + else: + coord = tvm.tirx.floordiv(remaining, inner) + remaining = tvm.tirx.floormod(remaining, inner) + term = coord * iterator.stride + offset = term if offset is None else offset + term + return offset + + +def _swizzle_offset_commutes(swizzle, extra_offset): + """Return whether an offset may be folded outside a swizzle permutation.""" + if swizzle is None or extra_offset is None: + return True + per_element, swizzle_len, atom_len, _ = swizzle + if swizzle_len == 0: + return True + period = 1 << (int(per_element) + int(atom_len) + swizzle_len) + offset_c = _concrete_int(extra_offset) + if offset_c is not None: + return offset_c % period == 0 + from ..arith import Analyzer # pylint: disable=import-outside-toplevel + + return Analyzer().can_prove_equal(tvm.tirx.floormod(extra_offset, period), 0) + + +def _rebuild_view(buf: Buffer, new_shape, new_shard, grouped, swizzle, extra_offset): + """Rebuild a derived view while preserving its physical addresses.""" + offset_map = dict(grouped.offset.items()) + is_tmem = ( + buf.scope() == "tmem" and buf.allocated_addr is not None and len(buf.allocated_addr) > 0 + ) + if ( + not is_tmem + and extra_offset is not None + and not _swizzle_offset_commutes(swizzle, extra_offset) + ): + m_axis = tvm.tirx.layout.Axis.get("m") + previous = offset_map.get(m_axis) + offset_map[m_axis] = extra_offset if previous is None else previous + extra_offset + extra_offset = None + new_layout = tvm.tirx.layout.TileLayout.from_iters(new_shard, list(grouped.replica), offset_map) + new_layout = _rewrap_swizzle(new_layout, swizzle) + if is_tmem: + addr_offset = _tmem_element_offset_to_column_offset(buf, extra_offset) + return _redecl(buf, new_shape, new_layout, addr_offset=addr_offset) + elem_offset = buf.elem_offset + if extra_offset is not None: + elem_offset = elem_offset + extra_offset + return _redecl(buf, new_shape, new_layout, elem_offset=elem_offset) + + +def _tmem_element_offset_to_column_offset(buf: Buffer, element_offset): + """Convert a TCol element offset to a physical 32-bit column offset.""" + if element_offset is None: + return None + + bits = tvm.DataType(buf.dtype).bits + bit_offset = element_offset * bits + bit_offset_c = _concrete_int(bit_offset) + if bit_offset_c is not None: + if bit_offset_c % 32 != 0: + raise ValueError( + "sub: tmem TCol offset must be aligned to a physical " + f"32-bit column; got {element_offset} elements of {buf.dtype}" + ) + return bit_offset_c // 32 + + from ..arith import Analyzer # pylint: disable=import-outside-toplevel + + analyzer = Analyzer() + if not analyzer.can_prove_equal(tvm.tirx.floormod(bit_offset, 32), 0): + raise ValueError( + "sub: tmem TCol offset must be provably aligned to a physical " + f"32-bit column; got {element_offset} elements of {buf.dtype}" + ) + return analyzer.simplify(tvm.tirx.floordiv(bit_offset, 32)) + + +def _tmem_offset_axis_check(buf: Buffer, group_iters, start_c, what): + """Check that a tmem view offset can move into ``allocated_addr``.""" + if buf.allocated_addr is None or len(buf.allocated_addr) == 0: + return + if start_c == 0: + return + for iterator in group_iters: + axis_name = getattr(iterator.axis, "name", None) if iterator.axis is not None else None + if axis_name != "TCol": + raise ValueError( + f"sub: tmem {what} with nonzero offset requires TCol-stride " + f"layout iters; dim iter has axis {axis_name!r} (a lane-axis " + f"offset cannot fold into the column allocated_addr)" + ) + + +def _view_drop(buf: Buffer, dim, index): + """Drop one logical dimension while preserving the selected storage offset.""" + dim = _normalized_dim(buf, dim, "index") + index_c = _concrete_int(index) + extent_c = _concrete_int(buf.shape[dim]) + if index_c is not None and (index_c < 0 or (extent_c is not None and index_c >= extent_c)): + raise ValueError( + f"index {index_c} out of range for dim {dim} " + f"of extent {extent_c if extent_c is not None else buf.shape[dim]}" + ) + layout, swizzle = _surgery_parts(buf) + grouped, seps = layout.group(list(buf.shape)) + iters = list(grouped.shard) + lo, hi = seps[dim], seps[dim + 1] + _tmem_offset_axis_check(buf, iters[lo:hi], index_c, "index") + offset = _dim_group_offset(iters[lo:hi], index) + new_shape = list(buf.shape[:dim]) + list(buf.shape[dim + 1 :]) + new_shard = iters[:lo] + iters[hi:] + return _rebuild_view(buf, new_shape, new_shard, grouped, swizzle, offset) + + +def _view_narrow(buf: Buffer, dim, start, length): + """Narrow one logical dimension while preserving its storage offset.""" + dim = _normalized_dim(buf, dim, "index") + start_c = _concrete_int(start) + length_c = _concrete_int(length) + extent_c = _concrete_int(buf.shape[dim]) + if start_c is not None and start_c < 0: + raise ValueError(f"sub: start {start_c} must be non-negative") + if length_c is not None and length_c < 1: + raise ValueError(f"sub: length {length_c} must be positive") + if ( + start_c is not None + and length_c is not None + and extent_c is not None + and start_c + length_c > extent_c + ): + raise ValueError( + f"sub: range [{start_c}, {start_c + length_c}) exceeds dim {dim} extent {extent_c}" + ) + layout, swizzle = _surgery_parts(buf) + grouped, seps = layout.group(list(buf.shape)) + iters = list(grouped.shard) + lo, hi = seps[dim], seps[dim + 1] + group = iters[lo:hi] + iter_type = tvm.tirx.layout.Iter + if len(group) == 1: + iterator = group[0] + _tmem_offset_axis_check(buf, [iterator], start_c, "narrow") + offset = start * iterator.stride + new_group = [iter_type(length, iterator.stride, iterator.axis)] + else: + inner = 1 + for iterator in group[1:]: + extent = iterator.extent + if not isinstance(extent, tvm.tirx.IntImm): + raise ValueError( + "sub: dim with multiple layout iters requires concrete iter extents" + ) + inner *= int(extent) + if not (isinstance(start, Integral) and isinstance(length, Integral)): + raise ValueError( + f"sub: dim {dim} is made of {len(group)} layout iters; " + f"start/length must be concrete multiples of {inner}" + ) + if start % inner != 0 or length % inner != 0: + raise ValueError( + f"sub: dim {dim} start/length must be multiples of the " + f"inner iter block {inner}; got start={start} length={length}" + ) + outer = group[0] + _tmem_offset_axis_check(buf, [outer], start // inner, "narrow") + offset = (start // inner) * outer.stride + new_group = [iter_type(length // inner, outer.stride, outer.axis), *group[1:]] + new_shape = list(buf.shape) + new_shape[dim] = length + new_shard = iters[:lo] + new_group + iters[hi:] + return _rebuild_view(buf, new_shape, new_shard, grouped, swizzle, offset) + + +class SubIndexer: + """Indexer returned by :attr:`Buffer.sub`.""" + + def __init__(self, buffer: Buffer): + self._buffer = buffer + + def __getitem__(self, indices) -> Buffer: + if not isinstance(indices, tuple): + indices = (indices,) + buf = self._buffer + if len(indices) > len(buf.shape): + raise ValueError(f"sub: {len(indices)} indices for buffer of rank {len(buf.shape)}") + dim = 0 + for item in indices: + if isinstance(item, slice): + step = item.step + if step is None or (isinstance(step, Integral) and step == 1): + if item.start is None and item.stop is None: + dim += 1 + continue + start = 0 if item.start is None else item.start + stop = buf.shape[dim] if item.stop is None else item.stop + buf = _view_narrow(buf, dim, start, stop - start) + dim += 1 + else: + if not isinstance(step, Integral) or step <= 0: + raise ValueError(f"sub: step must be a positive int, got {step!r}") + if item.stop is not None: + raise ValueError("sub: a stop bound with a step is not supported") + extent = buf.shape[dim] + if not isinstance(extent, tvm.tirx.IntImm): + raise ValueError("sub: a stepped slice requires a concrete dim extent") + extent = int(extent) + if extent % step != 0: + raise ValueError( + f"sub: dim extent {extent} is not divisible by step {step}" + ) + start = 0 if item.start is None else item.start + start_c = _concrete_int(start) + if start_c is not None and not 0 <= start_c < step: + raise ValueError( + f"sub: stepped-slice start {start_c} must be in [0, {step})" + ) + split = buf.view(*buf.shape[:dim], extent // step, step, *buf.shape[dim + 1 :]) + buf = _view_drop(split, dim + 1, start) + dim += 1 + else: + buf = _view_drop(buf, dim, item) + return buf + + +class ChunkIndexer: + """Indexer returned by :meth:`Buffer.chunk`.""" + + def __init__(self, buffer: Buffer, spec): + self._buffer = buffer + self._spec = spec + + def __getitem__(self, picks): + if not isinstance(picks, tuple): + picks = (picks,) + buf = self._buffer + if len(picks) > len(self._spec): + raise ValueError(f"chunk: {len(picks)} indices for a rank-{len(self._spec)} spec") + translated = [] + for dim, pick in enumerate(picks): + count = self._spec[dim] + if count is None: + translated.append(pick) + elif isinstance(pick, slice): + raise ValueError( + f"chunk: chunked dim {dim} takes a chunk index, not a slice; got {pick!r}" + ) + else: + size = buf.shape[dim] // int(count) + translated.append(slice(pick * size, (pick + 1) * size)) + return buf[tuple(translated)] + + +class TileIndexer: + """Indexer returned by :meth:`Buffer.tile`.""" + + def __init__(self, buffer: Buffer, specs): + self._buffer = buffer + self._specs = specs + + def __getitem__(self, picks) -> Buffer: + if not isinstance(picks, tuple): + picks = (picks,) + n_factors = sum(len(factors) for _, factors in self._specs) + if len(picks) != n_factors: + raise ValueError( + f"tile: split into {n_factors} factor(s) but {len(picks)} " + f"index(es) given (int/Expr picks, ':' keeps)" + ) + groups, pos = [], 0 + for dim, factors in self._specs: + groups.append((dim, factors, picks[pos : pos + len(factors)])) + pos += len(factors) + buf = self._buffer + for dim, factors, dim_picks in sorted(groups, key=lambda group: group[0], reverse=True): + buf = self._apply(buf, dim, factors, dim_picks) + return buf + + @staticmethod + def _is_keep(pick): + return ( + isinstance(pick, slice) + and pick.start is None + and pick.stop is None + and pick.step is None + ) + + @classmethod + def _apply(cls, buf, dim, factors, dim_picks): + for pick in dim_picks: + if isinstance(pick, slice) and not cls._is_keep(pick): + raise ValueError("tile: a factor slice must be ':' (keep whole); got a sub-slice") + n_keep = sum(1 for pick in dim_picks if cls._is_keep(pick)) + if n_keep == len(dim_picks): + raise ValueError( + f"tile: dim {dim} picks no factor (all ':'); a chunk must pick " + f"at least one factor. Use .view / .chunk for a pure reshape." + ) + buf = buf.view(*buf.shape[:dim], *factors, *buf.shape[dim + 1 :]) + for index in reversed(range(len(dim_picks))): + if not cls._is_keep(dim_picks[index]): + buf = _view_drop(buf, dim + index, dim_picks[index]) + if n_keep >= 2: + buf = buf.view(*buf.shape[:dim], -1, *buf.shape[dim + n_keep :]) + return buf diff --git a/python/tvm/tirx/bench.py b/python/tvm/tirx/bench.py index 519a823bf5cc..c4bab4f2d105 100644 --- a/python/tvm/tirx/bench.py +++ b/python/tvm/tirx/bench.py @@ -16,12 +16,19 @@ # under the License. import argparse +import gc +import inspect +import json +import math import os -import re -import subprocess +import statistics import sys +import tempfile import time -from collections.abc import Mapping +import uuid +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from dataclasses import dataclass from enum import Enum import numpy as np @@ -33,6 +40,39 @@ from tvm.script import tirx as T from tvm.support import nvcc +_DISTRIBUTED_KINETO_WARMUP_ITERATIONS = 5 +_DISTRIBUTED_KINETO_REPEAT_ITERATIONS = 30 +_DISTRIBUTED_FLUSH_L2_BYTES = int(8e9) +_DISTRIBUTED_GPU_SLEEP_CYCLES = int(2e7) + + +@dataclass(frozen=True) +class DistributedBenchContext: + """Rank-local synchronization needed by the distributed Kineto timer. + + Parameters + ---------- + rank : int + Rank of the current process. + world_size : int + Number of ranks participating in the benchmark. + barrier : callable + Host callback that blocks until every rank reaches the same point. + max_reduce : callable + Host callback that returns the maximum of one floating-point value over + all ranks. The distributed timer calls it once per measured launch. + stream : torch.cuda.Stream + The actual CUDA stream whose launch scope is being timed. Launch + closures that use auxiliary streams must make this stream wait for them + before returning. + """ + + rank: int + world_size: int + barrier: Callable[[], None] + max_reduce: Callable[[float], float] + stream: object + def is_running_under_pytest(): """Check if the code is being executed within a pytest session.""" @@ -57,580 +97,731 @@ def tvm_callback_cuda_compile(code, target): return args -_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +# --------------------------------------------------------------------------- +# Triton-standard benchmark path. +# +# Faithful in-repo port of triton.testing.do_bench / do_bench_proton / +# do_bench_cudagraph_proton +# (see https://github.com/triton-lang/triton, python/triton/testing.py). This is +# a torch CUDA timer port: Triton runtime-driver calls (get_device_interface / +# get_empty_cache_for_benchmark / clear_cache) are replaced with torch.cuda and +# a torch L2-flush buffer. The Proton variants still use triton.profiler for +# attribution. The timed function is a *pure no-arg launch closure* (inputs +# captured once, allocated outside the timed region) -- exactly how Triton +# times a function. +# --------------------------------------------------------------------------- + + +def _quantile(a, q): + # pure-Python np.quantile / torch.quantile (port of triton.testing._quantile) + n = len(a) + a = sorted(a) + + def get_quantile(qi): + if not (0 <= qi <= 1): + raise ValueError("Quantiles must be in the range [0, 1]") + point = qi * (n - 1) + lower = math.floor(point) + upper = math.ceil(point) + t = point - lower + return (1 - t) * a[lower] + t * a[upper] + + return [get_quantile(qi) for qi in q] + + +def _summarize_statistics(times, quantiles, return_mode): + # port of triton.testing._summarize_statistics + if quantiles is not None: + ret = _quantile(times, quantiles) + if len(ret) == 1: + ret = ret[0] + return ret + if return_mode == "all": + return times + elif return_mode == "min": + return min(times) + elif return_mode == "max": + return max(times) + elif return_mode == "mean": + return statistics.mean(times) + elif return_mode == "median": + return statistics.median(times) + + +@contextmanager +def _cuda_graph_without_gc(*args, **kwargs): + # port of triton.testing.cuda_graph_without_gc. A loaded kernel may be + # finalized by Python's cyclic GC; its destructor unloads the CUDA module, + # which is illegal during CUDA stream capture and invalidates the graph. + # Keep GC disabled only for the capture window and restore afterwards. + gc_was_enabled = gc.isenabled() + if gc_was_enabled: + gc.disable() + try: + with torch.cuda.graph(*args, **kwargs) as graph: + yield graph + finally: + if gc_was_enabled: + gc.enable() + + +def _empty_cache_for_benchmark(): + # torch equivalent of triton's driver.get_empty_cache_for_benchmark(): a + # 256 MiB buffer whose .zero_() evicts the L2 cache between measured iters. + return torch.empty(256 * 1024 * 1024 // 4, dtype=torch.int, device="cuda") + + +def _do_bench_event( + fn, + warmup=25, + rep=100, + grad_to_none=None, + quantiles=None, + return_mode="mean", +): + """Faithful port of triton.testing.do_bench (CUDA-event timing, per-iter L2 flush). + ``warmup`` and ``rep`` are millisecond time budgets, not iteration counts. + Returns the runtime in milliseconds (mean by default). + """ + assert return_mode in ["min", "max", "mean", "median", "all"] -# proton-viewer -m avg_time/us prints average kernel time in microseconds (see -# triton/profiler/viewer.py avg_time_factor_dict). Store microseconds as-is. -PROTON_AVG_TIME_METRIC = "avg_time/us" + fn() + torch.cuda.synchronize() + cache = _empty_cache_for_benchmark() -def _parse_proton_tree(text, *, kernel: str = ""): - """Parse proton-viewer tree output into {impl: time_us}. + # Estimate the runtime of the function. + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(5): + cache.zero_() + fn() + end_event.record() + torch.cuda.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 + + # Compute number of warmup and repeat iterations from the ms budgets. + # Keep this identical to triton.testing.do_bench. Unlike Triton's Proton + # and CUDA-graph helpers, do_bench deliberately has no zero-time fallback. + n_warmup = max(1, int(warmup / estimate_ms)) + n_repeat = max(1, int(rep / estimate_ms)) + start_event = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeat)] + end_event = [torch.cuda.Event(enable_timing=True) for _ in range(n_repeat)] + # Warm-up. + for _ in range(n_warmup): + fn() + # Benchmark. + for i in range(n_repeat): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + # Clear the L2 cache before each run. + cache.zero_() + start_event[i].record() + fn() + end_event[i].record() + torch.cuda.synchronize() + times = [s.elapsed_time(e) for s, e in zip(start_event, end_event)] + return _summarize_statistics(times, quantiles, return_mode) - Accepts ALL depth-1 nodes (no KNOWN_IMPLS filter). For each depth-1 impl, - takes the slowest depth-2 child kernel time. - Tree numbers are microseconds when ProtonContext uses avg_time/us. +def _collect_proton_scope_times(database, prefix): + """Port of triton.testing._collect_proton_scope_times. - Returns (impl_times, baseline_errors) where: - impl_times: {str: float} — impl name to avg time in microseconds - baseline_errors: {str: str} — impl name to error message - """ - _ = kernel # kept for callers; unit does not depend on workload - impl = None - results = {} - baseline_errors = {} - for raw in text.splitlines(): - line = _ANSI_RE.sub("", raw).rstrip() - if not line: - continue - if line.startswith("BASELINE_ERROR:"): - parts = line.split(":", 2) - if len(parts) >= 3: - baseline_errors[parts[1].strip()] = parts[2].strip() - continue - # Depth-1 impl header: starts with tree drawing chars - if line and line[0] in "\u251c\u2514": # ├ └ - parts = line.split("\u2500", 1)[-1].split() # split on ─ - if len(parts) >= 2: - impl = parts[1] - else: - impl = None - continue - # Depth-2 kernel: contains tree drawing chars at deeper indent - if impl and ("\u251c\u2500" in line or "\u2514\u2500" in line): # ├─ └─ - parts = line.split("\u2500", 1)[-1].split() - if len(parts) >= 2: - name = parts[1] - if ( - "vectorized_elementwise_kernel" in name - or "elementwise_kernel_with_index" in name - ): - continue - try: - t = float(parts[0]) - results[impl] = max(results.get(impl, 0), t) - except ValueError: - pass - return results, baseline_errors - - -class ProtonContext: - """Context manager for Proton profiling sessions. - - Always captures proton-viewer output and parses impl times so that - get_impl_times() / get_baseline_errors() work after exiting the context. - - The proton tree is printed to **stdout** by default (visible on screen - when running kernels interactively). When the environment variable - ``TIRX_BENCH_JSON=1`` is set (done automatically by ``--json`` mode), - the tree goes to **stderr** instead so it does not corrupt the JSON on - stdout. + Walk the Proton hatchet JSON tree and, for each scope whose frame name starts + with ``prefix``, sum the GPU ``time (ns)`` of all its leaf kernels. Returns the + per-scope times (ms), sorted by scope name. """ + scope_times = [] + + def kernel_time_ms(node): + children = node.get("children", []) + if len(children) == 0: + return node.get("metrics", {}).get("time (ns)", 0) / 1e6 + return sum(kernel_time_ms(child) for child in children) + + def visit(node): + name = node.get("frame", {}).get("name", "") + if name.startswith(prefix): + time_ms = kernel_time_ms(node) + if time_ms > 0: + scope_times.append((name, time_ms)) + return + for child in node.get("children", []): + visit(child) + + for node in database: + # The hatchet top-level list may carry a device_info dict (no "frame"/ + # "children"); the name-prefix walk simply ignores it. + if isinstance(node, dict): + visit(node) + return [t for _, t in sorted(scope_times)] + + +def _start_proton_session(profile_path, *, timer, explicit_alternative): + """Start Proton or fail without changing the requested measurement method.""" + session = proton.start(profile_path, context="shadow", data="tree") + if session is None: + raise RuntimeError( + f"{timer}: Proton profiler session could not be created. " + f"Use timer={explicit_alternative!r} explicitly if that timing is intended." + ) + return session - def __init__( - self, - name="kernel", - hook="triton", - debug=False, - nsight=False, - metric=PROTON_AVG_TIME_METRIC, - kernel="", - ): - self.name = name - self.hook = hook - self.debug = debug - self.nsight = nsight - self.metric = metric - self.kernel = kernel - self._impl_times = {} - self._baseline_errors = {} - - def __enter__(self): - if not is_running_under_pytest() and not self.debug and not self.nsight: - proton.start(self.name, hook=self.hook) - proton.deactivate() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if not is_running_under_pytest() and not self.debug and not self.nsight: - proton.finalize() - - hatchet = f"{self.name}.hatchet" - result = subprocess.run( - ["proton-viewer", "-m", self.metric, hatchet], - capture_output=True, - text=True, - check=False, - ) - if result.returncode == 0: - self._impl_times, self._baseline_errors = _parse_proton_tree( - result.stdout, kernel=self.kernel - ) - out = sys.stderr if os.environ.get("TIRX_BENCH_JSON") else sys.stdout - print(f"# proton {PROTON_AVG_TIME_METRIC} (microseconds)\n", file=out, end="") - print(result.stdout, file=out, end="") - else: - print( - f"proton-viewer failed (rc={result.returncode}): {result.stderr}", - file=sys.stderr, - ) - - if os.path.exists(hatchet): - os.remove(hatchet) - - def get_impl_times(self): - """Return {impl_name: avg_time_us} parsed from proton-viewer output.""" - return dict(self._impl_times) - - def get_baseline_errors(self): - """Return {impl_name: error_message} from BASELINE_ERROR lines.""" - return dict(self._baseline_errors) - - -def _get_l2_cache_bytes(): - """Query L2 cache size from the current CUDA device, fallback to 128MB.""" - try: - props = torch.cuda.get_device_properties(torch.cuda.current_device()) - if hasattr(props, "l2_cache_size") and props.l2_cache_size > 0: - return props.l2_cache_size - except Exception: - pass - return 128 * 1024 * 1024 # 128MB default (B200) - - -def _tensor_bytes(args, _seen=None): - """Sum the byte size of all torch/tvm tensors in a nested value.""" - if _seen is None: - _seen = set() - total = 0 - if isinstance(args, list | tuple): - for a in args: - total += _tensor_bytes(a, _seen) - elif isinstance(args, Mapping): - for a in args.values(): - total += _tensor_bytes(a, _seen) - elif isinstance(args, torch.Tensor): - key = ("torch", args.device.type, args.device.index, int(args.data_ptr())) - if key not in _seen: - _seen.add(key) - total += args.nelement() * args.element_size() - elif hasattr(args, "numpy"): # tvm.runtime.NDArray - try: - key = ("tvm", int(args.handle.value)) - except Exception: - key = ("tvm", id(args)) - if key not in _seen: - _seen.add(key) - try: - total += int(np.prod(args.shape)) * np.dtype(str(args.dtype)).itemsize - except Exception: - total += args.numpy().nbytes - return total +def _do_bench_proton( + fn, + warmup=25, + rep=100, + grad_to_none=None, + quantiles=None, + return_mode="mean", +): + """Port of triton.testing.do_bench_proton, aligned with ``_do_bench_event``. -def tensor_bytes(*values): - """Return unique torch/tvm tensor bytes for kernel-owned byte accounting. + IDENTICAL to ``_do_bench_event`` in everything -- warmup/rep millisecond budgets, + the 5-call estimate, per-iter L2 flush, the untimed warmup loop -- EXCEPT the + timing mechanism: each timed call runs inside a Proton scope and the per-kernel + GPU time (read from the hatchet tree) is used instead of the CUDA-event wall. + Cold cache. NVIDIA + Proton only. No CUDA graph (so it works for references that + can't be graph-captured, e.g. CuTeDSL flash-attention). - The benchmark driver does not use this implicitly. Kernel benchmark - factories may call it when their invocation footprint is exactly the set of - tensors in ``values``. + A missing Proton session is an error. Silently switching to event timing would + leave the result labelled ``proton`` while changing the measured quantity. """ - if len(values) == 1: - return _tensor_bytes(values[0]) - return _tensor_bytes(values) + assert return_mode in ["min", "max", "mean", "median", "all"] + fn() + torch.cuda.synchronize() -def _compute_group_count(input_bytes, l2_bytes=None): - """Return TK-style input-group count from one invocation's byte footprint.""" - if input_bytes <= 0: - return 1 - if l2_bytes is None: - l2_bytes = _get_l2_cache_bytes() - threshold = l2_bytes * 3 - if input_bytes >= threshold: - return 1 - return int(threshold // input_bytes) + 1 + cache = _empty_cache_for_benchmark() + # Estimate the runtime of the function (identical to _do_bench_event). + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + start_event.record() + for _ in range(5): + cache.zero_() + fn() + end_event.record() + torch.cuda.synchronize() + estimate_ms = start_event.elapsed_time(end_event) / 5 -def _make_bench_input(input_factory): - value = input_factory() - if not isinstance(value, tuple) or len(value) != 2: - raise TypeError("input_factory must return (case, input_bytes)") + if estimate_ms == 0: + n_warmup = 1000 + n_repeat = 1000 + else: + n_warmup = max(1, int(warmup / estimate_ms)) + n_repeat = max(1, int(rep / estimate_ms)) - case, input_bytes = value - try: - input_bytes = int(input_bytes) - except (TypeError, ValueError) as err: - raise TypeError("input_factory input_bytes must be an integer") from err - if input_bytes < 0: - raise ValueError("input_factory input_bytes must be non-negative") - return case, input_bytes + # Warm-up (untimed), same as _do_bench_event. + for _ in range(n_warmup): + fn() + torch.cuda.synchronize() + with tempfile.TemporaryDirectory(prefix=f"tirx-proton-{uuid.uuid4().hex}-") as tmpdir: + profile_path = os.path.join(tmpdir, "profile") + session = _start_proton_session(profile_path, timer="proton", explicit_alternative="event") + scope_prefix = f"proton.{uuid.uuid4().hex}." + # finalize() MUST run even if fn() raises mid-loop; otherwise the global + # Proton profiler stays active and the next session in this process starts + # dirty (corrupted attribution). Mirrors triton.testing._proton_bench_session + # (finalize in a finally), adapted to read the .hatchet finalize writes. + try: + for i in range(n_repeat): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + # Flush L2 OUTSIDE the scope so it is excluded from the measured time + # -- identical cold-cache behavior to _do_bench_event. + cache.zero_() + with proton.scope(f"{scope_prefix}{i:08d}"): + fn() + torch.cuda.synchronize() + finally: + proton.finalize(session) + with open(profile_path + ".hatchet") as f: + database = json.load(f) + times = _collect_proton_scope_times(database, scope_prefix) + + if not times: + raise RuntimeError( + "proton: Proton attributed no kernel time to the captured scopes. " + "Use timer='event' instead." + ) + return _summarize_statistics(times, quantiles, return_mode) -def prepare_input_groups(input_factory, l2_bytes=None): - """Materialize TK-style input groups from a single-group factory. - ``input_factory`` must return ``(case, input_bytes)``. ``case`` is passed - back to every benchmark function unchanged. ``input_bytes`` defines one - invocation's L2-eviction footprint and is intentionally owned by the kernel - benchmark harness instead of inferred here. - """ - if not callable(input_factory): - raise TypeError("input_factory must be callable") - if l2_bytes is None: - l2_bytes = _get_l2_cache_bytes() - - sample, input_bytes = _make_bench_input(input_factory) - num_groups = _compute_group_count(input_bytes, l2_bytes) - groups = [sample] - for _ in range(num_groups - 1): - case, _ = _make_bench_input(input_factory) - groups.append(case) - - return groups, { - "num_groups": num_groups, - "input_bytes": input_bytes, - "l2_bytes": l2_bytes, - "l2_eviction_factor": 3, - "flush_l2": False, - } +def _do_bench_cudagraph_proton(fn, rep=20, grad_to_none=None, quantiles=None, return_mode="mean"): + """Faithful port of triton.testing.do_bench_cudagraph_proton. + CUDA-graph replay (kills per-launch CPU overhead) + Proton per-kernel GPU time + + per-iter L2 flush. Best accuracy for short / multi-kernel workloads. NVIDIA + only (Proton cannot reliably attribute graph-replay launches to scopes on HIP). + ``rep`` (ms) sets the graph unroll count (``n_repeat = rep / estimate_ms``); the + measurement is 10 graph replays. Triton's default is ``rep=20``. Returns ms. -def _bench_event_groups(funcs, groups, warmup, repeat, cooldown_s): - num_groups = len(groups) - results = {} + Adapted to the installed proton (3.6.0): there is no ``proton.data.get`` / + ``deactivate(flushing=True)`` here, so we read the ``.hatchet`` JSON that + ``finalize`` writes (the same tree the in-memory getter would return). - for idx, (name, func) in enumerate(funcs.items()): - if idx > 0: - time.sleep(cooldown_s) + A missing Proton session is an error; this timer never changes measurement + method while retaining the ``cudagraph_proton`` label. + """ + assert return_mode in ["min", "max", "mean", "median", "all"] - for i in range(warmup): - func(groups[i % num_groups]) + if torch.version.hip is not None: + raise RuntimeError( + "cudagraph_proton requires the NVIDIA backend because Proton does not " + "reliably attribute CUDA graph replay launches to scopes on HIP." + ) + with torch.cuda.stream(torch.cuda.Stream()): + # warmup + fn() + if grad_to_none is not None: + for x in grad_to_none: + x.detach_() + x.requires_grad_(True) + x.grad = None + # estimate single-call time start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() - start_event.record() - for i in range(repeat): - func(groups[i % num_groups]) + for _ in range(5): + fn() end_event.record() - torch.cuda.synchronize() - results[name] = start_event.elapsed_time(end_event) / repeat * 1000.0 - - time.sleep(cooldown_s) - - return results - - -def _bench_proton_groups( - funcs, groups, warmup, repeat, cooldown_s, proton_name, debug, nsight, *, kernel="" -): - num_groups = len(groups) - with ProtonContext(proton_name, debug=debug, nsight=nsight, kernel=kernel) as ctx: - for idx, (name, func) in enumerate(funcs.items()): - if idx > 0: - time.sleep(cooldown_s) - - for i in range(warmup): - func(groups[i % num_groups]) - torch.cuda.synchronize() - - if not is_running_under_pytest() and not debug and not nsight: - proton.activate() - with proton.scope(name, metrics={}): - for i in range(repeat): - func(groups[i % num_groups]) - proton.deactivate() - else: - for i in range(repeat): - func(groups[i % num_groups]) - torch.cuda.synchronize() - - time.sleep(cooldown_s) - - return ctx.get_impl_times(), ctx.get_baseline_errors() + estimate_ms = start_event.elapsed_time(end_event) / 5 + n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms)) + + with tempfile.TemporaryDirectory(prefix=f"tirx-cgproton-{uuid.uuid4().hex}-") as tmpdir: + profile_path = os.path.join(tmpdir, "profile") + # shadow/CUPTI captures kernel GPU activity itself; the triton launch + # hook only adds flops/bytes metadata, so it is omitted here. + session = _start_proton_session( + profile_path, + timer="cudagraph_proton", + explicit_alternative="event", + ) + cache = _empty_cache_for_benchmark() + scope_prefix = f"proton.{uuid.uuid4().hex}." + g = torch.cuda.CUDAGraph() + n_retries = 10 + # finalize() MUST run even if capture/replay raises, or the global Proton + # profiler stays active and poisons the next session (see _do_bench_proton). + try: + with _cuda_graph_without_gc(g): + for i in range(n_repeat): + if grad_to_none is not None: + for x in grad_to_none: + x.grad = None + # Flush L2 OUTSIDE the scope so it is excluded from the timed span. + cache.zero_() + with proton.scope(f"{scope_prefix}{i:08d}"): + fn() + torch.cuda.synchronize() + for _ in range(n_retries): + g.replay() + torch.cuda.synchronize() + finally: + # finalize flushes the replay data and writes .hatchet. + proton.finalize(session) + with open(profile_path + ".hatchet") as f: + database = json.load(f) + times = [t / n_retries for t in _collect_proton_scope_times(database, scope_prefix)] + + if not times: + raise RuntimeError( + "cudagraph_proton: Proton attributed no kernel time to the captured scopes " + "(CUDA-graph replay scope attribution may be unsupported in this " + "environment). Use timer='event' or 'proton' instead." + ) + return _summarize_statistics(times, quantiles, return_mode) -def _flush_l2_legacy(flush_l2_size): - if flush_l2_size > 0: - torch.empty(flush_l2_size, dtype=torch.int, device="cuda").zero_() +def _sleep_before_impl(cooldown_s): + if cooldown_s > 0: + time.sleep(cooldown_s) -def _bench_legacy_callable(func, warmup, repeat, proton_name, debug, nsight, flush_l2_size): - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - def timed_loop(): - start_event.record() - for _ in range(repeat): - _flush_l2_legacy(flush_l2_size) +def _validate_distributed_context(distributed): + if not isinstance(distributed, DistributedBenchContext): + raise TypeError("distributed must be a DistributedBenchContext") + if ( + not isinstance(distributed.world_size, int) + or isinstance(distributed.world_size, bool) + or distributed.world_size < 1 + ): + raise ValueError("distributed.world_size must be a positive integer") + if ( + not isinstance(distributed.rank, int) + or isinstance(distributed.rank, bool) + or not 0 <= distributed.rank < distributed.world_size + ): + raise ValueError("distributed.rank must be in [0, world_size)") + if not callable(distributed.barrier): + raise TypeError("distributed.barrier must be callable") + if not callable(distributed.max_reduce): + raise TypeError("distributed.max_reduce must be callable") + if not hasattr(distributed.stream, "synchronize"): + raise TypeError("distributed.stream must be an actual CUDA stream") + + +def _distributed_cold_start(distributed): + """Apply DeepGEMM's cold-cache protocol before one distributed launch.""" + torch.empty( + _DISTRIBUTED_FLUSH_L2_BYTES // 4, + dtype=torch.int, + device="cuda", + ).zero_() + torch.cuda._sleep(_DISTRIBUTED_GPU_SLEEP_CYCLES) + distributed.barrier() + + +def _collect_kineto_span_samples(profiler, scope_names): + """Collect one cross-stream GPU activity span for every record-function scope.""" + scope_events = {scope_name: [] for scope_name in scope_names} + for event in profiler.events(): + if event.device_type == torch.autograd.DeviceType.CUDA and event.name in scope_events: + scope_events[event.name].append(event) + + samples = [] + stream_counts = [] + for scope_name in scope_names: + events = scope_events[scope_name] + if not events: + raise RuntimeError( + f"kineto: profiler attributed no CUDA activity to scope {scope_name!r}" + ) + start_us = min(event.time_range.start for event in events) + end_us = max(event.time_range.end for event in events) + elapsed_us = float(end_us - start_us) + if not math.isfinite(elapsed_us) or elapsed_us <= 0: + raise RuntimeError(f"kineto: scope {scope_name!r} has invalid GPU span {elapsed_us} us") + samples.append(elapsed_us) + stream_counts.append(len({event.device_resource_id for event in events})) + return samples, stream_counts + + +def _profile_distributed_kineto_span(name, func, prepare, distributed): + """Profile complete correlated GPU activity spans for one implementation.""" + prepare_fn = prepare.get(name) + with torch.cuda.stream(distributed.stream): + for _ in range(_DISTRIBUTED_KINETO_WARMUP_ITERATIONS): + _distributed_cold_start(distributed) + if prepare_fn is not None: + prepare_fn() func() - end_event.record() - - for _ in range(warmup): - _flush_l2_legacy(flush_l2_size) - func() - torch.cuda.synchronize() - if not is_running_under_pytest() and not debug and not nsight: - proton.activate() - with proton.scope(proton_name, metrics={}): - timed_loop() - proton.deactivate() - else: - timed_loop() - torch.cuda.synchronize() - - return start_event.elapsed_time(end_event) / repeat * 1000.0 - - -# Labels identifying our own kernel (vs external reference impls). Must match -# OUR_IMPLS in bench_suite's ratio_diff.py. Used by the TIRX_BENCH_IMPLS filter. -OURS_IMPLS = frozenset({"tir", "tirx"}) - - -def bench_impls_mode(): - """Current impl-selection mode: 'all' (default), 'ours', or 'baseline'. - - Set via the ``TIRX_BENCH_IMPLS`` env var (by bench_suite ``run.py --impls``). - A kernel's ``run_bench`` can use this to skip *building/warming* reference - impls (e.g. flashinfer autotune, deepgemm/cublas ext setup) that ``bench``'s - filter alone cannot avoid, since that setup runs before ``bench`` is called. - """ - mode = os.environ.get("TIRX_BENCH_IMPLS", "all").lower() - if mode not in {"all", "ours", "baseline"}: - raise ValueError(f"TIRX_BENCH_IMPLS must be 'all', 'ours', or 'baseline', got {mode!r}") - return mode - - -def bench_only_ours(): - """True when only our own kernel should be benched (reference setup skippable).""" - return bench_impls_mode() == "ours" - - -def filter_impls(funcs): - """Filter a ``{label: callable}`` impl map per the current ``bench_impls_mode``. + distributed.stream.synchronize() + + scope_names = [ + f"tirx.bench.{name}.{index:08d}" + for index in range(_DISTRIBUTED_KINETO_REPEAT_ITERATIONS) + ] + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) + with profiler: + for scope_name in scope_names: + _distributed_cold_start(distributed) + if prepare_fn is not None: + prepare_fn() + with torch.profiler.record_function(scope_name): + func() + distributed.stream.synchronize() + + return _collect_kineto_span_samples(profiler, scope_names) + + +def _distributed_stream_id(stream): + value = getattr(stream, "cuda_stream", None) + return int(value) if value is not None else repr(stream) + + +def _bench_distributed_kineto_span(funcs, prepare, distributed, rounds, cooldown_s): + """Time complete correlated GPU activity spans across all participating streams.""" + items = list(funcs.items()) + round_samples = {name: [] for name in funcs} + stream_counts = {name: [] for name in funcs} + for _ in range(rounds): + for name, func in items: + _sleep_before_impl(cooldown_s) + local_samples, local_stream_counts = _profile_distributed_kineto_span( + name, func, prepare, distributed + ) + samples = [] + for local_us in local_samples: + elapsed_us = float(distributed.max_reduce(local_us)) + if not math.isfinite(elapsed_us) or elapsed_us <= 0: + raise RuntimeError( + "kineto: distributed max-reduce returned invalid activity span " + f"for {name!r}: {elapsed_us}" + ) + samples.append(elapsed_us) + round_samples[name].append(statistics.median(samples)) + stream_counts[name].append( + {"min": min(local_stream_counts), "max": max(local_stream_counts)} + ) + distributed.barrier() - Call this right after building ``funcs`` so any subsequent - ``if "" in funcs:`` reference-setup blocks are skipped in 'ours' mode. - """ - mode = bench_impls_mode() - if mode == "ours": - return {n: f for n, f in funcs.items() if n in OURS_IMPLS} - if mode == "baseline": - return {n: f for n, f in funcs.items() if n not in OURS_IMPLS} - return funcs + return { + "impls": {name: statistics.mean(samples) for name, samples in round_samples.items()}, + "round_samples": round_samples, + "errors": {}, + "timer": "kineto", + "benchmark_protocol": { + "source": "torch.profiler Kineto correlated CUDA activities", + "timing_scope": "complete correlated GPU activity span", + "span_definition": "latest activity end minus earliest activity start", + "timing_stream": _distributed_stream_id(distributed.stream), + "all_correlated_streams": True, + "warmup_iterations": _DISTRIBUTED_KINETO_WARMUP_ITERATIONS, + "repeat_iterations": _DISTRIBUTED_KINETO_REPEAT_ITERATIONS, + "flush_l2": True, + "flush_l2_bytes": _DISTRIBUTED_FLUSH_L2_BYTES, + "gpu_sleep_cycles": _DISTRIBUTED_GPU_SLEEP_CYCLES, + "rank_barrier_before_each_launch": True, + "cold_start_outside_scope": True, + "prepare_outside_scope": True, + "rank_aggregate": "sample_wise_max", + "sample_aggregate": "median", + "round_aggregate": "mean", + "world_size": distributed.world_size, + "rounds": rounds, + "cooldown_s": cooldown_s, + "order": [name for name, _ in items], + "rank_local_scope_stream_counts": stream_counts, + }, + } def bench( funcs, - input_factory=None, - warmup=500, - repeat=100, - cooldown_s=1.0, - timer="proton", - proton_name="kernel", - l2_bytes=None, - debug=False, - nsight=False, - flush_l2_size=int(8e8 // 4), + *, + warmup=None, + repeat=None, + cudagraph_rep=None, + timer=None, references=None, + cooldown_s=1.0, rounds=1, - round_cooldown_s=1.0, - validate_case=None, + distributed=None, + prepare=None, ): - """Benchmark implementations with a factory-owned input footprint. + """Benchmark pure-launch implementations using Triton-standard timing. - This is the single TIRx benchmark API. It follows the ThunderKittens-style - multi-input protocol for L2 eviction and supports either Proton/CUPTI or - CUDA-event timing. The benchmark driver never infers which tensors belong - to a workload; ``input_factory`` owns that definition by returning - ``(case, input_bytes)``. + Each callable in ``funcs`` is a *no-arg launch closure* (inputs allocated + once and captured in the closure), timed exactly the way + ``triton.testing.do_bench`` / ``do_bench_proton`` / + ``do_bench_cudagraph_proton`` time a function. + The timing core is a faithful in-repo port (torch-only, no Triton runtime + driver dependency); see the corresponding ``_do_bench_*`` functions. Parameters ---------- funcs : dict[str, callable] - Map of implementation name to callable. Each callable receives one - ``case`` returned by ``input_factory``. This should hold only *our* - kernel(s); external baselines go in ``references``. + Map of implementation name to a no-arg callable that launches our + kernel. This should hold only *our* kernel(s); external baselines go in + ``references``. references : dict[str, callable], optional Map of reference-impl name to a no-arg *builder* that does the heavy - import/setup and returns the run callable. Builders run lazily and only - when that impl will actually be benched (skipped entirely under - ``--impls ours``); a builder that raises is recorded as a - ``BASELINE_ERROR`` instead of failing the workload. - input_factory : callable - Factory returning ``(case, input_bytes)`` for one benchmark group. - warmup : int - Number of untimed warmup iterations per implementation. - repeat : int - Number of timed iterations per round. - cooldown_s : float - Seconds to sleep between impls for thermal cooldown. + import/setup and returns the no-arg run callable. A builder that raises + is recorded as a ``BASELINE_ERROR`` instead of failing the workload. + warmup : int, optional + Warmup time budget (ms) for the ``event`` and ``proton`` timers. ``None`` + (default) defers to the selected timer's own default; pass a value only + to override. Ignored by ``cudagraph_proton``, which has no warmup argument. + repeat : int, optional + Rep time budget (ms) for the ``event`` and ``proton`` timers. ``None`` + (default) defers to the selected timer's own default. + cudagraph_rep : int, optional + ``rep`` (ms) for ``cudagraph_proton`` (graph unroll length). ``None`` + (default) defers to that timer's own default. Each timer default lives only + in its own signature (Triton: do_bench 25/100, graph 20); nothing is hardcoded + here. + timer : {None, "event", "proton", "cudagraph_proton", "kineto"} + ``None`` (default) resolves to ``proton`` for a local benchmark and to + ``kineto`` when ``distributed`` is present. The default is defined in + exactly one place so kernels pass ``None`` to inherit. + ``event`` -> ported ``do_bench`` (CUDA-event wall of each call). + ``proton`` -> ported ``do_bench_proton``: same setup as ``event``, differs + ONLY in timing -- Proton per-kernel GPU time instead of the event wall (so + launch/host overhead of the ref is excluded). No graph, so it works for + references that can't be CUDA-graph-captured (e.g. CuTeDSL flash-attention). + NVIDIA + Proton only. + ``cudagraph_proton`` -> ported ``do_bench_cudagraph_proton`` (graph replay + + Proton kernel time, with L2 flushes captured outside Proton scopes); like + ``proton`` it reports kernel time, but also removes launch overhead via + graph replay. Graph capture or attribution is not universal. + Prefer ``proton`` when the reference has heavy host dispatch (flashinfer, + CuTeDSL) -- ``event`` would over/under-credit us by measuring the wall, not the + kernel. Timer failures are propagated; a result is never silently relabelled + after falling back to a different measurement method. + A distributed benchmark supports only ``kineto``. It measures the complete + correlated GPU activity span across all streams, performs sample-wise + slowest-rank aggregation, and applies DeepGEMM's 8 GB L2 flush, GPU sleep, + and rank barrier before every launch. Distributed timing uses fixed + iteration counts and rejects local timer budget overrides. + distributed : DistributedBenchContext, optional + Rank-local barrier, max-reduce callback, and actual CUDA timing stream. + Launch closures that use auxiliary streams must make the timing stream + wait for them before returning. + prepare : dict[str, callable], optional + Per-implementation reset/preparation callbacks. Kineto scopes begin after + preparation, so preparation activities are excluded from the measured span. rounds : int - Independent benchmark rounds (compile + inputs once; each round runs - warmup + repeat for every selected impl). - round_cooldown_s : float - Seconds to sleep between rounds (ignored when ``rounds == 1``). - validate_case : callable, optional - Called once on the first prepared ``case`` (after ``prepare_input_groups``, - before warmup/repeat rounds). Under bench_suite, ``run_kernel_bench`` holds - the per-GPU lock for the whole ``run_bench()`` call. - timer : {"event", "proton"} - Timing backend. + Independent measurement rounds; per-impl times are averaged across + rounds. (Triton times a single fn with no rounds; this is our sampling + layer on top.) + cooldown_s : float + Seconds to sleep immediately before each implementation in every round. Returns ------- dict ``{"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}``. - Times are stored in microseconds (same unit as pinned bench_suite baselines). + Times are stored in microseconds (same unit as the pinned tir-bench baselines). """ - if repeat <= 0: - raise ValueError("repeat must be positive") - if warmup < 0: - raise ValueError("warmup must be non-negative") - if rounds < 1: + if not isinstance(rounds, int) or isinstance(rounds, bool) or rounds < 1: raise ValueError("rounds must be >= 1") - if round_cooldown_s < 0: - raise ValueError("round_cooldown_s must be non-negative") - if timer not in {"event", "proton"}: - raise ValueError(f"unsupported timer {timer!r}; expected event or proton") - - if callable(funcs) and input_factory is None: - return _bench_legacy_callable( - funcs, - warmup=warmup, - repeat=repeat, - proton_name=proton_name, - debug=debug, - nsight=nsight, - flush_l2_size=flush_l2_size, - ) - - if input_factory is None: - raise TypeError("input_factory is required when funcs is a mapping") + if cooldown_s < 0: + raise ValueError("cooldown_s must be non-negative") + + if distributed is None: + if prepare is not None: + raise ValueError("prepare is supported only with a distributed context") + if repeat is not None and repeat <= 0: + raise ValueError("repeat must be positive") + if warmup is not None and warmup < 0: + raise ValueError("warmup must be non-negative") + if timer is None: + timer = "proton" + if timer not in {"event", "proton", "cudagraph_proton"}: + raise ValueError( + f"unsupported timer {timer!r}; expected event, proton, or cudagraph_proton" + ) + else: + _validate_distributed_context(distributed) + if timer is None: + timer = "kineto" + elif timer != "kineto": + raise ValueError("a distributed context supports only timer='kineto'") + overrides = [ + name + for name, value in ( + ("warmup", warmup), + ("repeat", repeat), + ("cudagraph_rep", cudagraph_rep), + ) + if value is not None + ] + if overrides: + raise ValueError( + f"timer={timer!r} uses fixed iteration counts and rejects overrides: " + + ", ".join(overrides) + ) if not isinstance(funcs, Mapping) or not funcs: - raise TypeError("funcs must be a non-empty mapping of name to callable") + raise TypeError("funcs must be a non-empty mapping of name to no-arg callable") for name, func in funcs.items(): if not isinstance(name, str): raise TypeError("func names must be strings") if not callable(func): raise TypeError(f"funcs[{name!r}] must be callable") - # Select impls for this run. ``funcs`` holds our own kernel(s); external - # baselines are passed as ``references`` (name -> no-arg builder). A builder - # is invoked here ONLY when its impl will actually be benched, so --impls - # ours skips all reference setup, --impls baseline skips ours, and a builder - # that fails is recorded as a BASELINE_ERROR rather than failing the - # workload. Legacy kernels that put references directly in ``funcs`` are - # still handled by the label filter (filter_impls) below. - funcs = filter_impls(funcs) + # ``funcs`` holds our own kernel(s); external baselines are passed as + # ``references`` (name -> no-arg builder). A builder that fails is recorded + # as a BASELINE_ERROR rather than failing the workload. build_errors: dict[str, str] = {} - if bench_impls_mode() != "ours": - for ref_name, builder in (references or {}).items(): - if not isinstance(ref_name, str) or not callable(builder): - raise TypeError("references must map a name to a no-arg builder callable") - try: - ref_fn = builder() - except Exception as e: - build_errors[ref_name] = str(e) - print(f"BASELINE_ERROR: {ref_name}: {e}", file=sys.stderr) - continue - if ref_fn is None: - continue - if not callable(ref_fn): - raise TypeError(f"references[{ref_name!r}] builder must return a callable") - funcs = {**funcs, ref_name: ref_fn} - - if not funcs: - # Nothing to bench in this mode (e.g. 'ours' on a kernel with no tir - # impl, or 'baseline' with no references). Return an empty-but-valid - # result so the workload is a no-op. - return { - "impls": {}, - "round_samples": {}, - "errors": build_errors, - "timer": timer, - "benchmark_protocol": { - "warmup": warmup, - "repeat": repeat, - "cooldown_s": cooldown_s, - "rounds": rounds, - "round_cooldown_s": round_cooldown_s, - "order": [], - }, - } - - inputs, protocol = prepare_input_groups(input_factory, l2_bytes=l2_bytes) - num_groups = len(inputs) - if num_groups == 0: - return { - "impls": {}, - "round_samples": {}, - "errors": build_errors, - "timer": timer, - "benchmark_protocol": { - **protocol, - "warmup": warmup, - "repeat": repeat, - "cooldown_s": cooldown_s, - "rounds": rounds, - "round_cooldown_s": round_cooldown_s, - "order": list(funcs.keys()), - }, + for ref_name, builder in (references or {}).items(): + if not isinstance(ref_name, str) or not callable(builder): + raise TypeError("references must map a name to a no-arg builder callable") + try: + ref_fn = builder() + except Exception as e: + build_errors[ref_name] = str(e) + print(f"BASELINE_ERROR: {ref_name}: {e}", file=sys.stderr) + continue + if ref_fn is None: + continue + if not callable(ref_fn): + raise TypeError(f"references[{ref_name!r}] builder must return a callable") + funcs = {**funcs, ref_name: ref_fn} + + if distributed is not None: + if prepare is None: + prepare = {} + if not isinstance(prepare, Mapping): + raise TypeError("prepare must map implementation names to no-arg callables") + for name, prepare_fn in prepare.items(): + if name not in funcs: + raise ValueError(f"prepare contains unknown implementation {name!r}") + if not callable(prepare_fn): + raise TypeError(f"prepare[{name!r}] must be callable") + result = _bench_distributed_kineto_span( + funcs, prepare, distributed, rounds=rounds, cooldown_s=cooldown_s + ) + result["errors"] = build_errors + return result + + # Resolve the timer function once. Only forward warmup/repeat/cudagraph_rep when a + # caller explicitly overrode them; otherwise the _do_bench_* signature default + # applies, so each default lives in exactly ONE place (its timer signature). The + # effective value is read back via inspect, so the recorded protocol tracks that + # default automatically even if it later changes -- no value is duplicated here. + def _sig_default(fn, param): + return inspect.signature(fn).parameters[param].default + + if timer in ("event", "proton"): + # event and proton share the exact same warmup/rep setup; they differ only in + # how the timed calls are measured (CUDA-event wall vs Proton per-kernel time). + _timer_fn = _do_bench_event if timer == "event" else _do_bench_proton + _timer_kwargs = {} + if warmup is not None: + _timer_kwargs["warmup"] = warmup + if repeat is not None: + _timer_kwargs["rep"] = repeat + _eff = { + "warmup": _timer_kwargs.get("warmup", _sig_default(_timer_fn, "warmup")), + "repeat": _timer_kwargs.get("rep", _sig_default(_timer_fn, "rep")), } + else: # cudagraph_proton has no warmup; rep is the graph-unroll budget + _timer_fn = _do_bench_cudagraph_proton + _timer_kwargs = {} + if cudagraph_rep is not None: + _timer_kwargs["rep"] = cudagraph_rep + _eff = {"cudagraph_rep": _timer_kwargs.get("rep", _sig_default(_timer_fn, "rep"))} + + protocol = { + **_eff, + "cooldown_s": cooldown_s, + "rounds": rounds, + "round_aggregate": "mean", + "order": list(funcs.keys()), + } - if validate_case is not None: - validate_case(inputs[0]) - - errors = dict(build_errors) round_samples: dict[str, list[float]] = {} - for round_idx in range(rounds): - if round_idx > 0: - time.sleep(round_cooldown_s) - if timer == "event": - impls = _bench_event_groups(funcs, inputs, warmup, repeat, cooldown_s) - proton_errors = {} - else: - impls, proton_errors = _bench_proton_groups( - funcs, - inputs, - warmup, - repeat, - cooldown_s, - proton_name, - debug, - nsight, - kernel=proton_name, - ) - errors.update(proton_errors) - for impl, sec in impls.items(): - round_samples.setdefault(impl, []).append(sec) - - if not round_samples: - aggregated = {} - else: - import statistics + for _ in range(rounds): + for name, func in funcs.items(): + _sleep_before_impl(cooldown_s) + ms = _timer_fn(func, **_timer_kwargs) + # ms -> microseconds (matches pinned baseline units). + round_samples.setdefault(name, []).append(ms * 1000.0) - aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()} + aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()} return { "impls": aggregated, "round_samples": round_samples, - "errors": errors, + "errors": build_errors, "timer": timer, - "benchmark_protocol": { - **protocol, - "warmup": warmup, - "repeat": repeat, - "cooldown_s": cooldown_s, - "rounds": rounds, - "round_cooldown_s": round_cooldown_s, - "order": list(funcs.keys()), - }, + "benchmark_protocol": protocol, } diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index ef9cee6ee0d9..6246e5540b36 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -16,7 +16,6 @@ # under the License. """Abstraction for array data structures.""" -import functools from enum import IntEnum from numbers import Integral @@ -26,7 +25,9 @@ from tvm.ir import PointerType, PrimType, Range, Type from tvm.runtime import Object, convert -from . import _ffi_api +from . import _buffer_view, _ffi_api + +_REARRANGE_PATTERN_UNSET = object() @tvm_ffi.register_object("tirx.BufferType") @@ -326,90 +327,7 @@ def view(self, *args, **kwargs) -> "Buffer": The corresponding view buffer. """ - def _infer_shape(shape): - shape = list(shape) - if -1 in shape and shape.count(-1) == 1: - size = functools.reduce(lambda x, y: x * y, self.ty.shape) - n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1) - shape[shape.index(-1)] = size // n_size - else: - # Only validate the shape product when both old and new shapes - # are fully concrete: a Expr `==` returns an `EQ` node, not - # a Python bool, and `assert ` raises (no __bool__). - if all(isinstance(s, int) for s in shape) and all( - isinstance(s, int) for s in self.ty.shape - ): - assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce( - lambda x, y: x * y, self.ty.shape - ), ( - "The shape of the buffer " - + str(self.ty.shape) - + " and the new shape " - + str(shape) - + " are not compatible" - ) - return shape - - if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs: - cast_dtype = tvm.DataType(args[0]) - cur_dtype = tvm.DataType(self.ty.dtype) - if cast_dtype.bits > cur_dtype.bits: - # cast up - assert cast_dtype.bits % cur_dtype.bits == 0 - ratio = cast_dtype.bits // cur_dtype.bits - layout = self.ty.layout.pack(ratio) - shape = [s for s in self.ty.shape[:-1]] + [self.ty.shape[-1] // ratio] - new_elem_offset = self.ty.elem_offset // ratio - else: - # cast down - assert cur_dtype.bits % cast_dtype.bits == 0 - ratio = cur_dtype.bits // cast_dtype.bits - layout = self.ty.layout.unpack(ratio) - shape = [s for s in self.ty.shape[:-1]] + [self.ty.shape[-1] * ratio] - new_elem_offset = self.ty.elem_offset * ratio - return tvm.tirx.script.builder.decl_buffer( - shape, - cast_dtype, - buffer_data(self), - self.ty.strides, - new_elem_offset, - None, - self.ty.storage_scope, - self.ty.data_alignment, - self.ty.offset_factor, - layout, - ) - else: - # --- Signature 1: view(*shape, **opts) --- - # Check if all positional args are integers/PrimExprs with dtype int32 or int64 (the shape) # noqa: E501 - shape = args - assert all( - isinstance(arg, int) - or (tvm.ir.is_prim_expr(arg) and arg.ty.dtype in ["int32", "int64"]) - for arg in shape - ), "shape must be a list of integers or PrimExprs with dtype int32 or int64" - # Safely get optional keyword arguments - layout = kwargs.get("layout", None) - # Assert there are no other kwargs - assert set(kwargs.keys()).issubset({"layout"}), ( - f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}" - ) - - if layout is None: - shape = _infer_shape(shape) - - return tvm.tirx.script.builder.decl_buffer( - shape, - self.ty.dtype, - buffer_data(self), - self.ty.strides, - self.ty.elem_offset, - None, - self.ty.storage_scope, - self.ty.data_alignment, - self.ty.offset_factor, - self.ty.layout if layout is None else layout, - ) + return _buffer_view.view(self, *args, **kwargs) def local(self, *shape, layout=None) -> "Buffer": """Create a thread-local view of this buffer. @@ -432,24 +350,7 @@ def local(self, *shape, layout=None) -> "Buffer": local : DeclBufferFrame The corresponding local buffer. """ - if not shape: - local_layout = self.ty.layout.storage() - total = functools.reduce( - lambda x, y: x * y, [it.extent for it in local_layout.shard], 1 - ) - shape = (total,) - return tvm.tirx.script.builder.decl_buffer( - shape, - self.ty.dtype, - buffer_data(self), - self.ty.strides, - self.ty.elem_offset, - None, - self.ty.storage_scope, - self.ty.data_alignment, - self.ty.offset_factor, - self.ty.layout.storage() if layout is None else layout, - ) + return _buffer_view.local(self, *shape, layout=layout) def permute(self, *dims) -> "Buffer": """Permute the dimensions of the buffer. @@ -464,28 +365,94 @@ def permute(self, *dims) -> "Buffer": permuted : DeclBufferFrame The buffer with permuted dimensions. """ - new_shape = [self.ty.shape[d] for d in dims] - # Permute *logical* dims, not the layout's fine-grained shard iters: a - # tcgen05/atom layout maps several shard iters to each logical axis, so - # group by the current shape first and permute whole groups. ``group`` - # returns a regrouped layout (degenerate extent-1 iters folded away) - # plus seps over *that* layout — permute the regrouped one, not - # ``self.layout``. For a simple layout (one shard iter per axis) this - # reduces to ``permute_dims(dims)``. - grouped, seps = self.ty.layout.group(list(self.ty.shape)) - new_layout = grouped.permute_by_groups(seps, list(dims)) - return tvm.tirx.script.builder.decl_buffer( - new_shape, - self.ty.dtype, - buffer_data(self), - self.ty.strides, - self.ty.elem_offset, - None, - self.ty.storage_scope, - self.ty.data_alignment, - self.ty.offset_factor, - new_layout, - ) + return _buffer_view.permute(self, *dims) + + def rearrange(self, pattern: str = _REARRANGE_PATTERN_UNSET, /, **sizes) -> "Buffer": + """einops-style relayout in one line: ``buf.rearrange("b (2 r) -> 2 b r")``. + + A pure reshape+permute+reshape over the SAME physical bytes, spelled as + an einops pattern. Lowers to ``view`` (split lhs groups) → ``permute`` + (reorder to rhs atom order) → ``view`` (merge rhs groups), so it inherits + whatever the underlying axis machinery does: a plain (unswizzled) buffer + collapses to a flat layout, a swizzled buffer keeps its swizzle, + and a tmem buffer carries ``allocated_addr`` through. It therefore does + NOT flatten a swizzle atom — the same pattern on a swizzled SMEM buffer + vs an unswizzled TMEM buffer legitimately yields different physical + layouts (that is the point: rearrange acts on the operand, not a string). + + ``pattern`` is ``"lhs -> rhs"``; each side is space-separated axis names, + with ``(a b)`` grouping a product axis. Every lhs group's product must + equal that input dim; at most one axis per group may be unknown (inferred + from the dim), the rest supplied via ``**sizes``. Cannot express a + replica (``R[...]``), a stride-fiction/padded view, or a reshape crossing + a swizzle-atom boundary — keep those as explicit ``view(layout=...)``. + """ + if pattern is _REARRANGE_PATTERN_UNSET: + if "pattern" not in sizes: + raise TypeError("Buffer.rearrange() missing required argument: 'pattern'") + pattern = sizes.pop("pattern") + return _buffer_view.rearrange(self, pattern, **sizes) + + @property + def sub(self) -> "_buffer_view.SubIndexer": + """Numpy-style view indexer: ``buf.sub[2, 4:8, ::4]``. + + Unlike plain ``buf[...]`` (BufferLoad for scalar indices, extent-1 + BufferRegion dims for tile-primitive operands), ``sub`` follows numpy + basic-indexing semantics as a *view constructor*: an integer index + removes the dim (``select``), ``a:b`` narrows it, and ``a::s`` takes + every s-th element (requires the extent divisible by ``s`` and + ``a < s``). Trailing dims are kept whole. + """ + return _buffer_view.sub(self) + + def tile(self, *specs) -> "_buffer_view.TileIndexer": + """Chunk a dim: split it into factors, pick a chunk, keep the rest. + + Rank-preserving — the picked dim's remaining factors merge back into + that one dim, and every other dim is untouched, so N dims in gives N + dims out. Chunk multiple dims by chaining (dims never shift): + ``buf.tile(0, (nx, -1))[cx, :].tile(1, (-1, ny))[:, cy]``. + + Call as ``tile(dim, factors)`` for one dim, or pass several + ``(dim, factors)`` specs as sugar for a chain. ``factors`` is the + tuple the dim splits into (row-major, like :meth:`unflatten`; one + ``-1`` inferred). The indexer takes one entry per factor: an ``int`` / + ``Expr`` **picks** it (fixing the chunk, dropping the axis, folding + its offset) and ``:`` **keeps** it. At least one factor per dim must be + picked — a pure keep-everything split is :meth:`unflatten`, not a + chunk:: + + # 64 rows split into (stripe, warp, row) = (-1, WARPS, 4); this + # warp's 16 interleaved rows (stripe x row merged): + buf.tile(1, (-1, WARPS, 4))[:, warp, :] + + tile(d, (n, -1))[c, :] # contiguous block c + tile(d, (-1, n))[:, c] # round-robin chunk c + + A picked index may be a dynamic Expr (e.g. a warp id); picking + several factors of one dim is allowed. + """ + return _buffer_view.tile(self, *specs) + + def chunk(self, spec) -> "_buffer_view.ChunkIndexer": + """Split dims into equal contiguous chunks and pick a chunk per dim — + **rank-preserving**. Index the result with ``[picks]``. + + ``spec`` is a per-dim tuple (length = rank). Each entry is ``None`` + (leave the dim) or a positive int ``n`` (split that dim, extent ``E`` + with ``E % n == 0``, into ``n`` equal chunks of ``E // n``). Then + ``chunk(spec)[picks]`` takes one entry per dim: a chunked dim's pick is + the chunk index (int / Expr) and **narrows that dim** to the chunk's + ``[c*E//n : (c+1)*E//n)`` range — the dim is kept at ``E // n``, no + dimension is added; an unchunked dim's pick is a normal index (``:`` / + int / slice). The result is the *same BufferRegion* as the hand-written + slice — one line instead of the ``c*k : (c+1)*k`` arithmetic:: + + X[.., c * k : (c + 1) * k, ..] # before (k = E // n) + X.chunk((None, .., n, ..))[.., c, ..] # after (k inferred) + """ + return _buffer_view.chunk(self, spec) def __getitem__(self, indices): if not is_buffer_var(self): diff --git a/python/tvm/tirx/compilation_pipeline.py b/python/tvm/tirx/compilation_pipeline.py index 23dee416bbe2..d18186218c03 100644 --- a/python/tvm/tirx/compilation_pipeline.py +++ b/python/tvm/tirx/compilation_pipeline.py @@ -20,6 +20,7 @@ import tvm from tvm import tirx +from tvm.backend.cuda import transforms as cuda_transforms def default_tir_pipeline(): @@ -49,6 +50,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I tirx.transform.VerifyMemory(), tirx.transform.AnnotateEntryFunc(), tirx.transform.SplitHostDevice(), + cuda_transforms.LowerIket(), tirx.transform.MakePackedAPI(), tirx.transform.FP8StorageLegalize(), tirx.transform.BF16StorageLegalize(), @@ -88,6 +90,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I tirx.transform.VerifyMemory(), tirx.transform.AnnotateEntryFunc(), tirx.transform.SplitHostDevice(), + cuda_transforms.LowerIket(), tirx.transform.MakePackedAPI(), tirx.transform.FP8StorageLegalize(), tirx.transform.BF16StorageLegalize(), diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py index 3dcef98bc459..49b6ebbe6dbd 100644 --- a/python/tvm/tirx/layout.py +++ b/python/tvm/tirx/layout.py @@ -343,8 +343,8 @@ def _get_default_strides(data: list[int | Expr], stride: int = 1) -> tuple: return list(reversed(res)) def is_swizzle(self) -> bool: - """Check if the layout is swizzle.""" - return isinstance(self, SwizzleLayout) + """Check if the layout is a bare swizzle (ComposeLayout over a trivial tile).""" + return isinstance(self, ComposeLayout) and self.tile_layout.is_trivial() def is_trivial(self) -> bool: """Check if the layout is trivial.""" @@ -364,10 +364,14 @@ def storage(self) -> "Layout": exclude = {axis: offset for axis, offset in self.offset.items() if not axis.is_thread()} return TileLayout.from_iters(shard, replicate, exclude) # pylint: disable=no-member - elif isinstance(self, SwizzleLayout): - return self elif isinstance(self, ComposeLayout): - return ComposeLayout(self.swizzle.storage(), self.tile_layout.storage()) + return ComposeLayout( + self.per_element, + self.swizzle_len, + self.atom_len, + self.tile_layout.storage(), + self.swizzle_inner, + ) else: raise ValueError(f"Unsupported layout type: {type(self)}") @@ -388,16 +392,15 @@ def unpack(self, num: int) -> "Layout": shard = [Iter(iter.extent, iter.stride * num, iter.axis) for iter in self.shard] shard.append(Iter(num, 1, Axis.get("m"))) return TileLayout.from_iters(shard, self.replica, self.offset) - elif isinstance(self, SwizzleLayout): + elif isinstance(self, ComposeLayout): assert num & (num - 1) == 0, "num must be a power of 2" - return SwizzleLayout( + return ComposeLayout( self.per_element + (num.bit_length() - 1), self.swizzle_len, self.atom_len, + self.tile_layout.unpack(num), self.swizzle_inner, ) - elif isinstance(self, ComposeLayout): - return ComposeLayout(self.swizzle.unpack(num), self.tile_layout.unpack(num)) else: raise ValueError(f"Unsupported layout type: {type(self)}") @@ -421,7 +424,13 @@ def broadcast(self, num: int, position: int = -1, axis: '"Axis" | str' = "m") -> shard.insert(insert_at, new_iter) return TileLayout.from_iters(shard, self.replica, self.offset) elif isinstance(self, ComposeLayout): - return ComposeLayout(self.swizzle, self.tile_layout.broadcast(num, position, axis)) + return ComposeLayout( + self.per_element, + self.swizzle_len, + self.atom_len, + self.tile_layout.broadcast(num, position, axis), + self.swizzle_inner, + ) else: raise ValueError(f"broadcast not supported for {type(self)}") @@ -448,19 +457,18 @@ def pack(self, num: int) -> "Layout": shard = [Iter(iter.extent, iter.stride // num, iter.axis) for iter in self.shard[:-1]] shard.append(Iter(inner_iter.extent // num, 1, inner_iter.axis)) return TileLayout.from_iters(shard, self.replica, self.offset) - elif isinstance(self, SwizzleLayout): + elif isinstance(self, ComposeLayout): assert num & (num - 1) == 0, "num must be a power of 2" assert self.per_element >= num.bit_length() - 1, ( "per_element must be greater than or equal to num.bit_length() - 1" ) - return SwizzleLayout( + return ComposeLayout( self.per_element - (num.bit_length() - 1), self.swizzle_len, self.atom_len, + self.tile_layout.pack(num), self.swizzle_inner, ) - elif isinstance(self, ComposeLayout): - return ComposeLayout(self.swizzle.pack(num), self.tile_layout.pack(num)) else: raise ValueError(f"Unsupported layout type: {type(self)}") @@ -564,7 +572,12 @@ def __getattr__(name): __all__ = [] # type: ignore[var-annotated] __all__ += list(_AXIS_NAMES) __all__ += ["R", "S"] -__all__ += ["tcgen05_atom_layout", "tmem_datapath_layout", "wg_local_layout"] +__all__ += [ + "tcgen05_atom_layout", + "tmem_datapath_layout", + "tmem_mma_operand_layout", + "wg_local_layout", +] # ============================================================================ @@ -585,11 +598,10 @@ def __getattr__(name): # scatter directly. # # We surface this via the factory below. Callers pass the datapath letter -# (``"D"`` / ``"F"`` / ``"B"``) and the logical ``(rows, cols)``; the -# factory returns the appropriate TileLayout. ``tmem_pool.alloc(..., -# datapath="F")`` plumbs this into the buffer's layout so the dispatch can -# structurally verify atom ↔ datapath compatibility instead of silently -# accepting mismatches. +# (``"D"`` / ``"F"``) and the logical ``(rows, cols)``; the factory returns +# the appropriate TileLayout. ``tmem_pool.alloc(..., layout=...)`` plumbs +# this into the buffer's layout so the dispatch can structurally verify +# atom ↔ datapath compatibility instead of silently accepting mismatches. # # Supported today: # - ``"D"``: M=128, ``.cta_group::1``, full datapath. Identity row→lane. @@ -605,7 +617,7 @@ def __getattr__(name): # Layouts A / C / E / G are reserved for future expansion. -_TMEM_DATAPATH_ROWS = {"D": 128, "F": 64, "B": 64} +_TMEM_DATAPATH_ROWS = {"A": 128, "B": 64, "C": 64, "D": 128, "E": 64, "F": 64, "G": 32} def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) -> "TileLayout": @@ -620,13 +632,24 @@ def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) Parameters ---------- datapath : str - One of ``"D"`` (M=128, ``.cta_group::1``, full datapath), ``"F"`` - (M=64, non-``.ws``, half datapath), or ``"B"`` (per-CTA M=64, - ``.cta_group::2``, "2x2" datapath). Other layouts are not yet - supported by this factory. + PTX data path layout letter (all cta_group::2 layouts are the + per-CTA view of the M-rows this CTA of the pair owns): + + - ``"A"``: M=256, ``.cta_group::2`` — identity over 128 rows. + - ``"B"``: M=128, ``.cta_group::2``, dense A — 64 rows, column + halves folded across the two 64-lane halves. + - ``"C"``: M=128, ``.cta_group::2``, sparse A — 64 rows scattered + 16-per-warp, half datapath (same organization as F). + - ``"D"``: M=128, ``.cta_group::1`` — identity, full datapath. + - ``"E"``: M=64, ``.ws`` — 64 rows, column halves folded (same + organization as B). + - ``"F"``: M=64, non-``.ws`` — 64 rows scattered 16-per-warp, + half datapath. + - ``"G"``: M=32, ``.ws`` — 32 rows, column quarters folded across + the four 32-lane warp slabs. rows : int Logical row count of the TMEM buffer. Must match the datapath's M - dimension: 128 for D, 64 for F and B. + dimension: 128 for A/D, 64 for B/C/E/F, 32 for G. cols : int Logical column count. Datapath B requires an even count because its columns split into two equal lane halves. @@ -656,33 +679,45 @@ def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) raise ValueError(f"tmem_datapath_layout: sub_slab must be 0 or 1, got {sub_slab}") tlane = Axis.get("TLane") tcol = Axis.get("TCol") - if datapath == "D": - # M=128, identity row→lane: row r ∈ [0, 128) → physical lane r. D - # already spans both 16-lane sub-slabs of every warp partition. + if datapath in ("A", "D"): + # Identity row→lane (r → lane r). D is the full cta_group::1 datapath; + # A is the per-CTA view of M=256/cta_group::2 (this CTA's 128 M-rows). + # Both already span both 16-lane sub-slabs of every warp partition. if sub_slab != 0: raise ValueError( - "tmem_datapath_layout: datapath='D' (M=128) already spans both " + f"tmem_datapath_layout: datapath={datapath!r} (M=128) already spans both " "sub-slabs; sub_slab must be 0" ) return TileLayout(S[(rows, cols) : (1 @ tlane, 1 @ tcol)]) - if datapath == "B": - # Layout B: a per-CTA (64, N) accumulator uses the same physical - # footprint as a Layout D (128, N/2) accumulator. The high column bit - # selects one of the two 64-lane halves, while the low bits select TCol. - # B always spans all 128 lanes, so the F-only sub_slab selector does not - # apply. + if datapath == "G": + # Layout G (M=32, .ws): quarter datapath — column space splits into 4 + # quarters, lane = row + 32*q, physical col = col % (cols/4). + if sub_slab != 0: + raise ValueError( + "tmem_datapath_layout: datapath='G' folds column quarters across " + "full 32-lane slabs; sub_slab must be 0" + ) + if cols % 4 != 0: + raise ValueError( + f"tmem_datapath_layout: datapath='G' requires cols % 4 == 0, got {cols}" + ) + return TileLayout(S[(rows, 4, cols // 4) : (1 @ tlane, 32 @ tlane, 1 @ tcol)]) + if datapath in ("B", "E"): + # Layouts B (M=128 cta_group::2 dense) and E (M=64 .ws): half datapath, + # lane = row + 64 * (col >= cols/2), physical col = col % (cols/2). if sub_slab != 0: raise ValueError( - "tmem_datapath_layout: datapath='B' spans all 128 lanes; sub_slab must be 0" + f"tmem_datapath_layout: datapath={datapath!r} folds column halves across " + "full 64-lane halves; sub_slab must be 0" ) if cols % 2 != 0: raise ValueError( - f"tmem_datapath_layout: datapath='B' expects even cols (N), got {cols}" + f"tmem_datapath_layout: datapath={datapath!r} requires even cols, got {cols}" ) - n_half = cols // 2 - return TileLayout(S[(rows, 2, n_half) : (1 @ tlane, 64 @ tlane, 1 @ tcol)]) - # Layout F: M=64 scattered. Logical row r = wid * 16 + intra (wid ∈ [0,4), - # intra ∈ [0,16)) → physical lane wid * 32 + sub_slab * 16 + intra. + return TileLayout(S[(rows, 2, cols // 2) : (1 @ tlane, 64 @ tlane, 1 @ tcol)]) + # Layouts C (M=128 cta_group::2 sparse) and F (M=64 non-.ws): rows-per-CTA + # scattered, half datapath — logical row wid*16+intra → physical lane + # wid*32 + sub_slab*16 + intra. # ``TileLayout`` decomposes a scalar row index via ``SplitCoord`` # (src/tirx/ir/layout/utils.cc), which uses row-major ordering: with # shape ``(s0, s1)`` the FIRST iter receives ``coord // s1`` (the high @@ -696,6 +731,149 @@ def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) return TileLayout(spec) +def _mma_datapath_letter(M, cta_group, ws=False, sparse=False): + """Map (instruction ``M``, ``cta_group``, ``ws``, ``sparse``) to the PTX + tcgen05 datapath letter. Raises for families rejected in the v1 semantic + allocator API (sparse Layout C, M=32 Layout G). See + ``localdoc/claude_plan.txt`` S3a.""" + if sparse: + raise ValueError( + f"alloc_tcgen05_mma_AB: sparse A (Layout C) not supported in v1 " + f"(M={M}, cta_group={cta_group})" + ) + if ws and not (cta_group == 1 and M == 64): + raise ValueError( + f"alloc_tcgen05_mma_AB: ws=True only valid for cta_group=1, M=64 (Layout E); " + f"got M={M}, cta_group={cta_group}" + ) + if cta_group == 1: + if M == 128: + return "D" + if M == 64: + return "E" if ws else "F" + elif cta_group == 2: + if M == 256: + return "A" + if M == 128: + return "B" + raise ValueError( + f"alloc_tcgen05_mma_AB: no supported datapath for M={M}, cta_group={cta_group}, " + f"ws={ws} (supported: cta1 M in {{64,128}}, cta2 M in {{128,256}})" + ) + + +def tmem_mma_operand_layout( + operand, shape, dtype, *, M, cta_group, ws=False, sparse=False, group=None +): + """Resolve the TMEM ``TileLayout`` for a tcgen05 MMA operand from operand + role + instruction params. + + ``operand``: ``"A"`` (A-in-TMEM) or ``"D"`` (accumulator / C). + ``M`` is the PTX ``tcgen05.mma`` **instruction** M (256/128/64), NOT the + per-CTA buffer rows; ``per_cta_rows = M // cta_group``. Physical 32-bit + column count is left to the pool's ``_resolve_cols`` (layout-aware). + + Reproduces the hand-written layouts the three FlashMLA kernels use; the + supported/reject domain is spec'd in ``localdoc/claude_plan.txt`` S3a. + """ + tlane = Axis.get("TLane") + tcol = Axis.get("TCol") + datapath = _mma_datapath_letter(M, cta_group, ws, sparse) + per_cta_rows = M // cta_group + ext = [int(s) for s in shape] + + def _chk_rows(rows): + if rows != per_cta_rows: + raise ValueError( + f"alloc_mma_{operand}: M={M}, cta_group={cta_group} expects " + f"per-CTA rows {per_cta_rows}, got {rows} (shape {ext})" + ) + + if operand == "D": + if group is not None: + # flat buffer (m, N) + explicit grouping -> (m, s, 2, n) layout + # (codex plan option b): keeps kernel O buffer 2D, layout 4-axis. + if len(ext) != 2: + raise ValueError(f"alloc_tcgen05_mma_D group= requires 2D (m,N), got {ext}") + m, N = ext + gs = [int(g) for g in group] + if len(gs) != 3 or gs[1] != 2: + raise ValueError(f"alloc_tcgen05_mma_D group must be (s,2,n), got {group}") + s2, _, n2 = gs + if s2 * 2 * n2 != N: + raise ValueError(f"alloc_tcgen05_mma_D group {group} product != N={N}") + if datapath not in ("B", "E"): + raise ValueError(f"alloc_tcgen05_mma_D group= only for Layout B/E, got {datapath}") + _chk_rows(m) + return TileLayout( + S[(m, s2, 2, n2) : (1 @ tlane, n2 @ tcol, 64 @ tlane, 1 @ tcol)] + ).canonicalize() + # Grouped grammar (C3): (m,N) | (m,2,N/2) | (m,s,2,n); plus batched + # C[2,m,N/2] normalizing to packed. bank axis is fixed by position. + if len(ext) == 2: + m, N = ext + _chk_rows(m) + layout = tmem_datapath_layout(datapath, m, N) + elif len(ext) == 3 and ext[0] == 2 and datapath in ("B", "E"): + _, m, Nh = ext # batched C[2, m, N/2] + _chk_rows(m) + layout = TileLayout(S[(2, m, Nh) : (64 @ tlane, 1 @ tlane, 1 @ tcol)]) + elif len(ext) == 3 and ext[1] == 2 and datapath in ("B", "E"): + m, _, Nh = ext # explicit (m, 2, N/2) + _chk_rows(m) + layout = TileLayout(S[(m, 2, Nh) : (1 @ tlane, 64 @ tlane, 1 @ tcol)]) + elif len(ext) == 4 and ext[2] == 2 and datapath in ("B", "E"): + m, s, _, n = ext # column-segmented (m, s, 2, n): flatten N=s*2*n + _chk_rows(m) + layout = TileLayout(S[(m, s, 2, n) : (1 @ tlane, n @ tcol, 64 @ tlane, 1 @ tcol)]) + else: + raise ValueError( + f"alloc_tcgen05_mma_D: unsupported D shape {ext} for datapath {datapath}" + ) + elif operand == "A": + if datapath in ("A", "D"): + if len(ext) != 2: + raise ValueError( + f"alloc_tcgen05_mma_A: datapath {datapath} identity needs 2D (m,K), got {ext}" + ) + m, K = ext + _chk_rows(m) + layout = tmem_datapath_layout(datapath, m, K) # identity + elif datapath == "B": + if len(ext) == 2: # replica A + m, K = ext + _chk_rows(m) + layout = TileLayout(S[(m, K) : (1 @ tlane, 1 @ tcol)] + R[2 : 64 @ tlane]) + elif len(ext) == 3 and ext[0] == 2: # bank-batched A + _, m, K = ext + _chk_rows(m) + layout = TileLayout(S[(2, m, K) : (64 @ tlane, 1 @ tlane, 1 @ tcol)]) + else: + raise ValueError( + f"alloc_tcgen05_mma_A: datapath B needs (64,K) replica or (2,64,K) " + f"bank-batched, got {ext}" + ) + elif datapath == "E": # M=64 .ws: bank-batched only, flat rejected + if len(ext) == 3 and ext[0] == 2: + _, m, K = ext + _chk_rows(m) + layout = TileLayout(S[(2, m, K) : (64 @ tlane, 1 @ tlane, 1 @ tcol)]) + else: + raise ValueError( + f"alloc_tcgen05_mma_A: datapath E (M=64 .ws) requires bank-batched " + f"(2,64,K); flat {ext} rejected (hidden second-half read)" + ) + else: # F (M=64 non-.ws): A occupies lanes 0..63 identically + if len(ext) != 2: + raise ValueError(f"alloc_tcgen05_mma_A: datapath F needs 2D (m,K), got {ext}") + m, K = ext + _chk_rows(m) + layout = TileLayout(S[(m, K) : (1 @ tlane, 1 @ tcol)]) + else: + raise ValueError(f"tmem_mma_operand_layout: operand must be 'A' or 'D', got {operand!r}") + return layout.canonicalize() + + def wg_local_layout(cols, rows=128): """Return a warpgroup-local register layout. @@ -1056,14 +1234,14 @@ def __add__(self, other): replica=other.replica if other.replica else self.replica, offset=_merge_offset(self.offset, other.offset), ) - if isinstance(other, _OnAxis | _OffsetExpr | int): + if isinstance(other, _OnAxis | _OffsetExpr | Expr | int): return _LayoutSpec( shard=self.shard, replica=self.replica, offset=_merge_offset(self.offset, other) ) return NotImplemented def __radd__(self, other): - if isinstance(other, _OnAxis | _OffsetExpr | int): + if isinstance(other, _OnAxis | _OffsetExpr | Expr | int): return _LayoutSpec( shard=self.shard, replica=self.replica, offset=_merge_offset(other, self.offset) ) @@ -1223,6 +1401,28 @@ def group(self, shape: list[Expr]) -> tuple["Layout", list[int]]: """ return _ffi_api.TileLayoutGroup(self, shape) # pylint: disable=no-member + def group_many(self, shapes: Sequence[Sequence[Expr]]) -> tuple["TileLayout", list[list[int]]]: + """Group the layout by the minimal common refinement of several shapes. + + Repeated cumulative product boundaries are retained, so an extent-one + dimension in any input shape becomes a real unit iterator in the + refined layout. This operation only splits existing shard iterators; + it does not canonicalize or reorder the layout. + + Parameters + ---------- + shapes : Sequence[Sequence[Expr]] + Logical shapes with provably equal total products. + + Returns + ------- + Tuple[TileLayout, List[List[int]]] + The commonly refined layout and one separator list per input shape. + """ + return _ffi_api.TileLayoutGroupMany( # pylint: disable=no-member + self, [list(shape) for shape in shapes] + ) + def get_scope(self) -> tuple[ExecScope, ExecScope] | None: """Get the scope pair of the layout.""" return _ffi_api.TileLayoutGetScope(self) # pylint: disable=no-member @@ -1340,34 +1540,35 @@ def permute_by_groups(self, seps: list[int], perm: list[int]) -> "TileLayout": return self.permute_dims(flat) -@tvm_ffi.register_object("tirx.SwizzleLayout") -class SwizzleLayout(Layout): - """A memory layout that swizzles elements to improve memory access patterns.""" +@tvm_ffi.register_object("tirx.ComposeLayout") +class ComposeLayout(Layout): + """A memory layout that swizzles a tile layout. + + ``per_element`` / ``swizzle_len`` / ``atom_len`` / ``swizzle_inner`` carry the + swizzle (formerly the standalone ``SwizzleLayout``); ``tile_layout`` is the + tiled memory map the swizzle is applied to. A bare swizzle is a + ``ComposeLayout`` over a trivial identity tile. + """ per_element: int swizzle_len: int atom_len: int swizzle_inner: bool + tile_layout: "TileLayout" def __init__( - self, per_element: int, swizzle_len: int, atom_len: int, swizzle_inner: bool = True + self, + per_element: int, + swizzle_len: int, + atom_len: int, + tile_layout: "TileLayout", + swizzle_inner: bool = True, ): self.__init_handle_by_constructor__( - _ffi_api.SwizzleLayout, # pylint: disable=no-member + _ffi_api.ComposeLayout, # pylint: disable=no-member per_element, swizzle_len, atom_len, + tile_layout, swizzle_inner, ) - - -@tvm_ffi.register_object("tirx.ComposeLayout") -class ComposeLayout(Layout): - """A memory layout that composes 2 layouts.""" - - def __init__(self, layout_A: "SwizzleLayout", layout_B: "TileLayout"): - self.__init_handle_by_constructor__( - _ffi_api.ComposeLayout, # pylint: disable=no-member - layout_A, - layout_B, - ) diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 7812377a2056..5b142760bbcc 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -24,7 +24,7 @@ import tvm from tvm import tirx -from tvm.ir import Call, Expr, Op, PointerType +from tvm.ir import Call, Expr, Op, PointerType, PrimType from tvm.ir.base import Span from tvm.ir.type import TensorMapType from tvm.runtime import const @@ -958,7 +958,7 @@ def tvm_access_ptr(ptype, data, offset, extent, rw_mask): call : Expr The call expression. """ - if isinstance(ptype, str | tvm.ir.PrimType): + if isinstance(ptype, str | PrimType): ptype = type_annotation(ptype) data_type = getattr(data, "ty", None) storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global" @@ -979,7 +979,7 @@ def ptr_byte_offset(data, byte_offset, dtype): ``byte_offset`` is always in bytes. Use this when the source CUDA shape needs an explicitly typed local pointer derived from a byte-addressed base. """ - if isinstance(dtype, str | tvm.ir.PrimType): + if isinstance(dtype, str | PrimType): dtype = type_annotation(dtype) data_type = getattr(data, "ty", None) storage_scope = data_type.storage_scope if isinstance(data_type, PointerType) else "global" diff --git a/python/tvm/tirx/operator/intrinsics/_common.py b/python/tvm/tirx/operator/intrinsics/_common.py index f85fefc4c64d..a96bc8caf3a8 100644 --- a/python/tvm/tirx/operator/intrinsics/_common.py +++ b/python/tvm/tirx/operator/intrinsics/_common.py @@ -29,6 +29,12 @@ FENCE_SCOPE = ("cta", "cluster", "gpu", "sys") FENCE_PROXY_ASYNC_SPACE = ("", "global", "shared::cta", "shared::cluster") CLUSTER_BARRIER_SEM = ("", "release", "relaxed") +MBARRIER_COMPLETE_TX_SEM = ("relaxed",) +MBARRIER_COMPLETE_TX_SCOPE = ("cta", "cluster") +MBARRIER_COMPLETE_TX_SPACE = ("shared", "shared::cta", "shared::cluster") +MBARRIER_ARRIVE_SEM = ("", "release", "relaxed") +MBARRIER_ARRIVE_SCOPE = ("", "cta", "cluster") +MBARRIER_ARRIVE_SPACE = ("shared", "shared::cta", "shared::cluster") # CTA group (used by tcgen05 and TMA) ----------------------------------------- TCGEN05_CTA_GROUP = (1, 2) diff --git a/python/tvm/tirx/operator/tile_primitive/__init__.py b/python/tvm/tirx/operator/tile_primitive/__init__.py index 4f9aa93d0200..e8e21c3c52d2 100644 --- a/python/tvm/tirx/operator/tile_primitive/__init__.py +++ b/python/tvm/tirx/operator/tile_primitive/__init__.py @@ -25,6 +25,6 @@ # Dispatch infrastructure. Per-backend schedule registrations are loaded via # ``tvm.backend.load()``. from .dispatcher import fail, list_registered_schedules, predicate, register_dispatch -from .registry import DispatchContext +from ...tile_primitive import DispatchContext __all__ = ["DispatchContext", "fail", "list_registered_schedules", "predicate", "register_dispatch"] diff --git a/python/tvm/tirx/operator/tile_primitive/dispatch_context.py b/python/tvm/tirx/operator/tile_primitive/dispatch_context.py deleted file mode 100644 index d79b946f0270..000000000000 --- a/python/tvm/tirx/operator/tile_primitive/dispatch_context.py +++ /dev/null @@ -1,209 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -"""TIRx operator dispatch context.""" - -from tvm_ffi import register_object - -from tvm.ir import Range -from tvm.runtime import Object, Scriptable -from tvm.target import Target -from tvm.tirx import Buffer, IterVar, Stmt, Var, _ffi_api -from tvm.tirx.exec_scope import ExecScope - - -@register_object("tirx.DispatchContext") -class DispatchContext(Object, Scriptable): - """DispatchContext node. - - Parameters - ---------- - target : Target - The target of the dispatch context. - - exec_scope : ExecScope - The execution scope of the dispatch context. - - launch_params : Dict[str, Expr] - The launch parameters of the dispatch context. - - var_range_map : Dict[Var, Range] - A map from loop variables to their ranges. - - callbacks : Dict[str, Object] - The callbacks of the dispatch context. - - shared_state : Dict[str, Object] - Shared state persisting across dispatch calls within a single lowering pass. - """ - - target: Target - exec_scope: ExecScope - launch_params: dict[str, IterVar] - var_range_map: dict[Var, Range] - alloc_only: bool - callbacks: dict[str, Object] - shared_state: dict[str, Object] - inter: dict[str, list] - intra: dict[str, list] - scope_kind: str - - kPrivateAlloc = "private_alloc" - kDeviceInitStmt = "device_init_stmt" - kHostInitStmt = "host_init_stmt" - kPostBufferDefStmt = "post_buffer_def_stmt" - - def __init__( - self, - target: Target, - exec_scope: ExecScope, - launch_params: dict[str, IterVar], - var_range_map: dict[Var, Range], - alloc_only: bool = False, - callbacks: dict[str, Object] = {}, - shared_state: dict[str, Object] = {}, - inter: dict[str, list] | None = None, - intra: dict[str, list] | None = None, - scope_kind: str = "", - ) -> None: - self.__init_handle_by_constructor__( - _ffi_api.DispatchContext, # pylint: disable=no-member - target, - exec_scope, - launch_params, - var_range_map, - alloc_only, - callbacks, - shared_state, - inter or {}, - intra or {}, - scope_kind, - ) - - def add_alloc_buffer(self, buffer: Buffer) -> None: - """Add an allocated buffer to the dispatch context. - Can be called only if alloc_only is True. - The buffer will be added to the workspace of operator (the key in the workspace is the buffer name). - - Parameters - ---------- - buffer : Buffer - The buffer to be added. - """ # noqa: E501 - _ffi_api.DispatchContextAddAllocBuffer(self, buffer) # pylint: disable=no-member - - def add_init_stmt(self, stmt: Stmt, host: bool = False) -> None: - """Add an initialization statement to the dispatch context. - Device initialization statements is only allowed if alloc_only is True. - Host initialization statements will be ignored if alloc_only is True. - The statements will be added to the beginning of the kernel. - - Parameters - ---------- - stmt : Stmt - The initialization statement to be added. - host : bool - Whether the statement is a host statement. - If True, the statement will be added to the host code (before the kernel). - If False, the statement will be added to the kernel body (at the beginning of the kernel). - """ # noqa: E501 - _ffi_api.DispatchContextAddInitStmt(self, stmt, host) # pylint: disable=no-member - - def add_post_buffer_def_stmt(self, buffer: Buffer, stmt: Stmt) -> None: - """Add a statement to be inserted after a buffer's definition (DeclBuffer/AllocBuffer). - - Parameters - ---------- - buffer : Buffer - The buffer whose definition scope the statement should appear in. - stmt : Stmt - The statement to be inserted. - """ - _ffi_api.DispatchContextAddPostBufferDefStmt(self, buffer, stmt) # pylint: disable=no-member - - def cache_get(self, key: str) -> Object | None: - """Look up a cached value by key. - - Parameters - ---------- - key : str - Cache key (built by the caller from construction parameters). - - Returns - ------- - Optional[Object] - The cached value, or None on miss. - """ - return _ffi_api.DispatchContextSharedStateGet(self, key) - - def cache_set(self, key: str, value: Object) -> None: - """Store a value in the cross-dispatch cache. - - Parameters - ---------- - key : str - Cache key (built by the caller from construction parameters). - value : Object - The object to cache (e.g. a Buffer or Var). - """ - _ffi_api.DispatchContextSharedStateSet(self, key, value) - - def is_cuda(self) -> bool: - """Check if the target is CUDA.""" - return self.target.kind.name == "cuda" - - def is_trn(self) -> bool: - """Check if the target is Trainium.""" - return self.target.kind.name == "trn" - - def is_target(self, name: str) -> bool: - """Check if the target kind matches ``name``.""" - return self.target.kind.name == name - - # -- scope predicates ---------------------------------------------------- - # - # Each ``is_`` returns True iff the op site is at that scope kind. - # Backed by ``self.scope_kind``, which 1-1 maps to a canonical intra - # TileLayout shape: - # thread -> {} - # warp -> {laneid} - # warpgroup -> {laneid, wid_in_wg} - # cta -> {laneid, warpid} - # cluster -> {laneid, warpid, cta_id} - # - # Prefer these predicates over raw ``self.scope_kind == "..."`` comparisons - # so dispatchers that later need stricter intra/inter shape checks can - # tighten the predicate body without touching every call site. - - @property - def is_thread(self) -> bool: - return self.scope_kind == "thread" - - @property - def is_warp(self) -> bool: - return self.scope_kind == "warp" - - @property - def is_warpgroup(self) -> bool: - return self.scope_kind == "warpgroup" - - @property - def is_cta(self) -> bool: - return self.scope_kind == "cta" - - @property - def is_cluster(self) -> bool: - return self.scope_kind == "cluster" diff --git a/python/tvm/tirx/operator/tile_primitive/dispatcher.py b/python/tvm/tirx/operator/tile_primitive/dispatcher.py index 848951ade595..c25edeefcdc7 100644 --- a/python/tvm/tirx/operator/tile_primitive/dispatcher.py +++ b/python/tvm/tirx/operator/tile_primitive/dispatcher.py @@ -30,9 +30,7 @@ from tvm.ir import Op from tvm.tirx import PrimFunc from tvm.tirx.operator import get_tirx_op -from tvm.tirx.stmt import TilePrimitiveCall - -from .dispatch_context import DispatchContext +from tvm.tirx.tile_primitive import DispatchContext, TilePrimitiveCall class DispatchFail(RuntimeError): diff --git a/python/tvm/tirx/operator/tile_primitive/ops.py b/python/tvm/tirx/operator/tile_primitive/ops.py index 61a4312c820a..e4be09f45f73 100644 --- a/python/tvm/tirx/operator/tile_primitive/ops.py +++ b/python/tvm/tirx/operator/tile_primitive/ops.py @@ -19,7 +19,7 @@ from tvm.ir import Op from tvm.tirx import Expr -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def get_tirx_op(op_name: str): @@ -400,6 +400,15 @@ class Exp2(UnaryOpWithBiasScale): op = get_tirx_op("exp2") +class Log2(UnaryOpWithBiasScale): + """Compute base-2 logarithm of all elements in src and store to dst. + + If bias and scale are provided: dst = log2(src * scale + bias) + """ + + op = get_tirx_op("log2") + + class Select(BinaryOp): """Select elements from src1 or src2 based on the predicate. diff --git a/python/tvm/tirx/operator/tile_primitive/registry.py b/python/tvm/tirx/operator/tile_primitive/registry.py index c2f1d9d7f0d0..baf1f8cb7a63 100644 --- a/python/tvm/tirx/operator/tile_primitive/registry.py +++ b/python/tvm/tirx/operator/tile_primitive/registry.py @@ -23,8 +23,7 @@ from tvm_ffi import register_global_func -from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import DispatchContext, TilePrimitiveCall # Note: legacy `register_schedule` is intentionally removed. diff --git a/python/tvm/tirx/predicate.py b/python/tvm/tirx/predicate.py deleted file mode 100644 index cff1c1669c22..000000000000 --- a/python/tvm/tirx/predicate.py +++ /dev/null @@ -1,45 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# pylint: disable=no-member -"""Async structures for TIRX""" - -import inspect -from collections.abc import Callable - -from tvm_ffi import register_object - -from tvm.runtime import Object -from tvm.tirx import Expr, Var - -from . import _ffi_api - - -@register_object("tirx.Predicate") -class Predicate(Object): - """A predicate object for TIRX""" - - vars: list[Var] - pred: Expr - - def __init__(self, f_pred: Callable[..., Expr]): - vars = [Var(name, "int32") for name in inspect.signature(f_pred).parameters] - pred = f_pred(*vars) - self.__init_handle_by_constructor__(_ffi_api.Predicate, vars, pred) - - def apply(self, indices: list[Expr]) -> Expr: - """Apply the predicate to the given indices""" - return _ffi_api.PredicateApply(self, indices) diff --git a/python/tvm/tirx/script/__init__.py b/python/tvm/tirx/script/__init__.py index 8abbcda38781..aa4e53b6cff2 100644 --- a/python/tvm/tirx/script/__init__.py +++ b/python/tvm/tirx/script/__init__.py @@ -30,7 +30,7 @@ from .parser import macro except ImportError: macro = None -from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages +from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool from . import tile from .builder.ir import TensorMap, meta_class diff --git a/python/tvm/tirx/script/builder/__init__.py b/python/tvm/tirx/script/builder/__init__.py index 5cada7493a95..769d13c317e6 100644 --- a/python/tvm/tirx/script/builder/__init__.py +++ b/python/tvm/tirx/script/builder/__init__.py @@ -21,6 +21,6 @@ from .ir import boolean as bool # pylint: disable=redefined-builtin from .ir import buffer as Buffer from .utils import buffer_proxy, frame_scope, seq_scope -from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages +from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool from . import tirx as tile from .tirx import cluster, cta, thread, warp, warpgroup, wg diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index e6ebfba3daf7..a5dbfbc02aca 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -90,7 +90,6 @@ Layout, R, S, - SwizzleLayout, TileLayout, wg_local_layout, ) @@ -155,6 +154,9 @@ def _normalize_prim_type(dtype) -> ir.PrimType: ty = getattr(value, "ty", None) if isinstance(ty, ir.PrimType): return ty + type_annotation = getattr(value, "type_annotation", None) + if isinstance(type_annotation, ir.PrimType): + return type_annotation return ir.PrimType(dtype) @@ -3525,7 +3527,6 @@ def _register_tir_namespace_printer_names(): "R", "S", "ScopeIdDef", - "SwizzleLayout", "TensorMap", "TileLayout", "Var", diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index 5160cfb006c2..8ad43a96d992 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -19,13 +19,13 @@ import functools from collections.abc import Callable +import tvm import tvm.tirx.operator as tirx_op from tvm.ir import Op -from tvm.tirx import Buffer, BufferRegion, Expr, buffer_data, is_buffer_var +from tvm.tirx import Buffer, BufferRegion, Expr, LambdaExpr, buffer_data, is_buffer_var from tvm.tirx.exec_scope import _SCOPE_KIND_TO_NAME, ExecScope -from tvm.tirx.expr import FloatImm -from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool, TMEMStages -from tvm.tirx.predicate import Predicate +from tvm.tirx.expr import FloatImm, IntImm +from tvm.tirx.lang.alloc_pool import SMEMPool, TMEMPool from . import _ffi_api, frame from .ir import decl_buffer, meta_class @@ -489,6 +489,98 @@ def cast( ) +def _check_copy_regions_match(dst, src, config, name, dispatch=None): + """Enforce that plain copy operands cover identical regions. + + A plain copy is elementwise between two regions over one shared logical + iteration space; after dropping extent-1 dims the two region shapes must + be identical (same rank, same extent per dim). Unit dims are pure + padding; any real reshape is the caller's job via Buffer views + (unflatten/select/...), never the copy's. + + A region is buffer shape + slice: purely logical, independent of the + buffer's layout. The layout maps logical coords to physical addresses + (lanes, swizzle, broadcast) and is the right home for physical tiling; + baking that structure into a buffer's declared shape (e.g. declaring a + tmem tile ``(2, 32, C)`` instead of ``(64, C)`` with the lane split in + the layout) is what produces a spurious region mismatch. The fix is + always to declare the natural logical shape, never to exempt the copy. + + TMA is byte-oriented rather than elementwise. Its planners validate the + layout mapping, so the builder only proves equal payload bytes. A + ``tma_explicit`` gather4 source describes one row while the instruction + transfers the four rows named by ``gather4``. + """ + if dispatch in ("tma_auto", "tma_explicit"): + analyzer = None + + def _payload_bits(region): + nonlocal analyzer + value = tvm.DataType(region.buffer.dtype).bits + for axis in region.region: + value = value * axis.extent + if analyzer is None: + from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + + analyzer = Analyzer() + return analyzer.simplify(value) + + dst_bits = _payload_bits(dst) + src_bits = _payload_bits(src) + gather4 = config.get("gather4") + if gather4 is not None: + if dispatch != "tma_explicit": + raise ValueError("copy_async: gather4 is only supported by dispatch='tma_explicit'") + if not isinstance(gather4, list | tuple | tvm.ir.Array) or len(gather4) != 4: + raise ValueError("copy_async: gather4 must contain exactly four row coordinates") + if len(src.region) != 2: + raise ValueError("copy_async: gather4 requires a rank-2 global source") + row_extent = src.region[0].extent + if not analyzer.can_prove_equal(row_extent, 1): + raise ValueError("copy_async: gather4 global axis 0 must describe exactly one row") + src_bits = analyzer.simplify(src_bits * 4) + + if not analyzer.can_prove_equal(dst_bits, src_bits): + raise ValueError( + f"{name}: dst payload {dst_bits} bits != src payload {src_bits} bits; " + "TMA operands must transfer the same total number of bytes" + ) + return + + def _squeeze(region): + return [ + r.extent for r in region if not (isinstance(r.extent, IntImm) and r.extent.value == 1) + ] + + dst_extents = _squeeze(dst.region) + src_extents = _squeeze(src.region) + + def _fail(): + raise ValueError( + f"{name}: dst region shape {[str(e) for e in dst_extents]} != " + f"src region shape {[str(e) for e in src_extents]}; copy operands " + f"must cover identical regions - reshape via buffer views " + f"(unflatten/select/narrow/...) instead" + ) + + if len(dst_extents) != len(src_extents): + _fail() + analyzer = None + for d, s in zip(dst_extents, src_extents): + d_int = d.value if isinstance(d, IntImm) else None + s_int = s.value if isinstance(s, IntImm) else None + if d_int is not None and s_int is not None: + if d_int != s_int: + _fail() + continue + if analyzer is None: + from tvm.arith import Analyzer # pylint: disable=import-outside-toplevel + + analyzer = Analyzer() + if not analyzer.can_prove_equal(d, s): + _fail() + + @ScopedOp def copy( dst: BufferRegion | Buffer, @@ -516,6 +608,7 @@ def copy( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) + _check_copy_regions_match(dst, src, config, "copy", dispatch) return f_insert( tirx_op.Copy(dst, src, workspace=workspace, config=config, dispatch=dispatch, scope=scope) ) @@ -535,6 +628,7 @@ def copy_async( config = kwargs or {} dst = _to_region(dst) src = _to_region(src) + _check_copy_regions_match(dst, src, config, "copy_async", dispatch) return f_insert( tirx_op.CopyAsync( dst, src, workspace=workspace, config=config, dispatch=dispatch, scope=scope @@ -1180,6 +1274,46 @@ def exp2( ) +@ScopedOp +def log2( + dst: BufferRegion | Buffer, + src: BufferRegion | Buffer | None = None, + bias: BufferRegion | Buffer | FloatImm | None = None, + scale: FloatImm | None = None, + workspace: dict[str, Buffer] | None = None, + dispatch: str | None = None, + scope: ExecScope | None = None, + **kwargs, +): + """Compute base-2 logarithm of all elements in src and store to dst.""" + # Expression-form overload: ``log2(value)`` returns the underlying expression. + from tvm import tirx as _tirx + + if not _is_buffer_or_region(dst): + return _tirx.log2(dst) + if src is None: + src = dst + if workspace is None: + workspace = {} + config = kwargs or {} + dst = _to_region(dst) + src = _to_region(src) + if bias is not None and is_buffer_var(bias): + bias = _to_region(bias) + return f_insert( + tirx_op.Log2( + dst, + src, + bias, + scale, + workspace=workspace, + config=config, + dispatch=dispatch, + scope=scope, + ) + ) + + def compose_op( workspace: dict[str, Buffer] | None = None, dispatch: str | None = None, **kwargs ) -> frame.ComposeOpFrame: @@ -1513,7 +1647,7 @@ def select( dst: BufferRegion | Buffer, true_value: BufferRegion | Buffer | FloatImm, false_value: BufferRegion | Buffer | FloatImm, - pred: Predicate | Callable[..., Expr], + pred: LambdaExpr | Callable[..., Expr], scope: ExecScope | None = None, ): """Select between two values based on a predicate. @@ -1529,7 +1663,7 @@ def select( false_value : Union[BufferRegion, Buffer, FloatImm] The value to select if the predicate is false. - pred : Union[Predicate, Callable[..., Expr]] + pred : Union[LambdaExpr, Callable[..., Expr]] The predicate to evaluate. The callable should take the same number of arguments as the dimensions of the destination buffer. """ # noqa: E501 dst = _to_region(dst) @@ -1537,8 +1671,8 @@ def select( true_value = _to_region(true_value) if is_buffer_var(false_value): false_value = _to_region(false_value) - if not isinstance(pred, Predicate): - pred = Predicate(pred) + if not isinstance(pred, LambdaExpr): + pred = LambdaExpr(pred) return f_insert(tirx_op.Select(dst, true_value, false_value, pred, scope=scope)) @@ -1628,7 +1762,6 @@ def _to_region(b): "ScopeNamespace", "ScopedOp", "TMEMPool", - "TMEMStages", "add", "binary_chain", "binary_reduce", @@ -1645,6 +1778,7 @@ def _to_region(b): "fma", "gemm", "gemm_async", + "log2", "max", "maximum", "memset", diff --git a/python/tvm/tirx/script/builder/tmem_pool.py b/python/tvm/tirx/script/builder/tmem_pool.py index acd6278da3e6..58718f11d970 100644 --- a/python/tvm/tirx/script/builder/tmem_pool.py +++ b/python/tvm/tirx/script/builder/tmem_pool.py @@ -16,4 +16,4 @@ # under the License. """Re-export from canonical location.""" -from tvm.tirx.lang.alloc_pool import TMEMPool, TMEMStages # noqa: F401 +from tvm.tirx.lang.alloc_pool import TMEMPool # noqa: F401 diff --git a/python/tvm/tirx/script/parser/__init__.py b/python/tvm/tirx/script/parser/__init__.py index 7d59fc17545a..2b4301b0be31 100644 --- a/python/tvm/tirx/script/parser/__init__.py +++ b/python/tvm/tirx/script/parser/__init__.py @@ -24,7 +24,7 @@ from . import operation as _operation from . import parser as _parser -from .entry import Buffer, Ptr, constexpr +from .entry import Buffer, Optional, Ptr, constexpr if TYPE_CHECKING: # pylint: disable=invalid-name @@ -37,10 +37,10 @@ __all__ = _tir.__all__ + [ "Buffer", + "Optional", "Ptr", "SMEMPool", "TMEMPool", - "TMEMStages", "bool", "constexpr", "inline", diff --git a/python/tvm/tirx/script/parser/entry.py b/python/tvm/tirx/script/parser/entry.py index 9778a41794fd..d708aa621a5a 100644 --- a/python/tvm/tirx/script/parser/entry.py +++ b/python/tvm/tirx/script/parser/entry.py @@ -16,6 +16,7 @@ # under the License. """The entry point of TVM parser for tirx.""" +import ast import inspect from collections.abc import Callable from typing import Any @@ -212,7 +213,7 @@ def wrapper(*args, **kwargs): class TIRJit: - """Top-level kernel decorator with constexpr params + ``.specialize()``. + """Top-level kernel decorator with compile-time ``.specialize()`` params. Parses the function body lazily: parsing is deferred until ``.specialize()`` supplies concrete values for the params annotated as ``T.constexpr``. The @@ -240,7 +241,7 @@ def __init__( # Resolved closure vars (computed once; the function itself is the # capture point, so this never changes between specializations). self._closure_vars: dict[str, Any] = utils.inspect_function_capture(func) - # Detect which params are marked T.constexpr. With PEP 563 + # Detect which params are marked T.constexpr or T.Optional. With PEP 563 # (``from __future__ import annotations``), each annotation is a # string; we eval them one-by-one so a constexpr probe is not # blocked by sibling annotations that reference yet-undefined names @@ -250,8 +251,10 @@ def __init__( sig = inspect.signature(func) constexpr_names: set[str] = set() constexpr_defaults: dict[str, Any] = {} + optional_names: set[str] = set() for name, param in sig.parameters.items(): - ann = raw_anns.get(name) + raw_ann = raw_anns.get(name) + ann = raw_ann if isinstance(ann, str): try: ann = eval(ann, eval_globals) # pylint: disable=eval-used @@ -261,19 +264,23 @@ def __init__( constexpr_names.add(name) if param.default is not inspect.Parameter.empty: constexpr_defaults[name] = param.default + if isinstance(ann, _OptionalAnnotation) or _is_explicit_optional_annotation(raw_ann): + optional_names.add(name) self.constexpr_names: frozenset[str] = frozenset(constexpr_names) self.constexpr_defaults: dict[str, Any] = constexpr_defaults + self.optional_names: frozenset[str] = frozenset(optional_names) self._cache: dict[tuple, PrimFunc] = {} - def specialize(self, **constexpr_kwargs) -> PrimFunc: - """Build a concrete PrimFunc by binding the constexpr params. + def specialize(self, **specialization_kwargs) -> PrimFunc: + """Build a PrimFunc by binding constexprs and absent optional params. Parameters ---------- - **constexpr_kwargs - One value per ``T.constexpr``-annotated parameter. All such - parameters must be supplied; passing names that are not - constexpr-annotated is an error. + **specialization_kwargs + One value per ``T.constexpr``-annotated parameter. A + ``T.Optional`` parameter may additionally be supplied as ``None`` + to remove it from the resulting PrimFunc ABI. Omitting an + optional parameter keeps it as a normal runtime parameter. Returns ------- @@ -281,13 +288,29 @@ def specialize(self, **constexpr_kwargs) -> PrimFunc: A concrete TIRx PrimFunc, identical in type to the output of ``@T.prim_func``. """ - extra = constexpr_kwargs.keys() - self.constexpr_names + specializable_names = self.constexpr_names | self.optional_names + extra = specialization_kwargs.keys() - specializable_names if extra: raise TypeError( f"{self.func.__name__}.specialize() got unexpected arg(s): " - f"{sorted(extra)} (constexpr params are: {sorted(self.constexpr_names)})" + f"{sorted(extra)} (specializable params are: {sorted(specializable_names)})" ) - effective = {**self.constexpr_defaults, **constexpr_kwargs} + invalid_optional = { + name: specialization_kwargs[name] + for name in self.optional_names & specialization_kwargs.keys() + if specialization_kwargs[name] is not None + } + if invalid_optional: + raise TypeError( + f"{self.func.__name__}.specialize(): T.Optional parameters only accept None; " + f"pass actual tensors when calling the compiled kernel (got: {invalid_optional!r})" + ) + + supplied_constexprs = { + name: specialization_kwargs[name] + for name in self.constexpr_names & specialization_kwargs.keys() + } + effective = {**self.constexpr_defaults, **supplied_constexprs} missing = self.constexpr_names - effective.keys() if missing: raise TypeError( @@ -295,8 +318,9 @@ def specialize(self, **constexpr_kwargs) -> PrimFunc: f"(no default provided): {sorted(missing)}" ) + absent_params = {name: None for name in self.optional_names & specialization_kwargs.keys()} try: - cache_key = tuple(sorted(effective.items())) + cache_key = (tuple(sorted(effective.items())), tuple(sorted(absent_params))) cached = self._cache.get(cache_key) except TypeError as err: raise TypeError( @@ -312,6 +336,7 @@ def specialize(self, **constexpr_kwargs) -> PrimFunc: extra_vars, check_well_formed=self.check_well_formed, s_tir=self.is_stir, + absent_params=absent_params, ) setattr(prim_func, "__name__", self.func.__name__) self._cache[cache_key] = prim_func @@ -328,8 +353,9 @@ def jit( """Decorator: capture the kernel and defer parsing until ``.specialize()``. Use ``@T.jit`` (instead of ``@T.prim_func``) when the kernel takes - compile-time parameters annotated with ``T.constexpr``. The resulting - object exposes ``.specialize(**constexpr_kwargs)``, which returns a + compile-time parameters annotated with ``T.constexpr`` or runtime + parameters that may be removed with ``T.Optional``. The resulting object + exposes ``.specialize(**specialization_kwargs)``, which returns a ``tvm.tirx.PrimFunc``. Example:: @@ -346,6 +372,14 @@ def add( ... kernel = add.specialize(N=1024) # returns a PrimFunc + + @T.jit + def guarded(optional: T.Optional(T.handle), out: T.handle): + if optional is not None: + ... + + present = guarded.specialize() + absent = guarded.specialize(optional=None) """ def decorator_wrapper(func: Callable) -> TIRJit: @@ -511,6 +545,38 @@ def __ror__(self, other): return self +class _OptionalAnnotation: + """Wrapper around the runtime annotation of an optional JIT parameter.""" + + def __init__(self, annotation: Any) -> None: + self.annotation = annotation + + +class _OptionalProxy: + """Construct an explicit ``T.Optional(...)`` JIT parameter annotation.""" + + def __call__(self, annotation: Any) -> _OptionalAnnotation: + return _OptionalAnnotation(annotation) + + +def _is_explicit_optional_annotation(annotation: Any) -> bool: + """Recognize deferred ``T.Optional(...)`` without accepting typing unions.""" + + if not isinstance(annotation, str): + return False + try: + node = ast.parse(annotation, mode="eval").body + except SyntaxError: + return False + if not isinstance(node, ast.Call): + return False + func = node.func + return (isinstance(func, ast.Attribute) and func.attr == "Optional") or ( + isinstance(func, ast.Name) and func.id == "Optional" + ) + + Buffer = BufferProxy() # pylint: disable=invalid-name Ptr = PtrProxy() # pylint: disable=invalid-name constexpr = _ConstexprProxy() # pylint: disable=invalid-name +Optional = _OptionalProxy() # pylint: disable=invalid-name diff --git a/python/tvm/tirx/script/parser/parser.py b/python/tvm/tirx/script/parser/parser.py index 5004cdee2738..3e62e8b0559e 100644 --- a/python/tvm/tirx/script/parser/parser.py +++ b/python/tvm/tirx/script/parser/parser.py @@ -34,8 +34,8 @@ from tvm.tirx.script.builder.ir import name_meta_class_value from tvm.tirx.stmt import BufferRegion +from .entry import _OptionalAnnotation, inline from .entry import constexpr as _constexpr_sentinel -from .entry import inline def slice_buffer_from_region(br: BufferRegion) -> Buffer: @@ -228,20 +228,21 @@ def bind_assign_value(self: Parser, node: doc.expr, var_name: str, value: Any) - ann_var = T.Bind(value) IRBuilder.name(var_name, ann_var) return ann_var - if not tvm.ir.is_prim_expr(value): + if not tvm.ir.is_prim_expr(value) and not isinstance(value, Expr): + # Python scalar (int/float/bool) -> const prim expr value = tvm.tirx.const(value) - if not isinstance(value, tvm.tirx.StringImm): + if isinstance(value, tvm.tirx.StringImm) or not tvm.ir.is_prim_expr(value): + # StringImm or non-prim-expr (e.g. pointer Call): immutable Bind var + ann_var = tvm.tirx.Var(var_name, value.ty) + IRBuilder.name(var_name, ann_var) + T.Bind(value, var=ann_var) + return ann_var + else: # x = expr -> scalar (auto-typed from value) scalar = T.local_scalar(dtype=str(value.ty.dtype)) IRBuilder.name(var_name, scalar.scalar.buffer) T.buffer_store(scalar.scalar.buffer, value, [0]) return scalar.scalar - else: - # StringImm: x = expr -> immutable Bind var - ann_var = tvm.tirx.Var(var_name, value.ty) - IRBuilder.name(var_name, ann_var) - T.Bind(value, var=ann_var) - return ann_var def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool: @@ -259,6 +260,18 @@ def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: b return default +def _is_jit_function(node: doc.FunctionDef) -> bool: + """Return whether the parsed source function is decorated with ``T.jit``.""" + + for decorator in node.decorator_list: + target = decorator.func if isinstance(decorator, doc.Call) else decorator + if isinstance(target, doc.Attribute) and target.attr == "jit": + return True + if isinstance(target, doc.Name) and target.id == "jit": + return True + return False + + @dispatch.register(token="tirx", type_name="For") def visit_for(self: Parser, node: doc.For) -> None: """The for visiting method for tirx. @@ -271,7 +284,7 @@ def visit_for(self: Parser, node: doc.For) -> None: node : doc.For The doc AST for node. """ - # Intercept range() at AST level so it works with both Python ints and PrimExprs. + # Intercept range() at AST level so it works with both Python ints and Exprs. # In other contexts (e.g. list comprehensions), range remains Python's builtin. if ( isinstance(node.iter, doc.Call) @@ -643,16 +656,32 @@ def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: self.report_error(arg, "Type annotation required for function parameters.") try: ann = self.eval_expr(arg.annotation) - if ( - callable(ann) - and not isinstance(ann, Expr) - and ann is not _constexpr_sentinel - ): - ann = ann() except Exception: # pylint: disable=broad-except ann = func_annotation.get(arg.arg, None) if ann is None: raise + if isinstance(ann, _OptionalAnnotation): + if not _is_jit_function(node): + self.report_error( + arg.annotation, + "T.Optional parameter annotations are only supported by @T.jit", + ) + if arg.arg in self.absent_params: + self.var_table.add(arg.arg, None) + continue + ann = ann.annotation + elif arg.arg in self.absent_params: + self.report_error( + arg.annotation, + "Only T.Optional parameters may be absent, " + f"but {arg.arg!r} is not optional", + ) + if ( + callable(ann) + and not isinstance(ann, Expr) + and ann is not _constexpr_sentinel + ): + ann = ann() if ann is _constexpr_sentinel: # T.constexpr param: value was bound in extra_vars by # TIRJit.specialize() and lives in an outer var_table @@ -763,7 +792,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None: # different function Call representation. Convert to the TIR representation. T.evaluate(tvm.tirx.call_tir(res.op, *res.args)) else: - # Pointer-valued TIR calls are general Expr rather than PrimExpr, + # Pointer-valued TIR calls are general Expr rather than Expr, # but are still valid standalone Evaluate statements. T.evaluate(res) elif isinstance(res, str): @@ -792,29 +821,28 @@ def visit_if(self: Parser, node: doc.If) -> None: node : doc.If The doc AST if node. """ - with self.var_table.with_frame(): - predicate = self.eval_expr(node.test) - if tvm.ir.is_prim_expr(predicate) or isinstance(predicate, tvm.tirx.expr.ExprOp): - with T.If(self.eval_expr(node.test)): - with T.Then(): - with self.var_table.with_frame(): - self.visit_body(node.body) - if node.orelse: - with T.Else(): - with self.var_table.with_frame(): - self.visit_body(node.orelse) - elif isinstance(predicate, bool): - if predicate: + predicate = self.eval_expr(node.test) + if tvm.ir.is_prim_expr(predicate) or isinstance(predicate, tvm.tirx.expr.ExprOp): + with T.If(predicate): + with T.Then(): with self.var_table.with_frame(): self.visit_body(node.body) - elif node.orelse: - with self.var_table.with_frame(): - self.visit_body(node.orelse) - else: - self.report_error( - node.test, - f"If condition must be a boolean expression, but got {predicate}", - ) + if node.orelse: + with T.Else(): + with self.var_table.with_frame(): + self.visit_body(node.orelse) + elif isinstance(predicate, bool): + # Python ``if`` does not introduce a lexical scope. Bindings from the + # selected compile-time branch must therefore remain visible after it. + if predicate: + self.visit_body(node.body) + elif node.orelse: + self.visit_body(node.orelse) + else: + self.report_error( + node.test, + f"If condition must be a boolean expression, but got {predicate}", + ) @dispatch.register(token="tirx", type_name="Assert") diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py index 49ff2c0315dd..4aaf79113b86 100644 --- a/python/tvm/tirx/script/tile.py +++ b/python/tvm/tirx/script/tile.py @@ -44,7 +44,7 @@ def _validate_tile_call(op_name, args, kwargs): if op_name in {"cast", "max", "min", "permute_layout", "silu"}: src = _get_arg(args, kwargs, 1, "src") _require_buffer_arg(op_name, "src", src) - elif op_name in {"sqrt", "exp", "exp2", "reciprocal"}: + elif op_name in {"sqrt", "exp", "exp2", "log2", "reciprocal"}: src = _get_arg(args, kwargs, 1, "src") if src is not None: _require_buffer_arg(op_name, "src", src) @@ -70,6 +70,7 @@ def wrapper(*args, scope=None, **kwargs): "copy_async", "exp", "exp2", + "log2", "fdiv", "fill", "fma", @@ -96,6 +97,7 @@ def wrapper(*args, scope=None, **kwargs): for _op_name in _SCOPED_TILE_OP_NAMES: globals()[_op_name] = _tile_scoped_op(_op_name) + cluster = _builder.ScopeNamespace("cluster", "cluster") cta = _builder.ScopeNamespace("cta", "cta") wg = _builder.ScopeNamespace("warpgroup", "wg") @@ -105,6 +107,7 @@ def wrapper(*args, scope=None, **kwargs): compose_op = _builder.compose_op + __all__ = [ *_SCOPED_TILE_OP_NAMES, "cluster", diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 69fcf2c7542d..e68046774423 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -29,22 +29,19 @@ from collections.abc import Mapping from enum import IntEnum -from typing import TYPE_CHECKING, Any, ClassVar +from typing import Any import tvm_ffi -from tvm.ir import Expr, Op, Range, Span, is_prim_expr +from tvm.ir import Expr, Range, Span, is_prim_expr from tvm.runtime import Object, Scriptable, const -from tvm.tirx import FloatImm, IntImm +from tvm.tirx import IntImm from . import _ffi_api from .buffer import Buffer -from .exec_scope import ExecScope, ScopeIdDef +from .exec_scope import ScopeIdDef from .expr import IterVar, StringImm, Var -if TYPE_CHECKING: - from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext - @tvm_ffi.register_object("tirx.Stmt") class Stmt(Object, Scriptable): @@ -941,180 +938,3 @@ def stmt_list(stmt: Stmt) -> list[Stmt]: res += stmt_list(x) return res return [stmt] - - -def normalize_const_arg(arg) -> Expr: - if isinstance(arg, float): - return FloatImm("float32", arg) - return arg - - -@tvm_ffi.register_object("tirx.TilePrimitiveCall") -class TilePrimitiveCall(Stmt): - """TilePrimitiveCall node. - - Parameters - ---------- - op : Op - The operator. - - args : List[Expr] - The arguments. - - workspace : Map[str, Buffer] - The workspace. - - config : Map[str, ObjectRef] - The scheduler/config dictionary. - - dispatch : Optional[str] - The explicit variant name to dispatch to. - - scope : ExecScope - The cooperation scope of this call. Defaults to ``thread`` (an unscoped call). - """ - - args: list[Expr] - workspace: dict[str, Buffer] - config: dict[str, Any] - dispatch: str | None - scope: ExecScope - _registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {} - - def __init__( - self, - *args: list[Expr], - op: Op | None = None, - workspace: dict[str, Buffer] | None = None, - config: dict[str, Any] | None = None, - dispatch: str | None = None, - scope: ExecScope | None = None, - ) -> None: - if workspace is None: - workspace = {} - if config is None: - config = {} - if scope is None: - scope = ExecScope("thread") - if op is None: - assert self.__class__ != TilePrimitiveCall, ( - "Directly instantiating TilePrimitiveCall needs to specify the op" - ) - op = self.__class__.op - args = list(map(normalize_const_arg, args)) - self.__init_handle_by_constructor__( - _ffi_api.TilePrimitiveCall, - op, - args, - workspace, - config, - dispatch, - scope, # pylint: disable=no-member - ) - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - if hasattr(cls, "op"): - cls._registry[cls.op] = cls - - @classmethod - def downcast(cls, instance: "TilePrimitiveCall") -> "TilePrimitiveCall": - subclass = cls._registry.get(instance.op) - if subclass is None: - return instance # Unknown op: return as-is - new_instance = subclass.__new__(subclass) - new_instance.__init_handle_by_constructor__( - _ffi_api.TilePrimitiveCallCopyHandle, - instance, # pylint: disable=no-member - ) - return new_instance - - def replace(self, **changes: Any) -> "TilePrimitiveCall": - """Return a copy of this call with selected fields replaced. - - Every field that is not overridden in ``changes`` is preserved from - ``self`` (including ``scope``), so rebuilds never silently drop fields. - The returned node is downcast to the registered subclass for ``op``. - - Parameters - ---------- - **changes : Any - Field overrides; any of ``op``, ``args``, ``workspace``, ``config``, - ``dispatch``, ``scope``. - - Returns - ------- - new_call : TilePrimitiveCall - A new call with the requested fields replaced. - """ - unknown = set(changes) - {"op", "args", "workspace", "config", "dispatch", "scope"} - if unknown: - raise TypeError(f"Unknown field(s) for TilePrimitiveCall.replace: {sorted(unknown)}") - new_call = TilePrimitiveCall( - *changes.get("args", self.args), - op=changes.get("op", self.op), - workspace=changes.get("workspace", self.workspace), - config=changes.get("config", self.config), - dispatch=changes.get("dispatch", self.dispatch), - scope=changes.get("scope", self.scope), - ) - return TilePrimitiveCall.downcast(new_call) - - def with_workspace(self, workspace: dict[str, Buffer]) -> "TilePrimitiveCall": - """Return a copy with ``workspace`` replaced, preserving all other fields.""" - return self.replace(workspace=workspace) - - @property - def srcs(self) -> list[Expr]: - raise NotImplementedError("Subclass must implement this method") - - @property - def dsts(self) -> list[Expr]: - raise NotImplementedError("Subclass must implement this method") - - def get_private_buffers( - self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" - ) -> dict[str, Any]: - """ - Create private (intermediate) buffers needed in this operator. - - Parameters - ---------- - buffer_dict: Dict[Any, Tuple[Buffer, Optional[Stmt]]] - A dictionary containing private buffers (and their init stmts) in other operators. - Key can be anything to reference the buffer. - This is used to reuse private buffers in other operators (like identity tensor etc.). - If the buffer is not found in the buffer_dict, it will be created and added to - the buffer_dict. - If the buffer is found in the buffer_dict but smaller than required, it will be - enlarged and updated. - - sctx: DispatchContext - The dispatch context. - This is used to get the target and reuse op dispatch implementations. - - Returns: - private_buffer_refs: Dict[str, Any] - The references to private buffers created in this operator. - Key will be the name to add into workspace. - private buffer can be accessed by buffer_dict[private_buffer_refs[name]] - """ - if sctx.target.kind.name == "trn": - return self.get_private_buffers_trn(buffer_dict, sctx) - elif sctx.target.kind.name == "cuda": - return self.get_private_buffers_cuda(buffer_dict, sctx) - else: - raise ValueError(f"Unsupported target: {sctx.target.kind.name}") - - def get_private_buffers_trn( - self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" - ) -> dict[str, Any]: - return {} - - def get_private_buffers_cuda( - self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" - ) -> dict[str, Any]: - return {} - - def validate(self) -> None: - pass diff --git a/python/tvm/tirx/tile_primitive.py b/python/tvm/tirx/tile_primitive.py new file mode 100644 index 000000000000..c881b895ab02 --- /dev/null +++ b/python/tvm/tirx/tile_primitive.py @@ -0,0 +1,421 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""TIRx tile primitive IR nodes: LambdaExpr, DispatchContext, TilePrimitiveCall. + +Mirrors the C++ header ``include/tvm/tirx/tile_primitive.h``. +""" +# pylint: disable=no-member + +import inspect +from collections.abc import Callable +from typing import Any, ClassVar + +import tvm_ffi +from tvm_ffi import register_object + +from tvm.ir import Expr, Op, Range +from tvm.runtime import Object, Scriptable +from tvm.target import Target + +from . import _ffi_api +from .buffer import Buffer +from .exec_scope import ExecScope +from .expr import FloatImm, IterVar, Var +from .stmt import Stmt + + +@register_object("tirx.LambdaExpr") +class LambdaExpr(Object): + """A reified Python lambda: bound variables and a body over them. + + Used by tile primitive ops that take a per-element expression over the + destination axes (e.g. ``tirx.tile.select``). + """ + + vars: list[Var] + pred: Expr + + def __init__(self, f_pred: Callable[..., Expr]): + vars = [Var(name, "int32") for name in inspect.signature(f_pred).parameters] + pred = f_pred(*vars) + self.__init_handle_by_constructor__(_ffi_api.LambdaExpr, vars, pred) + + def apply(self, indices: list[Expr]) -> Expr: + """Substitute the bound variables with the given indices, returning the body.""" + return _ffi_api.LambdaExprApply(self, indices) + + +@register_object("tirx.DispatchContext") +class DispatchContext(Object, Scriptable): + """DispatchContext node. + + Parameters + ---------- + target : Target + The target of the dispatch context. + + exec_scope : ExecScope + The execution scope of the dispatch context. + + launch_params : Dict[str, Expr] + The launch parameters of the dispatch context. + + var_range_map : Dict[Var, Range] + A map from loop variables to their ranges. + + callbacks : Dict[str, Object] + The callbacks of the dispatch context. + + shared_state : Dict[str, Object] + Shared state persisting across dispatch calls within a single lowering pass. + """ + + target: Target + exec_scope: ExecScope + launch_params: dict[str, IterVar] + var_range_map: dict[Var, Range] + alloc_only: bool + callbacks: dict[str, Object] + shared_state: dict[str, Object] + inter: dict[str, list] + intra: dict[str, list] + scope_kind: str + + kPrivateAlloc = "private_alloc" + kDeviceInitStmt = "device_init_stmt" + kHostInitStmt = "host_init_stmt" + kPostBufferDefStmt = "post_buffer_def_stmt" + + def __init__( + self, + target: Target, + exec_scope: ExecScope, + launch_params: dict[str, IterVar], + var_range_map: dict[Var, Range], + alloc_only: bool = False, + callbacks: dict[str, Object] = {}, + shared_state: dict[str, Object] = {}, + inter: dict[str, list] | None = None, + intra: dict[str, list] | None = None, + scope_kind: str = "", + ) -> None: + self.__init_handle_by_constructor__( + _ffi_api.DispatchContext, # pylint: disable=no-member + target, + exec_scope, + launch_params, + var_range_map, + alloc_only, + callbacks, + shared_state, + inter or {}, + intra or {}, + scope_kind, + ) + + def add_alloc_buffer(self, buffer: Buffer) -> None: + """Add an allocated buffer to the dispatch context. + Can be called only if alloc_only is True. + The buffer will be added to the workspace of operator (the key in the workspace is the buffer name). + + Parameters + ---------- + buffer : Buffer + The buffer to be added. + """ # noqa: E501 + _ffi_api.DispatchContextAddAllocBuffer(self, buffer) # pylint: disable=no-member + + def add_init_stmt(self, stmt: Stmt, host: bool = False) -> None: + """Add an initialization statement to the dispatch context. + Device initialization statements is only allowed if alloc_only is True. + Host initialization statements will be ignored if alloc_only is True. + The statements will be added to the beginning of the kernel. + + Parameters + ---------- + stmt : Stmt + The initialization statement to be added. + host : bool + Whether the statement is a host statement. + If True, the statement will be added to the host code (before the kernel). + If False, the statement will be added to the kernel body (at the beginning of the kernel). + """ # noqa: E501 + _ffi_api.DispatchContextAddInitStmt(self, stmt, host) # pylint: disable=no-member + + def add_post_buffer_def_stmt(self, buffer: Buffer, stmt: Stmt) -> None: + """Add a statement to be inserted after a buffer's definition (DeclBuffer/AllocBuffer). + + Parameters + ---------- + buffer : Buffer + The buffer whose definition scope the statement should appear in. + stmt : Stmt + The statement to be inserted. + """ + _ffi_api.DispatchContextAddPostBufferDefStmt(self, buffer, stmt) # pylint: disable=no-member + + def cache_get(self, key: str) -> Object | None: + """Look up a cached value by key. + + Parameters + ---------- + key : str + Cache key (built by the caller from construction parameters). + + Returns + ------- + Optional[Object] + The cached value, or None on miss. + """ + return _ffi_api.DispatchContextSharedStateGet(self, key) + + def cache_set(self, key: str, value: Object) -> None: + """Store a value in the cross-dispatch cache. + + Parameters + ---------- + key : str + Cache key (built by the caller from construction parameters). + value : Object + The object to cache (e.g. a Buffer or Var). + """ + _ffi_api.DispatchContextSharedStateSet(self, key, value) + + def is_cuda(self) -> bool: + """Check if the target is CUDA.""" + return self.target.kind.name == "cuda" + + def is_trn(self) -> bool: + """Check if the target is Trainium.""" + return self.target.kind.name == "trn" + + def is_target(self, name: str) -> bool: + """Check if the target kind matches ``name``.""" + return self.target.kind.name == name + + # -- scope predicates ---------------------------------------------------- + # + # Each ``is_`` returns True iff the op site is at that scope kind. + # Backed by ``self.scope_kind``, which 1-1 maps to a canonical intra + # TileLayout shape: + # thread -> {} + # warp -> {laneid} + # warpgroup -> {laneid, wid_in_wg} + # cta -> {laneid, warpid} + # cluster -> {laneid, warpid, cta_id} + # + # Prefer these predicates over raw ``self.scope_kind == "..."`` comparisons + # so dispatchers that later need stricter intra/inter shape checks can + # tighten the predicate body without touching every call site. + + @property + def is_thread(self) -> bool: + return self.scope_kind == "thread" + + @property + def is_warp(self) -> bool: + return self.scope_kind == "warp" + + @property + def is_warpgroup(self) -> bool: + return self.scope_kind == "warpgroup" + + @property + def is_cta(self) -> bool: + return self.scope_kind == "cta" + + @property + def is_cluster(self) -> bool: + return self.scope_kind == "cluster" + + +def normalize_const_arg(arg) -> Expr: + if isinstance(arg, float): + return FloatImm("float32", arg) + return arg + + +@tvm_ffi.register_object("tirx.TilePrimitiveCall") +class TilePrimitiveCall(Stmt): + """TilePrimitiveCall node. + + Parameters + ---------- + op : Op + The operator. + + args : List[Expr] + The arguments. + + workspace : Map[str, Buffer] + The workspace. + + config : Map[str, ObjectRef] + The scheduler/config dictionary. + + dispatch : Optional[str] + The explicit variant name to dispatch to. + + scope : ExecScope + The cooperation scope of this call. Defaults to ``thread`` (an unscoped call). + """ + + args: list[Expr] + workspace: dict[str, Buffer] + config: dict[str, Any] + dispatch: str | None + scope: ExecScope + _registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {} + + def __init__( + self, + *args: list[Expr], + op: Op | None = None, + workspace: dict[str, Buffer] | None = None, + config: dict[str, Any] | None = None, + dispatch: str | None = None, + scope: ExecScope | None = None, + ) -> None: + if workspace is None: + workspace = {} + if config is None: + config = {} + if scope is None: + scope = ExecScope("thread") + if op is None: + assert self.__class__ != TilePrimitiveCall, ( + "Directly instantiating TilePrimitiveCall needs to specify the op" + ) + op = self.__class__.op + args = list(map(normalize_const_arg, args)) + self.__init_handle_by_constructor__( + _ffi_api.TilePrimitiveCall, + op, + args, + workspace, + config, + dispatch, + scope, # pylint: disable=no-member + ) + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if hasattr(cls, "op"): + cls._registry[cls.op] = cls + + @classmethod + def downcast(cls, instance: "TilePrimitiveCall") -> "TilePrimitiveCall": + subclass = cls._registry.get(instance.op) + if subclass is None: + return instance # Unknown op: return as-is + new_instance = subclass.__new__(subclass) + new_instance.__init_handle_by_constructor__( + _ffi_api.TilePrimitiveCallCopyHandle, + instance, # pylint: disable=no-member + ) + return new_instance + + def replace(self, **changes: Any) -> "TilePrimitiveCall": + """Return a copy of this call with selected fields replaced. + + Every field that is not overridden in ``changes`` is preserved from + ``self`` (including ``scope``), so rebuilds never silently drop fields. + The returned node is downcast to the registered subclass for ``op``. + + Parameters + ---------- + **changes : Any + Field overrides; any of ``op``, ``args``, ``workspace``, ``config``, + ``dispatch``, ``scope``. + + Returns + ------- + new_call : TilePrimitiveCall + A new call with the requested fields replaced. + """ + unknown = set(changes) - {"op", "args", "workspace", "config", "dispatch", "scope"} + if unknown: + raise TypeError(f"Unknown field(s) for TilePrimitiveCall.replace: {sorted(unknown)}") + new_call = TilePrimitiveCall( + *changes.get("args", self.args), + op=changes.get("op", self.op), + workspace=changes.get("workspace", self.workspace), + config=changes.get("config", self.config), + dispatch=changes.get("dispatch", self.dispatch), + scope=changes.get("scope", self.scope), + ) + return TilePrimitiveCall.downcast(new_call) + + def with_workspace(self, workspace: dict[str, Buffer]) -> "TilePrimitiveCall": + """Return a copy with ``workspace`` replaced, preserving all other fields.""" + return self.replace(workspace=workspace) + + @property + def srcs(self) -> list[Expr]: + raise NotImplementedError("Subclass must implement this method") + + @property + def dsts(self) -> list[Expr]: + raise NotImplementedError("Subclass must implement this method") + + def get_private_buffers( + self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" + ) -> dict[str, Any]: + """ + Create private (intermediate) buffers needed in this operator. + + Parameters + ---------- + buffer_dict: Dict[Any, Tuple[Buffer, Optional[Stmt]]] + A dictionary containing private buffers (and their init stmts) in other operators. + Key can be anything to reference the buffer. + This is used to reuse private buffers in other operators (like identity tensor etc.). + If the buffer is not found in the buffer_dict, it will be created and added to + the buffer_dict. + If the buffer is found in the buffer_dict but smaller than required, it will be + enlarged and updated. + + sctx: DispatchContext + The dispatch context. + This is used to get the target and reuse op dispatch implementations. + + Returns: + ------- + private_buffer_refs: Dict[str, Any] + The references to private buffers created in this operator. + Key will be the name to add into workspace. + private buffer can be accessed by buffer_dict[private_buffer_refs[name]] + """ + if sctx.target.kind.name == "trn": + return self.get_private_buffers_trn(buffer_dict, sctx) + elif sctx.target.kind.name == "cuda": + return self.get_private_buffers_cuda(buffer_dict, sctx) + else: + raise ValueError(f"Unsupported target: {sctx.target.kind.name}") + + def get_private_buffers_trn( + self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" + ) -> dict[str, Any]: + return {} + + def get_private_buffers_cuda( + self, buffer_dict: dict[Any, tuple[Buffer, Stmt | None]], sctx: "DispatchContext" + ) -> dict[str, Any]: + return {} + + def validate(self) -> None: + pass diff --git a/src/arith/analyzer.cc b/src/arith/analyzer.cc index 0332d23bdd56..2e6c2acda57b 100644 --- a/src/arith/analyzer.cc +++ b/src/arith/analyzer.cc @@ -337,6 +337,14 @@ TVM_FFI_STATIC_INIT_BLOCK() { [](Analyzer analyzer, const PrimExpr& expr, int strength) { return analyzer->CanProve(expr, static_cast(strength)); }) + .def("arith.EnterAllowUintAsIndex", []() { ++arith::uint_as_index::g_depth; }) + .def("arith.ExitAllowUintAsIndex", + []() { + TVM_FFI_ICHECK(arith::uint_as_index::g_depth > 0) + << "ExitAllowUintAsIndex without a matching Enter"; + --arith::uint_as_index::g_depth; + }) + .def("arith.GetAllowUintAsIndex", []() { return arith::uint_as_index::Enabled(); }) .def("arith.AnalyzerSetMaximumRewriteSteps", [](Analyzer analyzer, int64_t maximum) { analyzer->rewrite_simplify.SetMaximumRewriteSteps(maximum); diff --git a/src/arith/canonical_simplify.cc b/src/arith/canonical_simplify.cc index e72a2cc3c124..e80b7880576b 100644 --- a/src/arith/canonical_simplify.cc +++ b/src/arith/canonical_simplify.cc @@ -95,7 +95,14 @@ inline PrimExpr DivImpl(PrimExpr a, PrimExpr b, DivMode mode) { * \return whether value fits in dtype */ bool CastIsSafe(PrimType dtype, PrimExpr value, AnalyzerObj* analyzer) { - if (!IsIndexType(dtype->dtype)) { + // uint32/64 is not an index dtype by default; the no-overflow assertion + // (uint_as_index) admits it for bounded non-negative upcast elimination. + const bool is_index_dtype = + IsIndexType(dtype->dtype) || + (uint_as_index::Enabled() && + dtype->dtype.code == static_cast(DLDataTypeCode::kDLUInt) && + (dtype->dtype.bits == 32 || dtype->dtype.bits == 64) && dtype->dtype.lanes == 1); + if (!is_index_dtype) { return false; } ConstIntBound bound = analyzer->const_int_bound(value); @@ -1384,7 +1391,11 @@ Expr CanonicalSimplifier::Impl::VisitExpr_(const ReduceNode* op) { } Expr CanonicalSimplifier::Impl::VisitExpr_(const CastNode* op) { - if (!IsIndexTypedExpr(op)) { + // The cast reasoning below is index-centric; for unsigned operands it runs + // only under the caller's no-overflow assertion (uint_as_index). + if (!IsIndexTypedExpr(op) && + !(uint_as_index::Enabled() && op->ExprNode::ty.as()->dtype.code == + static_cast(DLDataTypeCode::kDLUInt))) { return Rewriter::VisitExpr_(op); } // normalize diff --git a/src/arith/const_fold.h b/src/arith/const_fold.h index f7f46fae78a4..0e324af275de 100644 --- a/src/arith/const_fold.h +++ b/src/arith/const_fold.h @@ -38,6 +38,33 @@ namespace tvm { namespace arith { +/*! + * \brief Scoped opt-in flag for unsigned index analysis (allow_u32). + * + * Unsigned arithmetic wraps modulo 2^bits, so most index rewrite rules are + * unsound for it in general. When the caller can guarantee values never + * approach the type limit (e.g. compile-time layout proofs over SMEM + * offsets), enabling this flag activates a small set of otherwise-unsafe + * rewrite rules for unsigned operands (see the unsigned blocks in + * rewrite_simplify.cc). It is OFF by default; thread-local and nested-safe. + */ +namespace uint_as_index { + +inline thread_local int g_depth = 0; + +inline bool Enabled() { return g_depth > 0; } + +} // namespace uint_as_index + +/*! \brief RAII guard enabling uint_as_index mode in the current scope. */ +class AllowUintAsIndexGuard { + public: + AllowUintAsIndexGuard() { ++uint_as_index::g_depth; } + ~AllowUintAsIndexGuard() { --uint_as_index::g_depth; } + AllowUintAsIndexGuard(const AllowUintAsIndexGuard&) = delete; + AllowUintAsIndexGuard& operator=(const AllowUintAsIndexGuard&) = delete; +}; + /*! * \brief Try to run binary compute with constant folding. * diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 461bb315431c..9c6897d83941 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -1196,6 +1196,48 @@ Expr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) { if (op_ty.MatchesCode(DLDataTypeCode::kDLUInt) && (op_ty.bits() == 32 || op_ty.bits() == 64)) { TVM_TRY_REWRITE(floordiv(x, x), OneWithTypeLike(x)); // x / x -> 1 (x != 0) TVM_TRY_REWRITE_IF(floordiv(x, c1), x, c1.Eval()->value == 1); // x / 1 -> x + // Unsigned x is always >= 0, so x / c2 == 0 when x < c2 is provable from + // the interval bound (const_int_bound is dtype-agnostic; CanProve's + // comparison machinery is signed-centric and can't be used here). + if (floordiv(x, c2).Match(ret) && c2.Eval()->value > 0) { + ConstIntBound x_bound = analyzer_->const_int_bound(x.Eval()); + if (x_bound->min_value >= 0 && x_bound->max_value < c2.Eval()->value) { + return ZeroWithTypeLike(x).Eval(); + } + } + + // Enabled only under the caller's no-overflow assertion (uint_as_index): + // these floordiv identities are unsound in general for unsigned + // wraparound (the quotient's high bits are lost when x*c1 overflows), but + // are exact in the asserted no-overflow domain. Each application is + // logged for audit. + if (uint_as_index::Enabled()) { + if (floordiv(x * c1, c2).Match(ret) && c2.Eval()->value > 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + return (x * floordiv(c1, c2)).Eval(); + } + if (floordiv(x * c1 + y, c2).Match(ret) && c2.Eval()->value > 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + return (x * floordiv(c1, c2) + floordiv(y, c2)).Eval(); + } + if (floordiv(y + x * c1, c2).Match(ret) && c2.Eval()->value > 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + return (x * floordiv(c1, c2) + floordiv(y, c2)).Eval(); + } + if (floordiv(x * c1 + y + z, c2).Match(ret) && c2.Eval()->value > 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + return (x * floordiv(c1, c2) + floordiv(y + z, c2)).Eval(); + } + if (floordiv(y + x * c1 + z, c2).Match(ret) && c2.Eval()->value > 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floordiv rule on unsigned expr: " << ret; + return (x * floordiv(c1, c2) + floordiv(y + z, c2)).Eval(); + } + } } return ret; } @@ -1208,7 +1250,7 @@ Expr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) { // Pattern var to match any expression PVar x, y, z, b1; // Pattern var match IntImm - PVar c1, c2; + PVar c1, c2, c3; // Pattern var for lanes in broadcast and ramp PVar lanes; @@ -1322,12 +1364,52 @@ Expr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) { // OVERFLOW-FREE identities are valid here. PrimType op_ty = op->ty.as_or_throw(); if (op_ty.MatchesCode(DLDataTypeCode::kDLUInt) && (op_ty.bits() == 32 || op_ty.bits() == 64)) { + const int64_t op_bits = op_ty.bits(); + const auto is_pow2_le_bits = [op_bits](int64_t v) { + return v > 0 && (v & (v - 1)) == 0 && + (op_bits < 64 ? v <= (int64_t{1} << op_bits) : v <= (int64_t{1} << 63)); + }; TVM_TRY_REWRITE(floormod(x, x), ZeroWithTypeLike(x)); // x % x -> 0 (x != 0) TVM_TRY_REWRITE_IF(floormod(x, c1), ZeroWithTypeLike(x), c1.Eval()->value == 1); // x % 1 -> 0 - // Overflow-free: when c1 is a multiple of c2, x * c1 is too (in Z and mod 2^bits). + // When c1 is a multiple of c2, x * c1 is too — but the identity is only + // sound when c2 | 2^bits: otherwise the wraparound subtraction of + // k*2^bits changes residues mod c2. (No non-pow2 c2 occurs in practice; + // probed across the arith/copy/kernel test suites.) TVM_TRY_REWRITE_IF(floormod(x * c1, c2), ZeroWithTypeLike(x), - c2.Eval()->value != 0 && c1.Eval()->value % c2.Eval()->value == 0); + c2.Eval()->value != 0 && c1.Eval()->value % c2.Eval()->value == 0 && + is_pow2_le_bits(c2.Eval()->value)); + + // When c2 divides 2^bits, reducing mod 2^bits does not change residues + // mod c2, so modular analysis is sound for unsigned operands too. + if (floormod(x, c2).Match(ret) && is_pow2_le_bits(c2.Eval()->value)) { + const int64_t c2val = c2.Eval()->value; + // x = coeff * q + base stays congruent mod c2 when c2 | 2^bits, so + // the signed block's modular-analysis branch is valid here as well. + ModularSet mod = analyzer_->modular_set(x.Eval()); + if (mod->coeff % c2val == 0) { + return floormod(mod->base, c2).Eval(); + } + } + + // Under the caller's no-overflow assertion (uint_as_index) the pow2 + // restriction above disappears: without wraparound the residues are the + // plain integer ones, so these hold for any c2 != 0. Each application is + // logged for audit. + if (uint_as_index::Enabled()) { + if (floormod(x * c1, c2).Match(ret) && c2.Eval()->value != 0 && + c1.Eval()->value % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floormod rule on unsigned expr: " << ret; + return ZeroWithTypeLike(x).Eval(); + } + if (floormod(x, c2).Match(ret) && c2.Eval()->value > 0) { + ModularSet mod = analyzer_->modular_set(x.Eval()); + if (mod->coeff % c2.Eval()->value == 0) { + LOG(WARNING) << "arith: no-overflow floormod rule on unsigned expr: " << ret; + return floormod(mod->base, c2).Eval(); + } + } + } } return ret; } diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 0245de7a4297..8db7f3d690e2 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -248,16 +248,32 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f, std::ostream& os) { return; } auto min_blocks_per_sm = f->GetAttr(tirx::attr::kLaunchBoundsMinBlocksPerSM); + auto max_blocks_per_cluster = f->GetAttr(tirx::attr::kLaunchBoundsMaxBlocksPerCluster); if (min_blocks_per_sm.has_value()) { TVM_FFI_ICHECK_GT(min_blocks_per_sm.value(), 0); - os << " __launch_bounds__(" << threadIdx_ext_int->value << ", " << min_blocks_per_sm.value() - << ")"; + os << " __launch_bounds__(" << threadIdx_ext_int->value << ", " << min_blocks_per_sm.value(); + if (max_blocks_per_cluster.has_value()) { + TVM_FFI_ICHECK_GT(max_blocks_per_cluster.value(), 0); + os << ", " << max_blocks_per_cluster.value(); + } + os << ")"; } else { + TVM_FFI_ICHECK(!max_blocks_per_cluster.has_value()) + << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires " + << tirx::attr::kLaunchBoundsMinBlocksPerSM; os << " __launch_bounds__(" << threadIdx_ext_int->value << ")"; } } } +void CodeGenCUDA::VisitStmt_(const ReturnNode* op) { + const auto* value = op->value.as(); + TVM_FFI_ICHECK(value && value->value == 0) + << "CUDA device kernel may only contain a successful early return, return 0"; + PrintIndent(); + stream << "return;\n"; +} + std::string CodeGenCUDA::Finish() { // Generate header auto header_generator = ffi::Function::GetGlobal("tirx.intrinsics.cuda.header_generator"); @@ -277,6 +293,15 @@ std::string CodeGenCUDA::Finish() { } void CodeGenCUDA::VisitStmt_(const tirx::ForNode* op) { + // Materialize the loop bounds before emitting an unroll pragma. PrintExpr + // may introduce temporaries (for example, for a Select expression). CUDA + // requires #pragma unroll to immediately precede the loop it controls; if + // those declarations are printed between the pragma and the for statement, + // nvcc is free to unroll the loop despite disable_unroll. + std::string begin_str = PrintExpr(op->min); + PrimExpr end = is_zero(op->min) ? op->extent : arith::Analyzer()->Simplify(op->min + op->extent); + std::string end_str = PrintExpr(end); + std::string step_str = op->step.has_value() ? PrintExpr(*op->step) : ""; if (op->annotations.count("disable_unroll")) { PrintIndent(); stream << "#pragma unroll 1\n"; @@ -284,10 +309,28 @@ void CodeGenCUDA::VisitStmt_(const tirx::ForNode* op) { PrintIndent(); stream << "#pragma unroll\n"; } - CodeGenC::VisitStmt_(op); + PrintIndent(); + std::string vid = AllocVarID(op->loop_var.get()); + stream << "for ("; + PrintType(op->loop_var.ty(), stream); + stream << ' ' << vid << " = " << begin_str << "; " << vid << " < " << end_str << "; "; + if (step_str.empty()) { + stream << "++" << vid; + } else { + stream << vid << " += " << step_str; + } + stream << ") {\n"; + int for_scope = BeginScope(); + PrintStmt(op->body); + this->EndScope(for_scope); + PrintIndent(); + stream << "}\n"; } void CodeGenCUDA::VisitStmt_(const WhileNode* op) { + PrintIndent(); + // Match CodeGenC: dynamic-trip-count loops must not be unrolled. + stream << "#pragma unroll 1\n"; PrintIndent(); stream << "while (1) {\n"; int while_scope = BeginScope(); @@ -980,7 +1023,6 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { static const Op& mma_store_legacy_op = Op::Get("tirx.mma_store_legacy"); static const Op& mma_fill_legacy_op = Op::Get("tirx.mma_fill_legacy"); static const Op& ptx_cp_async_bulk_op = Op::Get("tirx.ptx.cp_async_bulk"); - static const Op& ptx_cp_async_mbarrier_arrive_op = Op::Get("tirx.ptx.cp_async_mbarrier_arrive"); static const Op& ptx_ldg32_op = Op::Get("tirx.ptx.ldg32"); static const Op& cuda_func_call_op = Op::Get("tirx.cuda.func_call"); @@ -1289,16 +1331,6 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { std::string barrier_arr = barrier_name_ + "_" + std::to_string(barrier_arr_id); std::string barrier = barrier_arr + "[" + std::to_string(barrier_id) + "]"; this->stream << PrintCpAsyncBulkAsm(dst, dst_offset, src, src_offset, size, barrier); - } else if (IsOp(op, ptx_cp_async_mbarrier_arrive_op, "tirx.ptx.cp_async_mbarrier_arrive")) { - codegen_tags_.insert("cast_smem_ptr_to_int"); - int barrier_arr_id = op->args[0].as_or_throw()->value; - int barrier_id = op->args[1].as_or_throw()->value; - auto it = barrier_count_.find(barrier_arr_id); - TVM_FFI_ICHECK(it != barrier_count_.end()) << "Barrier array does not exist"; - TVM_FFI_ICHECK(barrier_id < it->second) << "Barrier id out of bounds"; - std::string barrier_arr = barrier_name_ + "_" + std::to_string(barrier_arr_id); - std::string barrier = barrier_arr + "[" + std::to_string(barrier_id) + "]"; - this->stream << PrintCpAsyncBarrierAsm(barrier); } else if (IsOp(op, ptx_ldg32_op, "tirx.ptx.ldg32")) { /* asm volatile ( @@ -1331,6 +1363,22 @@ void CodeGenCUDA::VisitExpr_(const CallNode* op, std::ostream& os) { << guard << ")\n"; stream << ");\n"; } else if (op->op.same_as(builtin::reinterpret())) { + // Compile-time pointer reinterpret of a literal (e.g. a tcgen05 descriptor + // template encoded at address 0): emit C++-style reinterpret_cast(...) + // to match the encoded template. Runtime pointer reinterprets fall through + // to CodeGenC, which emits the C-style (T*)... cast. + if (op->ty.as() && op->args[0].as()) { + os << "reinterpret_cast<"; + if (const auto* pt = op->ty.as()) { + if (const auto* et = pt->element_type.as()) { + this->PrintType(ffi::GetRef(et), os); + } else { + os << "void"; + } + } + os << "*>(" << PrintExpr(op->args[0]) << ")"; + return; + } auto tgt_prim_type = op->ty.as(); auto src_prim_type = op->args[0]->ty.as(); diff --git a/src/backend/cuda/codegen/codegen_cuda.h b/src/backend/cuda/codegen/codegen_cuda.h index 3435f0dc9c18..6b43cd890905 100644 --- a/src/backend/cuda/codegen/codegen_cuda.h +++ b/src/backend/cuda/codegen/codegen_cuda.h @@ -78,6 +78,7 @@ class CodeGenCUDA final : public CodeGenC { void VisitExpr_(const CallNode* op, std::ostream& os) final; void VisitExpr_(const CastNode* op, std::ostream& os) final; void VisitStmt_(const EvaluateNode* op) final; + void VisitStmt_(const ReturnNode* op) final; void VisitStmt_(const AllocBufferNode* op) final; void VisitStmt_(const AttrStmtNode* op) final; diff --git a/src/backend/cuda/codegen/ptx.cc b/src/backend/cuda/codegen/ptx.cc index 0d7f08b0d80a..62ed8ee71744 100644 --- a/src/backend/cuda/codegen/ptx.cc +++ b/src/backend/cuda/codegen/ptx.cc @@ -717,22 +717,5 @@ std::string PrintCpAsyncBulkAsm(const std::string& shared_ptr, return asm_code; } -std::string PrintCpAsyncBarrierAsm(const std::string& barrier) { - std::string predicated_asm_code = R"( - { - unsigned int barrier_addr_int = __cvta_generic_to_shared({barrier}); - __asm__ __volatile__( - "cp.async.mbarrier.arrive.shared.b64 [%0];" - :: "r" (barrier_addr_int) - ); - } -)"; - - Replacer replacer; - replacer.register_rule("{barrier}", "&" + barrier); - predicated_asm_code = replacer.rewrite(predicated_asm_code); - return predicated_asm_code; -} - } // namespace codegen } // namespace tvm diff --git a/src/backend/cuda/codegen/ptx.h b/src/backend/cuda/codegen/ptx.h index 3673795378a8..68e06d99085e 100644 --- a/src/backend/cuda/codegen/ptx.h +++ b/src/backend/cuda/codegen/ptx.h @@ -91,12 +91,6 @@ std::string PrintCpAsyncBulkAsm(const std::string& shared_ptr, const std::string& global_elem_offset, const std::string& bytes, const std::string& barrier); -/*! - * \brief Print ptx async copy barrier using cp.async.mbarrier.arrive - * \param barrier: The name of the barrier in shared memory. - */ -std::string PrintCpAsyncBarrierAsm(const std::string& barrier); - } // namespace codegen } // namespace tvm diff --git a/src/backend/cuda/op/iket.cc b/src/backend/cuda/op/iket.cc new file mode 100644 index 000000000000..faac989fea08 --- /dev/null +++ b/src/backend/cuda/op/iket.cc @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/*! + * \file backend/cuda/op/iket.cc + * \brief Frontend-only NVIDIA IKET annotation operators for CUDA. + */ + +#include +#include + +namespace tvm { +namespace tirx { + +#define TVM_REGISTER_CUDA_IKET_OP(OpName) \ + TVM_REGISTER_OP("tirx.cuda.iket_" #OpName) \ + .set_attr("TIRxOpCategory", ffi::String("device_intrin")) \ + .set_attr("TDeviceIntrinsicNamespace", ffi::String("cuda")) \ + .set_attr("TScriptPrinterName", ffi::String("cuda.iket." #OpName)) \ + .set_attr("TCallEffectKind", static_cast(CallEffectKind::kOpaque)) + +TVM_REGISTER_CUDA_IKET_OP(mark); +TVM_REGISTER_CUDA_IKET_OP(range_start); +TVM_REGISTER_CUDA_IKET_OP(range_end); +TVM_REGISTER_CUDA_IKET_OP(range_push); +TVM_REGISTER_CUDA_IKET_OP(range_pop); +TVM_REGISTER_CUDA_IKET_OP(sentinel_token); +TVM_REGISTER_CUDA_IKET_OP(official_event); + +#undef TVM_REGISTER_CUDA_IKET_OP + +} // namespace tirx +} // namespace tvm diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc index 852cab036a61..64a283cc2008 100644 --- a/src/backend/cuda/op/target_builtin.cc +++ b/src/backend/cuda/op/target_builtin.cc @@ -182,6 +182,7 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(cuda_cta_sync, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_cvta_generic_to_shared, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_fadd2_rn, cuda, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(cuda_fdividef, cuda, kPure), TIRX_DEVICE_INTRIN_ALIAS(cuda_ffs_u32, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_float22bfloat162_rn, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_float22bfloat162_rn_from_float2, cuda, kOpaque), @@ -207,7 +208,7 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(cuda_reduce_min_sync_u32, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_runtime_instr_desc, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_smem_addr_from_uint64, cuda, kOpaque), - TIRX_DEVICE_INTRIN_ALIAS(cuda_sm100_tma_2sm_mbarrier_addr, cuda, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(cuda_sm100_2sm_leader_smem_addr, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_syncthreads_and, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_syncthreads_or, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_thread_fence, cuda, kOpaque), @@ -241,6 +242,7 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(ptx_atom_scalar, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_bar_arrive, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_bar_sync, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_barrier_sync, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_barrier_cluster_arrive, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_barrier_cluster_wait, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_clc_query_cancel, ptx, kOpaque), @@ -253,15 +255,16 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_s2g, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_s2s_cluster, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_shared_to_cluster, ptx, kOpaque), - TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_global_to_cluster, ptx, kOpaque), - TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_global_to_cluster_prefetch, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_g2s_cluster, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_g2s_cta, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_prefetch, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_shared_to_global, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_shared_to_global_reduce, ptx, kOpaque), - TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_bulk_wait_group, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_commit_group, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_mbarrier_arrive, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_cp_async_wait_group, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_cvt, ptx, kPure), TIRX_DEVICE_INTRIN_ALIAS(ptx_elect_sync, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_exp2, ptx, kPure), TIRX_DEVICE_INTRIN_ALIAS(ptx_fence, ptx, kOpaque), @@ -277,6 +280,7 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(ptx_ld, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_acquire, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_global_acquire, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_global_nc, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_mmio, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_relaxed, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_volatile, ptx, kOpaque), @@ -287,6 +291,8 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(ptx_max_f32, ptx, kPure), TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_arrive, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_arrive_expect_tx, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_arrive_no_complete, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_complete_tx, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_init, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_test_wait_parity, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_mbarrier_try_wait, ptx, kOpaque), diff --git a/src/backend/cuda/runtime/cuda_device_api.cc b/src/backend/cuda/runtime/cuda_device_api.cc index 08c38a331c56..01222b408588 100644 --- a/src/backend/cuda/runtime/cuda_device_api.cc +++ b/src/backend/cuda/runtime/cuda_device_api.cc @@ -486,6 +486,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto oob_fill_kind = static_cast(args[arg_cnt++].cast()); int force_cu_dtype = (arg_cnt < static_cast(args.size())) ? args[arg_cnt++].cast() : -1; + TVM_FFI_ICHECK_GE(force_cu_dtype, -1) + << "force_cu_dtype must be -1 (derive from tensor_dtype) or a supported " + "CUtensorMapDataType value"; TVM_FFI_ICHECK_EQ(tensor_dtype.lanes, 1) << "Expect tensor_dtype to have lanes=1, but get " << tensor_dtype; @@ -583,6 +586,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { // bypass the dtype-derived mapping so the descriptor uses the requested // CUtensorMapDataType. Same byte size, different on-load rounding semantics. if (force_cu_dtype >= 0) { + TVM_FFI_ICHECK_EQ(force_cu_dtype, 11) + << "force_cu_dtype only supports CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 (11)"; + TVM_FFI_ICHECK(tensor_dtype.code == kDLFloat && tensor_dtype.bits == 32) + << "CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 requires a scalar float32 tensor_dtype"; cu_dtype = static_cast(force_cu_dtype); } @@ -669,6 +676,13 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_ICHECK_GE(tensor_rank, 3U) << "tensorRank must be greater than or equal to 3 when interleave is not NONE"; } + if (interleaved_kind == CU_TENSOR_MAP_INTERLEAVE_32B) { + TVM_FFI_ICHECK_EQ(swizzle_kind, CU_TENSOR_MAP_SWIZZLE_32B) + << "CU_TENSOR_MAP_INTERLEAVE_32B requires CU_TENSOR_MAP_SWIZZLE_32B"; + } + TVM_FFI_ICHECK_EQ(element_strides[0], 1U) + << "elementStrides[0] must be one because tiled TMA does not support an " + "innermost element stride"; if (interleaved_kind == CU_TENSOR_MAP_INTERLEAVE_32B || is_packed_align16) { TVM_FFI_ICHECK_EQ((reinterpret_cast(tensor_ptr) & 0b11111), 0) << "globalAddress must be 32-byte aligned"; @@ -685,6 +699,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_ICHECK_EQ(box_dim[0], 128U) << "boxDim[0] must be 128 for packed 16U4/16U6 align16 formats"; } + if (is_packed_16u4_align16) { + bool supported_swizzle = + swizzle_kind == CU_TENSOR_MAP_SWIZZLE_NONE || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B; +#ifdef CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B + supported_swizzle = supported_swizzle || swizzle_kind == CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; +#endif + TVM_FFI_ICHECK(supported_swizzle) + << "CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B supports only NONE, 128B, " + "or 128B_ATOM_32B swizzle"; + } if (is_packed_16u4_align8) { TVM_FFI_ICHECK_EQ(global_shape[0] % 2, 0) << "globalDim[0] must be a multiple of 2 for packed 16U4 align8 format"; diff --git a/src/backend/cuda/transforms/lower_iket.cc b/src/backend/cuda/transforms/lower_iket.cc new file mode 100644 index 000000000000..405174435a36 --- /dev/null +++ b/src/backend/cuda/transforms/lower_iket.cc @@ -0,0 +1,1581 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file backend/cuda/transforms/lower_iket.cc + * \brief Lower frontend-only TIRx IKET annotations to CUDA tracing helpers. + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace tvm { +namespace tirx { +namespace transform { + +namespace { + +constexpr int kMaxDeclarations = 30; +constexpr int kOfficialMetaInfoBytes = 48; +constexpr int kOfficialEventAttributesBytes = 60; +constexpr int kOfficialRangeAttributesBytes = 72; +constexpr size_t kMaxTokenAnalysisStates = 256; +constexpr size_t kMaxTokenAnalysisIterations = 256; +constexpr size_t kMaxConvergenceAnalysisIterations = 256; + +enum class DeclarationKind : int { + kRange = 1, + kPush = 2, + kMark = 3, +}; + +const Op& IketMarkOp() { + static const Op& op = Op::Get("tirx.cuda.iket_mark"); + return op; +} + +const Op& IketRangeStartOp() { + static const Op& op = Op::Get("tirx.cuda.iket_range_start"); + return op; +} + +const Op& IketRangeEndOp() { + static const Op& op = Op::Get("tirx.cuda.iket_range_end"); + return op; +} + +const Op& IketRangePushOp() { + static const Op& op = Op::Get("tirx.cuda.iket_range_push"); + return op; +} + +const Op& IketRangePopOp() { + static const Op& op = Op::Get("tirx.cuda.iket_range_pop"); + return op; +} + +const Op& IketSentinelOp() { + static const Op& op = Op::Get("tirx.cuda.iket_sentinel_token"); + return op; +} + +bool IsIketOp(const ffi::ObjectRef& value) { + auto op = value.as(); + if (!op.has_value()) return false; + return op.value().same_as(IketMarkOp()) || op.value().same_as(IketRangeStartOp()) || + op.value().same_as(IketRangeEndOp()) || op.value().same_as(IketRangePushOp()) || + op.value().same_as(IketRangePopOp()) || op.value().same_as(IketSentinelOp()); +} + +bool IsTokenProducer(const CallNode* call) { + return call->op.same_as(IketRangeStartOp()) || call->op.same_as(IketSentinelOp()); +} + +bool IsValidUTF8(const std::string& value) { + size_t i = 0; + while (i < value.size()) { + uint8_t first = static_cast(value[i]); + if (first <= 0x7f) { + ++i; + continue; + } + int continuation = 0; + uint32_t codepoint = 0; + if ((first & 0xe0) == 0xc0) { + continuation = 1; + codepoint = first & 0x1f; + } else if ((first & 0xf0) == 0xe0) { + continuation = 2; + codepoint = first & 0x0f; + } else if ((first & 0xf8) == 0xf0) { + continuation = 3; + codepoint = first & 0x07; + } else { + return false; + } + if (i + continuation >= value.size()) return false; + for (int j = 1; j <= continuation; ++j) { + uint8_t byte = static_cast(value[i + j]); + if ((byte & 0xc0) != 0x80) return false; + codepoint = (codepoint << 6) | (byte & 0x3f); + } + if ((continuation == 1 && codepoint < 0x80) || (continuation == 2 && codepoint < 0x800) || + (continuation == 3 && codepoint < 0x10000) || codepoint > 0x10ffff || + (codepoint >= 0xd800 && codepoint <= 0xdfff)) { + return false; + } + i += continuation + 1; + } + return true; +} + +std::string GetName(const CallNode* call, size_t index = 0) { + TVM_FFI_CHECK(call->args.size() > index, ValueError) + << call->op.as().value()->name << " requires a literal event name"; + const auto* name = call->args[index].as(); + TVM_FFI_CHECK(name != nullptr, TypeError) + << call->op.as().value()->name << " requires a string-literal event name"; + std::string result = name->value; + TVM_FFI_CHECK(!result.empty(), ValueError) << "IKET event names must not be empty"; + TVM_FFI_CHECK(result.size() <= 32, ValueError) + << "IKET event name exceeds 32 UTF-8 bytes: " << result; + TVM_FFI_CHECK(IsValidUTF8(result), ValueError) << "IKET event name is not valid UTF-8"; + return result; +} + +std::string ValidatePayload(const Expr& payload) { + DLDataType dtype = payload.as_or_throw().ty()->dtype; + TVM_FFI_CHECK_EQ(dtype.lanes, 1, TypeError) << "IKET payload must be a scalar value"; + TVM_FFI_CHECK_GT(dtype.bits, 0, TypeError) << "IKET payload must have a concrete bit width"; + TVM_FFI_CHECK_LE(dtype.bits, 64, TypeError) << "IKET payload must be at most 64 bits"; + bool supported = dtype.code == kDLInt || dtype.code == kDLUInt || dtype.code == kDLFloat || + dtype.code == kDLBfloat; + TVM_FFI_CHECK(supported, TypeError) << "IKET payload must be bool, int, uint, or float, but got " + << ffi::DLDataTypeToString(dtype); + CallEffectKind effect = SideEffect(payload.as_or_throw()); + TVM_FFI_CHECK(effect <= CallEffectKind::kReadState, ValueError) + << "IKET payload expressions may read state but must not update state; got " << effect; + return ffi::DLDataTypeToString(dtype); +} + +bool IsScalarBufferAccess(const BufferVar& buffer, const ffi::Array& indices) { + if (buffer->shape.size() != 1 || indices.size() != 1) return false; + const auto* extent = buffer->shape[0].as(); + const auto* index = indices[0].as(); + return extent && extent->value == 1 && index && index->value == 0; +} + +const char* DeclarationKindName(DeclarationKind kind) { + switch (kind) { + case DeclarationKind::kRange: + return "range"; + case DeclarationKind::kPush: + return "push_pop"; + case DeclarationKind::kMark: + return "mark"; + } + return "unknown"; +} + +struct DeclarationKey { + DeclarationKind kind; + std::string name; + + bool operator<(const DeclarationKey& other) const { + return std::tie(kind, name) < std::tie(other.kind, other.name); + } +}; + +struct Declaration { + DeclarationKey key; + bool payload_schema_set{false}; + bool has_payload{false}; + std::string payload_dtype; + bool end_payload_schema_set{false}; + bool end_has_payload{false}; + std::string end_payload_dtype; + uint32_t event_id{0}; +}; + +class AnnotationCollector : public StmtExprVisitor { + public: + std::map declarations; + bool has_annotations{false}; + + private: + void AddDeclaration(DeclarationKind kind, const CallNode* call, bool sentinel = false) { + std::string name = GetName(call); + TVM_FFI_CHECK(call->args.size() == 1 || call->args.size() == 2, TypeError) + << call->op.as().value()->name << " expects a name and optional payload"; + bool has_payload = call->args.size() == 2; + std::string dtype = has_payload ? ValidatePayload(call->args[1]) : std::string(); + auto [name_it, name_inserted] = declaration_kinds_.emplace(name, kind); + TVM_FFI_CHECK(name_inserted || name_it->second == kind, ValueError) + << "IKET declaration " << name << " changes event kind from " + << DeclarationKindName(name_it->second) << " to " << DeclarationKindName(kind); + DeclarationKey key{kind, std::move(name)}; + auto [it, inserted] = declarations.emplace(key, Declaration{key}); + Declaration& declaration = it->second; + if (!sentinel) { + if (declaration.payload_schema_set) { + TVM_FFI_CHECK_EQ(declaration.has_payload, has_payload, ValueError) + << "IKET declaration " << declaration.key.name + << " changes payload presence between sites"; + TVM_FFI_CHECK_EQ(declaration.payload_dtype, dtype, TypeError) + << "IKET declaration " << declaration.key.name + << " changes payload dtype between sites"; + } else { + declaration.payload_schema_set = true; + declaration.has_payload = has_payload; + declaration.payload_dtype = std::move(dtype); + } + } + } + + void VisitExpr_(const CallNode* call) final { + if (!IsIketOp(call->op)) { + StmtExprVisitor::VisitExpr_(call); + return; + } + has_annotations = true; + if (call->op.same_as(IketMarkOp())) { + AddDeclaration(DeclarationKind::kMark, call); + } else if (call->op.same_as(IketRangeStartOp())) { + TVM_FFI_CHECK(call->ty.as_or_throw()->dtype.code == kDLUInt && + call->ty.as_or_throw()->dtype.bits == 32, + TypeError) + << "IKET range_start must return uint32"; + AddDeclaration(DeclarationKind::kRange, call); + } else if (call->op.same_as(IketSentinelOp())) { + TVM_FFI_CHECK_EQ(call->args.size(), 1, TypeError) + << "IKET sentinel_token expects one event name"; + TVM_FFI_CHECK(call->ty.as_or_throw()->dtype.code == kDLUInt && + call->ty.as_or_throw()->dtype.bits == 32, + TypeError) + << "IKET sentinel_token must return uint32"; + AddDeclaration(DeclarationKind::kRange, call, true); + } else if (call->op.same_as(IketRangePushOp())) { + AddDeclaration(DeclarationKind::kPush, call); + } else if (call->op.same_as(IketRangeEndOp())) { + TVM_FFI_CHECK(call->args.size() == 1 || call->args.size() == 2, TypeError) + << "IKET range_end expects a token and optional payload"; + DLDataType token_dtype = call->args[0].as_or_throw().ty()->dtype; + TVM_FFI_CHECK(token_dtype.code == kDLUInt && token_dtype.bits == 32 && token_dtype.lanes == 1, + TypeError) + << "IKET RangeToken must have dtype uint32"; + if (call->args.size() == 2) ValidatePayload(call->args[1]); + } else if (call->op.same_as(IketRangePopOp())) { + TVM_FFI_CHECK_EQ(call->args.size(), 0, TypeError) << "IKET range_pop takes no arguments"; + } + } + + std::unordered_map declaration_kinds_; +}; + +using TokenBufferSet = std::unordered_set; + +class TokenBufferCollector : public StmtExprVisitor { + public: + explicit TokenBufferCollector(TokenBufferSet* buffers) : buffers_(buffers) {} + + bool changed{false}; + + private: + void VisitStmt_(const BufferStoreNode* store) final { + bool is_token_value = false; + if (const auto* call = store->value.as()) { + is_token_value = IsTokenProducer(call); + } else if (const auto* load = store->value.as()) { + is_token_value = buffers_->count(load->buffer.get()); + } + if (is_token_value && buffers_->insert(store->buffer.get()).second) changed = true; + StmtExprVisitor::VisitStmt_(store); + } + + TokenBufferSet* buffers_; +}; + +TokenBufferSet CollectTokenBuffers(const Stmt& body) { + TokenBufferSet buffers; + bool changed = true; + while (changed) { + TokenBufferCollector collector(&buffers); + collector(body); + changed = collector.changed; + } + return buffers; +} + +using TokenDeclarationMap = std::unordered_map>; + +class TokenDeclarationCollector : public StmtExprVisitor { + public: + explicit TokenDeclarationCollector(TokenDeclarationMap* declarations) + : declarations_(declarations) {} + + bool changed{false}; + + private: + void VisitStmt_(const BufferStoreNode* store) final { + std::set possible; + if (const auto* call = store->value.as(); call && IsTokenProducer(call)) { + possible.insert(DeclarationKey{DeclarationKind::kRange, GetName(call)}); + } else if (const auto* load = store->value.as()) { + auto it = declarations_->find(load->buffer.get()); + if (it != declarations_->end()) possible = it->second; + } + if (!possible.empty()) { + auto& target = (*declarations_)[store->buffer.get()]; + size_t old_size = target.size(); + target.insert(possible.begin(), possible.end()); + changed = changed || target.size() != old_size; + } + StmtExprVisitor::VisitStmt_(store); + } + + TokenDeclarationMap* declarations_; +}; + +TokenDeclarationMap CollectTokenDeclarations(const Stmt& body) { + TokenDeclarationMap declarations; + bool changed = true; + while (changed) { + TokenDeclarationCollector collector(&declarations); + collector(body); + changed = collector.changed; + } + return declarations; +} + +class RangeEndSchemaVerifier : public StmtExprVisitor { + public: + RangeEndSchemaVerifier(const TokenDeclarationMap& token_declarations, + std::map* declarations) + : token_declarations_(token_declarations), declarations_(declarations) {} + + private: + void VisitExpr_(const CallNode* call) final { + if (!call->op.same_as(IketRangeEndOp())) { + StmtExprVisitor::VisitExpr_(call); + return; + } + const auto* token = call->args[0].as(); + TVM_FFI_ICHECK(token != nullptr); + auto possible_it = token_declarations_.find(token->buffer.get()); + TVM_FFI_ICHECK(possible_it != token_declarations_.end()); + bool has_payload = call->args.size() == 2; + std::string dtype = has_payload ? ValidatePayload(call->args[1]) : std::string(); + for (const DeclarationKey& key : possible_it->second) { + auto declaration_it = declarations_->find(key); + TVM_FFI_ICHECK(declaration_it != declarations_->end()); + Declaration& declaration = declaration_it->second; + if (declaration.end_payload_schema_set) { + TVM_FFI_CHECK_EQ(declaration.end_has_payload, has_payload, ValueError) + << "range_end for " << key.name << " changes payload presence between sites"; + TVM_FFI_CHECK_EQ(declaration.end_payload_dtype, dtype, TypeError) + << "range_end for " << key.name << " changes payload dtype between sites"; + } else { + declaration.end_payload_schema_set = true; + declaration.end_has_payload = has_payload; + declaration.end_payload_dtype = dtype; + } + } + StmtExprVisitor::VisitExpr_(call); + } + + const TokenDeclarationMap& token_declarations_; + std::map* declarations_; +}; + +class TokenVerifier : public StmtExprVisitor { + public: + explicit TokenVerifier(const TokenBufferSet& token_buffers) : token_buffers_(token_buffers) {} + + private: + void VisitStmt_(const BufferStoreNode* store) final { + if (!token_buffers_.count(store->buffer.get())) { + StmtExprVisitor::VisitStmt_(store); + return; + } + TVM_FFI_CHECK(store->buffer->dtype->dtype.code == kDLUInt && + store->buffer->dtype->dtype.bits == 32 && + store->buffer->dtype->dtype.lanes == 1, + TypeError) + << "IKET RangeToken storage must have dtype uint32"; + TVM_FFI_CHECK( + store->buffer.scope() == "local" && IsScalarBufferAccess(store->buffer, store->indices), + ValueError) + << "IKET RangeToken must use element zero of a local uint32[1] buffer"; + bool valid_value = false; + if (const auto* call = store->value.as()) { + valid_value = IsTokenProducer(call); + allow_producer_ = valid_value; + VisitExpr(store->value); + allow_producer_ = false; + } else if (const auto* load = store->value.as()) { + valid_value = token_buffers_.count(load->buffer.get()); + allow_token_load_ = valid_value; + VisitExpr(store->value); + allow_token_load_ = false; + } + TVM_FFI_CHECK(valid_value, ValueError) + << "RangeToken may only be assigned another token, range_start, or sentinel_token"; + for (const PrimExpr& index : store->indices) VisitExpr(index); + } + + void VisitExpr_(const BufferLoadNode* load) final { + if (token_buffers_.count(load->buffer.get())) { + TVM_FFI_CHECK(allow_token_load_, ValueError) + << "RangeToken may only be assigned or passed directly to range_end"; + } + StmtExprVisitor::VisitExpr_(load); + } + + void VisitExpr_(const CallNode* call) final { + if (IsTokenProducer(call)) { + TVM_FFI_CHECK(allow_producer_, ValueError) + << "range_start and sentinel_token results must be assigned to a RangeToken"; + bool old_allow_producer = allow_producer_; + allow_producer_ = false; + for (const Expr& arg : call->args) VisitExpr(arg); + allow_producer_ = old_allow_producer; + return; + } + if (call->op.same_as(IketRangeEndOp())) { + TVM_FFI_CHECK_GE(call->args.size(), 1, TypeError) << "range_end requires a RangeToken"; + const auto* token = call->args[0].as(); + TVM_FFI_CHECK(token != nullptr && token_buffers_.count(token->buffer.get()), ValueError) + << "range_end requires a directly loaded RangeToken"; + allow_token_load_ = true; + VisitExpr(call->args[0]); + allow_token_load_ = false; + for (size_t i = 1; i < call->args.size(); ++i) VisitExpr(call->args[i]); + return; + } + StmtExprVisitor::VisitExpr_(call); + } + + const TokenBufferSet& token_buffers_; + bool allow_token_load_{false}; + bool allow_producer_{false}; +}; + +enum class TokenValueKind : uint8_t { + kSentinel, + kActiveRange, + kConsumed, +}; + +struct TokenValue { + TokenValueKind kind{TokenValueKind::kSentinel}; + std::string name; + + bool operator==(const TokenValue& other) const { + return kind == other.kind && name == other.name; + } + + bool operator<(const TokenValue& other) const { + return std::tie(kind, name) < std::tie(other.kind, other.name); + } +}; + +struct TokenAnalysisState { + std::map token_values; + std::set active_ranges; + + bool operator==(const TokenAnalysisState& other) const { + return token_values == other.token_values && active_ranges == other.active_ranges; + } + + bool operator<(const TokenAnalysisState& other) const { + return std::tie(token_values, active_ranges) < + std::tie(other.token_values, other.active_ranges); + } +}; + +using TokenAnalysisStates = std::set; + +class TokenOperationFinder : public StmtExprVisitor { + public: + bool found{false}; + + private: + void VisitExpr_(const CallNode* call) final { + if (IsTokenProducer(call) || call->op.same_as(IketRangeEndOp())) { + found = true; + return; + } + StmtExprVisitor::VisitExpr_(call); + } +}; + +/*! \brief Prove the strict token alternation required by NVIDIA IKET. + * + * Any state explosion, unsupported control flow, unknown token, second start of + * an active name, or end of an inactive name rejects the annotated module. + */ +class OfficialTokenAnalyzer { + public: + explicit OfficialTokenAnalyzer(const TokenBufferSet& token_buffers) + : token_buffers_(token_buffers) {} + + bool Prove(const Stmt& body) { + TokenAnalysisStates initial{TokenAnalysisState{}}; + TokenAnalysisStates output = Process(body, initial); + if (!valid_) return false; + for (const TokenAnalysisState& state : output) { + if (!state.active_ranges.empty()) return false; + } + return true; + } + + private: + TokenAnalysisStates Limit(TokenAnalysisStates states) { + if (states.size() > kMaxTokenAnalysisStates) valid_ = false; + return valid_ ? std::move(states) : TokenAnalysisStates{}; + } + + TokenAnalysisStates Union(const TokenAnalysisStates& lhs, const TokenAnalysisStates& rhs) { + TokenAnalysisStates result = lhs; + result.insert(rhs.begin(), rhs.end()); + return Limit(std::move(result)); + } + + TokenAnalysisStates ProcessLoop(const Stmt& body, const TokenAnalysisStates& input) { + TokenAnalysisStates closure = input; + for (size_t iteration = 0; valid_ && iteration < kMaxTokenAnalysisIterations; ++iteration) { + TokenAnalysisStates after_body = Process(body, closure); + if (!valid_) return {}; + TokenAnalysisStates next = Union(closure, after_body); + if (!valid_) return {}; + if (next == closure) return closure; + closure = std::move(next); + } + valid_ = false; + return {}; + } + + TokenAnalysisStates ProcessTokenStore(const BufferStoreNode* store, + const TokenAnalysisStates& input) { + TokenAnalysisStates result; + for (const TokenAnalysisState& old_state : input) { + TokenAnalysisState state = old_state; + if (const auto* call = store->value.as(); call && IsTokenProducer(call)) { + if (call->op.same_as(IketSentinelOp())) { + state.token_values[store->buffer.get()] = TokenValue{TokenValueKind::kSentinel, {}}; + } else { + std::string name = GetName(call); + if (state.active_ranges.count(name)) { + valid_ = false; + return {}; + } + state.active_ranges.insert(name); + state.token_values[store->buffer.get()] = + TokenValue{TokenValueKind::kActiveRange, std::move(name)}; + } + } else if (const auto* load = store->value.as(); + load && token_buffers_.count(load->buffer.get())) { + auto source = state.token_values.find(load->buffer.get()); + if (source == state.token_values.end()) { + valid_ = false; + return {}; + } + state.token_values[store->buffer.get()] = source->second; + } else { + valid_ = false; + return {}; + } + result.insert(std::move(state)); + if (store->predicate.has_value()) result.insert(old_state); + } + return Limit(std::move(result)); + } + + TokenAnalysisStates ProcessRangeEnd(const CallNode* call, const TokenAnalysisStates& input) { + const auto* load = call->args[0].as(); + if (!load || !token_buffers_.count(load->buffer.get())) { + valid_ = false; + return {}; + } + TokenAnalysisStates result; + for (const TokenAnalysisState& old_state : input) { + auto token = old_state.token_values.find(load->buffer.get()); + if (token == old_state.token_values.end()) { + valid_ = false; + return {}; + } + if (token->second.kind == TokenValueKind::kConsumed) { + valid_ = false; + return {}; + } + TokenAnalysisState state = old_state; + if (token->second.kind == TokenValueKind::kActiveRange) { + auto active = state.active_ranges.find(token->second.name); + if (active == state.active_ranges.end()) { + valid_ = false; + return {}; + } + state.active_ranges.erase(active); + } + state.token_values[load->buffer.get()] = TokenValue{TokenValueKind::kConsumed, {}}; + result.insert(std::move(state)); + } + return Limit(std::move(result)); + } + + TokenAnalysisStates Process(const Stmt& stmt, const TokenAnalysisStates& input) { + if (!valid_ || input.empty()) return input; + if (const auto* sequence = stmt.as()) { + TokenAnalysisStates states = input; + for (const Stmt& item : sequence->seq) { + states = Process(item, states); + if (!valid_ || states.empty()) break; + } + return states; + } + if (const auto* store = stmt.as()) { + if (token_buffers_.count(store->buffer.get())) { + return ProcessTokenStore(store, input); + } + return input; + } + if (const auto* evaluate = stmt.as()) { + if (const auto* call = evaluate->value.as()) { + if (call->op.same_as(IketRangeEndOp())) return ProcessRangeEnd(call, input); + if (call->op.same_as(builtin::thread_return())) { + for (const TokenAnalysisState& state : input) { + if (!state.active_ranges.empty()) { + valid_ = false; + break; + } + } + return {}; + } + } + return input; + } + if (const auto* branch = stmt.as()) { + TokenAnalysisStates then_states = Process(branch->then_case, input); + TokenAnalysisStates else_states = + branch->else_case.has_value() ? Process(branch->else_case.value(), input) : input; + return Union(then_states, else_states); + } + if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); + if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); + if (const auto* attr_stmt = stmt.as()) { + return Process(attr_stmt->body, input); + } + if (const auto* block = stmt.as()) { + TokenAnalysisStates states = input; + if (block->init.has_value()) states = Union(states, Process(block->init.value(), states)); + return Process(block->body, states); + } + if (const auto* realize = stmt.as()) { + return Union(input, Process(realize->block, input)); + } + if (stmt.as() || stmt.as()) { + valid_ = false; + return {}; + } + // Unknown statement forms are harmless only if they cannot hide token + // operations. Otherwise their control flow has not been proved. + TokenOperationFinder finder; + finder(stmt); + if (finder.found) { + valid_ = false; + return {}; + } + return input; + } + + const TokenBufferSet& token_buffers_; + bool valid_{true}; +}; + +using OfficialStackState = std::vector; +using OfficialStackStates = std::set; + +class StackOperationFinder : public StmtExprVisitor { + public: + bool found{false}; + + private: + void VisitExpr_(const CallNode* call) final { + if (call->op.same_as(IketRangePushOp()) || call->op.same_as(IketRangePopOp())) { + found = true; + return; + } + StmtExprVisitor::VisitExpr_(call); + } +}; + +/*! \brief Prove balanced LIFO push/pop behavior required by NVIDIA IKET. */ +class OfficialStackAnalyzer { + public: + bool Prove(const Stmt& body) { + OfficialStackStates output = Process(body, OfficialStackStates{OfficialStackState{}}); + if (!valid_) return false; + for (const OfficialStackState& state : output) { + if (!state.empty()) return false; + } + return true; + } + + private: + OfficialStackStates Limit(OfficialStackStates states) { + if (states.size() > kMaxTokenAnalysisStates) valid_ = false; + return valid_ ? std::move(states) : OfficialStackStates{}; + } + + OfficialStackStates Union(const OfficialStackStates& lhs, const OfficialStackStates& rhs) { + OfficialStackStates result = lhs; + result.insert(rhs.begin(), rhs.end()); + return Limit(std::move(result)); + } + + OfficialStackStates ProcessLoop(const Stmt& body, const OfficialStackStates& input) { + OfficialStackStates closure = input; + for (size_t iteration = 0; valid_ && iteration < kMaxTokenAnalysisIterations; ++iteration) { + OfficialStackStates after_body = Process(body, closure); + if (!valid_) return {}; + OfficialStackStates next = Union(closure, after_body); + if (!valid_) return {}; + if (next == closure) return closure; + closure = std::move(next); + } + valid_ = false; + return {}; + } + + OfficialStackStates Process(const Stmt& stmt, const OfficialStackStates& input) { + if (!valid_ || input.empty()) return input; + if (const auto* sequence = stmt.as()) { + OfficialStackStates states = input; + for (const Stmt& item : sequence->seq) { + states = Process(item, states); + if (!valid_ || states.empty()) break; + } + return states; + } + if (const auto* evaluate = stmt.as()) { + if (const auto* call = evaluate->value.as()) { + if (call->op.same_as(IketRangePushOp())) { + OfficialStackStates result; + for (OfficialStackState state : input) { + state.push_back(GetName(call)); + result.insert(std::move(state)); + } + return Limit(std::move(result)); + } + if (call->op.same_as(IketRangePopOp())) { + OfficialStackStates result; + for (OfficialStackState state : input) { + if (state.empty()) { + valid_ = false; + return {}; + } + state.pop_back(); + result.insert(std::move(state)); + } + return Limit(std::move(result)); + } + if (call->op.same_as(builtin::thread_return())) { + for (const OfficialStackState& state : input) { + if (!state.empty()) { + valid_ = false; + break; + } + } + return {}; + } + } + return input; + } + if (const auto* branch = stmt.as()) { + OfficialStackStates then_states = Process(branch->then_case, input); + OfficialStackStates else_states = + branch->else_case.has_value() ? Process(branch->else_case.value(), input) : input; + return Union(then_states, else_states); + } + if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); + if (const auto* loop = stmt.as()) return ProcessLoop(loop->body, input); + if (const auto* attr_stmt = stmt.as()) return Process(attr_stmt->body, input); + if (const auto* block = stmt.as()) { + OfficialStackStates states = input; + if (block->init.has_value()) states = Union(states, Process(block->init.value(), states)); + return Process(block->body, states); + } + if (const auto* realize = stmt.as()) { + return Union(input, Process(realize->block, input)); + } + if (stmt.as() || stmt.as()) { + valid_ = false; + return {}; + } + StackOperationFinder finder; + finder(stmt); + if (finder.found) { + valid_ = false; + return {}; + } + return input; + } + + bool valid_{true}; +}; + +class StripIket : public StmtExprMutator { + public: + explicit StripIket(TokenBufferSet token_buffers) : token_buffers_(std::move(token_buffers)) {} + + private: + Stmt VisitStmt_(const AllocBufferNode* alloc) final { + if (token_buffers_.count(alloc->buffer.get())) return Evaluate(0); + return StmtExprMutator::VisitStmt_(alloc); + } + + Stmt VisitStmt_(const BufferStoreNode* store) final { + if (token_buffers_.count(store->buffer.get())) return Evaluate(0); + return StmtExprMutator::VisitStmt_(store); + } + + Stmt VisitStmt_(const EvaluateNode* evaluate) final { + if (const auto* call = evaluate->value.as(); call && IsIketOp(call->op)) { + return Evaluate(0); + } + return StmtExprMutator::VisitStmt_(evaluate); + } + + Expr VisitExpr_(const CallNode* call) final { + if (call->op.same_as(IketRangeStartOp()) || call->op.same_as(IketSentinelOp())) { + return IntImm(PrimType::UInt(32), 0); + } + return StmtExprMutator::VisitExpr_(call); + } + + TokenBufferSet token_buffers_; +}; + +bool IsEvaluateZero(const Stmt& stmt) { + const auto* evaluate = stmt.as(); + const auto* value = evaluate ? evaluate->value.as() : nullptr; + return value && value->value == 0; +} + +class RemoveStrippedIketNoOps : public StmtExprMutator { + private: + Stmt PreserveConditionEffects(const PrimExpr& condition) { + return SideEffect(condition) > CallEffectKind::kReadState ? Evaluate(condition) : Evaluate(0); + } + + Stmt VisitStmt_(const SeqStmtNode* sequence) final { + ffi::Array items; + for (const Stmt& item : sequence->seq) { + Stmt rewritten = VisitStmt(item); + if (!IsEvaluateZero(rewritten)) items.push_back(std::move(rewritten)); + } + return SeqStmt::Flatten(items); + } + + Stmt VisitStmt_(const AttrStmtNode* attr_stmt) final { + Stmt body = VisitStmt(attr_stmt->body); + if (IsEvaluateZero(body)) return body; + if (body.same_as(attr_stmt->body)) return ffi::GetRef(attr_stmt); + return AttrStmt(attr_stmt->node, attr_stmt->attr_key, + VisitExpr(attr_stmt->value).as_or_throw(), body, attr_stmt->span); + } + + Stmt VisitStmt_(const ForNode* loop) final { + Stmt body = VisitStmt(loop->body); + if (IsEvaluateZero(body)) return body; + if (body.same_as(loop->body)) return ffi::GetRef(loop); + return For(loop->loop_var, VisitExpr(loop->min).as_or_throw(), + VisitExpr(loop->extent).as_or_throw(), loop->kind, body, + loop->thread_binding, loop->annotations, loop->step, loop->span); + } + + Stmt VisitStmt_(const WhileNode* loop) final { + Stmt body = VisitStmt(loop->body); + if (IsEvaluateZero(body)) return body; + if (body.same_as(loop->body)) return ffi::GetRef(loop); + return While(VisitExpr(loop->condition).as_or_throw(), body, loop->span); + } + + Stmt VisitStmt_(const IfThenElseNode* branch) final { + PrimExpr condition = VisitExpr(branch->condition).as_or_throw(); + Stmt then_case = VisitStmt(branch->then_case); + if (!branch->else_case.has_value()) { + if (IsEvaluateZero(then_case)) return PreserveConditionEffects(condition); + return IfThenElse(condition, then_case, std::nullopt, branch->span); + } + Stmt else_case = VisitStmt(branch->else_case.value()); + bool empty_then = IsEvaluateZero(then_case); + bool empty_else = IsEvaluateZero(else_case); + if (empty_then && empty_else) return PreserveConditionEffects(condition); + if (empty_else) return IfThenElse(condition, then_case, std::nullopt, branch->span); + if (empty_then) return IfThenElse(!condition, else_case, std::nullopt, branch->span); + return IfThenElse(condition, then_case, else_case, branch->span); + } +}; + +bool IsCudaDeviceFunction(const PrimFunc& function) { + auto target = function->GetAttr(tvm::attr::kTarget); + CallingConv calling_conv = + function->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); + return target.has_value() && target.value()->kind->name == "cuda" && + calling_conv == CallingConv::kDeviceKernelLaunch; +} + +bool IsSm90OrNewer(const PrimFunc& function) { + auto target = function->GetAttr(tvm::attr::kTarget); + if (!target.has_value()) return false; + auto arch = target.value()->GetAttr("arch"); + if (!arch.has_value()) return false; + std::string value = arch.value(); + if (!value.starts_with("sm_")) return false; + size_t end = 3; + while (end < value.size() && value[end] >= '0' && value[end] <= '9') ++end; + if (end == 3) return false; + return std::stoi(value.substr(3, end - 3)) >= 90; +} + +bool HasAnyPayload(const std::map& declarations) { + for (const auto& item : declarations) { + const Declaration& declaration = item.second; + if (declaration.has_payload || + (declaration.end_payload_schema_set && declaration.end_has_payload)) { + return true; + } + } + return false; +} + +std::string FunctionName(const GlobalVar& global_var, const PrimFunc& function) { + if (auto symbol = function->GetAttr(tvm::attr::kGlobalSymbol)) { + return symbol.value(); + } + return global_var->name_hint; +} + +struct KernelIketInfo { + GlobalVar global_var; + PrimFunc function; + std::string name; + std::map declarations; +}; + +void PutOfficialU32(std::vector* bytes, size_t offset, uint32_t value) { + TVM_FFI_ICHECK_LE(offset + sizeof(value), bytes->size()); + for (size_t index = 0; index < sizeof(value); ++index) { + (*bytes)[offset + index] = static_cast(value >> (index * 8)); + } +} + +uint32_t OfficialRangeId(const std::string& name) { + uint32_t value = 2166136261U; + for (uint8_t byte : name) { + value ^= byte; + value *= 16777619U; + } + return value; +} + +std::string OfficialSymbolName(const std::string& name) { + std::string result; + result.reserve(name.size()); + for (uint8_t byte : name) { + bool valid = (byte >= 'a' && byte <= 'z') || (byte >= 'A' && byte <= 'Z') || + (byte >= '0' && byte <= '9') || byte == '_'; + result.push_back(valid ? static_cast(byte) : '_'); + } + return result; +} + +std::string OfficialByteArray(const std::string& symbol, const std::vector& bytes) { + std::ostringstream os; + os << "__device__ __align__(1) unsigned char " << symbol << '[' << bytes.size() << "] = {"; + for (size_t index = 0; index < bytes.size(); ++index) { + if (index) os << ','; + os << static_cast(bytes[index]); + } + os << "};\n"; + return os.str(); +} + +std::map CollectOfficialDeclarations( + const std::vector& kernels) { + std::map declarations; + for (const KernelIketInfo& kernel : kernels) { + for (const auto& [key, declaration] : kernel.declarations) { + auto [it, inserted] = declarations.emplace(key, declaration); + TVM_FFI_ICHECK(inserted || it->second.event_id == declaration.event_id); + } + } + return declarations; +} + +std::string BuildOfficialDeviceSource(const std::vector& kernels) { + std::map declarations = CollectOfficialDeclarations(kernels); + std::ostringstream os; + os << R"IKET( +extern "C" { +)IKET"; + + std::vector meta(kOfficialMetaInfoBytes); + PutOfficialU32(&meta, 0, kOfficialMetaInfoBytes); + PutOfficialU32(&meta, 4, 0); + PutOfficialU32(&meta, 8, 5); + PutOfficialU32(&meta, 12, 31); + PutOfficialU32(&meta, 16, 32); + PutOfficialU32(&meta, 20, kOfficialEventAttributesBytes); + PutOfficialU32(&meta, 24, 0xbabef19dU); + PutOfficialU32(&meta, 28, 0); + PutOfficialU32(&meta, 32, 3); + os << OfficialByteArray("__iket_meta_info", meta); + + std::unordered_map range_ids; + for (const auto& [key, declaration] : declarations) { + uint32_t range_id = key.kind == DeclarationKind::kMark ? 0 : OfficialRangeId(key.name); + if (range_id != 0) { + auto [it, inserted] = range_ids.emplace(range_id, key.name); + TVM_FFI_CHECK(inserted || it->second == key.name, ValueError) + << "NVIDIA IKET range-id collision between " << it->second << " and " << key.name; + } + + std::vector event(kOfficialEventAttributesBytes); + PutOfficialU32(&event, 0, kOfficialEventAttributesBytes); + PutOfficialU32(&event, 4, declaration.event_id); + PutOfficialU32(&event, 8, 3); + PutOfficialU32(&event, 12, 0); + uint32_t event_position = + key.kind == DeclarationKind::kRange ? 4 : (key.kind == DeclarationKind::kPush ? 1 : 0); + PutOfficialU32(&event, 16, event_position); + PutOfficialU32(&event, 20, range_id); + PutOfficialU32(&event, 24, key.name.size()); + std::copy(key.name.begin(), key.name.end(), event.begin() + 28); + std::string event_symbol = "__iket_evt_decl_" + OfficialSymbolName(key.name) + "_" + + std::to_string(declaration.event_id) + "_attrs"; + os << OfficialByteArray(event_symbol, event); + + if (range_id != 0) { + std::vector range(kOfficialRangeAttributesBytes); + PutOfficialU32(&range, 0, kOfficialRangeAttributesBytes); + PutOfficialU32(&range, 4, 0); + PutOfficialU32(&range, 8, range_id); + PutOfficialU32(&range, 12, 0xffffffffU); + PutOfficialU32(&range, 16, key.kind == DeclarationKind::kRange ? 1 : 2); + PutOfficialU32(&range, 20, key.kind == DeclarationKind::kRange ? 1 : 0); + PutOfficialU32(&range, 32, key.name.size()); + std::copy(key.name.begin(), key.name.end(), range.begin() + 36); + std::string range_symbol = "__iket_range_decl_" + OfficialSymbolName(key.name) + "_" + + std::to_string(range_id) + "_attrs"; + os << OfficialByteArray(range_symbol, range); + } + } + os << R"IKET( +} + +template +__forceinline__ __device__ void tvm_builtin_iket_official_event_impl() { + asm volatile( + "{\n" + ".reg .b32 %%r, %%t;\n" + "mov.b32 %%r, %%cluster_ctarank;\n" + "mov.u32 %%t, %%globaltimer_lo;\n" + "or.b32 %%t, %%t, %0;\n" + "mad.lo.u32 %%r, %%r, 0x1000000, 0x20;\n" + "st.weak.shared.u32 [%%r], %%t;\n" + "pmevent.mask %0;\n" + "}\n" + : + : "n"(EventId) + : "memory"); +} + +__forceinline__ __device__ unsigned int tvm_builtin_iket_official_event( + unsigned int event_id) { + switch (event_id) { +)IKET"; + for (const auto& [key, declaration] : declarations) { + os << " case " << declaration.event_id << ":\n" + << " tvm_builtin_iket_official_event_impl<" << declaration.event_id << ">();\n" + << " break;\n"; + } + os << R"IKET( case 31: + tvm_builtin_iket_official_event_impl<31>(); + break; + default: + break; + } + return event_id; +} +)IKET"; + return os.str(); +} + +using UniformBufferSet = std::unordered_set; +using DivergentVarSet = std::unordered_set; + +UniformBufferSet IntersectUniformBuffers(const UniformBufferSet& lhs, const UniformBufferSet& rhs) { + UniformBufferSet result; + for (const VarNode* buffer : lhs) { + if (rhs.count(buffer)) result.insert(buffer); + } + return result; +} + +DivergentVarSet UnionDivergentVars(const DivergentVarSet& lhs, const DivergentVarSet& rhs) { + DivergentVarSet result = lhs; + result.insert(rhs.begin(), rhs.end()); + return result; +} + +class UniformExprChecker : public ExprVisitor { + public: + UniformExprChecker(const DivergentVarSet& divergent_vars, const UniformBufferSet& uniform_buffers) + : divergent_vars_(divergent_vars), uniform_buffers_(uniform_buffers) {} + + bool IsUniform(const Expr& expr) { + uniform_ = true; + operator()(expr); + return uniform_; + } + + private: + void VisitExpr_(const VarNode* var) final { + if (divergent_vars_.count(var)) uniform_ = false; + } + + void VisitExpr_(const BufferLoadNode* load) final { + if (!uniform_buffers_.count(load->buffer.get()) || + !IsScalarBufferAccess(load->buffer, load->indices)) { + uniform_ = false; + return; + } + for (const PrimExpr& index : load->indices) VisitExpr(index); + } + + void VisitExpr_(const CallNode* call) final { + // Calls may read lane-local or device state even when their explicit + // arguments are uniform. Treat them conservatively, except for likely(), + // which is only an annotation around its argument. + if (call->op.same_as(builtin::likely())) { + ExprVisitor::VisitExpr_(call); + } else if (IsTokenProducer(call)) { + // Both producers return a declaration id (or sentinel zero) that is + // independent of a potentially lane-varying payload. + return; + } else if (call->op.same_as(builtin::tvm_warp_shuffle()) && call->args.size() == 5) { + // A shuffle from one warp-uniform source lane is a broadcast. Its + // value may depend on threadIdx, but every active lane observes the + // selected lane's value. + VisitExpr(call->args[0]); + for (size_t i = 2; i < call->args.size(); ++i) VisitExpr(call->args[i]); + } else if (call->op.same_as(Op::Get("tirx.cuda.__shfl_sync")) && call->args.size() == 4) { + // CUDA's explicit __shfl_sync(mask, value, src_lane, width) has the + // same broadcast semantics when mask/src_lane/width are uniform. + // Ignore the lane-local value, but prove the control operands uniform. + VisitExpr(call->args[0]); + VisitExpr(call->args[2]); + VisitExpr(call->args[3]); + } else if (call->op.same_as(builtin::bitwise_and()) || + call->op.same_as(builtin::bitwise_or()) || + call->op.same_as(builtin::bitwise_xor()) || + call->op.same_as(builtin::bitwise_not())) { + // These integer/boolean operators are pure. A composed guard remains + // uniform exactly when each operand is uniform. + ExprVisitor::VisitExpr_(call); + } else { + uniform_ = false; + } + } + + const DivergentVarSet& divergent_vars_; + const UniformBufferSet& uniform_buffers_; + bool uniform_{true}; +}; + +class LoopControlFinder : public StmtExprVisitor { + public: + bool found{false}; + + private: + void VisitStmt_(const BreakNode* op) final { found = true; } + void VisitStmt_(const ContinueNode* op) final { found = true; } + void VisitExpr_(const CallNode* call) final { + if (call->op.same_as(builtin::break_loop()) || call->op.same_as(builtin::continue_loop())) { + found = true; + return; + } + StmtExprVisitor::VisitExpr_(call); + } +}; + +class AnnotationFinder : public StmtExprVisitor { + public: + bool found{false}; + + private: + void VisitExpr_(const CallNode* call) final { + if (IsIketOp(call->op)) { + found = true; + return; + } + StmtExprVisitor::VisitExpr_(call); + } +}; + +class IketConvergenceVerifier : public StmtExprVisitor { + private: + bool IsUniform(const Expr& expr) const { + return UniformExprChecker(divergent_vars_, uniform_buffers_).IsUniform(expr); + } + + void VisitStmt_(const BufferStoreNode* op) final { + bool uniform_store = !divergent_context_ && IsScalarBufferAccess(op->buffer, op->indices) && + IsUniform(op->value); + for (const PrimExpr& index : op->indices) { + uniform_store = uniform_store && IsUniform(index); + VisitExpr(index); + } + VisitExpr(op->value); + if (uniform_store) { + uniform_buffers_.insert(op->buffer.get()); + } else { + uniform_buffers_.erase(op->buffer.get()); + } + } + + void VisitStmt_(const AttrStmtNode* op) final { + const VarNode* thread_var = nullptr; + if (op->attr_key == attr::thread_extent) { + std::string thread_tag; + if (auto iter_var = op->node.as()) { + thread_tag = iter_var.value()->thread_tag; + thread_var = iter_var.value()->var.get(); + } else if (auto var = op->node.as()) { + thread_tag = var.value()->name; + thread_var = var.value().get(); + } + if (!thread_tag.starts_with("threadIdx.")) thread_var = nullptr; + } + bool inserted = thread_var && divergent_vars_.insert(thread_var).second; + VisitExpr(op->value); + VisitStmt(op->body); + if (inserted) divergent_vars_.erase(thread_var); + } + + void VisitStmt_(const BindNode* op) final { + if (!IsUniform(op->value)) divergent_vars_.insert(op->var.get()); + VisitExpr(op->value); + } + + void VisitStmt_(const IfThenElseNode* op) final { + VisitExpr(op->condition); + bool old_divergent = divergent_context_; + bool condition_uniform = IsUniform(op->condition); + divergent_context_ = divergent_context_ || !condition_uniform; + auto old_vars = divergent_vars_; + auto old_buffers = uniform_buffers_; + VisitStmt(op->then_case); + auto then_buffers = uniform_buffers_; + divergent_vars_ = old_vars; + uniform_buffers_ = old_buffers; + if (op->else_case.has_value()) { + VisitStmt(op->else_case.value()); + } + uniform_buffers_ = IntersectUniformBuffers(then_buffers, uniform_buffers_); + divergent_vars_ = std::move(old_vars); + divergent_context_ = old_divergent; + } + + void VisitStmt_(const ForNode* op) final { + VisitExpr(op->min); + VisitExpr(op->extent); + bool uniform_loop = IsUniform(op->min) && IsUniform(op->extent); + AnnotationFinder annotations; + annotations(op->body); + LoopControlFinder loop_control; + loop_control(op->body); + TVM_FFI_CHECK(!(annotations.found && loop_control.found), ValueError) + << "IKET event sites are not allowed in a loop containing break or continue"; + + bool old_divergent = divergent_context_; + DivergentVarSet old_vars = divergent_vars_; + UniformBufferSet old_buffers = uniform_buffers_; + DivergentVarSet loop_entry_vars = old_vars; + if (!uniform_loop) loop_entry_vars.insert(op->loop_var.get()); + + DivergentVarSet loop_head_vars = loop_entry_vars; + UniformBufferSet loop_head_buffers = old_buffers; + bool converged = false; + for (size_t iteration = 0; iteration < kMaxConvergenceAnalysisIterations; ++iteration) { + divergent_vars_ = loop_head_vars; + uniform_buffers_ = loop_head_buffers; + divergent_context_ = old_divergent || !uniform_loop; + VisitStmt(op->body); + + DivergentVarSet next_vars = UnionDivergentVars(loop_entry_vars, divergent_vars_); + UniformBufferSet next_buffers = IntersectUniformBuffers(old_buffers, uniform_buffers_); + if (next_vars == loop_head_vars && next_buffers == loop_head_buffers) { + converged = true; + break; + } + loop_head_vars = std::move(next_vars); + loop_head_buffers = std::move(next_buffers); + } + TVM_FFI_CHECK(converged, ValueError) + << "IKET convergence analysis did not reach a loop fixed point"; + uniform_buffers_ = std::move(loop_head_buffers); + divergent_vars_ = std::move(old_vars); + divergent_context_ = old_divergent; + } + + void VisitStmt_(const WhileNode* op) final { + AnnotationFinder annotations; + annotations(op->body); + LoopControlFinder loop_control; + loop_control(op->body); + TVM_FFI_CHECK(!(annotations.found && loop_control.found), ValueError) + << "IKET event sites are not allowed in a loop containing break or continue"; + bool old_divergent = divergent_context_; + DivergentVarSet old_vars = divergent_vars_; + UniformBufferSet old_buffers = uniform_buffers_; + DivergentVarSet loop_head_vars = old_vars; + UniformBufferSet loop_head_buffers = old_buffers; + bool converged = false; + for (size_t iteration = 0; iteration < kMaxConvergenceAnalysisIterations; ++iteration) { + divergent_vars_ = loop_head_vars; + uniform_buffers_ = loop_head_buffers; + VisitExpr(op->condition); + divergent_context_ = old_divergent || !IsUniform(op->condition); + VisitStmt(op->body); + + DivergentVarSet next_vars = UnionDivergentVars(old_vars, divergent_vars_); + UniformBufferSet next_buffers = IntersectUniformBuffers(old_buffers, uniform_buffers_); + if (next_vars == loop_head_vars && next_buffers == loop_head_buffers) { + converged = true; + break; + } + loop_head_vars = std::move(next_vars); + loop_head_buffers = std::move(next_buffers); + } + TVM_FFI_CHECK(converged, ValueError) + << "IKET convergence analysis did not reach a loop fixed point"; + uniform_buffers_ = std::move(loop_head_buffers); + divergent_vars_ = std::move(old_vars); + divergent_context_ = old_divergent; + } + + void VisitExpr_(const CallNode* call) final { + if (IsIketOp(call->op)) { + if (divergent_context_) { + TVM_FFI_THROW(ValueError) << "IKET event site may be reached by a divergent set of lanes: " + << GetRef(call); + } + if (call->op.same_as(IketRangeEndOp())) { + TVM_FFI_CHECK(IsUniform(call->args[0]), ValueError) + << "IKET range_end requires a warp-uniform RangeToken"; + } + } + StmtExprVisitor::VisitExpr_(call); + } + + DivergentVarSet divergent_vars_; + UniformBufferSet uniform_buffers_; + bool divergent_context_{false}; +}; + +class InstrumentOfficialKernel : public StmtExprMutator { + public: + InstrumentOfficialKernel(const KernelIketInfo& info, std::string device_source) + : info_(info), device_source_(std::move(device_source)) {} + + PrimFunc Run() { + PrimFunc result = info_.function; + result.CopyOnWrite()->body = operator()(info_.function->body); + return result; + } + + private: + const Declaration& Lookup(DeclarationKind kind, const CallNode* call) const { + DeclarationKey key{kind, GetName(call)}; + auto it = info_.declarations.find(key); + TVM_FFI_ICHECK(it != info_.declarations.end()) << "Missing official IKET declaration"; + return it->second; + } + + PrimExpr Event(PrimExpr event_id) const { + static const Op& event_op = Op::Get("tirx.cuda.iket_official_event"); + return Call(PrimType::UInt(32), event_op, + {cast(PrimType::UInt(32), event_id), StringImm(device_source_)}); + } + + Stmt VisitStmt_(const EvaluateNode* evaluate) final { + if (const auto* call = evaluate->value.as(); + call && call->op.same_as(IketRangeEndOp())) { + PrimExpr token = VisitExpr(call->args[0]).as_or_throw(); + return Evaluate(Event(token)); + } + return StmtExprMutator::VisitStmt_(evaluate); + } + + Expr VisitExpr_(const CallNode* call) final { + if (call->op.same_as(IketRangeStartOp())) { + return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kRange, call).event_id)); + } + if (call->op.same_as(IketSentinelOp())) return IntImm(PrimType::UInt(32), 0); + if (call->op.same_as(IketMarkOp())) { + return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kMark, call).event_id)); + } + if (call->op.same_as(IketRangePushOp())) { + return Event(IntImm(PrimType::UInt(32), Lookup(DeclarationKind::kPush, call).event_id)); + } + if (call->op.same_as(IketRangePopOp())) { + return Event(IntImm(PrimType::UInt(32), 31)); + } + if (call->op.same_as(IketRangeEndOp())) { + TVM_FFI_THROW(ValueError) << "range_end must be emitted in statement position"; + } + return StmtExprMutator::VisitExpr_(call); + } + + const KernelIketInfo& info_; + std::string device_source_; +}; + +bool IketEnabled(const IRModule& module) { return module->HasNonzeroAttr("tirx.iket.enabled"); } + +IRModule LowerIketImpl(IRModule module) { + if (!IketEnabled(module)) { + for (const auto& [global_var, base_function] : module->functions) { + const auto* prim_func = base_function.as(); + if (!prim_func) continue; + PrimFunc function = ffi::GetRef(prim_func); + AnnotationCollector collector; + collector(function->body); + TokenBufferSet tokens = CollectTokenBuffers(function->body); + if (!collector.has_annotations && tokens.empty()) continue; + if (collector.has_annotations) { + TokenVerifier verifier(tokens); + verifier(function->body); + TokenDeclarationMap token_declarations = CollectTokenDeclarations(function->body); + RangeEndSchemaVerifier schema_verifier(token_declarations, &collector.declarations); + schema_verifier(function->body); + } + StripIket strip(std::move(tokens)); + Stmt body = RemoveStrippedIketNoOps()(strip(function->body)); + if (!body.same_as(function->body)) { + function.CopyOnWrite()->body = body; + module->Update(global_var, function); + } + } + return module; + } + + std::vector kernels; + for (const auto& [global_var, base_function] : module->functions) { + const auto* prim_func = base_function.as(); + if (!prim_func) continue; + PrimFunc function = ffi::GetRef(prim_func); + AnnotationCollector collector; + collector(function->body); + if (!collector.has_annotations) continue; + + std::string function_name = FunctionName(global_var, function); + TVM_FFI_CHECK(IsCudaDeviceFunction(function), ValueError) + << "IKET annotations are only valid in a split CUDA device kernel"; + TVM_FFI_CHECK_LE(collector.declarations.size(), kMaxDeclarations, ValueError) + << "NVIDIA IKET supports at most " << kMaxDeclarations << " declarations per kernel"; + + TokenBufferSet tokens = CollectTokenBuffers(function->body); + TokenVerifier verifier(tokens); + verifier(function->body); + TokenDeclarationMap token_declarations = CollectTokenDeclarations(function->body); + RangeEndSchemaVerifier schema_verifier(token_declarations, &collector.declarations); + schema_verifier(function->body); + IketConvergenceVerifier convergence_verifier; + convergence_verifier(function->body); + + TVM_FFI_CHECK(IsSm90OrNewer(function), ValueError) + << "NVIDIA IKET requires SM90 or newer for kernel " << function_name; + TVM_FFI_CHECK(!HasAnyPayload(collector.declarations), ValueError) + << "NVIDIA IKET does not support payloads in kernel " << function_name; + TVM_FFI_CHECK(OfficialTokenAnalyzer(tokens).Prove(function->body), ValueError) + << "NVIDIA IKET requires token ranges to be provably strictly alternating " + "and closed on every exit in kernel " + << function_name; + TVM_FFI_CHECK(OfficialStackAnalyzer().Prove(function->body), ValueError) + << "NVIDIA IKET requires balanced range_push/range_pop paths in kernel " << function_name; + + kernels.push_back(KernelIketInfo{global_var, function, std::move(function_name), + std::move(collector.declarations)}); + } + + std::sort( + kernels.begin(), kernels.end(), + [](const KernelIketInfo& lhs, const KernelIketInfo& rhs) { return lhs.name < rhs.name; }); + for (size_t i = 1; i < kernels.size(); ++i) { + TVM_FFI_CHECK_NE(kernels[i - 1].name, kernels[i].name, ValueError) + << "IKET device kernels must have unique global symbols: " << kernels[i].name; + } + + std::map event_ids; + std::unordered_map event_kinds; + for (const KernelIketInfo& kernel : kernels) { + for (const auto& [key, declaration] : kernel.declarations) { + auto [kind_it, kind_inserted] = event_kinds.emplace(key.name, key.kind); + TVM_FFI_CHECK(kind_inserted || kind_it->second == key.kind, ValueError) + << "NVIDIA IKET declaration " << key.name << " changes event kind across kernels"; + event_ids.emplace(key, 0); + } + } + TVM_FFI_CHECK_LE(event_ids.size(), kMaxDeclarations, ValueError) + << "NVIDIA IKET supports at most " << kMaxDeclarations + << " distinct declarations in one CUDA module"; + + uint32_t event_id = 1; + for (auto& [key, id] : event_ids) id = event_id++; + for (KernelIketInfo& kernel : kernels) { + for (auto& [key, declaration] : kernel.declarations) { + declaration.event_id = event_ids.at(key); + } + } + + std::string device_source = BuildOfficialDeviceSource(kernels); + for (const KernelIketInfo& kernel : kernels) { + module->Update(kernel.global_var, InstrumentOfficialKernel(kernel, device_source).Run()); + } + return module; +} + +} // namespace + +Pass LowerIket() { + auto pass_func = [](IRModule module, tvm::transform::PassContext) { + return LowerIketImpl(std::move(module)); + }; + return tvm::transform::CreateModulePass(pass_func, 0, "tirx.backend.cuda.LowerIket", {}); +} + +TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; + refl::GlobalDef().def("tirx.backend.cuda.transforms.LowerIket", LowerIket); +} + +} // namespace transform +} // namespace tirx +} // namespace tvm diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 2241d22f5911..7c6f069d57ab 100644 --- a/src/backend/trn/transform/lower_trainium_layout.cc +++ b/src/backend/trn/transform/lower_trainium_layout.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc index b330fd3c0385..8626f96de149 100644 --- a/src/relax/script/printer/call.cc +++ b/src/relax/script/printer/call.cc @@ -240,9 +240,6 @@ ffi::Optional PrintRelaxPrint(const Call& n, const AccessPath& n_p, con } bool ShouldPrintAsTIR(const Call& call) { - if (!call->ty.as()) { - return false; - } if (call->op->ty.as() || call->op.as() || call->op.as() || call->op.as()) { return false; @@ -250,6 +247,9 @@ bool ShouldPrintAsTIR(const Call& call) { if (auto op = call->op.as()) { return op.value()->name.find("relax.") != 0; } + if (!call->ty.as()) { + return false; + } return true; } diff --git a/src/runtime/extra/contrib/nvshmem/dist_gemm.cu b/src/runtime/extra/contrib/nvshmem/dist_gemm.cu index 26cc55a44218..b772cd7eaa1e 100644 --- a/src/runtime/extra/contrib/nvshmem/dist_gemm.cu +++ b/src/runtime/extra/contrib/nvshmem/dist_gemm.cu @@ -22,7 +22,8 @@ #include #include #include -#include +#include +#include namespace tvm { namespace runtime { @@ -57,24 +58,28 @@ void copy_to_peer(void* dst, int dst_device, void* src, size_t size, TVMStreamHa cudaMemcpyAsync(remote_dst, src, size, cudaMemcpyDefault, CUstream(stream)); } +Device current_cuda_device() { + int device_id; + TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&device_id)); + return Device{DLDeviceType::kDLCUDA, device_id}; +} + TVMStreamHandle stream_create() { - DiscoWorker* worker = ThreadLocalDiscoWorker::Get()->worker; - TVM_FFI_ICHECK(worker != nullptr) << "NVSHMEM stream creation failed: worker is not initialized"; - cudaStream_t retval; - TVM_FFI_CHECK_CUDA_ERROR(cudaStreamCreateWithFlags(&retval, cudaStreamNonBlocking)); - return static_cast(retval); + Device device = current_cuda_device(); + return DeviceAPI::Get(device)->CreateStream(device); } void stream_sync(TVMStreamHandle from_stream, TVMStreamHandle to_stream) { - DiscoWorker* worker = ThreadLocalDiscoWorker::Get()->worker; - TVM_FFI_ICHECK(worker != nullptr) << "NVSHMEM stream sync failed: worker is not initialized"; - auto f_sync_stream = tvm::ffi::Function::GetGlobalRequired("runtime.Device_StreamSyncFromTo"); - f_sync_stream(worker->default_device, reinterpret_cast(from_stream), - reinterpret_cast(to_stream)); + Device device = current_cuda_device(); + DeviceAPI::Get(device)->SyncStreamFromTo(device, from_stream, to_stream); +} + +TVMStreamHandle stream_from_int(int64_t stream) { + return reinterpret_cast(stream); } -void set_streaming_policy(TVMStreamHandle stream, void* ptr, size_t size) { - cudaStream_t strm = static_cast(stream); +void set_streaming_policy(int64_t stream, void* ptr, size_t size) { + cudaStream_t strm = static_cast(stream_from_int(stream)); struct cudaAccessPolicyWindow accessPolicyWindow = {ptr, size, 0.0, cudaAccessPropertyStreaming, cudaAccessPropertyStreaming}; cudaStreamAttrValue streamAttrValue; @@ -83,11 +88,10 @@ void set_streaming_policy(TVMStreamHandle stream, void* ptr, size_t size) { } void transfer_to_peers_reduce_scatter(Tensor semaphore, Tensor gemm_out, Tensor staging_buffer, - TVMStreamHandle stream, int32_t M, int32_t N, int32_t BLK_M, + int64_t stream_handle, int32_t M, int32_t N, int32_t BLK_M, int32_t BLK_N, int32_t WORLD_SIZE) { - DiscoWorker* worker = ThreadLocalDiscoWorker::Get()->worker; - TVM_FFI_ICHECK(worker != nullptr) << "NVSHMEM transfer to peer failed: worker is not initialized"; - int my_rank = worker->worker_id; + TVMStreamHandle stream = stream_from_int(stream_handle); + int my_rank = nvshmem_my_pe(); int LOCAL_M = M / WORLD_SIZE; for (int i = 0; i < WORLD_SIZE; i++) { int to_rank = (my_rank + i + 1) % WORLD_SIZE; @@ -108,11 +112,10 @@ void transfer_to_peers_reduce_scatter(Tensor semaphore, Tensor gemm_out, Tensor } } -void transfer_to_peers_all_gather(Tensor semaphore, Tensor A, Tensor ag_out, TVMStreamHandle stream, +void transfer_to_peers_all_gather(Tensor semaphore, Tensor A, Tensor ag_out, int64_t stream_handle, int32_t M, int32_t K, int32_t WORLD_SIZE) { - DiscoWorker* worker = ThreadLocalDiscoWorker::Get()->worker; - TVM_FFI_ICHECK(worker != nullptr) << "NVSHMEM transfer to peer failed: worker is not initialized"; - int my_rank = worker->worker_id; + TVMStreamHandle stream = stream_from_int(stream_handle); + int my_rank = nvshmem_my_pe(); int LOCAL_M = M / WORLD_SIZE; for (int i = 0; i < WORLD_SIZE; i++) { int to_rank = (my_rank + WORLD_SIZE - i - 1) % WORLD_SIZE; @@ -132,10 +135,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("runtime.disco.stream_sync", stream_sync) .def("runtime.disco.transfer_to_peers_reduce_scatter", transfer_to_peers_reduce_scatter) .def("runtime.disco.transfer_to_peers_all_gather", transfer_to_peers_all_gather) - .def("runtime.disco.set_streaming_policy", - [](TVMStreamHandle stream, Tensor ptr, size_t size) { - set_streaming_policy(stream, ptr->data, size); - }); + .def("runtime.disco.set_streaming_policy", [](int64_t stream, Tensor ptr, size_t size) { + set_streaming_policy(stream, ptr->data, size); + }); } } // namespace runtime diff --git a/src/runtime/extra/contrib/nvshmem/init.cc b/src/runtime/extra/contrib/nvshmem/init.cc index aeb71bb6a19a..b110e08994eb 100644 --- a/src/runtime/extra/contrib/nvshmem/init.cc +++ b/src/runtime/extra/contrib/nvshmem/init.cc @@ -137,8 +137,8 @@ void NVSHMEMXCumoduleInit(void* cuModule) { } } -void NVSHMEMBarrierAllOnStream(TVMStreamHandle stream) { - CUstream strm = static_cast(stream); +void NVSHMEMBarrierAllOnStream(int64_t stream) { + CUstream strm = reinterpret_cast(stream); nvshmemx_barrier_all_on_stream(strm); } @@ -154,7 +154,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { int device_id; TVM_FFI_CHECK_CUDA_ERROR(cudaGetDevice(&device_id)); TVMStreamHandle stream = TVMFFIEnvGetStream(kDLCUDA, device_id); - NVSHMEMBarrierAllOnStream(stream); + NVSHMEMBarrierAllOnStream(reinterpret_cast(stream)); }); } diff --git a/src/runtime/extra/contrib/nvshmem/memory_allocator.cc b/src/runtime/extra/contrib/nvshmem/memory_allocator.cc index 1483563b6200..d92e905e09d2 100644 --- a/src/runtime/extra/contrib/nvshmem/memory_allocator.cc +++ b/src/runtime/extra/contrib/nvshmem/memory_allocator.cc @@ -88,7 +88,9 @@ class NVSHMEMAllocator final : public PooledAllocator { }; Tensor NVSHMEMEmpty(ffi::Shape shape, DLDataType dtype, ffi::Optional device) { - return NVSHMEMAllocator::Global()->Empty(shape, dtype, UseDefaultDeviceIfNone(device)); + Device resolved_device = + device.has_value() ? device.value() : UseDefaultDeviceIfNone(std::nullopt); + return NVSHMEMAllocator::Global()->Empty(shape, dtype, resolved_device); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/src/script/printer/doc_printer/python_doc_printer.cc b/src/script/printer/doc_printer/python_doc_printer.cc index 151125a20d3a..13ccbde1eabc 100644 --- a/src/script/printer/doc_printer/python_doc_printer.cc +++ b/src/script/printer/doc_printer/python_doc_printer.cc @@ -20,7 +20,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc index 423f239e9a31..da8a6cb84921 100644 --- a/src/target/source/codegen_c.cc +++ b/src/target/source/codegen_c.cc @@ -413,8 +413,6 @@ void CodeGenC::RegisterHandleType(const VarNode* buf_var, const PrimType& t) { void CodeGenC::RegisterHandleTypeFromPointer(const tirx::Var& var, const Expr* value) { if (value == nullptr) return; - auto* call = value->as(); - if (call == nullptr || !call->op.same_as(builtin::ptr_byte_offset())) return; std::optional value_dtype = [&]() { if (auto prim_value = value->as()) { return tirx::GetPointerType(GetType(prim_value.value())); @@ -422,8 +420,11 @@ void CodeGenC::RegisterHandleTypeFromPointer(const tirx::Var& var, const Expr* v return tirx::GetPointerType((*value)->ty); }(); if (!value_dtype.has_value()) return; + auto* call = value->as(); + if (call != nullptr && call->op.same_as(builtin::ptr_byte_offset())) { + pointer_offset_vars_.insert(var.get()); + } RegisterHandleType(var.get(), value_dtype.value()); - pointer_offset_vars_.insert(var.get()); } void CodeGenC::PrintVecElemLoad(const std::string& vec, const PrimType& t, int i, diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc b/src/tirx/analysis/verify_tirx_well_formed.cc index 67c8424cba41..889aba164f54 100644 --- a/src/tirx/analysis/verify_tirx_well_formed.cc +++ b/src/tirx/analysis/verify_tirx_well_formed.cc @@ -30,7 +30,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/tirx/ir/async_structs.cc b/src/tirx/ir/async_structs.cc deleted file mode 100644 index 0213468fb100..000000000000 --- a/src/tirx/ir/async_structs.cc +++ /dev/null @@ -1,87 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ - -/*! - * \file async_structs.cc - */ - -#include -#include -#include - -namespace tvm { -namespace tirx { - -TVM_FFI_STATIC_INIT_BLOCK() { - PipelineNode::RegisterReflection(); - CopyPipelineNode::RegisterReflection(); -} - -/*************************** Pipeline ***************************/ - -Pipeline::Pipeline(ExecScope thread_scope, size_t depth, bool separate_pc, ffi::String name_hint, - ffi::Map workspace, - ffi::Map schedule_config) { - auto n = ffi::make_object(); - n->thread_scope = std::move(thread_scope); - n->name_hint = std::move(name_hint); - n->depth = depth; - n->separate_pc = separate_pc; - n->workspace = std::move(workspace); - n->schedule_config = std::move(schedule_config); - data_ = std::move(n); -} - -TVM_FFI_STATIC_INIT_BLOCK() { - namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "tirx.Pipeline", [](ExecScope thread_scope, size_t depth, bool separate_pc, - ffi::String name_hint, ffi::Map workspace, - ffi::Map schedule_config) { - return Pipeline(thread_scope, depth, separate_pc, name_hint, workspace, schedule_config); - }); -} - -/*************************** CopyPipeline ***************************/ - -CopyPipeline::CopyPipeline(ExecScope thread_scope, size_t depth, bool separate_pc, - ffi::String name_hint, ffi::Map workspace, - ffi::Map schedule_config) { - auto n = ffi::make_object(); - n->thread_scope = std::move(thread_scope); - n->name_hint = std::move(name_hint); - n->depth = depth; - n->separate_pc = separate_pc; - n->workspace = std::move(workspace); - n->schedule_config = std::move(schedule_config); - data_ = std::move(n); -} - -TVM_FFI_STATIC_INIT_BLOCK() { - namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.CopyPipeline", [](ExecScope thread_scope, size_t depth, - bool separate_pc, ffi::String name_hint, - ffi::Map workspace, - ffi::Map schedule_config) { - return CopyPipeline(thread_scope, depth, separate_pc, name_hint, workspace, schedule_config); - }); -} - -} // namespace tirx -} // namespace tvm diff --git a/src/tirx/ir/predicate.cc b/src/tirx/ir/lambdaexpr.cc similarity index 64% rename from src/tirx/ir/predicate.cc rename to src/tirx/ir/lambdaexpr.cc index f687016c82c2..fea7191a174f 100644 --- a/src/tirx/ir/predicate.cc +++ b/src/tirx/ir/lambdaexpr.cc @@ -18,17 +18,19 @@ */ /*! - * \file predicate.cc + * \file lambdaexpr.cc + * \brief Implementation of LambdaExpr, a reified lambda used by tile primitive ops. */ -#include "tvm/tirx/predicate.h" +#include "tvm/tirx/stmt_functor.h" +#include "tvm/tirx/tile_primitive.h" namespace tvm { namespace tirx { -TVM_FFI_STATIC_INIT_BLOCK() { PredicateNode::RegisterReflection(); } +TVM_FFI_STATIC_INIT_BLOCK() { LambdaExprNode::RegisterReflection(); } -PrimExpr PredicateNode::Apply(const ffi::Array& indices) const { +PrimExpr LambdaExprNode::Apply(const ffi::Array& indices) const { TVM_FFI_ICHECK_EQ(indices.size(), vars.size()); ffi::Map vmap; @@ -37,12 +39,11 @@ PrimExpr PredicateNode::Apply(const ffi::Array& indices) const { vmap.Set(vars[i], indices[i]); } - return SubstituteWithDataTypeLegalization(std::move(pred), - [&](const Var& var) { return vmap.Get(var); }); + return Substitute(std::move(pred), vmap); } -Predicate::Predicate(ffi::Array vars, PrimExpr pred) { - auto n = ffi::make_object(); +LambdaExpr::LambdaExpr(ffi::Array vars, PrimExpr pred) { + auto n = ffi::make_object(); n->vars = std::move(vars); n->pred = std::move(pred); data_ = std::move(n); @@ -50,14 +51,13 @@ Predicate::Predicate(ffi::Array vars, PrimExpr pred) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.Predicate", [](ffi::Array vars, PrimExpr pred) { - return Predicate(vars, pred); - }); + refl::GlobalDef().def("tirx.LambdaExpr", + [](ffi::Array vars, PrimExpr pred) { return LambdaExpr(vars, pred); }); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.PredicateApply", [](Predicate pred, ffi::Array indices) { + refl::GlobalDef().def("tirx.LambdaExprApply", [](LambdaExpr pred, ffi::Array indices) { return pred->Apply(indices); }); } diff --git a/src/tirx/ir/layout/compose_layout.cc b/src/tirx/ir/layout/compose_layout.cc index 7ae3c1a2a35b..866ac483b73a 100644 --- a/src/tirx/ir/layout/compose_layout.cc +++ b/src/tirx/ir/layout/compose_layout.cc @@ -22,29 +22,36 @@ namespace tvm { namespace tirx { /**************** ComposeLayout ****************/ -ComposeLayout::ComposeLayout(SwizzleLayout layout_A, TileLayout layout_B) { +ComposeLayout::ComposeLayout(int per_element, int swizzle_len, int atom_len, TileLayout tile_layout, + bool swizzle_inner) { auto n = ffi::make_object(); - n->swizzle = layout_A; - n->tile_layout = layout_B; + n->per_element = per_element; + n->swizzle_len = swizzle_len; + n->atom_len = atom_len; + n->swizzle_inner = swizzle_inner; + n->tile_layout = std::move(tile_layout); TVM_FFI_ICHECK(n->VerifyWellFormed()) << "ValueError: The compose layout is not well-formed"; - + int swizzle_mask = (1 << swizzle_len) - 1; + n->inner_mask = swizzle_mask; + n->outer_mask = swizzle_mask << atom_len; data_ = std::move(n); } TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.ComposeLayout", [](SwizzleLayout layout_A, TileLayout layout_B) { - return ComposeLayout(layout_A, layout_B); + refl::GlobalDef().def("tirx.ComposeLayout", [](int per_element, int swizzle_len, int atom_len, + TileLayout tile_layout, bool swizzle_inner) { + return ComposeLayout(per_element, swizzle_len, atom_len, tile_layout, swizzle_inner); }); } bool ComposeLayoutNode::CompatibleWithShape(const Array& shape) const { return true; } bool ComposeLayoutNode::VerifyWellFormed() const { - if (!swizzle->VerifyWellFormed() || !tile_layout->VerifyWellFormed()) { + if (!(per_element >= 0 && swizzle_len >= 0 && atom_len >= swizzle_len)) { return false; } - return true; + return tile_layout->VerifyWellFormed(); } PrimExpr ComposeLayoutNode::GetSize(ffi::Optional axis_name) const { @@ -68,32 +75,60 @@ ffi::Map ComposeLayoutNode::Apply(PrimExpr coord) const { auto res = tile_layout->Apply(coord); TVM_FFI_ICHECK(res.size() == 1 && res.find("m") != res.end()); auto m = res["m"]; - auto swizzle_res = swizzle->Apply(m); - TVM_FFI_ICHECK(swizzle_res.size() == 1 && swizzle_res.find("m") != swizzle_res.end()); - return swizzle_res; + // Inline the swizzle XOR (formerly SwizzleLayoutNode::Apply): the swizzle + // operates on the tile-mapped coordinate ``m``. + auto f = [&](const PrimExpr& x) -> PrimExpr { + if (swizzle_inner) { + return x ^ ((x & outer_mask) >> atom_len); + } else { + return x ^ ((x & inner_mask) << atom_len); + } + }; + auto base = 1 << per_element; + arith::Analyzer analyzer; + return {{"m", analyzer->Simplify((f(floordiv(m, base)) << per_element) + floormod(m, base))}}; } Layout ComposeLayoutNode::Canonicalize() const { auto tile_normalized = tile_layout->Canonicalize().as().value(); - if (tile_normalized->IsTrivial()) { - return swizzle; - } - return ComposeLayout(swizzle, tile_normalized); + return ComposeLayout(per_element, swizzle_len, atom_len, tile_normalized, swizzle_inner); } Layout ComposeLayoutNode::Tile(const TileLayout& outer, const ffi::Array& outer_shape, const ffi::Array& inner_shape) const { - // layout_B is first tiled with `outer`, then compose with layout_A. - auto tiled_B = tile_layout->Tile(outer, outer_shape, inner_shape).as().value(); - return ComposeLayout(swizzle, tiled_B); + // A bare swizzle (ComposeLayout with a trivial tile) carries only the swizzle + // period, not a tile matching `inner_shape`; substitute an identity tile over + // the inner product first, matching the former SwizzleLayoutNode::Tile. + TileLayout base = this->tile_layout; + if (base->IsTrivial()) { + base = IdentityTileLayout(inner_shape); + } + auto tiled_B = base->Tile(outer, outer_shape, inner_shape).as().value(); + return ComposeLayout(per_element, swizzle_len, atom_len, tiled_B, swizzle_inner); } ffi::Optional ComposeLayoutNode::IsTileInner( const Layout& tile_layout, const ffi::Array& tiled_shape, const ffi::Array& inner_shape) const { if (auto comp = tile_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, this->swizzle)) { - return this->tile_layout->IsTileInner(comp.value()->tile_layout, tiled_shape, inner_shape); + if (comp.value()->per_element == this->per_element && + comp.value()->swizzle_len == this->swizzle_len && + comp.value()->atom_len == this->atom_len && + comp.value()->swizzle_inner == this->swizzle_inner) { + // A bare swizzle (ComposeLayout with a trivial tile) has no real tile to + // compare; its "tile" is the identity over the inner product, and a bare + // `tile_layout` argument contributes the identity over the tiled product. + // Substitute those identities so TileLayoutNode::IsTileInner sees the same + // inputs the former SwizzleLayoutNode::IsTileInner produced. + TileLayout this_tile = this->tile_layout; + if (this->tile_layout->IsTrivial()) { + this_tile = IdentityTileLayout(inner_shape); + } + TileLayout arg_tile = comp.value()->tile_layout; + if (comp.value()->tile_layout->IsTrivial()) { + arg_tile = IdentityTileLayout(tiled_shape); + } + return this_tile->IsTileInner(arg_tile, tiled_shape, inner_shape); } } return std::nullopt; @@ -107,11 +142,17 @@ ffi::Optional ComposeLayoutNode::IsTileOuter( ffi::Optional ComposeLayoutNode::Slice(const ffi::Array& shape, const Region& region) const { - // Slice applies to the tile layout then compose with swizzle. - auto sliced_opt = tile_layout->Slice(shape, region); + // A bare swizzle (ComposeLayout with a trivial tile) carries only the swizzle + // period, not a tile matching `shape`; substitute an identity tile over the + // buffer shape first, matching the former SwizzleLayoutNode::Slice. + TileLayout base = this->tile_layout; + if (base->IsTrivial()) { + base = IdentityTileLayout(shape); + } + auto sliced_opt = base->Slice(shape, region); if (!sliced_opt.has_value()) return std::nullopt; auto sliced = sliced_opt.value().as().value(); - return ComposeLayout(swizzle, sliced); + return ComposeLayout(per_element, swizzle_len, atom_len, sliced, swizzle_inner); } } // namespace tirx diff --git a/src/tirx/ir/layout/swizzle_layout.cc b/src/tirx/ir/layout/swizzle_layout.cc deleted file mode 100644 index aa80223085b0..000000000000 --- a/src/tirx/ir/layout/swizzle_layout.cc +++ /dev/null @@ -1,128 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -#include "utils.h" - -namespace tvm { -namespace tirx { - -/**************** SwizzleLayout ****************/ -SwizzleLayout::SwizzleLayout(int per_element, int swizzle_len, int atom_len, bool swizzle_inner) { - auto n = ffi::make_object(); - n->per_element = per_element; - n->swizzle_len = swizzle_len; - n->atom_len = atom_len; - n->swizzle_inner = swizzle_inner; - TVM_FFI_ICHECK(n->VerifyWellFormed()) << "ValueError: The swizzle layout is not well-formed"; - int swizzle_mask = (1 << swizzle_len) - 1; - n->inner_mask = swizzle_mask; - n->outer_mask = swizzle_mask << atom_len; - data_ = std::move(n); -} - -TVM_FFI_STATIC_INIT_BLOCK() { - namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.SwizzleLayout", - [](int per_element, int swizzle_len, int atom_len, bool swizzle_inner) { - return SwizzleLayout(per_element, swizzle_len, atom_len, swizzle_inner); - }); -} - -bool SwizzleLayoutNode::CompatibleWithShape(const Array& shape) const { return true; } - -bool SwizzleLayoutNode::VerifyWellFormed() const { - return per_element >= 0 && swizzle_len >= 0 && atom_len >= swizzle_len; -} - -PrimExpr SwizzleLayoutNode::GetSize(ffi::Optional axis_name) const { - TVM_FFI_ICHECK(!axis_name.has_value()) - << "ValueError: axis_name is not supported for swizzle layout"; - return 1 << (per_element + swizzle_len + atom_len); -} - -PrimExpr SwizzleLayoutNode::GetSpan(ffi::Optional axis_name) const { - TVM_FFI_ICHECK(!axis_name.has_value()) - << "ValueError: axis_name is not supported for swizzle layout"; - return GetSize(); -} - -ffi::Map SwizzleLayoutNode::Apply(ffi::Array coord) const { - LOG(FATAL) << "SwizzleLayoutNode::Apply(Array) is not implemented"; - return {}; -} - -ffi::Map SwizzleLayoutNode::Apply(PrimExpr coord) const { - PrimExpr input = coord; - auto f = [&](const PrimExpr& x) -> PrimExpr { - if (swizzle_inner) { - return x ^ ((x & outer_mask) >> atom_len); - } else { - return x ^ ((x & inner_mask) << atom_len); - } - }; - auto base = 1 << per_element; - arith::Analyzer analyzer; - // It takes more arithmetic operations to compute the result, but it is more friendly to the - // vectorization. We use "m" as the default axis name here. - return { - {"m", analyzer->Simplify((f(floordiv(input, base)) << per_element) + floormod(input, base))}}; -} - -Layout SwizzleLayoutNode::Canonicalize() const { return ffi::GetRef(this); } - -Layout SwizzleLayoutNode::Tile(const TileLayout& outer, const ffi::Array& outer_shape, - const ffi::Array& inner_shape) const { - // Compose(Swizzle, Identity) -> then tile with `outer`. - auto comp = ComposeLayout(ffi::GetRef(this), IdentityTileLayout(inner_shape)); - return comp->Tile(outer, outer_shape, inner_shape); -} - -ffi::Optional SwizzleLayoutNode::IsTileInner( - const Layout& tile_layout, const ffi::Array& tiled_shape, - const ffi::Array& inner_shape) const { - // We expect tile_layout to be Compose(SwizzleLayout(this), _). - if (auto comp = tile_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, ffi::GetRef(this))) { - auto identity = IdentityTileLayout(inner_shape); - return identity->IsTileInner(comp.value()->tile_layout, tiled_shape, inner_shape); - } - } else if (auto swizzle = tile_layout.as()) { - if (StructuralEqual()(swizzle.value(), ffi::GetRef(this))) { - auto inner_identity = IdentityTileLayout(inner_shape); - auto tile_identity = IdentityTileLayout(tiled_shape); - return inner_identity->IsTileInner(tile_identity, tiled_shape, inner_shape); - } - } - return std::nullopt; -} - -ffi::Optional SwizzleLayoutNode::IsTileOuter( - const Layout& tile_layout, const ffi::Array& tiled_shape, - const ffi::Array& outer_shape) const { - return std::nullopt; -} - -ffi::Optional SwizzleLayoutNode::Slice(const ffi::Array& shape, - const Region& region) const { - // Compose(Swizzle, Identity) -> then slice. - auto comp = ComposeLayout(ffi::GetRef(this), IdentityTileLayout(shape)); - return comp->Slice(shape, region); -} - -} // namespace tirx -} // namespace tvm diff --git a/src/tirx/ir/layout/tile_core.cc b/src/tirx/ir/layout/tile_core.cc index 518e392bf46a..63a8c7f91cd4 100644 --- a/src/tirx/ir/layout/tile_core.cc +++ b/src/tirx/ir/layout/tile_core.cc @@ -30,7 +30,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { AxisNode::RegisterReflection(); IterNode::RegisterReflection(); TileLayoutNode::RegisterReflection(); - SwizzleLayoutNode::RegisterReflection(); ComposeLayoutNode::RegisterReflection(); } diff --git a/src/tirx/ir/layout/tile_direct_sum_ops.cc b/src/tirx/ir/layout/tile_direct_sum_ops.cc index 33622453dc20..1a7c945bfc81 100644 --- a/src/tirx/ir/layout/tile_direct_sum_ops.cc +++ b/src/tirx/ir/layout/tile_direct_sum_ops.cc @@ -201,18 +201,38 @@ ffi::Optional TileLayoutNode::IsDirectSumLeft( Layout ComposeLayoutNode::DirectSum(const TileLayout& left, const Array& left_shape, const Array& right_shape) const { - // Direct-sum applies to the tile layout then compose with swizzle. - auto right_sum = tile_layout->DirectSum(left, left_shape, right_shape).as().value(); - return ComposeLayout(swizzle, right_sum); + // A bare swizzle (ComposeLayout with a trivial tile) carries only the swizzle + // period, not a tile matching `right_shape`; substitute an identity tile over + // the right shape first, matching the former SwizzleLayoutNode::DirectSum. + TileLayout base = this->tile_layout; + if (base->IsTrivial()) { + base = IdentityTileLayout(right_shape); + } + auto right_sum = base->DirectSum(left, left_shape, right_shape).as().value(); + return ComposeLayout(per_element, swizzle_len, atom_len, right_sum, swizzle_inner); } ffi::Optional ComposeLayoutNode::IsDirectSumRight( const Layout& sum_layout, const ffi::Array& interleaved_shape, const ffi::Array& right_shape) const { if (auto comp = sum_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, this->swizzle)) { - return this->tile_layout->IsDirectSumRight(comp.value()->tile_layout, interleaved_shape, - right_shape); + if (comp.value()->per_element == this->per_element && + comp.value()->swizzle_len == this->swizzle_len && + comp.value()->atom_len == this->atom_len && + comp.value()->swizzle_inner == this->swizzle_inner) { + // A bare swizzle (ComposeLayout with a trivial tile) carries only the swizzle + // period; substitute identity tiles so TileLayoutNode::IsDirectSumRight sees + // the same inputs the forward DirectSum produced (and the former SwizzleLayout + // path intended), instead of grouping the period tile against the full shape. + TileLayout this_tile = this->tile_layout; + if (this->tile_layout->IsTrivial()) { + this_tile = IdentityTileLayout(right_shape); + } + TileLayout sum_tile = comp.value()->tile_layout; + if (comp.value()->tile_layout->IsTrivial()) { + sum_tile = IdentityTileLayout(interleaved_shape); + } + return this_tile->IsDirectSumRight(sum_tile, interleaved_shape, right_shape); } } return std::nullopt; @@ -222,39 +242,23 @@ ffi::Optional ComposeLayoutNode::IsDirectSumLeft( const Layout& sum_layout, const ffi::Array& interleaved_shape, const ffi::Array& left_shape) const { if (auto comp = sum_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, this->swizzle)) { - return this->tile_layout->IsDirectSumLeft(comp.value()->tile_layout, interleaved_shape, - left_shape); - } - } - return std::nullopt; -} - -Layout SwizzleLayoutNode::DirectSum(const TileLayout& left, const Array& left_shape, - const Array& right_shape) const { - // Compose(Swizzle, Identity(right_shape)) then direct-sum with left. - auto comp = ComposeLayout(ffi::GetRef(this), IdentityTileLayout(right_shape)); - return comp->DirectSum(left, left_shape, right_shape); -} - -ffi::Optional SwizzleLayoutNode::IsDirectSumRight( - const Layout& sum_layout, const ffi::Array& interleaved_shape, - const ffi::Array& right_shape) const { - if (auto comp = sum_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, ffi::GetRef(this))) { - return comp.value()->tile_layout->IsDirectSumRight(sum_layout, interleaved_shape, - right_shape); - } - } - return std::nullopt; -} - -ffi::Optional SwizzleLayoutNode::IsDirectSumLeft( - const Layout& sum_layout, const ffi::Array& interleaved_shape, - const ffi::Array& left_shape) const { - if (auto comp = sum_layout.as()) { - if (StructuralEqual()(comp.value()->swizzle, ffi::GetRef(this))) { - return comp.value()->tile_layout->IsDirectSumLeft(sum_layout, interleaved_shape, left_shape); + if (comp.value()->per_element == this->per_element && + comp.value()->swizzle_len == this->swizzle_len && + comp.value()->atom_len == this->atom_len && + comp.value()->swizzle_inner == this->swizzle_inner) { + // A bare swizzle (ComposeLayout with a trivial tile) carries only the swizzle + // period; substitute identity tiles so TileLayoutNode::IsDirectSumLeft sees + // the same inputs the forward DirectSum produced (and the former SwizzleLayout + // path intended), instead of grouping the period tile against the full shape. + TileLayout this_tile = this->tile_layout; + if (this->tile_layout->IsTrivial()) { + this_tile = IdentityTileLayout(left_shape); + } + TileLayout sum_tile = comp.value()->tile_layout; + if (comp.value()->tile_layout->IsTrivial()) { + sum_tile = IdentityTileLayout(interleaved_shape); + } + return this_tile->IsDirectSumLeft(sum_tile, interleaved_shape, left_shape); } } return std::nullopt; diff --git a/src/tirx/ir/layout/tile_internal.h b/src/tirx/ir/layout/tile_internal.h index 0c0a55314dd6..f4caae0aee25 100644 --- a/src/tirx/ir/layout/tile_internal.h +++ b/src/tirx/ir/layout/tile_internal.h @@ -34,6 +34,10 @@ namespace tirx { std::pair> Group(TileLayout layout, const ffi::Array& shape); +// Group a tile layout by the minimal common refinement of several logical shapes. +std::pair>> GroupMany( + TileLayout layout, const ffi::Array>& shapes); + // Same as Group but returns std::nullopt instead of fatal-checking when the // layout cannot be regrouped by ``shape``. std::optional>> TryGroup( diff --git a/src/tirx/ir/layout/tile_slice.cc b/src/tirx/ir/layout/tile_slice.cc index 7e9a9571774d..b2c793a93d77 100644 --- a/src/tirx/ir/layout/tile_slice.cc +++ b/src/tirx/ir/layout/tile_slice.cc @@ -151,6 +151,12 @@ ffi::Optional TileLayoutNode::Slice(const Array& shape, auto [grouped_layout, seps] = Group(canon, shape); std::vector new_shard; ffi::Map new_offset; + // The buffer layout may already be a statement-local view with a physical + // base offset (for example a non-zero TMEM row/column). Group slicing adds + // the selected region's offset to that base; it must not replace it. + for (const auto& [axis, off] : grouped_layout->offset) { + new_offset.Set(axis, off); + } for (size_t i = 0; i < seps.size() - 1; ++i) { std::vector shard(grouped_layout->shard.begin() + seps[i], grouped_layout->shard.begin() + seps[i + 1]); diff --git a/src/tirx/ir/layout/tile_tile_ops.cc b/src/tirx/ir/layout/tile_tile_ops.cc index e6cabdd98aba..7a5c08ed1df7 100644 --- a/src/tirx/ir/layout/tile_tile_ops.cc +++ b/src/tirx/ir/layout/tile_tile_ops.cc @@ -20,6 +20,8 @@ /* * Tiling operations and helpers for TileLayout. */ +#include + #include "tile_internal.h" namespace tvm { @@ -71,6 +73,192 @@ std::pair> Group(TileLayout layout, return {ffi::GetRef(n), seps}; } +std::pair>> GroupMany( + TileLayout layout, const ffi::Array>& shapes) { + struct BoundaryClass { + PrimExpr value; + std::vector counts; + }; + + arith::Analyzer analyzer; + TVM_FFI_ICHECK(!shapes.empty()) << "group_many requires at least one shape"; + + std::vector> boundary_sequences; + boundary_sequences.reserve(shapes.size() + 1); + for (size_t shape_idx = 0; shape_idx < shapes.size(); ++shape_idx) { + const auto& shape = shapes[shape_idx]; + TVM_FFI_ICHECK(!shape.empty()) << "group_many shape " << shape_idx << " must not be empty"; + + std::vector boundaries{PrimExpr(1)}; + PrimExpr product = 1; + for (size_t dim_idx = 0; dim_idx < shape.size(); ++dim_idx) { + PrimExpr previous = product; + product = analyzer->Simplify(product * shape[dim_idx]); + TVM_FFI_ICHECK(analyzer->CanProve(previous <= product)) + << "group_many cannot prove cumulative boundary order for shape " << shape_idx + << " at dimension " << dim_idx << ": " << previous << " <= " << product; + TVM_FFI_ICHECK(analyzer->CanProveEqual(floormod(product, previous), 0)) + << "group_many cannot prove cumulative boundary divisibility for shape " << shape_idx + << " at dimension " << dim_idx << ": " << previous << " divides " << product; + boundaries.push_back(product); + } + boundary_sequences.push_back(std::move(boundaries)); + } + + std::vector layout_boundaries{PrimExpr(1)}; + PrimExpr layout_product = 1; + for (size_t iter_idx = 0; iter_idx < layout->shard.size(); ++iter_idx) { + PrimExpr previous = layout_product; + layout_product = analyzer->Simplify(layout_product * layout->shard[iter_idx]->extent); + TVM_FFI_ICHECK(analyzer->CanProve(previous <= layout_product)) + << "group_many cannot prove cumulative boundary order for layout iter " << iter_idx << ": " + << previous << " <= " << layout_product; + TVM_FFI_ICHECK(analyzer->CanProveEqual(floormod(layout_product, previous), 0)) + << "group_many cannot prove cumulative boundary divisibility for layout iter " << iter_idx + << ": " << previous << " divides " << layout_product; + layout_boundaries.push_back(layout_product); + } + + const PrimExpr& total = boundary_sequences[0].back(); + for (size_t shape_idx = 1; shape_idx < boundary_sequences.size(); ++shape_idx) { + TVM_FFI_ICHECK(analyzer->CanProveEqual(total, boundary_sequences[shape_idx].back())) + << "group_many cannot prove equal total products for shape 0 and shape " << shape_idx + << ": " << total << " vs " << boundary_sequences[shape_idx].back(); + } + TVM_FFI_ICHECK(analyzer->CanProveEqual(total, layout_boundaries.back())) + << "group_many cannot prove that the layout and shapes have equal total products: " + << layout_boundaries.back() << " vs " << total; + boundary_sequences.push_back(std::move(layout_boundaries)); + + std::vector classes; + for (size_t sequence_idx = 0; sequence_idx < boundary_sequences.size(); ++sequence_idx) { + for (const PrimExpr& boundary : boundary_sequences[sequence_idx]) { + size_t class_idx = classes.size(); + for (size_t idx = 0; idx < classes.size(); ++idx) { + if (analyzer->CanProveEqual(boundary, classes[idx].value)) { + class_idx = idx; + break; + } + } + + if (class_idx == classes.size()) { + size_t insert_at = classes.size(); + for (size_t idx = 0; idx < classes.size(); ++idx) { + if (analyzer->CanProve(boundary < classes[idx].value)) { + insert_at = idx; + break; + } + TVM_FFI_ICHECK(analyzer->CanProve(classes[idx].value < boundary)) + << "group_many cannot prove the order or equality of cumulative boundaries " + << boundary << " and " << classes[idx].value; + } + classes.insert(classes.begin() + insert_at, + BoundaryClass{boundary, std::vector(boundary_sequences.size(), 0)}); + class_idx = insert_at; + } + classes[class_idx].counts[sequence_idx]++; + } + } + + std::vector class_starts; + std::vector class_multiplicities; + std::vector common_boundaries; + class_starts.reserve(classes.size()); + class_multiplicities.reserve(classes.size()); + for (const BoundaryClass& boundary_class : classes) { + size_t multiplicity = + *std::max_element(boundary_class.counts.begin(), boundary_class.counts.end()); + class_starts.push_back(common_boundaries.size()); + class_multiplicities.push_back(multiplicity); + for (size_t idx = 0; idx < multiplicity; ++idx) { + common_boundaries.push_back(boundary_class.value); + } + } + + std::vector refined_extents; + refined_extents.reserve(common_boundaries.size() - 1); + for (size_t idx = 1; idx < common_boundaries.size(); ++idx) { + const PrimExpr& previous = common_boundaries[idx - 1]; + const PrimExpr& current = common_boundaries[idx]; + TVM_FFI_ICHECK(analyzer->CanProve(previous <= current)) + << "group_many cannot prove common boundary order: " << previous << " <= " << current; + TVM_FFI_ICHECK(analyzer->CanProveEqual(floormod(current, previous), 0)) + << "group_many has incompatible cumulative boundaries: " << previous << " does not divide " + << current; + refined_extents.push_back(analyzer->Simplify(floordiv(current, previous))); + } + + auto boundary_positions = [&](const std::vector& sequence, + size_t sequence_idx) -> std::vector { + std::vector seen(classes.size(), 0); + std::vector positions; + positions.reserve(sequence.size()); + size_t previous_position = 0; + for (size_t boundary_idx = 0; boundary_idx < sequence.size(); ++boundary_idx) { + size_t class_idx = classes.size(); + for (size_t idx = 0; idx < classes.size(); ++idx) { + if (analyzer->CanProveEqual(sequence[boundary_idx], classes[idx].value)) { + class_idx = idx; + break; + } + } + TVM_FFI_ICHECK(class_idx < classes.size()) + << "group_many lost cumulative boundary " << sequence[boundary_idx]; + + size_t position = class_starts[class_idx] + seen[class_idx]++; + if (boundary_idx == 0) { + position = class_starts[class_idx]; + } + if (boundary_idx + 1 == sequence.size()) { + position = class_starts[class_idx] + class_multiplicities[class_idx] - 1; + } + TVM_FFI_ICHECK(boundary_idx == 0 || position > previous_position) + << "group_many could not preserve repeated boundary " << sequence[boundary_idx] + << " for sequence " << sequence_idx; + previous_position = position; + positions.push_back(position); + } + TVM_FFI_ICHECK_EQ(positions.front(), 0); + TVM_FFI_ICHECK_EQ(positions.back(), common_boundaries.size() - 1); + return positions; + }; + + const auto& layout_sequence = boundary_sequences.back(); + std::vector layout_positions = + boundary_positions(layout_sequence, boundary_sequences.size() - 1); + std::vector refined_shard; + refined_shard.reserve(refined_extents.size()); + for (size_t iter_idx = 0; iter_idx < layout->shard.size(); ++iter_idx) { + size_t start = layout_positions[iter_idx]; + size_t end = layout_positions[iter_idx + 1]; + TVM_FFI_ICHECK_LT(start, end) << "group_many cannot preserve layout iter " << iter_idx; + TVM_FFI_ICHECK(analyzer->CanProveEqual( + common_boundaries[end], common_boundaries[start] * layout->shard[iter_idx]->extent)) + << "group_many refinement changed layout iter " << iter_idx << " extent " + << layout->shard[iter_idx]->extent; + for (size_t position = start; position < end; ++position) { + PrimExpr inner_product = + analyzer->Simplify(floordiv(common_boundaries[end], common_boundaries[position + 1])); + refined_shard.push_back( + Iter(refined_extents[position], + analyzer->Simplify(layout->shard[iter_idx]->stride * inner_product), + layout->shard[iter_idx]->axis)); + } + } + + auto* n = layout.CopyOnWrite(); + n->shard = refined_shard; + TileLayout grouped = ffi::GetRef(n); + + std::vector> separators; + separators.reserve(shapes.size()); + for (size_t shape_idx = 0; shape_idx < shapes.size(); ++shape_idx) { + std::vector positions = boundary_positions(boundary_sequences[shape_idx], shape_idx); + separators.emplace_back(positions.begin(), positions.end()); + } + return {grouped, separators}; +} + std::optional>> TryGroup( TileLayout layout, const ffi::Array& shape) { // Same algorithm as Group but returns std::nullopt instead of ICHECK-failing @@ -114,11 +302,23 @@ std::optional>> TryGroup( TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def( - "tirx.TileLayoutGroup", [](const TileLayout& layout, const Array& shape) { - auto [res, seps] = Group(layout, shape); - return Tuple>{res, Array(seps.begin(), seps.end())}; - }); + refl::GlobalDef() + .def("tirx.TileLayoutGroup", + [](const TileLayout& layout, const Array& shape) { + auto [res, seps] = Group(layout, shape); + return Tuple>{res, + Array(seps.begin(), seps.end())}; + }) + .def("tirx.TileLayoutGroupMany", + [](const TileLayout& layout, const Array>& shapes) { + auto [res, separators] = GroupMany(layout, shapes); + Array> ffi_separators; + ffi_separators.reserve(separators.size()); + for (const auto& seps : separators) { + ffi_separators.push_back(Array(seps.begin(), seps.end())); + } + return Tuple>>{res, ffi_separators}; + }); } Layout TileLayoutNode::Tile(const TileLayout& outer_in, const Array& outer_shape, @@ -373,7 +573,9 @@ ffi::Optional TileLayoutNode::IsTileOuter(const Layout& tile_layout, if (auto comp = tile_layout.as()) { auto inner_layout = IsTileOuter(comp.value()->tile_layout, tiled_shape, outer_shape); if (!inner_layout) return std::nullopt; - return ComposeLayout(comp.value()->swizzle, inner_layout.value().as().value()); + return ComposeLayout(comp.value()->per_element, comp.value()->swizzle_len, + comp.value()->atom_len, inner_layout.value().as().value(), + comp.value()->swizzle_inner); } return std::nullopt; } diff --git a/src/tirx/ir/layout/utils.cc b/src/tirx/ir/layout/utils.cc index abd114ddf97b..fe419050288e 100644 --- a/src/tirx/ir/layout/utils.cc +++ b/src/tirx/ir/layout/utils.cc @@ -41,16 +41,6 @@ PrimExpr FlattenCoord(const Array& coord, const Array& shape [&shape, i = 0](PrimExpr acc, const PrimExpr& c) mutable { return acc * shape[i++] + c; }); } -TileLayout IdentityTileLayout(const ffi::Array& shape) { - if (shape.empty()) { - // Degenerate identity: no shard dims. - return TileLayout({}, {}, {}); - } - PrimExpr extent = std::accumulate(shape.begin() + 1, shape.end(), shape[0], - [](PrimExpr a, PrimExpr b) { return a * b; }); - return TileLayout({Iter(extent, 1, Axis::Get("m"))}, {}, {}); -} - ffi::Map BuildSpanMap(const TileLayout& layout) { ffi::Map span_map; for (const auto& iter : layout->shard) { @@ -87,5 +77,14 @@ bool AxisMatchesFilter(const Axis& axis, const ffi::Optional& axis_ (axis_name.has_value() && axis->name == axis_name.value()); } +TileLayout IdentityTileLayout(const ffi::Array& shape) { + if (shape.empty()) { + return TileLayout({}, {}, {}); + } + PrimExpr extent = std::accumulate(shape.begin() + 1, shape.end(), shape[0], + [](PrimExpr a, PrimExpr b) { return a * b; }); + return TileLayout({Iter(extent, 1, Axis::Get("m"))}, {}, {}); +} + } // namespace tirx } // namespace tvm diff --git a/src/tirx/ir/layout/utils.h b/src/tirx/ir/layout/utils.h index b274339ed1a5..225b166fc1b7 100644 --- a/src/tirx/ir/layout/utils.h +++ b/src/tirx/ir/layout/utils.h @@ -55,14 +55,6 @@ Array SplitCoord(PrimExpr coord, const Array& shape); */ PrimExpr FlattenCoord(const Array& coord, const Array& shape); -/*! - * \brief Create a TileLayout that maps the given logical shape to itself on the memory axis. - * This is effectively an identity layout over axis "m" with unit stride. - * \param shape Logical shape to map. - * \return Identity TileLayout over the concatenated extent of `shape`. - */ -TileLayout IdentityTileLayout(const ffi::Array& shape); - /*! * \brief Build a map from axis name to span for the provided layout's shard axes. * If an axis appears multiple times, the first occurrence defines the span value. @@ -87,6 +79,14 @@ std::vector GetDefaultStrides(const ffi::Array& data, */ bool AxisMatchesFilter(const Axis& axis, const ffi::Optional& axis_name); +/*! + * \brief Build a 1D identity TileLayout whose single shard covers the product of + * `shape` (contiguous, stride 1, memory axis "m"). Used to represent the + * "tile" of a bare swizzle (a ComposeLayout carrying a trivial tile) when + * delegating tile-relationship queries to TileLayoutNode. + */ +TileLayout IdentityTileLayout(const ffi::Array& shape); + } // namespace tirx } // namespace tvm diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index 171f1b7eb305..be7364fb5329 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -179,14 +179,23 @@ void StmtVisitor::VisitStmt_(const ScopeIdDefStmtNode* op) { } void StmtVisitor::VisitStmt_(const tirx::TilePrimitiveCallNode* op) { - auto fvisit = [this](const ffi::Any& e) { + std::function fvisit; + fvisit = [this, &fvisit](const ffi::Any& e) { if (e == nullptr) return; if (auto buffer_region = e.as()) { - return; + this->VisitBufferUse(buffer_region.value()->buffer); + for (const auto& range : buffer_region.value()->region) { + this->VisitExpr(range->min); + this->VisitExpr(range->extent); + } + } else if (auto var = e.as(); var && var.value()->ty.as()) { + this->VisitBufferUse(BufferVar(var.value())); } else if (auto expr = e.as()) { this->VisitExpr(expr.value()); } else if (auto stmt = e.as()) { this->VisitStmt(stmt.value()); + } else if (auto array = e.as>()) { + for (const ffi::Any& item : array.value()) fvisit(item); } }; VisitArray(op->args, fvisit); @@ -677,14 +686,19 @@ Stmt StmtMutator::VisitStmt_(const ScopeIdDefStmtNode* op) { } Stmt StmtMutator::VisitStmt_(const tirx::TilePrimitiveCallNode* op) { - auto fmutate = [&](const ffi::Any& e) -> ffi::Any { + std::function fmutate; + fmutate = [&](const ffi::Any& e) -> ffi::Any { if (e == nullptr) return e; if (auto buffer_region = e.as()) { return Internal::Mutate(this, {buffer_region.value()})[0]; + } else if (auto var = e.as(); var && var.value()->ty.as()) { + return this->VisitBufferUse(BufferVar(var.value())); } else if (auto expr = e.as()) { return this->VisitPrimExpr(expr.value()); } else if (auto stmt = e.as()) { return this->VisitStmt(stmt.value()); + } else if (auto array = e.as>()) { + return Internal::MutateArray(this, array.value(), fmutate); } return e; }; diff --git a/src/tirx/ir/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc index d8be8942479d..b4f09619c56c 100644 --- a/src/tirx/ir/tirx_stmt.cc +++ b/src/tirx/ir/tirx_stmt.cc @@ -24,7 +24,7 @@ #include #include -#include +#include namespace tvm { namespace tirx { diff --git a/src/tirx/op/tirx.cc b/src/tirx/op/tirx.cc index 1f96ad3e0008..21249706f61f 100644 --- a/src/tirx/op/tirx.cc +++ b/src/tirx/op/tirx.cc @@ -24,8 +24,7 @@ #include #include -#include -#include +#include namespace tvm { namespace tirx { @@ -146,6 +145,7 @@ TIRX_DEFINE_TILE_OP(zero); TIRX_DEFINE_TILE_OP(sqrt); TIRX_DEFINE_TILE_OP(exp); TIRX_DEFINE_TILE_OP(exp2); +TIRX_DEFINE_TILE_OP(log2); TIRX_DEFINE_TILE_OP(add); TIRX_DEFINE_TILE_OP(sub); TIRX_DEFINE_TILE_OP(mul); diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index 60556ed647b8..97e24e0f6211 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -30,7 +30,7 @@ #include #include #include -#include +#include #include "./utils.h" diff --git a/src/tirx/script/printer/buffer.cc b/src/tirx/script/printer/buffer.cc index b90a0f569019..943b21100447 100644 --- a/src/tirx/script/printer/buffer.cc +++ b/src/tirx/script/printer/buffer.cc @@ -542,23 +542,19 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) // TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) // .set_dispatch( "", [](tirx::ComposeLayout layout, AccessPath p, IRDocsifier d) -> Doc { - auto layoutA = d->AsDoc(layout->swizzle, p->Attr("swizzle")); - auto layoutB = d->AsDoc(layout->tile_layout, p->Attr("tile_layout")); - return TIRx(d, "ComposeLayout")->Call({layoutA, layoutB}, {}, {}); - }); - -TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) // - .set_dispatch( - "", [](tirx::SwizzleLayout layout, AccessPath p, IRDocsifier d) -> Doc { - return TIRx(d, "SwizzleLayout") - ->Call( - { - LiteralDoc::Int(layout->per_element, p->Attr("per_element")), - LiteralDoc::Int(layout->swizzle_len, p->Attr("swizzle_len")), - LiteralDoc::Int(layout->atom_len, p->Attr("atom_len")), - }, - {"swizzle_inner"}, - {LiteralDoc::Boolean(layout->swizzle_inner, p->Attr("swizzle_inner"))}); + auto per_element = LiteralDoc::Int(layout->per_element, p->Attr("per_element")); + auto swizzle_len = LiteralDoc::Int(layout->swizzle_len, p->Attr("swizzle_len")); + auto atom_len = LiteralDoc::Int(layout->atom_len, p->Attr("atom_len")); + auto tile_doc = d->AsDoc(layout->tile_layout, p->Attr("tile_layout")); + ffi::Array kwargs_keys; + ffi::Array kwargs_values; + if (!layout->swizzle_inner) { + kwargs_keys.push_back("swizzle_inner"); + kwargs_values.push_back( + LiteralDoc::Boolean(layout->swizzle_inner, p->Attr("swizzle_inner"))); + } + return TIRx(d, "ComposeLayout") + ->Call({per_element, swizzle_len, atom_len, tile_doc}, kwargs_keys, kwargs_values); }); TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) @@ -586,7 +582,6 @@ TVM_SCRIPT_REPR(tirx::BufferTypeNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::IterNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::TileLayoutNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::ComposeLayoutNode, ReprPrintTIR); -TVM_SCRIPT_REPR(tirx::SwizzleLayoutNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::MatchBufferRegionNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::ProducerLoadNode, ReprPrintTIR); diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index 9d79955e282a..f0a46503e4db 100644 --- a/src/tirx/script/printer/expr.cc +++ b/src/tirx/script/printer/expr.cc @@ -257,9 +257,9 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) } }); -LambdaDoc PrintPredicate(const ffi::ObjectRef& pred, const ffi::Array& vs, - const AccessPath& vs_p, const PrimExpr& p, const AccessPath& p_p, - const IRDocsifier& d) { +LambdaDoc PrintLambda(const ffi::ObjectRef& pred, const ffi::Array& vs, + const AccessPath& vs_p, const PrimExpr& p, const AccessPath& p_p, + const IRDocsifier& d) { With f(d, pred); ffi::Array vars; for (int i = 0, l = vs.size(); i < l; ++i) { @@ -270,11 +270,11 @@ LambdaDoc PrintPredicate(const ffi::ObjectRef& pred, const ffi::Array("", - [](tirx::Predicate pred, AccessPath p, IRDocsifier d) -> Doc { - return PrintPredicate(pred, pred->vars, p->Attr("vars"), - pred->pred, p->Attr("pred"), d); - }); + .set_dispatch("", + [](tirx::LambdaExpr pred, AccessPath p, IRDocsifier d) -> Doc { + return PrintLambda(pred, pred->vars, p->Attr("vars"), + pred->pred, p->Attr("pred"), d); + }); TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) .set_dispatch("", [](tirx::Let let, AccessPath p, IRDocsifier d) -> Doc { @@ -520,7 +520,7 @@ TVM_SCRIPT_REPR(tirx::ShuffleNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::CommReducerNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::IndexMapNode, ReprPrintTIR); TVM_SCRIPT_REPR(tirx::ReduceNode, ReprPrintTIR); -TVM_SCRIPT_REPR(tirx::PredicateNode, ReprPrintTIR); +TVM_SCRIPT_REPR(tirx::LambdaExprNode, ReprPrintTIR); } // namespace printer } // namespace script diff --git a/src/tirx/script/printer/stmt.cc b/src/tirx/script/printer/stmt.cc index 9aa9822c10ae..93ef1724e0b2 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -296,7 +296,8 @@ bool IsDefaultLayout(const ffi::Optional& layout, const ffi::Array */ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child, const AccessPath& p, const IRDocsifier& d, - const tirx::BufferVar& parent) { + const tirx::BufferVar& parent, + bool require_same_layout) { ffi::Optional parent_doc = d->GetVarDoc(parent); if (!parent_doc.has_value()) return std::nullopt; ExprDoc pdoc = parent_doc.value(); @@ -321,67 +322,11 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child bool child_is_default = IsDefaultLayout(child->layout, child->shape); bool parent_is_default = IsDefaultLayout(parent->layout, parent->shape); - // --- (a) Slice (default layout, different elem_offset) --- - if (!same_elem_offset && same_dtype && !parent->shape.empty()) { - // Reconstruct start indices from elem_offset difference and parent strides (row-major) - // offset_diff = child->elem_offset - parent->elem_offset - // For row-major: strides[i] = prod(shape[i+1:]) - // start[i] = offset_diff / strides[i]; offset_diff %= strides[i] - // Build slice doc: parent[start:start+extent, ...] - // We only support this for IntImm offsets - auto* child_off = child->elem_offset.as(); - auto* parent_off = parent->elem_offset.as(); - if (child_off && parent_off) { - int64_t offset_diff = child_off->value - parent_off->value; - // Compute row-major strides - std::vector strides(parent->shape.size()); - int64_t stride = 1; - for (int i = static_cast(parent->shape.size()) - 1; i >= 0; --i) { - strides[i] = stride; - if (auto* s = parent->shape[i].as()) { - stride *= s->value; - } else { - return std::nullopt; // Non-constant shape, can't decompose - } - } - // Check child shape is also all IntImm - for (size_t i = 0; i < child->shape.size(); ++i) { - if (!child->shape[i].as()) return std::nullopt; - } - if (child->shape.size() != parent->shape.size()) return std::nullopt; - - ffi::Array slices; - int64_t remaining = offset_diff; - bool in_bounds = true; - for (size_t i = 0; i < parent->shape.size(); ++i) { - int64_t start_val = remaining / strides[i]; - remaining %= strides[i]; - int64_t extent_val = child->shape[i].as()->value; - int64_t parent_dim = parent->shape[i].as()->value; - int64_t stop_val = start_val + extent_val; - // Bounds check: start + extent must be within parent dim - if (stop_val > parent_dim) { - in_bounds = false; - break; - } - if (start_val == 0 && stop_val == parent_dim) { - // Full range: use 0:N slice - ExprDoc start_doc = LiteralDoc::Int(0, p->Attr("elem_offset")); - ExprDoc stop_doc = - d->AsDoc(parent->shape[i], p->Attr("buffer")->Attr("shape")->ArrayItem(i)); - slices.push_back(SliceDoc(start_doc, stop_doc, std::nullopt)); - } else { - ExprDoc start_doc = LiteralDoc::Int(start_val, p->Attr("elem_offset")); - ExprDoc stop_doc = LiteralDoc::Int(stop_val, p->Attr("elem_offset")); - slices.push_back(SliceDoc(start_doc, stop_doc, std::nullopt)); - } - } - if (remaining == 0 && in_bounds) { - return pdoc[slices]; - } - } - return std::nullopt; - } + // NOTE: an earlier sugar printed rank-preserving aliases with a different + // elem_offset as ``parent[slices]``. That print is not roundtrippable: it + // reparses as a BufferRegion, not a Buffer, so any later Buffer use of the + // alias (stores, views) breaks. Such aliases now print as plain + // T.decl_buffer, which reparses exactly. // --- (b) Local: parent has thread axes, child has storage layout (non-thread part) --- if (same_elem_offset && same_dtype && !parent_is_default && parent->layout.has_value()) { @@ -655,6 +600,9 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::BufferVar& child } else if (!child->layout.has_value() && !parent->layout.has_value()) { same_layout = true; } + // First pass prefers a parent whose layout matches structurally, so the + // sugar prints as a bare reshape instead of restating the layout. + if (require_same_layout && !same_layout) return std::nullopt; if (!same_layout && child->layout.has_value() && !child_is_default) { kwargs_keys.push_back("layout"); kwargs_values.push_back( @@ -673,7 +621,14 @@ ffi::Optional TryDeclBufferSugar(const tirx::BufferVar& child, const Ac const ffi::Optional& data, const IRDocsifier& d) { auto parents = FindParentBuffers(child, data, d); for (const auto& parent : parents) { - if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent)) { + if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent, + /*require_same_layout=*/true)) { + return sugar; + } + } + for (const auto& parent : parents) { + if (auto sugar = TryDeclBufferSugarWithParent(child, p, d, parent, + /*require_same_layout=*/false)) { return sugar; } } diff --git a/src/tirx/script/printer/utils.h b/src/tirx/script/printer/utils.h index f1f607ea8dff..296c24cd7ae9 100644 --- a/src/tirx/script/printer/utils.h +++ b/src/tirx/script/printer/utils.h @@ -29,10 +29,9 @@ #include #include #include -#include #include #include -#include +#include #include #include @@ -328,6 +327,14 @@ ExprDoc BufferAttn(const tirx::BufferVar& buffer, const AccessPath& p, const Fra */ ExprDoc PrintVarCreation(const tirx::Var& var, const AccessPath& var_p, const IRDocsifier& d); +/*! \brief Print a reified lambda ``(vars, body)`` as a ``LambdaDoc``. + +Used by the ``tirx.tile.select`` printer specialization. Defined in expr.cc. +*/ +LambdaDoc PrintLambda(const ffi::ObjectRef& pred, const ffi::Array& vs, + const AccessPath& vs_p, const PrimExpr& p, const AccessPath& p_p, + const IRDocsifier& d); + /*! \brief A Var occurrence counter visitor */ class OccurrenceCounter : public tirx::StmtExprVisitor { public: diff --git a/src/tirx/transform/common_subexpr_elim.cc b/src/tirx/transform/common_subexpr_elim.cc index 3cfe8635ab30..9ea63ea96a25 100644 --- a/src/tirx/transform/common_subexpr_elim.cc +++ b/src/tirx/transform/common_subexpr_elim.cc @@ -80,6 +80,7 @@ #include #include #include +#include #include #include @@ -762,12 +763,41 @@ class CSERewriter : public StmtExprMutator { * before recursing. If insertions are planned, wraps the visited result * in a SeqStmt with the Bind statements prepended. SeqStmt flattening * ensures correct structure regardless of context. + * + * The same statement object can occur at several tree positions (loop + * unrolling shares loop-invariant subtrees), and each occurrence hits the + * pointer-keyed plan entry. Re-emitting the planned Binds verbatim would + * define each cse var once per occurrence and the result would no longer + * be SSA, so only the first occurrence materializes the plan as-is; later + * occurrences bind fresh vars and substitute them through both the Bind + * values and their copy of the subtree. */ Stmt VisitStmt(const Stmt& stmt) override { auto it = insert_before_.find(stmt); Stmt visited = StmtExprMutator::VisitStmt(stmt); if (it != insert_before_.end()) { - ffi::Array new_stmts(it->second.begin(), it->second.end()); + ffi::Array new_stmts; + if (materialized_.insert(stmt.get()).second) { + new_stmts = ffi::Array(it->second.begin(), it->second.end()); + } else { + std::unordered_map remap; + auto lookup = [&remap](const Var& v) -> ffi::Optional { + auto rit = remap.find(v.get()); + if (rit != remap.end()) return Expr(rit->second); + return std::nullopt; + }; + for (const Stmt& s : it->second) { + const BindNode* bind = s.as(); + TVM_FFI_ICHECK(bind != nullptr); + // Deeper Bind values may reference shallower cse vars of this same + // insertion point; route them through the fresh vars as well. + Expr value = Substitute(bind->value, lookup); + Var fresh(bind->var->name, bind->var->ty.as_or_throw()); + remap[bind->var.get()] = fresh.as_or_throw(); + new_stmts.push_back(Bind(fresh, value)); + } + visited = Substitute(visited, lookup); + } new_stmts.push_back(visited); return SeqStmt(new_stmts); } @@ -775,10 +805,12 @@ class CSERewriter : public StmtExprMutator { } private: - /*! \brief Plan: stmts to insert before each target. */ + /*! \brief Plan: stmts to insert each target (keyed by object identity). */ InsertBeforeTable insert_before_; /*! \brief Plan: expressions to replace with CSE vars. */ ExprRemapTable expr_remap_; + /*! \brief Insertion targets whose plan has already been materialized once. */ + std::unordered_set materialized_; }; // ============================================================================ diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index d89687b954e5..b2e1ecd16ae2 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include #include diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index 1a895cbdd516..0251afd360fc 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -89,10 +89,16 @@ class LaunchBoundsAttrExtractor : public StmtMutator { public: Stmt Extract(Stmt stmt) { min_blocks_per_sm_.reset(); - return operator()(std::move(stmt)); + max_blocks_per_cluster_.reset(); + Stmt result = operator()(std::move(stmt)); + TVM_FFI_ICHECK(!max_blocks_per_cluster_.has_value() || min_blocks_per_sm_.has_value()) + << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " requires " + << tirx::attr::kLaunchBoundsMinBlocksPerSM; + return result; } std::optional min_blocks_per_sm() const { return min_blocks_per_sm_; } + std::optional max_blocks_per_cluster() const { return max_blocks_per_cluster_; } private: Stmt VisitStmt_(const AttrStmtNode* op) final { @@ -108,11 +114,24 @@ class LaunchBoundsAttrExtractor : public StmtMutator { } min_blocks_per_sm_ = min_blocks_per_sm->value; return VisitStmt(op->body); + } else if (op->attr_key == tirx::attr::kLaunchBoundsMaxBlocksPerCluster) { + const auto* max_blocks_per_cluster = op->value.as(); + TVM_FFI_ICHECK(max_blocks_per_cluster) + << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " expects an integer value"; + TVM_FFI_ICHECK_GT(max_blocks_per_cluster->value, 0) + << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " must be positive"; + if (max_blocks_per_cluster_.has_value()) { + TVM_FFI_ICHECK_EQ(max_blocks_per_cluster_.value(), max_blocks_per_cluster->value) + << "Conflicting " << tirx::attr::kLaunchBoundsMaxBlocksPerCluster << " values"; + } + max_blocks_per_cluster_ = max_blocks_per_cluster->value; + return VisitStmt(op->body); } return StmtMutator::VisitStmt_(op); } std::optional min_blocks_per_sm_; + std::optional max_blocks_per_cluster_; }; class HostDeviceSplitter : public StmtMutator { @@ -222,9 +241,20 @@ class HostDeviceSplitter : public StmtMutator { if (is_stir) { device_func = WithAttr(std::move(device_func), tvm::attr::kSTir, true); } - if (device_target->kind->name == "cuda" && launch_bounds_attr.min_blocks_per_sm().has_value()) { - device_func = WithAttr(std::move(device_func), tirx::attr::kLaunchBoundsMinBlocksPerSM, - launch_bounds_attr.min_blocks_per_sm().value()); + if (auto launch_params = + cur_func_->GetAttr>(tirx::attr::kKernelLaunchParams)) { + device_func = + WithAttr(std::move(device_func), tirx::attr::kKernelLaunchParams, launch_params.value()); + } + if (device_target->kind->name == "cuda") { + if (launch_bounds_attr.min_blocks_per_sm().has_value()) { + device_func = WithAttr(std::move(device_func), tirx::attr::kLaunchBoundsMinBlocksPerSM, + launch_bounds_attr.min_blocks_per_sm().value()); + } + if (launch_bounds_attr.max_blocks_per_cluster().has_value()) { + device_func = WithAttr(std::move(device_func), tirx::attr::kLaunchBoundsMaxBlocksPerCluster, + launch_bounds_attr.max_blocks_per_cluster().value()); + } } auto num_inputs = cur_func_->GetAttr(tvm::attr::kNumInputs); if (num_inputs.has_value()) { @@ -303,8 +333,25 @@ class DeviceInfoCollector : public StmtVisitor { collector.info_.target = func->GetAttr(tvm::attr::kTarget).value().WithoutHost(); collector.info_.params = func->params; + if (auto requested = func->GetAttr>(tirx::attr::kKernelLaunchParams)) { + for (const ffi::String& tag : requested.value()) { + if (tag == tvm::runtime::launch_param::kUseProgramaticDependentLaunch) { + collector.use_programmatic_dependent_launch_ = true; + } else if (tag == tvm::runtime::launch_param::kUseCooperativeLaunch) { + collector.use_cooperative_launch_ = true; + } + } + } + collector(func->body); + if (collector.use_programmatic_dependent_launch_) { + collector.info_.launch_params.push_back( + tvm::runtime::launch_param::kUseProgramaticDependentLaunch); + } + if (collector.use_cooperative_launch_) { + collector.info_.launch_params.push_back(tvm::runtime::launch_param::kUseCooperativeLaunch); + } // The dynamic shared memory is required to be the last of the // kernel launch parameters. if (collector.dyn_shmem_size) { @@ -315,8 +362,13 @@ class DeviceInfoCollector : public StmtVisitor { collector.info_.global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol).value_or(gvar->name_hint); - collector.info_.launch_args = collector.info_.launch_params.Map( - [&](const auto& param) { return collector.GetArgument(param); }); + for (const ffi::String& param : collector.info_.launch_params) { + if (param == tvm::runtime::launch_param::kUseProgramaticDependentLaunch || + param == tvm::runtime::launch_param::kUseCooperativeLaunch) { + continue; + } + collector.info_.launch_args.push_back(collector.GetArgument(param)); + } return collector.info_; } @@ -412,24 +464,31 @@ class DeviceInfoCollector : public StmtVisitor { ffi::Map thread_extent; // The amount of dynamic shared memory used. ffi::Optional dyn_shmem_size{std::nullopt}; + // Flag-only launch attributes requested by the original PrimFunc. + bool use_programmatic_dependent_launch_{false}; + bool use_cooperative_launch_{false}; // Accumulated Bind definitions for inlining into extent/size expressions. ffi::Map bind_map_; }; class ReturnRemover : public StmtExprMutator { public: - static Stmt Apply(const Stmt& stmt) { - ReturnRemover mutator; + static Stmt Apply(const Stmt& stmt, bool remove) { + ReturnRemover mutator(remove); return mutator(stmt); } private: + explicit ReturnRemover(bool remove) : remove_(remove) {} + Stmt VisitStmt_(const ReturnNode* op) override { auto as_int = op->value.as(); TVM_FFI_ICHECK(as_int && as_int->value == 0) << "Device kernel may only contain a successful return, return 0"; - return Evaluate(0); + return remove_ ? Evaluate(0) : ffi::GetRef(op); } + + bool remove_; }; class GlobalVarCallCollector : public StmtExprVisitor { @@ -516,7 +575,9 @@ class DeviceKernelMutator : public StmtExprMutator { { auto write_ptr = func.CopyOnWrite(); write_ptr->ret_type = VoidType(); - write_ptr->body = ReturnRemover::Apply(write_ptr->body); + Target target = func->GetAttr(tvm::attr::kTarget).value(); + bool preserve_early_returns = target->kind->name == "cuda"; + write_ptr->body = ReturnRemover::Apply(write_ptr->body, !preserve_early_returns); } func = WithAttrs(std::move(func), diff --git a/src/tirx/transform/tile_primitive_dispatch.cc b/src/tirx/transform/tile_primitive_dispatch.cc index 1f192999bd43..dfd00c39e638 100644 --- a/src/tirx/transform/tile_primitive_dispatch.cc +++ b/src/tirx/transform/tile_primitive_dispatch.cc @@ -35,7 +35,7 @@ #include #include #include -#include +#include #include #include @@ -469,6 +469,20 @@ class TilePrimitiveDispatcher : public StmtExprMutator { return SeqStmt::Flatten(rebuilt); } + Stmt VisitStmt_(const BindNode* op) final { + Stmt stmt = StmtExprMutator::VisitStmt_(op); + const auto* bind = stmt.as(); + TVM_FFI_ICHECK(bind); + if (auto value = bind->value.as()) { + // Bind is flat: the definition is visible to subsequent statements in + // its enclosing scope. Under SSA, an inner-scope Var cannot be + // referenced after leaving that scope or rebound elsewhere, so stale + // entries are never consulted and no scope-based cleanup is needed. + var_range_map_.Set(bind->var, Range::FromMinExtent(value.value(), 1)); + } + return stmt; + } + Stmt VisitStmt_(const ForNode* op) final { // Collect the loop variables auto loop_var = op->loop_var.as_or_throw(); diff --git a/tests/python/arith/test_arith_rewrite_simplify.py b/tests/python/arith/test_arith_rewrite_simplify.py index 4620d0a7fdea..8c64cb9781ad 100644 --- a/tests/python/arith/test_arith_rewrite_simplify.py +++ b/tests/python/arith/test_arith_rewrite_simplify.py @@ -758,6 +758,56 @@ def test_uint_floormod_const_fold(): assert analyzer.can_prove_equal(expr, 0) +class TestFloorModUnsigned(BaseCompare): + """Modular-analysis simplifications for unsigned operands (uint32/uint64). + + Only sound when the divisor is a power of two (c2 | 2^bits): reducing + mod 2^bits does not change residues mod c2 in that case, so the + modular-decomposition rules carry over from the signed block. + """ + + x = tirx.Var("x", "uint32") + y = tirx.Var("y", "uint32") + z = tirx.Var("z", "int32") + test_case = tvm.testing.parameter( + TestCase( + flm(x * tirx.const(8, "uint32"), tirx.const(4, "uint32")), tirx.const(0, "uint32") + ), + TestCase(flm((z * 4).astype("uint32"), tirx.const(4, "uint32")), tirx.const(0, "uint32")), + TestCase( + flm(x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), tirx.const(4, "uint32")), + tirx.const(0, "uint32"), + ), + TestCase( + flm( + x * tirx.const(32, "uint32") + y * tirx.const(8, "uint32"), tirx.const(8, "uint32") + ), + tirx.const(0, "uint32"), + ), + # Non-power-of-two divisor must NOT simplify (unsigned wraparound + # changes residues mod c2). + TestCase( + flm(x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), tirx.const(12, "uint32")), + flm(x * tirx.const(16, "uint32") + (z * 4).astype("uint32"), tirx.const(12, "uint32")), + ), + # The pre-existing x * c1 % c2 -> 0 rule is pow2-guarded: c2=12 must + # not fold to 0 (e.g. x=2^29 wraps: uint32(x*12)=2^31, and 2^31 % 12 = 8). + TestCase( + flm(x * tirx.const(12, "uint32"), tirx.const(12, "uint32")), + flm(x * tirx.const(12, "uint32"), tirx.const(12, "uint32")), + ), + TestCase( + flm(x * tirx.const(8, "uint32"), tirx.const(8, "uint32")), tirx.const(0, "uint32") + ), + # Unsigned floordiv to zero by const bound. + TestCase( + tvm.tirx.floordiv(x, tirx.const(512, "uint32")), + tirx.const(0, "uint32"), + [x < 512], + ), + ) + + class TestMinIndex(BaseCompare): x, y, z = tvm.tirx.Var("x", "int32"), tvm.tirx.Var("y", "int32"), tvm.tirx.Var("z", "int32") test_case = tvm.testing.parameter( @@ -1329,3 +1379,36 @@ class TestCLZ(BaseCompare): if __name__ == "__main__": tvm.testing.main() + + +def test_allow_uint_as_index(): + """Opt-in no-overflow domain: unsigned index terms are not analyzed by + default; within allow_uint_as_index() the signed rule set applies and the + swizzle-style floordiv-quotient bit checks over uint32 offsets discharge.""" + w = tirx.Var("w", "int32") + lane = tirx.Var("lane", "int32") + x = tirx.Var("x", "uint32") + analyzer = tvm.arith.Analyzer() + analyzer.bind(lane, tvm.ir.Range.from_min_extent(0, 32)) + analyzer.bind(w, tvm.ir.Range.from_min_extent(0, 4)) + s_off = ( + (w * 1024 + lane // 8 * 8).astype("uint32") + + x * T.uint32(4096) + + (lane % 8 * 64).astype("uint32") + ) + check = flm(fld(s_off, T.uint32(512)), T.uint32(2)) + assert not analyzer.can_prove_equal(check, 0) + # a non-pow2 divisor never folds, inside or outside the mode... unless + # the flag is on: under the no-overflow assertion any c2 != 0 works. + np2 = flm(x * T.uint32(12), T.uint32(12)) + assert not analyzer.can_prove_equal(np2, 0) + with tvm.arith.allow_uint_as_index(): + assert analyzer.can_prove_equal(check, 0) + assert analyzer.can_prove_equal(flm(fld(s_off, T.uint32(32)), T.uint32(2)), 0) + assert analyzer.can_prove_equal(np2, 0) + # nested scopes compose + with tvm.arith.allow_uint_as_index(): + assert analyzer.can_prove_equal(check, 0) + assert analyzer.can_prove_equal(check, 0) + # the mode is scoped: disabled again afterwards + assert not analyzer.can_prove_equal(check, 0) diff --git a/tests/python/tirx-base/test_tir_op_types.py b/tests/python/tirx-base/test_tir_op_types.py index d51ed9db8111..3e3633828110 100644 --- a/tests/python/tirx-base/test_tir_op_types.py +++ b/tests/python/tirx-base/test_tir_op_types.py @@ -116,16 +116,14 @@ def test_tir_op_type_annotation(): def test_tir_op_tvm_access_ptr(): buffer = tirx.decl_buffer((128), "float32") - expr = tirx.tvm_access_ptr("float32", buffer.data, 0, 1, 2) - assert expr.op.name == "tirx.tvm_access_ptr" - assert expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("float32")) - offset_expr = tirx.ptr_byte_offset(buffer.data, 16, "uint8") - assert offset_expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8")) - prim_type = tvm.ir.PrimType("uint8") - typed_access_expr = tirx.tvm_access_ptr(prim_type, buffer.data, 0, 1, 2) - assert typed_access_expr.ty == tvm.ir.PointerType(prim_type) - typed_offset_expr = tirx.ptr_byte_offset(buffer.data, 16, prim_type) - assert typed_offset_expr.ty == tvm.ir.PointerType(prim_type) + for ptype in ("float32", tvm.ir.PrimType("float32")): + expr = tirx.tvm_access_ptr(ptype, buffer.data, 0, 1, 2) + assert expr.op.name == "tirx.tvm_access_ptr" + assert expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("float32")) + + for dtype in ("uint8", tvm.ir.PrimType("uint8")): + offset_expr = tirx.ptr_byte_offset(buffer.data, 16, dtype) + assert offset_expr.ty == tvm.ir.PointerType(tvm.ir.PrimType("uint8")) def test_tir_op_tvm_throw_last_error(): diff --git a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py index 052ed5668e14..f83331c1225a 100644 --- a/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py +++ b/tests/python/tirx-transform/test_tir_transform_common_subexpr_elim.py @@ -754,6 +754,30 @@ def main(B: T.Buffer((50,), "int32"), a: T.bool, b: T.bool, x: T.int32): assert "cse_v" not in after["main"].script() +# T24: Shared subtree stays SSA. UnrollLoop reuses one Stmt object across +# positions; the identity-keyed rewriter must bind fresh cse vars per occurrence. +def test_shared_subtree_stays_ssa(): + @tvm.script.ir_module + class Payload: + @T.prim_func(s_tir=True) + def main(B: T.Buffer((50,), "int32"), i1: T.int32, i2: T.int32): + B[i1] = (i1 + i2) * 2 + B[i2] = (i1 + i2) * 3 + + f = Payload["main"] + shared = f.body # one Stmt object, placed at two positions below + func = f.with_body(tvm.tirx.SeqStmt([shared, shared])) + after = tvm.tirx.transform.CommonSubexprElim()(tvm.IRModule({"main": func}))["main"] + + binds = [s for s in after.body if isinstance(s, tvm.tirx.Bind)] + assert len(binds) == 6, after.script() # (i1+i2, *2, *3) per occurrence + bound_vars = [b.var for b in binds] + for i, va in enumerate(bound_vars): + for vb in bound_vars[i + 1 :]: + assert not va.same_as(vb), f"duplicate var definitions:\n{after.script()}" + assert tvm.tirx.analysis.verify_ssa(after), after.script() + + if __name__ == "__main__": test_basic() test_if_single_branch() @@ -778,3 +802,4 @@ def main(B: T.Buffer((50,), "int32"), a: T.bool, b: T.bool, x: T.int32): test_let_floordiv_pattern() test_no_lift_bool_predicate() test_no_lift_bool_logical() + test_shared_subtree_stays_ssa() diff --git a/tests/python/tirx-transform/test_tir_transform_split_host_device.py b/tests/python/tirx-transform/test_tir_transform_split_host_device.py index 69bcf5a95f0d..5379b2e3b5ed 100644 --- a/tests/python/tirx-transform/test_tir_transform_split_host_device.py +++ b/tests/python/tirx-transform/test_tir_transform_split_host_device.py @@ -405,6 +405,38 @@ def main_kernel(A_data: T.handle("float32")): tvm.ir.assert_structural_equal(After, Expected) +def test_cuda_launch_preserves_flag_metadata(): + @I.ir_module + class Before: + @T.prim_func(s_tir=True) + def main(A: T.Buffer(16, "float32")): + T.func_attr( + { + "target": T.target("cuda", host="llvm"), + "tirx.kernel_launch_params": [ + "threadIdx.x", + "tirx.use_programtic_dependent_launch", + ], + } + ) + T.attr(T.target("cuda"), "target", 0) + tx = T.launch_thread("threadIdx.x", 16) + A[tx] = 0.0 + + after = tvm.tirx.transform.SplitHostDevice()(Before) + kernel = after["main_kernel"] + assert list(kernel.attrs["tirx.kernel_launch_params"]) == [ + "threadIdx.x", + "tirx.use_programtic_dependent_launch", + ] + + launch = after["main"].body.value + assert isinstance(launch, tvm.ir.Call) + # Programmatic launch is flag-only and therefore adds no packed operand. + assert len(launch.args) == 3 + assert int(launch.args[-1]) == 16 + + def test_device_scope_region_extracted_as_device_kernel(): """A bare device_scope is annotated and extracted as a device kernel.""" diff --git a/tests/python/tirx/codegen/conftest.py b/tests/python/tirx/codegen/conftest.py index 69e14a36c6ef..e6efc9eee924 100644 --- a/tests/python/tirx/codegen/conftest.py +++ b/tests/python/tirx/codegen/conftest.py @@ -16,6 +16,8 @@ # under the License. """Hardware requirements for TIRx codegen tests.""" +import gc +import os from pathlib import Path import pytest @@ -23,7 +25,39 @@ from tvm.testing import env -def pytest_collection_modifyitems(items): +def _set_cuda_device_for_xdist_worker(): + try: + import torch + except ImportError: + return + + if not torch.cuda.is_available(): + return + + worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0") + worker_index = int(worker[2:]) if worker.startswith("gw") and worker[2:].isdigit() else 0 + torch.cuda.set_device(worker_index % torch.cuda.device_count()) + + +def pytest_configure(config): + del config + _set_cuda_device_for_xdist_worker() + + +@pytest.fixture(autouse=True) +def _release_cuda_cache_between_tests(): + _set_cuda_device_for_xdist_worker() + yield + gc.collect() + try: + import torch + except ImportError: + return + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def pytest_collection_modifyitems(config, items): if env.has_cuda_compute(10): return suite_root = Path(__file__).resolve().parent diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py index 99bd4f57de1d..2d2991261f0d 100644 --- a/tests/python/tirx/codegen/test_codegen_blackwell.py +++ b/tests/python/tirx/codegen/test_codegen_blackwell.py @@ -33,7 +33,25 @@ def _get_source(func: tvm.tirx.PrimFunc) -> str: return src, mod -def _assert_remote_mbarrier_ir(func, arrive_op_name, cta_arg_index): +def _remote_operand_index(call): + """Index of the ``remote`` operand in a tirx mbarrier arrive call. + + Both arrive ops pack their optional operands positionally and record which + ones are present in trailing int flags, so the remote operand's index moves + with what the caller actually passed: + + - ``mbarrier_arrive``: ``[bar, count?, remote?, pred?]`` followed by + ``sem, scope, space, has_count, has_remote, has_pred``. + - ``mbarrier_arrive_expect_tx``: ``[bar, byte_count, remote?, pred?]`` + followed by ``sem, scope, space, has_remote, has_pred`` (``byte_count`` + is mandatory, so the index is fixed). + """ + if call.op.name == "tirx.ptx.mbarrier_arrive": + return 1 + int(call.args[-3]) + return 2 + + +def _assert_remote_mbarrier_ir(func, arrive_op_name): bindings = [] buffers = [] mapa_calls = [] @@ -62,7 +80,9 @@ def visit(node): assert buffers[0].scope() == "shared" for arrive in arrive_calls: tvm.ir.assert_structural_equal(arrive.args[0], mapa_calls[0].args[0]) - tvm.ir.assert_structural_equal(arrive.args[cta_arg_index], mapa_calls[0].args[1]) + tvm.ir.assert_structural_equal( + arrive.args[_remote_operand_index(arrive)], mapa_calls[0].args[1] + ) @pytest.mark.gpu @@ -141,12 +161,14 @@ def test_remote_view(): remote_bar.arrive(0, count=2) # fmt: on - _assert_remote_mbarrier_ir(test_remote_view, "tirx.ptx.mbarrier_arrive", 1) + _assert_remote_mbarrier_ir(test_remote_view, "tirx.ptx.mbarrier_arrive") with tvm.target.Target("cuda"): src, _ = _get_source(test_remote_view) assert "tvm_builtin_ptx_mapa_u64" in src - assert "tvm_builtin_ptx_mbarrier_arrive_remote" in src - assert "tvm_builtin_ptx_mbarrier_arrive_remote_count" in src + # The emitted helper name spells out the full attribute set it was + # specialized for: ___. + assert "tvm_builtin_ptx_mbarrier_arrive_shared_cluster_remote_pred" in src + assert "tvm_builtin_ptx_mbarrier_arrive_shared_cluster_count_remote_pred" in src assert "mbarrier.arrive.shared::cluster.b64" in src assert 'asm volatile("mbarrier.arrive.shared.b64' not in src @@ -170,10 +192,10 @@ def test_remote_view(): remote_bar.arrive(0, tx_count=128) # fmt: on - _assert_remote_mbarrier_ir(test_remote_view, "tirx.ptx.mbarrier_arrive_expect_tx", 2) + _assert_remote_mbarrier_ir(test_remote_view, "tirx.ptx.mbarrier_arrive_expect_tx") with tvm.target.Target("cuda"): src, _ = _get_source(test_remote_view) - assert "tvm_builtin_ptx_mbarrier_arrive_expect_tx_remote" in src + assert "tvm_builtin_ptx_mbarrier_arrive_expect_tx_shared_cluster_remote_pred" in src assert "mbarrier.arrive.expect_tx.shared::cluster.b64" in src assert 'asm volatile("mbarrier.arrive.expect_tx.shared.b64' not in src @@ -205,7 +227,7 @@ def invalid_wait(): bar.remote_view(0).wait(0, 0) # fmt: on - with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify cta_id"): + with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify remote"): # fmt: off @T.prim_func def ambiguous_mbarrier_arrive(): @@ -214,10 +236,10 @@ def ambiguous_mbarrier_arrive(): T.thread_id([128]) pool = T.SMEMPool() bar = MBarrier(pool, 1) - bar.remote_view(0).arrive(0, cta_id=1) + bar.remote_view(0).arrive(0, remote=1) # fmt: on - with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify cta_id"): + with pytest.raises(tvm.error.DiagnosticError, match="cannot also specify remote"): # fmt: off @T.prim_func def ambiguous_tma_arrive(): @@ -226,7 +248,7 @@ def ambiguous_tma_arrive(): T.thread_id([128]) pool = T.SMEMPool() bar = TMABar(pool, 1) - bar.remote_view(0).arrive(0, tx_count=128, cta_id=1) + bar.remote_view(0).arrive(0, tx_count=128, remote=1) # fmt: on with pytest.raises( @@ -440,34 +462,16 @@ def test_tcgen05_mma_ss_no_tma(swizzle): B_layout = T.TileLayout(T.S[(N, K // 8, 8) : (8, N * 8, 1)]) ldo, sdo = 128, 8 elif SWIZZLE == 1: - A_layout = T.ComposeLayout( - T.SwizzleLayout(3, 1, 3, swizzle_inner=True), - T.TileLayout(T.S[(M, K // 16, 16) : (16, M * 16, 1)]), - ) - B_layout = T.ComposeLayout( - T.SwizzleLayout(3, 1, 3, swizzle_inner=True), - T.TileLayout(T.S[(N, K // 16, 16) : (16, N * 16, 1)]), - ) + A_layout = T.ComposeLayout(3, 1, 3, T.TileLayout(T.S[(M, K // 16, 16) : (16, M * 16, 1)])) + B_layout = T.ComposeLayout(3, 1, 3, T.TileLayout(T.S[(N, K // 16, 16) : (16, N * 16, 1)])) ldo, sdo = 256, 16 elif SWIZZLE == 2: - A_layout = T.ComposeLayout( - T.SwizzleLayout(3, 2, 3, swizzle_inner=True), - T.TileLayout(T.S[(M, K // 32, 32) : (32, M * 32, 1)]), - ) - B_layout = T.ComposeLayout( - T.SwizzleLayout(3, 2, 3, swizzle_inner=True), - T.TileLayout(T.S[(N, K // 32, 32) : (32, N * 32, 1)]), - ) + A_layout = T.ComposeLayout(3, 2, 3, T.TileLayout(T.S[(M, K // 32, 32) : (32, M * 32, 1)])) + B_layout = T.ComposeLayout(3, 2, 3, T.TileLayout(T.S[(N, K // 32, 32) : (32, N * 32, 1)])) ldo, sdo = 512, 32 elif SWIZZLE == 3: - A_layout = T.ComposeLayout( - T.SwizzleLayout(3, 3, 3, swizzle_inner=True), - T.TileLayout(T.S[(M, 1, 64) : (64, M * 64, 1)]), - ) - B_layout = T.ComposeLayout( - T.SwizzleLayout(3, 3, 3, swizzle_inner=True), - T.TileLayout(T.S[(N, 1, 64) : (64, N * 64, 1)]), - ) + A_layout = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(M, 1, 64) : (64, M * 64, 1)])) + B_layout = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(N, 1, 64) : (64, N * 64, 1)])) ldo, sdo = 1, 64 else: raise ValueError(f"Invalid swizzle: {SWIZZLE}") @@ -565,5 +569,46 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_tcgen05_mma_pred_codegen(): + # fmt: off + @T.prim_func + def test_mma_pred(): + T.device_entry() + T.thread_id([1]) + tmem_addr = T.alloc_buffer((1,), "uint32", scope="local") + desc_a = T.alloc_buffer((1,), "uint64", scope="local") + desc_b = T.alloc_buffer((1,), "uint64", scope="local") + desc_i = T.alloc_buffer((1,), "uint32", scope="local") + pred = T.alloc_buffer((1,), "uint32", scope="local") + + tmem_addr[0] = T.uint32(0) + desc_a[0] = T.uint64(0) + desc_b[0] = T.uint64(0) + desc_i[0] = T.uint32(0) + pred[0] = T.uint32(1) + T.ptx.tcgen05.mma( + tmem_addr[0], + desc_a[0], + desc_b[0], + desc_i[0], + d_dtype="float32", + a_dtype="float16", + b_dtype="float16", + use_a_tmem=False, + cta_group=1, + pred=pred[0], + ) + # fmt: on + + target = tvm.target.Target("cuda") + with target: + src, _ = _get_source(test_mma_pred) + assert "ptx_tcgen05_mma_cta_1_kind_f16_SS_pred" in src + assert "setp.ne.b32 p_issue" in src + assert "@p_issue tcgen05.mma.cta_group::1.kind::f16" in src + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py b/tests/python/tirx/codegen/test_codegen_cuda.py index 1b31c2b0ee03..1418d0dd3f62 100644 --- a/tests/python/tirx/codegen/test_codegen_cuda.py +++ b/tests/python/tirx/codegen/test_codegen_cuda.py @@ -15,14 +15,38 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring +import re + import numpy as np import pytest import tvm import tvm.testing +from tvm import tirx as tir from tvm.script import tirx as T from tvm.testing import env +_CUDA_LDG_SCALAR_CASES = [ + ("int8", "i8", "signed char"), + ("uint8", "u8", "unsigned char"), + ("int16", "i16", "short"), + ("uint16", "u16", "unsigned short"), + ("int32", "i32", "int"), + ("uint32", "u32", "unsigned int"), + ("int64", "i64", "long long"), + ("uint64", "u64", "unsigned long long"), + ("float16", "f16", "half"), + ("bfloat16", "bf16", "nv_bfloat16"), + ("float32", "f32", "float"), + ("float64", "f64", "double"), +] + +_CUDA_LDG_VECTOR_CASES = [ + ("int32", "i32", "int", "int"), + ("uint32", "u32", "unsigned int", "uint"), + ("float32", "f32", "float", "float"), +] + def _get_source(func: tvm.tirx.PrimFunc, target=None) -> tuple[str, tvm.IRModule]: if target is None: @@ -36,7 +60,11 @@ def _get_source(func: tvm.tirx.PrimFunc, target=None) -> tuple[str, tvm.IRModule def _helper_source(src: str, helper_name: str) -> str: - start = src.index(helper_name) + pattern = rf"__forceinline__ __device__ [^{{;]+ {re.escape(helper_name)}\(" + match = re.search(pattern, src) + if match is None: + raise ValueError(f"helper {helper_name!r} not found") + start = match.start() next_helper = src.find("__device__", start + len(helper_name)) if next_helper == -1: return src[start:] @@ -69,6 +97,63 @@ def test_vector_access_ptr_preserves_packed_offset(monkeypatch): assert " + 8 / 4" in call +def _cuda_ldg_scalar_kernel(dtype: str): + @T.prim_func + def main(src: T.Buffer((1,), dtype), out: T.Buffer((1,), dtype)): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + out[0] = T.cuda.ldg(src.data, dtype) + + return main + + +def _cuda_ldg_vector_kernel(dtype: str, vec: str): + vec_len = int(vec[1:]) + + if vec_len == 2: + + @T.prim_func + def main(src: T.Buffer((2,), dtype), out: T.Buffer((2,), dtype)): + T.device_entry() + tx = T.thread_id([32]) + tmp0 = T.alloc_local((1,), dtype) + tmp1 = T.alloc_local((1,), dtype) + if tx == 0: + T.cuda.ldg(src.data, dtype, dst=(tmp0.ptr_to([0]), tmp1.ptr_to([0])), vec=vec) + out[0] = tmp0[0] + out[1] = tmp1[0] + + return main + + @T.prim_func + def main(src: T.Buffer((4,), dtype), out: T.Buffer((4,), dtype)): + T.device_entry() + tx = T.thread_id([32]) + tmp0 = T.alloc_local((1,), dtype) + tmp1 = T.alloc_local((1,), dtype) + tmp2 = T.alloc_local((1,), dtype) + tmp3 = T.alloc_local((1,), dtype) + if tx == 0: + T.cuda.ldg( + src.data, + dtype, + dst=( + tmp0.ptr_to([0]), + tmp1.ptr_to([0]), + tmp2.ptr_to([0]), + tmp3.ptr_to([0]), + ), + vec=vec, + ) + out[0] = tmp0[0] + out[1] = tmp1[0] + out[2] = tmp2[0] + out[3] = tmp3[0] + + return main + + def test_tirx_launch_bounds_omits_min_blocks_without_persistent_schedule(): @T.prim_func def main(A: T.Buffer((4,), "int32")): @@ -98,6 +183,45 @@ def main(A: T.Buffer((4,), "int32")): assert "tirx.launch_bounds_min_blocks_per_sm" not in src +def test_tirx_launch_bounds_max_blocks_per_cluster_emits_third_operand(): + @T.prim_func + def main(A: T.Buffer((4,), "int32")): + T.device_entry() + T.attr( + { + "tirx.launch_bounds_min_blocks_per_sm": 1, + "tirx.launch_bounds_max_blocks_per_cluster": 1, + } + ) + bx = T.cta_id([4]) + tx = T.thread_id([384]) + if tx == 0: + A[bx] = A[bx] + 1 + + src, _ = _get_source(main) + assert 'extern "C" __global__ void __launch_bounds__(384, 1, 1) main_kernel' in src + assert "tirx.launch_bounds_max_blocks_per_cluster" not in src + + +def test_tirx_cuda_kernel_return_zero_codegen_is_void_early_return(): + @T.prim_func + def main(A: T.Buffer((4,), "int32")): + T.device_entry() + bx = T.cta_id([4]) + tx = T.thread_id([32]) + if bx >= 3: + return 0 + if tx == 0: + A[bx] = A[bx] + 1 + + src, _ = _get_source(main) + # The bounded blockIdx.x domain simplifies ``blockIdx.x >= 3`` to the + # equivalent final-point predicate, including CUDA's explicit index cast. + assert re.search(r"if \(\(\(int\)blockIdx\.x\) == 3\)", src) + assert "return;" in src + assert "return 0;" not in src + + def test_serial_pragma_unroll_codegen(): @T.prim_func def main(A: T.Buffer((4,), "int32")): @@ -115,6 +239,21 @@ def main(A: T.Buffer((4,), "int32")): assert "break;" in src +def test_serial_disable_unroll_pragma_immediately_precedes_dynamic_for(): + @T.prim_func + def main(A: T.Buffer((4,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + begin: T.let = T.if_then_else(A[0] > 0, A[1], A[2]) + end: T.let = T.if_then_else(A[0] > 1, A[2], A[3]) + for i in T.serial(begin, end, unroll=False): + A[0] = A[0] + i + + src, _ = _get_source(main) + assert re.search(r"#pragma unroll 1\s*for \(", src) + + def test_cluster_cta_id_codegen_uses_coordinate_sregs(): @T.prim_func def main(A: T.Buffer((1,), "int32")): @@ -196,6 +335,68 @@ def main(A: T.Buffer((1,), "uint64"), B: T.Buffer((1,), "int32"), C: T.Buffer((1 assert "ld.volatile.global.u64" in src +def test_ptx_f32x2_value_codegen(): + @T.prim_func + def main(A: T.Buffer((2,), "uint64"), B: T.Buffer((2,), "float32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + lhs: T.let = T.cuda.make_float2(B[0], B[1]) + rhs: T.let = T.cuda.make_float2(B[1], B[0]) + prod: T.let = T.ptx.mul_f32x2(lhs, rhs, dps=False) + A[0] = T.ptx.fma_f32x2(lhs, rhs, prod, dps=False) + sum_pair: T.let = T.ptx.add_f32x2(lhs, rhs, dps=False, return_dtype="float32x2") + A[1] = T.ptx.sub_f32x2(sum_pair, rhs, dps=False) + + src, _ = _get_source(main) + assert "tirx.ptx.fma_f32x2_val" not in src + assert "tvm_builtin_ptx_mul_f32x2_ret_u64" in src + assert "tvm_builtin_ptx_fma_f32x2_ret_u64_rn" in src + assert "tvm_builtin_ptx_add_f32x2_ret_f32x2_rn" in src + assert "tvm_builtin_ptx_sub_f32x2_ret_u64_rn" in src + assert src.count("tvm_builtin_ptx_mul_f32x2_ret_u64") == 2 + assert src.count("tvm_builtin_ptx_fma_f32x2_ret_u64_rn") == 2 + assert "A_ptr[0] = tvm_builtin_ptx_fma_f32x2_ret_u64_rn(lhs, rhs, prod);" in src + assert "float2 tvm_builtin_ptx_add_f32x2_ret_f32x2_rn" in src + assert "float2 sum_pair" in src + assert "return *reinterpret_cast(&result);" in src + assert "mul.f32x2 %0, %1, %2;" in src + assert "mul.rn.f32x2 %0, %1, %2;" not in src + assert "fma.rn.f32x2 %0, %1, %2, %3;" in src + + +@pytest.mark.skipif( + not (env.has_cuda_compute(10, 0) and env.has_nvcc_version(13, 2)), + reason="PTX 9.2 packed bf16 conversion requires sm_100 and CUDA 13.2", +) +def test_sparse_decode_conversion_intrinsics_codegen(monkeypatch): + monkeypatch.setenv("TVM_CUDA_COMPILE_MODE", "nvcc") + + @T.prim_func + def main( + U16: T.Buffer((1,), "uint16"), + U32: T.Buffer((2,), "uint32"), + U64: T.Buffer((1,), "uint64"), + F32: T.Buffer((2,), "float32"), + ): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + pair: T.let = T.cuda.make_float2(F32[0], F32[1]) + U16[0] = T.ptx.cvt(F32[1], F32[0], dtype="ue8m0x2", atype="f32", rounding="rz") + U32[0] = T.ptx.cvt(U16[0], dtype="bf16x2", atype="ue8m0x2", rounding="rn") + U32[1] = T.ptx.cvt(U16[0], dtype="bf16x2", atype="e4m3x2", rounding="rn") + U64[0] = T.ptx.add_f32x2(pair, pair, rounding="", dps=False) + + src, _ = _get_source(main) + assert "cvt.rz.ue8m0x2.f32 %0, %1, %2;" in src + assert "cvt.rn.bf16x2.ue8m0x2 %0, %1;" in src + assert "cvt.rn.bf16x2.e4m3x2 %0, %1;" in src + assert "__nv_cvt_" not in src + assert "add.f32x2 %0, %1, %2;" in src + assert "add.rn.f32x2 %0, %1, %2;" not in src + + @pytest.mark.gpu def test_megamoe_extracted_intrinsics_codegen(): @T.prim_func @@ -266,6 +467,8 @@ def main( F32[1] = T.cuda.uint_as_float(U32[0]) F32[2] = T.ptx.ld(F32.data, "float32", "f32", space="global") + F32[2] = T.cuda.ldg(T.handle_add_byte_offset(F32.data, 4), "float32") + F32[3] = T.cuda.fdividef(F32[0], F32[1]) U32[3] = T.cuda.float_as_uint(F32[1]) F32[0] = T.ptx.add_rn_f32_bf16(F32[0], T.cast(U32[0], "uint16")) U64[0] = T.reinterpret("uint64", U32.data) @@ -289,6 +492,9 @@ def main( "fns.b32", "stmatrix.sync.aligned.m16n8.x1.trans.shared.b8", "ld.global.f32", + "tvm_builtin_cuda_ldg_f32", + "reinterpret_cast", + "__fdividef", "add.rn.f32.bf16", "__uint_as_float", "__float_as_uint", @@ -302,6 +508,34 @@ def main( assert snippet in src +@pytest.mark.parametrize(("dtype", "suffix", "c_type"), _CUDA_LDG_SCALAR_CASES) +def test_cuda_ldg_scalar_dtype_codegen(dtype, suffix, c_type): + src, _ = _get_source(_cuda_ldg_scalar_kernel(dtype)) + helper = f"tvm_builtin_cuda_ldg_{suffix}" + helper_src = _helper_source(src, helper) + assert f"__forceinline__ __device__ {c_type} {helper}(void* src)" in src + assert f"__ldg(reinterpret_cast(src))" in helper_src + + +@pytest.mark.parametrize(("dtype", "suffix", "c_type", "vec_base"), _CUDA_LDG_VECTOR_CASES) +@pytest.mark.parametrize("vec", ["v2", "v4"]) +def test_cuda_ldg_vector_dtype_codegen(dtype, suffix, c_type, vec_base, vec): + vec_len = int(vec[1:]) + src, _ = _get_source(_cuda_ldg_vector_kernel(dtype, vec)) + helper = f"tvm_builtin_cuda_ldg_{suffix}_{vec}_to_dst{vec_len}" + helper_src = _helper_source(src, helper) + assert ( + f"{vec_base}{vec_len} v = __ldg(reinterpret_cast(src));" + in helper_src + ) + assert f"*reinterpret_cast<{c_type}*>(dst{vec_len - 1}) = v." in helper_src + + +def test_cuda_ldg_vector_rejects_unsupported_dtype(): + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="Unsupported vector CUDA"): + _get_source(_cuda_ldg_vector_kernel("float16", "v2")) + + def test_ptx_cp_async_bulk_non_tma_form_codegen(): @T.prim_func def main( @@ -320,12 +554,196 @@ def main( smem.ptr_to([0]), A.data, T.uint32(64), smem.ptr_to([0]), cache_policy=C[0] ) T.ptx.cp_async_bulk_s2g(B.data, smem.ptr_to([0]), T.uint32(64), cache_policy=C[0]) + T.ptx.cp_async.bulk.commit_group() + T.ptx.cp_async.bulk.wait_group(0, read=True) + T.ptx.cp_async.bulk.wait_group(1, read=False) src, _ = _get_source(main) assert "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes.L2::cache_hint" in src assert "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint" in src assert "cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint" in src assert "unsigned long long cache_policy" in src + assert 'asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory");' in src + assert 'asm volatile("cp.async.bulk.wait_group 1;" ::: "memory");' in src + + +def test_ptx_sync_and_clc_codegen(): + @T.prim_func + def main(A: T.Buffer((1,), "uint32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + bar = T.alloc_buffer((5,), "uint64", scope="shared", align=16) + response = T.alloc_buffer((4,), "uint32", scope="shared", align=16) + T.ptx.cp_async.mbarrier.arrive(bar.ptr_to([0])) + T.ptx.cp_async.mbarrier.arrive.noinc(bar.ptr_to([0])) + T.ptx.mbarrier.complete_tx(bar.ptr_to([0]), T.uint32(16)) + T.ptx.mbarrier.complete_tx( + bar.ptr_to([1]), T.uint32(24), scope="cta", space="shared::cta" + ) + T.ptx.mbarrier.complete_tx( + bar.ptr_to([2]), + T.uint32(28), + scope="cta", + space="shared::cta", + pred=T.uint32(1), + ) + T.ptx.mbarrier.complete_tx(bar.ptr_to([3]), T.uint32(32), remote=T.int32(1)) + T.ptx.mbarrier.complete_tx( + bar.ptr_to([4]), T.uint32(40), remote=T.int32(1), pred=T.uint32(1) + ) + T.ptx.clc_try_cancel(response.ptr_to([0]), bar.ptr_to([0])) + A[0] = T.ptx.clc_query_cancel(response.ptr_to([0])) + A[0] = T.ptx.clc_query_cancel(response.ptr_to([0]), use_ld_acquire=False) + T.ptx.griddepcontrol.launch_dependents() + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile(tvm.IRModule({"main": main}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + assert "cp.async.mbarrier.arrive.shared.b64" in src + assert "cp.async.mbarrier.arrive.noinc.shared::cta.b64" in src + assert "tirx.ptx.cp_async_mbarrier_arrive_noinc" not in src + assert "mbarrier.complete_tx.relaxed.cluster.shared::cluster.b64" in src + assert "mbarrier.complete_tx.relaxed.cta.shared::cta.b64" in src + assert "@p mbarrier.complete_tx.relaxed.cta.shared::cta.b64" in src + assert "@p mbarrier.complete_tx.relaxed.cluster.shared::cluster.b64" in src + pred_only = _helper_source( + src, + "tvm_builtin_ptx_mbarrier_complete_tx_relaxed_cta_shared_cta_pred", + ) + remote_only = _helper_source( + src, + "tvm_builtin_ptx_mbarrier_complete_tx_relaxed_cluster_shared_cluster_remote", + ) + remote_pred = _helper_source( + src, + "tvm_builtin_ptx_mbarrier_complete_tx_relaxed_cluster_shared_cluster_remote_pred", + ) + assert "@p mbarrier.complete_tx.relaxed.cta.shared::cta.b64" in pred_only + assert "mapa.shared::cluster.u32" not in pred_only + assert "mapa.shared::cluster.u32" in remote_only + assert "@p mbarrier.complete_tx" not in remote_only + assert "mapa.shared::cluster.u32" in remote_pred + assert "@p mbarrier.complete_tx.relaxed.cluster.shared::cluster.b64" in remote_pred + assert "mbarrier.complete_tx.shared::cluster.relaxed.cluster.b64" not in src + assert "mapa.shared::cluster.u32" in src + assert "clusterlaunchcontrol.try_cancel.async.shared::cta" in src + assert "ld.acquire.cta.shared.b128" in src + assert "ld.shared.b128" in src + assert "clusterlaunchcontrol.query_cancel.get_first_ctaid::x.b32.b128" in src + assert "griddepcontrol.launch_dependents" in src + + +def test_ptx_mbarrier_complete_tx_rejects_remote_non_cluster_space(): + bar = tir.Var("bar", "handle") + with pytest.raises(ValueError, match="requires space='shared::cluster'"): + T.ptx.mbarrier.complete_tx(bar, T.uint32(16), space="shared::cta", remote=T.int32(0)) + + +def test_ptx_mbarrier_arrive_remote_codegen(): + @T.prim_func + def main(Pred: T.Buffer((1,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + bar = T.alloc_buffer((3,), "uint64", scope="shared", align=16) + T.ptx.mbarrier.arrive(bar.ptr_to([0]), remote=T.int32(0), pred=True) + T.ptx.mbarrier.arrive(bar.ptr_to([1]), remote=T.int32(0), pred=Pred[0]) + T.ptx.mbarrier.arrive(bar.ptr_to([2]), remote=T.int32(0), pred=True, count=T.int32(2)) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile(tvm.IRModule({"main": main}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + assert "tvm_builtin_ptx_mbarrier_arrive_remote_unpred" not in src + assert "tvm_builtin_ptx_mbarrier_arrive_remote_count_unpred" not in src + assert "tvm_builtin_ptx_mbarrier_arrive_shared_cluster_remote_pred(" in src + assert "tvm_builtin_ptx_mbarrier_arrive_shared_cluster_count_remote_pred(" in src + assert src.count("@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32];") == 1 + assert "@p mbarrier.arrive.shared::cluster.b64 _, [remAddr32], %1;" in src + + +def test_ptx_mbarrier_arrive_new_forms_codegen(): + @T.prim_func + def main(Pred: T.Buffer((1,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + bar = T.alloc_buffer((6,), "uint64", scope="shared", align=16) + T.ptx.mbarrier.arrive(bar.ptr_to([0]), sem="relaxed", scope="cta", space="shared::cta") + T.ptx.mbarrier.arrive( + bar.ptr_to([1]), sem="relaxed", scope="cluster", remote=T.int32(1) + ) + T.ptx.mbarrier.arrive(bar.ptr_to([2]), sem="release", scope="cta", pred=Pred[0]) + T.ptx.mbarrier.arrive.expect_tx( + bar.ptr_to([3]), T.int32(128), sem="relaxed", scope="cluster", remote=T.int32(1) + ) + T.ptx.mbarrier.arrive.no_complete(bar.ptr_to([4]), T.int32(2)) + T.ptx.mbarrier.arrive.no_complete( + bar.ptr_to([5]), T.int32(3), space="shared::cta", pred=Pred[0] + ) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile(tvm.IRModule({"main": main}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + assert "mbarrier.arrive.relaxed.cta.shared::cta.b64 _, [%0];" in src + assert "mapa.shared::cluster.u32 remAddr32, %0, %2;" in src + assert "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [remAddr32];" in src + assert "@p mbarrier.arrive.release.cta.shared.b64 _, [%0];" in src + assert "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cluster.b64" in src + assert "mbarrier.arrive.noComplete.release.cta.shared.b64 _, [%0], %1;" in src + assert "@p mbarrier.arrive.noComplete.release.cta.shared::cta.b64 _, [%0], %1;" in src + + +def test_ptx_mbarrier_arrive_rejects_removed_and_invalid_forms(): + bar = tir.Var("bar", "handle") + with pytest.raises(ValueError, match="remote=.*cta_id"): + T.ptx.mbarrier.arrive(bar, cta_id=T.int32(0), pred=True) + with pytest.raises(ValueError, match="remote=.*cta_id"): + T.ptx.mbarrier.arrive.expect_tx(bar, T.int32(128), cta_id=T.int32(0)) + assert not hasattr(T.ptx.mbarrier.arrive, "cluster_count") + with pytest.raises(ValueError, match="sem and scope"): + T.ptx.mbarrier.arrive(bar, sem="relaxed") + with pytest.raises(ValueError, match="requires space='shared::cluster'"): + T.ptx.mbarrier.arrive(bar, remote=T.int32(0), space="shared::cta") + with pytest.raises(ValueError, match="does not support remote"): + T.ptx.mbarrier.arrive.no_complete(bar, T.int32(1), remote=T.int32(0)) + with pytest.raises(ValueError, match="space must"): + T.ptx.mbarrier.arrive.no_complete(bar, T.int32(1), space="shared::cluster") + + +def test_cuda_ldg_vector_scatter_codegen(): + @T.prim_func + def main(src: T.Buffer((4,), "int32"), out: T.Buffer((4,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + tmp0 = T.alloc_local((1,), "int32") + tmp1 = T.alloc_local((1,), "int32") + tmp2 = T.alloc_local((1,), "int32") + tmp3 = T.alloc_local((1,), "int32") + if tx == 0: + T.cuda.ldg( + src.data, + "int32", + dst=( + tmp0.ptr_to([0]), + tmp1.ptr_to([0]), + tmp2.ptr_to([0]), + tmp3.ptr_to([0]), + ), + vec="v4", + ) + out[0] = tmp0[0] + out[1] = tmp1[0] + out[2] = tmp2[0] + out[3] = tmp3[0] + + src, _ = _get_source(main) + assert "int4 v = __ldg(reinterpret_cast(src));" in src + assert "tvm_builtin_cuda_ldg_i32_v4_to_dst4" in src + assert "*reinterpret_cast(dst3) = v.w" in src def test_tensor_map_param_codegen(): @@ -352,7 +770,7 @@ def main(Cache: T.Buffer((1,), "uint64")): if tx == 0: smem = T.alloc_buffer((128,), "float32", scope="shared", align=128) bar = T.shared_scalar("uint64") - T.ptx.cp_async.bulk.tensor.g2c( + T.ptx.cp_async.bulk.tensor.g2s_cluster( 2, smem.data, T.address_of(bar), @@ -364,7 +782,7 @@ def main(Cache: T.Buffer((1,), "uint64")): 0, cache_policy=Cache[0], ) - T.ptx.cp_async.bulk.tensor.g2c( + T.ptx.cp_async.bulk.tensor.g2s_cluster( 2, smem.data, T.address_of(bar), @@ -379,51 +797,70 @@ def main(Cache: T.Buffer((1,), "uint64")): T.ptx.cp_async.bulk.tensor.s2g( 2, smem.data, T.address_of(A_map), "", 0, 0, cache_policy=Cache[0] ) - masked_bar = T.cuda.sm100_tma_2sm_mbarrier_addr(T.address_of(bar)) - T.ptx.cp_async.bulk.tensor.g2c_bar_addr( + T.ptx.cp_async.bulk.tensor.g2s_cta( 2, smem.data, - masked_bar, + T.address_of(bar), + T.address_of(A_map), + 2, + "", + 0, + 1, + 2, + 3, + 4, + load_mode="tile_gather4", + cache_policy=Cache[0], + ) + leader_mbar_addr = T.cuda.sm100_2sm_leader_smem_addr(T.address_of(bar)) + T.ptx.cp_async.bulk.tensor.g2s_cluster( + 2, + smem.data, + leader_mbar_addr, T.address_of(A_map), 1, 2, "", 0, 0, + mbar_is_shared_addr=True, cache_policy=Cache[0], ) if tx == 0: - T.ptx.cp_async.bulk.tensor.g2c_bar_addr( + T.ptx.cp_async.bulk.tensor.g2s_cluster( 2, smem.data, - masked_bar, + leader_mbar_addr, T.address_of(A_map), 1, 2, "", 0, 0, + mbar_is_shared_addr=True, cache_policy=Cache[0], ) else: - T.ptx.cp_async.bulk.tensor.g2c_bar_addr( + T.ptx.cp_async.bulk.tensor.g2s_cluster( 2, smem.data, - masked_bar, + leader_mbar_addr, T.address_of(B_map), 1, 2, "", 0, 0, + mbar_is_shared_addr=True, cache_policy=Cache[0], ) - src, _ = _get_source(main, {"kind": "cuda", "arch": "sm_100a"}) - assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_cache_hint" in src - assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_multicast_cache_hint" in src - assert "g2cluster_unicast" not in src - assert "ptx_cp_async_bulk_tensor_g2cta" not in src + src, _ = _get_source(main) + assert "ptx_cp_async_bulk_tensor_g2s_cluster_tile_2d_cache_hint" in src + assert "ptx_cp_async_bulk_tensor_g2s_cta_tile_gather4_2d_cache_hint" in src + assert "ptx_cp_async_bulk_tensor_g2s_cluster_tile_2d_multicast_cache_hint" in src + assert "g2cluster" not in src + assert "cta_group2_unicast" not in src assert ( "cp.async.bulk.tensor.2d.shared::cluster.global" ".mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint" @@ -433,10 +870,15 @@ def main(Cache: T.Buffer((1,), "uint64")): ".mbarrier::complete_tx::bytes.multicast::cluster" ".cta_group::2.L2::cache_hint" ) in src + assert ( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4" + ".mbarrier::complete_tx::bytes.cta_group::2.L2::cache_hint" + ) in src assert "cp.async.bulk.tensor.2d.global.shared::cta.tile.bulk_group.L2::cache_hint" in src - assert "tvm_builtin_cp_async_bulk_tensor_2d_g2c_cta_group2" not in src + assert "bar_addr &= 0xFEFFFFFFu;" not in src + assert "mbar_addr &= 0xFEFFFFFFu;" not in src assert "tvm_builtin_cuda_cvta_generic_to_shared((&(bar_ptr[0]))) & (uint)4278190079" in src - assert "ptx_cp_async_bulk_tensor_g2cluster_tile_2d_cache_hint_bar_addr" in src + assert "ptx_cp_async_bulk_tensor_g2s_cluster_tile_2d_cache_hint_mbar_addr" in src assert "unsigned long long cache_policy" in src diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py index e3405ce3fe28..ac5edb872afb 100644 --- a/tests/python/tirx/codegen/test_codegen_hopper.py +++ b/tests/python/tirx/codegen/test_codegen_hopper.py @@ -316,6 +316,23 @@ def func(A: T.Buffer(1)): assert 'bar.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory"' in src +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_barrier_sync_unaligned(): + # fmt: off + @T.prim_func + def func(A: T.Buffer(1)): + T.device_entry() + cta_id = T.cta_id([1]) + tid = T.thread_id([128]) + T.ptx.barrier.sync(0, 128) + # fmt: on + + src, mod = _get_source(func) + assert "tvm_builtin_ptx_barrier_sync(0, 128)" in src + assert 'barrier.sync %0, %1;" : : "r"(name_bar_id), "r"(thread_count) : "memory"' in src + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") def test_fence_mbarrier_init_release_clsuter(): @@ -430,7 +447,7 @@ def main(A_ptr: T.handle, B_ptr: T.handle): if threadIdx == 0: T.ptx.mbarrier.init(T.address_of(bar), 1) T.ptx.fence.proxy_async("shared::cta") - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501 T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes) T.ptx.mbarrier.try_wait(T.address_of(bar), phase) phase = phase ^ 1 @@ -522,6 +539,12 @@ def run_and_check(): [16, 16, 32, 16, 16, 9, 1, 0, 0, 0, 0], r"elementStrides\[0\] must be less than or equal to 8", ), + ( + (16, 16), + "float16", + [16, 16, 32, 16, 16, 2, 1, 0, 0, 0, 0], + r"elementStrides\[0\] must be one", + ), ( (16, 16), "float16", @@ -531,9 +554,15 @@ def run_and_check(): ( (8, 8, 8), "float16", - [8, 8, 8, 16, 128, 8, 8, 8, 1, 1, 1, 2, 0, 0, 0], + [8, 8, 8, 16, 128, 8, 8, 8, 1, 1, 1, 2, 1, 0, 0], r"globalStrides\[0\] must be a multiple of 32", ), + ( + (8, 8, 8), + "float16", + [8, 8, 8, 32, 256, 8, 8, 8, 1, 1, 1, 2, 2, 0, 0], + r"CU_TENSOR_MAP_INTERLEAVE_32B requires CU_TENSOR_MAP_SWIZZLE_32B", + ), ( (16, 16), "int32", @@ -543,6 +572,18 @@ def run_and_check(): r"floating-point tensorDataType" ), ), + ( + (16, 16), + "float32", + [16, 16, 64, 4, 16, 1, 1, 0, 0, 0, 0, 7], + r"force_cu_dtype only supports CU_TENSOR_MAP_DATA_TYPE_TFLOAT32", + ), + ( + (16, 16), + "float16", + [16, 16, 32, 8, 16, 1, 1, 0, 0, 0, 0, 11], + r"CU_TENSOR_MAP_DATA_TYPE_TFLOAT32 requires a scalar float32", + ), ], ) def test_tensormap_encode_tiled_runtime_validation(shape, dtype, encode_args, error_msg): @@ -582,9 +623,9 @@ def main(A_ptr: T.handle, B_ptr: T.handle): B = T.match_buffer(B_ptr, total_elems, dtype=dtype, align=16) A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1) - T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, dtype, len(shape), A.data, *load_args) # noqa: E501 + T.call_packed("runtime.cuTensorMapEncodeTiled", A_map, str(dtype), len(shape), A.data, *load_args) # noqa: E501 B_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1) - T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, dtype, len(shape), B.data, *store_args) # noqa: E501 + T.call_packed("runtime.cuTensorMapEncodeTiled", B_map, str(dtype), len(shape), B.data, *store_args) # noqa: E501 T.device_entry() for blockIdx in T.thread_binding(1, thread="blockIdx.x"): @@ -597,7 +638,7 @@ def main(A_ptr: T.handle, B_ptr: T.handle): if threadIdx == 0: T.ptx.mbarrier.init(T.address_of(bar), 1) T.ptx.fence.proxy_async("shared::cta") - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coord) # noqa: E501 T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes) T.ptx.mbarrier.try_wait(T.address_of(bar), phase) phase = phase ^ 1 @@ -625,8 +666,11 @@ def main(A_ptr: T.handle, B_ptr: T.handle): A_np = np.array(A_np).astype(dtype) B_np = np.zeros((total_elems,)).astype(dtype) dtype = tvm.DataType(dtype) - layout = T.SwizzleLayout( - per_element=int(math.log2(128 // dtype.bits)), swizzle_len=swizzle, atom_len=3 + layout = T.ComposeLayout( + int(math.log2(128 // dtype.bits)), + swizzle, + 3, + T.TileLayout(T.S[(1 << (int(math.log2(128 // dtype.bits)) + swizzle + 3),)]), ) def run_and_check(): @@ -690,7 +734,7 @@ def main(A_ptr: T.handle, B_ptr: T.handle): T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes) if clusterCtaIdx == 0: # only the first CTA in the cluster does the copy, and then multicast # noqa: E501 - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.data, T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord) # noqa: E501 # wait for the copy to finish T.ptx.mbarrier.try_wait(T.address_of(bar), phase) phase = phase ^ 1 @@ -949,8 +993,8 @@ def main(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle): T.cuda.cta_sync() # load A and B to smem if tx == 0: - T.ptx.cp_async.bulk.tensor.g2c(len(shapeA), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coordA) # noqa: E501 - T.ptx.cp_async.bulk.tensor.g2c(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shapeA), A_smem.data, T.address_of(bar), T.address_of(A_map), 0, 1, "", *coordA) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501 T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), A_bytes + B_bytes) T.ptx.mbarrier.try_wait(T.address_of(bar), phase) phase = phase ^ 1 @@ -1109,7 +1153,7 @@ def main(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle): T.cuda.cta_sync() # load B to smem if tx == 0: - T.ptx.cp_async.bulk.tensor.g2c(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shapeB), B_smem.data, T.address_of(bar), T.address_of(B_map), 0, 1, "", *coordB) # noqa: E501 T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), B_bytes) T.ptx.mbarrier.try_wait(T.address_of(bar), 0) T.cuda.cta_sync() diff --git a/tests/python/tirx/codegen/test_ptx_cvt.py b/tests/python/tirx/codegen/test_ptx_cvt.py new file mode 100644 index 000000000000..e53c762cf8e6 --- /dev/null +++ b/tests/python/tirx/codegen/test_ptx_cvt.py @@ -0,0 +1,151 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for the PTX ``cvt`` instruction-family registration.""" + +import pytest + +from tvm.backend.cuda.operator.intrinsics.cvt import _CVT_FORMS, _cvt_form_parts +from tvm.backend.cuda.operator.intrinsics.registry import CODEGEN_REGISTRY + +# One concrete expansion for every syntax form in the PTX ``cvt`` grammar. +# Args are register placeholders followed by the form-specific attrs. +_FORM_CASES = [ + ("scalar", (0, "f32", "f32", 0, 0), "cvt.f32.f32 %0, %1;"), + ("irnd", (0, "f32", "f32", "rni", 0, 0), "cvt.rni.f32.f32 %0, %1;"), + ("frnd", (0, "f32", "f64", "rn", 0, 0), "cvt.rn.f32.f64 %0, %1;"), + ("frnd2_f16_f32", (0, "rn", 0, 0), "cvt.rn.f16.f32 %0, %1;"), + ("frnd2_f16x2_f32", (0, 0, "rz", 1, 1), "cvt.rz.relu.satfinite.f16x2.f32"), + ("rs_f16x2_f32", (0, 0, 0, 1, 0), "cvt.rs.relu.f16x2.f32"), + ("frnd2_bf16_f32", (0, "rn", 0, 1), "cvt.rn.satfinite.bf16.f32"), + ("frnd2_bf16x2_f32", (0, 0, "rz", 0, 0), "cvt.rz.bf16x2.f32"), + ("rs_bf16x2_f32", (0, 0, 0, 0, 1), "cvt.rs.satfinite.bf16x2.f32"), + ("rna_tf32_f32", (0, 1), "cvt.rna.satfinite.tf32.f32"), + ("frnd2_tf32_f32", (0, "rn", 1, 1), "cvt.rn.satfinite.relu.tf32.f32"), + ( + "rn_satfinite_f8x2_f32", + (0, 0, "e4m3x2", 1), + "cvt.rn.satfinite.relu.e4m3x2.f32", + ), + ( + "rn_satfinite_f8x2_fp16x2", + (0, "e5m2x2", "bf16x2", 0), + "cvt.rn.satfinite.e5m2x2.bf16x2", + ), + ("rn_f16x2_f8x2", (0, "e4m3x2", 1), "cvt.rn.relu.f16x2.e4m3x2"), + ( + "rn_bf16x2_f8x2", + (0, 0, "e4m3x2", 1, 1, "n2::ue8m0"), + "cvt.rn.relu.satfinite.scaled::n2::ue8m0.bf16x2.e4m3x2", + ), + ( + "rs_satfinite_f8x4_f32", + (0, 0, 0, 0, 0, "e4m3x4", 1), + "cvt.rs.relu.satfinite.e4m3x4.f32", + ), + ( + "rn_satfinite_f4x2_f32", + (0, 0, "e2m1x2", 0), + "cvt.rn.satfinite.e2m1x2.f32", + ), + ( + "rn_satfinite_f4x2_fp16x2", + (0, "e2m1x2", "f16x2", 1), + "cvt.rn.satfinite.relu.e2m1x2.f16x2", + ), + ("rn_f16x2_f4x2", (0, "e2m1x2", 0), "cvt.rn.f16x2.e2m1x2"), + ( + "rn_bf16x2_f4x2", + (0, 0, "e2m1x2", 0, 0, "n2::ue8m0"), + "cvt.rn.scaled::n2::ue8m0.bf16x2.e2m1x2", + ), + ( + "rs_satfinite_f4x4_f32", + (0, 0, 0, 0, 0, "e2m1x4", 0), + "cvt.rs.satfinite.e2m1x4.f32", + ), + ( + "rn_satfinite_f6x2_f32", + (0, 0, "e2m3x2", 1), + "cvt.rn.satfinite.relu.e2m3x2.f32", + ), + ( + "rn_satfinite_f6x2_fp16x2", + (0, "e3m2x2", "bf16x2", 0), + "cvt.rn.satfinite.e3m2x2.bf16x2", + ), + ("rn_f16x2_f6x2", (0, "e2m3x2", 1), "cvt.rn.relu.f16x2.e2m3x2"), + ( + "rn_bf16x2_f6x2", + (0, "e3m2x2", 0, 1, ""), + "cvt.rn.satfinite.bf16x2.e3m2x2", + ), + ( + "rs_satfinite_f6x4_f32", + (0, 0, 0, 0, 0, "e2m3x4", 1), + "cvt.rs.relu.satfinite.e2m3x4.f32", + ), + ("frnd3_ue8m0x2_f32", (0, 0, "rz", 0), "cvt.rz.ue8m0x2.f32"), + ( + "frnd3_ue8m0x2_bf16x2", + (0, "rp", 1), + "cvt.rp.satfinite.ue8m0x2.bf16x2", + ), + ("rn_bf16x2_ue8m0x2", (0,), "cvt.rn.bf16x2.ue8m0x2"), + ( + "rn_satfinite_s2f6x2_f32", + (0, 0, 0, 1, "n2::ue8m0"), + "cvt.rn.satfinite.relu.scaled::n2::ue8m0.s2f6x2.f32", + ), + ( + "rn_satfinite_s2f6x2_bf16x2", + (0, 0, ""), + "cvt.rn.satfinite.s2f6x2.bf16x2", + ), + ( + "rn_bf16x2_s2f6x2", + (0, 0, 1, 1, "n2::ue8m0"), + "cvt.rn.satfinite.relu.scaled::n2::ue8m0.bf16x2.s2f6x2", + ), +] + + +@pytest.mark.parametrize(("form_name", "args", "instruction"), _FORM_CASES) +def test_ptx_cvt_form_registration(form_name, args, instruction): + helper_name, _signature, _c_type, _tvm_type, body = _cvt_form_parts(form_name, *args) + + assert instruction in body + assert helper_name.startswith("tvm_builtin_ptx_cvt_") + assert f"tirx._ptx_cvt_{form_name}" in CODEGEN_REGISTRY + + +def test_ptx_cvt_test_matrix_covers_every_registered_form(): + assert {case[0] for case in _FORM_CASES} == set(_CVT_FORMS) + + +def test_ptx_cvt_e2m1x2_uses_b8_staging(): + _name, signature, c_type, tvm_type, body = _cvt_form_parts( + "rn_satfinite_f4x2_f32", 0, 0, "e2m1x2", 0 + ) + assert signature == "(float a, float b)" + assert (c_type, tvm_type) == ("uint8_t", "uint8") + assert ".reg .b8 raw_result;" in body + assert "cvt.u16.u8 %0, raw_result;" in body + + _name, signature, _c_type, _tvm_type, body = _cvt_form_parts("rn_f16x2_f4x2", 0, "e2m1x2", 0) + assert signature == "(uint16_t a)" + assert ".reg .b8 raw_a;" in body + assert "cvt.u8.u16 raw_a, %1;" in body diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py index 784996db01f4..aadef8a9a784 100644 --- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py +++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Unit tests for generic PTX ``T.ptx.ld`` / ``T.ptx.st`` vector copy ops.""" +"""Unit tests for PTX ``T.ptx.ld`` / ``T.ptx.ld_global_nc`` / ``T.ptx.st`` ops.""" import numpy as np import pytest @@ -40,6 +40,18 @@ 1: {"nelems": 1, "smem_dtype": "uint8", "tmp_dtype": "uint32", "fill_u8": 255}, } +_PTX_ST_SRC_VECTOR_CASES = [ + ("uint8", "u8", "v4", "uchar4 src_", "unsigned int r3 = static_cast(src_.w)"), + ("int8", "s8", "v4", "char4 src_", "int r3 = static_cast(src_.w)"), + ("uint16", "u16", "v4", "ushort4 src_", "unsigned short r3 = src_.w"), + ("int16", "s16", "v4", "short4 src_", "short r3 = src_.w"), + ("uint32", "b32", "v4", "uint4 src_", "unsigned int r3 = src_.w"), + ("uint32", "u32", "v4", "uint4 src_", "unsigned int r3 = src_.w"), + ("int32", "s32", "v4", "int4 src_", "int r3 = src_.w"), + ("float32", "f32", "v4", "float4 src_", "float r3 = src_.w"), + ("float64", "f64", "v2", "double2 src_", "double r1 = src_.y"), +] + def _build_and_run(func, *np_args): mod = tvm.compile(tvm.IRModule({"main": func}), target=TARGET, tir_pipeline="tirx") @@ -123,13 +135,34 @@ def func(out_ptr: T.handle): return func +def _ptx_st_src_vector_kernel(dtype: str, ptx_type: str, vec: str): + vec_len = int(vec[1:]) + + @T.prim_func + def func(out_ptr: T.handle) -> None: + out = T.match_buffer(out_ptr, (1,), "uint32") + T.device_entry() + T.cta_id([1]) + tx = T.thread_id([32]) + smem = T.alloc_buffer((vec_len,), dtype, scope="shared") + reg = T.alloc_local((vec_len,), dtype) + if tx == 0: + T.ptx.st( + smem.ptr_to([0]), src=reg.ptr_to([0]), space="shared", vec=vec, ptx_type=ptx_type + ) + out[0] = T.uint32(0) + + return func + + def test_ptx_ld_st_ops_registered(): """PTX ld/st must be registered TIR ops and exposed on the T.ptx namespace.""" - for name in ("tirx.ptx.ld", "tirx.ptx.st"): + for name in ("tirx.ptx.ld", "tirx.ptx.ld_global_nc", "tirx.ptx.st"): Op.get(name) # raises if unregistered for attr in ( "ld", + "ld_global_nc", "st", "ld_acquire", "st_release", @@ -175,6 +208,276 @@ def copy_kernel(d_ptr: T.handle) -> None: assert "tvm_builtin_ptx_st" in src +def test_ptx_ld_st_raw_shared_address_codegen(): + @T.prim_func + def main(out: T.Buffer((2,), "uint64")): + T.device_entry() + tx = T.thread_id([32]) + smem = T.alloc_buffer((2,), "uint64", scope="shared") + values = T.alloc_local((4,), "uint32") + if tx == 0: + raw_addr: T.uint32 = T.cuda.cvta_generic_to_shared(smem.data) + out[0] = T.ptx.ld(raw_addr, "uint64", "u64", space="shared") + out[1] = T.ptx.ld(smem.data, "uint64", "u64", space="shared") + T.ptx.st( + raw_addr, + src=values.ptr_to([0]), + weak=True, + space="shared::cta", + ptx_type="b128", + ) + + with TARGET: + mod = tvm.compile(tvm.IRModule({"main": main}), target=TARGET, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source("cuda") + assert "unsigned int address" in src + assert "ld.shared.u64 %0, [%1];" in src + assert src.count("__cvta_generic_to_shared") == 2 + assert 'asm volatile("st.weak.shared::cta.b128 [%0], %1;"' in src + assert '"q"(value));' in src + b128_helper = src[src.index("tvm_builtin_ptx_st_weak_shared_cta_b128_from_src") :] + b128_helper = b128_helper[: b128_helper.index("\n}")] + assert '"memory"' not in b128_helper + + +def test_ptx_ld_st_reject_invalid_integer_addresses_and_b128_forms(): + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="uint32 address"): + + @T.prim_func + def bad_global_raw(out: T.Buffer((1,), "uint32")): + T.device_entry() + out[0] = T.ptx.ld(T.uint32(0), "uint32", "u32", space="global") + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="integer address"): + + @T.prim_func + def bad_shared_u64(out: T.Buffer((1,), "uint64")): + T.device_entry() + out[0] = T.ptx.ld(T.uint64(0), "uint64", "u64", space="shared") + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="b128 requires"): + + @T.prim_func + def bad_b128(out: T.Buffer((1,), "uint32")): + T.device_entry() + T.ptx.st( + T.uint32(0), + out[0], + weak=True, + space="shared::cta", + ptx_type="b128", + ) + + +def test_ptx_ld_global_nc_v8_codegen(): + """FlashMLA index loads need ``ld.global.nc`` with a 256B prefetch.""" + + @T.prim_func + def copy_kernel(src_ptr: T.handle, out_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (8,), "int32") + out = T.match_buffer(out_ptr, (8,), "int32") + T.device_entry() + tx = T.thread_id([32]) + tmp = T.alloc_local((8,), "int32") + if tx == 0: + T.ptx.ld_global_nc( + src.data, + "int32", + "s32", + dst=tmp.ptr_to([0]), + vec="v8", + l1_evict="L1::no_allocate", + l2_evict="L2::evict_first", + prefetch_size="L2::256B", + ) + for i in T.unroll(8): + out[i] = tmp[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source("cuda") + assert "ld.global.nc.L1::no_allocate.L2::evict_first.L2::256B.v8.s32" in src + assert "tvm_builtin_ptx_ld_global_nc_v8_s32_to_dst" in src + assert "tvm_builtin_ptx_ld_plain_global_nc" not in src + assert "dst[7] = r7" in src + + +def test_ptx_ld_global_nc_v4_u64_256b_codegen(): + """FlashMLA 32-byte index loads may use four 64-bit PTX outputs.""" + + @T.prim_func + def copy_kernel(src_ptr: T.handle, out_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (4,), "uint64") + out = T.match_buffer(out_ptr, (4,), "uint64") + T.device_entry() + tx = T.thread_id([32]) + tmp = T.alloc_local((4,), "uint64") + if tx == 0: + T.ptx.ld_global_nc( + src.data, + "uint64", + "u64", + dst=tmp.ptr_to([0]), + vec="v4", + l1_evict="L1::no_allocate", + l2_evict="L2::evict_normal", + prefetch_size="L2::256B", + ) + for i in T.unroll(4): + out[i] = tmp[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source("cuda") + assert "ld.global.nc.L1::no_allocate.L2::evict_normal.L2::256B.v4.u64" in src + assert "tvm_builtin_ptx_ld_global_nc_v4_u64_to_dst" in src + assert "tvm_builtin_ptx_ld_plain_global_nc" not in src + assert "unsigned long long r3" in src + + +def test_ptx_ld_global_nc_scalar_cop_cache_hint_codegen(): + """The ``ld.global{.cop}.nc`` syntax form supports scalar cache-policy operands.""" + + @T.prim_func + def copy_kernel(src_ptr: T.handle, out_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1,), "int32") + out = T.match_buffer(out_ptr, (1,), "int32") + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + out[0] = T.ptx.ld_global_nc( + src.data, + "int32", + "s32", + cop="ca", + cache_hint="evict_last", + prefetch_size="L2::128B", + ) + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source("cuda") + assert "ld.global.ca.nc.L2::cache_hint.L2::128B.s32" in src + assert "tvm_builtin_ptx_ld_global_nc_ca_s32_int32_cache_hint_L2_128B" in src + assert "unsigned long long cache_policy" in src + + +def test_ptx_ld_rejects_legacy_global_nc_cop(): + """``ld.global.nc`` must use the dedicated ``T.ptx.ld_global_nc`` wrapper.""" + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="ld_global_nc"): + + @T.prim_func + def copy_kernel(src_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1,), "int32") + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + value: T.int32 = T.ptx.ld(src.data, "int32", "s32", cop="nc") + T.evaluate(value) + + +def test_ptx_ld_global_nc_rejects_cop_with_eviction(): + """The ``.cop`` syntax form cannot mix with L1/L2 eviction modifiers.""" + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="l1_evict"): + + @T.prim_func + def copy_kernel(src_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1,), "int32") + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + value: T.int32 = T.ptx.ld_global_nc( + src.data, + "int32", + "s32", + cop="ca", + l1_evict="L1::evict_first", + ) + T.evaluate(value) + + +def test_ptx_ld_global_nc_rejects_non_nc_cop(): + """``ld.global.nc`` has a narrower ``.cop`` set than generic ``ld.global``.""" + + with pytest.raises((ValueError, tvm.error.DiagnosticError), match="cop"): + + @T.prim_func + def copy_kernel(src_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1,), "int32") + T.device_entry() + tx = T.thread_id([32]) + if tx == 0: + value: T.int32 = T.ptx.ld_global_nc(src.data, "int32", "s32", cop="lu") + T.evaluate(value) + + +def test_ptx_ld_vector_scatter_dst_codegen(): + """Vector loads may write independent destination pointers.""" + + @T.prim_func + def copy_kernel(src_ptr: T.handle, out_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (4,), "int32") + out = T.match_buffer(out_ptr, (4,), "int32") + T.device_entry() + tx = T.thread_id([32]) + tmp0 = T.alloc_local((1,), "int32") + tmp1 = T.alloc_local((1,), "int32") + tmp2 = T.alloc_local((1,), "int32") + tmp3 = T.alloc_local((1,), "int32") + if tx == 0: + T.ptx.ld_global_nc( + src.data, + "int32", + "s32", + dst=( + tmp0.ptr_to([0]), + tmp1.ptr_to([0]), + tmp2.ptr_to([0]), + tmp3.ptr_to([0]), + ), + vec="v4", + ) + out[0] = tmp0[0] + out[1] = tmp1[0] + out[2] = tmp2[0] + out[3] = tmp3[0] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": copy_kernel}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source("cuda") + assert "ld.global.nc.v4.s32" in src + assert "tvm_builtin_ptx_ld_global_nc_v4_s32_to_dst4" in src + assert "tvm_builtin_ptx_ld_plain_global_nc" not in src + assert "void* dst0, void* dst1, void* dst2, void* dst3, void* src_ptr" in src + assert "*reinterpret_cast(dst3) = r3" in src + + +@pytest.mark.parametrize( + ("dtype", "ptx_type", "vec", "vector_decl", "last_load"), + _PTX_ST_SRC_VECTOR_CASES, +) +def test_ptx_st_from_src_vector_typed_load_codegen(dtype, ptx_type, vec, vector_decl, last_load): + """``T.ptx.st(src=...)`` loads vector source pointers with the PTX element type.""" + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _ptx_st_src_vector_kernel(dtype, ptx_type, vec)}), + target=target, + tir_pipeline="tirx", + ) + src = mod.mod.imports[0].inspect_source("cuda") + assert vector_decl in src + assert last_load in src + assert f"st.shared.{vec}.{ptx_type}" in src + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") @pytest.mark.parametrize( diff --git a/tests/python/tirx/iket/iket_profile_workload.py b/tests/python/tirx/iket/iket_profile_workload.py new file mode 100644 index 000000000000..9e2218ba84a8 --- /dev/null +++ b/tests/python/tirx/iket/iket_profile_workload.py @@ -0,0 +1,111 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Canonical NVIDIA IKET workload. + +Run as a normal Python program; :func:`iket.run` performs the replay:: + + python tests/python/tirx/iket/iket_profile_workload.py \ + --output-dir /tmp/tvm-iket-workload +""" + +import argparse +import json +import os +from functools import partial +from pathlib import Path + +import numpy as np + +import tvm +from tvm.script import tirx as T +from tvm.tirx.cuda import iket + + +@T.prim_func +def canonical_iket_workload(out: T.Buffer((32,), "int32")): + T.device_entry() + profiler = iket.IketProfiler() + tx = T.thread_id([32]) + token = profiler.sentinel_token("token") + profiler.range_end(token) + token = profiler.range_start("token") + profiler.mark("checkpoint") + profiler.range_end(token) + profiler.range_push("stack") + profiler.mark("inside_stack") + profiler.range_pop() + out[tx] = tx + 1 + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", default="/tmp/tvm-iket-workload") + parser.add_argument( + "--postprocess", + choices=("perfetto", "json", "html", "none", "all"), + default="all", + ) + parser.add_argument("--clobber", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--keep", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--timeout", type=float, default=600.0) + parser.add_argument("--max-ts-cnt-per-warp", type=int, default=None) + parser.add_argument( + "--fail-capture", + action="store_true", + help="Fail only in the real capture pass to verify error propagation", + ) + return parser.parse_args() + + +def _injection_tool_name(): + config_path = os.environ.get("SMODEL_INJECTION_CONFIG") + if not config_path: + return None + return json.loads(Path(config_path).read_text(encoding="utf-8")).get("toolName") + + +def _profile_workload(args): + if args.fail_capture and _injection_tool_name() == "iket": + raise RuntimeError("intentional capture-only IKET workload failure") + executable = iket.IketProfiler().compile( + canonical_iket_workload, + target=tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}), + tir_pipeline="tirx", + ) + out = tvm.runtime.empty((32,), "int32", device=tvm.cuda()) + executable["canonical_iket_workload"](out) + tvm.cuda().sync() + np.testing.assert_array_equal(out.numpy(), np.arange(32, dtype=np.int32) + 1) + + +def main(): + args = _parse_args() + result = iket.run( + partial(_profile_workload, args), + output_dir=args.output_dir, + postprocess=args.postprocess, + clobber=args.clobber, + keep=args.keep, + timeout=args.timeout, + max_ts_cnt_per_warp=args.max_ts_cnt_per_warp, + ) + print(f"IKET output directory: {result.output_dir}") + for path in (*result.json_traces, *result.perfetto_traces, *result.html_reports): + print(f"IKET artifact: {path}") + + +if __name__ == "__main__": + main() diff --git a/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py new file mode 100644 index 000000000000..acc5826381f1 --- /dev/null +++ b/tests/python/tirx/iket/oracle/generate_iket_official_oracle.py @@ -0,0 +1,214 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Generate the normalized CUTLASS DSL 4.6.1 IKET oracle manifest. + +This helper intentionally writes NVIDIA-generated binary artifacts only to the +requested output directory. The repository stores the normalized manifest, +not the proprietary artifacts themselves. +""" + +import argparse +import dataclasses +import hashlib +import json +import subprocess +from enum import IntEnum +from importlib import metadata +from pathlib import Path + +import cutlass +import cutlass.cute as cute +import iket +from cutlass.cute.experimental import iket as cute_iket + + +@cute.kernel +def oracle_kernel(): + """Exercise every no-payload operation supported by the official backend.""" + cute_iket.mark("mark") + token = cute_iket.range_start("token") + cute_iket.range_end(token) + cute_iket.range_push("stack") + cute_iket.range_pop() + cute_iket.mark("slash/name") + cute_iket.mark("é") + + +@cute.jit +def oracle_launch(): + oracle_kernel().launch(grid=(1, 1, 1), block=(32, 1, 1)) + + +def _normalize(value): + if dataclasses.is_dataclass(value): + return { + field.name: _normalize(getattr(value, field.name)) + for field in dataclasses.fields(value) + } + if isinstance(value, IntEnum): + return int(value) + if isinstance(value, dict): + return {str(key): _normalize(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_normalize(item) for item in value] + return value + + +def _sha256(value): + if isinstance(value, Path): + value = value.read_bytes() + if isinstance(value, str): + value = value.encode() + return hashlib.sha256(value).hexdigest() + + +def _artifact_bytes(value): + if isinstance(value, bytes): + return value + if isinstance(value, Path): + return value.read_bytes() + if isinstance(value, str): + path = Path(value) + return path.read_bytes() if "\n" not in value and path.is_file() else value.encode() + return bytes(value) + + +def _command_output(command): + return subprocess.run(command, check=True, capture_output=True, text=True).stdout.strip() + + +def _normalized_metadata(context): + events = sorted(_normalize(context.get_all_events()), key=lambda item: item["event_id"]) + ranges = sorted(_normalize(context.get_all_ranges()), key=lambda item: item["range_id"]) + kernels = [] + for kernel in _normalize(context.get_kernels_instrument_info()): + kernels.append( + { + "kernel_name": kernel["kernel_name"], + "event_sequence": [point["event_attr"] for point in kernel["instrument_points"]], + } + ) + return { + "info": _normalize(context.get_meta_info()), + "ranges": ranges, + "events": events, + "kernels": kernels, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--wheel", type=Path, action="append", default=[]) + for artifact_name in ( + "unpatched-kernel", + "patched-kernel", + "unpatched-disassembly", + "patched-disassembly", + ): + parser.add_argument(f"--{artifact_name}", type=Path) + args = parser.parse_args() + patch_artifacts = { + name: getattr(args, name) + for name in ( + "unpatched_kernel", + "patched_kernel", + "unpatched_disassembly", + "patched_disassembly", + ) + } + if any(patch_artifacts.values()) and not all(patch_artifacts.values()): + parser.error("all four patched/unpatched artifact paths must be supplied together") + args.output_dir.mkdir(parents=True, exist_ok=True) + + cutlass.cuda.initialize_cuda_context() + compiled = cute.compile( + oracle_launch, + options=(f"iket --dump-dir={args.output_dir} --keep-ptx --keep-cubin --keep-sass"), + ) + artifacts = { + name.lower(): _artifact_bytes(getattr(compiled.artifacts, name)) + for name in ("PTX", "CUBIN", "SASS") + } + cubin = artifacts["cubin"] + context = iket.Context(cubin) + if not context.is_instrumented(): + raise RuntimeError("CUTLASS DSL oracle CUBIN is not IKET-instrumented") + + normalized_metadata = _normalized_metadata(context) + manifest = { + "schema_version": 2, + "profile": { + "cutlass_dsl": metadata.version("nvidia-cutlass-dsl"), + "instrument_method": "NativeDump", + "driver": _command_output( + [ + "nvidia-smi", + "--query-gpu=name,driver_version", + "--format=csv,noheader", + "--id=0", + ] + ), + "nvdisasm": _command_output(["nvdisasm", "--version"]), + "compiler_flags": [ + "iket", + "--dump-dir=", + "--keep-ptx", + "--keep-cubin", + "--keep-sass", + ], + }, + "wheels": {wheel.name: _sha256(wheel) for wheel in sorted(args.wheel)}, + "artifact_sha256": {name: _sha256(value) for name, value in artifacts.items()}, + "metadata": normalized_metadata, + "metadata_sha256": _sha256( + json.dumps(normalized_metadata, sort_keys=True, separators=(",", ":")) + ), + "native_dump_abi": { + "meta_info_bytes": 48, + "event_attributes_bytes": 60, + "range_attributes_bytes": 72, + "sentinel_event_id": 0, + "range_pop_event_id": 31, + "placeholder": [ + "READ_CLUSTER_CTARANK", + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "COMPUTE_SHARED_PLACEHOLDER_ADDRESS", + "STORE_SHARED_WEAK_32", + "PMEVENT", + ], + "patched_hot_path": [ + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "STORE_GLOBAL_32", + "ADD_WRITE_PTR_64_4", + ], + }, + } + if all(patch_artifacts.values()): + manifest["patch_artifact_sha256"] = { + name: _sha256(path) for name, path in patch_artifacts.items() + } + manifest_path = args.output_dir / "iket_official_oracle.json" + manifest_path.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(manifest_path) + + +if __name__ == "__main__": + main() diff --git a/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json new file mode 100644 index 000000000000..d84d38115d73 --- /dev/null +++ b/tests/python/tirx/iket/oracle/iket_official_cutlass_4_6_1_oracle.json @@ -0,0 +1,207 @@ +{ + "artifact_sha256": { + "cubin": "ddd0cbfc511914d453cfc9ba770bddc3939a6ce77f13488eae1aa31bd867a1f5", + "ptx": "9a679b1ceefe35bd4f783d13bcf7b8a4c1490f94b606323cd9ffb7b1c5aa7554", + "sass": "5d28eac4ad110174ab8fb99d1437ebf0e99b4de4df6c1cee2341e2ef210f5af7" + }, + "metadata": { + "events": [ + { + "event_id": 0, + "event_name": "reserved_event_id_0", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 1, + "event_name": "mark", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 3, + "event_name": "stack", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 0, + "range_id": 1649501183 + }, + { + "event_id": 4, + "event_name": "slash/name", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 5, + "event_name": "\u00e9", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + } + ], + "info": { + "legacy_iket_cubin": 0, + "magic_number": 3133075869, + "max_event_id": 31, + "max_event_name_size": 32, + "supported_features": 3, + "version_major": 0, + "version_minor": 5 + }, + "kernels": [ + { + "event_sequence": [ + { + "event_id": 1, + "event_name": "mark", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 2, + "event_name": "token", + "event_pos": 4, + "instrument_method": 3, + "payload_type": 0, + "range_id": 2491017778 + }, + { + "event_id": 3, + "event_name": "stack", + "event_pos": 1, + "instrument_method": 3, + "payload_type": 0, + "range_id": 1649501183 + }, + { + "event_id": 31, + "event_name": "pop_range", + "event_pos": 2, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 4, + "event_name": "slash/name", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + }, + { + "event_id": 5, + "event_name": "\u00e9", + "event_pos": 0, + "instrument_method": 3, + "payload_type": 0, + "range_id": 0 + } + ], + "kernel_name": "kernel_cutlass_oracle_kernel_0" + } + ], + "ranges": [ + { + "color": -1, + "evt_pair_mode": 0, + "range_id": 1649501183, + "range_name": "stack", + "range_scope": 0, + "range_type": 2 + }, + { + "color": -1, + "evt_pair_mode": 1, + "range_id": 2491017778, + "range_name": "token", + "range_scope": 0, + "range_type": 1 + } + ] + }, + "metadata_sha256": "85c13d2584ca6a37b6bfd17dc17d70ded6696f63fae728cf9940f7ecf2c90dc2", + "native_dump_abi": { + "event_attributes_bytes": 60, + "meta_info_bytes": 48, + "patched_hot_path": [ + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "STORE_GLOBAL_32", + "ADD_WRITE_PTR_64_4" + ], + "placeholder": [ + "READ_CLUSTER_CTARANK", + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "COMPUTE_SHARED_PLACEHOLDER_ADDRESS", + "STORE_SHARED_WEAK_32", + "PMEVENT" + ], + "range_attributes_bytes": 72, + "range_pop_event_id": 31, + "sentinel_event_id": 0 + }, + "patch_artifact_sha256": { + "patched_disassembly": "7a7689265c4d59fff9799f5a283e22a3ae9eb3ed91a08743106f6213931dd1c3", + "patched_kernel": "b6c57768c8bf1a1d20ad52d8987c571318340f6d2b006cb62ae58c1d934a44f4", + "unpatched_disassembly": "5d28eac4ad110174ab8fb99d1437ebf0e99b4de4df6c1cee2341e2ef210f5af7", + "unpatched_kernel": "823305d494ef1cc26146443ac9d5108da2c13684b5a07aeddc72ac04909550af" + }, + "profile": { + "compiler_flags": [ + "iket", + "--dump-dir=", + "--keep-ptx", + "--keep-cubin", + "--keep-sass" + ], + "cutlass_dsl": "4.6.1", + "driver": "NVIDIA B200, 595.58.03", + "instrument_method": "NativeDump", + "nvdisasm": "nvdisasm: NVIDIA (R) CUDA disassembler\nCopyright (c) 2005-2026 NVIDIA Corporation\nBuilt on Tue_Jun_09_02:42:28_PM_PDT_2026\nCuda compilation tools, release 13.3, V13.3.73\nBuild cuda_13.3.r13.3/compiler.38244171_0" + }, + "schema_version": 2, + "wheels": { + "nvidia_cuda_nvdisasm-13.3.73-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl": "fa17084b07c0dca68a42892f771b4b1b40fbe9b91660209623e61cea611cae8c", + "nvidia_cuda_nvrtc-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl": "82530788b8c6164a54d3fd9ae8bcca8893d397c4aeb998861982a03bbe41e204", + "nvidia_cutlass_dsl-4.6.1-py3-none-any.whl": "93135a9d48e1bedf584828e0a021f174ac31591ef21f6ea53c206169ccbfab26", + "nvidia_cutlass_dsl_libs_base-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl": "dcbbf471839801501030f1097ce13dbf5082ad34f26d6b503459c6fed078e4e9", + "nvidia_cutlass_dsl_libs_core-4.6.1-py3-none-any.whl": "f1d895ee24b1ba711b2b9d4a43c62fa1f6fe3a50634e25eeffe88fa17f6c5e47", + "nvidia_cutlass_dsl_libs_cu13-4.6.1-cp311-cp311-manylinux_2_28_x86_64.whl": "a8483778ad75ae50efd5b8981adfe8d7bc0c9cbcf5fd16a02b3bd9a062b98e2d" + } +} diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tests/python/tirx/iket/test_iket_orchestration.py new file mode 100644 index 000000000000..5e66c6b91479 --- /dev/null +++ b/tests/python/tirx/iket/test_iket_orchestration.py @@ -0,0 +1,502 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for NVIDIA IKET command orchestration and replay semantics.""" + +from __future__ import annotations + +import importlib.machinery +import os +import sys +import textwrap +import time +import types +from pathlib import Path + +import pytest + +from tvm.tirx.cuda import iket + + +def _write_artifacts(output_dir: Path, postprocess: str, *, count: int = 1) -> None: + for index in range(count): + if postprocess in ("json", "all"): + (output_dir / f"iket_pid_{index}.trace.json").write_text( + '{"launches": [{"index": %d}]}' % index, encoding="utf-8" + ) + if postprocess in ("perfetto", "html", "all"): + (output_dir / f"iket_pid_{index}.pftrace").write_bytes(b"perfetto") + if postprocess in ("html", "all"): + (output_dir / f"iket_pid_{index}.html").write_text("", encoding="utf-8") + + +def _fake_profile_process(captured, *, artifact_count: int = 1): + def run_process(argv, *, cwd, env, timeout): + captured.update(argv=list(argv), cwd=cwd, env=env, timeout=timeout) + output_dir = Path(argv[argv.index("--output-dir") + 1]) + postprocess = argv[argv.index("--postprocess") + 1] + _write_artifacts(output_dir, postprocess, count=artifact_count) + + return run_process + + +def test_public_namespace_and_old_bench_path_is_absent(): + import tvm.tirx.bench as bench + + assert iket.__all__ == [ + "IketProfileError", + "IketProfileResult", + "IketProfiler", + "profile", + "run", + ] + assert not hasattr(bench, "IketProfiler") + + +@pytest.mark.parametrize("injection_variable", ["CUDA_INJECTION64_PATH", "SMODEL_INJECTION_CONFIG"]) +def test_profile_rejects_nested_injection_before_validation( + injection_variable, tmp_path, monkeypatch +): + validation_started = False + + def validate(_name): + nonlocal validation_started + validation_started = True + + monkeypatch.setenv(injection_variable, "injected") + monkeypatch.setattr(iket, "_validate_official_installation", validate) + target = tmp_path / "result" + + with pytest.raises(iket.IketProfileError, match="nested profiling"): + iket.profile(("python", "workload.py"), output_dir=target) + + assert not validation_started + assert not target.exists() + + +@pytest.mark.parametrize(("keep", "keep_arg"), [(True, "--keep"), (False, "--no-keep")]) +def test_build_run_iket_argv_is_exact(keep, keep_arg, tmp_path): + output_dir = tmp_path / "staging" + assert iket._build_run_iket_argv( # pylint: disable=protected-access + "/opt/bin/run-iket", + ("python", "workload.py", "--size", "4"), + output_dir=output_dir, + postprocess="html", + max_ts_cnt_per_warp=23, + keep=keep, + ) == [ + "/opt/bin/run-iket", + "--output-dir", + str(output_dir), + "--clobber", + "profile", + "--postprocess", + "html", + keep_arg, + "--max-ts-cnt-per-warp", + "23", + "--", + "python", + "workload.py", + "--size", + "4", + ] + + +def test_profile_forwards_cwd_environment_timeout_and_publishes(tmp_path, monkeypatch): + target = tmp_path / "published" + cwd = tmp_path / "work" + cwd.mkdir() + captured = {} + monkeypatch.setattr( + iket, "_validate_official_installation", lambda profile_name: "/opt/bin/run-iket" + ) + monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured)) + monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "inherited-profile") + + with pytest.warns(RuntimeWarning, match="takes precedence"): + result = iket.profile( + (sys.executable, "workload.py"), + output_dir=target, + profile_name="cutlass-4.6.1", + postprocess="all", + clobber=False, + cwd=cwd, + env={"IKET_TEST_ENV": "present"}, + max_ts_cnt_per_warp=17, + keep=True, + timeout=12.5, + ) + + staging = Path(captured["argv"][2]) + assert staging.parent == target.parent + assert staging != target + assert captured["cwd"] == cwd + assert captured["timeout"] == 12.5 + assert captured["env"]["IKET_TEST_ENV"] == "present" + assert captured["env"]["TVM_IKET_OFFICIAL_PROFILE"] == "cutlass-4.6.1" + assert os.environ["TVM_IKET_OFFICIAL_PROFILE"] == "inherited-profile" + assert result.output_dir == target + assert result.command == (sys.executable, "workload.py") + assert result.trace == {"launches": [{"index": 0}]} + assert result.launches == [{"index": 0}] + assert result.trace_json.parent == target + assert result.perfetto.parent == target + assert result.html.parent == target + + +@pytest.mark.parametrize("postprocess", ["perfetto", "json", "html", "none", "all"]) +def test_postprocess_artifact_contract(postprocess, tmp_path, monkeypatch): + target = tmp_path / postprocess + captured = {} + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: "run-iket") + monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured)) + + result = iket.profile( + ("python", "workload.py"), output_dir=target, postprocess=postprocess, timeout=None + ) + assert result.postprocess == postprocess + assert captured["timeout"] is None + + requested = iket._requested_formats(postprocess) # pylint: disable=protected-access + properties = { + "json": lambda: result.trace_json, + "perfetto": lambda: result.perfetto, + "html": lambda: result.html, + } + for format_name, getter in properties.items(): + if format_name in requested: + assert getter().is_file() + else: + with pytest.raises(iket.IketProfileError, match="was not requested"): + getter() + + +def test_multi_process_singular_properties_point_to_plural(tmp_path, monkeypatch): + target = tmp_path / "multi" + captured = {} + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: "run-iket") + monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured, artifact_count=2)) + result = iket.profile(("python", "workload.py"), output_dir=target) + + assert len(result.json_traces) == 2 + assert len(result.perfetto_traces) == 2 + assert len(result.html_reports) == 2 + with pytest.raises(iket.IketProfileError, match="json_traces"): + _ = result.trace_json + with pytest.raises(iket.IketProfileError, match="perfetto_traces"): + _ = result.perfetto + with pytest.raises(iket.IketProfileError, match="html_reports"): + _ = result.html + + +def test_missing_artifact_rolls_back_existing_output(tmp_path, monkeypatch): + target = tmp_path / "result" + target.mkdir() + sentinel = target / "old.txt" + sentinel.write_text("old", encoding="utf-8") + + def incomplete(argv, **_kwargs): + staging = Path(argv[argv.index("--output-dir") + 1]) + (staging / "iket_pid_1.trace.json").write_text("{}", encoding="utf-8") + + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: "run-iket") + monkeypatch.setattr(iket, "_run_process", incomplete) + with pytest.raises(iket.IketProfileError, match="Perfetto trace"): + iket.profile(("python", "workload.py"), output_dir=target, clobber=True) + + assert sentinel.read_text(encoding="utf-8") == "old" + assert not list(tmp_path.glob(".result.staging-*")) + + +def test_successful_clobber_replaces_only_after_profile_success(tmp_path, monkeypatch): + target = tmp_path / "result" + target.mkdir() + (target / "old.txt").write_text("old", encoding="utf-8") + captured = {} + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: "run-iket") + monkeypatch.setattr(iket, "_run_process", _fake_profile_process(captured)) + + result = iket.profile(("python", "workload.py"), output_dir=target, clobber=True) + assert not (target / "old.txt").exists() + assert result.trace_json.is_file() + assert not list(tmp_path.glob(".result.backup-*")) + + +def test_existing_output_fails_before_installation_validation(tmp_path, monkeypatch): + target = tmp_path / "result" + target.mkdir() + called = False + + def validate(_name): + nonlocal called + called = True + + monkeypatch.setattr(iket, "_validate_official_installation", validate) + with pytest.raises(FileExistsError, match="already exists"): + iket.profile(("python", "workload.py"), output_dir=target) + assert not called + + +@pytest.mark.parametrize("timeout", [0, -1, float("inf"), float("nan")]) +def test_invalid_timeout_is_rejected_before_validation(timeout, tmp_path, monkeypatch): + called = False + + def validate(_name): + nonlocal called + called = True + + monkeypatch.setattr(iket, "_validate_official_installation", validate) + with pytest.raises(ValueError, match="timeout"): + iket.profile(("python", "workload.py"), output_dir=tmp_path / "result", timeout=timeout) + assert not called + + +def test_command_must_be_an_argument_sequence(tmp_path): + with pytest.raises(TypeError, match="shell command string"): + iket.profile("python workload.py", output_dir=tmp_path / "result") + + +def _make_executable(path: Path, source: str) -> Path: + path.write_text("#!/usr/bin/env python3\n" + textwrap.dedent(source), encoding="utf-8") + path.chmod(0o755) + return path + + +def test_nonzero_exit_keeps_last_100_output_lines(tmp_path, monkeypatch): + executable = _make_executable( + tmp_path / "run-iket", + """ + import sys + for index in range(105): + print(f"line-{index}", flush=True) + sys.exit(7) + """, + ) + target = tmp_path / "result" + target.mkdir() + sentinel = target / "old.txt" + sentinel.write_text("old", encoding="utf-8") + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: str(executable)) + + with pytest.raises(iket.IketProfileError) as error_info: + iket.profile( + ("python", "workload.py"), + output_dir=target, + postprocess="none", + clobber=True, + keep=True, + ) + error = error_info.value + assert error.returncode == 7 + assert len(error.output_tail.splitlines()) == 100 + assert error.output_tail.splitlines()[0] == "line-5" + assert error.output_tail.splitlines()[-1] == "line-104" + assert sentinel.read_text(encoding="utf-8") == "old" + assert not list(tmp_path.glob(".result.staging-*")) + + +def _pid_is_running(pid: int) -> bool: + stat_path = Path(f"/proc/{pid}/stat") + if not stat_path.exists(): + return False + try: + state = stat_path.read_text(encoding="utf-8").split()[2] + except (FileNotFoundError, IndexError): + return False + return state != "Z" + + +def test_timeout_kills_run_iket_workload_and_grandchild(tmp_path, monkeypatch): + parent_pid = tmp_path / "parent.pid" + child_pid = tmp_path / "child.pid" + grandchild_pid = tmp_path / "grandchild.pid" + executable = _make_executable( + tmp_path / "run-iket", + """ + import os + import signal + import subprocess + import sys + import time + from pathlib import Path + + Path(os.environ["PARENT_PID_FILE"]).write_text(str(os.getpid())) + grandchild_source = """ + + repr( + textwrap.dedent( + """ + import os + import signal + import time + from pathlib import Path + Path(os.environ["GRANDCHILD_PID_FILE"]).write_text(str(os.getpid())) + signal.signal(signal.SIGTERM, signal.SIG_IGN) + while True: + time.sleep(1) + """ + ) + ) + + "\n" + + """ + child_source = "import os, signal, subprocess, sys, time; " \\ + "from pathlib import Path; " \\ + "Path(os.environ['CHILD_PID_FILE']).write_text(str(os.getpid())); " \\ + f"subprocess.Popen([sys.executable, '-c', {grandchild_source!r}]); " \\ + "signal.signal(signal.SIGTERM, signal.SIG_IGN); " \\ + "time.sleep(3600)" + child = subprocess.Popen([sys.executable, "-c", child_source]) + signal.signal(signal.SIGTERM, signal.SIG_IGN) + deadline = time.monotonic() + 5 + while not Path(os.environ["GRANDCHILD_PID_FILE"]).exists(): + if time.monotonic() > deadline: + raise RuntimeError("grandchild did not start") + time.sleep(0.01) + print("fake run-iket ready", flush=True) + time.sleep(3600) + """, + ) + monkeypatch.setattr(iket, "_validate_official_installation", lambda _name: str(executable)) + monkeypatch.setattr(iket, "_TERMINATION_GRACE_SECONDS", 0.2) + + with pytest.raises(iket.IketProfileError) as error_info: + iket.profile( + ("python", "workload.py"), + output_dir=tmp_path / "result", + postprocess="none", + env={ + "PARENT_PID_FILE": str(parent_pid), + "CHILD_PID_FILE": str(child_pid), + "GRANDCHILD_PID_FILE": str(grandchild_pid), + }, + timeout=1.0, + keep=True, + ) + error = error_info.value + assert error.returncode is None + assert error.timeout == 1.0 + assert "fake run-iket ready" in error.output_tail + + pids = [ + int(path.read_text(encoding="utf-8")) for path in (parent_pid, child_pid, grandchild_pid) + ] + deadline = time.monotonic() + 2 + while any(_pid_is_running(pid) for pid in pids) and time.monotonic() < deadline: + time.sleep(0.02) + assert not any(_pid_is_running(pid) for pid in pids) + assert not (tmp_path / "result").exists() + assert not list(tmp_path.glob(".result.staging-*")) + + +def test_injected_run_preserves_main_exit_semantics(tmp_path, monkeypatch): + monkeypatch.setenv("CUDA_INJECTION64_PATH", "injected") + monkeypatch.setattr(iket, "_validate_official_environment", lambda: None) + + with pytest.raises(SystemExit) as normal_exit: + iket.run(lambda: None, output_dir=tmp_path / "normal") + assert normal_exit.value.code == 0 + + def fail(): + raise RuntimeError("capture failed") + + with pytest.raises(RuntimeError, match="capture failed"): + iket.run(fail, output_dir=tmp_path / "failed") + + def explicit_exit(): + raise SystemExit(9) + + with pytest.raises(SystemExit) as nonzero_exit: + iket.run(explicit_exit, output_dir=tmp_path / "nonzero") + assert nonzero_exit.value.code == 9 + + def interrupt(): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + iket.run(interrupt, output_dir=tmp_path / "interrupt") + + +def test_imported_callable_is_rejected_before_profile_starts(tmp_path, monkeypatch): + started = False + + def should_not_start(*_args, **_kwargs): + nonlocal started + started = True + + monkeypatch.delenv("CUDA_INJECTION64_PATH", raising=False) + monkeypatch.delenv("SMODEL_INJECTION_CONFIG", raising=False) + monkeypatch.setattr(iket, "profile", should_not_start) + with pytest.raises(ValueError, match=r"iket\.profile\(command\)"): + iket.run(test_imported_callable_is_rejected_before_profile_starts, output_dir=tmp_path) + assert not started + + +def _artificial_main(source_path: Path): + module = types.ModuleType("__main__") + module.__file__ = str(source_path) + exec(compile("def entry():\n return None\n", str(source_path), "exec"), module.__dict__) + return module + + +def test_replay_guard_accepts_matching_script_file(tmp_path, monkeypatch): + source_path = tmp_path / "workload.py" + source_path.write_text("def entry():\n return None\n", encoding="utf-8") + main_module = _artificial_main(source_path) + monkeypatch.setitem(sys.modules, "__main__", main_module) + monkeypatch.setattr(sys, "orig_argv", [sys.executable, str(source_path), "--size", "4"]) + + assert iket._replay_command(main_module.entry) == ( # pylint: disable=protected-access + sys.executable, + str(source_path), + "--size", + "4", + ) + + +def test_replay_guard_accepts_matching_python_module(tmp_path, monkeypatch): + source_path = tmp_path / "workload.py" + source_path.write_text("def entry():\n return None\n", encoding="utf-8") + main_module = _artificial_main(source_path) + main_module.__spec__ = importlib.machinery.ModuleSpec( + "package.workload", loader=None, origin=str(source_path) + ) + monkeypatch.setitem(sys.modules, "__main__", main_module) + monkeypatch.setattr( + sys, "orig_argv", [sys.executable, "-X", "dev", "-m", "package.workload", "--size", "4"] + ) + + assert iket._replay_command(main_module.entry) == ( # pylint: disable=protected-access + sys.executable, + "-X", + "dev", + "-m", + "package.workload", + "--size", + "4", + ) + + +def test_replay_guard_rejects_file_mismatch(tmp_path, monkeypatch): + source_path = tmp_path / "workload.py" + other_path = tmp_path / "other.py" + source_path.write_text("def entry():\n return None\n", encoding="utf-8") + other_path.write_text("pass\n", encoding="utf-8") + main_module = _artificial_main(source_path) + monkeypatch.setitem(sys.modules, "__main__", main_module) + monkeypatch.setattr(sys, "orig_argv", [sys.executable, str(other_path)]) + + with pytest.raises(ValueError, match="different files"): + iket._replay_command(main_module.entry) # pylint: disable=protected-access diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py new file mode 100644 index 000000000000..2771b1c8bbbf --- /dev/null +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -0,0 +1,670 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for IKET lowering, metadata, installation locks, and trace contracts.""" + +import hashlib +import importlib +import inspect +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +import tvm +from tvm.backend.cuda import transforms as cuda_transforms +from tvm.script import tirx as T +from tvm.tirx.cuda.iket import IketProfiler + +TARGET = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) +ORACLE_PATH = Path(__file__).parent / "oracle" / "iket_official_cutlass_4_6_1_oracle.json" + + +@T.prim_func +def serial_a(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("a") + out[tx] = tx + 1 + + +@T.prim_func +def serial_b(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("b") + out[tx] = tx + 2 + + +@T.prim_func +def plain_entry(out: T.Buffer((32,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + out[tx] = tx + 7 + + +@T.prim_func +def explicit_cuda_shuffle_guard(out: T.Buffer((64,), "int32")): + T.device_entry() + iket = IketProfiler() + bx = T.cta_id([2]) + tx = T.thread_id([64]) + warp = T.cuda.__shfl_sync(T.uint32(0xFFFFFFFF), tx // 32, 0, 32) + if (bx == 0) & (warp == 0): + iket.mark("warp-zero") + out[tx] = tx + + +@T.prim_func +def push_pop_kernel(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.range_push("outer") + iket.range_push("inner") + iket.mark("point") + iket.range_pop() + iket.range_pop() + out[tx] = tx + + +@T.prim_func +def token_loop(n: T.int32, out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + token = iket.sentinel_token("sentinel") + for i in T.serial(n, unroll=False): + iket.range_end(token) + if i % 2 == 0: + token = iket.range_start("even") + else: + token = iket.range_start("odd") + iket.range_end(token) + out[tx] = tx + n + + +@T.prim_func +def overlapping_ranges(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + outer = iket.range_start("overlap") + inner = iket.range_start("overlap") + iket.range_end(inner) + iket.range_end(outer) + out[tx] = tx + + +@T.prim_func +def repeated_range_end(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + token = iket.range_start("twice") + iket.range_end(token) + iket.range_end(token) + out[tx] = tx + + +@T.prim_func +def unbalanced_stack(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.range_pop() + out[tx] = tx + + +@T.prim_func +def payload_kernel(out: T.Buffer((32,), "int32")): + T.device_entry() + tx = T.thread_id([32]) + T.evaluate(tvm.tirx.call_intrin("", "tirx.cuda.iket_mark", "payload", tx)) + out[tx] = tx + + +@T.prim_func +def loop_carried_divergent_token(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + guard = T.alloc_local((1,), "int32") + guard[0] = 0 + token = iket.sentinel_token("loop") + for _i in T.serial(2, unroll=False): + token = iket.sentinel_token("loop") + if guard[0] == 0: + token = iket.range_start("loop") + iket.range_end(token) + guard[0] = tx + out[tx] = guard[0] + + +@T.prim_func +def while_carried_divergent_token(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + guard = T.alloc_local((1,), "int32") + iteration = T.alloc_local((1,), "int32") + guard[0] = 0 + iteration[0] = 0 + token = iket.sentinel_token("while-loop") + while iteration[0] < 2: + token = iket.sentinel_token("while-loop") + if guard[0] == 0: + token = iket.range_start("while-loop") + iket.range_end(token) + guard[0] = tx + iteration[0] = iteration[0] + 1 + out[tx] = guard[0] + + +@T.prim_func +def marks_30(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("e00") + iket.mark("e01") + iket.mark("e02") + iket.mark("e03") + iket.mark("e04") + iket.mark("e05") + iket.mark("e06") + iket.mark("e07") + iket.mark("e08") + iket.mark("e09") + iket.mark("e10") + iket.mark("e11") + iket.mark("e12") + iket.mark("e13") + iket.mark("e14") + iket.mark("e15") + iket.mark("e16") + iket.mark("e17") + iket.mark("e18") + iket.mark("e19") + iket.mark("e20") + iket.mark("e21") + iket.mark("e22") + iket.mark("e23") + iket.mark("e24") + iket.mark("e25") + iket.mark("e26") + iket.mark("e27") + iket.mark("e28") + iket.mark("e29") + out[tx] = tx + + +@T.prim_func +def marks_31(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + iket.mark("e00") + iket.mark("e01") + iket.mark("e02") + iket.mark("e03") + iket.mark("e04") + iket.mark("e05") + iket.mark("e06") + iket.mark("e07") + iket.mark("e08") + iket.mark("e09") + iket.mark("e10") + iket.mark("e11") + iket.mark("e12") + iket.mark("e13") + iket.mark("e14") + iket.mark("e15") + iket.mark("e16") + iket.mark("e17") + iket.mark("e18") + iket.mark("e19") + iket.mark("e20") + iket.mark("e21") + iket.mark("e22") + iket.mark("e23") + iket.mark("e24") + iket.mark("e25") + iket.mark("e26") + iket.mark("e27") + iket.mark("e28") + iket.mark("e29") + iket.mark("e30") + out[tx] = tx + + +def _compile(func, *, target=TARGET): + return IketProfiler().compile(func, target=target, tir_pipeline="tirx") + + +def _cuda_source(executable): + modules = executable.mod._collect_from_import_tree( # pylint: disable=protected-access + lambda module: module.kind == "cuda" + ) + assert len(modules) == 1 + return modules[0].inspect_source("cuda") + + +def _official_global_bytes(source, symbol): + match = re.search(rf"unsigned char {re.escape(symbol)}\[(\d+)\] = \{{([^}}]*)\}};", source) + assert match is not None, symbol + size = int(match.group(1)) + values = [int(value) for value in match.group(2).split(",") if value] + assert len(values) == size + return bytes(values) + + +def test_public_interface_is_official_only(): + signature = inspect.signature(IketProfiler.compile) + assert "backend" not in signature.parameters + profiler = IketProfiler() + assert not hasattr(profiler, "capture") + assert not hasattr(profiler, "export") + assert not hasattr(tvm.tirx.transform, "LowerIket") + assert callable(cuda_transforms.LowerIket) + + script = serial_a.script() + assert 'T.cuda.iket.mark("a")' in script + assert "T.tirx.iket" not in script + assert tvm.script.from_source(script).script() == script + + +@pytest.mark.parametrize( + "name", + ("mark", "range_start", "range_end", "range_push", "range_pop", "sentinel_token"), +) +def test_annotation_ops_are_cuda_owned(name): + op = tvm.ir.Op.get(f"tirx.cuda.iket_{name}") + assert op.get_attr("TIRxOpCategory") == "device_intrin" + assert op.get_attr("TDeviceIntrinsicNamespace") == "cuda" + + +def test_regular_lowering_strips_annotations_and_tokens(): + stripped = cuda_transforms.LowerIket()(tvm.IRModule({"main": token_loop})) + script = stripped.script() + assert "cuda.iket" not in script + assert "sentinel" not in script + assert "token:" not in script + + def make_kernel(with_annotation): + @T.prim_func + def main(out: T.Buffer((32,), "int32")): + T.device_entry() + iket = IketProfiler() + tx = T.thread_id([32]) + if with_annotation and tx % 2 == 0: + iket.mark("strip-me") + out[tx] = tx + 1 + + return main + + sources = [] + for with_annotation in (False, True): + executable = tvm.compile( + tvm.IRModule({"main": make_kernel(with_annotation)}), + target=TARGET, + tir_pipeline="tirx", + ) + sources.append(executable.mod.imports[0].inspect_source("cuda")) + assert sources[0] == sources[1] + assert "iket" not in sources[1].lower() + + +@pytest.mark.parametrize( + ("module_name", "factory_name"), + [ + ("tvm.tirx.compilation_pipeline", "default_tir_pipeline"), + ("tvm.tirx.compilation_pipeline", "tirx_pipeline"), + ], +) +def test_tirx_pipelines_immediately_lower_iket(module_name, factory_name): + factory = getattr(importlib.import_module(module_name), factory_name) + source = inspect.getsource(factory) + assert re.search( + r"tirx\.transform\.SplitHostDevice\(\),\s+cuda_transforms\.LowerIket\(\)", source + ) + + +@pytest.mark.parametrize( + ("module_name", "factory_name"), + [ + ("tvm.s_tir.pipeline", "default_s_tir_pipeline"), + ("tvm.s_tir.backend.adreno.pipeline", "default_tir_pipeline"), + ("tvm.backend.trn.pipeline", "trn_pipeline"), + ], +) +def test_non_tirx_pipelines_do_not_lower_iket(module_name, factory_name): + factory = getattr(importlib.import_module(module_name), factory_name) + assert "LowerIket" not in inspect.getsource(factory) + + +def test_lowering_emits_native_dump_metadata_without_control_abi(): + source = _cuda_source(_compile(push_pop_kernel)) + + meta = _official_global_bytes(source, "__iket_meta_info") + assert [int.from_bytes(meta[offset : offset + 4], "little") for offset in range(0, 36, 4)] == [ + 48, + 0, + 5, + 31, + 32, + 60, + 0xBABEF19D, + 0, + 3, + ] + point = _official_global_bytes(source, "__iket_evt_decl_point_3_attrs") + assert [int.from_bytes(point[offset : offset + 4], "little") for offset in range(0, 28, 4)] == [ + 60, + 3, + 3, + 0, + 0, + 0, + 5, + ] + assert point[28:33] == b"point" + outer = _official_global_bytes(source, "__iket_range_decl_outer_2718680436_attrs") + assert [int.from_bytes(outer[offset : offset + 4], "little") for offset in range(0, 24, 4)] == [ + 72, + 0, + 2718680436, + 0xFFFFFFFF, + 2, + 0, + ] + assert outer[36:41] == b"outer" + + assert "mov.u32 %%t, %%globaltimer_lo" in source + assert "st.weak.shared.u32 [%%r], %%t" in source + assert "pmevent.mask %0" in source + assert source.count("tvm_builtin_iket_official_event((uint)31)") == 2 + for removed_symbol in ( + "__tvm_iket_get_metadata", + "__tvm_iket_set_context", + "__tvm_iket_clear_context", + "tvm_builtin_iket_prologue", + "tvm_builtin_iket_finalize", + ): + assert removed_symbol not in source + + +def test_token_sentinel_and_dynamic_alternation_lowering(): + source = _cuda_source(_compile(token_loop)) + + assert "token_ptr[0] = (uint)0" in source + assert source.count("tvm_builtin_iket_official_event(token_ptr[0])") == 2 + assert "case 1:" in source + assert "case 2:" in source + assert "case 3:" in source + assert "case 31:" in source + for name, event_id in (("even", 1), ("odd", 2), ("sentinel", 3)): + event = _official_global_bytes(source, f"__iket_evt_decl_{name}_{event_id}_attrs") + assert int.from_bytes(event[4:8], "little") == event_id + assert int.from_bytes(event[16:20], "little") == 4 + + +def test_explicit_cuda_shuffle_broadcast_is_warp_uniform(): + source = _cuda_source(_compile(explicit_cuda_shuffle_guard)) + assert "__iket_evt_decl_warp_zero_1_attrs" in source + + +@pytest.mark.parametrize( + ("kernel_func", "message"), + [ + pytest.param(payload_kernel, "does not support payloads", id="payload"), + pytest.param(overlapping_ranges, "strictly alternating", id="overlap"), + pytest.param(repeated_range_end, "strictly alternating", id="repeated-end"), + pytest.param(unbalanced_stack, "balanced range_push/range_pop", id="unbalanced-stack"), + pytest.param( + loop_carried_divergent_token, + "divergent set of lanes", + id="loop-fixed-point-divergence", + ), + pytest.param( + while_carried_divergent_token, + "divergent set of lanes", + id="while-fixed-point-divergence", + ), + ], +) +def test_rejects_unsupported_or_divergent_semantics(kernel_func, message): + with pytest.raises(ValueError, match=message): + _compile(kernel_func) + + +def test_declaration_module_and_architecture_boundaries(): + source = _cuda_source(_compile(marks_30)) + assert len(re.findall(r"__iket_evt_decl_e\d\d_\d+_attrs", source)) == 30 + + with pytest.raises(ValueError, match="at most 30 declarations per kernel"): + _compile(marks_31) + with pytest.raises(ValueError, match="at most 30 distinct declarations"): + IketProfiler().compile( + tvm.IRModule({"marks_30": marks_30, "serial_a": serial_a}), + target=TARGET, + tir_pipeline="tirx", + ) + with pytest.raises(ValueError, match="requires SM90 or newer"): + _compile(serial_a, target=tvm.target.Target({"kind": "cuda", "arch": "sm_80"})) + + +def test_multi_kernel_module_has_no_tvm_control_plane(): + executable = IketProfiler().compile( + tvm.IRModule( + { + "plain_entry": plain_entry, + "serial_a": serial_a, + "serial_b": serial_b, + } + ), + target=TARGET, + tir_pipeline="tirx", + ) + source = _cuda_source(executable) + assert "__iket_evt_decl_a_1_attrs" in source + assert "__iket_evt_decl_b_2_attrs" in source + assert "plain_entry_kernel" in source + for module in executable.mod._collect_from_import_tree(lambda _module: True): + for control_name in ( + "__tvm_iket_get_metadata", + "__tvm_iket_set_context", + "__tvm_iket_clear_context", + ): + assert not module.implements_function(control_name, query_imports=False) + + +def test_proxy_fails_closed_and_forbids_export(tmp_path, monkeypatch): + executable = _compile(serial_a) + with pytest.raises(RuntimeError, match="cannot be exported"): + executable.export_library(tmp_path / "official.so") + + monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE", raising=False) + with pytest.raises(RuntimeError, match="TVM_IKET_OFFICIAL_PROFILE must be set"): + executable.jit() + assert executable._executable._jitted_mod is None # pylint: disable=protected-access + + +def test_environment_validation_is_not_process_cached(tmp_path, monkeypatch): + from tvm.tirx.cuda import iket as _iket_official + + injection_path = tmp_path / "libsmodel_injection.so" + injection_path.write_bytes(b"locked") + injection_relative = "nvidia_cutlass_dsl/dsl_packages/iket/profiler/libsmodel_injection.so" + profile = { + "versions": {"nvidia-cutlass-dsl-libs-base": "4.6.1"}, + "files": { + "nvidia-cutlass-dsl-libs-base": { + injection_relative: hashlib.sha256(b"locked").hexdigest() + } + }, + } + + class FakeDistribution: + version = "4.6.1" + + @staticmethod + def locate_file(_relative_path): + return injection_path + + monkeypatch.setitem(_iket_official._OFFICIAL_PROFILES, "cutlass-4.6.1", profile) + monkeypatch.setattr(_iket_official.metadata, "distribution", lambda _name: FakeDistribution()) + monkeypatch.setattr(_iket_official, "_validate_run_iket_entrypoint", lambda: None) + monkeypatch.setattr( + _iket_official, "_validate_injection_environment", lambda _expected_digest: None + ) + monkeypatch.setattr(_iket_official, "_validate_nvrtc_13_3", lambda: None) + monkeypatch.setenv("TVM_IKET_OFFICIAL_PROFILE", "cutlass-4.6.1") + + _iket_official.validate_official_environment() + monkeypatch.delenv("TVM_IKET_OFFICIAL_PROFILE") + with pytest.raises(RuntimeError, match="TVM_IKET_OFFICIAL_PROFILE must be set"): + _iket_official.validate_official_environment() + + +def test_injection_environment_accepts_run_iket_two_passes(tmp_path, monkeypatch): + from tvm.tirx.cuda import iket as _iket_official + + injection = tmp_path / "libsmodel_injection.so" + injection.write_bytes(b"locked-injection") + expected_digest = hashlib.sha256(injection.read_bytes()).hexdigest() + config_path = tmp_path / "config.json" + monkeypatch.setenv("CUDA_INJECTION64_PATH", str(injection)) + monkeypatch.setenv("SMODEL_INJECTION_CONFIG", str(config_path)) + + config_path.write_text(json.dumps({"toolName": "tracker", "toolConfig": {}})) + _iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access + + instrument = tmp_path / "instrument.config.json" + instrument.write_text("{}", encoding="utf-8") + config_path.write_text( + json.dumps( + { + "toolName": "iket", + "toolConfig": {"appInstrument": str(instrument)}, + } + ), + encoding="utf-8", + ) + _iket_official._validate_injection_environment(expected_digest) # pylint: disable=protected-access + + with pytest.raises(RuntimeError, match="locked run-iket binary"): + _iket_official._validate_injection_environment("0" * 64) # pylint: disable=protected-access + config_path.write_text(json.dumps({"toolName": "other"}), encoding="utf-8") + with pytest.raises(RuntimeError, match="not generated by run-iket profile"): + _iket_official._validate_injection_environment( # pylint: disable=protected-access + expected_digest + ) + + +def test_cutlass_4_6_1_oracle_manifest_integrity(): + oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8")) + assert oracle["schema_version"] == 2 + assert oracle["profile"]["cutlass_dsl"] == "4.6.1" + assert oracle["profile"]["instrument_method"] == "NativeDump" + assert "--dump-dir=" in oracle["profile"]["compiler_flags"] + metadata_bytes = json.dumps(oracle["metadata"], sort_keys=True, separators=(",", ":")).encode() + assert hashlib.sha256(metadata_bytes).hexdigest() == oracle["metadata_sha256"] + abi = oracle["native_dump_abi"] + assert abi["sentinel_event_id"] == 0 + assert abi["range_pop_event_id"] == 31 + assert ( + abi["meta_info_bytes"], + abi["event_attributes_bytes"], + abi["range_attributes_bytes"], + ) == (48, 60, 72) + assert abi["patched_hot_path"] == [ + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "STORE_GLOBAL_32", + "ADD_WRITE_PTR_64_4", + ] + all_hashes = [ + *oracle["artifact_sha256"].values(), + *oracle["patch_artifact_sha256"].values(), + *oracle["wheels"].values(), + oracle["metadata_sha256"], + ] + assert all(re.fullmatch(r"[0-9a-f]{64}", value) for value in all_hashes) + + +def test_external_trace_contract(): + trace_path = os.environ.get("TVM_IKET_OFFICIAL_TRACE_JSON") + if trace_path is None: + pytest.skip("set TVM_IKET_OFFICIAL_TRACE_JSON after the locked run-iket workload") + trace = json.loads(Path(trace_path).read_text(encoding="utf-8")) + assert len(trace["launches"]) == 1 + launch = trace["launches"][0] + assert launch["kernelName"] == "canonical_iket_workload_kernel" + strings = trace["stringTable"] + assert [strings[marker["markerNameIdx"]] for marker in launch["markers"]] == [ + "checkpoint", + "inside_stack", + ] + ranges = {strings[item["rangeNameIdx"]]: item for item in launch["ranges"]} + assert set(ranges) == {"token", "stack"} + assert ranges["token"]["rangeType"] == 1 + assert [item["eventId"] for item in ranges["token"]["internalEvents"]] == [1, 1] + assert ranges["stack"]["rangeType"] == 2 + assert [item["eventId"] for item in ranges["stack"]["internalEvents"]] == [2, 31] + for item in ranges.values(): + assert item["startTs"] <= item["endTs"] + assert [event["timestamp"] for event in item["internalEvents"]] == [ + item["startTs"], + item["endTs"], + ] + + +def test_external_patch_contract(tmp_path): + run_dir = os.environ.get("TVM_IKET_OFFICIAL_PATCH_RUN_DIR") + if run_dir is None: + pytest.skip("set TVM_IKET_OFFICIAL_PATCH_RUN_DIR after a trace-level run-iket workload") + nvdisasm = os.environ.get("TVM_IKET_OFFICIAL_NVDISASM") or shutil.which("nvdisasm") + if nvdisasm is None: + pytest.skip("nvdisasm is required for official patch verification") + + verifier = Path(__file__).with_name("verify_iket_official_patch.py") + subprocess.run( + [ + sys.executable, + str(verifier), + "--run-dir", + run_dir, + "--nvdisasm", + nvdisasm, + "--output-dir", + str(tmp_path), + ], + check=True, + ) + report = json.loads((tmp_path / "verification.json").read_text(encoding="utf-8")) + oracle = json.loads(ORACLE_PATH.read_text(encoding="utf-8")) + assert report["schema_version"] == 1 + assert report["site_count"] > 0 + assert report["normalized_signature"] == oracle["native_dump_abi"]["patched_hot_path"] + assert all(re.fullmatch(r"[0-9a-f]{64}", value) for value in report["sha256"].values()) diff --git a/tests/python/tirx/iket/verify_iket_official_patch.py b/tests/python/tirx/iket/verify_iket_official_patch.py new file mode 100644 index 000000000000..7376f674b323 --- /dev/null +++ b/tests/python/tirx/iket/verify_iket_official_patch.py @@ -0,0 +1,219 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Verify the locked run-iket NativeDump patch from retained trace artifacts. + +Run ``run-iket`` with ``--log-level trace --keep`` so its per-kernel patched +and unpatched images are retained. This helper reconstructs a disassemblable +patched CUBIN outside the repository and emits only a normalized JSON report. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import re +import shutil +import subprocess +from pathlib import Path + +INSTRUCTION = re.compile(r"/\*([0-9a-f]+)\*/\s+(.*?)\s*;") +PATCH_OFFSET = re.compile( + r"Patching SASS => name: (?P.*?), .*?driver_patch_offset: (?P\d+)" +) +STORE_REGISTER = re.compile(r"\bSTS\s+\[[^]]+\],\s+(R\d+)$") +GLOBAL_STORE_REGISTER = re.compile(r"\bSTG\.E\s+\[RZ\.U32\+UR0\],\s+(R\d+)$") +POINTER_INCREMENT = re.compile(r"^UIADD3\.64\s+UR0,\s+UPT,\s+UPT,\s+UR0,\s+0x4,\s+URZ$") +NORMALIZED_SIGNATURE = [ + "GLOBALTIMERLO", + "ENCODE_EVENT_ID", + "STORE_GLOBAL_32", + "ADD_WRITE_PTR_64_4", +] + + +def _sha256(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def _single(paths, description: str) -> Path: + paths = list(paths) + if len(paths) != 1: + raise RuntimeError(f"expected one {description}, found {len(paths)}") + return paths[0] + + +def _discover_artifacts(run_dir: Path, kernel: str | None): + unpatched_paths = list(run_dir.glob("iket/pid_*/unpatched_*")) + if kernel is not None: + unpatched_paths = [path for path in unpatched_paths if path.name == f"unpatched_{kernel}"] + unpatched = _single(unpatched_paths, "unpatched kernel image") + kernel = unpatched.name.removeprefix("unpatched_") + patched = unpatched.with_name(kernel) + if not patched.is_file(): + raise RuntimeError("patched kernel image is missing; rerun with --log-level trace --keep") + injection_log = unpatched.with_name("smodel.injection.log") + offsets = [ + int(match.group("offset")) + for match in PATCH_OFFSET.finditer(injection_log.read_text(encoding="utf-8")) + if match.group("kernel") == kernel + ] + if len(offsets) != 1: + raise RuntimeError(f"expected one driver patch offset for {kernel}, found {offsets}") + + unpatched_bytes = unpatched.read_bytes() + probe = unpatched_bytes[offsets[0] : offsets[0] + 256] + matching_cubins = [] + for cubin in run_dir.glob("tracker/pid_*/module_*.cubin"): + data = cubin.read_bytes() + position = data.find(probe) + if position >= 0 and data.find(probe, position + 1) < 0: + matching_cubins.append((cubin, position)) + if len(matching_cubins) != 1: + raise RuntimeError( + f"expected one tracker CUBIN containing {kernel}, found {len(matching_cubins)}" + ) + cubin, cubin_offset = matching_cubins[0] + return kernel, cubin, unpatched, patched, offsets[0], cubin_offset + + +def _disassemble(nvdisasm: str, cubin: Path) -> str: + return subprocess.run( + [nvdisasm, "-c", str(cubin)], check=True, capture_output=True, text=True + ).stdout + + +def _function_instructions(disassembly: str, kernel: str) -> dict[int, str]: + marker = f"\n.text.{kernel}:\n" + if marker not in disassembly: + raise RuntimeError(f"nvdisasm output does not contain .text.{kernel}") + function = disassembly.split(marker, 1)[1].split("\n//---------------------", 1)[0] + return {int(address, 16): text.strip() for address, text in INSTRUCTION.findall(function)} + + +def _without_predicate(instruction: str) -> str: + return re.sub(r"^@!?[A-Z0-9]+\s+", "", instruction) + + +def _verify_patch(unpatched_sass: str, patched_sass: str, kernel: str) -> int: + before = _function_instructions(unpatched_sass, kernel) + after = _function_instructions(patched_sass, kernel) + if before.keys() != after.keys(): + raise RuntimeError("patched kernel changed the instruction address layout") + changed = [address for address in before if before[address] != after[address]] + if not changed or len(changed) % 2: + raise RuntimeError(f"unexpected changed instruction count: {len(changed)}") + + addresses = sorted(before) + address_index = {address: index for index, address in enumerate(addresses)} + sites = [] + for index in range(0, len(changed), 2): + store_address, increment_address = changed[index : index + 2] + if increment_address not in (store_address + 16, store_address + 32): + raise RuntimeError( + f"patch changes are not a NativeDump store/increment pair at {store_address:#x}" + ) + if increment_address == store_address + 32: + middle = store_address + 16 + if before[middle] != after[middle]: + raise RuntimeError(f"interleaved instruction changed at {middle:#x}") + old_store = _without_predicate(before[store_address]) + new_store = _without_predicate(after[store_address]) + old_register = STORE_REGISTER.search(old_store) + new_register = GLOBAL_STORE_REGISTER.search(new_store) + if old_register is None or new_register is None: + raise RuntimeError(f"unexpected store rewrite at {store_address:#x}") + if old_register.group(1) != new_register.group(1): + raise RuntimeError(f"store source register changed at {store_address:#x}") + if not _without_predicate(before[increment_address]).startswith("PMTRIG "): + raise RuntimeError(f"expected PMTRIG placeholder at {increment_address:#x}") + if POINTER_INCREMENT.fullmatch(_without_predicate(after[increment_address])) is None: + raise RuntimeError(f"unexpected pointer increment at {increment_address:#x}") + + position = address_index[store_address] + prefix = [after[address] for address in addresses[max(0, position - 12) : position]] + if not any("SR_GLOBALTIMERLO" in instruction for instruction in prefix): + raise RuntimeError(f"event site lacks GLOBALTIMERLO at {store_address:#x}") + if not any("LOP3.LUT" in instruction for instruction in prefix): + raise RuntimeError(f"event site lacks event-id encoding at {store_address:#x}") + sites.append(store_address) + + placeholder_count = sum("PMTRIG " in instruction for instruction in before.values()) + if placeholder_count != len(sites): + raise RuntimeError( + f"patched {len(sites)} of {placeholder_count} NativeDump placeholder sites" + ) + if any("PMTRIG " in instruction for instruction in after.values()): + raise RuntimeError("patched kernel still contains PMTRIG placeholders") + return len(sites) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument("--kernel") + parser.add_argument("--nvdisasm", default=shutil.which("nvdisasm")) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + if args.nvdisasm is None: + parser.error("nvdisasm is required") + + kernel, cubin, unpatched, patched, patch_offset, cubin_offset = _discover_artifacts( + args.run_dir, args.kernel + ) + unpatched_bytes = unpatched.read_bytes() + patched_bytes = patched.read_bytes() + if len(unpatched_bytes) != len(patched_bytes) or unpatched_bytes == patched_bytes: + raise RuntimeError("patched and unpatched kernel images have invalid sizes or contents") + + cubin_bytes = bytearray(cubin.read_bytes()) + patched_code = patched_bytes[patch_offset:] + end = cubin_offset + len(patched_code) + if end > len(cubin_bytes): + raise RuntimeError("patched kernel image extends beyond the tracker CUBIN") + cubin_bytes[cubin_offset:end] = patched_code + + args.output_dir.mkdir(parents=True, exist_ok=True) + patched_cubin = args.output_dir / "patched.cubin" + patched_cubin.write_bytes(cubin_bytes) + unpatched_sass = _disassemble(args.nvdisasm, cubin) + patched_sass = _disassemble(args.nvdisasm, patched_cubin) + (args.output_dir / "unpatched.sass").write_text(unpatched_sass, encoding="utf-8") + (args.output_dir / "patched.sass").write_text(patched_sass, encoding="utf-8") + site_count = _verify_patch(unpatched_sass, patched_sass, kernel) + + report = { + "schema_version": 1, + "kernel": kernel, + "site_count": site_count, + "driver_patch_offset": patch_offset, + "cubin_code_offset": cubin_offset, + "normalized_signature": NORMALIZED_SIGNATURE, + "sha256": { + "tracker_cubin": _sha256(cubin), + "unpatched_kernel": _sha256(unpatched), + "patched_kernel": _sha256(patched), + "unpatched_disassembly": hashlib.sha256(unpatched_sass.encode()).hexdigest(), + "patched_disassembly": hashlib.sha256(patched_sass.encode()).hexdigest(), + }, + } + report_path = args.output_dir / "verification.json" + report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(report_path) + + +if __name__ == "__main__": + main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py index 1b8f0239c242..d1a227d0f89e 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py @@ -18,8 +18,8 @@ """Tests for the priority=0 ``copy/fallback`` dispatch — scalar single-thread emit picked when every higher-priority variant rejects. -The cases here are *intentionally* shaped so ``gmem_smem`` rejects (region -element count doesn't divide ``thread_cnt``) and ``reg`` / ``ld_stmatrix`` / +The cases here are *intentionally* shaped so the ``vec_auto`` global ↔ shared +path rejects (region element count doesn't divide ``thread_cnt``) and register / ``ld_stmatrix`` / ... don't apply (scope pair mismatch). The dispatcher should land on fallback, the emit should pick one active thread, and the round-trip ``A_gmem → A_smem → B_gmem`` should match. @@ -42,7 +42,7 @@ def _round_trip_shapes_and_threads(): - """Cases where ``gmem_smem`` rejects on ``n_elements % thread_cnt``. + """Cases where the vec_auto global ↔ shared path rejects on ``n_elements % thread_cnt``. Per task: ``(scope, n_threads, shape, why_fallback)``. ``shape`` is small enough that scalar emit is fine, and chosen so the higher-priority @@ -168,7 +168,7 @@ def run_and_check(): @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") def test_fallback_thread_scope(): - """``T.thread()`` — single thread, no gate. Either ``gmem_smem`` picks + """``T.thread()`` — single thread, no gate. Either the vec_auto path picks it up (n_elements % 1 == 0) or ``fallback`` does — both end up emitting a sensible single-thread copy. We only check the round trip is correct, not which variant fired.""" diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py index 6d650e3610e3..f83486bd4c76 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring -"""Round-trip tests for the ``gmem_smem`` copy dispatch (synthesized partition). +"""Round-trip tests for the ``vec_auto`` global ↔ shared path. Pipeline: A_gmem --G2S--> A_smem --S2G--> B_gmem. If either direction is wrong the round trip leaves B mismatched against A. @@ -29,7 +29,7 @@ from tvm.script import tirx as T from tvm.script.tirx import tile as Tx from tvm.testing import env -from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import ComposeLayout, S, TileLayout def _build_kernel(scope, n_threads, shape, dtype): @@ -134,7 +134,7 @@ def run_and_check(): # ---------------------------------------------------------------------------- # Migrated from test_copy_sync.py: sync G↔S copy via the user-facing -# Tx.copy() (which dispatches to gmem_smem). +# Tx.copy() (which dispatches to the vec_auto global ↔ shared path). # ---------------------------------------------------------------------------- @@ -189,7 +189,7 @@ def run_and_check(): 32, TileLayout(S[96, 512]), TileLayout(S[96, 512]), - ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[8, 64])) + ComposeLayout(3, 3, 3, TileLayout(S[8, 64])) .tile_to((16, 128), (8, 64)) .tile_to((32, 256), (16, 128)), ), @@ -287,17 +287,17 @@ def _align( ) @pytest.mark.parametrize("per_element,expected_max_vec", [(2, 4), (1, 2), (0, 1)]) def test_swizzled_smem_vec_len_must_fit_chunk(per_element, expected_max_vec): - """``SwizzleLayout(per_element, ...)`` keeps the bottom ``per_element`` + """A swizzled ``ComposeLayout`` keeps the bottom ``per_element`` bits unswizzled. vec must stay within that chunk or it crosses an XOR boundary and reads/writes the wrong physical bytes.""" shape = (32, 32) # 1024 fp16 elements total g_layout = TileLayout(S[shape]) - s_layout = ComposeLayout(SwizzleLayout(per_element, 3, 3), TileLayout(S[shape])) + s_layout = ComposeLayout(per_element, 3, 3, TileLayout(S[shape])) _g, _s, vec_len = _align(g_layout, shape, s_layout, shape, elem_bits=16, thread_cnt=32) chunk_elems = 1 << per_element assert vec_len <= chunk_elems, ( f"vec_len={vec_len} crosses swizzle chunk size={chunk_elems} " - f"(SwizzleLayout(per_element={per_element}, ...))" + f"(swizzle per_element={per_element}, ...)" ) @@ -348,16 +348,16 @@ def test_unaligned_region_offset_must_clamp_vec_len(): def test_swizzled_smem_emit_must_be_swizzle_aware(): - """Codegen-level: emitted S address should go through the SwizzleLayout's + """Codegen-level: emitted S address should go through the swizzle's Apply so the XOR scrambling is honored. Currently emit uses ``s_buf.ptr_to([0,..,0]) + linear_offset`` which only matches a non-swizzled storage layout.""" import tvm from tvm.script import tirx as T - from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout + from tvm.tirx.layout import ComposeLayout, S, TileLayout shape = (128, 32) - s_layout = ComposeLayout(SwizzleLayout(3, 3, 3), TileLayout(S[shape])) + s_layout = ComposeLayout(3, 3, 3, TileLayout(S[shape])) @T.prim_func def kernel(A_ptr: T.handle) -> None: @@ -506,7 +506,7 @@ def test_layout_permute_copy_preserves_smem_strides(): # ---------------------------------------------------------------------------- # Fast-path firing test (positive). Pairs with the var_bounds wiring inside -# ``gmem_smem._emit_gmem_smem``. +# ``vec_auto_gmem_smem._emit_gmem_smem``. # # Setup: warp-scope 32x64 fp16 G2S/S2G with 128b swizzled SMEM. The outer # iter stride is ``thread_cnt * vec_len = 32 * 8 = 256``, which puts the @@ -525,10 +525,16 @@ def test_gmem_smem_swizzle_fast_path_fires_with_var_bounds(): per outer iter, no per-iter ``swizzle.apply`` XOR splice in the hot path.""" import re - swizzle = SwizzleLayout(3, 3, 3) + swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) shape = (32, 64) g_layout = TileLayout(S[shape]) - s_layout = ComposeLayout(swizzle, TileLayout(S[shape])) + s_layout = ComposeLayout( + swizzle.per_element, + swizzle.swizzle_len, + swizzle.atom_len, + TileLayout(S[shape]), + swizzle.swizzle_inner, + ) @T.prim_func def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py index d5ae8130a764..64e66cd5e089 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py @@ -40,7 +40,7 @@ from tvm.script import tirx as T from tvm.script.tirx import tile as Tx from tvm.testing import env -from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout, laneid, tid_in_wg, tx +from tvm.tirx.layout import ComposeLayout, S, TileLayout, laneid, tid_in_wg, tx def _compile_src(kernel): @@ -79,13 +79,19 @@ def _s_layout_warpgroup_or_cta(num, trans): # 128b swizzle for fp16 (p=3 ⇒ 8 fp16 chunk; sw=at=3 ⇒ 8-row swizzle period). -_SWIZZLE_128B = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) +_SWIZZLE_128B = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) def _maybe_wrap_swizzle(tile_layout, enable: bool): if not enable: return tile_layout - return ComposeLayout(_SWIZZLE_128B, tile_layout) + return ComposeLayout( + _SWIZZLE_128B.per_element, + _SWIZZLE_128B.swizzle_len, + _SWIZZLE_128B.atom_len, + tile_layout, + _SWIZZLE_128B.swizzle_inner, + ) # --------------------------------------------------------------------------- @@ -405,7 +411,7 @@ def _build_multi_iter_kernel(outer_ext: int): = outer_ext*16 - 1, matching extent product = 16*outer_ext.""" shape = (outer_ext, 8, 2, 4, 4, 2) r_layout = TileLayout(S[shape : (16, 4 @ laneid, 8, 2, 1 @ laneid, 1)]) - s_layout = SwizzleLayout(3, 3, 3) + s_layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) full = tuple(slice(0, e) for e in shape) @T.prim_func diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py index a19ca76d1272..4b7ecc26d2db 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=missing-function-docstring -"""Round-trip tests for the ``reg`` copy dispatch. +"""Round-trip tests for the ``vec_auto`` register copy path. R = per-thread local (register). The dispatch handles round-trips between R and any non-R buffer (``shared*`` or ``global``); ``non_r_scope`` parametrize @@ -274,7 +274,7 @@ def run_and_check(): # ---------------------------------------------------------------------------- # Migrated from test_copy_sync.py: sync G↔L copy via Tx.copy() (L = local = -# per-thread register, so it dispatches to the reg variant). +# per-thread register, so it dispatches to the vec_auto register path). # ---------------------------------------------------------------------------- @@ -338,6 +338,65 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +# vec_auto reg path must honor cache="nc": a strided [4,4]:(16,1) global->reg +# copy vectorizes to 4x 128b either way, but cache="nc" must emit ld.global.nc. + + +@T.prim_func +def _nc_strided_reg_copy(src_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1024,), "int32") + T.device_entry() + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + dst = T.alloc_local((4, 4), "int32") + if tid == 0: + # view as (blk, row, warp, j); pick warp=1 -> [4,4]:(16,1) + blk = src.view(1024 // 64, 4, 4, 4).sub[0, :, 1, :] + Tx.copy(dst[:, :], blk[:, :], cache="nc") + for i in T.unroll(4): + for j in T.unroll(4): + T.evaluate(dst[i, j]) + + +@T.prim_func +def _plain_strided_reg_copy(src_ptr: T.handle) -> None: + src = T.match_buffer(src_ptr, (1024,), "int32") + T.device_entry() + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + dst = T.alloc_local((4, 4), "int32") + if tid == 0: + blk = src.view(1024 // 64, 4, 4, 4).sub[0, :, 1, :] + Tx.copy(dst[:, :], blk[:, :]) + for i in T.unroll(4): + for j in T.unroll(4): + T.evaluate(dst[i, j]) + + +def test_vec_auto_reg_honors_cache_nc(): + target = tvm.target.Target("cuda") + + def _src(f): + with target: + mod = tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx") + return mod.mod.imports[0].inspect_source() + + nc_src = _src(_nc_strided_reg_copy) + plain_src = _src(_plain_strided_reg_copy) + + # cache="nc" -> vectorized 128b ld.global.nc; no cache -> plain 128b ld. + # Both must vectorize; only the cache qualifier differs. + assert "ld_global_nc_v4_u32" in nc_src, ( + "vec_auto reg path must vectorize AND honor cache='nc' (emit " + "ld.global.nc.v4), got:\n" + + "\n".join(line for line in nc_src.splitlines() if "ld_" in line and "v4" in line) + ) + assert "ld_plain_None_global_v4_u32" in plain_src, ( + "no cache hint must vectorize to a plain 128b ld (ld.global.v4)" + ) + assert "ld_global_nc_v4_u32" not in plain_src, "no cache hint must NOT be nc" + + @pytest.mark.gpu def test_reg_copy_wg_local_to_swizzled_shared_uses_swizzle_fastpath(): """Regression: R→S copy where R has a ``wg_local_layout`` (thread iter @@ -360,13 +419,13 @@ def test_reg_copy_wg_local_to_swizzled_shared_uses_swizzle_fastpath(): ``tvm_builtin_pointer_offset`` swizzle XOR ends up recomputed every iteration. Loop must be ``T.unroll``. """ - from tvm.tirx.layout import SwizzleLayout, wg_local_layout + from tvm.tirx.layout import ComposeLayout, wg_local_layout N_THREADS, EPI_N = 128, 64 g_shape = (N_THREADS, EPI_N) g_layout = TileLayout(S[g_shape]) # 128b swizzle on the SMEM side (per_element=3 ⇒ 8 fp16 atom width). - smem_layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + smem_layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) @T.prim_func def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: @@ -419,6 +478,248 @@ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: ) +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_ptx_st_from_src_f32_vector_preserves_values(): + """The generic ``T.ptx.st(src=...)`` helper must preserve f32 values.""" + + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_local((4,), "float32") + out = T.alloc_local((4,), "float32") + for i in range(4): + reg[i] = T.cast(i + 1, "float32") + T.ptx.st(smem.ptr_to([0]), src=reg.ptr_to([0]), space="shared", vec="v4", ptx_type="f32") + T.ptx.ld(smem.ptr_to([0]), "float32", "f32", dst=out.ptr_to([0]), space="shared", vec="v4") + for i in range(4): + B[i] = out[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = ex.mod.imports[0].inspect_source() + + assert "float4 src_" in src + assert "st.shared.v4.f32" in src + + dev = tvm.cuda(0) + out = tvm.runtime.tensor(np.zeros((4,), dtype="float32"), dev) + ex(out) + np.testing.assert_equal(out.numpy(), np.array([1, 2, 3, 4], dtype="float32")) + + +def test_copy_fallback_handles_scalar_regions(): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (1,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + src = T.alloc_local((1,), "float32") + dst = T.alloc_local((4,), "float32") + src[0] = T.cast(7, "float32") + Tx.copy(dst[2:3], src[:]) + B[0] = dst[2] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = ex.mod.imports[0].inspect_source() + + assert "dst_ptr[2] = src_ptr[0];" in src + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +@pytest.mark.parametrize( + "variant,dtype,n_elements,expected_st,expected_ld", + [ + ("vec_16b", "uint16", 1, "st.shared.u16", "ld.shared.u16"), + ("vec_32b", "uint32", 1, "st.shared.u32", "ld.shared.u32"), + ("vec_64b", "uint32", 2, "st.shared.v2.u32", "ld.shared.v2.u32"), + ("vec_128b", "uint32", 4, "st.shared.v4.u32", "ld.shared.v4.u32"), + ], +) +def test_copy_forced_vec_width_codegen(variant, dtype, n_elements, expected_st, expected_ld): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (n_elements,), dtype) + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((n_elements,), dtype, scope="shared") + reg = T.alloc_local((n_elements,), dtype) + out = T.alloc_local((n_elements,), dtype) + for i in range(n_elements): + reg[i] = T.cast(i + 1, dtype) + Tx.copy(smem[:], reg[:], dispatch=variant) + T.cuda.cta_sync() + Tx.copy(out[:], smem[:], dispatch=variant) + for i in range(n_elements): + B[i] = out[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = ex.mod.imports[0].inspect_source() + + assert expected_st in src + assert expected_ld in src + + dev = tvm.cuda(0) + np_dtype = tvm.testing.np_dtype_from_str(dtype) + out = tvm.runtime.tensor(np.zeros((n_elements,), dtype=np_dtype), dev) + ex(out) + np.testing.assert_equal(out.numpy(), np.arange(1, n_elements + 1, dtype=np_dtype)) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_forced_vec_dynamic_swizzled_shared_uses_vector_ptx(): + from tvm.tirx.layout import ComposeLayout + + smem_layout = ComposeLayout(2, 3, 3, TileLayout(S[(64, 8, 32) : (32, 2048, 1)])) + + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (128, 4), "float32") + T.device_entry() + T.cta_id([1]) + tid = T.thread_id([128]) + smem = T.alloc_buffer((64, 256), "float32", scope="shared", layout=smem_layout) + reg = T.alloc_local((4,), "float32") + out = T.alloc_local((4,), "float32") + for i in range(4): + reg[i] = T.cast(tid * 16 + i + 1, "float32") + row: T.let = tid % 64 + col: T.let = (tid // 64) * 4 + Tx.copy(smem[row, col : col + 4], reg[:], dispatch="vec_128b") + T.cuda.cta_sync() + Tx.copy(out[:], smem[row, col : col + 4], dispatch="vec_128b") + for i in range(4): + B[tid, i] = out[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = ex.mod.imports[0].inspect_source() + + assert "copy/fallback" not in src + assert "st.shared.v4.u32" in src + assert "ld.shared.v4.u32" in src + + dev = tvm.cuda(0) + out = tvm.runtime.tensor(np.zeros((128, 4), dtype="float32"), dev) + ex(out) + expected = np.array([[tid * 16 + i + 1 for i in range(4)] for tid in range(128)], "float32") + np.testing.assert_equal(out.numpy(), expected) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_explicit_vec_auto_uses_auto_family(): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_buffer((4,), "float32", scope="local", layout=TileLayout(S[4])) + out = T.alloc_buffer((4,), "float32", scope="local", layout=TileLayout(S[4])) + for i in range(4): + reg[i] = T.cast(i + 1, "float32") + Tx.copy(smem[:], reg[:], dispatch="vec_auto") + T.cuda.cta_sync() + Tx.copy(out[:], smem[:], dispatch="vec_auto") + for i in range(4): + B[i] = out[i] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = ex.mod.imports[0].inspect_source() + + assert "tvm_builtin_copy_" not in src + assert "st.shared.v4.u32" in src + assert "ld.shared.v4.u32" in src + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +@pytest.mark.parametrize("dispatch", ["reg", "gmem_smem"]) +def test_copy_old_dispatch_names_are_not_registered(dispatch): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_local((4,), "float32") + Tx.copy(smem[:], reg[:], dispatch=dispatch) + B[0] = T.cast(0, "float32") + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + with pytest.raises(RuntimeError, match=f"no variant named '{dispatch}' is registered"): + tvm.compile(mod, target=target, tir_pipeline="tirx") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_forced_vec_rejects_size_mismatch(): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_local((4,), "float32") + Tx.copy(smem[:], reg[:], dispatch="vec_64b") + B[0] = T.cast(0, "float32") + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + with pytest.raises(RuntimeError, match="src region does not contain exactly 2 elements"): + tvm.compile(mod, target=target, tir_pipeline="tirx") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_forced_vec_rejects_non_thread_scope(): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.lane_id([32]) + T.thread_id([32]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_buffer((4,), "float32", scope="local", layout=TileLayout(S[4])) + Tx.warp.copy(smem[:], reg[:], dispatch="vec_128b") + B[0] = T.cast(0, "float32") + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + with pytest.raises(RuntimeError, match="expected thread exec_scope"): + tvm.compile(mod, target=target, tir_pipeline="tirx") + + # --- tcgen05 D epilogue deposit (tf32_hc_prenorm_gemm) ----------------------- # Production op: ``Tx.warpgroup.copy(smem_cd_mma, d_reg)`` after ``tcgen05.ld`` # pulls the M=64 accumulator fragment from TMEM into ``d_reg``, then deposits it @@ -480,7 +781,7 @@ def test_reg_copy_tcgen05_d_epilogue_deposit_layout_pairing(): ``d_reg``: ``(64,64)`` fp32 ``tcgen05_atom_layout("16x256b", ...)``. ``smem_cd_mma``: ``(64,64)`` fp32 ``mma_shared_layout(..., swizzle=128B)``. """ - from tvm.backend.cuda.operator.tile_primitive.copy.reg import ( + from tvm.backend.cuda.operator.tile_primitive.copy.vec_auto_reg import ( _split_thread_loop, align_layouts_raw, ) @@ -525,8 +826,8 @@ def test_reg_copy_tcgen05_d_epilogue_deposit_codegen(): ex = _compile_tcgen05_d_epilogue_deposit() src = ex.mod.imports[0].inspect_source() - assert "copy/fallback" not in src, "reg dispatch must not fall back to scalar copy" - assert "tvm_builtin_copy_" not in src, "reg dispatch should emit PTX ld/st only" + assert "copy/fallback" not in src, "vec_auto register path must not fall back to scalar copy" + assert "tvm_builtin_copy_" not in src, "vec_auto register path should emit PTX ld/st only" assert "tvm_builtin_ptx_st" in src assert "st.shared.v2.u32" in src, "fp32 vec=2 → 8B shared store per outer iter" diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py index 1231b338ba08..98bf68e61fb1 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py @@ -42,16 +42,16 @@ try_recognize, ) from tvm.tirx.expr import IntImm as _IntImm -from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, TileLayout +from tvm.tirx.layout import ComposeLayout, S, TileLayout # ---------------------------------------------------------------------------- -# Pure-Python reference: SwizzleLayout's Apply, plus the proof's formula. +# Pure-Python reference: the bare-swizzle ComposeLayout's Apply, plus the proof's formula. # Used as ground truth — both must agree for the proof to hold. # ---------------------------------------------------------------------------- def py_swizzle_apply(M: int, p: int, sw: int, at: int) -> int: - """Pure-Python reimplementation of SwizzleLayoutNode::Apply (swizzle_inner=True): + """Pure-Python reimplementation of the swizzle Apply (swizzle_inner=True): phys = swz_q * C + (M mod C) q = M / C; swz_q = q XOR ((q & outer_mask) >> at) """ @@ -111,18 +111,29 @@ def py_outer_ds(k: int, iter_extents: list[int], iter_strides: list[int]) -> int def test_get_swizzle_extracts_from_compose(): - sw = SwizzleLayout(3, 3, 3) + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert get_swizzle(sw) is not None - assert get_swizzle(ComposeLayout(sw, TileLayout(S[(64, 64)]))) is not None + assert ( + get_swizzle( + ComposeLayout( + sw.per_element, + sw.swizzle_len, + sw.atom_len, + TileLayout(S[(64, 64)]), + sw.swizzle_inner, + ) + ) + is not None + ) assert get_swizzle(TileLayout(S[(64, 64)])) is None def test_recognize_nvfp4_case(): - """nvfp4's epilogue: SwizzleLayout(3,3,3), iter extents [2,2,2] strides + """nvfp4's epilogue: swizzle(3,3,3), iter extents [2,2,2] strides [8,16,32], M0 = tid * 64 (each thread starts at col 0 of one row; row_stride 64 = 8 chunks, ensures chunk bits of M0/C are zero for all iter bit positions).""" - sw = SwizzleLayout(3, 3, 3) + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) tid = _TirVar("tid", "int32") # M0 = tid * 64 → M0/C = tid * 8 → bits 0,1,2 are 0 (since multiplied by 8). M0 = tid * _IntImm("int32", 64) @@ -136,7 +147,7 @@ def test_recognize_nvfp4_case(): def test_recognize_binary_split(): """A single outer iter with extent=4 stride=8 splits into two binary iters with strides 16 and 8 (outermost first, matching _flat_outer_coords).""" - sw = SwizzleLayout(3, 3, 3) + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) tid = _TirVar("tid", "int32") M0 = tid * _IntImm("int32", 64) pat = try_recognize(sw, [4], [8], M0) @@ -147,10 +158,10 @@ def test_recognize_binary_split(): def test_recognize_mid_bits(): - """SwizzleLayout(p=4, sw=2, at=4): chunk bits [0,2), mid bits [2,4), + """swizzle(p=4, sw=2, at=4): chunk bits [0,2), mid bits [2,4), row bits [4,6). An iter at bj=2 lives in mid_bits → sigma is always +1 (i.e., the recognizer accepts and the sign formula won't read row bits).""" - sw = SwizzleLayout(4, 2, 4) # C=16, mid_bits cover bits 2..3 + sw = ComposeLayout(4, 2, 4, TileLayout(S[(1024,)])) # C=16, mid_bits cover bits 2..3 tid = _TirVar("tid", "int32") # M0/C must have bit 2 == 0. Pick row_stride = 64 (= 4*C) so M0/C = tid*4 # which has zeros at bit 0,1, and bit 2 is bit 0 of tid... hmm that varies. @@ -167,7 +178,7 @@ def test_recognize_mid_bits(): def test_reject_not_chunk_aligned(): """Condition (a): stride must be a multiple of C.""" - sw = SwizzleLayout(3, 3, 3) # C=8 + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) # C=8 tid = _TirVar("tid", "int32") M0 = tid * _IntImm("int32", 64) # stride 4 is not a multiple of C=8 → reject. @@ -177,7 +188,7 @@ def test_reject_not_chunk_aligned(): def test_reject_carries_into_row_bits(): """Condition (b): bj < at. A binary iter with stride C * 2^at lands at bj=at, which would change the row bits → reject.""" - sw = SwizzleLayout(3, 3, 3) # at=3, so max bj = 2 + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) # at=3, so max bj = 2 tid = _TirVar("tid", "int32") M0 = tid * _IntImm("int32", 64) # Strides 8,16,32 OK (bj=0,1,2); 64 → bj=3 → reject. @@ -188,7 +199,7 @@ def test_reject_chunk_overlap(): """Condition (c): (M0/C) must have 0 bits at all iter-bit positions per thread. If M0 = tid * 8 (so M0/C = tid), then bit 0 of M0/C is bit 0 of tid — analyzer can't prove this is 0 across all threads, so reject.""" - sw = SwizzleLayout(3, 3, 3) # C=8 + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) # C=8 tid = _TirVar("tid", "int32") # M0 = tid * C = tid * 8 → M0/C = tid → bit 0 NOT provably zero. M0 = tid * _IntImm("int32", 8) @@ -198,7 +209,7 @@ def test_reject_chunk_overlap(): def test_recognize_no_outer_iters(): """Degenerate case: no outer iter at all. Recognizer returns a trivial pattern (empty bit_positions). Emit will use base_off alone.""" - sw = SwizzleLayout(3, 3, 3) + sw = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) tid = _TirVar("tid", "int32") M0 = tid * _IntImm("int32", 64) pat = try_recognize(sw, [], [], M0) @@ -245,7 +256,7 @@ def test_formula_matches_apply_under_conditions( """For every (M0, k) sample, the signed-strides formula must equal py_swizzle_apply(M0 + ds_k). Sweeps multiple per-thread M0 values to catch any per-thread-sign bug (a constant-sign impl would fail here).""" - swizzle = SwizzleLayout(p, sw, at) + swizzle = ComposeLayout(p, sw, at, TileLayout(S[(1 << (p + sw + at),)])) tid = _TirVar("tid", "int32") M0_template = tid * _IntImm("int32", row_stride) pat = try_recognize(swizzle, iter_extents, iter_strides, M0_template) @@ -318,7 +329,7 @@ def test_recognize_linear_iter_pure_case_1d(): ) p, sw, at = 3, 3, 3 - swizzle = SwizzleLayout(p, sw, at) + swizzle = ComposeLayout(p, sw, at, TileLayout(S[(1 << (p + sw + at),)])) period = 1 << (p + at + sw) # 512 # Outer iter (ext=3, stride=period) — non-pow2 but pure Case 1.D. # Inner iter (ext=2, stride=8) — pow2, Case 1.A (bj=0). @@ -343,7 +354,7 @@ def test_reject_non_pow2_ext_not_case_1d(): """Non-pow2 ext where stride is NOT in pure Case 1.D regime — reject. stride=64 = 2^(p+at) = one atom row, which is Case 1.C (in [at, at+sw)) territory and the XOR depends on M0, so the linear path is unsafe.""" - swizzle = SwizzleLayout(3, 3, 3) + swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) pat = try_recognize(swizzle, [3], [64], _IntImm("int32", 0)) assert pat is None @@ -357,7 +368,7 @@ def test_emit_mixed_linear_bit_correctness(): ) p, sw, at = 3, 3, 3 - swizzle = SwizzleLayout(p, sw, at) + swizzle = ComposeLayout(p, sw, at, TileLayout(S[(1 << (p + sw + at),)])) period = 1 << (p + at + sw) # 512 iter_extents, iter_strides = [3, 2], [period, 8] tid = _TirVar("tid", "int32") @@ -414,7 +425,7 @@ def test_fallback_path_when_recognizer_rejects(): wrong answer on at least one (tid, k) sample. The fallback emit, by construction, delegates to swizzle.apply and is thus correct.""" p, sw, at = 3, 3, 3 - swizzle = SwizzleLayout(p, sw, at) + swizzle = ComposeLayout(p, sw, at, TileLayout(S[(1 << (p + sw + at),)])) tid = _TirVar("tid", "int32") M0_template = tid * _IntImm("int32", 8) # (c) fails: bit 0 of M0/C = bit 0 of tid pat = try_recognize(swizzle, [2], [8], M0_template) diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_vec_forced_cache.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_vec_forced_cache.py new file mode 100644 index 000000000000..3dcf9f9ccad2 --- /dev/null +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_vec_forced_cache.py @@ -0,0 +1,172 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=missing-function-docstring +"""Tests for the forced-vec copy cache-semantics config. + +Covers the ``vec_256b`` (32-byte / PTX ``.v8``) variant and the +``cache="nc"`` + L1/L2 hint config on the forced-vec copy dispatches: +``Tx.copy(dst, src, dispatch="vec_*", cache="nc", l1_evict=..., ...)`` +must lower the global load to ``ld.global.nc`` with the hint suffixes +instead of a plain ``ld.global``. +""" + +import numpy as np +import pytest + +import tvm +import tvm.testing +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.testing import env + + +def _build_g2l2g_kernel(n_elements, dtype, dispatch, **copy_config): + """One thread: A (global) → reg (local) via forced-vec copy (with the + given config), then reg → B (global) via a plain forced-vec copy.""" + + @T.prim_func + def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (n_elements,), dtype) + B = T.match_buffer(B_ptr, (n_elements,), dtype) + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + reg = T.alloc_local((n_elements,), dtype) + Tx.copy(reg[:], A[:], dispatch=dispatch, **copy_config) + Tx.copy(B[:], reg[:], dispatch=dispatch) + + return kernel + + +def _build_g2s2g_kernel(n_elements, dtype, dispatch, **copy_config): + """One thread: A (global) → smem via forced-vec copy (ld + st through a + local tmp inside the dispatch), then smem → B elementwise.""" + + @T.prim_func + def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (n_elements,), dtype) + B = T.match_buffer(B_ptr, (n_elements,), dtype) + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((n_elements,), dtype, scope="shared") + Tx.copy(smem[:], A[:], dispatch=dispatch, **copy_config) + for i in range(n_elements): + B[i] = smem[i] + + return kernel + + +def _compile(kernel): + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + ex = tvm.compile(mod, target=target, tir_pipeline="tirx") + return ex, ex.mod.imports[0].inspect_source() + + +def _run_roundtrip(ex, n_elements, dtype): + dev = tvm.cuda(0) + np_dtype = tvm.testing.np_dtype_from_str(dtype) + a_np = np.arange(1, n_elements + 1, dtype=np_dtype) * 3 + a = tvm.runtime.tensor(a_np, dev) + b = tvm.runtime.tensor(np.zeros((n_elements,), dtype=np_dtype), dev) + ex(a, b) + np.testing.assert_array_equal(b.numpy(), a_np) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="PTX .v8 ld/st needs cuda compute >= 10.0") +def test_copy_vec_256b_plain(): + kernel = _build_g2l2g_kernel(8, "int32", "vec_256b") + ex, src = _compile(kernel) + assert "ld.global.v8.u32" in src + assert "st.global.v8.u32" in src + assert "ld.global.nc" not in src + _run_roundtrip(ex, 8, "int32") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_vec_128b_nc_with_hints(): + # NB: ptxas only accepts an L2 eviction-priority qualifier on 256-bit + # loads (.v8.b32 / .v4.b64), so the 128-bit case sticks to L1 + prefetch. + kernel = _build_g2l2g_kernel( + 4, + "int32", + "vec_128b", + cache="nc", + l1_evict="L1::no_allocate", + prefetch_size="L2::256B", + ) + ex, src = _compile(kernel) + assert "ld.global.nc.L1::no_allocate.L2::256B.v4.u32" in src + assert "st.global.v4.u32" in src + _run_roundtrip(ex, 4, "int32") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="PTX .v8 ld/st needs cuda compute >= 10.0") +def test_copy_vec_256b_nc_with_hints(): + kernel = _build_g2l2g_kernel( + 8, + "int32", + "vec_256b", + cache="nc", + l1_evict="L1::no_allocate", + l2_evict="L2::evict_first", + prefetch_size="L2::256B", + ) + ex, src = _compile(kernel) + assert "ld.global.nc.L1::no_allocate.L2::evict_first.L2::256B.v8.u32" in src + assert "v8" in src + _run_roundtrip(ex, 8, "int32") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_vec_128b_nc_global_to_shared(): + kernel = _build_g2s2g_kernel(4, "float32", "vec_128b", cache="nc") + ex, src = _compile(kernel) + assert "ld.global.nc.v4.u32" in src + assert "st.shared.v4.u32" in src + _run_roundtrip(ex, 4, "float32") + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_copy_vec_nc_rejects_non_global_src(): + @T.prim_func + def kernel(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (4,), "float32") + T.device_entry() + T.cta_id([1]) + T.thread_id([1]) + smem = T.alloc_buffer((4,), "float32", scope="shared") + reg = T.alloc_local((4,), "float32") + Tx.copy(reg[:], smem[:], dispatch="vec_128b", cache="nc") + B[0] = reg[0] + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + with pytest.raises(RuntimeError, match="requires src scope 'global'"): + tvm.compile(mod, target=target, tir_pipeline="tirx") + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py index 4a6e046939f7..ce468c41adbb 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py @@ -35,10 +35,10 @@ from tvm.tirx.cuda.operator.tile_primitive.copy_async.dsmem import copy_dsmem_impl from tvm.tirx.exec_scope import ExecScope from tvm.tirx.layout import S, TileLayout -from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext from tvm.tirx.operator.tile_primitive.dispatcher import DispatchFail from tvm.tirx.operator.tile_primitive.ops import CopyAsync from tvm.tirx.stmt_functor import StmtExprVisitor +from tvm.tirx.tile_primitive import DispatchContext def _make_dsmem_dispatch_call(shape, dtype, src_layout, dst_layout): diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_ldgsts.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_ldgsts.py index 080743933083..6b1b22cb9078 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_ldgsts.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_ldgsts.py @@ -119,5 +119,37 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +def test_copy_ldgsts_predicate_zero_fill_codegen(): + """ldgsts direct mode forwards predicate/zero-fill/prefetch without partition temps.""" + + @T.prim_func + def copy_async(A_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (32, 16), "uint8", layout=TileLayout(S[32, 16])) + + T.device_entry() + tid = T.thread_id([32]) + A_smem = T.alloc_buffer((32, 16), "uint8", scope="shared", layout=TileLayout(S[32, 16])) + + Tx.copy_async( + A_smem[tid, :], + A[tid, :], + dispatch="ldgsts", + direct=True, + prefetch_size=128, + predicate=tid < 16, + fill_mode="zero", + ) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile(tvm.IRModule({"main": copy_async}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + assert "cp.async.cg.shared.global.L2::128B" in src + assert "with_src_size" in src + assert "condval" in src + assert "s_ptr_ptr" not in src + assert "g_ptr_ptr" not in src + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py deleted file mode 100644 index e58f4b2423b4..000000000000 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py +++ /dev/null @@ -1,481 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# pylint: disable=invalid-name, missing-function-docstring -"""End-to-end tests for the smem->tmem (tcgen05.cp.32x128b.warpx4) dispatch. - -The new dispatch requires the user to declare the t buffer with an -explicit ``R[4 : 32@TLane]`` indicating warpx4 broadcast — i.e., t.shape[lane] = 32 -with replica 4 → 128 physical lanes. - -Run with: pytest test_smem_tmem_dispatch.py -n 8 -v -""" - -import numpy as np -import pytest - -import tvm -import tvm.testing -from tvm.script import tirx as T -from tvm.script.tirx import tile as Tx -from tvm.testing import env -from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode, mma_shared_layout -from tvm.tirx.layout import R, S, TCol, TileLayout, TLane - -T_LAY_BASIC = TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) - - -def _make_2d_kernel( - s_full, - t_full, - s_full_shape, - t_full_shape, - s_r0, - s_r1, - s_c0, - s_c1, - t_r0, - t_r1, - t_c0, - t_c1, - dtype, - cta_group=1, -): - """2D variant: SMEM/TMEM are both 2D; copy a rectangular sub-region.""" - n_tmem_cols_total = max(32, t_full_shape[-1]) - OUT_LANES = 32 - OUT_BYTES = 16 - - @T.prim_func(check_well_formed=False) - def kernel(A_ptr: T.handle, B_ptr: T.handle): - A = T.match_buffer(A_ptr, s_full_shape, dtype) - B = T.match_buffer(B_ptr, (OUT_LANES, OUT_BYTES), dtype) - T.device_entry() - warp_id = T.warp_id([4]) - wg_id = T.warpgroup_id([1]) - tid_in_wg = T.thread_id_in_wg([128]) - lane_id = T.lane_id([32]) - A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024) - tmem_addr = T.alloc_shared([1], "uint32") - cp_mbar = T.alloc_shared([1], "uint64") - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc( - T.address_of(tmem_addr), - n_cols=n_tmem_cols_total, - cta_group=cta_group, - ) - if tid_in_wg == 0: - T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - Tx.cta.copy(A_smem[:, :], A[:, :]) - T.cuda.cta_sync() - tmem = T.decl_buffer( - t_full_shape, - dtype, - scope="tmem", - allocated_addr=tmem_addr[0], - layout=t_full, - ) - if tid_in_wg == 0: - Tx.copy_async( - tmem[t_r0:t_r1, t_c0:t_c1], - A_smem[s_r0:s_r1, s_c0:s_c1], - cta_group=cta_group, - ) - T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group) - T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) - T.cuda.cta_sync() - T.ptx.tcgen05.fence.after_thread_sync() - if warp_id == 0: - reg = T.alloc_buffer((4,), "uint32", scope="local") - for i in range(4): - T.ptx.tcgen05.ld( - tmem.allocated_addr[0], - reg[i], - shape="32x32b", - num=1, - row=0, - col=i, - ) - T.ptx.tcgen05.wait.ld() - B_bytes = reg.view(dtype) - for i in range(OUT_BYTES): - B[lane_id, i] = B_bytes[i] - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group) - - return kernel - - -def _make_3d_4tile_kernel(s_full, t_full, s_full_shape, t_full_shape, dtype, cta_group=1): - """3D variant: 4 stacked tiles (NVFP4-style multi-cp test).""" - n_tmem_cols_total = max(32, t_full_shape[-1]) - - @T.prim_func(check_well_formed=False) - def kernel(A_ptr: T.handle, B_ptr: T.handle): - A = T.match_buffer(A_ptr, s_full_shape, dtype) - B = T.match_buffer(B_ptr, (32, 16), dtype) - T.device_entry() - warp_id = T.warp_id([4]) - wg_id = T.warpgroup_id([1]) - tid_in_wg = T.thread_id_in_wg([128]) - lane_id = T.lane_id([32]) - A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024) - tmem_addr = T.alloc_shared([1], "uint32") - cp_mbar = T.alloc_shared([1], "uint64") - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc( - T.address_of(tmem_addr), - n_cols=n_tmem_cols_total, - cta_group=cta_group, - ) - if tid_in_wg == 0: - T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - Tx.cta.copy(A_smem[:, :, :], A[:, :, :]) - T.cuda.cta_sync() - tmem = T.decl_buffer( - t_full_shape, - dtype, - scope="tmem", - allocated_addr=tmem_addr[0], - layout=t_full, - ) - if tid_in_wg == 0: - Tx.copy_async( - tmem[:, :, :], - A_smem[:, :, :], - cta_group=cta_group, - ) - T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group) - T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) - T.cuda.cta_sync() - T.ptx.tcgen05.fence.after_thread_sync() - if warp_id == 0: - reg = T.alloc_buffer((4,), "uint32", scope="local") - for i in range(4): - T.ptx.tcgen05.ld( - tmem.allocated_addr[0], - reg[i], - shape="32x32b", - num=1, - row=0, - col=i, - ) - T.ptx.tcgen05.wait.ld() - B_bytes = reg.view(dtype) - for i in range(16): - B[lane_id, i] = B_bytes[i] - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group) - - return kernel - - -def _run_2d(s_full, t_full, s_full_shape, s_region, dtype, A_init, expected): - s_r0, s_r1 = s_region[0] - s_c0, s_c1 = s_region[1] - kernel = _make_2d_kernel( - s_full, t_full, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, dtype - ) - return _execute(kernel, A_init, expected) - - -def _run_3d_4tile(s_full, t_full, s_full_shape, dtype, A_init, expected): - kernel = _make_3d_4tile_kernel(s_full, t_full, s_full_shape, s_full_shape, dtype) - return _execute(kernel, A_init, expected) - - -def _execute(kernel, A_init, expected): - target = tvm.target.Target("cuda") - with target: - mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") - B_np = np.zeros((32, 16), dtype=A_init.dtype) - - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_init, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - B_out = B.numpy() - assert np.array_equal(B_out, expected), ( - f"mismatch:\nlane 0 expected={expected[0].tolist()}\n" - f" got ={B_out[0].tolist()}" - ) - - tvm.testing.run_with_gpu_lock(run_and_check) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -@pytest.mark.parametrize( - "name,s_full,s_full_shape,s_region", - [ - ("sw0_plain_atom_aligned", TileLayout(S[(32, 16) : (16, 1)]), [32, 16], [(0, 32), (0, 16)]), - ( - "sw1_32B_atom", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_32B_ATOM, [32, 32]), - [32, 32], - [(0, 32), (0, 16)], - ), - ( - "sw2_64B_atom", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 64]), - [32, 64], - [(0, 32), (0, 16)], - ), - ( - "sw3_128B_atom", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [32, 128]), - [32, 128], - [(0, 32), (0, 16)], - ), - ( - "sw3_64x128_corner", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), - [64, 128], - [(0, 32), (0, 16)], - ), - ( - "sw3_64x128_atom_row_8", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), - [64, 128], - [(8, 40), (0, 16)], - ), - ( - "sw2_32x256_col_64", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 256]), - [32, 256], - [(0, 32), (64, 80)], - ), - ( - "sw0_M_atom_major_4_0", - TileLayout(S[(8, 8, 2, 16) : (128, 16, 1024, 1)]), - [64, 32], - [(4, 36), (0, 16)], - ), - ], -) -def test_single_cp(name, s_full, s_full_shape, s_region): - A_np = np.arange(int(np.prod(s_full_shape)), dtype=np.uint8).reshape(s_full_shape) - r0, r1 = s_region[0] - c0, c1 = s_region[1] - expected = A_np[r0:r1, c0:c1] - _run_2d(s_full, T_LAY_BASIC, s_full_shape, s_region, "uint8", A_np, expected) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -def test_multi_cp_sw0_4tiles(): - s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)]) - t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) - A_np = (np.arange(4 * 32 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(4, 32, 16) - expected = A_np[0] - _run_3d_4tile(s_full, t_full, [4, 32, 16], "uint8", A_np, expected) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -def test_align_middle_2_to_1_nvfp4_sfb(): - """SFB-style nvfp4 case: TMEM mid canonicalizes to single iter - (16@TCol + 4@TCol merge), but SMEM mid stays as 2 iters - (stride 512 + stride 2048 — outer/inner reversed so canon can't merge). - Exercises ``_align_middles`` union-cut algorithm. - - Layout shapes mirror SFB nvfp4 with PIPE=1, SFB_n_chunks=2, - MMA_K_BLOCKS=4, sf_mma_k=4. - """ - # SMEM: (2, 4, 32, 4, 4) extents, strides (2048, 4, 16, 512, 1) - # — N_chunk outer (stride 2048), then sub-warp tile (4, stride 4), lane - # (32, stride 16), K_block (4, stride 512), sf_mma_k (4, stride 1). - # Mid post-canon = [(4, 512), (2, 2048)] — non-mergeable in this order. - s_full = TileLayout(S[(2, 4, 32, 4, 4) : (2048, 4, 16, 512, 1)]) - # TMEM: SFB-style 5-axis layout. K_outer (4, 4@TCol) and N_chunk - # (2, 16@TCol) merge into single mid iter (8, 4@TCol). - t_full = TileLayout( - S[(2, 4, 32, 4, 4) : (16 @ TCol, 4 @ TCol, 1 @ TLane, 32 @ TCol, 1 @ TCol)] - + R[4 : 32 @ TLane] - ) - s_full_shape = [256, 16] - t_full_shape = [256, 16] - n_tmem_cols_total = max(32, 32) # SFB occupies 32 cols total (8*4 elements / 4 epc) - - @T.prim_func(check_well_formed=False) - def kernel(A_ptr: T.handle, B_ptr: T.handle): - A = T.match_buffer(A_ptr, s_full_shape, "uint8") - B = T.match_buffer(B_ptr, (32, 16), "uint8") - T.device_entry() - warp_id = T.warp_id([4]) - wg_id = T.warpgroup_id([1]) - tid_in_wg = T.thread_id_in_wg([128]) - lane_id = T.lane_id([32]) - A_smem = T.alloc_buffer(s_full_shape, "uint8", scope="shared", layout=s_full, align=1024) - tmem_addr = T.alloc_shared([1], "uint32") - cp_mbar = T.alloc_shared([1], "uint64") - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_tmem_cols_total, cta_group=1) - if tid_in_wg == 0: - T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - Tx.cta.copy(A_smem[:, :], A[:, :]) - T.cuda.cta_sync() - tmem = T.decl_buffer( - t_full_shape, - "uint8", - scope="tmem", - allocated_addr=tmem_addr[0], - layout=t_full, - ) - if tid_in_wg == 0: - Tx.copy_async(tmem[:, :], A_smem[:, :], cta_group=1) - T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1) - T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) - T.cuda.cta_sync() - T.ptx.tcgen05.fence.after_thread_sync() - if warp_id == 0: - reg = T.alloc_buffer((4,), "uint32", scope="local") - for i in range(4): - T.ptx.tcgen05.ld( - tmem.allocated_addr[0], - reg[i], - shape="32x32b", - num=1, - row=0, - col=i, - ) - T.ptx.tcgen05.wait.ld() - B_bytes = reg.view("uint8") - for i in range(16): - B[lane_id, i] = B_bytes[i] - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=1) - - A_np = (np.arange(256 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(256, 16) - - # Compute expected: for each (lane=L in 0..32, byte b in 0..15), the - # tcgen05.ld reads physical (TLane=L, TCol=b). We must invert the TMEM - # layout to find which logical (m, k) is at that physical position, then - # expected[L, b] = A[m, k]. - # Layout shard iters (i0..i4) with extents (2, 4, 32, 4, 4) and TMEM - # strides (16, 4, 1@TLane, 32, 1) — only TLane and TCol contribute. - # For (TLane=L, TCol=p) with L in 0..32, replica r=0: - # i2 = L; remaining iters (i0, i1, i3, i4) contribute to TCol: - # p = 16*i0 + 4*i1 + 32*i3 + i4 - # For p in 0..15 only i1 and i4 vary (i0 = i3 = 0): - # i1 = p // 4, i4 = p % 4 - # Logical buffer index: rev row-major over iter coords following shard order. - # Shard order outer→inner: (i0, i1, i2, i3, i4) with extents (2, 4, 32, 4, 4). - # Logical buffer index = i0*(4*32*4*4) + i1*(32*4*4) + i2*(4*4) + i3*4 + i4 - expected = np.zeros((32, 16), dtype=np.uint8) - for L in range(32): - for p in range(16): - i0 = 0 - i3 = 0 - i1 = p // 4 - i4 = p % 4 - i2 = L - logical = i0 * (4 * 32 * 4 * 4) + i1 * (32 * 4 * 4) + i2 * (4 * 4) + i3 * 4 + i4 - m, k = divmod(logical, 16) - expected[L, p] = A_np[m, k] - - _execute(kernel, A_np, expected) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -@pytest.mark.parametrize( - "bad", - [ - pytest.param( - ( - "sw3_mid_atom_row", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), - [64, 128], - [(4, 36), (0, 16)], - ), - id="sw3_mid_atom_row", - ), - pytest.param( - ( - "sw2_mid_atom_col", - mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 128]), - [32, 128], - [(0, 32), (32, 48)], - ), - id="sw2_mid_atom_col", - ), - pytest.param( - ("sw0_row_stride_64", TileLayout(S[(64, 64) : (64, 1)]), [64, 64], [(4, 36), (0, 16)]), - id="sw0_row_stride_64", - ), - ], -) -def test_dispatch_rejects_bad_inputs(bad): - """Configurations where cp 32x128b cannot read the user's intended sub-tile. - Compilation should fail with a clear ValueError from the dispatch.""" - name, s_full, s_full_shape, s_region = bad - s_r0, s_r1 = s_region[0] - s_c0, s_c1 = s_region[1] - kernel = _make_2d_kernel( - s_full, T_LAY_BASIC, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, "uint8" - ) - with pytest.raises(Exception): - target = tvm.target.Target("cuda") - with target: - tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") - - -@pytest.mark.gpu -def test_multi_cp_encodes_descriptor_once_and_patches_addr(): - """Compile-only regression for the shared-descriptor cp path. - - A multi-tile smem->tmem copy encodes ONE SMEM matrix descriptor template at - SMEM base 0 (so the cache key no longer depends on the buffer identity) and - patches its 14-bit address field per cp via ``cvta(addr) >> 4 & 0x3FFF``, - instead of re-encoding a descriptor per tile. Verifies the 4-tile copy emits - a single ``encode_matrix_descriptor`` reused across four - ``tcgen05.cp.32x128b.warpx4`` issues, each with the address-field patch. - """ - s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)]) - t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) - kernel = _make_3d_4tile_kernel(s_full, t_full, [4, 32, 16], [4, 32, 16], "uint8") - - target = tvm.target.Target("cuda") - with target: - mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") - src = mod.mod.imports[0].inspect_source() - - assert "tcgen05.cp.cta_group::1.32x128b.warpx4" in src, f"cp not emitted; src=\n{src}" - # Descriptor encoded once (single matrix-descriptor encode call), then - # reused with a per-cp 14-bit SMEM address patch (0x3FFF == 16383 mask). - assert "16383" in src, "expected 14-bit SMEM address-field patch (0x3FFF mask)" - assert src.count("cp_desc[0] &") == 4, ( - f"expected 4 address-patched cp's reusing one cp_desc; got " - f"{src.count('cp_desc[0] &')}\nsrc=\n{src}" - ) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py new file mode 100644 index 000000000000..de914da54367 --- /dev/null +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py @@ -0,0 +1,1306 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name, missing-function-docstring +"""End-to-end tests for the ``tcgen05.cp`` smem->tmem dispatch. + +The bulk is bit-exact GPU round-trips through the generic planner; the +"legacy 32x128b.warpx4" section at the bottom keeps the original +single-shape tests (uint8 sub-region copies, NVFP4 SFB middle alignment, +shared-descriptor emission) that predate the generalization. + +For each ``(shape, multicast, swizzle, dtype, tile size)`` combination we: + +1. Fill a host buffer ``A`` with random values and stage it into shared + memory with an MMA-style (optionally swizzled) layout. +2. Issue a generic ``Tx.copy_async(tmem_region, smem_region, shape=..., + multicast=...)`` — no ``desc_*`` fields, so the generic planner derives + the matrix descriptor and cp loop. +3. Read every TMEM cell back via per-warp ``tcgen05.ld.32x32b`` (each of the + 4 warps reads its own 32-lane slab at taddr lane 0) into a + ``(128, W32)`` uint32 dump. +4. Rebuild the expected dump on the host purely from the *destination* TMEM + layout: logical element ``i`` of the copied region carries ``A``'s value + at the matching source coordinate, lands at the (TLane, TCol) position + the tmem layout assigns, and is replicated at the multicast lane offsets. + Only cells covered by the copy are compared. + +The multicast row→lane mappings this asserts (verified on B200 hardware, +PTX ISA 8.8 §9.7.16.9.2 p675 + Layout E/B figures 213/207): + +============= ============================== ==================== +multicast one copy (row → lane) replica lane offsets +============= ============================== ==================== +warpx4 rows 0-31 → lanes 0-31 +0 / +32 / +64 / +96 +warpx2::02_13 rows 0-63 → lanes 0-63 +0 / +64 +warpx2::01_23 rows 0-31 → lanes 0-31, +0 / +32 + rows 32-63 → lanes 64-95 +(4x256b) rows 0-3 → lanes 0/32/64/96 (none) +============= ============================== ==================== +""" + +import itertools + +import numpy as np +import pytest + +import tvm +import tvm.testing +from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx +from tvm.testing import env +from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode, mma_shared_layout +from tvm.tirx.layout import ComposeLayout, R, S, TCol, TileLayout, TLane + +# Multicast replica lane offsets: (extent, stride) pairs on TLane. These are +# the warp-slab offsets receiving copies of the data (see module docstring). +_REPLICA_PATTERN = { + None: [], + "warpx4": [(4, 32)], + "warpx2::02_13": [(2, 64)], + "warpx2::01_23": [(2, 32)], +} + +# The full (shape, multicast) matrix of the generic planner. +_SHAPE_MULTICAST = [ + ("128x256b", None), + ("128x128b", None), + ("64x128b", "warpx2::02_13"), + ("64x128b", "warpx2::01_23"), + ("32x128b", "warpx4"), + ("4x256b", None), +] + + +def _shape_dims(shape): + rows, bits = shape.split("x") + return int(rows), int(bits[:-1]) + + +def _tmem_layout_for(shape, multicast, C): + """Destination TMEM layout + logical buffer shape for one cp footprint. + + The layout encodes the hardware row→lane mapping of ONE copy; the + multicast replicas are declared via ``R[...]`` so the planner can verify + them (and the host expectation applies them explicitly). + """ + if shape in ("128x256b", "128x128b"): + return TileLayout(S[(128, C) : (1 @ TLane, 1 @ TCol)]), [128, C] + if shape == "4x256b": + # One row per warp-quadrant datapath: rows 0..3 at lanes 0/32/64/96. + return TileLayout(S[(4, C) : (32 @ TLane, 1 @ TCol)]), [4, C] + if shape == "32x128b": + return TileLayout(S[(32, C) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]), [32, C] + if shape == "64x128b": + if multicast == "warpx2::02_13": + # Rows 0-63 contiguous on lanes 0-63; mirrored at +64. + return ( + TileLayout(S[(64, C) : (1 @ TLane, 1 @ TCol)] + R[2 : 64 @ TLane]), + [64, C], + ) + # 01_23: rows 0-31 at lanes 0-31, rows 32-63 at lanes 64-95; +32 mirror. + # Lane split lives in the iters; buffer stays the natural (64, C) tile. + return ( + TileLayout(S[(2, 32, C) : (64 @ TLane, 1 @ TLane, 1 @ TCol)] + R[2 : 32 @ TLane]), + [64, C], + ) + raise ValueError(shape) + + +def _next_pow2(x): + return 1 if x <= 1 else 1 << (x - 1).bit_length() + + +def _make_cp_kernel( + s_full, + s_full_shape, + s_region, + t_full, + t_full_shape, + t_region, + dtype, + shape, + multicast, + W32, + n_tmem_cols, + extra_cfg=None, + infer=False, + pre_zero=False, +): + s_sl = tuple(slice(a, b) for a, b in s_region) + t_sl = tuple(slice(a, b) for a, b in t_region) + s_full_sl = tuple(slice(0, e) for e in s_full_shape) + cfg = {"cta_group": 1} if infer else {"shape": shape, "cta_group": 1} + if not infer and multicast is not None: + cfg["multicast"] = multicast + if extra_cfg: + cfg.update(extra_cfg) + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, s_full_shape, dtype) + B = T.match_buffer(B_ptr, (128, W32), "uint32") + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + T.lane_id([32]) + A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + cp_mbar = T.alloc_shared([1], "uint64") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_tmem_cols, cta_group=1) + if tid_in_wg == 0: + T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + Tx.cta.copy(A_smem[s_full_sl], A[s_full_sl]) + T.cuda.cta_sync() + tmem = T.decl_buffer( + t_full_shape, dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=t_full + ) + if pre_zero: + zero_reg = T.alloc_buffer((W32,), "uint32", scope="local") + for i in range(W32): + zero_reg[i] = T.uint32(0) + for i in range(W32): + T.ptx.tcgen05.st(tmem_addr[0], zero_reg[i], shape="32x32b", num=1, row=0, col=i) + T.ptx.tcgen05.wait.st() + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + if tid_in_wg == 0: + Tx.copy_async(tmem[t_sl], A_smem[s_sl], **cfg) + T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1) + T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + # Each of the 4 warps reads its own 32-lane slab (taddr lane 0 is + # warp-slab-relative for .32x32b), covering all 128 TMEM lanes. + reg = T.alloc_buffer((W32,), "uint32", scope="local") + for i in range(W32): + T.ptx.tcgen05.ld(tmem_addr[0], reg[i], shape="32x32b", num=1, row=0, col=i) + T.ptx.tcgen05.wait.ld() + for i in range(W32): + B[tid_in_wg, i] = reg[i] + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols, cta_group=1) + + return kernel + + +def _expected_readback(A_bits, s_region, t_full, t_full_shape, t_region, multicast, bits, W32): + """Host reconstruction of the (128, W32) uint32 TMEM dump (+ mask).""" + expected = np.zeros((128, W32), np.uint32) + mask = np.zeros((128, W32), bool) + rep_offs = [0] + for e, st in _REPLICA_PATTERN[multicast]: + rep_offs = [o + k * st for o in rep_offs for k in range(e)] + t_coords = itertools.product(*[range(a, b) for a, b in t_region]) + s_coords = itertools.product(*[range(a, b) for a, b in s_region]) + per32 = 32 // bits + for tc, sc in zip(t_coords, s_coords): + val = int(A_bits[sc]) + av = t_full.apply(*tc, shape=t_full_shape) + lane = int(av.get("TLane", 0)) + ce = int(av.get("TCol", 0)) + c32, sub = divmod(ce, per32) + for off in rep_offs: + expected[lane + off, c32] |= np.uint32(val << (sub * bits)) + mask[lane + off, c32] = True + return expected, mask + + +def _build_case( + shape, multicast, sw, dtype, n_mid, s_row_off=0, t_col_off_e=0, extra_cfg=None, infer=False +): + """Assemble (kernel, host-check closure) for one cp configuration.""" + bits = tvm.runtime.DataType(dtype).bits + rows, atom_bits = _shape_dims(shape) + epa = atom_bits // bits # elems per lane per cp instruction + C = epa * n_mid # copied cols (elems) + t_C = C + t_col_off_e # tmem buffer cols (elems) + W32 = t_C * bits // 32 + n_tmem_cols = max(32, _next_pow2(W32)) + + s_rows = max(rows + s_row_off, 8) + if sw == 0: + if atom_bits == 128: + # Contiguous rows of exactly one 16B unit: atom_K = 16B → sw0. + s_shape = [s_rows, epa] + s_full = TileLayout(S[(s_rows, epa) : (epa, 1)]) + else: + # 256b rows with sw0 need a real descriptor LDO: 16B unit (R, c) + # sits at R%8 * 16B + c * LDO + R//8 * SDO. + u = 128 // bits # elems per 16B unit + n_grp = s_rows // 8 + LDO_e = 8 * u + SDO_e = 16 * u + s_shape = [s_rows, 2 * u] + s_full = TileLayout(S[(n_grp, 8, 2, u) : (SDO_e, u, LDO_e, 1)]) + else: + atom_row_elems = (16 << sw) // (bits // 8) + s_cols = max(C, atom_row_elems) + s_shape = [s_rows, s_cols] + s_full = mma_shared_layout(dtype, SwizzleMode(sw), s_shape) + s_region = [(s_row_off, s_row_off + rows), (0, C)] + + t_full, t_full_shape = _tmem_layout_for(shape, multicast, t_C) + t_region = [(0, t_full_shape[0]), (t_col_off_e, t_C)] + + kernel = _make_cp_kernel( + s_full, + s_shape, + s_region, + t_full, + t_full_shape, + t_region, + dtype, + shape, + multicast, + W32, + n_tmem_cols, + extra_cfg=extra_cfg, + infer=infer, + ) + + def check(mod): + A_np = tvm.testing.generate_random_array(dtype, tuple(s_shape)) + dev = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, dev) + B = tvm.runtime.tensor(np.zeros((128, W32), np.uint32), dev) + mod(A, B) + B_out = B.numpy() + A_bits = np.asarray(A_np).view(np.uint16 if bits == 16 else np.uint32) + exp, mask = _expected_readback( + A_bits, s_region, t_full, t_full_shape, t_region, multicast, bits, W32 + ) + np.testing.assert_array_equal(np.where(mask, B_out, 0), np.where(mask, exp, 0)) + + return kernel, check + + +def _compile(kernel): + target = tvm.target.Target("cuda") + with target: + return tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") + + +def _run_case(shape, multicast, sw, dtype, n_mid, s_row_off=0, t_col_off_e=0): + kernel, check = _build_case( + shape, multicast, sw, dtype, n_mid, s_row_off=s_row_off, t_col_off_e=t_col_off_e + ) + check(_compile(kernel)) + + +# GPU round-trip: full shape x multicast x swizzle x dtype x size matrix + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize("shape,multicast", _SHAPE_MULTICAST) +@pytest.mark.parametrize("sw", [1, 2, 3], ids=["SW32", "SW64", "SW128"]) +@pytest.mark.parametrize("dtype", ["bfloat16", "float32"]) +@pytest.mark.parametrize("n_mid", [1, 4], ids=["single_cp", "multi_cp"]) +def test_cp_shape_roundtrip_swizzled(shape, multicast, sw, dtype, n_mid): + """Bit-exact smem→tmem round-trip through the generic planner.""" + _run_case(shape, multicast, sw, dtype, n_mid) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize( + "shape,multicast,dtype", + [ + ("128x128b", None, "bfloat16"), + ("128x256b", None, "bfloat16"), # exercises the derived non-trivial LDO + ("64x128b", "warpx2::02_13", "bfloat16"), + ("64x128b", "warpx2::01_23", "float32"), + ("32x128b", "warpx4", "float32"), + ("4x256b", None, "float32"), + ], +) +def test_cp_shape_roundtrip_nonswizzled(shape, multicast, dtype): + """sw=0 sources: 128b rows are single 16B units; 256b rows carry LDO.""" + _run_case(shape, multicast, 0, dtype, 1) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize( + "shape,multicast,sw,dtype,n_mid,s_row_off,t_col_off_e", + [ + ("32x128b", "warpx4", 3, "bfloat16", 1, 8, 0), # smem row offset + ("64x128b", "warpx2::02_13", 3, "bfloat16", 2, 0, 8), # tmem col offset + ("128x256b", None, 2, "float32", 2, 0, 8), + ], +) +def test_cp_shape_roundtrip_offsets(shape, multicast, sw, dtype, n_mid, s_row_off, t_col_off_e): + """Sub-region copies: non-zero smem row offset / TMEM column offset.""" + _run_case(shape, multicast, sw, dtype, n_mid, s_row_off=s_row_off, t_col_off_e=t_col_off_e) + + +# Compile-level checks + + +def test_cp_4x256b_compile_emits_shape_and_count(): + """4x256b, 2 middle iterations → exactly 2 cp instructions of that shape.""" + kernel, _ = _build_case("4x256b", None, 1, "bfloat16", 2) + mod = _compile(kernel) + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.cp.cta_group::1.4x256b" in src, f"cp asm missing; src=\n{src}" + helper_refs = src.count("ptx_tcgen05_cp_cta_group_1_shape_4x256b") + assert helper_refs - 1 == 2, f"expected 2 cp calls, got {helper_refs - 1}; src=\n{src}" + + +def test_cp_shape_config_routes_to_generic_planner(): + """``shape=`` without desc_* must use the generic planner, not the + explicit path: one hoisted descriptor encoded at SMEM base 0 with the + 14-bit address patch per cp.""" + kernel, _ = _build_case("128x256b", None, 3, "bfloat16", 4) + mod = _compile(kernel) + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.cp.cta_group::1.128x256b" in src + # Generic-planner signature: descriptor template encoded at address 0... + assert "reinterpret_cast((uint64_t)0)" in src + # ...and patched per cp via the 0x3FFF address-field mask. + assert src.count("cp_desc[0] &") == 4, f"expected 4 patched cps; src=\n{src}" + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize("shape,multicast", [sm for sm in _SHAPE_MULTICAST if sm[0] != "128x128b"]) +def test_cp_shape_inferred_from_layouts_matches_explicit(shape, multicast): + """A bare copy_async infers (shape, multicast) from the buffer layouts and + must lower byte-identically to the explicitly configured call. (128x128b + is excluded: its layouts also fit the wider 128x256b atom, which inference + prefers — covered by the dedicated test below.)""" + explicit, _ = _build_case(shape, multicast, 2, "bfloat16", 2) + inferred, _ = _build_case(shape, multicast, 2, "bfloat16", 2, infer=True) + src_explicit = _compile(explicit).mod.imports[0].inspect_source() + src_inferred = _compile(inferred).mod.imports[0].inspect_source() + assert src_explicit == src_inferred + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_cp_128x128b_layout_infers_wider_256b_atom(): + """A 128x128b-compatible source also fits the 128x256b atom at half the + instruction count; inference must pick the wider atom and stay bit-exact.""" + inferred, check = _build_case("128x128b", None, 2, "bfloat16", 2, infer=True) + mod = _compile(inferred) + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.cp.cta_group::1.128x256b" in src, f"expected wider atom; src=\n{src}" + check(mod) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_cp_4x256b_lane_tiled_partial_window(): + """A region row offset folds into the taddr lane half-word: after + pre-zeroing tmem, copying into rows [4:12) of the (16, 4, C) Layout-F + scatter view lands on EXACTLY lanes {32q + 4..11} — those match the smem + rows bit-for-bit and every other lane reads back all-zero.""" + dtype, bits = "bfloat16", 16 + C = 16 + W32 = C * bits // 32 + lo, hi = 4, 12 + s_full = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_32B_ATOM, [64, C]) + t_full = TileLayout(S[(16, 4, C) : (1 @ TLane, 32 @ TLane, 1 @ TCol)]) + kernel = _make_cp_kernel( + s_full, + [64, C], + [(0, (hi - lo) * 4), (0, C)], + t_full, + [64, C], + [(lo * 4, hi * 4), (0, C)], + dtype, + None, + None, + W32, + 32, + infer=True, + pre_zero=True, + ) + mod = _compile(kernel) + A_np = tvm.testing.generate_random_array(dtype, (64, C)) + dev = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, dev) + B = tvm.runtime.tensor(np.zeros((128, W32), np.uint32), dev) + mod(A, B) + B_u16 = B.numpy().view(np.uint16).reshape(128, W32 * 2) + A_u16 = np.asarray(A_np).view(np.uint16) + written = {32 * q + lo + il for q in range(4) for il in range(hi - lo)} + for lane in range(128): + if lane in written: + q, il = lane // 32, lane % 32 - lo + np.testing.assert_array_equal( + B_u16[lane, :C], A_u16[il * 4 + q], err_msg=f"lane {lane}" + ) + else: + assert not B_u16[lane].any(), f"lane {lane} must stay zero" + + +def _make_cp_kernel_cta2(s_full, s_shape, t_full, t_shape, dtype, cfg, W32, n_cols): + """2-CTA cluster harness: each CTA stages its own A[cbx] slice into smem; + only the even CTA issues the cp (cta_group=2, commit cta_mask=3); both CTAs + dump their own tmem into B[cbx*128 + lane].""" + s_sl = tuple(slice(0, e) for e in s_shape) + t_sl = tuple(slice(0, e) for e in t_shape) + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, (2, *s_shape), dtype) + B = T.match_buffer(B_ptr, (256, W32), "uint32") + T.device_entry() + warp_id = T.warp_id([4]) + cbx, cby = T.cta_id_in_cluster([2, 1]) + cta_id = T.cta_id([2]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer(s_shape, dtype, scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + cp_mbar = T.alloc_shared([1], "uint64") + if tid_in_wg == 0: + T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_cols, cta_group=2) + tmem = T.decl_buffer( + t_shape, dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=t_full + ) + T.ptx.fence.mbarrier_init() + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + T.cuda.cluster_sync() + # Pre-zero both CTAs' tmem: alloc does not clear it, and the + # 128x256b test asserts the odd CTA stays untouched. + zero_reg = T.alloc_buffer((W32,), "uint32", scope="local") + for i in range(W32): + zero_reg[i] = T.uint32(0) + for i in range(W32): + T.ptx.tcgen05.st(tmem_addr[0], zero_reg[i], shape="32x32b", num=1, row=0, col=i) + T.ptx.tcgen05.wait.st() + Tx.cta.copy(A_smem[s_sl], A[(cbx, *s_sl)]) + T.cuda.cta_sync() + T.cuda.cluster_sync() + if cbx == 0: + if tid_in_wg == 0: + Tx.copy_async(tmem[t_sl], A_smem[s_sl], **cfg) + T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=2, cta_mask=3) + T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + reg = T.alloc_buffer((W32,), "uint32", scope="local") + for i in range(W32): + T.ptx.tcgen05.ld(tmem_addr[0], reg[i], shape="32x32b", num=1, row=0, col=i) + T.ptx.tcgen05.wait.ld() + for i in range(W32): + B[cbx * 128 + tid_in_wg, i] = reg[i] + T.cuda.cluster_sync() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=2) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_cols, cta_group=2) + + return kernel + + +def _run_cta2(s_full, s_shape, t_full, t_shape, cfg, W32): + kernel = _make_cp_kernel_cta2(s_full, s_shape, t_full, t_shape, "bfloat16", cfg, W32, 32) + mod = _compile(kernel) + A_np = tvm.testing.generate_random_array("bfloat16", (2, *s_shape)) + dev = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, dev) + B = tvm.runtime.tensor(np.zeros((256, W32), np.uint32), dev) + mod(A, B) + B_u16 = B.numpy().view(np.uint16).reshape(2, 128, W32 * 2) + A_u16 = np.asarray(A_np).view(np.uint16) + return A_u16, B_u16 + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_cp_cta_group_2_128x256b_pair_collective(): + """Production small_topk q-cp form: 128x256b + cta_group=2. One even-CTA + issue makes EACH CTA copy from its own smem into its own tmem (descriptor + smem address resolves per CTA rank). B200-pinned.""" + C = 16 + s_full = mma_shared_layout("bfloat16", SwizzleMode.SWIZZLE_32B_ATOM, [128, C]) + t_full = TileLayout(S[(128, C) : (1 @ TLane, 1 @ TCol)]) + A, B = _run_cta2( + s_full, [128, C], t_full, [128, C], {"shape": "128x256b", "cta_group": 2}, C * 16 // 32 + ) + for cta in range(2): + for lane in range(128): + np.testing.assert_array_equal( + B[cta, lane, :C], A[cta, lane], err_msg=f"cta{cta} lane {lane}" + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_cp_cta_group_2_64x128b_warpx2_pair_collective(): + """Production head128 q-cp form: 64x128b.warpx2::02_13 + cta_group=2 is a + pair-collective — one even-CTA issue makes EACH CTA copy from its own smem + into its own tmem (descriptor smem address resolves per CTA rank, like the + cta2 MMA B descriptor), with the warpx2 mirror at lane +64. B200-pinned.""" + C = 8 + s_full = TileLayout(S[(64, C) : (C, 1)]) + t_full = TileLayout(S[(64, C) : (1 @ TLane, 1 @ TCol)] + R[2 : 64 @ TLane]) + A, B = _run_cta2( + s_full, + [64, C], + t_full, + [64, C], + {"shape": "64x128b", "multicast": "warpx2::02_13", "cta_group": 2}, + C * 16 // 32, + ) + for cta in range(2): + for lane in range(128): + np.testing.assert_array_equal( + B[cta, lane, :C], A[cta, lane % 64], err_msg=f"cta{cta} lane {lane}" + ) + + +def test_cp_4x256b_lane_tiled_layout_f_scatter(): + """4x256b atoms tiled through the taddr lane half-word: 16 cps land 16 + distinct rows on each warp's first 16 lanes (the M=64 Layout-F scatter). + Smem holds the rows quadrant-interleaved (row i*4+q = logical (q, i)).""" + dtype, bits = "bfloat16", 16 + C = 16 # one 256b atom row = 16 bf16 elements + W32 = C * bits // 32 + s_full = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_32B_ATOM, [64, C]) + t_full = TileLayout(S[(16, 4, C) : (1 @ TLane, 32 @ TLane, 1 @ TCol)]) + kernel = _make_cp_kernel( + s_full, + [64, C], + [(0, 64), (0, C)], + t_full, + [64, C], + [(0, 64), (0, C)], + dtype, + None, + None, + W32, + 32, + infer=True, + ) + mod = _compile(kernel) + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.cp.cta_group::1.4x256b" in src, f"expected 4x256b atoms; src=\n{src}" + + A_np = tvm.testing.generate_random_array(dtype, (64, C)) + dev = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, dev) + B = tvm.runtime.tensor(np.zeros((128, W32), np.uint32), dev) + mod(A, B) + B_out = B.numpy() + A_bits = np.asarray(A_np).view(np.uint16) + for i in range(16): + for q in range(4): + lane = 32 * q + i + exp = A_bits[i * 4 + q].copy().view(np.uint32) + np.testing.assert_array_equal(B_out[lane, :W32], exp, err_msg=f"lane {lane}") + + +def test_cp_default_32x128b_instruction_sequence_unchanged(): + """Back-compat pin: a config-less smem->tmem copy_async must emit the + exact legacy 32x128b.warpx4 sequence (hardcoded from the pre-generalization + planner output): one descriptor encode with (ldo=0, sdo=8, swizzle=0) and + four cps at tmem cols 0/4/8/12 reading smem bytes 0/512/1024/1536.""" + s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)]) + t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (4, 32, 16), "uint8") + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + T.lane_id([32]) + A_smem = T.alloc_buffer((4, 32, 16), "uint8", scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=32, cta_group=1) + T.cuda.cta_sync() + Tx.cta.copy(A_smem[:, :, :], A[:, :, :]) + T.cuda.cta_sync() + tmem = T.decl_buffer( + (4, 32, 16), "uint8", scope="tmem", allocated_addr=tmem_addr[0], layout=t_full + ) + if tid_in_wg == 0: + # NOTE: no shape/multicast config — the legacy default route. + Tx.copy_async(tmem[:, :, :], A_smem[:, :, :], cta_group=1) + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=32, cta_group=1) + + mod = _compile(kernel) + src = mod.mod.imports[0].inspect_source() + + encode_lines = [ + line + for line in src.splitlines() + if "tvm_builtin_ptx_tcgen05_encode_matrix_descriptor(" in line + and "__forceinline__" not in line + ] + assert len(encode_lines) == 1, encode_lines + # (ldo, sdo, swizzle) = (0, 8, 0), template encoded at smem address 0. + assert "reinterpret_cast((uint64_t)0), 0, 8, 0);" in encode_lines[0], encode_lines[0] + + cp_lines = [ + line + for line in src.splitlines() + if "ptx_tcgen05_cp_cta_group_1_shape_32x128b_multicast_warpx4_decompress_(" in line + and "__forceinline__" not in line + ] + assert len(cp_lines) == 4, f"expected 4 cp calls, got {len(cp_lines)}" + for i, (t_off, s_off) in enumerate([(0, 0), (4, 512), (8, 1024), (12, 1536)]): + t_tok = "[0], 0, 0," if t_off == 0 else f"[0] + (uint){t_off}), 0, 0," + assert t_tok in cp_lines[i], f"cp[{i}] tmem col: want {t_tok!r} in {cp_lines[i]!r}" + s_tok = f"+ {s_off}))" + assert s_tok in cp_lines[i], f"cp[{i}] smem byte off: want {s_tok!r} in {cp_lines[i]!r}" + assert "cp_desc[0] &" in cp_lines[i] and "16383" in cp_lines[i], cp_lines[i] + + +# Negative tests: readable ValueErrors from the planner + + +def _assert_compile_raises(kernel, fragment): + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match=fragment): + tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") + + +def test_cp_rejects_wrong_replica_for_multicast(): + """A 02_13-declared TMEM layout used with 01_23 (and vice versa) must + fail the replica router with a readable message.""" + bits = 16 + C = 128 // bits + t_full, t_full_shape = _tmem_layout_for("64x128b", "warpx2::02_13", C) + s_full = mma_shared_layout("bfloat16", SwizzleMode(3), [64, 64]) + kernel = _make_cp_kernel( + s_full, + [64, 64], + [(0, 64), (0, C)], + t_full, + t_full_shape, + [(0, 64), (0, C)], + "bfloat16", + "64x128b", + "warpx2::01_23", # wrong variant for the declared replica structure + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "replica mismatch") + + +def test_cp_rejects_illegal_shape_multicast_combo(): + kernel, _ = _build_case("128x256b", None, 3, "bfloat16", 1, extra_cfg={"multicast": "warpx4"}) + _assert_compile_raises(kernel, "illegal multicast") + + +def test_cp_rejects_missing_multicast_for_64x128b(): + bits = 16 + C = 128 // bits + t_full, t_full_shape = _tmem_layout_for("64x128b", "warpx2::02_13", C) + s_full = mma_shared_layout("bfloat16", SwizzleMode(3), [64, 64]) + kernel = _make_cp_kernel( + s_full, + [64, 64], + [(0, 64), (0, C)], + t_full, + t_full_shape, + [(0, 64), (0, C)], + "bfloat16", + "64x128b", + None, # 64x128b has two legal multicasts: must be explicit + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "requires an explicit multicast") + + +def test_cp_rejects_non_128b_tcol_offset(): + """A whole 32-bit TMEM cell is still below tcgen05.cp's 128-bit alignment.""" + bits = 16 + epa = 128 // bits + C = 2 * epa + t_C = C + 8 # headroom so the misaligned region stays in range + t_full, t_full_shape = _tmem_layout_for("32x128b", "warpx4", t_C) + s_full = mma_shared_layout("bfloat16", SwizzleMode(3), [32, 64]) + kernel = _make_cp_kernel( + s_full, + [32, 64], + [(0, 32), (0, C)], + t_full, + t_full_shape, + [(0, 32), (2, C + 2)], # TCol offset 2 bf16 elems = one 32-bit cell + "bfloat16", + "32x128b", + "warpx4", + t_C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "not provably 128b-aligned") + + +def test_cp_rejects_unknown_shape(): + # Valid 32x128b setup, but the shape string is overridden with a bogus one. + kernel, _ = _build_case("32x128b", None, 3, "bfloat16", 1, extra_cfg={"shape": "96x128b"}) + _assert_compile_raises(kernel, "unknown tcgen05.cp shape") + + +def test_cp_rejects_nonzero_tmem_lane_offset(): + """A lane offset whose footprint overflows lane space must be rejected: + warpx2's +64 mirror spans all 128 lanes, so any nonzero offset would + write past lane 127.""" + bits = 16 + C = 128 // bits + t_full = TileLayout(S[(128, C) : (1 @ TLane, 1 @ TCol)] + R[2 : 64 @ TLane]) + s_full = mma_shared_layout("bfloat16", SwizzleMode(3), [64, 64]) + kernel = _make_cp_kernel( + s_full, + [64, 64], + [(0, 64), (0, C)], + t_full, + [128, C], + [(64, 128), (0, C)], # lane offset 64 + "bfloat16", + "64x128b", + "warpx2::02_13", + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "overflows the 128-lane space") + + +def test_cp_rejects_decompress_on_generic_path(): + """decompress is a real PTX feature the planner cannot lower; it must be + rejected loudly instead of silently copying without decompression.""" + kernel, _ = _build_case( + "128x256b", None, 3, "bfloat16", 1, extra_cfg={"decompress": "b8x16.b6x16_p32"} + ) + _assert_compile_raises(kernel, "does not support decompress") + + +def test_cp_rejects_non_16b_aligned_row_group_stride(): + """The descriptor SBO field is encoded in 16B units: an smem layout whose + 8-row-group stride is not 16B-divisible must be rejected.""" + bits = 16 + C = 128 // bits # 8 bf16 = one 16B unit per row + # Non-swizzled (32, 8) tile: 8-row blocks are atom_K=8 elems (16B) apart, + # but the group stride is 100 elems = 200B, not a multiple of 16B. + s_full = TileLayout(S[(4, 8, C) : (100, C, 1)]) + t_full = TileLayout(S[(32, C) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) + kernel = _make_cp_kernel( + s_full, + [32, C], + [(0, 32), (0, C)], + t_full, + [32, C], + [(0, 32), (0, C)], + "bfloat16", + "32x128b", + "warpx4", + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "not 16B-aligned") + + +def test_cp_rejects_non_canonical_swizzle_family(): + """A swizzled layout outside the canonical mma atom family (wrong + per_element for the dtype) must be rejected.""" + bits = 16 + C = 128 // bits + # per_element=2 is the f32 family; with a bf16 source it disagrees with the + # descriptor permutation. Linear part is valid, so only family check rejects. + s_full = ComposeLayout(2, 3, 3, TileLayout(S[(16, 8, 64) : (512, 64, 1)])) + t_full = TileLayout(S[(128, C) : (1 @ TLane, 1 @ TCol)]) + kernel = _make_cp_kernel( + s_full, + [128, 64], + [(0, 128), (0, C)], + t_full, + [128, C], + [(0, 128), (0, C)], + "bfloat16", + "128x128b", + None, + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "swizzle family mismatch") + + +def test_cp_rejects_flipped_swizzle_inner(): + """A canonical-family swizzled layout with ``swizzle_inner=False`` must be + rejected: the descriptor's hardware walk implements the + ``swizzle_inner=True`` permutation ``x ^ ((x & outer_mask) >> atom_len)`` + (pinned bit-exactly on B200 by the round-trip tests above), while + ``swizzle_inner=False`` selects the mirrored + ``x ^ ((x & inner_mask) << atom_len)`` — a different permutation on any + chunk whose swizzle bits are nonzero.""" + bits = 16 + C = 128 // bits + # Same linear tiling and correct (per_element=3, atom_len=3) bf16 family; + # only the permutation direction is flipped, so only swizzle_inner rejects. + s_full = ComposeLayout(3, 3, 3, TileLayout(S[(16, 8, 64) : (512, 64, 1)]), swizzle_inner=False) + t_full = TileLayout(S[(128, C) : (1 @ TLane, 1 @ TCol)]) + kernel = _make_cp_kernel( + s_full, + [128, 64], + [(0, 128), (0, C)], + t_full, + [128, C], + [(0, 128), (0, C)], + "bfloat16", + "128x128b", + None, + C * bits // 32, + 32, + ) + _assert_compile_raises(kernel, "swizzle_inner") + + +# Legacy 32x128b.warpx4 tests (predate the generic planner) + +# warpx4 requires the t buffer to declare the broadcast explicitly: +# R[4 : 32@TLane] — t.shape[lane] = 32 with replica 4 → 128 physical lanes. +T_LAY_BASIC = TileLayout(S[(32, 16) : (1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) + + +def _make_2d_kernel( + s_full, + t_full, + s_full_shape, + t_full_shape, + s_r0, + s_r1, + s_c0, + s_c1, + t_r0, + t_r1, + t_c0, + t_c1, + dtype, + cta_group=1, +): + """2D variant: SMEM/TMEM are both 2D; copy a rectangular sub-region.""" + n_tmem_cols_total = max(32, t_full_shape[-1]) + OUT_LANES = 32 + OUT_BYTES = 16 + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, s_full_shape, dtype) + B = T.match_buffer(B_ptr, (OUT_LANES, OUT_BYTES), dtype) + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + lane_id = T.lane_id([32]) + A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + cp_mbar = T.alloc_shared([1], "uint64") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc( + T.address_of(tmem_addr), + n_cols=n_tmem_cols_total, + cta_group=cta_group, + ) + if tid_in_wg == 0: + T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + Tx.cta.copy(A_smem[:, :], A[:, :]) + T.cuda.cta_sync() + tmem = T.decl_buffer( + t_full_shape, + dtype, + scope="tmem", + allocated_addr=tmem_addr[0], + layout=t_full, + ) + if tid_in_wg == 0: + Tx.copy_async( + tmem[t_r0:t_r1, t_c0:t_c1], + A_smem[s_r0:s_r1, s_c0:s_c1], + cta_group=cta_group, + ) + T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group) + T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + if warp_id == 0: + reg = T.alloc_buffer((4,), "uint32", scope="local") + for i in range(4): + T.ptx.tcgen05.ld( + tmem.allocated_addr[0], + reg[i], + shape="32x32b", + num=1, + row=0, + col=i, + ) + T.ptx.tcgen05.wait.ld() + B_bytes = reg.view(dtype) + for i in range(OUT_BYTES): + B[lane_id, i] = B_bytes[i] + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group) + + return kernel + + +def _make_3d_4tile_kernel(s_full, t_full, s_full_shape, t_full_shape, dtype, cta_group=1): + """3D variant: 4 stacked tiles (NVFP4-style multi-cp test).""" + n_tmem_cols_total = max(32, t_full_shape[-1]) + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, s_full_shape, dtype) + B = T.match_buffer(B_ptr, (32, 16), dtype) + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + lane_id = T.lane_id([32]) + A_smem = T.alloc_buffer(s_full_shape, dtype, scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + cp_mbar = T.alloc_shared([1], "uint64") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc( + T.address_of(tmem_addr), + n_cols=n_tmem_cols_total, + cta_group=cta_group, + ) + if tid_in_wg == 0: + T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + Tx.cta.copy(A_smem[:, :, :], A[:, :, :]) + T.cuda.cta_sync() + tmem = T.decl_buffer( + t_full_shape, + dtype, + scope="tmem", + allocated_addr=tmem_addr[0], + layout=t_full, + ) + if tid_in_wg == 0: + Tx.copy_async( + tmem[:, :, :], + A_smem[:, :, :], + cta_group=cta_group, + ) + T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=cta_group) + T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + if warp_id == 0: + reg = T.alloc_buffer((4,), "uint32", scope="local") + for i in range(4): + T.ptx.tcgen05.ld( + tmem.allocated_addr[0], + reg[i], + shape="32x32b", + num=1, + row=0, + col=i, + ) + T.ptx.tcgen05.wait.ld() + B_bytes = reg.view(dtype) + for i in range(16): + B[lane_id, i] = B_bytes[i] + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=cta_group) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=cta_group) + + return kernel + + +def _run_2d(s_full, t_full, s_full_shape, s_region, dtype, A_init, expected): + s_r0, s_r1 = s_region[0] + s_c0, s_c1 = s_region[1] + kernel = _make_2d_kernel( + s_full, t_full, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, dtype + ) + return _execute(kernel, A_init, expected) + + +def _run_3d_4tile(s_full, t_full, s_full_shape, dtype, A_init, expected): + kernel = _make_3d_4tile_kernel(s_full, t_full, s_full_shape, s_full_shape, dtype) + return _execute(kernel, A_init, expected) + + +def _execute(kernel, A_init, expected): + mod = _compile(kernel) + dev = tvm.cuda(0) + A = tvm.runtime.tensor(A_init, dev) + B_np = np.zeros((32, 16), dtype=A_init.dtype) + B = tvm.runtime.tensor(B_np, dev) + mod(A, B) + B_out = B.numpy() + assert np.array_equal(B_out, expected), ( + f"mismatch:\nlane 0 expected={expected[0].tolist()}\n got ={B_out[0].tolist()}" + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize( + "name,s_full,s_full_shape,s_region", + [ + ("sw0_plain_atom_aligned", TileLayout(S[(32, 16) : (16, 1)]), [32, 16], [(0, 32), (0, 16)]), + ( + "sw1_32B_atom", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_32B_ATOM, [32, 32]), + [32, 32], + [(0, 32), (0, 16)], + ), + ( + "sw2_64B_atom", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 64]), + [32, 64], + [(0, 32), (0, 16)], + ), + ( + "sw3_128B_atom", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [32, 128]), + [32, 128], + [(0, 32), (0, 16)], + ), + ( + "sw3_64x128_corner", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), + [64, 128], + [(0, 32), (0, 16)], + ), + ( + "sw3_64x128_atom_row_8", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), + [64, 128], + [(8, 40), (0, 16)], + ), + ( + "sw2_32x256_col_64", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 256]), + [32, 256], + [(0, 32), (64, 80)], + ), + ( + "sw0_M_atom_major_4_0", + TileLayout(S[(8, 8, 2, 16) : (128, 16, 1024, 1)]), + [64, 32], + [(4, 36), (0, 16)], + ), + ], +) +def test_single_cp(name, s_full, s_full_shape, s_region): + A_np = np.arange(int(np.prod(s_full_shape)), dtype=np.uint8).reshape(s_full_shape) + r0, r1 = s_region[0] + c0, c1 = s_region[1] + expected = A_np[r0:r1, c0:c1] + _run_2d(s_full, T_LAY_BASIC, s_full_shape, s_region, "uint8", A_np, expected) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_multi_cp_sw0_4tiles(): + s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)]) + t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) + A_np = (np.arange(4 * 32 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(4, 32, 16) + expected = A_np[0] + _run_3d_4tile(s_full, t_full, [4, 32, 16], "uint8", A_np, expected) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_align_middle_2_to_1_nvfp4_sfb(): + """SFB-style nvfp4 case: TMEM mid canonicalizes to single iter + (16@TCol + 4@TCol merge), but SMEM mid stays as 2 iters + (stride 512 + stride 2048 — outer/inner reversed so canon can't merge). + Exercises ``_align_middles`` union-cut algorithm. + + Layout shapes mirror SFB nvfp4 with PIPE=1, SFB_n_chunks=2, + MMA_K_BLOCKS=4, sf_mma_k=4. + """ + # SMEM: (2, 4, 32, 4, 4) extents, strides (2048, 4, 16, 512, 1). + # Mid post-canon = [(4, 512), (2, 2048)] — non-mergeable in this order. + s_full = TileLayout(S[(2, 4, 32, 4, 4) : (2048, 4, 16, 512, 1)]) + # TMEM: SFB-style 5-axis layout. K_outer (4, 4@TCol) and N_chunk + # (2, 16@TCol) merge into single mid iter (8, 4@TCol). + t_full = TileLayout( + S[(2, 4, 32, 4, 4) : (16 @ TCol, 4 @ TCol, 1 @ TLane, 32 @ TCol, 1 @ TCol)] + + R[4 : 32 @ TLane] + ) + s_full_shape = [256, 16] + t_full_shape = [256, 16] + n_tmem_cols_total = max(32, 32) # SFB occupies 32 cols total (8*4 elements / 4 epc) + + @T.prim_func(check_well_formed=False) + def kernel(A_ptr: T.handle, B_ptr: T.handle): + A = T.match_buffer(A_ptr, s_full_shape, "uint8") + B = T.match_buffer(B_ptr, (32, 16), "uint8") + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + lane_id = T.lane_id([32]) + A_smem = T.alloc_buffer(s_full_shape, "uint8", scope="shared", layout=s_full, align=1024) + tmem_addr = T.alloc_shared([1], "uint32") + cp_mbar = T.alloc_shared([1], "uint64") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=n_tmem_cols_total, cta_group=1) + if tid_in_wg == 0: + T.ptx.mbarrier.init(cp_mbar.ptr_to([0]), 1) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + Tx.cta.copy(A_smem[:, :], A[:, :]) + T.cuda.cta_sync() + tmem = T.decl_buffer( + t_full_shape, + "uint8", + scope="tmem", + allocated_addr=tmem_addr[0], + layout=t_full, + ) + if tid_in_wg == 0: + Tx.copy_async(tmem[:, :], A_smem[:, :], cta_group=1) + T.ptx.tcgen05.commit(cp_mbar.ptr_to([0]), cta_group=1) + T.ptx.mbarrier.try_wait(cp_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + if warp_id == 0: + reg = T.alloc_buffer((4,), "uint32", scope="local") + for i in range(4): + T.ptx.tcgen05.ld( + tmem.allocated_addr[0], + reg[i], + shape="32x32b", + num=1, + row=0, + col=i, + ) + T.ptx.tcgen05.wait.ld() + B_bytes = reg.view("uint8") + for i in range(16): + B[lane_id, i] = B_bytes[i] + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=n_tmem_cols_total, cta_group=1) + + A_np = (np.arange(256 * 16, dtype=np.int32) & 0xFF).astype(np.uint8).reshape(256, 16) + + # Invert TMEM layout to map physical (TLane=L, TCol=p) → logical index, then + # expected[L, b] = A[m, k]. i2 = L; TCol p: i1 = p // 4, i4 = p % 4 (i0=i3=0). + expected = np.zeros((32, 16), dtype=np.uint8) + for L in range(32): + for p in range(16): + i0 = 0 + i3 = 0 + i1 = p // 4 + i4 = p % 4 + i2 = L + logical = i0 * (4 * 32 * 4 * 4) + i1 * (32 * 4 * 4) + i2 * (4 * 4) + i3 * 4 + i4 + m, k = divmod(logical, 16) + expected[L, p] = A_np[m, k] + + _execute(kernel, A_np, expected) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize( + "bad", + [ + pytest.param( + ( + "sw3_mid_atom_row", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_128B_ATOM, [64, 128]), + [64, 128], + [(4, 36), (0, 16)], + ), + id="sw3_mid_atom_row", + ), + pytest.param( + ( + "sw2_mid_atom_col", + mma_shared_layout("uint8", SwizzleMode.SWIZZLE_64B_ATOM, [32, 128]), + [32, 128], + [(0, 32), (32, 48)], + ), + id="sw2_mid_atom_col", + ), + pytest.param( + ("sw0_row_stride_64", TileLayout(S[(64, 64) : (64, 1)]), [64, 64], [(4, 36), (0, 16)]), + id="sw0_row_stride_64", + ), + ], +) +def test_dispatch_rejects_bad_inputs(bad): + """Configurations where cp 32x128b cannot read the user's intended sub-tile. + Compilation should fail with a clear ValueError from the dispatch.""" + name, s_full, s_full_shape, s_region = bad + s_r0, s_r1 = s_region[0] + s_c0, s_c1 = s_region[1] + kernel = _make_2d_kernel( + s_full, T_LAY_BASIC, s_full_shape, [32, 16], s_r0, s_r1, s_c0, s_c1, 0, 32, 0, 16, "uint8" + ) + with pytest.raises(Exception): + _compile(kernel) + + +def test_multi_cp_encodes_descriptor_once_and_patches_addr(): + """Compile-only regression for the shared-descriptor cp path. + + A multi-tile smem->tmem copy encodes ONE SMEM matrix descriptor template at + SMEM base 0 (so the cache key no longer depends on the buffer identity) and + patches its 14-bit address field per cp via ``cvta(addr) >> 4 & 0x3FFF``, + instead of re-encoding a descriptor per tile. Verifies the 4-tile copy emits + a single ``encode_matrix_descriptor`` reused across four + ``tcgen05.cp.32x128b.warpx4`` issues, each with the address-field patch. + """ + s_full = TileLayout(S[(4, 32, 16) : (512, 16, 1)]) + t_full = TileLayout(S[(4, 32, 16) : (16 @ TCol, 1 @ TLane, 1 @ TCol)] + R[4 : 32 @ TLane]) + kernel = _make_3d_4tile_kernel(s_full, t_full, [4, 32, 16], [4, 32, 16], "uint8") + + mod = _compile(kernel) + src = mod.mod.imports[0].inspect_source() + + assert "tcgen05.cp.cta_group::1.32x128b.warpx4" in src, f"cp not emitted; src=\n{src}" + # Descriptor encoded once (single matrix-descriptor encode call), then + # reused with a per-cp 14-bit SMEM address patch (0x3FFF == 16383 mask). + assert "16383" in src, "expected 14-bit SMEM address-field patch (0x3FFF mask)" + assert src.count("cp_desc[0] &") == 4, ( + f"expected 4 address-patched cp's reusing one cp_desc; got " + f"{src.count('cp_desc[0] &')}\nsrc=\n{src}" + ) + + +if __name__ == "__main__": + tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py similarity index 77% rename from tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py rename to tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py index ac35f4c62869..57b3c9707724 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem_16xnb.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_ldst.py @@ -15,9 +15,11 @@ # specific language governing permissions and limitations # under the License. # pylint: disable=invalid-name, missing-function-docstring -"""Bit-exact tests for the ``.16x{64,128,256}b`` ``tcgen05.{ld,st}`` dispatch. +"""Tests for the tcgen05 TMEM load/store (ld/st) dispatch. -For each ``(shape, rep, dtype, direction)`` we: +Covers ``.32x32b`` tmem<->reg / smem<->tmem copy_async and bit-exact +``.16x{64,128,256}b`` atom verification. For each ``(shape, rep, dtype, +direction)`` in the ``.16x*`` sweeps we: 1. Fill a (128, FULL_W) host buffer ``A`` with random values. 2. Stage ``A`` into TMEM via the existing ``.32x32b`` ld/st round-trip. @@ -1201,6 +1203,113 @@ def kernel(A_ptr: T.handle) -> None: ) +def test_tcgen05_32x32b_float32_keeps_typed_register_operands(): + """The .32x32b float32 lowering should pass the local fragment's typed + registers to the PTX helper. The helper reinterprets operands as b32, so a + call-site uint32 view is unnecessary and makes generated CUDA diverge from + handwritten ``T.ptx.tcgen05`` calls.""" + + K_cols = 32 + + @T.prim_func + def kernel(A_ptr: T.handle) -> None: + T.match_buffer(A_ptr, (128, K_cols), "float32") + T.device_entry() + warp_id = T.warp_id([4]) + T.cta_id([2]) + wg_id = T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + T.thread_id([128]) + + tmem_addr = T.alloc_shared([1], "uint32") + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=K_cols, cta_group=1) + T.tvm_storage_sync("shared") + tmem = T.decl_buffer( + (128, K_cols), + "float32", + scope="tmem", + allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, K_cols) : (1 @ TLane, 1 @ TCol)]), + ) + frag = T.alloc_tcgen05_ldst_frag("32x32b", (128, K_cols), "float32") + Tx.wg.copy_async(frag[:, :], tmem[:, :]) + T.ptx.tcgen05.wait.ld() + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=K_cols, cta_group=1) + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + call_lines = [ + line + for line in src.splitlines() + if "tvm_builtin_ptx_tcgen05_ld_32x32b_x32((" in line and "__forceinline__" not in line + ] + assert call_lines + assert "((uint*)" not in call_lines[0] + + +def test_tcgen05_ldst_constant_tmem_address_is_uint32(): + """A constant TMEM base address must still be passed as uint32. + + Sparse FlashMLA uses pool-relative constant TMEM columns. Handwritten + ``T.ptx.tcgen05`` calls pass ``T.uint32(0)`` for the base address; the tile + primitive should emit the same ABI shape instead of a bare integer literal. + """ + + K_cols = 32 + + @T.prim_func + def kernel() -> None: + T.device_entry() + T.warp_id([4]) + T.cta_id([1]) + wg_id = T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + T.thread_id([128]) + + if wg_id == 0: + tmem = T.decl_buffer( + (128, K_cols), + "float32", + scope="tmem", + allocated_addr=0, + layout=TileLayout(S[(128, K_cols) : (1 @ TLane, 1 @ TCol)]), + ) + frag = T.alloc_tcgen05_ldst_frag("32x32b", (128, K_cols), "float32") + Tx.wg.copy_async(frag[:, :], tmem[:, :]) + T.ptx.tcgen05.wait.ld() + Tx.wg.copy_async(tmem[:, :], frag[:, :]) + T.ptx.tcgen05.wait.st() + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": kernel}) + mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + ld_lines = [ + line + for line in src.splitlines() + if "tvm_builtin_ptx_tcgen05_ld_32x32b_x32((" in line and "__forceinline__" not in line + ] + st_lines = [ + line + for line in src.splitlines() + if "tvm_builtin_ptx_tcgen05_st_32x32b_x32(" in line and "__forceinline__" not in line + ] + assert ld_lines + assert st_lines + assert "(uint)0" in ld_lines[0] + assert "(uint)0" in st_lines[0] + + # -------------------------------------------------------------------------- # Test 3: column-slice loads of a wider frag # @@ -1336,5 +1445,304 @@ def test_tcgen05_ld_16x256b_sliced_matches_full_M128(full_rep, n_chunks): _run_sliced_vs_full_load("16x256b", full_rep, n_chunks) +# .32x32b tmem<->reg and smem<->tmem copy_async (migrated from test_tmem.py). + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize("dtype", ["float16", "float32"]) +@pytest.mark.parametrize("width_32b", [4, 8, 16, 32]) +def test_copy_tmem2reg_async(dtype, width_32b): + """Test async tmem<->local copy using copy_async instead of copy. + + This tests the new copy_async dispatch for tmem<->local that doesn't + immediately wait after the operation, allowing for pipelining. + """ + + def next_power_of_2(x): + """Return the smallest power of 2 greater than or equal to x.""" + if x <= 1: + return 1 + return 1 << (x - 1).bit_length() + + bits = tvm.runtime.DataType(dtype).bits + if 128 % bits != 0 or 32 % bits != 0: + pytest.skip(f"dtype {dtype} is not supported") + + WIDTH = width_32b * (32 // bits) + VEC_LEN = 128 // bits + if WIDTH % VEC_LEN != 0: + pytest.skip(f"dtype {dtype} + width {width_32b} is not supported") + + g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) + local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)]) + + # fmt: off + @T.prim_func + def copy_async_test(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (128, WIDTH), dtype) + B = T.match_buffer(B_ptr, (128, WIDTH), dtype) + + A_flat = A.view(-1) + B_flat = B.view(-1) + + T.device_entry() + warp_id = T.warp_id([(128) // 32]) + cta_id = T.cta_id([2]) + wg_id = T.warpgroup_id([1]) + warp_id_in_wg = T.warp_id_in_wg([4]) + lane_id = T.lane_id([32]) + tid_in_wg = T.thread_id([128]) + + tmem_addr = T.alloc_shared([1], "uint32") + + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 + + T.tvm_storage_sync("shared") + + tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)])) + + A_reg = T.alloc_local((WIDTH), dtype) + B_reg = T.alloc_local((WIDTH), dtype) + A_local = A_reg.view(128, WIDTH, layout=local_view) + B_local = B_reg.view(128, WIDTH, layout=local_view) + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 + for i in range(WIDTH): + B_reg[i] = T.cast(0, dtype) + T.cuda.cta_sync() + + # A_local -> tmem (async) + Tx.wg.copy_async(tmem[:, :], A_local[:, :]) + T.ptx.tcgen05.wait.st() # explicit wait + T.cuda.cta_sync() + + # tmem -> B_local (async) + Tx.wg.copy_async(B_local[:, :], tmem[:, :]) + T.ptx.tcgen05.wait.ld() # explicit wait + T.cuda.cta_sync() + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501 + + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 + # fmt: on + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": copy_async_test}) + mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) + B_np = np.zeros((128, WIDTH), dtype=dtype) + DEV = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, DEV) + B = tvm.runtime.tensor(B_np, DEV) + mod(A, B) + np.testing.assert_allclose(B.numpy(), A_np) + + +# tmem<->reg round-trip via T.copy_async (migrated from test_copy_sync.py): +# the kernels are the real async tmem dispatch tests; G<->L copies just stage. + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize("dtype", ["uint8", "float16", "float32"]) +@pytest.mark.parametrize("width_32b", [2, 4, 8, 16, 32, 64, 128]) +@pytest.mark.parametrize("offset_32b", [0, 3, 10]) +def test_copy_tmem2reg(dtype, width_32b, offset_32b): + def next_power_of_2(x): + if x <= 1: + return 1 + return 1 << (x - 1).bit_length() + + bits = tvm.runtime.DataType(dtype).bits + if 128 % bits != 0 or 32 % bits != 0: + pytest.skip(f"dtype {dtype} is not supported") + + WIDTH = width_32b * (32 // bits) + OFFSET = offset_32b * (32 // bits) + VEC_LEN = 128 // bits + if WIDTH % VEC_LEN != 0: + pytest.skip(f"dtype {dtype} + width {width_32b} is not supported") + + g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) + local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)]) + + # fmt: off + @T.prim_func + def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (128, WIDTH), dtype) + B = T.match_buffer(B_ptr, (128, WIDTH), dtype) + + A_flat = A.view(-1) + B_flat = B.view(-1) + + T.device_entry() + warp_id = T.warp_id([(128) // 32]) + T.cta_id([2]) + wg_id = T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + tid_in_wg = T.thread_id([128]) + + tmem_addr = T.alloc_shared([1], "uint32") + + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501 + + T.tvm_storage_sync("shared") + + tmem = T.decl_buffer((128, OFFSET + WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], # noqa: E501 + layout=TileLayout(S[(128, OFFSET + WIDTH) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 + + A_reg = T.alloc_local((WIDTH), dtype) + B_reg = T.alloc_local((WIDTH), dtype) + A_local = A_reg.view(128, WIDTH, layout=local_view) + B_local = B_reg.view(128, WIDTH, layout=local_view) + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 + for i in range(WIDTH): + B_reg[i] = T.cast(0, dtype) + T.cuda.cta_sync() + + # A_local -> tmem + Tx.wg.copy_async(tmem[:, OFFSET: OFFSET + WIDTH], A_local[:, :]) + T.ptx.tcgen05.wait.st() + T.cuda.cta_sync() + + # tmem -> B_local + Tx.wg.copy_async(B_local[:, :], tmem[:, OFFSET: OFFSET + WIDTH]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501 + + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501 + # fmt: on + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": copy_sync}) + mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) + B_np = np.zeros((128, WIDTH), dtype=dtype) + DEV = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, DEV) + B = tvm.runtime.tensor(B_np, DEV) + mod(A, B) + np.testing.assert_allclose(B.numpy(), A_np) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.parametrize("dtype", ["float16", "float32"]) +@pytest.mark.parametrize("width_32b", [4, 8, 16, 32]) +@pytest.mark.parametrize("local_offset_32b", [0, 2, 4]) +def test_copy_tmem2reg_sliced_local(dtype, width_32b, local_offset_32b): + """tmem<->local copy with a sliced local buffer region.""" + + def next_power_of_2(x): + if x <= 1: + return 1 + return 1 << (x - 1).bit_length() + + bits = tvm.runtime.DataType(dtype).bits + if 128 % bits != 0 or 32 % bits != 0: + pytest.skip(f"dtype {dtype} is not supported") + + WIDTH = width_32b * (32 // bits) + LOCAL_OFFSET = local_offset_32b * (32 // bits) + TOTAL_LOCAL_WIDTH = WIDTH + LOCAL_OFFSET + VEC_LEN = 128 // bits + if WIDTH % VEC_LEN != 0 or TOTAL_LOCAL_WIDTH % VEC_LEN != 0: + pytest.skip( + f"dtype {dtype} + width {width_32b} + offset {local_offset_32b} is not supported" + ) + + g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) + local_view = TileLayout(S[(128, TOTAL_LOCAL_WIDTH) : (1 @ axis_tid_in_wg, 1)]) + + # fmt: off + @T.prim_func + def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (128, WIDTH), dtype) + B = T.match_buffer(B_ptr, (128, WIDTH), dtype) + + A_flat = A.view(-1) + B_flat = B.view(-1) + + T.device_entry() + warp_id = T.warp_id([(128) // 32]) + T.cta_id([2]) + wg_id = T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + tid_in_wg = T.thread_id([128]) + + tmem_addr = T.alloc_shared([1], "uint32") + + if wg_id == 0: + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 + + T.tvm_storage_sync("shared") + + tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)])) + + A_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype) + B_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype) + A_local = A_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view) + B_local = B_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view) + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(A_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 + for i in range(TOTAL_LOCAL_WIDTH): + B_reg[i] = T.cast(0, dtype) + T.cuda.cta_sync() + + # A_local[sliced] -> tmem (use sliced region) + Tx.wg.copy_async(tmem[:, 0:WIDTH], A_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH]) + T.ptx.tcgen05.wait.st() + T.cuda.cta_sync() + + # tmem -> B_local[sliced] (use sliced region) + Tx.wg.copy_async(B_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH], tmem[:, 0:WIDTH]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + for i in range(WIDTH // VEC_LEN): + g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) + Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN]) # noqa: E501 + + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 + # fmt: on + + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": copy_sync}) + mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) + B_np = np.zeros((128, WIDTH), dtype=dtype) + DEV = tvm.cuda(0) + A = tvm.runtime.tensor(A_np, DEV) + B = tvm.runtime.tensor(B_np, DEV) + mod(A, B) + np.testing.assert_allclose(B.numpy(), A_np) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 2096e2537edb..f2431ab15fdf 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py @@ -8,26 +8,44 @@ # # http://www.apache.org/licenses/LICENSE-2.0 # -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# pylint: disable=invalid-name, missing-function-docstring -import functools +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the independent ``tma_auto`` and ``tma_explicit`` planners.""" + +from dataclasses import dataclass, replace +from functools import reduce +from operator import mul import numpy as np import pytest import tvm import tvm.testing -from tvm.ir import PointerType, PrimType -from tvm.ir.type import TensorMapType +from tvm.arith import Analyzer +from tvm.ir import PointerType, PrimType, Range from tvm.script import tirx as T from tvm.script.tirx import tile as Tx from tvm.testing import env from tvm.tirx import IntImm, StringImm, Var +from tvm.tirx.cuda.operator.tile_primitive.copy_async.tma import ( + AutoIssueAxis, + IssueCoord, + ProofStatus, + TensorMapSpec, + TMAPlan, + _build_auto_plan, + _build_explicit_plan, + _promote_auto_once, + _selector_compatibility, + _validation_failures, + copy_tma_auto_impl, + copy_tma_explicit_impl, + validate_tensor_map_spec, +) from tvm.tirx.cuda.operator.tile_primitive.tma_utils import ( mma_atom_layout, mma_atom_shape, @@ -35,1628 +53,1962 @@ ) from tvm.tirx.exec_scope import ExecScope from tvm.tirx.layout import S, TileLayout -from tvm.tirx.operator.tile_primitive.dispatch_context import DispatchContext from tvm.tirx.operator.tile_primitive.ops import CopyAsync -from tvm.tirx.stmt import DeclBuffer +from tvm.tirx.stmt import BufferRegion from tvm.tirx.stmt_functor import StmtExprVisitor +from tvm.tirx.tile_primitive import DispatchContext -# =========================================================================== -# Helpers -# =========================================================================== +_TMA_OPS = { + "tirx.ptx.cp_async_bulk_tensor_g2s_cluster", + "tirx.ptx.cp_async_bulk_tensor_shared_to_global", + "tirx.ptx.cp_async_bulk_tensor_shared_to_global_reduce", +} -class TMACounter(StmtExprVisitor): - """Visitor to count total TMA operations including loop iterations. - - This verifies that TMA copy operations are optimized correctly, - resulting in minimal TMA instructions instead of multiple iterations. - """ - +class _TMACounter(StmtExprVisitor): def __init__(self): super().__init__() - self.loop_extents = [] # Stack of loop extents - self.total_tma_ops = 0 + self.loop_extents = [] + self.total = 0 + self.calls = [] + self.weighted_calls = [] def visit_for_(self, op): - extent = op.extent - self.loop_extents.append(extent) + self.loop_extents.append(op.extent) self.visit_stmt(op.body) self.loop_extents.pop() def visit_evaluate_(self, op): - if isinstance(op.value, tvm.ir.Call): - if op.value.op.name in ( - "tirx.ptx.cp_async_bulk_tensor_global_to_cluster", - "tirx.ptx.cp_async_bulk_tensor_shared_to_global", - "tirx.ptx.cp_async_bulk_tensor_shared_to_global_reduce", - ): - # Multiply all enclosing loop extents - iters = 1 - for ext in self.loop_extents: - iters *= ext - self.total_tma_ops += iters - - -def _make_tma_call( - g_shape, - g_region, - s_shape, - s_region, - gmem_layout, - smem_layout, - dtype="float16", - direction="g2s", - config=None, -): - """Construct TilePrimitiveCall + DispatchContext and call copy_tma_impl. + if isinstance(op.value, tvm.ir.Call) and op.value.op.name in _TMA_OPS: + multiplier = reduce(mul, (int(extent) for extent in self.loop_extents), 1) + self.total += multiplier + self.calls.append(op.value) + self.weighted_calls.append((op.value, multiplier)) + super().visit_evaluate_(op) - Returns (impl, host_init_stmts) on success, raises DispatchFail on failure. - impl is the device-side PrimFunc, host_init_stmts is a list of Stmt - for host-side tensor map creation. - """ - from tvm.ir import Range - from tvm.tirx import Var - from tvm.tirx.cuda.operator.tile_primitive.copy_async.tma import copy_tma_impl - from tvm.tirx.stmt import BufferRegion - g_buf = tvm.tirx.decl_buffer(g_shape, dtype, "A", layout=gmem_layout) - s_buf = tvm.tirx.decl_buffer(s_shape, dtype, "A_smem", scope="shared.dyn", layout=smem_layout) +class _EncodeCollector(StmtExprVisitor): + def __init__(self): + super().__init__() + self.calls = [] - g_ranges = [Range.from_min_extent(r[0], r[1] - r[0]) for r in g_region] - s_ranges = [Range.from_min_extent(r[0], r[1] - r[0]) for r in s_region] + def visit_call_(self, op): + if ( + isinstance(op.op, tvm.ir.Op) + and op.op.name == "tirx.tvm_call_packed" + and isinstance(op.args[0], StringImm) + and op.args[0].value == "runtime.cuTensorMapEncodeTiled" + ): + self.calls.append(op) + super().visit_call_(op) - config = dict(config or {}) - if direction == "g2s": - mbar_ptr = Var("mbar_ptr", "handle") - config.setdefault("mbar", mbar_ptr) - config.setdefault("cta_group", 1) - dst_br = BufferRegion(s_buf, s_ranges) - src_br = BufferRegion(g_buf, g_ranges) - else: # s2g - config.setdefault("cta_group", 1) - dst_br = BufferRegion(g_buf, g_ranges) - src_br = BufferRegion(s_buf, s_ranges) - - op_call = CopyAsync(dst_br, src_br, config=config) - target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"}) - sctx = DispatchContext(target, ExecScope("thread"), {}, {}) +class _SelectCollector(StmtExprVisitor): + def __init__(self): + super().__init__() + self.nodes = [] - impl = copy_tma_impl(op_call, sctx) - host_init_stmts = list(sctx.callbacks.get("host_init_stmt", [])) - return impl, host_init_stmts + def visit_select_(self, op): + self.nodes.append(op) + super().visit_select_(op) -def _count_tma_ops(impl): - """Count total TMA ops in a PrimFunc (including loop multiplier).""" - counter = TMACounter() - counter.visit_stmt(impl.body) - return counter.total_tma_ops +class _PrefetchCollector(StmtExprVisitor): + def __init__(self): + super().__init__() + self.names = [] + def visit_call_(self, op): + if ( + isinstance(op.op, tvm.ir.Op) + and op.op.name == "tirx.ptx.prefetch_tensormap" + and isinstance(op.args[0], tvm.ir.Call) + and op.args[0].op.name == "tirx.address_of" + and isinstance(op.args[0].args[0], Var) + ): + self.names.append(op.args[0].args[0].name) + super().visit_call_(op) -def _build_expected_host_init(dtype, encode_args): - """Build expected host_init Bind+SeqStmt for cuTensorMapEncodeTiled. - encode_args is a list of ints: the numeric arguments to cuTensorMapEncodeTiled - after (tensormap, dtype_str, ndim, A_ptr). The full call is: - runtime.cuTensorMapEncodeTiled(tensormap, dtype_str, ndim, A_ptr, *encode_args) - where ndim = encode_args[0] and the rest are the tensor map parameters. - """ - A_tensormap = Var("A_tensormap", PointerType(TensorMapType(), "global")) - stack_alloca = tvm.ir.Call( - tvm.ir.Op.get("tirx.tvm_stack_alloca"), - [StringImm("tensormap"), IntImm("int32", 1)], - ret_ty="handle", - ) - A_var = Var("A", PointerType(PrimType(dtype), "global")) - call_args = ( - [ - StringImm("runtime.cuTensorMapEncodeTiled"), - A_tensormap, - StringImm(dtype), - IntImm("int32", encode_args[0]), # ndim - A_var, - ] - + [IntImm("int32", v) for v in encode_args[1:]] - ) - encode_call = tvm.ir.Call(tvm.ir.Op.get("tirx.tvm_call_packed"), call_args, ret_ty="int32") - replace_point = tvm.tirx.Evaluate(tvm.tirx.op.tvm_kernel_replace_point()) - return tvm.tirx.SeqStmt( - [tvm.tirx.Bind(A_tensormap, stack_alloca), tvm.tirx.Evaluate(encode_call), replace_point] - ) +def _plain_layout(shape, strides=None): + shape = tuple(shape) + if strides is None: + return TileLayout(S[shape]) + return TileLayout(S[shape : tuple(strides)]) -def _build_expected_impl(direction, dtype, s_shape, s_layout, impl_spec): - """Build expected impl PrimFunc. - - impl_spec is a dict with: - loop_extents: list[int] — e.g. [1], [2, 2], [8] - dim: int — TMA rank (number of coordinates, also the dim arg to PTX call) - elem_offset_fn: callable(loop_vars) -> Expr (or None for 0) - coord_fn: callable(loop_vars) -> list[Expr] (dim coordinate args) - s_start: optional list[int] — starting index for address_of (default all zeros) - """ - from tvm.tirx.layout import ComposeLayout, SwizzleLayout - - loop_extents = impl_spec["loop_extents"] - dim = impl_spec["dim"] - elem_offset_fn = impl_spec.get("elem_offset_fn") - coord_fn = impl_spec["coord_fn"] - - # Mirror _to_tile_layout() in copy_async/tma.py: - # ComposeLayout → tile_layout - # SwizzleLayout → identity TileLayout(S[shape]) - # TileLayout → as-is - if isinstance(s_layout, ComposeLayout): - buf_layout = s_layout.tile_layout - elif isinstance(s_layout, SwizzleLayout): - buf_layout = TileLayout(S[tuple(s_shape)]) - else: - buf_layout = s_layout +def _ranges(region): + return [Range.from_min_extent(start, extent) for start, extent in region] - # Create loop vars - n_loops = len(loop_extents) - if n_loops == 1: - loop_vars = [Var("loop_vars", "int32")] - else: - loop_vars = [Var(f"loop_vars_{i}", "int32") for i in range(n_loops)] - # Buffer - s_buf_ptr = Var("s_buf_w_offset_ptr", PointerType(PrimType(dtype), "shared.dyn")) - elem_offset = elem_offset_fn(loop_vars) if elem_offset_fn else IntImm("int32", 0) +def _make_op( + *, + g_shape, + s_shape=None, + g_region=None, + s_region=None, + g_layout=None, + s_layout=None, + dtype="float16", + s_dtype=None, + direction="g2s", + config=None, + g_data=None, + g_elem_offset=0, + s_elem_offset=0, + target_arch="sm_90a", + sctx=None, +): + g_shape = tuple(g_shape) + s_shape = tuple(s_shape or g_shape) + g_region = tuple(g_region or ((0, extent) for extent in g_shape)) + s_region = tuple(s_region or ((0, extent) for extent in s_shape)) + g_layout = g_layout or _plain_layout(g_shape) + s_layout = s_layout or _plain_layout(s_shape) + s_dtype = s_dtype or dtype + g_buf = tvm.tirx.decl_buffer( + g_shape, + dtype, + "A", + data=g_data, + elem_offset=g_elem_offset, + layout=g_layout, + ) s_buf = tvm.tirx.decl_buffer( s_shape, - dtype, - "s_buf_w_offset", - data=s_buf_ptr, - elem_offset=elem_offset, + s_dtype, + "A_smem", scope="shared.dyn", - layout=buf_layout, + elem_offset=s_elem_offset, + layout=s_layout, ) - - # Free variables - mbar_ptr = Var("mbar_ptr", "handle") - A_tensormap = Var("A_tensormap", PointerType(TensorMapType(), "global")) - - # address_of(s_buf[s_start...]) - s_start = impl_spec.get("s_start") - if s_start: - buf_indices = [IntImm("int32", v) for v in s_start] + config = dict(config or {}) + if direction == "g2s": + config.setdefault("mbar", Var("mbar", "handle")) + dst = BufferRegion(s_buf, _ranges(s_region)) + src = BufferRegion(g_buf, _ranges(g_region)) else: - buf_indices = [IntImm("int32", 0)] * len(s_shape) - addr_of = tvm.ir.Call( - tvm.ir.Op.get("tirx.address_of"), - [tvm.tirx.BufferLoad(s_buf, buf_indices)], - ret_ty="handle", + dst = BufferRegion(g_buf, _ranges(g_region)) + src = BufferRegion(s_buf, _ranges(s_region)) + op = CopyAsync(dst, src, config=config) + if sctx is None: + target = tvm.target.Target({"kind": "cuda", "arch": target_arch}) + sctx = DispatchContext(target, ExecScope("thread"), {}, {}) + return op, sctx, g_buf, s_buf + + +def _lower_direct(variant, **kwargs): + op, sctx, _, _ = _make_op(**kwargs) + if variant == "tma_auto": + impl = copy_tma_auto_impl(op, sctx) + else: + impl = copy_tma_explicit_impl(op, sctx) + return impl, list(sctx.callbacks.get("host_init_stmt", [])), sctx + + +def _auto_plan(**kwargs): + op, sctx, _, _ = _make_op(**kwargs) + return _build_auto_plan(op, sctx) + + +def _direct_plan(variant, **kwargs): + op, sctx, _, _ = _make_op(**kwargs) + if variant == "tma_auto": + return _build_auto_plan(op, sctx) + plan, _ = _build_explicit_plan(op, sctx) + return plan + + +def _count_tma(stmt): + counter = _TMACounter() + counter.visit_stmt(stmt.body if isinstance(stmt, tvm.tirx.PrimFunc) else stmt) + return counter + + +def _collect_encodes(stmts): + collector = _EncodeCollector() + for stmt in stmts: + collector.visit_stmt(stmt) + return collector.calls + + +def _encode_signature(call): + rank = int(call.args[3]) + cursor = 5 + dims = tuple(call.args[cursor : cursor + rank]) + cursor += rank + strides = tuple(call.args[cursor : cursor + rank - 1]) + cursor += rank - 1 + boxes = tuple(call.args[cursor : cursor + rank]) + cursor += rank + element_strides = tuple(call.args[cursor : cursor + rank]) + cursor += rank + enums = tuple(call.args[cursor : cursor + 4]) + cursor += 4 + forced_dtype = call.args[cursor] if cursor < len(call.args) else None + return { + "dtype": call.args[2].value, + "rank": rank, + "base": call.args[4], + "dims": dims, + "strides": strides, + "boxes": boxes, + "element_strides": element_strides, + "enums": enums, + "forced_dtype": forced_dtype, + } + + +def _ints(values): + return tuple(int(value) for value in values) + + +def _make_spec(**overrides): + g_data = Var("A", PointerType(PrimType("float16"), "global")) + s_buf = tvm.tirx.decl_buffer( + (4, 8), + "float16", + "A_smem", + scope="shared.dyn", + layout=_plain_layout((4, 8)), ) - - # Coordinate args (must have exactly `dim` entries) - coords = coord_fn(loop_vars) - tensormap_addr = tvm.ir.Call(tvm.ir.Op.get("tirx.address_of"), [A_tensormap], ret_ty="uint64") - - # Build PTX call based on direction - if direction == "g2s": - # g2c(dim, addr, mbar, tensormap, cta_mask, cta_group, - # cache_policy, has_cache_policy, *coords) - ptx_op = tvm.ir.Op.get("tirx.ptx.cp_async_bulk_tensor_global_to_cluster") - ptx_args = [ - IntImm("int32", dim), - addr_of, - mbar_ptr, - tensormap_addr, - IntImm("int32", 0), - IntImm("int32", 1), - IntImm("uint64", 0), - IntImm("int32", 0), - *coords, - ] - else: # s2g - # s2g(dim, addr, tensormap, cache_policy, has_cache_policy, *coords) - ptx_op = tvm.ir.Op.get("tirx.ptx.cp_async_bulk_tensor_shared_to_global") - ptx_args = [ - IntImm("int32", dim), - addr_of, - tensormap_addr, - IntImm("uint64", 0), - IntImm("int32", 0), - *coords, - ] - - eval_stmt = tvm.tirx.Evaluate(tvm.ir.Call(ptx_op, ptx_args)) - - # Wrap: DeclBuffer -> nested For loops (skipped when total extent is 1, - # matching the implementation's always-unroll single-loop emission). - body = DeclBuffer(s_buf, eval_stmt, data=s_buf_ptr) - for i in range(n_loops - 1, -1, -1): - body = tvm.tirx.For( - loop_vars[i], - IntImm("int32", 0), - IntImm("int32", loop_extents[i]), - tvm.tirx.ForKind.UNROLLED, - body, - ) - - func = tvm.tirx.PrimFunc([], body, ret_type=None, buffer_map={}) - func = func.with_attr("global_symbol", "impl") - # default s_tir=False is implicit; nothing to set here - return func - - -def _zeros(n): - """Return n zero IntImm coords.""" - return [IntImm("int32", 0)] * n - - -def _atom_rank5_elem_offset(lvs): - """elem_offset for the structural 5D atom plan: lv * 8192.""" - return lvs[0] * 8192 - - -def _atom_rank5_coords(lvs): - """coord_fn for the structural 5D atom plan: [0, 0, 0, lv*2, 0].""" - return [ - IntImm("int32", 0), - IntImm("int32", 0), - IntImm("int32", 0), - lvs[0] * 2, - IntImm("int32", 0), - ] - - -def _stride_gap_elem_offset(lvs): - """elem_offset for stride-gap-outer: lv * 4096.""" - return lvs[0] * 4096 + values = dict( + descriptor_dtype="float16", + descriptor_bits=16, + effective_bytes=2, + packed_kind=None, + force_cu_dtype=-1, + base=g_data, + base_key="A", + descriptor_name="A", + base_byte_offset=0, + global_dims=(16, 64), + global_strides=(32,), + inner_stride=1, + box_dims=(8, 4), + element_strides=(1, 1), + interleave=0, + swizzle=0, + l2_promotion=2, + oob_fill=0, + direction="g2s", + load_mode="tile", + target_arch="sm_90a", + coordinates=(0, 0), + gather4=(), + smem_buffer=s_buf, + smem_start=(0, 0), + smem_base_offset=0, + mbar=Var("mbar", "handle"), + mbar_is_shared_addr=False, + cta_group=1, + cta_mask=0, + cache_hint="", + cache_policy=None, + use_tma_reduce=None, + payload_bits=1024, + transaction_bits=1024, + ) + values.update(overrides) + return TensorMapSpec(**values) -def _stride_gap_3d_coords(lvs): - """coord_fn for stride-gap-outer (rank=3): [0, 0, lv].""" - return [IntImm("int32", 0), IntImm("int32", 0), lvs[0]] +def _finding(spec, rule): + return [item for item in validate_tensor_map_spec(spec) if item.rule == rule] -def _atom_multiphase_rank5_elem_offset(lvs): - """elem_offset for the multiphase 5D atom plan: lv * 4096.""" - return lvs[0] * 4096 +def _from_source(source): + return tvm.script.from_source(source, {"T": T, "Tx": Tx}) -def _atom_multiphase_rank5_coords(lvs): - """coord_fn for multiphase rank-5 atom: [0, 0, lv%2*4, lv//2*2, 0].""" - return [ - IntImm("int32", 0), - IntImm("int32", 0), - (lvs[0] % 2) * 4, - (lvs[0] // 2) * 2, - IntImm("int32", 0), - ] +def _lower_module(func, arch="sm_100a"): + target = tvm.target.Target({"kind": "cuda", "arch": arch}) + with target: + return tvm.tirx.transform.LowerTIRx()(tvm.IRModule({"main": func})) -# fmt: off -# Expected parameters for each TMA test case. -# Each entry maps case_id -> (impl_spec_dict, encode_args_list). -# -# impl_spec keys: -# loop_extents: list[int] — iteration counts for nested loops -# dim: int — TMA rank = number of coordinates = dim arg to PTX call -# coord_fn: callable(loop_vars) -> list[Expr] — coordinate arguments (len == dim) -# elem_offset_fn: optional callable(loop_vars) -> Expr — buffer offset -# -# encode_args: list[int] — all numeric args to cuTensorMapEncodeTiled -# [ndim, global_strides..., global_dims..., box_dims..., elem_strides..., -# interleave, swizzle_mode, l2_promotion, oob_fill] +def _compile_module(func, arch="sm_100a"): + target = tvm.target.Target({"kind": "cuda", "arch": arch}) + with target: + return tvm.compile( + tvm.IRModule({"main": func}), + target=target, + tir_pipeline="tirx", + ) -# =========================================================================== -# Section 2: TMA unit tests — single parametrized structural-golden driver -# =========================================================================== +def _tma_case(case_id, **kwargs): + return pytest.param(case_id, kwargs, id=case_id) -def _tma_case( - *, - id, - g_shape, - g_region, - s_shape, - s_region, - gmem_layout, - smem_layout, - dtype="float16", - direction="g2s", - config=None, - impl_spec=None, - encode_args=None, - raises=None, -): - """Build a pytest.param carrying a dict-form case for ``test_copy_tma_codegen``. - - Required: ``g_shape``, ``g_region``, ``s_shape``, ``s_region``, ``gmem_layout``, - ``smem_layout``, ``id``. - - Optional: - ``dtype``: element dtype (default ``"float16"``). - ``direction``: ``"g2s"`` or ``"s2g"`` (default ``"g2s"``). - ``config``: op config dict forwarded to ``copy_tma_impl`` (e.g. - ``{"oob": "nan"}``). - ``impl_spec``: kwargs for ``_build_expected_impl``. ``None`` skips the - device-impl structural check. - ``encode_args``: list for ``_build_expected_host_init``. ``None`` skips - the host-init structural check. - ``raises``: ``(ExceptionClass, regex_str)`` to expect instead of a - successful dispatch. - """ - return pytest.param( - dict( - g_shape=g_shape, g_region=g_region, - s_shape=s_shape, s_region=s_region, - gmem_layout=gmem_layout, smem_layout=smem_layout, - dtype=dtype, direction=direction, config=config, - impl_spec=impl_spec, encode_args=encode_args, raises=raises, - ), - id=id, - ) +_BASELINE_AUTO_CASES = [ + ("g2s-2d-8x256", "float16", 3), + ("g2s-2d-8x256-swizzle2", "float16", 2), + ("g2s-2d-8x256-swizzle1", "float16", 1), + ("g2s-2d-8x256-swizzle0", "float16", 0), + ("g2s-2d-8x256-int8", "int8", 3), + ("g2s-2d-8x256-bf16", "bfloat16", 3), + ("g2s-2d-8x256-fp32", "float32", 3), + ("g2s-2d-8x256-uint8", "uint8", 3), + ("g2s-2d-8x256-fp8e4m3", "float8_e4m3fn", 3), + ("g2s-2d-8x256-fp8e5m2", "float8_e5m2", 3), +] -# fmt: off TMA_CASES = [ - # ====================================================================== - # G2S — 2D baseline (swizzle + dtype variants sharing (8, 256) shape) - # ====================================================================== - _tma_case( - id="g2s-2d-8x256", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float16", 3, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 8, 4, 512, 128, 64, 8, 4, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-swizzle2", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float16", 2, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 32, 8, 8, 512, 64, 32, 8, 8, 1, 1, 1, 0, 2, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-swizzle1", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float16", 1, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 16, 8, 16, 512, 32, 16, 8, 16, 1, 1, 1, 0, 1, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-swizzle0", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float16", 0, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=2, coord_fn=lambda lv: _zeros(2)), - encode_args=[2, 256, 8, 512, 256, 8, 1, 1, 0, 0, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-int8", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("int8", 3, (8, 256)), - dtype="int8", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 128, 8, 2, 256, 128, 128, 8, 2, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-bf16", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("bfloat16", 3, (8, 256)), - dtype="bfloat16", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 8, 4, 512, 128, 64, 8, 4, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-fp32", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float32", 3, (8, 256)), - dtype="float32", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 32, 8, 8, 1024, 128, 32, 8, 8, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-uint8", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("uint8", 3, (8, 256)), - dtype="uint8", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 128, 8, 2, 256, 128, 128, 8, 2, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-fp8e4m3", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float8_e4m3fn", 3, (8, 256)), - dtype="float8_e4m3fn", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 128, 8, 2, 256, 128, 128, 8, 2, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="g2s-2d-8x256-fp8e5m2", - g_shape=(8, 256), g_region=((0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=mma_shared_layout("float8_e5m2", 3, (8, 256)), - dtype="float8_e5m2", - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 128, 8, 2, 256, 128, 128, 8, 2, 1, 1, 1, 0, 3, 2, 0], - ), - # ====================================================================== - # G2S — 3D / partial / edge / multidim layouts - # ====================================================================== + *[ + _tma_case( + case_id, + g_shape=(8, 256), + g_region=((0, 8), (0, 256)), + s_shape=(8, 256), + s_region=((0, 8), (0, 256)), + g_layout=_plain_layout((8, 256)), + s_layout=mma_shared_layout(dtype, swizzle, (8, 256)), + dtype=dtype, + ) + for case_id, dtype, swizzle in _BASELINE_AUTO_CASES + ], _tma_case( - id="g2s-3d-shared-64x256", - g_shape=(64, 256), g_region=((0, 64), (0, 256)), - s_shape=(3, 64, 256), s_region=((1, 2), (0, 64), (0, 256)), - gmem_layout=TileLayout(S[64, 256]), - smem_layout=mma_shared_layout("float16", 3, (3, 64, 256)), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3), s_start=[1, 0, 0]), - encode_args=[3, 64, 64, 4, 512, 128, 64, 64, 4, 1, 1, 1, 0, 3, 2, 0], + "g2s-3d-shared-64x256", + g_shape=(64, 256), + g_region=((0, 64), (0, 256)), + s_shape=(3, 64, 256), + s_region=((1, 1), (0, 64), (0, 256)), + g_layout=_plain_layout((64, 256)), + s_layout=mma_shared_layout("float16", 3, (3, 64, 256)), ), _tma_case( - id="g2s-2d-32x512-atom", - g_shape=(32, 512), g_region=((0, 32), (0, 512)), - s_shape=(32, 512), s_region=((0, 32), (0, 512)), - gmem_layout=TileLayout(S[32, 512]), - smem_layout=( + "g2s-2d-32x512-atom", + g_shape=(32, 512), + g_region=((0, 32), (0, 512)), + s_shape=(32, 512), + s_region=((0, 32), (0, 512)), + g_layout=_plain_layout((32, 512)), + s_layout=( mma_atom_layout("float16", 3) .tile_to((16, 256), mma_atom_shape("float16", 3)) .tile_to((32, 512), (16, 256)) ), - impl_spec=dict( - loop_extents=[2], dim=5, - coord_fn=_atom_rank5_coords, elem_offset_fn=_atom_rank5_elem_offset, - ), - encode_args=[5, 64, 8, 4, 4, 2, 1024, 128, 8192, 512, 64, 8, 4, 2, 2, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 ), _tma_case( - id="g2s-2d-partial-8192", - g_shape=(8192, 8192), g_region=((0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[8192, 8192]), - smem_layout=mma_shared_layout("float16", 3, (128, 64)), - impl_spec=dict(loop_extents=[1], dim=2, coord_fn=lambda lv: _zeros(2)), - encode_args=[2, 8192, 8192, 16384, 64, 128, 1, 1, 0, 3, 2, 0], + "g2s-2d-partial-8192", + g_shape=(8192, 8192), + g_region=((0, 128), (0, 64)), + s_shape=(128, 64), + s_region=((0, 128), (0, 64)), + g_layout=_plain_layout((8192, 8192)), + s_layout=mma_shared_layout("float16", 3, (128, 64)), ), _tma_case( - id="g2s-edge-4d-shared-128x64", - g_shape=(128, 64), g_region=((0, 128), (0, 64)), - s_shape=(2, 2, 128, 64), s_region=((0, 1), (0, 1), (0, 128), (0, 64)), - gmem_layout=TileLayout(S[128, 64]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (2, 2, 128, 64)).canonicalize(), - impl_spec=dict(loop_extents=[1], dim=2, coord_fn=lambda lv: _zeros(2)), - encode_args=[2, 64, 128, 128, 64, 128, 1, 1, 0, 3, 2, 0], + "g2s-edge-4d-shared-128x64", + g_shape=(128, 64), + g_region=((0, 128), (0, 64)), + s_shape=(2, 2, 128, 64), + s_region=((0, 1), (0, 1), (0, 128), (0, 64)), + g_layout=_plain_layout((128, 64)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (2, 2, 128, 64)).canonicalize(), ), _tma_case( - id="g2s-edge-partial-offset", - g_shape=(128, 64), g_region=((64, 64 + 24), (0, 64)), - s_shape=(2, 2, 24, 64), s_region=((0, 1), (0, 1), (0, 24), (0, 64)), - gmem_layout=TileLayout(S[128, 64]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (2, 2, 24, 64)).canonicalize(), - impl_spec=dict( - loop_extents=[1], dim=2, - coord_fn=lambda lv: [IntImm("int32", 0), IntImm("int32", 64)], - ), - encode_args=[2, 64, 128, 128, 64, 24, 1, 1, 0, 3, 2, 0], + "g2s-edge-partial-offset", + g_shape=(128, 64), + g_region=((64, 24), (0, 64)), + s_shape=(2, 2, 24, 64), + s_region=((0, 1), (0, 1), (0, 24), (0, 64)), + g_layout=_plain_layout((128, 64)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (2, 2, 24, 64)).canonicalize(), ), _tma_case( - id="g2s-edge-large-region", - g_shape=(256, 64), g_region=((128, 256), (0, 64)), - s_shape=(256, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[256, 64]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (256, 64)).canonicalize(), - impl_spec=dict( - loop_extents=[1], dim=2, - coord_fn=lambda lv: [IntImm("int32", 0), IntImm("int32", 128)], - ), - encode_args=[2, 64, 256, 128, 64, 128, 1, 1, 0, 3, 2, 0], + "g2s-edge-large-region", + g_shape=(256, 64), + g_region=((128, 128), (0, 64)), + s_shape=(256, 64), + s_region=((0, 128), (0, 64)), + g_layout=_plain_layout((256, 64)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (256, 64)).canonicalize(), ), _tma_case( - id="g2s-partial-3d-shared-a", - g_shape=(128, 256), g_region=((0, 32), (0, 64)), - s_shape=(6, 128, 64), s_region=((0, 1), (0, 32), (0, 64)), - gmem_layout=TileLayout(S[128, 256]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (6, 128, 64)).canonicalize(), - impl_spec=dict(loop_extents=[1], dim=2, coord_fn=lambda lv: _zeros(2)), - encode_args=[2, 256, 128, 512, 64, 32, 1, 1, 0, 3, 2, 0], + "g2s-partial-3d-shared-a", + g_shape=(128, 256), + g_region=((0, 32), (0, 64)), + s_shape=(6, 128, 64), + s_region=((0, 1), (0, 32), (0, 64)), + g_layout=_plain_layout((128, 256)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (6, 128, 64)).canonicalize(), ), _tma_case( - id="g2s-partial-3d-shared-b", - g_shape=(256, 512), g_region=((0, 64), (0, 64)), - s_shape=(4, 256, 64), s_region=((1, 2), (0, 64), (0, 64)), - gmem_layout=TileLayout(S[256, 512]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (4, 256, 64)).canonicalize(), - impl_spec=dict(loop_extents=[1], dim=2, coord_fn=lambda lv: _zeros(2), s_start=[1, 0, 0]), - encode_args=[2, 512, 256, 1024, 64, 64, 1, 1, 0, 3, 2, 0], + "g2s-partial-3d-shared-b", + g_shape=(256, 512), + g_region=((0, 64), (0, 64)), + s_shape=(4, 256, 64), + s_region=((1, 1), (0, 64), (0, 64)), + g_layout=_plain_layout((256, 512)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (4, 256, 64)).canonicalize(), ), _tma_case( - id="g2s-3d-full-contiguous", - g_shape=(4, 32, 64), g_region=((0, 4), (0, 32), (0, 64)), - s_shape=(4, 32, 64), s_region=((0, 4), (0, 32), (0, 64)), - gmem_layout=TileLayout(S[4, 32, 64]), - smem_layout=TileLayout(S[4, 32, 64]), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 32, 4, 128, 4096, 64, 32, 4, 1, 1, 1, 0, 0, 2, 0], + "g2s-3d-full-contiguous", + g_shape=(4, 32, 64), + g_region=((0, 4), (0, 32), (0, 64)), + s_shape=(4, 32, 64), + s_region=((0, 4), (0, 32), (0, 64)), + g_layout=_plain_layout((4, 32, 64)), + s_layout=_plain_layout((4, 32, 64)), ), _tma_case( - id="g2s-3d-partial-contiguous", - g_shape=(8, 16, 128), g_region=((0, 4), (0, 16), (0, 128)), - s_shape=(4, 16, 128), s_region=((0, 4), (0, 16), (0, 128)), - gmem_layout=TileLayout(S[8, 16, 128]), - smem_layout=TileLayout(S[4, 16, 128]), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 128, 16, 8, 256, 4096, 128, 16, 4, 1, 1, 1, 0, 0, 2, 0], + "g2s-3d-partial-contiguous", + g_shape=(8, 16, 128), + g_region=((0, 4), (0, 16), (0, 128)), + s_shape=(4, 16, 128), + s_region=((0, 4), (0, 16), (0, 128)), + g_layout=_plain_layout((8, 16, 128)), + s_layout=_plain_layout((4, 16, 128)), ), _tma_case( - id="g2s-3d-stride-gap-outer", - g_shape=(8, 32, 64), g_region=((0, 8), (0, 32), (0, 64)), - s_shape=(8, 32, 64), s_region=((0, 8), (0, 32), (0, 64)), - gmem_layout=TileLayout(S[8, 32, 64]), - smem_layout=TileLayout(S[(8, 32, 64):(4096, 64, 1)]), - impl_spec=dict( - loop_extents=[8], dim=3, - coord_fn=_stride_gap_3d_coords, elem_offset_fn=_stride_gap_elem_offset, - s_start=[0, 0, 0], - ), - encode_args=[3, 64, 32, 8, 128, 4096, 64, 32, 1, 1, 1, 1, 0, 0, 2, 0], + "g2s-3d-stride-gap-outer", + g_shape=(8, 32, 64), + g_region=((0, 8), (0, 32), (0, 64)), + s_shape=(8, 32, 64), + s_region=((0, 8), (0, 32), (0, 64)), + g_layout=_plain_layout((8, 32, 64)), + s_layout=_plain_layout((8, 32, 64), (4096, 64, 1)), ), _tma_case( - id="g2s-4d-reorder-a", - g_shape=(2, 128, 8, 64), g_region=((0, 1), (0, 128), (0, 1), (0, 64)), - s_shape=(1, 1, 128, 64), s_region=((0, 1), (0, 1), (0, 128), (0, 64)), - gmem_layout=TileLayout(S[2, 128, 8, 64]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (1, 1, 128, 64)).canonicalize(), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4), s_start=[0, 0, 0, 0]), # noqa: E501 - encode_args=[4, 64, 128, 8, 2, 1024, 128, 131072, 64, 128, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], + "g2s-4d-reorder-a", + g_shape=(2, 128, 8, 64), + g_region=((0, 1), (0, 128), (0, 1), (0, 64)), + s_shape=(1, 1, 128, 64), + s_region=((0, 1), (0, 1), (0, 128), (0, 64)), + g_layout=_plain_layout((2, 128, 8, 64)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (1, 1, 128, 64)).canonicalize(), ), _tma_case( - id="g2s-4d-reorder-b", - g_shape=(4, 64, 4, 128), g_region=((0, 1), (0, 64), (0, 1), (0, 128)), - s_shape=(1, 1, 64, 128), s_region=((0, 1), (0, 1), (0, 64), (0, 128)), - gmem_layout=TileLayout(S[4, 64, 4, 128]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (1, 1, 64, 128)).canonicalize(), - impl_spec=dict(loop_extents=[1], dim=5, coord_fn=lambda lv: _zeros(5), s_start=[0, 0, 0, 0]), # noqa: E501 - encode_args=[5, 64, 64, 2, 4, 4, 1024, 128, 256, 65536, 64, 64, 2, 1, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 + "g2s-4d-reorder-b", + g_shape=(4, 64, 4, 128), + g_region=((0, 1), (0, 64), (0, 1), (0, 128)), + s_shape=(1, 1, 64, 128), + s_region=((0, 1), (0, 1), (0, 64), (0, 128)), + g_layout=_plain_layout((4, 64, 4, 128)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (1, 1, 64, 128)).canonicalize(), ), _tma_case( - id="g2s-multidim-4d-a", - g_shape=(2, 2, 128, 64), g_region=((0, 1), (0, 1), (0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[2, 2, 128, 64]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (128, 64)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 64, 128, 2, 2, 128, 16384, 32768, 64, 128, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], + "g2s-multidim-4d-a", + g_shape=(2, 2, 128, 64), + g_region=((0, 1), (0, 1), (0, 128), (0, 64)), + s_shape=(128, 64), + s_region=((0, 128), (0, 64)), + g_layout=_plain_layout((2, 2, 128, 64)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (128, 64)), ), _tma_case( - id="g2s-multidim-4d-b", - g_shape=(4, 64, 4, 128), g_region=((0, 1), (0, 64), (0, 1), (0, 128)), - s_shape=(64, 128), s_region=((0, 64), (0, 128)), - gmem_layout=TileLayout(S[4, 64, 4, 128]).canonicalize(), - smem_layout=mma_shared_layout("float16", 3, (64, 128)), - impl_spec=dict(loop_extents=[1], dim=5, coord_fn=lambda lv: _zeros(5)), - encode_args=[5, 64, 64, 2, 4, 4, 1024, 128, 256, 65536, 64, 64, 2, 1, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 + "g2s-multidim-4d-b", + g_shape=(4, 64, 4, 128), + g_region=((0, 1), (0, 64), (0, 1), (0, 128)), + s_shape=(64, 128), + s_region=((0, 64), (0, 128)), + g_layout=_plain_layout((4, 64, 4, 128)).canonicalize(), + s_layout=mma_shared_layout("float16", 3, (64, 128)), ), - # ====================================================================== - # G2S — per-phase slices (multiphase) - # ====================================================================== _tma_case( - id="g2s-multiphase-3x8x256", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float16", 3, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 64, 8, 4, 3, 512, 128, 4096, 64, 8, 4, 1, 1, 1, 1, 1, 0, 3, 2, 0], + "g2s-multiphase-3x8x256", + g_shape=(3, 8, 256), + g_region=((0, 1), (0, 8), (0, 256)), + s_shape=(8, 256), + s_region=((0, 8), (0, 256)), + g_layout=_plain_layout((3, 8, 256)), + s_layout=mma_shared_layout("float16", 3, (8, 256)), ), _tma_case( - id="g2s-multiphase-5x64x256", - g_shape=(5, 64, 256), g_region=((0, 1), (0, 64), (0, 256)), - s_shape=(64, 256), s_region=((0, 64), (0, 256)), - gmem_layout=TileLayout(S[5, 64, 256]), - smem_layout=mma_shared_layout("float16", 3, (64, 256)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 64, 64, 4, 5, 512, 128, 32768, 64, 64, 4, 1, 1, 1, 1, 1, 0, 3, 2, 0], + "g2s-multiphase-5x64x256", + g_shape=(5, 64, 256), + g_region=((0, 1), (0, 64), (0, 256)), + s_shape=(64, 256), + s_region=((0, 64), (0, 256)), + g_layout=_plain_layout((5, 64, 256)), + s_layout=mma_shared_layout("float16", 3, (64, 256)), ), _tma_case( - id="g2s-multiphase-7x32x512-atom", - g_shape=(7, 32, 512), g_region=((0, 1), (0, 32), (0, 512)), - s_shape=(32, 512), s_region=((0, 32), (0, 512)), - gmem_layout=TileLayout(S[7, 32, 512]), - smem_layout=( + "g2s-multiphase-7x32x512-atom", + g_shape=(7, 32, 512), + g_region=((0, 1), (0, 32), (0, 512)), + s_shape=(32, 512), + s_region=((0, 32), (0, 512)), + g_layout=_plain_layout((7, 32, 512)), + s_layout=( mma_atom_layout("float16", 3) .tile_to((16, 256), mma_atom_shape("float16", 3)) .tile_to((32, 512), (16, 256)) ), - impl_spec=dict( - loop_extents=[4], dim=5, - coord_fn=_atom_multiphase_rank5_coords, elem_offset_fn=_atom_multiphase_rank5_elem_offset, # noqa: E501 - ), - encode_args=[5, 64, 8, 8, 4, 7, 1024, 128, 8192, 32768, 64, 8, 4, 2, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 ), - # ====================================================================== - # G2S — transpose-like permuted layouts - # ====================================================================== _tma_case( - id="g2s-transpose-32x64", - g_shape=(32, 64), g_region=((0, 32), (0, 64)), - s_shape=(32, 64), s_region=((0, 32), (0, 64)), - gmem_layout=TileLayout(S[32, 64]), - smem_layout=TileLayout(S[(32, 64):(1, 32)]), - impl_spec=dict( - loop_extents=[2048], dim=2, - coord_fn=lambda lv: [lv[0] % 64, lv[0] // 64], - elem_offset_fn=lambda lv: lv[0] % 64 * 32 + lv[0] // 64, - ), - encode_args=[2, 64, 32, 128, 1, 1, 1, 1, 0, 0, 2, 0], + "g2s-transpose-32x64", + g_shape=(32, 64), + g_region=((0, 32), (0, 64)), + s_shape=(32, 64), + s_region=((0, 32), (0, 64)), + g_layout=_plain_layout((32, 64)), + s_layout=_plain_layout((32, 64), (1, 32)), ), _tma_case( - id="g2s-transpose-64x32", - g_shape=(64, 32), g_region=((0, 64), (0, 32)), - s_shape=(64, 32), s_region=((0, 64), (0, 32)), - gmem_layout=TileLayout(S[64, 32]), - smem_layout=TileLayout(S[(64, 32):(1, 64)]), - impl_spec=dict( - loop_extents=[2048], dim=2, - coord_fn=lambda lv: [lv[0] % 32, lv[0] // 32], - elem_offset_fn=lambda lv: lv[0] % 32 * 64 + lv[0] // 32, - ), - encode_args=[2, 32, 64, 64, 1, 1, 1, 1, 0, 0, 2, 0], + "g2s-transpose-64x32", + g_shape=(64, 32), + g_region=((0, 64), (0, 32)), + s_shape=(64, 32), + s_region=((0, 64), (0, 32)), + g_layout=_plain_layout((64, 32)), + s_layout=_plain_layout((64, 32), (1, 64)), ), _tma_case( - id="g2s-transpose-partial-region", - g_shape=(128, 64), g_region=((0, 64), (0, 64)), - s_shape=(64, 64), s_region=((0, 64), (0, 64)), - gmem_layout=TileLayout(S[128, 64]), - smem_layout=TileLayout(S[(64, 64):(1, 64)]), - impl_spec=dict( - loop_extents=[4096], dim=2, - coord_fn=lambda lv: [lv[0] % 64, lv[0] // 64], - elem_offset_fn=lambda lv: lv[0] % 64 * 64 + lv[0] // 64, - ), - encode_args=[2, 64, 128, 128, 1, 1, 1, 1, 0, 0, 2, 0], + "g2s-transpose-partial-region", + g_shape=(128, 64), + g_region=((0, 64), (0, 64)), + s_shape=(64, 64), + s_region=((0, 64), (0, 64)), + g_layout=_plain_layout((128, 64)), + s_layout=_plain_layout((64, 64), (1, 64)), ), _tma_case( - id="g2s-transpose-partial-offset", - g_shape=(128, 64), g_region=((64, 128), (0, 32)), - s_shape=(64, 32), s_region=((0, 64), (0, 32)), - gmem_layout=TileLayout(S[128, 64]), - smem_layout=TileLayout(S[(64, 32):(1, 64)]), - impl_spec=dict( - loop_extents=[2048], dim=2, - coord_fn=lambda lv: [lv[0] % 32, lv[0] // 32 + 64], - elem_offset_fn=lambda lv: lv[0] % 32 * 64 + lv[0] // 32, - ), - encode_args=[2, 64, 128, 128, 1, 1, 1, 1, 0, 0, 2, 0], + "g2s-transpose-partial-offset", + g_shape=(128, 64), + g_region=((64, 64), (0, 32)), + s_shape=(64, 32), + s_region=((0, 64), (0, 32)), + g_layout=_plain_layout((128, 64)), + s_layout=_plain_layout((64, 32), (1, 64)), ), - # ====================================================================== - # G2S — non-prefix compact (4D gmem collapses to one TMA tile) - # ====================================================================== _tma_case( - id="g2s-non-prefix-compact-elides", - g_shape=(16, 16, 128, 128), g_region=((3, 4), (4, 5), (0, 128), (0, 128)), - s_shape=(128, 128), s_region=((0, 128), (0, 128)), - gmem_layout=TileLayout(S[(16, 16, 128, 128):(1024 * 128, 128, 1024, 1)]), - smem_layout=TileLayout(S[128, 128]), - impl_spec=dict( - loop_extents=[1], dim=4, - coord_fn=lambda lv: [ - IntImm("int32", 0), IntImm("int32", 0), - IntImm("int32", 4), IntImm("int32", 3), - ], - ), - encode_args=[4, 128, 128, 16, 16, 2048, 256, 262144, 128, 128, 1, 1, 1, 1, 1, 1, 0, 0, 2, 0], # noqa: E501 + "g2s-non-prefix-compact-elides", + g_shape=(16, 16, 128, 128), + g_region=((3, 1), (4, 1), (0, 128), (0, 128)), + s_shape=(128, 128), + s_region=((0, 128), (0, 128)), + g_layout=_plain_layout((16, 16, 128, 128), (1024 * 128, 128, 1024, 1)), + s_layout=_plain_layout((128, 128)), ), - # ====================================================================== - # G2S — oob contract (config={"oob": ...}); fill kind is encoded in - # encode_args[-1]. ``None`` and ``"zero"`` both map to fill_kind=0. - # ====================================================================== _tma_case( - id="g2s-oob-zero", - g_shape=(128, 64), g_region=((120, 136), (0, 64)), - s_shape=(16, 64), s_region=((0, 16), (0, 64)), - gmem_layout=TileLayout(S[128, 64]), - smem_layout=mma_shared_layout("float16", 3, (16, 64)), + "g2s-oob-zero", + g_shape=(128, 64), + g_region=((120, 16), (0, 64)), + s_shape=(16, 64), + s_region=((0, 16), (0, 64)), + g_layout=_plain_layout((128, 64)), + s_layout=mma_shared_layout("float16", 3, (16, 64)), config={"oob": "zero"}, - impl_spec=dict( - loop_extents=[1], dim=2, - coord_fn=lambda lv: [IntImm("int32", 0), IntImm("int32", 120)], - ), - encode_args=[2, 64, 128, 128, 64, 16, 1, 1, 0, 3, 2, 0], ), _tma_case( - id="g2s-oob-nan", - g_shape=(128, 64), g_region=((120, 136), (0, 64)), - s_shape=(16, 64), s_region=((0, 16), (0, 64)), - gmem_layout=TileLayout(S[128, 64]), - smem_layout=mma_shared_layout("float16", 3, (16, 64)), + "g2s-oob-nan", + g_shape=(128, 64), + g_region=((120, 16), (0, 64)), + s_shape=(16, 64), + s_region=((0, 16), (0, 64)), + g_layout=_plain_layout((128, 64)), + s_layout=mma_shared_layout("float16", 3, (16, 64)), config={"oob": "nan"}, - impl_spec=dict( - loop_extents=[1], dim=2, - coord_fn=lambda lv: [IntImm("int32", 0), IntImm("int32", 120)], - ), - encode_args=[2, 64, 128, 128, 64, 16, 1, 1, 0, 3, 2, 1], - ), - # ====================================================================== - # G2S — flash_attention4 Q/K/V regression baselines - # Representative config: batch=1, seq_len=2048, num_qo_heads=32, - # num_kv_heads=8, head_dim=128 → GQA_RATIO=4, SEQ_Q_PER_TILE=32, - # BLK_M=BLK_N=128, SMEM_PIPE_DEPTH_Q=2, SMEM_PIPE_DEPTH_KV=3. Each case - # lowers to exactly one cp_async_bulk_tensor; structural golden locks - # rank / shape / coord / box. - # ====================================================================== - _tma_case( - id="g2s-fa4-q", - g_shape=(1, 2048, 32, 128), g_region=((0, 1), (0, 32), (0, 4), (0, 128)), - s_shape=(2, 128, 128), s_region=((0, 1), (0, 128), (0, 128)), - gmem_layout=TileLayout(S[1, 2048, 32, 128]), - smem_layout=mma_shared_layout("float16", 3, (2, 128, 128)), - impl_spec=dict(loop_extents=[1], dim=5, coord_fn=lambda lv: _zeros(5)), - encode_args=[5, 64, 32, 2048, 2, 1, 256, 8192, 128, 0, 64, 4, 32, 2, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 ), _tma_case( - id="g2s-fa4-k", - g_shape=(1, 2048, 8, 128), g_region=((0, 1), (0, 128), (0, 1), (0, 128)), - s_shape=(3, 128, 128), s_region=((0, 1), (0, 128), (0, 128)), - gmem_layout=TileLayout(S[1, 2048, 8, 128]), - smem_layout=mma_shared_layout("float16", 3, (3, 128, 128)), - impl_spec=dict(loop_extents=[1], dim=5, coord_fn=lambda lv: _zeros(5)), - encode_args=[5, 64, 2048, 2, 8, 1, 2048, 128, 256, 0, 64, 128, 2, 1, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 + "g2s-fa4-q", + g_shape=(1, 2048, 32, 128), + g_region=((0, 1), (0, 32), (0, 4), (0, 128)), + s_shape=(2, 128, 128), + s_region=((0, 1), (0, 128), (0, 128)), + g_layout=TileLayout(S[(1, 2048, 32, 128)]), + s_layout=mma_shared_layout("float16", 3, (2, 128, 128)), ), _tma_case( - id="g2s-fa4-v", - g_shape=(1, 2048, 8, 128), g_region=((0, 1), (0, 128), (0, 1), (0, 128)), - s_shape=(3, 128, 128), s_region=((0, 1), (0, 128), (0, 128)), - gmem_layout=TileLayout(S[1, 2048, 8, 128]), - smem_layout=mma_shared_layout("float16", 3, (3, 128, 128)), - impl_spec=dict(loop_extents=[1], dim=5, coord_fn=lambda lv: _zeros(5)), - encode_args=[5, 64, 2048, 2, 8, 1, 2048, 128, 256, 0, 64, 128, 2, 1, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 + "g2s-fa4-k", + g_shape=(1, 2048, 8, 128), + g_region=((0, 1), (0, 128), (0, 1), (0, 128)), + s_shape=(3, 128, 128), + s_region=((0, 1), (0, 128), (0, 128)), + g_layout=TileLayout(S[(1, 2048, 8, 128)]), + s_layout=mma_shared_layout("float16", 3, (3, 128, 128)), ), - # ====================================================================== - # S2G — per-phase slices (swizzle + dtype variants) - # ====================================================================== _tma_case( - id="s2g-multiphase-3x8x256", - direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float16", 3, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 64, 8, 4, 3, 512, 128, 4096, 64, 8, 4, 1, 1, 1, 1, 1, 0, 3, 2, 0], + "g2s-fa4-v", + g_shape=(1, 2048, 8, 128), + g_region=((0, 1), (0, 128), (0, 1), (0, 128)), + s_shape=(3, 128, 128), + s_region=((0, 1), (0, 128), (0, 128)), + g_layout=TileLayout(S[(1, 2048, 8, 128)]), + s_layout=mma_shared_layout("float16", 3, (3, 128, 128)), ), + *[ + _tma_case( + case_id, + g_shape=(3, 8, 256), + g_region=((0, 1), (0, 8), (0, 256)), + s_shape=(8, 256), + s_region=((0, 8), (0, 256)), + g_layout=_plain_layout((3, 8, 256)), + s_layout=mma_shared_layout(dtype, swizzle, (8, 256)), + dtype=dtype, + direction="s2g", + ) + for case_id, dtype, swizzle in [ + ("s2g-multiphase-3x8x256", "float16", 3), + ("s2g-multiphase-3x8x256-swizzle2", "float16", 2), + ("s2g-multiphase-3x8x256-swizzle0", "float16", 0), + ("s2g-multiphase-3x8x256-int8", "int8", 3), + ("s2g-multiphase-3x8x256-fp32", "float32", 3), + ] + ], _tma_case( - id="s2g-multiphase-5x64x256", + "s2g-multiphase-5x64x256", + g_shape=(5, 64, 256), + g_region=((0, 1), (0, 64), (0, 256)), + s_shape=(64, 256), + s_region=((0, 64), (0, 256)), + g_layout=_plain_layout((5, 64, 256)), + s_layout=mma_shared_layout("float16", 3, (64, 256)), direction="s2g", - g_shape=(5, 64, 256), g_region=((0, 1), (0, 64), (0, 256)), - s_shape=(64, 256), s_region=((0, 64), (0, 256)), - gmem_layout=TileLayout(S[5, 64, 256]), - smem_layout=mma_shared_layout("float16", 3, (64, 256)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 64, 64, 4, 5, 512, 128, 32768, 64, 64, 4, 1, 1, 1, 1, 1, 0, 3, 2, 0], ), _tma_case( - id="s2g-multiphase-7x32x512-atom", - direction="s2g", - g_shape=(7, 32, 512), g_region=((0, 1), (0, 32), (0, 512)), - s_shape=(32, 512), s_region=((0, 32), (0, 512)), - gmem_layout=TileLayout(S[7, 32, 512]), - smem_layout=( + "s2g-multiphase-7x32x512-atom", + g_shape=(7, 32, 512), + g_region=((0, 1), (0, 32), (0, 512)), + s_shape=(32, 512), + s_region=((0, 32), (0, 512)), + g_layout=_plain_layout((7, 32, 512)), + s_layout=( mma_atom_layout("float16", 3) .tile_to((16, 256), mma_atom_shape("float16", 3)) .tile_to((32, 512), (16, 256)) ), - impl_spec=dict( - loop_extents=[4], dim=5, - coord_fn=_atom_multiphase_rank5_coords, elem_offset_fn=_atom_multiphase_rank5_elem_offset, # noqa: E501 - ), - encode_args=[5, 64, 8, 8, 4, 7, 1024, 128, 8192, 32768, 64, 8, 4, 2, 1, 1, 1, 1, 1, 1, 0, 3, 2, 0], # noqa: E501 - ), - _tma_case( - id="s2g-multiphase-3x8x256-swizzle2", - direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float16", 2, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 32, 8, 8, 3, 512, 64, 4096, 32, 8, 8, 1, 1, 1, 1, 1, 0, 2, 2, 0], - ), - _tma_case( - id="s2g-multiphase-3x8x256-swizzle0", - direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float16", 0, (8, 256)), - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 256, 8, 3, 512, 4096, 256, 8, 1, 1, 1, 1, 0, 0, 2, 0], - ), - _tma_case( - id="s2g-multiphase-3x8x256-int8", direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("int8", 3, (8, 256)), - dtype="int8", - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 128, 8, 2, 3, 256, 128, 2048, 128, 8, 2, 1, 1, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="s2g-multiphase-3x8x256-fp32", - direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float32", 3, (8, 256)), - dtype="float32", - impl_spec=dict(loop_extents=[1], dim=4, coord_fn=lambda lv: _zeros(4)), - encode_args=[4, 32, 8, 8, 3, 1024, 128, 8192, 32, 8, 8, 1, 1, 1, 1, 1, 0, 3, 2, 0], ), - # ====================================================================== - # S2G — retain multi-dim coords without linear-carry (bf16, custom layout) - # ====================================================================== _tma_case( - id="s2g-keeps-multidim-coords", - direction="s2g", - g_shape=(1024, 4, 1024), g_region=((128, 128 + 128), (1, 1 + 1), (32, 32 + 32)), - s_shape=(128, 32), s_region=((0, 128), (0, 32)), - gmem_layout=TileLayout(S[(1024, 4, 1024):(4 * 1024, 1024, 1)]), - smem_layout=TileLayout(S[(128, 32):(32, 1)]), + "s2g-keeps-multidim-coords", + g_shape=(1024, 4, 1024), + g_region=((128, 128), (1, 1), (32, 32)), + s_shape=(128, 32), + s_region=((0, 128), (0, 32)), + g_layout=_plain_layout((1024, 4, 1024), (4 * 1024, 1024, 1)), + s_layout=_plain_layout((128, 32), (32, 1)), dtype="bfloat16", - impl_spec=dict( - loop_extents=[1], dim=3, - coord_fn=lambda lv: [ - IntImm("int32", 32), - IntImm("int32", 128), - IntImm("int32", 1), - ], - ), - ), - # ====================================================================== - # S2G — oob contract variants over the same (2, 128, 64) shape. ``None`` - # and ``"zero"`` map to fill_kind=0; ``"nan"`` maps to fill_kind=1. The - # descriptor geometry is identical across the three variants. - # ====================================================================== - _tma_case( - id="s2g-oob-none", direction="s2g", - g_shape=(2, 128, 64), g_region=((0, 1), (0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[(2, 128, 64)]), - smem_layout=mma_shared_layout("float16", 3, (128, 64)), - config=None, - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 128, 2, 128, 16384, 64, 128, 1, 1, 1, 1, 0, 3, 2, 0], ), + *[ + _tma_case( + case_id, + g_shape=(2, 128, 64), + g_region=((0, 1), (0, 128), (0, 64)), + s_shape=(128, 64), + s_region=((0, 128), (0, 64)), + g_layout=_plain_layout((2, 128, 64)), + s_layout=mma_shared_layout("float16", 3, (128, 64)), + direction="s2g", + config=config, + ) + for case_id, config in [ + ("s2g-oob-none", None), + ("s2g-oob-zero", {"oob": "zero"}), + ("s2g-oob-nan", {"oob": "nan"}), + ] + ], _tma_case( - id="s2g-oob-zero", - direction="s2g", - g_shape=(2, 128, 64), g_region=((0, 1), (0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[(2, 128, 64)]), - smem_layout=mma_shared_layout("float16", 3, (128, 64)), - config={"oob": "zero"}, - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 128, 2, 128, 16384, 64, 128, 1, 1, 1, 1, 0, 3, 2, 0], - ), - _tma_case( - id="s2g-oob-nan", - direction="s2g", - g_shape=(2, 128, 64), g_region=((0, 1), (0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[(2, 128, 64)]), - smem_layout=mma_shared_layout("float16", 3, (128, 64)), - config={"oob": "nan"}, - impl_spec=dict(loop_extents=[1], dim=3, coord_fn=lambda lv: _zeros(3)), - encode_args=[3, 64, 128, 2, 128, 16384, 64, 128, 1, 1, 1, 1, 0, 3, 2, 1], - ), - # ====================================================================== - # Rejection cases — oob contract validation - # ====================================================================== - _tma_case( - id="reject-unknown-oob", - direction="s2g", - g_shape=(3, 8, 256), g_region=((0, 1), (0, 8), (0, 256)), - s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[3, 8, 256]), - smem_layout=mma_shared_layout("float16", 3, (8, 256)), + "reject-unknown-oob", + g_shape=(3, 8, 256), + g_region=((0, 1), (0, 8), (0, 256)), + s_shape=(8, 256), + s_region=((0, 8), (0, 256)), + g_layout=_plain_layout((3, 8, 256)), + s_layout=mma_shared_layout("float16", 3, (8, 256)), config={"oob": "bogus"}, - raises=(Exception, "Unsupported TMA oob mode"), ), _tma_case( - id="reject-g2s-nan-on-non-float", - g_shape=(128, 64), g_region=((120, 136), (0, 64)), - s_shape=(16, 64), s_region=((0, 16), (0, 64)), - gmem_layout=TileLayout(S[128, 64]), - smem_layout=TileLayout(S[16, 64]), + "reject-g2s-nan-on-non-float", + g_shape=(128, 64), + g_region=((120, 16), (0, 64)), + s_shape=(16, 64), + s_region=((0, 16), (0, 64)), + g_layout=_plain_layout((128, 64)), + s_layout=_plain_layout((16, 64)), dtype="int8", config={"oob": "nan"}, - raises=(Exception, "requires a floating-point dtype"), ), _tma_case( - id="reject-s2g-nan-on-non-float", - direction="s2g", - g_shape=(2, 128, 64), g_region=((0, 1), (0, 128), (0, 64)), - s_shape=(128, 64), s_region=((0, 128), (0, 64)), - gmem_layout=TileLayout(S[2, 128, 64]), - smem_layout=TileLayout(S[128, 64]), + "reject-s2g-nan-on-non-float", + g_shape=(2, 128, 64), + g_region=((0, 1), (0, 128), (0, 64)), + s_shape=(128, 64), + s_region=((0, 128), (0, 64)), + g_layout=_plain_layout((2, 128, 64)), + s_layout=_plain_layout((128, 64)), dtype="int8", + direction="s2g", config={"oob": "nan"}, - raises=(Exception, "requires a floating-point dtype"), ), ] -# fmt: on -@pytest.mark.gpu -@pytest.mark.parametrize("case", TMA_CASES) -def test_copy_tma_codegen(case): - """Unified structural-golden driver for every TMA unit test case. - - See ``_tma_case`` for the dict-form input. When ``raises`` is set, the - test expects ``_make_tma_call`` to raise; otherwise it compares the - emitted device impl and host tensormap-init against the inlined - ``impl_spec`` / ``encode_args`` goldens. - """ - call_kwargs = dict( - g_shape=case["g_shape"], - g_region=case["g_region"], - s_shape=case["s_shape"], - s_region=case["s_region"], - gmem_layout=case["gmem_layout"], - smem_layout=case["smem_layout"], - dtype=case["dtype"], - direction=case["direction"], - config=case["config"], +@dataclass(frozen=True) +class _TMAGolden: + dtype: str + dims: tuple + strides: tuple + boxes: tuple + coordinates: tuple + enums: tuple + + +def _tma_golden( + dtype, + dims, + strides, + boxes, + coordinates=None, + enums=(0, 3, 2, 0), +): + return _TMAGolden( + dtype=dtype, + dims=dims, + strides=strides, + boxes=boxes, + coordinates=(0,) * len(dims) if coordinates is None else coordinates, + enums=enums, ) - if case["raises"] is not None: - exc, match = case["raises"] - with pytest.raises(exc, match=match): - _make_tma_call(**call_kwargs) + + +_TMA_CASE_GOLDENS = { + "g2s-2d-8x256": _tma_golden("float16", (64, 8, 4), (512, 128), (64, 8, 4)), + "g2s-2d-8x256-swizzle2": _tma_golden( + "float16", (32, 8, 8), (512, 64), (32, 8, 8), enums=(0, 2, 2, 0) + ), + "g2s-2d-8x256-swizzle1": _tma_golden( + "float16", (16, 8, 16), (512, 32), (16, 8, 16), enums=(0, 1, 2, 0) + ), + "g2s-2d-8x256-swizzle0": _tma_golden( + "float16", (8, 8, 32), (512, 16), (8, 8, 32), enums=(0, 0, 2, 0) + ), + "g2s-2d-8x256-int8": _tma_golden("int8", (128, 8, 2), (256, 128), (128, 8, 2)), + "g2s-2d-8x256-bf16": _tma_golden("bfloat16", (64, 8, 4), (512, 128), (64, 8, 4)), + "g2s-2d-8x256-fp32": _tma_golden("float32", (32, 8, 8), (1024, 128), (32, 8, 8)), + "g2s-2d-8x256-uint8": _tma_golden("uint8", (128, 8, 2), (256, 128), (128, 8, 2)), + "g2s-2d-8x256-fp8e4m3": _tma_golden("float8_e4m3fn", (128, 8, 2), (256, 128), (128, 8, 2)), + "g2s-2d-8x256-fp8e5m2": _tma_golden("float8_e5m2", (128, 8, 2), (256, 128), (128, 8, 2)), + "g2s-3d-shared-64x256": _tma_golden("float16", (64, 64, 4), (512, 128), (64, 64, 4)), + "g2s-2d-partial-8192": _tma_golden("float16", (8192, 8192), (16384,), (64, 128)), + "g2s-edge-4d-shared-128x64": _tma_golden("float16", (64, 128), (128,), (64, 128)), + "g2s-edge-partial-offset": _tma_golden("float16", (64, 128), (128,), (64, 24), (0, 64)), + "g2s-edge-large-region": _tma_golden("float16", (64, 256), (128,), (64, 128), (0, 128)), + "g2s-partial-3d-shared-a": _tma_golden("float16", (256, 128), (512,), (64, 32)), + "g2s-partial-3d-shared-b": _tma_golden("float16", (512, 256), (1024,), (64, 64)), + "g2s-3d-full-contiguous": _tma_golden( + "float16", (64, 128), (128,), (64, 128), enums=(0, 0, 2, 0) + ), + "g2s-3d-partial-contiguous": _tma_golden( + "float16", (128, 128), (256,), (128, 64), enums=(0, 0, 2, 0) + ), + "g2s-4d-reorder-a": _tma_golden("float16", (64, 256, 8), (1024, 128), (64, 128, 1)), + "g2s-4d-reorder-b": _tma_golden( + "float16", + (64, 64, 2, 4, 4), + (1024, 128, 65536, 256), + (64, 64, 2, 1, 1), + ), + "g2s-multidim-4d-a": _tma_golden( + "float16", (64, 128, 2, 2), (128, 32768, 16384), (64, 128, 1, 1) + ), + "g2s-multidim-4d-b": _tma_golden( + "float16", + (64, 64, 2, 4, 4), + (1024, 128, 65536, 256), + (64, 64, 2, 1, 1), + ), + "g2s-multiphase-3x8x256": _tma_golden( + "float16", (64, 8, 4, 3), (512, 128, 4096), (64, 8, 4, 1) + ), + "g2s-multiphase-5x64x256": _tma_golden( + "float16", (64, 64, 4, 5), (512, 128, 32768), (64, 64, 4, 1) + ), + "g2s-non-prefix-compact-elides": _tma_golden( + "float16", + (128, 2048, 16), + (2048, 256), + (128, 128, 1), + (0, 384, 4), + (0, 0, 2, 0), + ), + "g2s-oob-zero": _tma_golden("float16", (64, 128), (128,), (64, 16), (0, 120)), + "g2s-oob-nan": _tma_golden("float16", (64, 128), (128,), (64, 16), (0, 120), (0, 3, 2, 1)), + "g2s-fa4-q": _tma_golden("float16", (64, 32, 2048, 2), (256, 8192, 128), (64, 4, 32, 2)), + "g2s-fa4-k": _tma_golden("float16", (64, 2048, 16), (2048, 128), (64, 128, 2)), + "g2s-fa4-v": _tma_golden("float16", (64, 2048, 16), (2048, 128), (64, 128, 2)), + "s2g-multiphase-3x8x256": _tma_golden( + "float16", (64, 8, 4, 3), (512, 128, 4096), (64, 8, 4, 1) + ), + "s2g-multiphase-3x8x256-swizzle2": _tma_golden( + "float16", + (32, 8, 8, 3), + (512, 64, 4096), + (32, 8, 8, 1), + enums=(0, 2, 2, 0), + ), + "s2g-multiphase-3x8x256-swizzle0": _tma_golden( + "float16", + (8, 8, 32, 3), + (512, 16, 4096), + (8, 8, 32, 1), + enums=(0, 0, 2, 0), + ), + "s2g-multiphase-3x8x256-int8": _tma_golden( + "int8", (128, 8, 2, 3), (256, 128, 2048), (128, 8, 2, 1) + ), + "s2g-multiphase-3x8x256-fp32": _tma_golden( + "float32", (32, 8, 8, 3), (1024, 128, 8192), (32, 8, 8, 1) + ), + "s2g-multiphase-5x64x256": _tma_golden( + "float16", (64, 64, 4, 5), (512, 128, 32768), (64, 64, 4, 1) + ), + "s2g-keeps-multidim-coords": _tma_golden( + "bfloat16", + (1024, 1024, 4), + (8192, 2048), + (32, 128, 1), + (32, 128, 1), + (0, 0, 2, 0), + ), + "s2g-oob-none": _tma_golden("float16", (64, 256), (128,), (64, 128)), +} + + +_TMA_EXPLICIT_CASES = { + "g2s-oob-zero", + "g2s-oob-nan", + "s2g-oob-zero", + "s2g-oob-nan", + "reject-unknown-oob", + "reject-g2s-nan-on-non-float", + "reject-s2g-nan-on-non-float", +} + + +_TMA_CASE_ERRORS = { + "g2s-2d-32x512-atom": r"stage=prefix-search: rank: .*got 6", + "g2s-3d-stride-gap-outer": r"stage=shared-chain:", + "g2s-multiphase-7x32x512-atom": r"stage=prefix-search: rank: .*got 6", + "g2s-transpose-32x64": r"stage=prefix-search: global_stride_alignment:", + "g2s-transpose-64x32": r"stage=prefix-search: global_stride_alignment:", + "g2s-transpose-partial-region": r"stage=prefix-search: global_stride_alignment:", + "g2s-transpose-partial-offset": r"stage=prefix-search: global_stride_alignment:", + "s2g-multiphase-7x32x512-atom": r"stage=prefix-search: rank: .*got 6", + "s2g-oob-zero": r"oob is only valid for explicit global-to-shared", + "s2g-oob-nan": r"oob is only valid for explicit global-to-shared", + "reject-unknown-oob": r"unsupported TensorMap oob='bogus'", + "reject-g2s-nan-on-non-float": r"nan_oob_dtype:", + "reject-s2g-nan-on-non-float": r"oob is only valid for explicit global-to-shared", +} + + +def test_tma_case_matrix_is_complete(): + case_ids = [case.values[0] for case in TMA_CASES] + expected_ids = set(_TMA_CASE_GOLDENS) | set(_TMA_CASE_ERRORS) + assert len(case_ids) == 52 + assert len(set(case_ids)) == len(case_ids) + assert set(case_ids) == expected_ids + assert not set(_TMA_CASE_GOLDENS) & set(_TMA_CASE_ERRORS) + assert _TMA_EXPLICIT_CASES <= expected_ids + + +@pytest.mark.parametrize(("case_id", "kwargs"), TMA_CASES) +def test_tma_codegen(case_id, kwargs): + variant = "tma_explicit" if case_id in _TMA_EXPLICIT_CASES else "tma_auto" + if case_id in _TMA_CASE_ERRORS: + with pytest.raises(Exception, match=_TMA_CASE_ERRORS[case_id]): + _lower_direct(variant, **kwargs) return - impl, host_init_stmts = _make_tma_call(**call_kwargs) - if case["impl_spec"] is not None: - expected_impl = _build_expected_impl( - case["direction"], - case["dtype"], - case["s_shape"], - case["smem_layout"], - case["impl_spec"], + expected = _TMA_CASE_GOLDENS[case_id] + plan = _direct_plan(variant, **kwargs) + spec = plan.spec + assert spec.descriptor_dtype == expected.dtype + assert _ints(spec.global_dims) == expected.dims + assert _ints(spec.global_strides) == expected.strides + assert _ints(spec.box_dims) == expected.boxes + assert _ints(spec.coordinates) == expected.coordinates + assert _ints(spec.element_strides) == (1,) * len(expected.dims) + assert (spec.interleave, spec.swizzle, spec.l2_promotion, spec.oob_fill) == expected.enums + assert plan.issue_axes == () + + impl, host, _ = _lower_direct(variant, **kwargs) + counter = _count_tma(impl) + assert counter.total == 1 + assert len(counter.calls) == 1 + call = counter.calls[0] + assert int(call.args[0]) == len(expected.dims) + coord_start = 11 if kwargs.get("direction", "g2s") == "g2s" else 5 + assert _ints(call.args[coord_start:]) == expected.coordinates + + encodes = _collect_encodes(host) + assert len(encodes) == 1 + signature = _encode_signature(encodes[0]) + assert signature["dtype"] == expected.dtype + assert _ints(signature["dims"]) == expected.dims + assert _ints(signature["strides"]) == expected.strides + assert _ints(signature["boxes"]) == expected.boxes + assert _ints(signature["element_strides"]) == (1,) * len(expected.dims) + assert _ints(signature["enums"]) == expected.enums + + +def test_auto_canonicalization_does_not_create_an_oversized_box(): + plan = _auto_plan(g_shape=(8, 64)) + assert plan.spec.descriptor_dtype == "float16" + assert plan.spec.rank == 2 + assert _ints(plan.spec.global_dims) == (64, 8) + assert _ints(plan.spec.box_dims) == (64, 8) + assert plan.issue_axes == () + + +def test_auto_sw128_gt_grouping_and_cuda_order(): + shape = (8, 256) + impl, host, _ = _lower_direct( + "tma_auto", + g_shape=shape, + s_layout=mma_shared_layout("float16", 3, shape), + ) + assert _count_tma(impl).total == 1 + signature = _encode_signature(_collect_encodes(host)[0]) + assert signature["dtype"] == "float16" + assert signature["rank"] == 3 + assert _ints(signature["dims"]) == (64, 8, 4) + assert _ints(signature["strides"]) == (512, 128) + assert _ints(signature["boxes"]) == (64, 8, 4) + assert _ints(signature["enums"]) == (0, 3, 2, 0) + + +def test_auto_retains_coordinate_only_global_dimension(): + g_shape = (64, 128, 2, 4, 4096) + s_shape = (64, 64, 2, 4) + plan = _auto_plan( + g_shape=g_shape, + s_shape=s_shape, + g_region=((0, 64), (32, 64), (0, 2), (0, 4), (7, 1)), + g_layout=_plain_layout(g_shape, (1, 512, 256, 64, 65536)), + s_layout=_plain_layout(s_shape, (1, 64, 4096, 8192)), + ) + assert plan.spec.rank == 5 + assert _ints(plan.spec.global_dims) == g_shape + assert _ints(plan.spec.global_strides) == (1024, 512, 128, 131072) + assert _ints(plan.spec.box_dims) == (64, 64, 2, 4, 1) + assert _ints(plan.spec.coordinates) == (0, 32, 0, 0, 7) + assert plan.issue_axes == () + + +def test_auto_retains_unit_copy_dimensions_with_coordinates_and_strides(): + plan = _auto_plan( + g_shape=(4, 256, 8, 512), + s_shape=(1, 128, 1, 128), + g_region=((2, 1), (64, 128), (3, 1), (128, 128)), + ) + assert plan.spec.rank == 4 + assert _ints(plan.spec.global_dims) == (512, 256, 4, 8) + assert _ints(plan.spec.global_strides) == (8192, 2097152, 1024) + assert _ints(plan.spec.box_dims) == (128, 128, 1, 1) + assert _ints(plan.spec.coordinates) == (128, 64, 2, 3) + assert int(plan.spec.base_byte_offset) == 0 + assert plan.issue_axes == () + + +def test_auto_canonicalizes_dynamic_stage_slice_before_global_grouping(): + kv_row = Var("kv_row", "int32") + head = Var("head", "int32") + stage = Var("stage", "int32") + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + sctx = DispatchContext( + target, + ExecScope("thread"), + {}, + { + kv_row: Range.from_min_extent(0, 3969), + head: Range.from_min_extent(0, 16), + stage: Range.from_min_extent(0, 3), + }, + ) + plan = _auto_plan( + g_shape=(1, 4096, 16, 128), + s_shape=(3, 128, 128), + g_region=((0, 1), (kv_row, 128), (head, 1), (0, 128)), + s_region=((stage, 1), (0, 128), (0, 128)), + g_layout=TileLayout(S[(1, 4096, 16, 128)]), + s_layout=TileLayout(S[(3, 128, 2, 64) : (16384, 64, 8192, 1)]), + sctx=sctx, + ) + assert plan.spec.rank == 3 + assert _ints(plan.spec.global_dims) == (64, 4096, 32) + assert _ints(plan.spec.global_strides) == (4096, 128) + assert _ints(plan.spec.box_dims) == (64, 128, 2) + analyzer = Analyzer() + assert analyzer.can_prove_equal(plan.spec.coordinates[0], 0) + assert analyzer.can_prove_equal(plan.spec.coordinates[1], kv_row) + assert analyzer.can_prove_equal(plan.spec.coordinates[2], head * 2) + assert analyzer.can_prove_equal(plan.spec.smem_base_offset, stage * 16384) + assert int(plan.spec.transaction_bits) // 8 == 32 * 1024 + assert plan.issue_axes == () + + +def test_auto_shared_pointer_counts_each_offset_once(): + warp_offset = Var("warp_offset", "int32") + plan = _auto_plan( + g_shape=(2, 64), + s_shape=(4, 64), + s_region=((2, 2), (0, 64)), + s_layout=TileLayout(S[(4, 64)] + warp_offset * 64), + s_elem_offset=32, + ) + assert Analyzer().can_prove_equal( + plan.spec.smem_base_offset, + 32 + warp_offset * 64 + 128, + ) + + +def test_auto_defers_dynamic_global_dimension_bounds_to_runtime(): + seq_len = Var("seq_len", "uint32") + row = Var("row", "uint32") + plan = _auto_plan( + g_shape=(seq_len * T.uint32(64), 64), + s_shape=(128, 64), + g_region=((row, 128), (0, 64)), + g_layout=TileLayout(S[(seq_len * T.uint32(64), 64)]), + s_layout=TileLayout(S[(128, 64)]), + dtype="uint8", + ) + assert plan.spec.rank == 2 + assert Analyzer().can_prove_equal(plan.spec.global_dims[0], 64) + assert Analyzer().can_prove_equal(plan.spec.global_dims[1], seq_len * T.uint32(64)) + assert _ints(plan.spec.box_dims) == (64, 128) + assert Analyzer().can_prove_equal(plan.spec.coordinates[0], 0) + assert Analyzer().can_prove_equal(plan.spec.coordinates[1], row) + assert plan.issue_axes == () + + +def test_dispatch_propagates_flat_bind_to_auto_coordinate_proof(): + func = _from_source( + """ +@T.prim_func +def bind_coordinate(D_ptr: T.handle): + D = T.match_buffer(D_ptr, (33360, 6144), "bfloat16") + T.device_entry() + block = T.cta_id([192]) + tid = T.thread_id([1]) + tile_index = T.alloc_local((1,), "int32") + D_smem = T.alloc_buffer( + (2, 16, 128), + "bfloat16", + scope="shared.dyn", + layout=T.TileLayout(T.S[(2, 16, 2, 64) : (2048, 64, 1024, 1)]), + ) + tile_index[0] = block + d_n: T.let = tile_index[0] * 128 + if tid == 0: + Tx.copy_async( + D[0:16, d_n : d_n + 128], + D_smem[0, :, :], + dispatch="tma_auto", ) - tvm.ir.assert_structural_equal(impl, expected_impl, map_free_vars=True) - if case["encode_args"] is not None: - expected_host = _build_expected_host_init(case["dtype"], case["encode_args"]) - assert len(host_init_stmts) == 1 - tvm.ir.assert_structural_equal(host_init_stmts[0], expected_host, map_free_vars=True) +""" + ) + lowered = _lower_module(func) + encodes = _collect_encodes([lowered["main"].body]) + assert len(encodes) == 1 + signature = _encode_signature(encodes[0]) + assert _ints(signature["dims"]) == (64, 33360, 96) + assert _ints(signature["strides"]) == (12288, 128) + assert _ints(signature["boxes"]) == (64, 16, 2) + + counter = _count_tma(lowered["main"]) + assert counter.total == 1 + call = counter.calls[0] + assert int(call.args[0]) == 3 + assert int(call.args[5]) == 0 + assert int(call.args[6]) == 0 + assert str(call.args[7]) == "tile_index[0] * 2" + + +def test_auto_recovers_stride_from_split_coordinate_only_dimension(): + plan = _auto_plan( + g_shape=(24576, 7168), + s_shape=(128, 128), + g_region=((128, 128), (0, 128)), + g_layout=TileLayout(S[(4, 6144, 7168) : (44040192, 7168, 1)]), + ) + assert plan.spec.rank == 2 + assert _ints(plan.spec.global_dims) == (7168, 24576) + assert _ints(plan.spec.global_strides) == (14336,) + assert _ints(plan.spec.box_dims) == (128, 128) + assert _ints(plan.spec.coordinates) == (0, 128) + assert plan.issue_axes == () + + +def test_auto_preserves_coordinate_stride_for_interleaved_tensor_dimensions(): + plan = _auto_plan( + g_shape=(2048, 128), + s_shape=(5, 128, 16), + g_region=((128, 128), (16, 16)), + s_region=((0, 1), (0, 128), (0, 16)), + g_layout=TileLayout(S[(16, 4, 32, 32, 4) : (16384, 4, 16, 512, 1)]), + s_layout=TileLayout(S[(5, 4, 32, 4, 4) : (2048, 4, 16, 512, 1)]), + dtype="uint8", + ) + assert plan.spec.descriptor_dtype == "uint16" + assert plan.spec.rank == 3 + assert _ints(plan.spec.global_dims) == (256, 32, 16) + assert _ints(plan.spec.global_strides) == (512, 16384) + assert _ints(plan.spec.box_dims) == (256, 4, 1) + assert _ints(plan.spec.coordinates) == (0, 4, 1) + assert plan.issue_axes == () + + +def test_auto_recovers_leading_unit_dimension_stride(): + plan = _auto_plan( + g_shape=(1, 8192), + s_shape=(256,), + g_region=((0, 1), (512, 256)), + g_layout=_plain_layout((1, 8192), (8192, 1)), + ) + assert plan.spec.rank == 1 + assert _ints(plan.spec.global_dims) == (8192,) + assert _ints(plan.spec.global_strides) == () + assert _ints(plan.spec.box_dims) == (256,) + assert _ints(plan.spec.coordinates) == (512,) + assert plan.issue_axes == () + + +def test_auto_maximum_prefix_and_mixed_radix_issue_pointer(): + plan = _auto_plan(g_shape=(512, 64)) + assert _ints(plan.spec.box_dims) == (64, 1) + assert len(plan.issue_axes) == 1 + assert int(plan.issue_axes[0].extent) == 512 + assert int(plan.issue_axes[0].smem_stride) == 64 + offset, coords = plan.offsets_and_coords(IntImm("int32", 7)) + analyzer = Analyzer() + assert int(analyzer.simplify(offset)) == 448 + assert _ints(analyzer.simplify(coord) for coord in coords) == (0, 7) + + impl, _, _ = _lower_direct("tma_auto", g_shape=(512, 64)) + assert _count_tma(impl).total == 512 def test_copy_tma_host_init_dtype_is_string(): - _, host_init_stmts = _make_tma_call( + """The host-init encode call must carry the dtype as a StringImm, not a + packed enum -- ``_encode_signature`` reads ``args[2].value`` as a str.""" + _, host_init_stmts, _ = _lower_direct( + "tma_auto", g_shape=(8, 256), g_region=((0, 8), (0, 256)), s_shape=(8, 256), s_region=((0, 8), (0, 256)), - gmem_layout=TileLayout(S[8, 256]), - smem_layout=TileLayout(S[8, 256]), + dtype="float16", ) - encode_call = host_init_stmts[0].seq[1].value + encode_call = _collect_encodes(host_init_stmts)[0] assert isinstance(encode_call.args[2], StringImm) assert encode_call.args[2].value == "float16" -# Section 3: TMA special cases (symbolic dimension, buffer view) -# =========================================================================== - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") -@pytest.mark.parametrize("swizzle_len", [3]) -@pytest.mark.parametrize("dtype", ["float16"]) -def test_copy_tma_symbolic_dimension(dtype, swizzle_len): - """Test TMA copy with symbolic dimension in global buffer (like hgemm pattern). - - This tests the pattern: - Tx.copy_async(A_smem[ks, :, :], A[m_st : m_st + BLK_M, k_start : k_start + BLK_K], **tma_copy) # noqa: E501 - - Where M is a symbolic dimension in the global buffer. - """ # noqa: E501 - # Fixed dimensions - K = 256 - BLK_M = 64 - BLK_K = 64 - SMEM_PIPE_DEPTH = 2 - M_CONCRETE = 128 # Concrete value for testing - thread_cnt = 128 - - # Shared memory layout with swizzle - shared_layout = T.ComposeLayout( - T.SwizzleLayout(3, swizzle_len, 3, swizzle_inner=True), - T.TileLayout(T.S[(SMEM_PIPE_DEPTH, BLK_M, BLK_K) : (BLK_M * BLK_K, BLK_K, 1)]), +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ( + { + "g_shape": (8, 32, 64), + "s_layout": _plain_layout((8, 32, 64), (4096, 64, 1)), + }, + "stage=shared-chain", + ), + ( + { + "g_shape": (128, 64), + "g_region": ((4, 24), (0, 64)), + "s_shape": (3, 64, 8), + "s_layout": _plain_layout((3, 64, 8), (512, 1, 64)), + }, + "stage=coordinate", + ), + ( + { + "g_shape": (130, 64), + "g_region": ((0, 24), (0, 64)), + "s_shape": (3, 64, 8), + "s_layout": _plain_layout((3, 64, 8), (512, 1, 64)), + }, + "stage=global-shape", + ), + ], +) +def test_auto_requires_complete_chain_and_exact_division(kwargs, error): + with pytest.raises(Exception, match=error): + _auto_plan(**kwargs) + + +def test_auto_canonicalization_repairs_rank_and_unaligned_stride(): + rank_plan = _auto_plan(g_shape=(2, 2, 2, 2, 2, 8)) + assert rank_plan.spec.descriptor_dtype == "float16" + assert rank_plan.spec.rank == 1 + assert _ints(rank_plan.spec.global_dims) == (256,) + + stride_plan = _auto_plan(g_shape=(32, 8), dtype="uint8") + assert stride_plan.spec.descriptor_dtype == "uint8" + assert stride_plan.spec.rank == 1 + assert _ints(stride_plan.spec.global_dims) == (256,) + + +def test_auto_promotion_preserves_payload_and_shared_pointer(): + plan = _auto_plan(g_shape=(64, 8), dtype="uint8") + assert plan.spec.descriptor_dtype == "uint16" + assert plan.spec.rank == 1 + assert _ints(plan.spec.global_dims) == (256,) + assert _ints(plan.spec.box_dims) == (256,) + assert int(plan.spec.payload_bits) == 4096 + assert plan.spec.smem_start == (0, 0) + assert str(plan.spec.smem_buffer.dtype) == "uint8" + + +def test_auto_promotes_before_crossing_box_blocked_inner_chain_boundary(): + region = ((0, 1), (0, 4), (0, 32), (0, 4), (0, 4)) + plan = _auto_plan( + g_shape=(8, 16, 32, 4, 4), + g_region=region, + s_region=region, + dtype="uint8", + ) + assert plan.spec.descriptor_dtype == "uint16" + assert _ints(plan.spec.global_dims) == (256, 16, 8) + assert _ints(plan.spec.global_strides) == (512, 8192) + assert _ints(plan.spec.box_dims) == (256, 4, 1) + assert _ints(plan.spec.coordinates) == (0, 0, 0) + assert plan.issue_axes == () + assert int(plan.spec.payload_bits) == 16384 + + +def test_auto_promotion_rejects_unsafe_modes(): + spec = _make_spec( + descriptor_dtype="uint8", + descriptor_bits=8, + effective_bytes=1, + global_dims=(8, 64), + global_strides=(8,), + box_dims=(8, 64), + coordinates=(0, 0), + payload_bits=4096, + transaction_bits=4096, + ) + plan = TMAPlan(spec) + assert _promote_auto_once(plan) is not None + assert _promote_auto_once(replace(plan, spec=replace(spec, box_dims=(7, 64)))) is None + assert _promote_auto_once(replace(plan, spec=replace(spec, coordinates=(1, 0)))) is None + assert ( + _promote_auto_once( + replace( + plan, + issue_axes=(AutoIssueAxis(2, 8, (IssueCoord(0, 1, 2),)),), + ) + ) + is None + ) + assert _promote_auto_once(replace(plan, spec=replace(spec, use_tma_reduce="add"))) is None + assert _promote_auto_once(replace(plan, spec=replace(spec, force_cu_dtype=11))) is None + assert _promote_auto_once(replace(plan, spec=replace(spec, element_strides=(1, 2)))) is None + assert _promote_auto_once(replace(plan, spec=replace(spec, inner_stride=2))) is None + assert ( + _promote_auto_once( + replace( + plan, + spec=replace( + spec, + descriptor_dtype="float4_e2m1fn", + descriptor_bits=4, + packed_kind="16u4_align16", + ), + ) + ) + is None ) - # Compute bytes for mbarrier - smem_bytes = SMEM_PIPE_DEPTH * BLK_M * BLK_K * tvm.DataType(dtype).bits // 8 - copy_bytes = BLK_M * BLK_K * tvm.DataType(dtype).bits // 8 - # fmt: off - @T.prim_func - def copy_async(A_ptr: T.handle, B_ptr: T.handle) -> None: - M = T.int32() - A = T.match_buffer(A_ptr, [M, K], dtype) - B = T.match_buffer(B_ptr, [SMEM_PIPE_DEPTH, BLK_M, BLK_K], dtype) +def test_auto_odd_partial_and_unrepairable_box_fail_loudly(): + with pytest.raises(Exception, match="inner_box_bytes|global_stride_alignment"): + _auto_plan( + g_shape=(8, 64, 8), + g_region=((0, 8), (0, 64), (0, 5)), + s_shape=(8, 64, 5), + dtype="uint8", + ) + with pytest.raises(Exception, match="boxDim"): + _auto_plan(g_shape=(4096,), dtype="uint8") - T.device_entry() - cta_id = T.cta_id([1]) - tid = T.thread_id([thread_cnt]) - dyn = T.alloc_buffer([smem_bytes + 64], "uint8", scope="shared.dyn") - A_smem = T.decl_buffer( - [SMEM_PIPE_DEPTH, BLK_M, BLK_K], dtype, dyn.data, elem_offset=0, layout=shared_layout + +def test_auto_descriptor_cache_key_includes_promoted_dtype(): + data = Var("A", PointerType(PrimType("uint8"), "global")) + target = tvm.target.Target({"kind": "cuda", "arch": "sm_90a"}) + sctx = DispatchContext(target, ExecScope("thread"), {}, {}) + for rows in (32, 64): + op, _, _, _ = _make_op( + g_shape=(rows, 8), + dtype="uint8", + g_data=data, + sctx=sctx, + ) + copy_tma_auto_impl(op, sctx) + signatures = [ + _encode_signature(call) for call in _collect_encodes(sctx.callbacks["host_init_stmt"]) + ] + assert [item["dtype"] for item in signatures] == ["uint8", "uint16"] + assert [_ints(item["dims"]) for item in signatures] == [(256,), (256,)] + + +@pytest.mark.parametrize("rank", range(1, 6)) +def test_explicit_direct_rank_1_through_5(rank): + shape = (8,) * rank + impl, host, _ = _lower_direct("tma_explicit", g_shape=shape) + assert _count_tma(impl).total == 1 + signature = _encode_signature(_collect_encodes(host)[0]) + assert signature["rank"] == rank + assert _ints(signature["dims"]) == shape + assert _ints(signature["boxes"]) == shape + + +def test_explicit_retains_symbolic_outer_stride_for_runtime_encode(): + row_stride = Var("row_stride", "int32") + shape = (64, 64) + impl, host, _ = _lower_direct( + "tma_explicit", + g_shape=shape, + dtype="bfloat16", + g_layout=_plain_layout(shape, (row_stride, 1)), + ) + assert _count_tma(impl).total == 1 + signature = _encode_signature(_collect_encodes(host)[0]) + tvm.ir.assert_structural_equal(signature["strides"][0], row_stride * 2) + + +def test_explicit_rejects_unknown_or_nonunit_implicit_inner_stride(): + inner_stride = Var("inner_stride", "int32") + with pytest.raises(Exception, match="innermost memory stride must be provably one"): + _lower_direct( + "tma_explicit", + g_shape=(64, 64), + g_layout=_plain_layout((64, 64), (64, inner_stride)), + ) + with pytest.raises(Exception, match="innermost memory stride must be provably one"): + _lower_direct( + "tma_explicit", + g_shape=(64, 64), + g_layout=_plain_layout((64, 64), (64, 2)), ) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - mbar_ptr = T.meta_var(mbarrier.ptr_to([0])) - if tid == 0: - T.ptx.mbarrier.init(mbar_ptr, 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - # Copy with pipeline index (like hgemm pattern) - for ks in range(SMEM_PIPE_DEPTH): - if tid == 0: - Tx.copy_async( - A_smem[ks, :, :], - A[0:BLK_M, ks * BLK_K:(ks + 1) * BLK_K], - dispatch="tma", - mbar=mbar_ptr - ) - T.ptx.mbarrier.arrive.expect_tx(mbar_ptr, copy_bytes) +def test_explicit_view_base_offset_is_encoded_without_coordinate_rewrite(): + _, host, _ = _lower_direct( + "tma_explicit", + g_shape=(64, 64), + g_elem_offset=8, + ) + signature = _encode_signature(_collect_encodes(host)[0]) + assert isinstance(signature["base"], tvm.ir.Call) + assert signature["base"].op.name == "tirx.handle_add_byte_offset" + assert int(signature["base"].args[1]) == 16 + assert _ints(signature["dims"]) == (64, 64) - T.ptx.mbarrier.try_wait(mbar_ptr, ks % 2) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - for ks in range(SMEM_PIPE_DEPTH): - Tx.cta.copy( - B[ks, :, :], - A_smem[ks, :, :] - ) - # fmt: on +def test_explicit_oob_reduce_and_tf32_configs(): + _, load_host, _ = _lower_direct( + "tma_explicit", + g_shape=(64, 64), + config={"oob": "nan"}, + ) + assert _ints(_encode_signature(_collect_encodes(load_host)[0])["enums"])[-1] == 1 - np_dtype = tvm.testing.np_dtype_from_str(dtype) - target = tvm.target.Target("cuda") + store_impl, _, _ = _lower_direct( + "tma_explicit", + g_shape=(64, 64), + direction="s2g", + config={"use_tma_reduce": "add"}, + ) + counter = _count_tma(store_impl) + assert counter.total == 1 + assert counter.calls[0].op.name.endswith("shared_to_global_reduce") - with target: - mod = tvm.IRModule({"main": copy_async}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + _, tf32_host, _ = _lower_direct( + "tma_explicit", + g_shape=(64, 32), + dtype="float32", + config={"tma_dtype": "tf32"}, + ) + signature = _encode_signature(_collect_encodes(tf32_host)[0]) + assert int(signature["forced_dtype"]) == 11 - np.random.seed(0) - A_np = tvm.testing.generate_random_array(dtype, (M_CONCRETE, K)) - B_np = np.zeros((SMEM_PIPE_DEPTH, BLK_M, BLK_K), dtype=np_dtype) - # Verify: B[ks, :, :] should equal A[0:BLK_M, ks*BLK_K:(ks+1)*BLK_K] - B_ref = np.zeros((SMEM_PIPE_DEPTH, BLK_M, BLK_K), dtype=np_dtype) - for ks in range(SMEM_PIPE_DEPTH): - B_ref[ks, :, :] = A_np[0:BLK_M, ks * BLK_K : (ks + 1) * BLK_K] +def test_explicit_never_repairs_rank_or_descriptor_units(): + with pytest.raises(Exception, match="rank"): + _lower_direct("tma_explicit", g_shape=(2, 2, 2, 2, 2, 8)) + with pytest.raises(Exception, match="global_stride_alignment"): + _lower_direct("tma_explicit", g_shape=(32, 8), dtype="uint8") - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(B_ref, B.numpy()) - tvm.testing.run_with_gpu_lock(run_and_check) +def test_explicit_requires_box_linear_shared_slice(): + with pytest.raises(Exception, match="shared-layout"): + _lower_direct( + "tma_explicit", + g_shape=(4, 64), + s_layout=_plain_layout((4, 64), (128, 1)), + ) -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") -@pytest.mark.parametrize("swizzle_len", [3]) -@pytest.mark.parametrize("dtype", ["float16"]) -def test_copy_tma_3d_with_view(dtype, swizzle_len): - """Test 3D TMA copy using buffer view and swizzle layout (like flash attention pattern). - - This tests the pattern from FA4: - Q_smem allocated as 4D: (SMEM_PIPE_DEPTH, NUM_BLK_K, BLK_M, BLK_K) - Q_smem_3d = Q_smem.view(SMEM_PIPE_DEPTH, NUM_BLK_K, SEQ_TILE, GQA_RATIO, BLK_K) - Tx.copy_async(Q_smem_3d[pipe_idx, blk_k_idx, :, :, :], - Q[batch, seq_start:seq_end, head_start:head_end, k_start:k_end], ...) - """ - smem_bytes = 2 * 2 * 128 * 64 * tvm.DataType(dtype).bits // 8 - copy_bytes_per_blk = 32 * 4 * 64 * tvm.DataType(dtype).bits // 8 - - # Shared memory layout with swizzle - shared_layout = T.ComposeLayout( - T.SwizzleLayout(3, swizzle_len, 3, swizzle_inner=True), - T.TileLayout(T.S[(2, 128, 128) : (128 * 128, 128, 1)]), +def test_explicit_gather_uses_extracted_swizzled_slice_pointer(): + impl, _, _ = _lower_direct( + "tma_explicit", + g_shape=(8, 64), + s_shape=(8, 64), + g_region=((0, 1), (0, 64)), + s_region=((4, 4), (0, 64)), + s_layout=mma_shared_layout("float16", 3, (8, 64)), + config={"gather4": [0, 1, 2, 3]}, + target_arch="sm_100a", + ) + shared_ptr = _count_tma(impl).calls[0].args[1] + assert shared_ptr.op.name == "tirx.address_of" + assert int(shared_ptr.args[0].buffer.elem_offset) == 0 + assert int(shared_ptr.args[0].indices[0]) == 256 + + +def test_explicit_shared_pointer_counts_each_offset_once(): + warp_offset = Var("warp_offset", "int32") + layout = TileLayout(S[(4, 64)] + warp_offset * 64) + impl, _, _ = _lower_direct( + "tma_explicit", + g_shape=(2, 64), + s_shape=(4, 64), + s_region=((2, 2), (0, 64)), + s_layout=layout, + s_elem_offset=32, ) + shared_ptr = _count_tma(impl).calls[0].args[1] + assert shared_ptr.op.name == "tirx.address_of" + pointer_index = shared_ptr.args[0].indices[0] + assert Analyzer().can_prove_equal(pointer_index, 32 + warp_offset * 64 + 128) + + +def test_explicit_allows_different_operand_ranks_with_equal_payload_bytes(): + source = _from_source( + """ +@T.prim_func +def rank_change(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (8, 8), "float16") + T.device_entry() + T.cta_id([1]) + tid = T.thread_id([1]) + dyn = T.alloc_buffer((65,), "uint64", scope="shared.dyn") + A_smem = T.decl_buffer((64,), "float16", dyn.data, layout=T.TileLayout(T.S[64])) + mbar = T.decl_buffer((1,), "uint64", dyn.data, elem_offset=16) + if tid == 0: + Tx.copy_async( + A_smem[:], A[:, :], dispatch="tma_explicit", mbar=mbar.ptr_to([0]) + ) +""" + ) + lowered = _lower_module(source, "sm_90a") + assert _count_tma(lowered["main"]).total == 1 - # fmt: off - @T.prim_func - def copy_async(Q_ptr: T.handle, B_ptr: T.handle) -> None: - Q = T.match_buffer(Q_ptr, (2, 128, 8, 128), dtype) - B = T.match_buffer(B_ptr, (32, 4, 64), dtype) - T.device_entry() - cta_id = T.cta_id([1]) - tid = T.thread_id([128]) - dyn = T.alloc_buffer([smem_bytes + 64], "uint8", scope="shared.dyn") - # Allocate as 4D like FA4: (SMEM_PIPE_DEPTH, NUM_BLK_K, BLK_M, BLK_K) - Q_smem = T.decl_buffer( - (2, 2, 128, 64), - dtype, dyn.data, elem_offset=0, layout=shared_layout +_SELECTOR_SOURCE = """ +@T.prim_func +def selector_gather( + A_ptr: T.handle, + B_ptr: T.handle, + flag: T.int32, +): + A = T.match_buffer(A_ptr, (256, 64), "bfloat16") + B = T.match_buffer(B_ptr, (512, 80), "bfloat16") + B_view = B.sub[16:512, 8:72] + T.device_entry() + T.cta_id([1]) + tid = T.thread_id([128]) + dyn = T.alloc_buffer((520,), "uint64", scope="shared.dyn") + A_smem = T.decl_buffer( + (4, 64), "bfloat16", dyn.data, layout=T.TileLayout(T.S[4, 64]) + ) + mbar = T.decl_buffer((1,), "uint64", dyn.data, elem_offset=64) + if tid == 0: + T.ptx.mbarrier.init(mbar.ptr_to([0]), 1) + Tx.copy_async( + A_smem[:, :], + A[0:1, :], + dispatch="tma_explicit", + mbar=mbar.ptr_to([0]), + gather4=[1, 2, 3, 4], + src_selector=[(flag != 0, B_view)], + ) +""" + + +def test_explicit_gather_selector_encodes_each_map_selects_address_and_issues_once(): + func = _from_source(_SELECTOR_SOURCE) + lowered = _lower_module(func) + encodes = _collect_encodes([lowered["main"].body]) + assert len(encodes) == 2 + signatures = [_encode_signature(call) for call in encodes] + by_dims = {_ints(item["dims"]): item for item in signatures} + assert set(by_dims) == {(64, 256), (64, 496)} + assert _ints(by_dims[(64, 256)]["strides"]) == (128,) + assert _ints(by_dims[(64, 496)]["strides"]) == (160,) + assert [_ints(item["boxes"]) for item in signatures] == [(64, 1), (64, 1)] + + selects = _SelectCollector() + selects.visit_stmt(lowered["main"].body) + assert len(selects.nodes) == 1 + assert selects.nodes[0].ty == PrimType("uint64") + assert _count_tma(lowered["main"]).total == 1 + + executable = _compile_module(func) + cuda_source = executable.mod.imports[0].inspect_source() + assert "uint64_t selected_tensormap" in cuda_source + assert "flag != 0) ?" in cuda_source + assert "if (flag" not in cuda_source + assert cuda_source.count("cp_async_bulk_tensor_g2s_cluster_tile_gather4_2d(") == 2 + + +def test_explicit_selector_uses_first_true_and_requires_static_compatibility(): + main = _make_spec() + compatible = replace(main, global_dims=(16, 128), global_strides=(64,), base_key="B") + _selector_compatibility(main, compatible, 0) + with pytest.raises(Exception, match="descriptor dtype"): + _selector_compatibility( + main, + replace( + compatible, + descriptor_dtype="uint16", + ), + 0, + ) + symbolic_box = Var("candidate_box", "int32") + with pytest.raises(Exception, match="incompatible box"): + _selector_compatibility(main, replace(compatible, box_dims=(symbolic_box, 4)), 0) + + +def test_explicit_gather_validates_target_source_and_destination_layouts(): + common = dict( + g_shape=(256, 96), + g_region=((0, 1), (0, 64)), + s_shape=(4, 64), + g_layout=_plain_layout((256, 96), (96, 1)), + config={"gather4": [1, 2, 3, 4]}, + target_arch="sm_100a", + ) + impl, host, _ = _lower_direct("tma_explicit", **common) + assert _count_tma(impl).total == 1 + assert _ints(_encode_signature(_collect_encodes(host)[0])["strides"]) == (192,) + + with pytest.raises(Exception, match="SM100"): + _lower_direct("tma_explicit", **{**common, "target_arch": "sm_90a"}) + with pytest.raises(Exception, match="innermost memory stride"): + _lower_direct( + "tma_explicit", + **{ + **common, + "g_layout": _plain_layout((256, 96), (1, 256)), + }, + ) + with pytest.raises(Exception, match="shared-layout|box-linear"): + _lower_direct( + "tma_explicit", + **{ + **common, + "s_layout": _plain_layout((4, 64), (128, 1)), + }, ) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - mbar_ptr = T.meta_var(mbarrier.ptr_to([0])) - - # Create 5D view for 3D copy pattern - Q_smem_5d = Q_smem.view(2, 2, 32, 4, 64) - if tid == 0: - T.ptx.mbarrier.init(mbar_ptr, 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - if tid == 0: - # 3D copy: [SEQ_Q_PER_TILE, GQA_RATIO, BLK_K] - Tx.copy_async( - Q_smem_5d[0, 0, :, :, :], - Q[0, 0:32, 0:4, 0:64], - dispatch="tma", - mbar=mbar_ptr - ) - T.ptx.mbarrier.arrive.expect_tx(mbar_ptr, copy_bytes_per_blk) +def test_nested_selector_parser_printer_roundtrip_and_lowering_remap(): + func = _from_source(_SELECTOR_SOURCE) + script = func.script(extra_config={"tirx.prefix": "T"}) + reparsed = _from_source(script) + tvm.ir.assert_structural_equal(func, reparsed, map_free_vars=True) + lowered = _lower_module(reparsed) + assert len(_collect_encodes([lowered["main"].body])) == 2 + assert _count_tma(lowered["main"]).total == 1 - T.ptx.mbarrier.try_wait(mbar_ptr, 0) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - Tx.cta.copy( - B[:, :, :], - Q_smem_5d[0, 0, :, :, :] +def test_prefetch_only_main_selector_descriptor(): + func = _from_source( + _SELECTOR_SOURCE.replace( + "gather4=[1, 2, 3, 4],", + "gather4=[1, 2, 3, 4],\n prefetch_tensormap=True,", ) - # fmt: on + ) + lowered = _lower_module(func) + collector = _PrefetchCollector() + collector.visit_stmt(lowered["main"].body) + assert collector.names == ["A_ptr_tensormap"] - np_dtype = tvm.testing.np_dtype_from_str(dtype) - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_async}) +def _build_sparse_decode_qo_tma_regression(): + q_layout = mma_shared_layout("bfloat16", 3, (64, 512)) + q_tail_layout = mma_shared_layout("bfloat16", 2, (64, 64)) + o_layout = mma_shared_layout("bfloat16", 3, (64, 512)) + q_elements = 64 * 512 + q_tail_elements = 64 * 64 + o_elements = 64 * 512 + shared_bytes = (q_elements + q_tail_elements + o_elements) * 2 - # Verify that LowerTIRx generates exactly 1 TMA instruction - lowered = tvm.tirx.transform.LowerTIRx()(mod) - counter = TMACounter() - counter.visit_stmt(lowered["main"].body) - - assert counter.total_tma_ops == 1, ( - f"Expected exactly 1 TMA operation, got {counter.total_tma_ops}. " - "This indicates the 3D TMA copy with view is not generating optimal code." + # fmt: off + @T.prim_func + def kernel( + Q_ptr: T.handle, + O_ptr: T.handle, + q_stride_b: T.int64, + q_stride_s: T.int64, + q_stride_h: T.int64, + o_stride_b: T.int64, + o_stride_s: T.int64, + o_stride_h: T.int64, + ): + Q_storage = T.match_buffer(Q_ptr, (64 * 576,), "bfloat16") + O_storage = T.match_buffer(O_ptr, (64 * 512,), "bfloat16") + Q = Q_storage.view( + 1, + 1, + 64, + 576, + layout=T.TileLayout( + T.S[(1, 1, 64, 576) : (q_stride_b, q_stride_s, q_stride_h, 1)] + ), + ) + O_view = O_storage.view( + 1, + 1, + 64, + 512, + layout=T.TileLayout( + T.S[(1, 1, 64, 512) : (o_stride_b, o_stride_s, o_stride_h, 1)] + ), + ) + T.device_entry() + T.cta_id([1]) + tid = T.thread_id([128]) + dyn = T.alloc_buffer((shared_bytes + 8,), "uint8", scope="shared.dyn") + q_smem = T.decl_buffer( + (64, 512), "bfloat16", dyn.data, scope="shared.dyn", layout=q_layout + ) + q_tail_smem = T.decl_buffer( + (64, 64), + "bfloat16", + dyn.data, + elem_offset=q_elements, + scope="shared.dyn", + layout=q_tail_layout, ) + o_smem = T.decl_buffer( + (64, 512), + "bfloat16", + dyn.data, + elem_offset=q_elements + q_tail_elements, + scope="shared.dyn", + layout=o_layout, + ) + mbar = T.decl_buffer( + (1,), "uint64", dyn.data, elem_offset=shared_bytes // 8, scope="shared.dyn" + ) + q_tail_smem_tma = q_tail_smem.view(64, 2, 32).permute(1, 0, 2) + q_tail_gmem_tma = Q.sub[0, 0, :, 512:576].view(64, 2, 32).permute(1, 0, 2) + if tid == 0: + for q_tile in T.unroll(8): + Tx.copy_async( + q_smem[:, q_tile * 64 : (q_tile + 1) * 64], + Q[0, 0, :, q_tile * 64 : (q_tile + 1) * 64], + dispatch="tma_explicit", + mbar=mbar.ptr_to([0]), + cache_hint="evict_first", + tensormap_l2_promotion="L2::128B", + ) + Tx.copy_async( + q_tail_smem_tma[:, :, :], + q_tail_gmem_tma[:, :, :], + dispatch="tma_explicit", + mbar=mbar.ptr_to([0]), + cache_hint="evict_first", + tensormap_l2_promotion="L2::128B", + ) + for o_tile in T.unroll(8): + store_tile: T.let = o_tile + Tx.copy_async( + O_view[0, 0, :, store_tile * 64 : (store_tile + 1) * 64], + o_smem[:, store_tile * 64 : (store_tile + 1) * 64], + dispatch="tma_explicit", + cache_hint="evict_first", + tensormap_l2_promotion="L2::128B", + ) + # fmt: on - # Now compile and verify correctness - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + return kernel + + +def test_sparse_decode_runtime_stride_q_tail_o_counts_and_descriptor_dedup(): + kernel = _build_sparse_decode_qo_tma_regression() + lowered = _lower_module(kernel) + counter = _count_tma(lowered["main"]) + load_op = "tirx.ptx.cp_async_bulk_tensor_g2s_cluster" + store_op = "tirx.ptx.cp_async_bulk_tensor_shared_to_global" + assert sum(weight for call, weight in counter.weighted_calls if call.op.name == load_op) == 9 + assert sum(weight for call, weight in counter.weighted_calls if call.op.name == store_op) == 8 + + encodes = _collect_encodes([lowered["main"].body]) + signatures = [_encode_signature(call) for call in encodes] + assert len(signatures) == 3 + swizzles = [int(signature["enums"][1]) for signature in signatures] + assert swizzles.count(3) == 2 + assert swizzles.count(2) == 1 + assert all( + any(not isinstance(stride, IntImm) for stride in signature["strides"]) + for signature in signatures + ) + _compile_module(kernel) - np.random.seed(0) - Q_np = tvm.testing.generate_random_array(dtype, (2, 128, 8, 128)) - B_np = np.zeros((32, 4, 64), dtype=np_dtype) - B_ref = np.zeros((32, 4, 64), dtype=np_dtype) - B_ref[:, :, :] = Q_np[0, 0:32, 0:4, 0:64] +@pytest.mark.parametrize( + ("dtype", "cols", "swizzle"), + [ + ("int64", 64, 0), + ("bfloat16", 32, 2), + ("bfloat16", 64, 3), + ], +) +def test_sparse_decode_nope_rope_descriptor_dtypes_and_swizzles(dtype, cols, swizzle): + s_layout = ( + _plain_layout((8, cols)) if swizzle == 0 else mma_shared_layout(dtype, swizzle, (8, cols)) + ) + _, host, _ = _lower_direct( + "tma_explicit", + g_shape=(256, cols), + s_shape=(8, cols), + g_region=((0, 1), (0, cols)), + s_region=((0, 4), (0, cols)), + config={"gather4": [0, 1, 2, 3]}, + s_layout=s_layout, + dtype=dtype, + target_arch="sm_100a", + ) + signature = _encode_signature(_collect_encodes(host)[0]) + assert signature["dtype"] == dtype + assert int(signature["enums"][1]) == swizzle - def run_and_check(): - dev = tvm.cuda(0) - Q = tvm.runtime.tensor(Q_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(Q, B) - np.testing.assert_allclose(B_ref, B.numpy()) - tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.parametrize( + ("cols", "row_stride", "expected_bytes"), + [ + (64, 82, 656), + (56, 72, 576), + ], +) +def test_sparse_decode_kv_descriptor_uses_tma_row_stride(cols, row_stride, expected_bytes): + _, host, _ = _lower_direct( + "tma_explicit", + g_shape=(256, cols), + s_shape=(4, cols), + g_region=((0, 1), (0, cols)), + g_layout=_plain_layout((256, cols), (row_stride, 1)), + dtype="int64", + config={"gather4": [0, 1, 2, 3]}, + target_arch="sm_100a", + ) + signature = _encode_signature(_collect_encodes(host)[0]) + assert _ints(signature["strides"]) == (expected_bytes,) + assert 584 not in _ints(signature["strides"]) -# =========================================================================== -# Section 4: TMA GPU smoke tests (end-to-end compilation + correctness) -# =========================================================================== +def test_sparse_decode_absent_extra_map_has_no_dummy_descriptor_or_select(): + source = _SELECTOR_SOURCE.replace( + " src_selector=[(flag != 0, B_view)],\n", + " prefetch_tensormap=True,\n", + ) + lowered = _lower_module(_from_source(source)) + assert len(_collect_encodes([lowered["main"].body])) == 1 + selects = _SelectCollector() + selects.visit_stmt(lowered["main"].body) + assert selects.nodes == [] + prefetches = _PrefetchCollector() + prefetches.visit_stmt(lowered["main"].body) + assert prefetches.names == ["A_ptr_tensormap"] -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") @pytest.mark.parametrize( - "task", + ("mutate", "rule"), [ - # (a) Basic 2D G2S: (8,256) full region - pytest.param( - ( - (8, 256), # g_shape - ((0, 8), (0, 256)), # g_region - (8, 256), # s_shape - ((0, 8), (0, 256)), # s_region - 8, # thread count per CTA - TileLayout(S[8, 256]), # A_layout - TileLayout(S[8, 256]), # B_layout - lambda dtype: mma_shared_layout(dtype, 3, (8, 256)), - ), - id="g2s-2d-basic", - ), - # (b) 3D pipeline G2S: (3,8,256) → (8,256) per-phase - pytest.param( - ( - (3, 8, 256), - None, # multi-phase: region computed per-phase - (8, 256), - None, # multi-phase - 8, - TileLayout(S[3, 8, 256]), - TileLayout(S[3, 8, 256]), - lambda dtype: mma_shared_layout(dtype, 3, (8, 256)), - ), - id="g2s-3d-pipeline", + (lambda spec: replace(spec, global_dims=(0, 64)), "global_dim"), + (lambda spec: replace(spec, global_dims=((1 << 32) + 1, 64)), "global_dim"), + (lambda spec: replace(spec, global_strides=(-16,)), "global_stride_range"), + (lambda spec: replace(spec, global_strides=(8,)), "global_stride_alignment"), + (lambda spec: replace(spec, global_strides=(1 << 40,)), "global_stride_range"), + (lambda spec: replace(spec, base_byte_offset=2), "global_base_alignment"), + (lambda spec: replace(spec, box_dims=(0, 4)), "box_dim"), + (lambda spec: replace(spec, box_dims=(257, 4)), "box_dim"), + (lambda spec: replace(spec, box_dims=(7, 4)), "inner_box_bytes"), + (lambda spec: replace(spec, element_strides=(0, 1)), "element_stride"), + (lambda spec: replace(spec, element_strides=(9, 1)), "element_stride"), + (lambda spec: replace(spec, element_strides=(2, 1)), "inner_stride"), + (lambda spec: replace(spec, inner_stride=2), "inner_stride"), + (lambda spec: replace(spec, interleave=3), "interleave"), + (lambda spec: replace(spec, swizzle=9), "swizzle"), + (lambda spec: replace(spec, l2_promotion=4), "l2_promotion"), + (lambda spec: replace(spec, oob_fill=2), "oob"), + (lambda spec: replace(spec, target_arch="sm_80"), "target"), + (lambda spec: replace(spec, cta_group=3), "cta_group"), + (lambda spec: replace(spec, coordinates=(0,)), "array_lengths"), + (lambda spec: replace(spec, force_cu_dtype=7), "forced_cuda_dtype"), + (lambda spec: replace(spec, descriptor_bits=8), "dtype_bits"), + ( + lambda spec: replace(spec, swizzle=1, box_dims=(24, 4)), + "swizzle_inner_box", ), - # (c) 4D with unit dims: (2,2,128,64), copy (1,1,128,64) → 2D shared (128,64) - pytest.param( - ( - (2, 2, 128, 64), - ((0, 1), (0, 1), (0, 128), (0, 64)), - (128, 64), - ((0, 128), (0, 64)), - 128, - TileLayout(S[2, 2, 128, 64]).canonicalize(), - TileLayout(S[2, 2, 128, 64]).canonicalize(), - lambda dtype: mma_shared_layout(dtype, 3, (128, 64)), + ( + lambda spec: replace( + spec, + descriptor_dtype="int32", + descriptor_bits=32, + effective_bytes=4, + global_strides=(64,), + box_dims=(4, 4), + oob_fill=1, ), - id="g2s-4d-unit-dims", + "nan_oob_dtype", ), ], ) -@pytest.mark.parametrize("dtype", ["float16"]) -def test_copy_tma_gpu_smoke_g2s(task, dtype): - """Smoke test: compile and run TMA G2S copy on GPU to verify end-to-end correctness.""" - g_shape, g_region, s_shape, s_region, thread_cnt, layoutA, layoutB, layoutS_fn = task - shared_layout = layoutS_fn(dtype) - is_pipeline = g_region is None - - if is_pipeline: - n = g_shape[0] - smem_bytes = functools.reduce(lambda acc, e: acc * e, s_shape, 1) - smem_bytes = smem_bytes * tvm.DataType(dtype).bits // 8 - - r_smem = [slice(0, s) for s in s_shape] - - def r_gmem(stage): - return [ - slice(stage, stage + 1), - *[slice(0, g_shape[i]) for i in range(1, len(g_shape))], - ] - - # fmt: off - @T.prim_func - def copy_async(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA) - B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB) - - T.device_entry() - cta_id = T.cta_id([1]) - tid = T.thread_id([thread_cnt]) - dyn = T.alloc_buffer([smem_bytes + 8], "uint8", scope="shared.dyn") - A_smem = T.decl_buffer(s_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - phase: T.int32 - - phase = 0 - if tid == 0: - T.ptx.mbarrier.init(mbarrier.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - - for stage in range(n): - if tid == 0: - Tx.copy_async(A_smem[tuple(r_smem)], A[tuple(r_gmem(stage))], dispatch="tma", mbar=mbarrier.ptr_to([0])) # noqa: E501 - T.ptx.mbarrier.arrive.expect_tx(mbarrier.ptr_to([0]), smem_bytes) - - T.ptx.mbarrier.try_wait(mbarrier.ptr_to([0]), phase) - phase = phase ^ 1 - - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - Tx.cta.copy(B[tuple(r_gmem(stage))], A_smem[tuple(r_smem)]) - # fmt: on - - np_dtype = tvm.testing.np_dtype_from_str(dtype) - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_async}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - - np.random.seed(0) - A_np = tvm.testing.generate_random_array(dtype, g_shape) - B_np = np.zeros(g_shape, dtype=np_dtype) - - B_ref = A_np - else: - total_bytes = functools.reduce( - lambda acc, region: acc * (region[1] - region[0]), s_region, 1 - ) - total_bytes = total_bytes * tvm.DataType(dtype).bits // 8 - - smem_bytes = functools.reduce(lambda acc, e: acc * e, s_shape, 1) - smem_bytes = smem_bytes * tvm.DataType(dtype).bits // 8 - - r_smem = [slice(s_region[i][0], s_region[i][1]) for i in range(len(s_shape))] - r_gmem = [slice(g_region[i][0], g_region[i][1]) for i in range(len(g_shape))] - - # fmt: off - @T.prim_func - def copy_async(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA) - B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB) - - T.device_entry() - cta_id = T.cta_id([1]) - tid = T.thread_id([thread_cnt]) - dyn = T.alloc_buffer([smem_bytes + 64], "uint8", scope="shared.dyn") - A_smem = T.decl_buffer(s_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - mbar_ptr = T.meta_var(mbarrier.ptr_to([0])) - - if tid == 0: - T.ptx.mbarrier.init(mbar_ptr, 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - - if tid == 0: - Tx.copy_async(A_smem[tuple(r_smem)], A[tuple(r_gmem)], dispatch="tma", mbar=mbar_ptr) # noqa: E501 - T.ptx.mbarrier.arrive.expect_tx(mbar_ptr, total_bytes) - T.ptx.mbarrier.try_wait(mbar_ptr, 0) - T.cuda.cta_sync() - Tx.cta.copy(B[tuple(r_gmem)], A_smem[tuple(r_smem)]) - # fmt: on - - np_dtype = tvm.testing.np_dtype_from_str(dtype) - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_async}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - - np.random.seed(0) - A_np = tvm.testing.generate_random_array(dtype, g_shape) - B_np = np.zeros(g_shape, dtype=np_dtype) - - B_ref = np.zeros(g_shape, dtype=np_dtype) - B_ref[tuple(r_gmem)] = A_np[tuple(r_gmem)] - - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(B_ref, B.numpy()) - - tvm.testing.run_with_gpu_lock(run_and_check) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") -@pytest.mark.parametrize("dtype", ["float16"]) -def test_copy_tma_gpu_smoke_s2g(dtype): - """Smoke test: compile and run TMA S2G store on GPU.""" - g_shape = (3, 8, 256) - s_shape = (8, 256) - thread_cnt = 8 - n = g_shape[0] - - shared_layout = mma_shared_layout(dtype, 3, s_shape) - - smem_bytes = functools.reduce(lambda acc, e: acc * e, s_shape, 1) - smem_bytes = smem_bytes * tvm.DataType(dtype).bits // 8 - - r_smem = [slice(0, s) for s in s_shape] - - def r_gmem(stage): - return [slice(stage, stage + 1), *[slice(0, g_shape[i]) for i in range(1, len(g_shape))]] - - layoutA = TileLayout(S[3, 8, 256]) - layoutB = TileLayout(S[3, 8, 256]) - - # fmt: off - @T.prim_func - def copy_async(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, g_shape, dtype, layout=layoutA) - B = T.match_buffer(B_ptr, g_shape, dtype, layout=layoutB) - - T.device_entry() - cta_id = T.cta_id([1]) - tid = T.thread_id([thread_cnt]) - dyn = T.alloc_buffer([smem_bytes], "uint8", scope="shared.dyn") - A_smem = T.decl_buffer(s_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout) - - for stage in range(n): - Tx.copy(A_smem[tuple(r_smem)], A[tuple(r_gmem(stage))]) - T.cuda.cta_sync() - T.ptx.fence.proxy_async("shared::cta") - if tid == 0: - Tx.copy_async(B[tuple(r_gmem(stage))], A_smem[tuple(r_smem)], dispatch="tma") - T.ptx.cp_async.bulk.commit_group() - T.ptx.cp_async.bulk.wait_group() - T.cuda.cta_sync() - # fmt: on - - np_dtype = tvm.testing.np_dtype_from_str(dtype) - target = tvm.target.Target("cuda") +def test_shared_validator_static_boundary_matrix(mutate, rule): + findings = _finding(mutate(_make_spec()), rule) + assert any(item.status == ProofStatus.DISPROVEN for item in findings) + + +def test_shared_validator_unknown_policy_differs_for_auto_and_explicit(): + stride = Var("runtime_stride", "int64") + spec = _make_spec(global_strides=(stride,)) + assert _validation_failures(spec, auto=False) == [] + auto_failures = _validation_failures(spec, auto=True) + assert {item.rule for item in auto_failures} >= { + "global_stride_range", + "global_stride_alignment", + } + + +def test_auto_defers_only_unknown_global_dimension_bounds_to_runtime(): + dynamic_dim = Var("runtime_global_dim", "int64") + dynamic = _make_spec(global_dims=(dynamic_dim, 64)) + assert _finding(dynamic, "global_dim")[0].status == ProofStatus.UNKNOWN + assert not any( + finding.rule == "global_dim" for finding in _validation_failures(dynamic, auto=True) + ) - with target: - mod = tvm.IRModule({"main": copy_async}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") + invalid = replace(dynamic, global_dims=(0, 64)) + assert any(finding.rule == "global_dim" for finding in _validation_failures(invalid, auto=True)) - np.random.seed(0) - A_np = tvm.testing.generate_random_array(dtype, g_shape) - B_np = np.zeros(g_shape, dtype=np_dtype) - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(A_np, B.numpy()) +def test_shared_validator_packed_and_interleave_rules(): + packed = _make_spec( + descriptor_dtype="float4_e2m1fn", + descriptor_bits=4, + effective_bytes=1, + packed_kind="16u4_align16", + global_dims=(128, 64), + global_strides=(64,), + box_dims=(128, 4), + swizzle=3, + ) + assert _validation_failures(packed, auto=False) == [] + assert _finding(replace(packed, global_dims=(64, 64)), "packed_shape")[0].status == ( + ProofStatus.DISPROVEN + ) + assert _finding(replace(packed, box_dims=(64, 4)), "packed_box")[0].status == ( + ProofStatus.DISPROVEN + ) + assert _finding(replace(packed, swizzle=2), "packed_swizzle")[0].status == ( + ProofStatus.DISPROVEN + ) + assert _finding(replace(packed, direction="s2g", mbar=None), "packed_direction")[0].status == ( + ProofStatus.DISPROVEN + ) - tvm.testing.run_with_gpu_lock(run_and_check) + interleaved = _make_spec( + global_dims=(16, 8, 4), + global_strides=(32, 256), + box_dims=(8, 4, 2), + element_strides=(1, 1, 1), + coordinates=(0, 0, 0), + interleave=2, + swizzle=1, + ) + assert _validation_failures(interleaved, auto=False) == [] + assert _finding(replace(interleaved, swizzle=2), "interleave_swizzle")[0].status == ( + ProofStatus.DISPROVEN + ) -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") -@pytest.mark.parametrize("dtype", ["float16"]) -def test_copy_tma_dynamic_cta_mask(dtype): - """Regression test for B00004: dynamic cta_mask expression in TMA multicast. - - Verifies that a TIR expression (depending on T.cta_id) used as cta_mask in - copy_async compiles through the full TIRX pipeline without crashing. - Previously, lower_tirx_scope_ids replaced scope-ID vars via Substitute, - but Substitute didn't visit TilePrimitiveCall.config values, leaving stale var - references that caused MakePackedAPI to fail with: - "variables [...] are used, but are not passed in as API arguments" - """ - CLUSTER_SIZE = 4 - CTA_GROUP = 2 - BLK_M = 64 - BLK_K = 64 - thread_cnt = 128 - - smem_shape = (BLK_M, BLK_K) - shared_layout = T.ComposeLayout( - T.SwizzleLayout(3, 3, 3, swizzle_inner=True), T.TileLayout(T.S[smem_shape : (BLK_K, 1)]) - ) - smem_bytes = BLK_M * BLK_K * tvm.DataType(dtype).bits // 8 - copy_bytes = smem_bytes +def _build_selector_gather_gpu_kernel(dtype="float16"): + rows = 256 + cols = 64 + shared_bytes = 4 * cols * tvm.DataType(dtype).bits // 8 # fmt: off @T.prim_func - def copy_async_dynamic_mask(A_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, [BLK_M, BLK_K], dtype) - + def kernel( + A_ptr: T.handle, + B_ptr: T.handle, + flag: T.int32, + Out_ptr: T.handle, + ): + A = T.match_buffer(A_ptr, (rows, cols), dtype) + B = T.match_buffer(B_ptr, (rows, cols), dtype) + Out = T.match_buffer(Out_ptr, (4, cols), dtype) T.device_entry() - cbx = T.cta_id_in_cluster([CLUSTER_SIZE]) - cta_id = T.cta_id([CLUSTER_SIZE]) - tid = T.thread_id([thread_cnt]) - - # Dynamic cta_mask: exact expression from B00004 bug report - cta_mask = T.meta_var(5 + 5 * cbx) - dyn = T.alloc_buffer([smem_bytes + 64], "uint8", scope="shared.dyn") + T.cta_id([1]) + tid = T.thread_id([128]) + dyn = T.alloc_buffer((shared_bytes + 64,), "uint8", scope="shared.dyn") A_smem = T.decl_buffer( - smem_shape, dtype, dyn.data, elem_offset=0, layout=shared_layout, + (4, cols), dtype, dyn.data, layout=T.TileLayout(T.S[4, cols]) ) - mbarrier = T.decl_buffer([1], "uint64", dyn.data, elem_offset=smem_bytes // 8) - mbar_ptr = T.meta_var(mbarrier.ptr_to([0])) - + mbar = T.decl_buffer((1,), "uint64", dyn.data, elem_offset=shared_bytes // 8) + mbar_ptr = T.meta_var(mbar.ptr_to([0])) if tid == 0: T.ptx.mbarrier.init(mbar_ptr, 1) T.ptx.fence.proxy_async("shared::cta") T.cuda.cta_sync() - if tid == 0: Tx.copy_async( A_smem[:, :], - A[:, :], - dispatch="tma", + A[0:1, :], + dispatch="tma_explicit", mbar=mbar_ptr, - cta_mask=cta_mask, - cta_group=CTA_GROUP, + gather4=[7, 3, 19, 5], + src_selector=[(flag != 0, B)], ) - T.ptx.mbarrier.arrive.expect_tx(mbar_ptr, copy_bytes) - + T.ptx.mbarrier.arrive.expect_tx(mbar_ptr, shared_bytes) T.ptx.mbarrier.try_wait(mbar_ptr, 0) - # fmt: on - - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_async_dynamic_mask}) - # This compilation crashed before the B00004 fix with: - # "variables [...] are used, but are not passed in as API arguments" - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - - # Verify multicast instruction was generated - src = mod.mod.imports[0].inspect_source() - assert "multicast" in src, "Expected multicast TMA instruction in generated code" - - -@pytest.mark.gpu -def test_copy_tma_uint32_shape_extent(): - BK = 64 - A_layout = mma_shared_layout("float16", 3, (128, BK)) - - @T.prim_func - def tma_load(n: T.uint32, a_ptr: T.handle, o_ptr: T.handle) -> None: - A = T.match_buffer(a_ptr, (n, BK), "float16") - Out = T.match_buffer(o_ptr, (128, BK), "float16") - T.device_entry() - T.warp_id([4]) - T.cta_id([1]) - T.warpgroup_id([1]) - tid = T.thread_id_in_wg([128]) - sm = T.alloc_buffer((128, BK), "float16", scope="shared", layout=A_layout) - mb = T.alloc_shared([1], "uint64") - if tid == 0: - T.ptx.mbarrier.init(mb.ptr_to([0]), 1) T.ptx.fence.proxy_async("shared::cta") T.cuda.cta_sync() - if tid == 0: - Tx.copy_async(sm[:, :], A[0:128, 0:BK], dispatch="tma", mbar=mb.ptr_to([0])) - T.ptx.mbarrier.arrive.expect_tx(mb.ptr_to([0]), 128 * BK * 2) - T.ptx.mbarrier.try_wait(mb.ptr_to([0]), 0) - T.cuda.cta_sync() - reg = T.alloc_local(BK, "float16") - Tx.copy(reg[:], sm[tid, 0:BK]) - Tx.copy(Out[tid, 0:BK], reg[:]) + Tx.cta.copy(Out[:, :], A_smem[:, :]) + # fmt: on - target = tvm.target.Target("cuda") - with target: - tvm.compile(tvm.IRModule({"main": tma_load}), target=target, tir_pipeline="tirx") + return kernel @pytest.mark.gpu -def test_copy_tma_uint32_slice_base(): - BK = 64 - A_layout = mma_shared_layout("float16", 3, (128, BK)) - - @T.prim_func - def tma_off(off: T.uint32, a_ptr: T.handle, o_ptr: T.handle) -> None: - A = T.match_buffer(a_ptr, (4096, BK), "float16") - Out = T.match_buffer(o_ptr, (128, BK), "float16") - T.device_entry() - T.warp_id([4]) - T.cta_id([1]) - T.warpgroup_id([1]) - tid = T.thread_id_in_wg([128]) - sm = T.alloc_buffer((128, BK), "float16", scope="shared", layout=A_layout) - mb = T.alloc_shared([1], "uint64") - if tid == 0: - T.ptx.mbarrier.init(mb.ptr_to([0]), 1) - T.ptx.fence.proxy_async("shared::cta") - T.cuda.cta_sync() - if tid == 0: - Tx.copy_async(sm[:, :], A[off : off + 128, 0:BK], dispatch="tma", mbar=mb.ptr_to([0])) - T.ptx.mbarrier.arrive.expect_tx(mb.ptr_to([0]), 128 * BK * 2) - T.ptx.mbarrier.try_wait(mb.ptr_to([0]), 0) - T.cuda.cta_sync() - reg = T.alloc_local(BK, "float16") - Tx.copy(reg[:], sm[tid, 0:BK]) - Tx.copy(Out[tid, 0:BK], reg[:]) - - target = tvm.target.Target("cuda") - with target: - tvm.compile(tvm.IRModule({"main": tma_off}), target=target, tir_pipeline="tirx") +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_explicit_gather_selector_gpu_roundtrip(): + dtype = "float16" + kernel = _build_selector_gather_gpu_kernel(dtype) + executable = _compile_module(kernel) + dev = tvm.cuda(0) + rng = np.random.default_rng(0) + a_np = rng.standard_normal((256, 64)).astype(dtype) + b_np = rng.standard_normal((256, 64)).astype(dtype) + indices = np.array([7, 3, 19, 5]) + + def run(): + a = tvm.runtime.tensor(a_np, dev) + b = tvm.runtime.tensor(b_np, dev) + out = tvm.runtime.tensor(np.zeros((4, 64), dtype=dtype), dev) + executable(a, b, 0, out) + np.testing.assert_allclose(out.numpy(), a_np[indices]) + executable(a, b, 1, out) + np.testing.assert_allclose(out.numpy(), b_np[indices]) + + tvm.testing.run_with_gpu_lock(run) if __name__ == "__main__": diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem.py deleted file mode 100644 index ede3a646a2c6..000000000000 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tmem.py +++ /dev/null @@ -1,344 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# pylint: disable=invalid-name, missing-function-docstring -"""Tests for the TMEM copy_async dispatch (tcgen05-based tmem<->reg and smem<->tmem).""" - -import numpy as np -import pytest - -import tvm -import tvm.testing -from tvm.script import tirx as T -from tvm.script.tirx import tile as Tx -from tvm.testing import env -from tvm.tirx.layout import S, TCol, TileLayout, TLane -from tvm.tirx.layout import tid_in_wg as axis_tid_in_wg - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -@pytest.mark.parametrize("dtype", ["float16", "float32"]) -@pytest.mark.parametrize("width_32b", [4, 8, 16, 32]) -def test_copy_tmem2reg_async(dtype, width_32b): - """Test async tmem<->local copy using copy_async instead of copy. - - This tests the new copy_async dispatch for tmem<->local that doesn't - immediately wait after the operation, allowing for pipelining. - """ - - def next_power_of_2(x): - """Return the smallest power of 2 greater than or equal to x.""" - if x <= 1: - return 1 - return 1 << (x - 1).bit_length() - - bits = tvm.runtime.DataType(dtype).bits - if 128 % bits != 0 or 32 % bits != 0: - pytest.skip(f"dtype {dtype} is not supported") - - WIDTH = width_32b * (32 // bits) - VEC_LEN = 128 // bits - if WIDTH % VEC_LEN != 0: - pytest.skip(f"dtype {dtype} + width {width_32b} is not supported") - - g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) - local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)]) - - # fmt: off - @T.prim_func - def copy_async_test(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, (128, WIDTH), dtype) - B = T.match_buffer(B_ptr, (128, WIDTH), dtype) - - A_flat = A.view(-1) - B_flat = B.view(-1) - - T.device_entry() - warp_id = T.warp_id([(128) // 32]) - cta_id = T.cta_id([2]) - wg_id = T.warpgroup_id([1]) - warp_id_in_wg = T.warp_id_in_wg([4]) - lane_id = T.lane_id([32]) - tid_in_wg = T.thread_id([128]) - - tmem_addr = T.alloc_shared([1], "uint32") - - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 - - T.tvm_storage_sync("shared") - - tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], - layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)])) - - A_reg = T.alloc_local((WIDTH), dtype) - B_reg = T.alloc_local((WIDTH), dtype) - A_local = A_reg.view(128, WIDTH, layout=local_view) - B_local = B_reg.view(128, WIDTH, layout=local_view) - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 - for i in range(WIDTH): - B_reg[i] = T.cast(0, dtype) - T.cuda.cta_sync() - - # A_local -> tmem (async) - Tx.wg.copy_async(tmem[:, :], A_local[:, :]) - T.ptx.tcgen05.wait.st() # explicit wait - T.cuda.cta_sync() - - # tmem -> B_local (async) - Tx.wg.copy_async(B_local[:, :], tmem[:, :]) - T.ptx.tcgen05.wait.ld() # explicit wait - T.cuda.cta_sync() - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501 - - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 - # fmt: on - - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_async_test}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) - B_np = np.zeros((128, WIDTH), dtype=dtype) - - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(B.numpy(), A_np) - - tvm.testing.run_with_gpu_lock(run_and_check) - - -# ---------------------------------------------------------------------------- -# Migrated from test_copy_sync.py: tmem<->reg round-trip via T.copy_async -# (the kernels themselves are the actual async tmem dispatch tests; the -# G↔L copies bookending them just stage data). -# ---------------------------------------------------------------------------- - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -@pytest.mark.parametrize("dtype", ["uint8", "float16", "float32"]) -@pytest.mark.parametrize("width_32b", [2, 4, 8, 16, 32, 64, 128]) -@pytest.mark.parametrize("offset_32b", [0, 3, 10]) -def test_copy_tmem2reg(dtype, width_32b, offset_32b): - def next_power_of_2(x): - if x <= 1: - return 1 - return 1 << (x - 1).bit_length() - - bits = tvm.runtime.DataType(dtype).bits - if 128 % bits != 0 or 32 % bits != 0: - pytest.skip(f"dtype {dtype} is not supported") - - WIDTH = width_32b * (32 // bits) - OFFSET = offset_32b * (32 // bits) - VEC_LEN = 128 // bits - if WIDTH % VEC_LEN != 0: - pytest.skip(f"dtype {dtype} + width {width_32b} is not supported") - - g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) - local_view = TileLayout(S[(128, WIDTH) : (1 @ axis_tid_in_wg, 1)]) - - # fmt: off - @T.prim_func - def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, (128, WIDTH), dtype) - B = T.match_buffer(B_ptr, (128, WIDTH), dtype) - - A_flat = A.view(-1) - B_flat = B.view(-1) - - T.device_entry() - warp_id = T.warp_id([(128) // 32]) - T.cta_id([2]) - wg_id = T.warpgroup_id([1]) - T.warp_id_in_wg([4]) - T.lane_id([32]) - tid_in_wg = T.thread_id([128]) - - tmem_addr = T.alloc_shared([1], "uint32") - - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501 - - T.tvm_storage_sync("shared") - - tmem = T.decl_buffer((128, OFFSET + WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], # noqa: E501 - layout=TileLayout(S[(128, OFFSET + WIDTH) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 - - A_reg = T.alloc_local((WIDTH), dtype) - B_reg = T.alloc_local((WIDTH), dtype) - A_local = A_reg.view(128, WIDTH, layout=local_view) - B_local = B_reg.view(128, WIDTH, layout=local_view) - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(A_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 - for i in range(WIDTH): - B_reg[i] = T.cast(0, dtype) - T.cuda.cta_sync() - - # A_local -> tmem - Tx.wg.copy_async(tmem[:, OFFSET: OFFSET + WIDTH], A_local[:, :]) - T.ptx.tcgen05.wait.st() - T.cuda.cta_sync() - - # tmem -> B_local - Tx.wg.copy_async(B_local[:, :], tmem[:, OFFSET: OFFSET + WIDTH]) - T.ptx.tcgen05.wait.ld() - T.cuda.cta_sync() - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[i * VEC_LEN: i * VEC_LEN + VEC_LEN]) # noqa: E501 - - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(offset_32b + width_32b)), cta_group=1) # noqa: E501 - # fmt: on - - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_sync}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) - B_np = np.zeros((128, WIDTH), dtype=dtype) - - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(B.numpy(), A_np) - - tvm.testing.run_with_gpu_lock(run_and_check) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") -@pytest.mark.parametrize("dtype", ["float16", "float32"]) -@pytest.mark.parametrize("width_32b", [4, 8, 16, 32]) -@pytest.mark.parametrize("local_offset_32b", [0, 2, 4]) -def test_copy_tmem2reg_sliced_local(dtype, width_32b, local_offset_32b): - """tmem<->local copy with a sliced local buffer region.""" - - def next_power_of_2(x): - if x <= 1: - return 1 - return 1 << (x - 1).bit_length() - - bits = tvm.runtime.DataType(dtype).bits - if 128 % bits != 0 or 32 % bits != 0: - pytest.skip(f"dtype {dtype} is not supported") - - WIDTH = width_32b * (32 // bits) - LOCAL_OFFSET = local_offset_32b * (32 // bits) - TOTAL_LOCAL_WIDTH = WIDTH + LOCAL_OFFSET - VEC_LEN = 128 // bits - if WIDTH % VEC_LEN != 0 or TOTAL_LOCAL_WIDTH % VEC_LEN != 0: - pytest.skip( - f"dtype {dtype} + width {width_32b} + offset {local_offset_32b} is not supported" - ) - - g_layout = TileLayout(S[(128, WIDTH // VEC_LEN, VEC_LEN) : (WIDTH, VEC_LEN, 1)]) - local_view = TileLayout(S[(128, TOTAL_LOCAL_WIDTH) : (1 @ axis_tid_in_wg, 1)]) - - # fmt: off - @T.prim_func - def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None: - A = T.match_buffer(A_ptr, (128, WIDTH), dtype) - B = T.match_buffer(B_ptr, (128, WIDTH), dtype) - - A_flat = A.view(-1) - B_flat = B.view(-1) - - T.device_entry() - warp_id = T.warp_id([(128) // 32]) - T.cta_id([2]) - wg_id = T.warpgroup_id([1]) - T.warp_id_in_wg([4]) - T.lane_id([32]) - tid_in_wg = T.thread_id([128]) - - tmem_addr = T.alloc_shared([1], "uint32") - - if wg_id == 0: - if warp_id == 0: - T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 - - T.tvm_storage_sync("shared") - - tmem = T.decl_buffer((128, WIDTH), dtype, scope="tmem", allocated_addr=tmem_addr[0], - layout=TileLayout(S[(128, WIDTH) : (1 @ TLane, 1 @ TCol)])) - - A_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype) - B_reg = T.alloc_local((TOTAL_LOCAL_WIDTH), dtype) - A_local = A_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view) - B_local = B_reg.view(128, TOTAL_LOCAL_WIDTH, layout=local_view) - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(A_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN], A_flat[g_offset: g_offset + VEC_LEN]) # noqa: E501 - for i in range(TOTAL_LOCAL_WIDTH): - B_reg[i] = T.cast(0, dtype) - T.cuda.cta_sync() - - # A_local[sliced] -> tmem (use sliced region) - Tx.wg.copy_async(tmem[:, 0:WIDTH], A_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH]) - T.ptx.tcgen05.wait.st() - T.cuda.cta_sync() - - # tmem -> B_local[sliced] (use sliced region) - Tx.wg.copy_async(B_local[:, LOCAL_OFFSET:LOCAL_OFFSET + WIDTH], tmem[:, 0:WIDTH]) - T.ptx.tcgen05.wait.ld() - T.cuda.cta_sync() - for i in range(WIDTH // VEC_LEN): - g_offset = T.meta_var(g_layout.apply(tid_in_wg, i, 0)["m"]) - Tx.copy(B_flat[g_offset: g_offset + VEC_LEN], B_reg[LOCAL_OFFSET + i * VEC_LEN: LOCAL_OFFSET + i * VEC_LEN + VEC_LEN]) # noqa: E501 - - if warp_id == 0: - T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) - T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=max(32, next_power_of_2(width_32b)), cta_group=1) # noqa: E501 - # fmt: on - - target = tvm.target.Target("cuda") - with target: - mod = tvm.IRModule({"main": copy_sync}) - mod = tvm.compile(mod, target=target, tir_pipeline="tirx") - A_np = tvm.testing.generate_random_array(dtype, (128, WIDTH)) - B_np = np.zeros((128, WIDTH), dtype=dtype) - - def run_and_check(): - dev = tvm.cuda(0) - A = tvm.runtime.tensor(A_np, dev) - B = tvm.runtime.tensor(B_np, dev) - mod(A, B) - np.testing.assert_allclose(B.numpy(), A_np) - - tvm.testing.run_with_gpu_lock(run_and_check) - - -if __name__ == "__main__": - tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py index efef5bbad967..62fb0c6db88c 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py @@ -451,7 +451,7 @@ def run_and_check(): ) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -@pytest.mark.parametrize("op_type", ["reciprocal", "exp", "exp2"]) +@pytest.mark.parametrize("op_type", ["reciprocal", "exp", "exp2", "log2"]) @pytest.mark.parametrize( "src_dtype,dst_dtype", [("float16", "float16"), ("float32", "float16"), ("float32", "bfloat16")] ) @@ -510,6 +510,8 @@ def test_unary(A_ptr: T.handle, B_ptr: T.handle) -> None: Tx.warp.exp(res_view, acc_view) elif op_type == "exp2": Tx.warp.exp2(res_view, acc_view) + elif op_type == "log2": + Tx.warp.log2(res_view, acc_view) # write res into B for i in T.serial(NUM_COL // 8): @@ -539,6 +541,8 @@ def test_unary(A_ptr: T.handle, B_ptr: T.handle) -> None: B_ref = np.exp(A_np).astype(dst_dtype) elif op_type == "exp2": B_ref = np.exp2(A_np).astype(dst_dtype) + elif op_type == "log2": + B_ref = np.log2(A_np).astype(dst_dtype) else: raise ValueError(f"op_type={op_type} is not supported") diff --git a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py index 716c0c02e3fa..02f7fc212d01 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py @@ -18,6 +18,7 @@ import copy import functools import operator +import re import numpy as np import pytest @@ -35,11 +36,13 @@ from tvm.testing import env from tvm.tirx.cuda.operator.tile_primitive.gemm_async import sf_tmem_layout from tvm.tirx.cuda.operator.tile_primitive.tma_utils import ( + SwizzleMode, mma_atom_layout, mma_atom_shape, mma_shared_layout, ) from tvm.tirx.layout import ( + R, S, TCol, TileLayout, @@ -74,7 +77,7 @@ def _mn_major_layout(dtype, swizzle_mode, shape): For shape (..., M, K), the standard K-major atom is [8, T*s] with K contiguous. MN-major swaps this: atom becomes [T*s, 8] with M contiguous. - This is achieved by composing the SwizzleLayout with a stride-reversed TileLayout. + This is achieved by composing the swizzle with a stride-reversed TileLayout. """ from tvm.tirx.layout import ComposeLayout @@ -83,7 +86,13 @@ def _mn_major_layout(dtype, swizzle_mode, shape): swapped = [base_shape[1], base_shape[0]] # [T*s, 8] # Stride-reversed tile: first dim (T*s) contiguous, second dim (8) has stride T*s mn_tile = TileLayout(S[tuple(swapped) : (1, swapped[0])]) - mn_atom = ComposeLayout(swizzle_atom, mn_tile) + mn_atom = ComposeLayout( + swizzle_atom.per_element, + swizzle_atom.swizzle_len, + swizzle_atom.atom_len, + mn_tile, + swizzle_atom.swizzle_inner, + ) # Tile up: first expand penultimate dim, then full shape tile_step = [1] * (len(shape) - 2) + [shape[-2], swapped[1]] atom_nd = [1] * (len(shape) - 2) + swapped @@ -251,7 +260,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: tmem = T.decl_buffer((128, C_shape[1]), C_dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(128, C_shape[1]) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[tuple(r_gmem_A)], A[tuple(r_gmem_A)], **tma_args) Tx.copy_async(B_smem[tuple(r_gmem_B)], B[tuple(r_gmem_B)], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -259,7 +268,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.cuda.cta_sync() if tid_in_wg == 0: - Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05") # noqa: E501 + Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", mma_m=128, mma_n=64) # noqa: E501 T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=1) T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) T.cuda.cta_sync() @@ -308,7 +317,7 @@ def run_and_check(): @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_gemm_tcgen05_cta_group_1_layout_f_m64(): - """M=64 MMA with C operand allocated as Layout F (datapath="F"). + """M=64 MMA with C operand allocated as Layout F. Exercises the new ``gemm_async`` path that accepts C buffers tagged Layout F — written by an M=64 MMA in their canonical scattered @@ -362,7 +371,7 @@ def gemm_layout_f(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: tmem = T.decl_buffer((64, N), C_dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=c_layout) # noqa: E501 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[:, :], A[:, :], **tma_args) Tx.copy_async(B_smem[:, :], B[:, :], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), (M * K + N * K) * 2) @@ -505,7 +514,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.cuda.cta_sync() T.cuda.cluster_sync() - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 if tid_in_wg == 0: Tx.copy_async(A_smem[tuple(r_smem_A_in)], A[tuple(get_global_region(A_shape_per_cta, transA, cbx))], **tma_args) # noqa: E501 Tx.copy_async(B_smem[tuple(r_smem_B_in)], B[tuple(get_global_region(B_shape_per_cta, transB, cbx))], **tma_args) # noqa: E501 @@ -517,7 +526,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.ptx.tcgen05.fence.after_thread_sync() T.cuda.cta_sync() if tid_in_wg == 0: - Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", cta_group=2) # noqa: E501 + Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", cta_group=2, mma_m=256, mma_n=128) # noqa: E501 T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=2, cta_mask=3) # signal cta 1's mbarrier # noqa: E501 T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) # both cta 0 and cta 1 have done mma T.ptx.tcgen05.fence.after_thread_sync() @@ -639,7 +648,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.cuda.cta_sync() T.cuda.cluster_sync() - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 if tid_in_wg == 0: # CTA cbx loads its portion of A and B Tx.copy_async(A_smem[0:M_per_cta, 0:K], A[cbx * M_per_cta:(cbx + 1) * M_per_cta, 0:K], **tma_args) # noqa: E501 @@ -652,7 +661,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.ptx.tcgen05.fence.after_thread_sync() T.cuda.cta_sync() if tid_in_wg == 0: - Tx.gemm_async(tmem[0:M_per_cta, 0:N_logical], A_smem[0:M_per_cta, 0:K], B_smem[0:N_half, 0:K], dispatch="tcgen05", cta_group=2) # noqa: E501 + Tx.gemm_async(tmem[0:M_per_cta, 0:N_logical], A_smem[0:M_per_cta, 0:K], B_smem[0:N_half, 0:K], dispatch="tcgen05", cta_group=2, mma_m=128, mma_n=128) # noqa: E501 T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=2, cta_mask=3) T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) T.ptx.tcgen05.fence.after_thread_sync() @@ -943,7 +952,7 @@ def gemm_async_fn(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, SFA_ptr: T. # TMA load A and B from global to shared if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[tuple(r_gmem_A)], A[tuple(r_gmem_A)], **tma_args) Tx.copy_async(B_smem[tuple(r_gmem_B)], B[tuple(r_gmem_B)], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -1147,7 +1156,7 @@ def gemm_async_fn(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, SFA_ptr: T. T.cuda.cluster_sync() # TMA load A and B (both CTAs issue with multicast) - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 if tid_in_wg == 0: Tx.copy_async(A_smem[tuple(r_smem_A_in)], A[tuple(get_global_region(A_shape_per_cta, transA, cbx))], **tma_args) # noqa: E501 Tx.copy_async(B_smem[tuple(r_smem_B_in)], B[tuple(get_global_region(B_shape_per_cta, transB, cbx))], **tma_args) # noqa: E501 @@ -1341,7 +1350,7 @@ def gemm_async_fn(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, SFA_ptr: T. # TMA load A and B as uint8 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem_packed[:, :], A_packed[:, :], **tma_args) Tx.copy_async(B_smem_packed[:, :], B_packed[:, :], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -1529,7 +1538,7 @@ def gemm_async_fn(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, SFA_ptr: T. T.cuda.cluster_sync() # TMA load A and B with multicast (each CTA loads its portion) - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar_cta_0.ptr_to([0]), "cta_group": 2}) # noqa: E501 if tid_in_wg == 0: Tx.copy_async(A_smem_packed[:, :], A_packed[cbx * M_per_cta:(cbx + 1) * M_per_cta, :], **tma_args) # noqa: E501 Tx.copy_async(B_smem_packed[:, :], B_packed[cbx * N_per_cta:(cbx + 1) * N_per_cta, :], **tma_args) # noqa: E501 @@ -1728,7 +1737,7 @@ def gemm_async_fn(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle, SFA_ptr: T. # TMA load A and B from global to shared if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[0:M, 0:K], A[0:M, 0:K], **tma_args) Tx.copy_async(B_smem[0:N, 0:K], B[0:N, 0:K], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -2069,7 +2078,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: tmem = T.decl_buffer((M, C_shape[1]), C_dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(M, C_shape[1]) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[tuple(r_gmem_A)], A[tuple(r_gmem_A)], **tma_args) Tx.copy_async(B_smem[tuple(r_gmem_B)], B[tuple(r_gmem_B)], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -2138,6 +2147,291 @@ def run_and_check(): tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.parametrize("a_layout_kind", ["column_major", "packed_16b"]) +def test_gemm_tcgen05_no_swizzle_smem_descriptor_codegen(a_layout_kind): + M, K, N = 64, 64, 256 + B_N = N // 2 + dtype = "bfloat16" + if a_layout_kind == "column_major": + A_layout = _col_major_layout((M, K)) + else: + A_layout = TileLayout(S[(M, K // 8, 8) : (8, M * 8, 1)]) + B_layout = _mn_major_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, B_N)) + C_layout = TileLayout(S[(M, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + + @T.prim_func + def gemm_async_no_swizzle(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (M, K), dtype, layout=A_layout) + B = T.match_buffer(B_ptr, (K, B_N), dtype, layout=B_layout) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid_in_wg = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer((M, K), dtype, scope="shared", layout=A_layout) + B_smem = T.alloc_buffer((K, B_N), dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=256, cta_group=2) + T.cuda.cta_sync() + tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=tmem_addr[0], + layout=C_layout, + ) + if tid_in_wg == 0: + Tx.gemm_async( + tmem[:, :], + A_smem[:, :], + B_smem[:, :], + transB=True, + dispatch="tcgen05", + cta_group=2, + ) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile( + tvm.IRModule({"main": gemm_async_no_swizzle}), target=target, tir_pipeline="tirx" + ) + + src = mod.mod.imports[0].inspect_source() + assert "descA" in src + assert ", 64, 8, 0)" in src + assert "encode_instr_descriptor" not in src + + +def test_gemm_tcgen05_cta_group_2_accepts_replicated_tmem_a_codegen(): + """A-in-TMEM for cta_group::2 may declare the physical +64 lane mirror. + + FlashMLA head128 copies Qt with 64x128b.warpx2::02_13, which writes rows + 0..63 and mirrors them at lane +64. The QK GEMM still addresses the anchor + tile, but the A buffer layout should be allowed to make that mirror explicit. + """ + + M = 64 + N = 128 + N_half = N // 2 + K = 128 + dtype = "bfloat16" + C_layout = TileLayout(S[(M, 2, N_half) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + A_layout = TileLayout(S[(M, K) : (1 @ TLane, 1 @ TCol)] + R[2 : 64 @ TLane]) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (N_half, K)) + + @T.prim_func + def gemm_async_replicated_a() -> None: + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid_in_wg = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N_half, K), dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=128, cta_group=2) + T.cuda.cta_sync() + C_tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=tmem_addr[0], + layout=C_layout, + ) + A_tmem = T.decl_buffer( + (M, K), + dtype, + scope="tmem", + allocated_addr=tmem_addr[0] + T.uint32(64), + layout=A_layout, + ) + if tid_in_wg == 0: + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=2, + ) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + mod = tvm.compile( + tvm.IRModule({"main": gemm_async_replicated_a}), + target=target, + tir_pipeline="tirx", + ) + + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.mma.cta_group::2" in src + assert "tcgen05.mma.ws" not in src + assert "_TS" in src + + +def test_gemm_tcgen05_cta_group_2_rejects_flat_tmem_a_codegen(): + """cta_group::2 Layout-B A-in-TMEM must declare its +64-lane footprint.""" + + M = 64 + N = 128 + N_half = N // 2 + K = 128 + dtype = "bfloat16" + C_layout = TileLayout(S[(M, 2, N_half) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + A_layout = TileLayout(S[(M, K) : (1 @ TLane, 1 @ TCol)]) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (N_half, K)) + + @T.prim_func + def gemm_async_flat_a() -> None: + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid_in_wg = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N_half, K), dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=128, cta_group=2) + T.cuda.cta_sync() + C_tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=tmem_addr[0], + layout=C_layout, + ) + A_tmem = T.decl_buffer( + (M, K), + dtype, + scope="tmem", + allocated_addr=tmem_addr[0] + T.uint32(64), + layout=A_layout, + ) + if tid_in_wg == 0: + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=2, + ) + + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + with target: + with pytest.raises(Exception, match="TMEM A layout"): + tvm.compile( + tvm.IRModule({"main": gemm_async_flat_a}), + target=target, + tir_pipeline="tirx", + ) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +@pytest.mark.skipif(ml_dtypes is None, reason="Requires ml_dtypes") +def test_gemm_tcgen05_no_swizzle_col_major_a_ws_local_idesc(): + """A column-major-viewed unswizzled SMEM A under kind::f16 + ws. + + This is the FlashMLA head64 O-GEMM: A = S tile [M=64, K=64] bf16 whose + GEMM operand view is column-major (M, K):(1, M) with no swizzle. The view + is a stride fiction: the S tile is physically stored 16B-line packed along + K (elem offset 8*m + 8*M*(k//8) + k%8), and the view's strides reproduce + that K-major layout's byte offsets. The dispatcher's no-swizzle + descriptor (ldo=64, sdo=8) is therefore a K-major encoding: hardware + consumes it with instruction-descriptor bit15 (a_major) CLEAR. + Historically callers masked this by hand-passing ``descI`` with + trans_a=0; with the locally encoded descI the dispatcher's returned + majorness must itself be consistent with the descriptor it constructed. + + Asserts (a) the generated idesc literal equals the hand-validated + 0x04410490 (bit15=0) and not the MN-major mis-encoding 0x04418490, with + the unchanged descA fields (64, 8, 0); and (b) the GEMM is numerically + correct on GPU. + """ + M, K, N = 64, 64, 256 + dtype = "bfloat16" + A_layout = TileLayout(S[(M, K) : (1, M)]) # column-major, unswizzled + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + + # fmt: off + @T.prim_func + def gemm_ws(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (M, K), dtype) + B = T.match_buffer(B_ptr, (K, N), dtype) + C = T.match_buffer(C_ptr, (128, N // 2), "float32") + T.device_entry() + warp_id = T.warp_id([4]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer((M, K), dtype, scope="shared", layout=A_layout) + B_smem = T.alloc_buffer((K, N), dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + mma_mbar = T.alloc_shared([1], "uint64") + if tid_in_wg == 0: + T.ptx.mbarrier.init(mma_mbar.ptr_to([0]), 1) + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=128, cta_group=1) + T.cuda.cta_sync() + # M=64 .ws accumulates via datapath E (FlashMLA head64's tmem_o layout): + # lane = m + 64*(n >= N/2), col = n % (N/2). + tmem = T.decl_buffer((M, N), "float32", scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(M, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)])) # noqa: E501 + # Identity overlay of the physical 128x128 TMEM footprint for readback. + tmem_ldst = T.decl_buffer((128, N // 2), "float32", scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(128, N // 2) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 + # Plain generic-proxy stores; A's bytes follow the FlashMLA S-tile ABI + # (phys = 8*m + 8*M*(k//8) + k%8), col-major A_smem is descriptor fiction. + for i in range(M * K // 128): + a_idx = i * 128 + tid_in_wg + a_m = a_idx % M + a_k = a_idx // M + a_phys = 8 * a_m + 8 * M * (a_k // 8) + (a_k % 8) + A_smem[a_phys % M, a_phys // M] = A[a_m, a_k] + for i in range(K * N // 128): + b_idx = i * 128 + tid_in_wg + B_smem[b_idx // N, b_idx % N] = B[b_idx // N, b_idx % N] + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + if tid_in_wg == 0: + Tx.gemm_async(tmem[:, :], A_smem[:, :], B_smem[:, :], transB=True, dispatch="tcgen05", cta_group=1, weight_stationary=True) # noqa: E501 + T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=1) + T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + C_reg = T.alloc_local(N // 2, dtype="float32") + C_view = C_reg.view(128, N // 2, layout=TileLayout(S[(128, N // 2) : (1@axis_tid_in_wg, 1)])) # noqa: E501 + if wg_id == 0: + Tx.wg.copy_async(C_view[:, :], tmem_ldst[:, :]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + Tx.copy(C[tid_in_wg, 0 : N // 2], C_reg[:]) + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=128, cta_group=1) + # fmt: on + + with tvm.target.Target("cuda"): + mod = tvm.compile(tvm.IRModule({"main": gemm_ws}), target="cuda", tir_pipeline="tirx") + + src = mod.mod.imports[0].inspect_source() + # Descriptor construction: no-swizzle col-major A -> (ldo=64, sdo=8, swizzle=0). + assert ", 64, 8, 0)" in src + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in src + # idesc literal: M=64, N=256, f32 <- bf16 x bf16, a_major=0 (K-major, bit15 + # clear), b_major=1 — hand-encoded value FlashMLA head64 validated bit-exactly. + assert str(0x04410490) in src, "expected K-major (bit15=0) idesc literal" + assert str(0x04418490) not in src, "MN-major idesc (bit15=1) mis-pairs the K-major descA" + + dev = tvm.cuda(0) + np.random.seed(0) + A_np = np.random.randn(M, K).astype(ml_dtypes.bfloat16) + B_np = np.random.randn(K, N).astype(ml_dtypes.bfloat16) + C_np = np.zeros((128, N // 2), "float32") + A_t, B_t, C_t = (tvm.runtime.tensor(x, dev) for x in (A_np, B_np, C_np)) + mod["main"](A_t, B_t, C_t) + C_ref = A_np.astype("float32") @ B_np.astype("float32") + C_out = C_t.numpy() + # Datapath E: lanes 0-63 hold columns [0, N/2), lanes 64-127 hold [N/2, N). + np.testing.assert_allclose(C_out[:M], C_ref[:, : N // 2], atol=1e-2, rtol=1e-2) + np.testing.assert_allclose(C_out[M:], C_ref[:, N // 2 :], atol=1e-2, rtol=1e-2) + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") @pytest.mark.parametrize("k_lo,k_hi", [(0, 16), (0, 32), (16, 32), (16, 48), (32, 64)]) @@ -2185,7 +2479,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.cuda.cta_sync() tmem = T.decl_buffer((128, N), "float32", scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(128, N) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[0:M, 0:K_alloc], A[0:M, 0:K_alloc], **tma_args) Tx.copy_async(B_smem[0:N, 0:K_alloc], B[0:N, 0:K_alloc], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) @@ -2246,7 +2540,7 @@ def _run_dense_gemm( gemm_kw = {"dispatch": "tcgen05"} if is_AB_tf32: gemm_kw["is_AB_tf32"] = True - b_tma_kw = {"dispatch": "tma"} + b_tma_kw = {"dispatch": "tma_auto"} if tma_dtype_B is not None: b_tma_kw["tma_dtype"] = tma_dtype_B @@ -2281,7 +2575,7 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: layout=TileLayout(S[(128, N) : (1 @ TLane, 1 @ TCol)]), ) if tid_in_wg == 0: - Tx.copy_async(A_smem[:, :], A[:, :], dispatch="tma", mbar=tma_mbar.ptr_to([0])) + Tx.copy_async(A_smem[:, :], A[:, :], dispatch="tma_auto", mbar=tma_mbar.ptr_to([0])) Tx.copy_async(B_smem[:, :], B[:, :], mbar=tma_mbar.ptr_to([0]), **b_tma_kw) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) T.ptx.mbarrier.try_wait(tma_mbar.ptr_to([0]), 0) @@ -2348,8 +2642,103 @@ def test_gemm_tf32_with_tfloat32_tma(): ) -def _build_smem_desc_kernel(smem_desc): +def _run_dense_gemm( + A_dtype, B_dtype, C_dtype, K, *, is_AB_tf32=False, tma_dtype_B=None, atol=1e-3, rtol=1e-3 +): + M, N = 128, 128 + A_shape = (M, K) + B_shape = (N, K) + C_shape = (M, N) + A_swizzle, B_swizzle = 3, 3 + A_layout = mma_shared_layout(A_dtype, A_swizzle, A_shape) + B_layout = mma_shared_layout(B_dtype, B_swizzle, B_shape) + C_elem_32b = 4 // (tvm.runtime.DataType(C_dtype).bits // 8) + cols_alloc = max(32, next_power_of_2(N // C_elem_32b)) + total_bytes = functools.reduce(operator.mul, A_shape, 1) * ( + tvm.runtime.DataType(A_dtype).bits // 8 + ) + functools.reduce(operator.mul, B_shape, 1) * (tvm.runtime.DataType(B_dtype).bits // 8) + gemm_kw = {"dispatch": "tcgen05"} + if is_AB_tf32: + gemm_kw["is_AB_tf32"] = True + b_tma_kw = {"dispatch": "tma_auto"} + if tma_dtype_B is not None: + b_tma_kw["tma_dtype"] = tma_dtype_B + + @T.prim_func + def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, A_shape, A_dtype) + B = T.match_buffer(B_ptr, B_shape, B_dtype) + C = T.match_buffer(C_ptr, C_shape, C_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.cta_id([1]) + wg_id = T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer(A_shape, A_dtype, scope="shared", layout=A_layout) + B_smem = T.alloc_buffer(B_shape, B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + tma_mbar = T.alloc_shared([1], "uint64") + mma_mbar = T.alloc_shared([1], "uint64") + if tid_in_wg == 0: + T.ptx.mbarrier.init(tma_mbar.ptr_to([0]), 1) + T.ptx.mbarrier.init(mma_mbar.ptr_to([0]), 1) + T.ptx.fence.proxy_async("shared::cta") + T.cuda.cta_sync() + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr), n_cols=cols_alloc, cta_group=1) + T.cuda.cta_sync() + tmem = T.decl_buffer( + (128, N), + C_dtype, + scope="tmem", + allocated_addr=tmem_addr[0], + layout=TileLayout(S[(128, N) : (1 @ TLane, 1 @ TCol)]), + ) + if tid_in_wg == 0: + Tx.copy_async(A_smem[:, :], A[:, :], dispatch="tma_auto", mbar=tma_mbar.ptr_to([0])) + Tx.copy_async(B_smem[:, :], B[:, :], mbar=tma_mbar.ptr_to([0]), **b_tma_kw) + T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) + T.ptx.mbarrier.try_wait(tma_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + if tid_in_wg == 0: + Tx.gemm_async(tmem[:, :], A_smem[:, :], B_smem[:, :], **gemm_kw) + T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=1) + T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) + T.cuda.cta_sync() + T.ptx.tcgen05.fence.after_thread_sync() + C_reg = T.alloc_local(N, dtype=C_dtype) + C_view = C_reg.view(128, N, layout=TileLayout(S[(128, N) : (1 @ axis_tid_in_wg, 1)])) + if wg_id == 0: + Tx.wg.copy_async(C_view[:, :], tmem[:, :]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + Tx.copy(C[tid_in_wg, 0:N], C_reg[:]) + if warp_id == 0: + T.ptx.tcgen05.relinquish_alloc_permit(cta_group=1) + T.ptx.tcgen05.dealloc(tmem_addr[0], n_cols=cols_alloc, cta_group=1) + + dev = tvm.cuda(0) + np.random.seed(0) + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": gemm_async}), target=target, tir_pipeline="tirx") + + def _rand(shape, dtype): + f = np.random.randn(*shape).astype("float32") + return f.astype(dtype) if ml_dtypes is not None or "float8" not in dtype else f + + A_np = _rand(A_shape, A_dtype) + B_np = _rand(B_shape, B_dtype) + C_np = np.zeros(C_shape, dtype=C_dtype) + A_t, B_t, C_t = (tvm.runtime.tensor(x, dev) for x in (A_np, B_np, C_np)) + mod["main"](A_t, B_t, C_t) + C_ref = A_np.astype("float32") @ B_np.astype("float32").T + np.testing.assert_allclose(C_t.numpy().astype("float32"), C_ref, atol=atol, rtol=rtol) + + +def _build_smem_desc_kernel(smem_desc, weight_stationary=False, pass_descI=False, mma_config=None): """Minimal cta_group=1 fp16 gemm_async kernel parametrized on ``smem_desc``.""" + mma_cfg = {} if mma_config is None else mma_config C_shape, C_dtype, C_region = (128, 512), "float32", [(0, 128), (256, 384)] A_shape, A_dtype, A_sw = (3, 128, 64), "float16", 3 B_shape, B_dtype, B_sw = (3, 128, 64), "float16", 3 @@ -2392,14 +2781,30 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: T.cuda.cta_sync() tmem = T.decl_buffer((128, C_shape[1]), C_dtype, scope="tmem", allocated_addr=tmem_addr[0], layout=TileLayout(S[(128, C_shape[1]) : (1 @ TLane, 1 @ TCol)])) # noqa: E501 if tid_in_wg == 0: - tma_args = T.meta_var({"dispatch": "tma", "mbar": tma_mbar.ptr_to([0])}) + tma_args = T.meta_var({"dispatch": "tma_auto", "mbar": tma_mbar.ptr_to([0])}) Tx.copy_async(A_smem[tuple(r_gmem_A)], A[tuple(r_gmem_A)], **tma_args) Tx.copy_async(B_smem[tuple(r_gmem_B)], B[tuple(r_gmem_B)], **tma_args) T.ptx.mbarrier.arrive.expect_tx(tma_mbar.ptr_to([0]), total_bytes) T.ptx.mbarrier.try_wait(tma_mbar.ptr_to([0]), 0) T.cuda.cta_sync() if tid_in_wg == 0: - Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", smem_desc=smem_desc) # noqa: E501 + if pass_descI: + desc_i: T.uint32 + T.ptx.tcgen05.encode_instr_descriptor( + T.address_of(desc_i), # noqa: F821 + d_dtype=C_dtype, + a_dtype=A_dtype, + b_dtype=B_dtype, + M=128, + N=width, + K=16, + trans_a=False, + trans_b=False, + n_cta_groups=1, + ) + Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", smem_desc=smem_desc, weight_stationary=weight_stationary, descI=desc_i, **mma_cfg) # noqa: E501, F821 + else: + Tx.gemm_async(tmem[tuple(r_tmem_C)], A_smem[tuple(r_smem_A)], B_smem[tuple(r_smem_B)], dispatch="tcgen05", smem_desc=smem_desc, weight_stationary=weight_stationary, **mma_cfg) # noqa: E501 T.ptx.tcgen05.commit(mma_mbar.ptr_to([0]), cta_group=1) T.ptx.mbarrier.try_wait(mma_mbar.ptr_to([0]), 0) T.cuda.cta_sync() @@ -2419,18 +2824,174 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: return gemm_async -@pytest.mark.parametrize("smem_desc", ["hoist", "recompute"]) -@pytest.mark.gpu -def test_gemm_smem_desc_hoist_vs_recompute(smem_desc): +def _build_explicit_cta2_dense_kernel(M_per_cta, mma_m): + """Minimal cta_group=2 fp16 kernel with an explicit descriptor M.""" + N, K = 64, 16 + A_shape = (M_per_cta, K) + B_shape = (N // 2, K) + A_layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_NONE, A_shape) + B_layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_NONE, B_shape) + if M_per_cta == 64: + C_layout = TileLayout(S[(M_per_cta, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + else: + C_layout = TileLayout(S[(M_per_cta, N) : (1 @ TLane, 1 @ TCol)]) + + # fmt: off + @T.prim_func + def kernel() -> None: + T.device_entry() + cta_id = T.cta_id([2]) + cbx, cby = T.cta_id_in_cluster([2, 1]) + thread_id = T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer(A_shape, "float16", scope="shared", layout=A_layout) + B_smem = T.alloc_buffer(B_shape, "float16", scope="shared", layout=B_layout) + C_tmem = T.decl_buffer((M_per_cta, N), "float32", scope="tmem", allocated_addr=0, layout=C_layout) # noqa: E501 + if tid == 0: + Tx.gemm_async(C_tmem[:, :], A_smem[:, :], B_smem[:, :], dispatch="tcgen05", cta_group=2, mma_m=mma_m, mma_n=N) # noqa: E501 + # fmt: on + + return kernel + + +def _build_explicit_block_scaled_split_n_kernel(): + """Minimal cta_group=2 block-scaled fp8 kernel split into two N instructions.""" + M, N, K = 128, 128, 32 + A_shape = (M, K) + B_shape = (N // 2, K) + A_layout = mma_shared_layout("float8_e4m3fn", SwizzleMode.SWIZZLE_32B_ATOM, A_shape) + B_layout = mma_shared_layout("float8_e4m3fn", SwizzleMode.SWIZZLE_32B_ATOM, B_shape) + C_layout = TileLayout(S[(M, N) : (1 @ TLane, 1 @ TCol)]) + sf_layout = sf_tmem_layout(M, SF_K=1, sf_per_mma=1) + + # fmt: off + @T.prim_func + def kernel() -> None: + T.device_entry() + cta_id = T.cta_id([2]) + cbx, cby = T.cta_id_in_cluster([2, 1]) + thread_id = T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + A_smem = T.alloc_buffer(A_shape, "float8_e4m3fn", scope="shared", layout=A_layout) + B_smem = T.alloc_buffer(B_shape, "float8_e4m3fn", scope="shared", layout=B_layout) + C_tmem = T.decl_buffer((M, N), "float32", scope="tmem", allocated_addr=0, layout=C_layout) + SFA_tmem = T.decl_buffer((M, 1), "float8_e8m0fnu", scope="tmem", allocated_addr=N, layout=sf_layout) # noqa: E501 + SFB_tmem = T.decl_buffer((N, 1), "float8_e8m0fnu", scope="tmem", allocated_addr=N + 4, layout=sf_layout) # noqa: E501 + if tid == 0: + Tx.gemm_async(C_tmem[:, :], A_smem[:, :], B_smem[:, :], SFA=SFA_tmem[:, :], SFB=SFB_tmem[:, :], dispatch="tcgen05", cta_group=2, mma_m=256, mma_n=64) # noqa: E501 + # fmt: on + + return kernel + + +def _compile_cuda_source(func): + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": func}), + target=target, + tir_pipeline="tirx", + ) + return mod.mod.imports[0].inspect_source() + + +def _dense_mma_call_lines(src): + helper = "ptx_tcgen05_mma_cta_1_kind_f16_SS(" + return [line for line in src.splitlines() if line.lstrip().startswith(helper)] + + +def _mma_descriptor_values(lines): + return [int(re.search(r", \(uint\)(\d+), \(bool\)", line).group(1)) for line in lines] + + +def test_gemm_tcgen05_explicit_mma_tile_changes_physical_instructions(): + default_src = _compile_cuda_source(_build_smem_desc_kernel("hoist")) + explicit_src = _compile_cuda_source( + _build_smem_desc_kernel("hoist", mma_config={"mma_m": 128, "mma_n": 64}) + ) + + default_calls = _dense_mma_call_lines(default_src) + explicit_calls = _dense_mma_call_lines(explicit_src) + assert len(default_calls) == 4 + assert len(explicit_calls) == 8 + + assert {(desc >> 17) & 0x3F for desc in _mma_descriptor_values(default_calls)} == {128 // 8} + assert {(desc >> 17) & 0x3F for desc in _mma_descriptor_values(explicit_calls)} == {64 // 8} + + +@pytest.mark.parametrize( + ("M_per_cta", "mma_m"), + [ + (64, 128), + (128, 256), + ], +) +def test_gemm_tcgen05_explicit_mma_m_is_cluster_total_for_cta_group_2(M_per_cta, mma_m): + src = _compile_cuda_source(_build_explicit_cta2_dense_kernel(M_per_cta, mma_m)) + assert "tcgen05.mma.cta_group::2.kind::f16" in src + helper = "ptx_tcgen05_mma_cta_2_kind_f16_SS(" + calls = [line for line in src.splitlines() if line.lstrip().startswith(helper)] + assert len(calls) == 1 + assert {(desc >> 24) & 0x1F for desc in _mma_descriptor_values(calls)} == {mma_m // 16} + + +def test_gemm_tcgen05_block_scaled_explicit_mma_tile_splits_n(): + src = _compile_cuda_source(_build_explicit_block_scaled_split_n_kernel()) + helper = "ptx_tcgen05_mma_block_scaled_cta_2_kind_mxf8f6f4" + calls = [line for line in src.splitlines() if line.lstrip().startswith(helper)] + assert len(calls) == 2 + + +@pytest.mark.parametrize( + ("mma_config", "message"), + [ + ({"mma_m": 128}, "must be provided together"), + ({"mma_n": 64}, "must be provided together"), + ({"mma_m": 0, "mma_n": 64}, "mma_m must be a positive integer"), + ({"mma_m": 128, "mma_n": -1}, "mma_n must be a positive integer"), + ({"mma_m": True, "mma_n": 64}, "mma_m must be a positive integer"), + ({"mma_m": 128.0, "mma_n": 64}, "mma_m must be a positive integer"), + ({"mma_m": 64, "mma_n": 64}, "explicit mma_m must match"), + ({"mma_m": 128, "mma_n": 96}, "explicit mma_n must divide"), + ({"mma_m": 128, "mma_n": 8}, "Invalid matrix shape"), + ], +) +def test_gemm_tcgen05_explicit_mma_tile_rejects_invalid_config(mma_config, message): + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match=message): + tvm.compile( + tvm.IRModule( + { + "main": _build_smem_desc_kernel( + "hoist", + mma_config=mma_config, + ) + } + ), + target=target, + tir_pipeline="tirx", + ) + + +@pytest.mark.parametrize("smem_desc", ["hoist", "local_hoist", "encode", "recompute"]) +def test_gemm_smem_desc_modes_codegen(smem_desc): """Compile-only: the SMEM matrix descriptor is built per-MMA from the buffer base address, selected by the ``smem_desc`` config. - - ``hoist`` (default): allocate + encode one warp-uniform descriptor per - operand (``descA`` / ``descB`` + ``smem_desc_make_lo_uniform``) and add the - per-MMA 16B offset via ``smem_desc_add_16B_offset``. + - ``hoist`` (default): allocate + encode one descriptor per operand + (``descA`` / ``descB``), then add the per-MMA 16B offset via + ``smem_desc_add_16B_offset``. This builder uses a single-thread scope, + where the encoding thread is also the consumer and no warp shuffle is + valid or needed. - ``recompute``: build the full descriptor inline per MMA (``_uniform_desc``) with no allocated/encoded descriptor cell — trades a few ALU ops for one fewer live register on the hot path. + - ``local_hoist``: encode the descriptor at the ``gemm_async`` call site, + inside the caller's elected-thread control flow, then use 16B-offset adds + like hoist mode without an extra warp shuffle. + - ``encode``: encode an exact shared pointer for every MMA, then pass the + resulting descriptor directly to the instruction. Both must emit the MMA; the descriptor-construction fingerprints differ. """ @@ -2445,13 +3006,773 @@ def test_gemm_smem_desc_hoist_vs_recompute(smem_desc): assert "tcgen05.mma" in src, f"mma not emitted; src=\n{src}" if smem_desc == "hoist": - assert "smem_desc_make_lo_uniform" in src, "hoist mode must encode a uniform descriptor" + assert "encode_matrix_descriptor" in src, "hoist mode must encode a descriptor" + assert "smem_desc_make_lo_uniform" not in src, "hoist mode must not warp-shuffle" assert "smem_desc_add_16B_offset" in src, "hoist mode must add the per-MMA 16B offset" + elif smem_desc == "local_hoist": + assert "descA_local" in src and "descB_local" in src + assert "smem_desc_add_16B_offset" in src + assert "encode_matrix_descriptor" in src + elif smem_desc == "encode": + assert "encode_matrix_descriptor" in src + assert "smem_desc_add_16B_offset" not in src else: assert "smem_desc_make_lo_uniform" not in src, "recompute mode must not hoist a descriptor" assert "smem_desc_add_16B_offset" not in src, "recompute mode must not add a 16B offset" assert "encode_matrix_descriptor" not in src, "recompute mode must not encode a descriptor" +def test_gemm_tcgen05_weight_stationary_codegen(): + """FlashMLA head64 requires the tcgen05.mma.ws PTX form.""" + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _build_smem_desc_kernel("hoist", weight_stationary=True)}), + target=target, + tir_pipeline="tirx", + ) + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in src + assert "tvm_builtin_cuda_get_tmem_addr" not in src + + +def test_gemm_tcgen05_dense_descI_rejected(): + """Dense ``descI=`` was removed: the dispatcher self-encodes. + + A hand-passed dense descI performed zero cross-checks against the + dispatcher-constructed descA/descB majorness — historically this masked + the col-major-view majorness desync. Passing it must now raise loudly + (block-scaled gemm_async still accepts descI for the hoisted-encode + + per-ki sf_id rotation pattern; see deepgemm/mqa_logits_fp4). + """ + + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match="descI was removed"): + tvm.compile( + tvm.IRModule( + { + "main": _build_smem_desc_kernel( + "local_hoist", weight_stationary=True, pass_descI=True + ) + } + ), + target=target, + tir_pipeline="tirx", + ) + + +def _build_cta1_m64_packed_c_kernel(weight_stationary=None, mma_config=None): + """cta_group::1 M=64 GEMM whose C uses the packed (M, 2, N//2) TMEM layout. + + ``weight_stationary=None`` omits the flag entirely (the dispatch infers + .ws from the Layout-E C); True/False pass it explicitly. + """ + ws_cfg = {} if weight_stationary is None else {"weight_stationary": weight_stationary} + mma_cfg = {} if mma_config is None else mma_config + + M, N, K = 64, 128, 16 + A_dtype = B_dtype = "bfloat16" + C_dtype = "float32" + B_layout = mma_shared_layout(B_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)) + C_layout = TileLayout(S[(M, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + + @T.prim_func + def gemm_packed_c(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (N, K), B_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N, K), B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=512, cta_group=1) + T.cuda.cta_sync() + # A-in-TMEM .ws reads A from both 64-lane halves, so A is declared in + # the honest batched A[2, M, K] fold (the A-side). + A_tmem = T.decl_buffer( + (2, M, K), + A_dtype, + scope="tmem", + allocated_addr=256, + layout=TileLayout(S[(2, M, K) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]), + ) + C_tmem = T.decl_buffer( + (M, N), + C_dtype, + scope="tmem", + allocated_addr=400, + layout=C_layout, + ) + if tid == 0: + Tx.copy(B_smem[:, :], B[:, :]) + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=1, + **ws_cfg, + **mma_cfg, + ) + + return gemm_packed_c + + +def test_gemm_tcgen05_cta1_m64_accepts_packed_c_layout_ws(): + """FlashMLA head64 stores logical N=128 in 64 physical TMEM columns. + + The packed (M, 2, N//2):(1@TLane, 64@TLane, 1@TCol) C layout is the M=64 + ``.ws`` datapath organization (PTX ISA 8.8 §9.7.16.10.5 Layout E, + cta_group::1), so it is accepted with ``weight_stationary=True``.""" + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _build_cta1_m64_packed_c_kernel(weight_stationary=True)}), + target=target, + tir_pipeline="tirx", + ) + + src = mod.mod.imports[0].inspect_source() + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in src + assert "ptx_tcgen05_mma_cta_1_kind_f16_TS_ws(400," in src + assert "get_tmem_addr(400, 0, 0)" not in src + assert "get_tmem_addr(400, 0, 64)" not in src + + +def test_gemm_tcgen05_explicit_mma_tile_preserves_inferred_ws_datapath(): + src = _compile_cuda_source( + _build_cta1_m64_packed_c_kernel( + mma_config={"mma_m": 64, "mma_n": 64}, + ) + ) + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in src + helper = "ptx_tcgen05_mma_cta_1_kind_f16_TS_ws(" + calls = [line for line in src.splitlines() if line.lstrip().startswith(helper)] + assert len(calls) == 2 + + +def _build_cta1_m64_batched_c_kernel(): + """cta_group::1 M=64 GEMM whose C is the honest batched form C[2, M, N//2] + (the leading dim is the Layout-E lane fold), instead of the packed + C[M, 2, N//2].""" + + M, N, K = 64, 128, 16 + B_dtype = "bfloat16" + B_layout = mma_shared_layout(B_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)) + C_layout = TileLayout(S[(2, M, N // 2) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]) + + @T.prim_func + def gemm_batched_c(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (N, K), B_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N, K), B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=512, cta_group=1) + T.cuda.cta_sync() + # Honest batched A[2, M, K] fold (A-side), matching the + # batched C below: both banks of the M=64 .ws are explicit. + A_tmem = T.decl_buffer( + (2, M, K), + B_dtype, + scope="tmem", + allocated_addr=256, + layout=TileLayout(S[(2, M, K) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]), + ) + C_tmem = T.decl_buffer( + (2, M, N // 2), + "float32", + scope="tmem", + allocated_addr=400, + layout=C_layout, + ) + if tid == 0: + Tx.copy(B_smem[:, :], B[:, :]) + Tx.gemm_async( + # No weight_stationary: the dispatch infers .ws from the + # batched C[2, M, N] fold layout. + C_tmem[:, :, :], + A_tmem[:, :, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=1, + ) + + return gemm_batched_c + + +def _build_cta1_m64_identity_c_ws_kernel(): + """M=64 cta_group::1 gemm forcing .ws but declaring C in the identity + (Layout-D) layout instead of the Layout-E fold — the dual + error the dispatch must reject.""" + + M, N, K = 64, 128, 16 + B_dtype = "bfloat16" + B_layout = mma_shared_layout(B_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)) + + @T.prim_func + def gemm_identity_c(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (N, K), B_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N, K), B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=512, cta_group=1) + T.cuda.cta_sync() + A_tmem = T.decl_buffer( + (M, K), + B_dtype, + scope="tmem", + allocated_addr=256, + layout=TileLayout(S[(M, K) : (1 @ TLane, 1 @ TCol)]), + ) + C_tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=400, + layout=TileLayout(S[(M, N) : (1 @ TLane, 1 @ TCol)]), + ) + if tid == 0: + Tx.copy(B_smem[:, :], B[:, :]) + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=1, + weight_stationary=True, + ) + + return gemm_identity_c + + +def test_gemm_tcgen05_cta1_m64_ws_requires_layout_e_c(): + """A forced .ws (M=64, cta_group::1) with an identity/Layout-F C is + rejected: .ws writes the Layout-E fold, so an identity C would read the + two banks from the wrong lanes (dual error).""" + + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match="Layout-E"): + tvm.compile( + tvm.IRModule({"main": _build_cta1_m64_identity_c_ws_kernel()}), + target=target, + tir_pipeline="tirx", + ) + + +def _build_cta1_m64_flat_a_ws_kernel(): + """M=64 cta_group::1 .ws with A in TMEM declared as a flat 2D [M, K] + identity (lanes 0-63 only) instead of the batched A[2, M, K] fold — the + A-side error the dispatch must reject (the A-side dual of the + identity-C reject above).""" + + M, N, K = 64, 128, 16 + B_dtype = "bfloat16" + B_layout = mma_shared_layout(B_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)) + C_layout = TileLayout(S[(M, 2, N // 2) : (1 @ TLane, 64 @ TLane, 1 @ TCol)]) + + @T.prim_func + def gemm_flat_a(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (N, K), B_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N, K), B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=512, cta_group=1) + T.cuda.cta_sync() + A_tmem = T.decl_buffer( + (M, K), + B_dtype, + scope="tmem", + allocated_addr=256, + layout=TileLayout(S[(M, K) : (1 @ TLane, 1 @ TCol)]), + ) + C_tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=400, + layout=C_layout, + ) + if tid == 0: + Tx.copy(B_smem[:, :], B[:, :]) + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=1, + weight_stationary=True, + ) + + return gemm_flat_a + + +def test_gemm_tcgen05_cta1_m64_ws_requires_batched_a(): + """A forced .ws (M=64, cta_group::1) with A in TMEM declared as a flat 2D + [M, K] identity is rejected: the .ws reads A from both 64-lane halves, so a + flat A (lanes 0-63 only) cannot express — nor let the dispatch verify — the + second-half occupancy. A must use the batched A[2, M, K] fold.""" + + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match="both 64-lane halves"): + tvm.compile( + tvm.IRModule({"main": _build_cta1_m64_flat_a_ws_kernel()}), + target=target, + tir_pipeline="tirx", + ) + + +def _build_m128_batched_a_kernel(): + """M=128 cta_group::1 NON-ws gemm whose A is (illegitimately) declared in the + batched A[2, M, K] fold. The batched fold is defined *only* for the M=64 + cta_group::1 .ws datapath, so it must be rejected here (the converse of the + requires-batched-A reject: batched A only for M=64 .ws).""" + + M, N, K = 128, 128, 16 + B_dtype = "bfloat16" + B_layout = mma_shared_layout(B_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)) + + @T.prim_func + def gemm_m128_batched_a(B_ptr: T.handle) -> None: + B = T.match_buffer(B_ptr, (N, K), B_dtype) + T.device_entry() + warp_id = T.warp_id([4]) + T.thread_id([128]) + tid = T.thread_id_in_wg([128]) + B_smem = T.alloc_buffer((N, K), B_dtype, scope="shared", layout=B_layout) + tmem_addr = T.alloc_shared([1], "uint32") + if warp_id == 0: + T.ptx.tcgen05.alloc(T.address_of(tmem_addr[0]), n_cols=512, cta_group=1) + T.cuda.cta_sync() + A_tmem = T.decl_buffer( + (2, M, K), + B_dtype, + scope="tmem", + allocated_addr=256, + layout=TileLayout(S[(2, M, K) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]), + ) + C_tmem = T.decl_buffer( + (M, N), + "float32", + scope="tmem", + allocated_addr=400, + layout=TileLayout(S[(M, N) : (1 @ TLane, 1 @ TCol)]), + ) + if tid == 0: + Tx.copy(B_smem[:, :], B[:, :]) + Tx.gemm_async( + C_tmem[:, :], + A_tmem[:, :, :], + B_smem[:, :], + dispatch="tcgen05", + cta_group=1, + ) + + return gemm_m128_batched_a + + +def test_gemm_tcgen05_batched_a_rejects_unproven_datapath(): + """A batched A[2, M, K] fold outside Layout E / Layout B is rejected.""" + + target = tvm.target.Target("cuda") + with target: + with pytest.raises(Exception, match="only valid for Layout E"): + tvm.compile( + tvm.IRModule({"main": _build_m128_batched_a_kernel()}), + target=target, + tir_pipeline="tirx", + ) + + +def test_gemm_tcgen05_cta1_m64_accepts_batched_c_layout_ws(): + """The honest batched C[2, M, N//2] .ws output form is accepted and emits + byte-identically to the packed C[M, 2, N//2] form: the two + describe the same physical Layout-E tile.""" + + target = tvm.target.Target("cuda") + + def _compile(kernel): + with target: + mod = tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") + return mod.mod.imports[0].inspect_source() + + batched_src = _compile(_build_cta1_m64_batched_c_kernel()) + packed_src = _compile(_build_cta1_m64_packed_c_kernel(weight_stationary=True)) + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in batched_src + + # identical up to the kernel entry name (gemm_batched_c vs gemm_packed_c) + def norm(s): + return s.replace("gemm_batched_c", "K").replace("gemm_packed_c", "K") + + assert norm(batched_src) == norm(packed_src) + + +def test_gemm_tcgen05_cta1_m64_packed_c_infers_weight_stationary(): + """The packed (M, 2, N//2):(1@TLane, 64@TLane, 1@TCol) Layout-E C is + uniquely the M=64 .ws datapath (PTX ISA 8.8 §9.7.16.10.5), so .ws is + inferred with no weight_stationary flag; and an explicit + weight_stationary=False contradicts the layout and is rejected.""" + + target = tvm.target.Target("cuda") + with target: + # omitted flag -> .ws inferred, compiles to a .ws MMA + mod = tvm.compile( + tvm.IRModule({"main": _build_cta1_m64_packed_c_kernel()}), + target=target, + tir_pipeline="tirx", + ) + assert "tcgen05.mma.ws.cta_group::1.kind::f16" in mod.mod.imports[0].inspect_source() + # explicit False contradicts the Layout-E fold -> rejected + with pytest.raises(Exception, match="weight_stationary=False"): + tvm.compile( + tvm.IRModule({"main": _build_cta1_m64_packed_c_kernel(weight_stationary=False)}), + target=target, + tir_pipeline="tirx", + ) + + +# Dispatch-level regression tests: call gemm_async_tcgen05_impl directly on a +# constructed TilePrimitiveCall to pin rejection paths without full compilation. + + +def _make_gemm_tcgen05_call( + M, + N, + K, + dtype, + A_layout, + B_layout, + transA=False, + transB=True, + config=None, + scope_kind="thread", + return_context=False, + C_layout=None, + C_allocated_addr=0, + A_scope="shared.dyn", + A_allocated_addr=0, +): + """Construct a GemmAsync TilePrimitiveCall and run the tcgen05 dispatch. + + Buffer-shape convention follows the dispatcher: transA=False -> A is + [M, K]; transB=True -> B is [K, N], transB=False -> B is [N, K]. + C is a full-region (M, N) float32 TMEM buffer with the identity + (1@TLane, 1@TCol) layout. + """ + from tvm.ir import Range + from tvm.tirx.cuda.operator.tile_primitive.gemm_async.tcgen05 import ( + gemm_async_tcgen05_impl, + ) + from tvm.tirx.exec_scope import ExecScope + from tvm.tirx.operator.tile_primitive.ops import GemmAsync + from tvm.tirx.stmt import BufferRegion + from tvm.tirx.tile_primitive import DispatchContext + + def full_region(buf): + return BufferRegion(buf, [Range.from_min_extent(0, s) for s in buf.shape]) + + A_shape = (M, K) if not transA else (K, M) + B_shape = (K, N) if transB else (N, K) + A_buf = tvm.tirx.decl_buffer(A_shape, dtype, "A", scope=A_scope, layout=A_layout) + if A_scope == "tmem": + A_buf = A_buf.with_allocated_addr([tvm.tirx.IntImm("uint32", A_allocated_addr)]) + B_buf = tvm.tirx.decl_buffer(B_shape, dtype, "B_smem", scope="shared.dyn", layout=B_layout) + if C_layout is None: + C_layout = TileLayout(S[(M, N) : (1 @ TLane, 1 @ TCol)]) + C_buf = tvm.tirx.decl_buffer((M, N), "float32", "C_tmem", scope="tmem", layout=C_layout) + C_buf = C_buf.with_allocated_addr([tvm.tirx.IntImm("uint32", C_allocated_addr)]) + call = GemmAsync( + full_region(C_buf), + full_region(A_buf), + full_region(B_buf), + transA, + transB, + False, + config=dict(config or {}), + ) + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + sctx = DispatchContext(target, ExecScope(scope_kind), {}, {}, scope_kind=scope_kind) + impl = gemm_async_tcgen05_impl(call, sctx) + return (impl, sctx) if return_context else impl + + +def test_gemm_tcgen05_preserves_explicit_tmem_lane_bases(): + """D and TMEM-A row offsets are part of their physical taddr operands.""" + M, N, K = 64, 64, 16 + dtype = "bfloat16" + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_32B_ATOM, (K, N)) + + d_base = tmem_datapath_layout("F", M, N) + d_offset = TileLayout.from_iters(d_base.shard, d_base.replica, {TLane: 1}) + A_smem_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_32B_ATOM, (M, K)) + d_impl = _make_gemm_tcgen05_call( + M, + N, + K, + dtype, + A_smem_layout, + B_layout, + C_layout=d_offset, + C_allocated_addr=400, + config={"mma_m": M, "mma_n": N}, + ) + assert "T.cuda.get_tmem_addr(T.uint32(400), 1, ni * 64)" in d_impl.script() + + a_base = TileLayout(S[(M, K) : (1 @ TLane, 1 @ TCol)]) + a_offset = TileLayout.from_iters(a_base.shard, a_base.replica, {TLane: 1}) + a_impl = _make_gemm_tcgen05_call( + M, + N, + K, + dtype, + a_offset, + B_layout, + C_layout=d_base, + C_allocated_addr=400, + A_scope="tmem", + A_allocated_addr=256, + config={"mma_m": M, "mma_n": N}, + ) + assert "T.cuda.get_tmem_addr(T.uint32(256), 1, ki * 8)" in a_impl.script() + + +def test_gemm_tcgen05_preserves_block_scale_tmem_lane_bases(): + """SFA/SFB row offsets must reach the encoded TMEM address operands.""" + from tvm.ir import Range + from tvm.tirx.cuda.operator.tile_primitive.gemm_async.tcgen05 import ( + gemm_async_tcgen05_impl, + ) + from tvm.tirx.exec_scope import ExecScope + from tvm.tirx.operator.tile_primitive.ops import GemmAsync + from tvm.tirx.stmt import BufferRegion + from tvm.tirx.tile_primitive import DispatchContext + + def full_region(buf): + return BufferRegion(buf, [Range.from_min_extent(0, s) for s in buf.shape]) + + M, N, K = 128, 64, 64 + data_dtype = "float4_e2m1fn" + sf_dtype = "float8_e4m3fn" + A = tvm.tirx.decl_buffer( + (M, K), + data_dtype, + "A_smem", + scope="shared.dyn", + layout=mma_shared_layout(data_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (M, K)), + ) + B = tvm.tirx.decl_buffer( + (N, K), + data_dtype, + "B_smem", + scope="shared.dyn", + layout=mma_shared_layout(data_dtype, SwizzleMode.SWIZZLE_32B_ATOM, (N, K)), + ) + C = tvm.tirx.decl_buffer( + (M, N), + "float32", + "C_tmem", + scope="tmem", + layout=tmem_datapath_layout("D", M, N), + ).with_allocated_addr([tvm.tirx.IntImm("uint32", 0)]) + + sf_per_mma = 4 + sfa_base = sf_tmem_layout(M, SF_K=sf_per_mma, sf_per_mma=sf_per_mma) + sfb_base = sf_tmem_layout(N, SF_K=sf_per_mma, sf_per_mma=sf_per_mma) + sfa_layout = TileLayout.from_iters(sfa_base.shard, sfa_base.replica, {TLane: 1}) + sfb_layout = TileLayout.from_iters(sfb_base.shard, sfb_base.replica, {TLane: 2}) + SFA = tvm.tirx.decl_buffer( + (M, sf_per_mma), sf_dtype, "SFA_tmem", scope="tmem", layout=sfa_layout + ).with_allocated_addr([tvm.tirx.IntImm("uint32", 256)]) + SFB = tvm.tirx.decl_buffer( + (N, sf_per_mma), sf_dtype, "SFB_tmem", scope="tmem", layout=sfb_layout + ).with_allocated_addr([tvm.tirx.IntImm("uint32", 320)]) + + call = GemmAsync( + full_region(C), + full_region(A), + full_region(B), + full_region(SFA), + full_region(SFB), + False, + False, + False, + config={"cta_group": 1, "mma_m": M, "mma_n": N}, + ) + target = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) + sctx = DispatchContext(target, ExecScope("thread"), {}, {}, scope_kind="thread") + script = gemm_async_tcgen05_impl(call, sctx).script() + + assert "T.cuda.get_tmem_addr(T.uint32(256), 1, 0)" in script + assert "T.cuda.get_tmem_addr(T.uint32(320), 2, 0)" in script + + +@pytest.mark.parametrize( + ("scope_kind", "expect_uniform"), + [("thread", False), ("warp", True)], +) +def test_gemm_tcgen05_hoisted_descriptor_uniformization(scope_kind, expect_uniform): + M, N, K = 64, 256, 64 + dtype = "bfloat16" + A_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (M, K)) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + impl, sctx = _make_gemm_tcgen05_call( + M, + N, + K, + dtype, + A_layout, + B_layout, + config={"smem_desc": "hoist"}, + scope_kind=scope_kind, + return_context=True, + ) + + callback_text = str(sctx.callbacks["post_buffer_def_stmt"]) + assert "encode_matrix_descriptor" in callback_text + assert ("smem_desc_make_lo_uniform" in callback_text) == expect_uniform + assert ("elect_sync" in impl.script()) == expect_uniform + if expect_uniform: + assert callback_text.index("smem_desc_make_lo_uniform") < callback_text.index( + "tvm_kernel_replace_point" + ) + + +def test_gemm_tcgen05_no_swizzle_col_major_rejects_non_128B_contiguous(): + """the col-major-view no-swizzle branch encodes the SBO field as + the literal 8 (128B row-group pitch), which only matches the packed ABI + when the contiguous dim spans exactly 128B (64 bf16 elements). A larger + contiguous dim (M=128 here) used to be accepted with a silently wrong + ``sdo = shape // elem_per_16B = 16`` — it must now be rejected.""" + M, N, K = 128, 256, 64 + dtype = "bfloat16" + A_layout = TileLayout(S[(M, K) : (1, M)]) # column-major view, contiguous dim 128 + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + with pytest.raises(ValueError, match="span exactly 128B"): + _make_gemm_tcgen05_call(M, N, K, dtype, A_layout, B_layout) + + # In-domain sanity: M=64 col-major view still dispatches. + A_ok = TileLayout(S[(64, K) : (1, 64)]) + impl = _make_gemm_tcgen05_call(64, N, K, dtype, A_ok, B_layout) + assert impl is not None + + +def test_gemm_tcgen05_no_swizzle_packed_16b_rejects_non_16bit_dtype(): + """the packed-16B no-swizzle branch used to encode + ``sdo = elem_per_16B``, which equals the true SBO field (8 = 128B row + pitch, PTX ISA §9.7.16.3.2) only for 16-bit dtypes. An fp8 packed-16B + layout (elem_per_16B = 16) used to be accepted with sdo=16 — it must now + be rejected (that domain is not hardware-validated).""" + M, N, K = 64, 256, 64 + dtype = "float8_e4m3fn" # 16 elements per 16B line + A_layout = TileLayout(S[(M, K // 16, 16) : (16, M * 16, 1)]) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + with pytest.raises(ValueError, match="16-bit dtypes"): + _make_gemm_tcgen05_call(M, N, K, dtype, A_layout, B_layout) + + # In-domain sanity: the bf16 packed-16B layout still dispatches. + bf16_A = TileLayout(S[(M, K // 8, 8) : (8, M * 8, 1)]) + bf16_B = mma_shared_layout("bfloat16", SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + impl = _make_gemm_tcgen05_call(M, N, K, "bfloat16", bf16_A, bf16_B) + assert impl is not None + + +def test_gemm_tcgen05_instr_desc_fold_mirrors_runtime_shape_rules(): + """the compile-time descI fold must run the runtime encoder's + shape validation. cta_group=1 with descriptor M=128 requires N % 16 == 0; + the tile chooser alone only guarantees N % 8, so M=128/N=24 used to fold + a descriptor the runtime encoder would reject.""" + M, N, K = 128, 24, 64 + dtype = "bfloat16" + A_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (M, K)) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (N, K)) + with pytest.raises(ValueError, match="Invalid matrix shape"): + _make_gemm_tcgen05_call(M, N, K, dtype, A_layout, B_layout, transB=False) + + # N=32 (divisible by 16) with the same M=128 tile is accepted. + B_ok = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (32, K)) + impl = _make_gemm_tcgen05_call(M, 32, K, dtype, A_layout, B_ok, transB=False) + assert impl is not None + + +def test_gemm_tcgen05_rejects_tf32_mn_major(): + """tf32 MN-major operands are PTX-illegal without the + 128B-32B-atomicity swizzle, which the atom matcher never produces. A + tf32 (is_AB_tf32) B operand matching MN-major used to be silently + encoded; it must now be rejected.""" + M, N, K = 64, 256, 32 + dtype = "float32" + A_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (M, K)) + # transB=True means B is [K, N]; a row-major swizzled layout then has the + # MN (=N) dim contiguous, i.e. it matches as MN-major. + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + with pytest.raises(ValueError, match="PTX-illegal"): + _make_gemm_tcgen05_call( + M, N, K, dtype, A_layout, B_layout, transB=True, config={"is_AB_tf32": True} + ) + + # MN-major stays accepted for a dtype where it is PTX-legal (bf16). + bf16_A = mma_shared_layout("bfloat16", SwizzleMode.SWIZZLE_128B_ATOM, (M, 64)) + bf16_B = mma_shared_layout("bfloat16", SwizzleMode.SWIZZLE_128B_ATOM, (64, N)) + impl = _make_gemm_tcgen05_call(M, N, 64, "bfloat16", bf16_A, bf16_B, transB=True) + assert impl is not None + + +def test_gemm_tcgen05_rejects_non_uniform_atom_grid(): + """``_try_atom`` reads only the innermost iter of each tiler + dimension for the LBO/SBO fields, so each dimension must group into a + single iter. A swizzled layout whose M-direction atom tiling has a + stride gap (atoms (2,4) with outer stride 4096 instead of 2048) used to + match with fields taken from the local inner stride, silently dropping + the gap; it must now be rejected.""" + from tvm.tirx.layout import ComposeLayout + + M, N, K = 64, 256, 64 + dtype = "bfloat16" + exotic_A = ComposeLayout(3, 3, 3, TileLayout(S[(2, 4, 8, 64) : (4096, 512, 64, 1)])) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + with pytest.raises(ValueError, match="no MMA SMEM descriptor matches"): + _make_gemm_tcgen05_call(M, N, K, dtype, exotic_A, B_layout) + + # The uniform version of the same tiling (outer stride 2048) is accepted. + uniform_A = ComposeLayout(3, 3, 3, TileLayout(S[(2, 4, 8, 64) : (2048, 512, 64, 1)])) + impl = _make_gemm_tcgen05_call(M, N, K, dtype, uniform_A, B_layout) + assert impl is not None + + +def test_gemm_tcgen05_dense_descI_rejected_at_dispatch(): + """Dispatch-level twin of the compile-path test above.""" + M, N, K = 64, 256, 64 + dtype = "bfloat16" + A_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (M, K)) + B_layout = mma_shared_layout(dtype, SwizzleMode.SWIZZLE_128B_ATOM, (K, N)) + with pytest.raises(ValueError, match="descI was removed"): + _make_gemm_tcgen05_call( + M, + N, + K, + dtype, + A_layout, + B_layout, + config={"descI": tvm.tirx.const(0x04410490, "uint32")}, + ) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py b/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py index 4e589a7031f0..b1fc86bd62cf 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py @@ -51,7 +51,7 @@ _check_bijection, _choose_xor_k, ) -from tvm.tirx.layout import S, SwizzleLayout, TileLayout +from tvm.tirx.layout import S, TileLayout # --------------------------------------------------------------------------- # Algorithm-only tests (no CUDA needed). @@ -374,12 +374,18 @@ def f(A: T.handle, B: T.handle): def test_reject_swizzle_layout(): - """ComposeLayout(SwizzleLayout, TileLayout) is not supported by the warp variant.""" + """A swizzled ComposeLayout tile is not supported by the warp variant.""" from tvm.tirx.layout import ComposeLayout inner = TileLayout(S[(4, 32) : (32, 1)]) - sw = SwizzleLayout(per_element=2, swizzle_len=2, atom_len=4) - swizzled = ComposeLayout(sw, inner) + sw = ComposeLayout(2, 2, 4, TileLayout(S[(256,)])) + swizzled = ComposeLayout( + sw.per_element, + sw.swizzle_len, + sw.atom_len, + inner, + sw.swizzle_inner, + ) plain = TileLayout(S[(4, 32) : (1, 4)]) # fmt: off diff --git a/tests/python/tirx/operator/tile_primitive/test_dispatcher.py b/tests/python/tirx/operator/tile_primitive/test_dispatcher.py index 5ff5fe6caeaf..d5bc3b6fe7f8 100644 --- a/tests/python/tirx/operator/tile_primitive/test_dispatcher.py +++ b/tests/python/tirx/operator/tile_primitive/test_dispatcher.py @@ -134,7 +134,7 @@ def test_dispatch_prints_real_opcall_ir(): from tvm.ir import Op from tvm.tirx.buffer import decl_buffer from tvm.tirx.operator.tile_primitive.dispatcher import run_dispatch - from tvm.tirx.stmt import TilePrimitiveCall + from tvm.tirx.tile_primitive import TilePrimitiveCall # Build a real TIRx TilePrimitiveCall: tirx.tile.copy(A[0:64], B[0:64]) A = decl_buffer((64,), "float32", scope="global") diff --git a/tests/python/tirx/test_alloc_pool.py b/tests/python/tirx/test_alloc_pool.py index 41ea560ed706..ebc5f751a6f5 100644 --- a/tests/python/tirx/test_alloc_pool.py +++ b/tests/python/tirx/test_alloc_pool.py @@ -22,7 +22,7 @@ from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode # --------------------------------------------------------------------------- -# alloc_mma shape validation: bad inputs raise actionable ValueError instead of +# alloc_tcgen05_mma_AB shape validation: bad inputs raise actionable ValueError instead of # the opaque "Divide by zero" diagnostic that ``Layout.tile_to`` would emit. # --------------------------------------------------------------------------- diff --git a/tests/python/tirx/test_bench_utils.py b/tests/python/tirx/test_bench_utils.py index ccf5e2122653..898240263644 100644 --- a/tests/python/tirx/test_bench_utils.py +++ b/tests/python/tirx/test_bench_utils.py @@ -16,229 +16,322 @@ # under the License. """Tests for tvm.tirx.bench utilities.""" +import importlib +import inspect +from types import SimpleNamespace + import pytest import torch -import tvm.testing - pytest.importorskip("triton") # tvm.tirx.bench imports triton.profiler from tvm.testing import env -from tvm.tirx.bench import _compute_group_count, _parse_proton_tree, bench, tensor_bytes +from tvm.tirx.bench import DistributedBenchContext, bench -# ── _parse_proton_tree ────────────────────────────────────────────────────── +bench_module = importlib.import_module("tvm.tirx.bench") -SAMPLE_TREE = """\ -├─ 1.500 tir -│ ├─ 1.500 my_kernel_fn -│ └─ 0.001 vectorized_elementwise_kernel -└─ 0.800 cublas - └─ 0.800 sm90_xmma_gemm_f16f16 -""" +class _FakeStream: + cuda_stream = 123 + def synchronize(self): + pass -def test_parse_proton_tree_basic(): - impls, errors = _parse_proton_tree(SAMPLE_TREE) - assert impls == {"tir": 1.5, "cublas": 0.8} - assert errors == {} +def _distributed_context(max_reduce=lambda value: value): + return DistributedBenchContext( + rank=0, + world_size=4, + barrier=lambda: None, + max_reduce=max_reduce, + stream=_FakeStream(), + ) -def test_parse_proton_tree_filters_elementwise(): - """vectorized_elementwise_kernel and elementwise_kernel_with_index are skipped.""" - tree = """\ -├─ 0.500 tir -│ ├─ 0.500 real_kernel -│ └─ 0.001 elementwise_kernel_with_index -""" - impls, _ = _parse_proton_tree(tree) - assert impls == {"tir": 0.5} +def test_bench_cooldown_precedes_every_impl(monkeypatch): + """cooldown_s sleeps immediately before each impl's warmup+measurement. -def test_parse_proton_tree_slowest_child(): - """Takes the slowest depth-2 child per impl.""" - tree = """\ -├─ 2.000 tir -│ ├─ 0.300 kernel_a -│ └─ 0.700 kernel_b -""" - impls, _ = _parse_proton_tree(tree) - assert impls == {"tir": 0.7} + 2 impls x 2 rounds = 4 timed calls, so 4 sleeps (the first impl in the + first round is included). Pins the #29 per-impl cooldown semantics. + """ + calls = [] + sleeps = [] + def fake_timer(fn, warmup=25, rep=100): + del warmup, rep + fn() + return 0.001 -def test_parse_proton_tree_baseline_errors(): - tree = """\ -BASELINE_ERROR: cublas: CUDA OOM -├─ 1.000 tir -│ └─ 1.000 my_kernel -""" - impls, errors = _parse_proton_tree(tree) - assert impls == {"tir": 1.0} - assert errors == {"cublas": "CUDA OOM"} + monkeypatch.setattr(bench_module, "_do_bench_event", fake_timer) + monkeypatch.setattr(bench_module.time, "sleep", sleeps.append) + results = bench( + {"a": lambda: calls.append("a"), "b": lambda: calls.append("b")}, + warmup=0, + repeat=1, + timer="event", + cooldown_s=1.0, + rounds=2, + ) -def test_parse_proton_tree_ansi_stripped(): - """ANSI color codes are stripped before parsing.""" - tree = "\x1b[32m├─ 1.000 tir\x1b[0m\n│ └─ 1.000 k\n" - impls, _ = _parse_proton_tree(tree) - assert impls == {"tir": 1.0} + assert calls == ["a", "b", "a", "b"] + assert sleeps == [1.0, 1.0, 1.0, 1.0] + assert results["benchmark_protocol"]["cooldown_s"] == 1.0 + assert results["benchmark_protocol"]["round_aggregate"] == "mean" -def test_parse_proton_tree_empty(): - impls, errors = _parse_proton_tree("") - assert impls == {} - assert errors == {} +def test_bench_retains_round_samples_and_uses_arithmetic_mean(monkeypatch): + values = iter([0.001, 0.002, 0.100]) + def fake_timer(_fn, warmup=25, rep=100): + del warmup, rep + return next(values) -# ── bench ─────────────────────────────────────────────────────────────────── + monkeypatch.setattr(bench_module, "_do_bench_event", fake_timer) + results = bench({"tir": lambda: None}, timer="event", cooldown_s=0, rounds=3) -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_basic(): - """bench returns positive times for each impl.""" - M, N = 256, 256 + assert results["round_samples"] == {"tir": [1.0, 2.0, 100.0]} + assert results["impls"] == {"tir": 103.0 / 3.0} - funcs = {"matmul": lambda case: torch.mm(case[0], case[1])} - def make_input(): - A = torch.randn(M, N, device="cuda", dtype=torch.float16) - B = torch.randn(M, N, device="cuda", dtype=torch.float16) - return (A, B), tensor_bytes(A, B) +def test_bench_l2_flush_buffer_matches_triton_256_mib(monkeypatch): + captured = {} - def run_and_check(): - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") - assert "matmul" in results["impls"] - assert results["impls"]["matmul"] > 0 + def fake_empty(size, *, dtype, device): + captured.update(size=size, dtype=dtype, device=device) + return object() - tvm.testing.run_with_gpu_lock(run_and_check) + monkeypatch.setattr(bench_module.torch, "empty", fake_empty) + bench_module._empty_cache_for_benchmark() -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_multiple_impls(): - """Multiple impls each get their own timing.""" - M, N = 128, 128 - funcs = { - "mm": lambda case: torch.mm(case[0], case[1]), - "addmm": lambda case: torch.addmm( - torch.zeros(M, N, device="cuda", dtype=torch.float16), case[0], case[1] - ), + assert captured == { + "size": 256 * 1024 * 1024 // 4, + "dtype": torch.int, + "device": "cuda", } - def make_input(): - A = torch.randn(M, N, device="cuda", dtype=torch.float16) - B = torch.randn(M, N, device="cuda", dtype=torch.float16) - return (A, B), tensor_bytes(A, B) - def run_and_check(): - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") - assert set(results["impls"].keys()) == {"mm", "addmm"} - assert all(v > 0 for v in results["impls"].values()) +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_bench_event_pure_launch(): + """New Triton-standard bench(): no-arg launch closures, event timer.""" + M, N = 256, 256 + A = torch.randn(M, N, device="cuda", dtype=torch.float16) + B = torch.randn(M, N, device="cuda", dtype=torch.float16) + + funcs = {"mm": lambda: torch.mm(A, B)} + results = bench(funcs, warmup=5, repeat=10, timer="event") + assert "mm" in results["impls"] + assert results["impls"]["mm"] > 0 + assert results["timer"] == "event" + + +def test_bench_default_timer_is_proton(monkeypatch): + """Omitting timer resolves to Proton and invokes only the Proton timer.""" + calls = [] + + def fake_proton(fn, warmup=25, rep=100): + calls.append((warmup, rep)) + fn() + return 0.001 + + monkeypatch.setattr(bench_module, "_do_bench_proton", fake_proton) + + results = bench({"noop": lambda: None}, cooldown_s=0) + + assert results["timer"] == "proton" + assert results["impls"] == {"noop": 1.0} + assert calls == [(25, 100)] + + +def test_distributed_kineto_span_uses_cross_stream_samples_and_rank_max(monkeypatch): + reductions = [] + local_samples = [float(index + 1) for index in range(30)] + + def fake_profile(name, func, prepare, distributed): + assert name == "impl" + assert callable(func) + assert callable(prepare["impl"]) + assert distributed.world_size == 4 + return local_samples, [2] * len(local_samples) + + def max_reduce(value): + reductions.append(value) + return value + 100.0 + + monkeypatch.setattr(bench_module, "_profile_distributed_kineto_span", fake_profile) + + result = bench( + {"impl": lambda: None}, + distributed=_distributed_context(max_reduce), + prepare={"impl": lambda: None}, + cooldown_s=0, + ) + + assert result["timer"] == "kineto" + assert result["impls"] == {"impl": 115.5} + assert reductions == local_samples + protocol = result["benchmark_protocol"] + assert protocol["timing_scope"] == "complete correlated GPU activity span" + assert protocol["span_definition"] == "latest activity end minus earliest activity start" + assert protocol["rank_aggregate"] == "sample_wise_max" + assert protocol["rank_local_scope_stream_counts"] == {"impl": [{"min": 2, "max": 2}]} + + +def test_kineto_span_collector_uses_earliest_start_and_latest_end(): + def event(name, device_type, stream, start, end): + return SimpleNamespace( + name=name, + device_type=device_type, + device_resource_id=stream, + time_range=SimpleNamespace(start=start, end=end), + ) - tvm.testing.run_with_gpu_lock(run_and_check) + profiler = SimpleNamespace( + events=lambda: [ + event("sample.0", torch.autograd.DeviceType.CPU, 0, 0.0, 10.0), + event("sample.0", torch.autograd.DeviceType.CUDA, 11, 2.0, 5.0), + event("sample.0", torch.autograd.DeviceType.CUDA, 17, 3.0, 8.0), + event("sample.1", torch.autograd.DeviceType.CUDA, 11, 20.0, 21.5), + ] + ) + + samples, stream_counts = bench_module._collect_kineto_span_samples( + profiler, ["sample.0", "sample.1"] + ) + + assert samples == [6.0, 1.5] + assert stream_counts == [2, 1] + + +def test_distributed_kineto_span_keeps_fixed_order_without_ab_ba(): + source = inspect.getsource(bench_module._bench_distributed_kineto_span) + assert "reversed" not in source + assert "round_orders" not in source + + +@pytest.mark.parametrize( + "kwargs, match", + [ + ({"timer": "event"}, "only timer='kineto'"), + ({"timer": "proton"}, "only timer='kineto'"), + ({"warmup": 1}, "rejects overrides: warmup"), + ({"repeat": 1}, "rejects overrides: repeat"), + ({"cudagraph_rep": 1}, "rejects overrides: cudagraph_rep"), + ], +) +def test_distributed_timers_reject_invalid_timer_and_budgets(kwargs, match): + with pytest.raises(ValueError, match=match): + bench( + {"noop": lambda: None}, + distributed=_distributed_context(), + cooldown_s=0, + **kwargs, + ) -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_multiple_input_groups(): - """Multiple input groups cycle correctly (L2 eviction).""" - M, N = 128, 128 - call_count = [0] - - def make_input(): - call_count[0] += 1 - A = torch.randn(M, N, device="cuda", dtype=torch.float16) - B = torch.randn(M, N, device="cuda", dtype=torch.float16) - return (A, B), tensor_bytes(A, B) - - funcs = {"mm": lambda case: torch.mm(case[0], case[1])} - - def run_and_check(): - results = bench( - funcs, - make_input, - warmup=5, - repeat=20, - cooldown_s=0.0, - timer="event", - l2_bytes=64 * 1024, +@pytest.mark.parametrize( + ("overrides", "error", "match"), + [ + ({"world_size": 0}, ValueError, "world_size must be a positive integer"), + ({"world_size": True}, ValueError, "world_size must be a positive integer"), + ({"rank": -1}, ValueError, "rank must be in"), + ({"rank": 4}, ValueError, "rank must be in"), + ({"barrier": None}, TypeError, "barrier must be callable"), + ({"max_reduce": None}, TypeError, "max_reduce must be callable"), + ({"stream": object()}, TypeError, "stream must be an actual CUDA stream"), + ], +) +def test_distributed_timer_rejects_invalid_context(overrides, error, match): + values = { + "rank": 0, + "world_size": 4, + "barrier": lambda: None, + "max_reduce": lambda value: value, + "stream": _FakeStream(), + } + values.update(overrides) + + with pytest.raises(error, match=match): + bench( + {"noop": lambda: None}, + distributed=DistributedBenchContext(**values), + cooldown_s=0, ) - assert results["impls"]["mm"] > 0 - assert call_count[0] > 1 - tvm.testing.run_with_gpu_lock(run_and_check) +def test_prepare_requires_distributed_context(): + with pytest.raises(ValueError, match="only with a distributed context"): + bench({"noop": lambda: None}, prepare={"noop": lambda: None}, cooldown_s=0) -# ── _compute_group_count ─────────────────────────────────────────────────── +def test_bench_cudagraph_proton_wiring(monkeypatch): + calls = [] -def test_compute_groups_small_tensors(): - """Small tensors need many groups to fill 3x L2.""" - # 128x128 fp16 = 32KB. 3*128MB / 32KB = 12288, +1 = 12289 - input_bytes = tensor_bytes(torch.empty(128, 128, dtype=torch.float16)) - n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024) - assert n == 12289 + def fake_cudagraph_proton(fn, rep=20): + calls.append(rep) + fn() + return 0.002 + monkeypatch.setattr(bench_module, "_do_bench_cudagraph_proton", fake_cudagraph_proton) -def test_compute_groups_large_tensors(): - """Inputs >= 3x L2 need only 1 group.""" - # 16384x16384 fp32 = 1GB >> 3*128MB = 384MB - input_bytes = tensor_bytes(torch.empty(16384, 16384, dtype=torch.float32)) - n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024) - assert n == 1 + results = bench({"noop": lambda: None}, timer="cudagraph_proton", cudagraph_rep=7, cooldown_s=0) + assert results["impls"] == {"noop": 2.0} + assert results["timer"] == "cudagraph_proton" + assert calls == [7] -def test_compute_groups_moderate_tensors(): - """Moderate tensors: floor(3*L2 / input) + 1.""" - # 8192x8192 bf16 = 128MB. floor(384M / 128M) + 1 = 4 - input_bytes = tensor_bytes(torch.empty(8192, 8192, dtype=torch.bfloat16)) - n = _compute_group_count(input_bytes, l2_bytes=128 * 1024 * 1024) - assert n == 4 +def test_bench_never_silently_falls_back_from_proton(monkeypatch): + def unavailable(_fn, warmup=25, rep=100): + del warmup, rep + raise RuntimeError("Proton profiler session could not be created") -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_legacy_callable_api(): - """bench still accepts the existing single-callable API used by TIRx tests.""" - M, N = 128, 128 + monkeypatch.setattr(bench_module, "_do_bench_proton", unavailable) - def run_and_check(): - A = torch.randn(M, N, device="cuda", dtype=torch.float16) - B = torch.randn(M, N, device="cuda", dtype=torch.float16) - result = bench( - lambda: torch.mm(A, B), warmup=1, repeat=2, proton_name="legacy", flush_l2_size=1 - ) - assert result > 0 + with pytest.raises(RuntimeError, match="Proton profiler session"): + bench({"noop": lambda: None}, timer="proton", cooldown_s=0) - tvm.testing.run_with_gpu_lock(run_and_check) + +@pytest.mark.parametrize( + ("timer", "alternative"), + [("proton", "event"), ("cudagraph_proton", "event")], +) +def test_missing_proton_session_is_an_explicit_error(monkeypatch, timer, alternative): + monkeypatch.setattr(bench_module.proton, "start", lambda *_args, **_kwargs: None) + + with pytest.raises(RuntimeError, match=rf"{timer}.*timer='{alternative}'"): + bench_module._start_proton_session("profile", timer=timer, explicit_alternative=alternative) @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_callable_inputs(): - """bench accepts a factory callable and auto-computes groups.""" - M, N = 256, 256 +def test_bench_references_pure_launch(): + """New bench(): reference builders return no-arg callables and get timed.""" + M, N = 128, 128 + A = torch.randn(M, N, device="cuda", dtype=torch.float16) + B = torch.randn(M, N, device="cuda", dtype=torch.float16) - call_count = [0] + funcs = {"tir": lambda: torch.mm(A, B)} - def make_input(): - call_count[0] += 1 - case = ( - torch.randn(M, N, device="cuda", dtype=torch.float16), - torch.randn(M, N, device="cuda", dtype=torch.float16), - ) - return case, tensor_bytes(*case) + def _addmm(): + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + return lambda: torch.addmm(C, A, B) - funcs = {"mm": lambda case: torch.mm(case[0], case[1])} + results = bench(funcs, warmup=5, repeat=10, timer="event", references={"addmm": _addmm}) + assert set(results["impls"].keys()) == {"tir", "addmm"} + assert all(v > 0 for v in results["impls"].values()) - def run_and_check(): - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") - assert "mm" in results["impls"] - assert results["impls"]["mm"] > 0 - assert call_count[0] >= 2 # at least 2 groups created - tvm.testing.run_with_gpu_lock(run_and_check) +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_bench_rejects_unknown_timer(): + """Unknown timer names fail instead of changing measurement method.""" + A = torch.randn(8, 8, device="cuda", dtype=torch.float16) + with pytest.raises(ValueError): + bench({"mm": lambda: torch.mm(A, A)}, timer="unknown") if __name__ == "__main__": diff --git a/tests/python/tirx/test_hint.py b/tests/python/tirx/test_hint.py index 88daad9b188a..1b3ca5648784 100644 --- a/tests/python/tirx/test_hint.py +++ b/tests/python/tirx/test_hint.py @@ -162,7 +162,7 @@ def func(A_ptr: T.handle) -> None: def test_hint_keyword_arg_on_tx_op(): """Tx.op(..., hint="msg") stores hint in TilePrimitiveCall.config.""" from tvm.tirx.buffer import decl_buffer - from tvm.tirx.stmt import TilePrimitiveCall + from tvm.tirx.tile_primitive import TilePrimitiveCall A = decl_buffer((64, 64), "float32", scope="global") A_sm = decl_buffer((64, 64), "float32", scope="shared") diff --git a/tests/python/tirx/test_jit.py b/tests/python/tirx/test_jit.py index 393d640b237b..d3cd32d13480 100644 --- a/tests/python/tirx/test_jit.py +++ b/tests/python/tirx/test_jit.py @@ -15,15 +15,18 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: F821 -"""Tests for ``@T.jit`` + ``T.constexpr``.""" +"""Tests for ``@T.jit`` compile-time specialization.""" from __future__ import annotations +import typing + import pytest import tvm from tvm.ir import assert_structural_equal from tvm.script import tirx as T +from tvm.script.tirx import tile as Tx def test_int_constexpr_specializes_loop_bound(): @@ -221,5 +224,230 @@ def k( assert len(spec.params) == 1 +def test_constexpr_specializes_nested_selector_condition(): + @T.jit(private=True) + def k( + A: T.Buffer((8,), "float16"), + B: T.Buffer((8,), "float16"), + C: T.Buffer((8,), "float16"), + flag: T.int32, + *, + LIMIT: T.constexpr, + ): + Tx.copy_async( + C[:], + A[:], + dispatch="tma_explicit", + mbar=C.data, + src_selector=[(flag < LIMIT, B)], + ) + + specialized = k.specialize(LIMIT=4) + op_call = specialized.body + condition, candidate = op_call.config["src_selector"][0] + assert isinstance(condition, tvm.tirx.LT) + assert int(condition.b) == 4 + assert candidate.same_as(specialized.buffer_map[specialized.params[1]]) + + +def test_optional_param_present_and_absent_ir(): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), out_h: T.handle): + out = T.match_buffer(out_h, (1,), "int32") + if a is not None: + A = T.match_buffer(a, (1,), "int32") + out[0] = A[0] + else: + out[0] = -1 + + @T.prim_func(private=True) + def expected_present(a: T.handle, out_h: T.handle): + A = T.match_buffer(a, (1,), "int32") + out = T.match_buffer(out_h, (1,), "int32") + out[0] = A[0] + + @T.prim_func(private=True) + def expected_absent(out_h: T.handle): + out = T.match_buffer(out_h, (1,), "int32") + out[0] = -1 + + present = kernel.specialize() + absent = kernel.specialize(a=None) + assert_structural_equal(present, expected_present, map_free_vars=True) + assert_structural_equal(absent, expected_absent, map_free_vars=True) + assert [param.name for param in present.params] == ["a", "out_h"] + assert [param.name for param in absent.params] == ["out_h"] + assert {param.name for param in present.buffer_map} == {"a", "out_h"} + assert {param.name for param in absent.buffer_map} == {"out_h"} + + +def test_optional_specialization_cache_includes_presence(): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), out_h: T.handle): + out = T.match_buffer(out_h, (1,), "int32") + if a is not None: + A = T.match_buffer(a, (1,), "int32") + out[0] = A[0] + else: + out[0] = 0 + + present = kernel.specialize() + absent = kernel.specialize(a=None) + assert present is kernel.specialize() + assert absent is kernel.specialize(a=None) + assert present is not absent + + +def test_multiple_optional_params_preserve_runtime_order(): + @T.jit(private=True) + def kernel( + first_h: T.handle, + a: T.Optional(T.handle), + scale: T.int32, + b: T.Optional(T.handle), + out_h: T.handle, + ): + first = T.match_buffer(first_h, (1,), "int32") + out = T.match_buffer(out_h, (1,), "int32") + out[0] = first[0] * scale + if a is not None: + A = T.match_buffer(a, (1,), "int32") + out[0] = out[0] + A[0] + if b is not None: + B = T.match_buffer(b, (1,), "int32") + out[0] = out[0] + B[0] + + assert [param.name for param in kernel.specialize().params] == [ + "first_h", + "a", + "scale", + "b", + "out_h", + ] + assert [param.name for param in kernel.specialize(a=None).params] == [ + "first_h", + "scale", + "b", + "out_h", + ] + assert [param.name for param in kernel.specialize(a=None, b=None).params] == [ + "first_h", + "scale", + "out_h", + ] + + +def test_optional_only_accepts_none_at_specialization_time(): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), out_h: T.handle): + if a is not None: + T.match_buffer(a, (1,), "int32") + T.match_buffer(out_h, (1,), "int32") + + with pytest.raises(TypeError, match="only accept None"): + kernel.specialize(a=object()) + + +def test_only_explicit_t_optional_is_specializable(): + @T.jit(private=True) + def typing_optional(a: typing.Optional[T.handle], out_h: T.handle): # noqa: UP045 + T.evaluate(0) + + @T.jit(private=True) + def union_optional(a: T.handle | None, out_h: T.handle): + T.evaluate(0) + + with pytest.raises(TypeError, match="unexpected"): + typing_optional.specialize(a=None) + with pytest.raises(TypeError, match="unexpected"): + union_optional.specialize(a=None) + + +def test_t_optional_is_restricted_to_jit(): + with pytest.raises(tvm.error.DiagnosticError, match="only supported by @T.jit"): + + @T.prim_func(private=True) + def invalid(a: T.Optional(T.handle)): + T.evaluate(0) + + +def test_compile_time_if_binding_uses_python_scope(): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), out_h: T.handle): + if a is None: + selected = T.match_buffer(out_h, (1,), "int32") + else: + selected = T.match_buffer(a, (1,), "int32") + selected[0] = 1 + + present = kernel.specialize() + absent = kernel.specialize(a=None) + assert len(present.params) == 2 + assert len(absent.params) == 1 + assert len(present.buffer_map) == 1 + assert len(absent.buffer_map) == 1 + + +def test_compile_time_bool_ops_and_if_expression_short_circuit(): + def fail_if_evaluated(): + raise RuntimeError("dead expression was evaluated") + + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), out_h: T.handle): + out = T.match_buffer(out_h, (1,), "int32") + if a is None or fail_if_evaluated(): + out[0] = 1 + if a is not None and fail_if_evaluated(): + out[0] = 2 + out[0] = 3 if a is None else fail_if_evaluated() + + absent = kernel.specialize(a=None) + assert [param.name for param in absent.params] == ["out_h"] + + +def test_runtime_tir_if_cannot_guard_absent_optional_param(): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle), flag: T.int32): + if flag != 0: + T.match_buffer(a, (1,), "int32") + + with pytest.raises(tvm.error.DiagnosticError, match="match_buffer"): + kernel.specialize(a=None) + + +@pytest.mark.parametrize( + ("operation", "source_text"), + [ + ("subscript", "a[10]"), + ("attribute", "a.ptr_to"), + ("match_buffer", "T.match_buffer"), + ], +) +def test_unguarded_absent_optional_param_reports_source(operation, source_text): + @T.jit(private=True) + def kernel(a: T.Optional(T.handle)): + if operation == "subscript": + a[10] + elif operation == "attribute": + a.ptr_to([0]) + else: + T.match_buffer(a, (1,), "int32") + + with pytest.raises(tvm.error.DiagnosticError) as exc_info: + kernel.specialize(a=None) + assert source_text in str(exc_info.value) + + +def test_present_optional_param_still_rejects_ffi_none(): + @T.jit + def kernel(a: T.Optional(T.handle)): + A = T.match_buffer(a, (1,), "int32") + A[0] = 0 + + executable = tvm.compile(kernel.specialize(), target="llvm", tir_pipeline="tirx") + with pytest.raises(TypeError, match="expected Tensor"): + executable(None) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py index 998eab97d0c1..1ad21c9b9d88 100644 --- a/tests/python/tirx/test_layout.py +++ b/tests/python/tirx/test_layout.py @@ -32,7 +32,6 @@ from tvm.tirx.cuda.operator.tile_primitive.tma_utils import ( SwizzleMode, mma_shared_layout, - tma_shared_layout, ) from tvm.tirx.layout import ( Axis, @@ -42,11 +41,13 @@ P, R, S, - SwizzleLayout, + TCol, TileLayout, + TLane, laneid, m, tid_in_wg, + tmem_mma_operand_layout, tx, warpid, wg_local_layout, @@ -758,15 +759,10 @@ def case_dims_mismatch(): def case_tile_compose_layout(): # tile(TileLayout, ComposeLayout) - compose = ComposeLayout( - layout_A=SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - layout_B=TileLayout(S[(8, 64) : (64, 1)]), - ) + compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) layout = TileLayout(S[(8, 1) : (1, 1)]) layout_tile = compose.tile(layout, (8, 1), (8, 64)) - layout_expected = ComposeLayout( - SwizzleLayout(3, 3, 3, swizzle_inner=True), TileLayout(S[4096:1]) - ) + layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[4096:1])) assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) outer_res = compose.is_tile_inner(layout_tile, (4096,), (512,)) @@ -784,12 +780,10 @@ def case_tile_compose_layout(): def case_tile_swizzle_layout(): # swizzle_128B_atom - swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) layout = TileLayout(S[(8, 4) : (1, 8)]) layout_tile = swizzle.tile(layout, (8, 4), (8, 64)) - layout_expected = ComposeLayout( - SwizzleLayout(3, 3, 3, swizzle_inner=True), TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) - ) + layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(64, 4, 64) : (64, 4096, 1)])) assert_structural_equal(layout_tile.canonicalize(), layout_expected) outer_res = swizzle.is_tile_inner(layout_tile, (64, 256), (8, 64)) @@ -804,11 +798,15 @@ def case_tile_swizzle_layout(): def case_tile_swizzle_layout2(): # swizzle_128B_atom - swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) tile = TileLayout(S[(3, 8, 4) : (8 * 4, 1, 8)]) layout_tile = swizzle.tile(tile, (3, 8, 4), (1, 8, 64)) layout_expected = ComposeLayout( - swizzle, TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]) + swizzle.per_element, + swizzle.swizzle_len, + swizzle.atom_len, + TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]), + swizzle.swizzle_inner, ) assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) @@ -824,10 +822,16 @@ def case_tile_swizzle_layout2(): def case_tile_swizzle_layout3(): # swizzle_64B_atom - swizzle = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) + swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) tile = TileLayout(S[(8, 8) : (1, 8)]) layout_tile = swizzle.tile(tile, (8, 8), (8, 32)) - layout_expected = ComposeLayout(swizzle, TileLayout(S[(64, 8, 32) : (32, 2048, 1)])) + layout_expected = ComposeLayout( + swizzle.per_element, + swizzle.swizzle_len, + swizzle.atom_len, + TileLayout(S[(64, 8, 32) : (32, 2048, 1)]), + swizzle.swizzle_inner, + ) assert_structural_equal(layout_tile.canonicalize(), layout_expected.canonicalize()) outer_res = swizzle.is_tile_inner(layout_tile, (64, 256), (8, 32)) @@ -842,7 +846,7 @@ def case_tile_swizzle_layout3(): def case_tile_swizzle_layout4(): # swizzle_64B_atom - swizzle = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) + swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) outer = swizzle.is_tile_inner(swizzle, (64, 256), (8, 32)) assert outer is None @@ -855,7 +859,7 @@ def case_tile_swizzle_layout4(): def case_tile_swizzle_layout5(): # swizzle_128B_atom - swizzle = SwizzleLayout(per_element=3, swizzle_len=2, atom_len=3) + swizzle = ComposeLayout(3, 2, 3, TileLayout(S[(256,)])) tile1 = TileLayout(S[(8, 8) : (1, 8)]) tile2 = TileLayout(S[(2, 2) : (1, 2)]) layout_tile = swizzle.tile(tile1, (8, 8), (8, 32)) @@ -988,18 +992,15 @@ def tile_layout_size(): tile_layout_size() def swizzle_layout_size(): - layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert layout.size() == 512 - layout = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) + layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) assert layout.size() == 1024 swizzle_layout_size() def compose_layout_size(): - layout = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) assert layout.size() == 512 compose_layout_size() @@ -1015,18 +1016,15 @@ def tile_layout_span(): tile_layout_span() def swizzle_layout_span(): - layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert layout.span() == 512 - layout = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) + layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) assert layout.span() == 1024 swizzle_layout_span() def compose_layout_span(): - layout = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) assert layout.span() == 512 compose_layout_span() @@ -1144,7 +1142,7 @@ def test_tile_layout_4(): ################ Swizzle Layout def test_swizzle_layout_0(): - layout = SwizzleLayout(per_element=0, swizzle_len=3, atom_len=3) + layout = ComposeLayout(0, 3, 3, TileLayout(S[(64,)])) # assert layout.size == 64 for i, j in itertools.product(range(8), range(8)): assert layout.apply(i * 8 + j)["m"] == i * 8 + i ^ j @@ -1152,7 +1150,7 @@ def test_swizzle_layout_0(): test_swizzle_layout_0() def test_swizzle_layout_1(): - layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert layout.size() == 512 for i, j, k in itertools.product(range(8), range(8), range(8)): assert layout.apply((i * 8 + j) * 8 + k)["m"] == (i * 8 + (i ^ j)) * 8 + k @@ -1166,7 +1164,7 @@ def test_swizzle_layout_1(): test_swizzle_layout_1() def test_swizzle_layout_2(): - layout = SwizzleLayout(per_element=0, swizzle_len=3, atom_len=3, swizzle_inner=False) + layout = ComposeLayout(0, 3, 3, TileLayout(S[(64,)]), swizzle_inner=False) assert layout.size() == 64 for i, j in itertools.product(range(8), range(8)): assert layout.apply(i * 8 + j)["m"] == (i ^ j) * 8 + j @@ -1174,7 +1172,7 @@ def test_swizzle_layout_2(): test_swizzle_layout_2() def test_swizzle_layout_3(): - layout = SwizzleLayout(per_element=0, swizzle_len=2, atom_len=3) + layout = ComposeLayout(0, 2, 3, TileLayout(S[(32,)])) for i, j in itertools.product(range(8), range(8)): _outer_i, inner_i = i // 4, i % 4 outer_j, inner_j = j // 4, j % 4 @@ -1184,9 +1182,15 @@ def test_swizzle_layout_3(): ################ Compose Layout def test_compose_layout_0(): - layoutA = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) layoutB = TileLayout(S[(8, 64) : (64, 1)]) - layout = ComposeLayout(layoutA, layoutB) + layout = ComposeLayout( + layoutA.per_element, + layoutA.swizzle_len, + layoutA.atom_len, + layoutB, + layoutA.swizzle_inner, + ) assert layout.size() == 512 assert layout.span() == 512 for i, j in itertools.product(range(8), range(64)): @@ -1197,9 +1201,15 @@ def test_compose_layout_0(): test_compose_layout_0() def test_compose_layout_1(): - layoutA = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) layoutB = TileLayout(S[(16, 64, 8) : (64, 1, 1024)]) - layout = ComposeLayout(layoutA, layoutB) + layout = ComposeLayout( + layoutA.per_element, + layoutA.swizzle_len, + layoutA.atom_len, + layoutB, + layoutA.swizzle_inner, + ) assert layout.size() == 16 * 64 * 8 assert layout.span() == 16 * 64 * 8 for i, j, k in itertools.product(range(16), range(64), range(8)): @@ -1254,17 +1264,29 @@ def test_trainium_psum_layout_0(): def test_normalize_compose_layout(): def case1(): - layoutA = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) layoutB = TileLayout(S[(8, 64) : (64, 1)]) - layout = ComposeLayout(layoutA, layoutB.canonicalize()) + layout = ComposeLayout( + layoutA.per_element, + layoutA.swizzle_len, + layoutA.atom_len, + layoutB.canonicalize(), + layoutA.swizzle_inner, + ) assert_structural_equal(layout.canonicalize(), layoutA) case1() def case2(): - layoutA = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layoutA = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) layoutB = TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) - layout = ComposeLayout(layoutA, layoutB.canonicalize()) + layout = ComposeLayout( + layoutA.per_element, + layoutA.swizzle_len, + layoutA.atom_len, + layoutB.canonicalize(), + layoutA.swizzle_inner, + ) assert_structural_equal(layout.canonicalize(), layout) case2() @@ -1341,6 +1363,63 @@ def case1(): case1() +def test_group_many_uses_minimal_common_refinement(): + layout = TileLayout(S[30:1]) + grouped, separators = layout.group_many(([6, 5], [2, 15])) + + expected = TileLayout(S[(2, 3, 5) : (15, 5, 1)]) + assert_structural_equal(grouped, expected) + assert [list(seps) for seps in separators] == [[0, 2, 3], [0, 1, 3]] + + +def test_group_many_preserves_repeated_boundaries_as_unit_iters(): + layout = TileLayout(S[30:1]) + grouped, separators = layout.group_many( + ( + [1, 6, 1, 5, 1], + [2, 1, 3, 5], + ) + ) + + expected = TileLayout(S[(1, 2, 1, 3, 1, 5, 1) : (30, 15, 15, 5, 5, 1, 1)]) + assert_structural_equal(grouped, expected) + assert [list(seps) for seps in separators] == [ + [0, 1, 4, 5, 6, 7], + [0, 2, 3, 4, 7], + ] + + +def test_group_many_preserves_existing_unit_iter_stride(): + layout = TileLayout(S[(6, 1, 5) : (5, 97, 1)]) + grouped, separators = layout.group_many(([6, 1, 5], [30])) + + assert_structural_equal(grouped, layout) + assert [list(seps) for seps in separators] == [[0, 1, 2, 3], [0, 3]] + + +def test_group_many_rejects_incompatible_boundaries(): + layout = TileLayout(S[12:1]) + with pytest.raises(tvm.error.InternalError, match="incompatible cumulative boundaries"): + layout.group_many(([3, 4], [2, 6])) + + +def test_group_many_accepts_only_provably_ordered_symbolic_boundaries(): + n = Var("n", "int32") + extent = tvm.tirx.floormod(n, 16) + 2 + layout = TileLayout(S[2 * extent : 1]) + + grouped, separators = layout.group_many(([2, extent], [1, 2, extent])) + expected = TileLayout(S[(1, 2, extent) : (Analyzer().simplify(2 * extent), extent, 1)]) + assert_structural_equal(grouped, expected) + assert [list(seps) for seps in separators] == [[0, 2, 3], [0, 1, 2, 3]] + + with pytest.raises( + tvm.error.InternalError, + match="cannot prove cumulative boundary divisibility|order or equality", + ): + layout.group_many(([2, extent], [extent, 2])) + + def test_permute_by_groups(): def case_swap_two_groups(): # Two groups, each with 2 shard iters: swap them. @@ -1385,9 +1464,7 @@ def case1(): def test_mma_shared_layout(): def case1(): layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (64, 256)) - layout_expected = ComposeLayout( - SwizzleLayout(3, 3, 3, swizzle_inner=True), TileLayout(S[(64, 4, 64) : (64, 4096, 1)]) - ) + layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(64, 4, 64) : (64, 4096, 1)])) assert_structural_equal(layout, layout_expected) case1() @@ -1395,8 +1472,7 @@ def case1(): def case2(): layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, (3, 64, 256)) layout_expected = ComposeLayout( - SwizzleLayout(3, 3, 3, swizzle_inner=True), - TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]), + 3, 3, 3, TileLayout(S[(3, 64, 4, 64) : (16384, 64, 4096, 1)]) ) assert_structural_equal(layout, layout_expected) @@ -1405,27 +1481,19 @@ def case2(): def case3(): layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_64B_ATOM, (3, 64, 256)) layout_expected = ComposeLayout( - SwizzleLayout(3, 2, 3, swizzle_inner=True), - TileLayout(S[(3, 64, 8, 32) : (16384, 32, 2048, 1)]), + 3, 2, 3, TileLayout(S[(3, 64, 8, 32) : (16384, 32, 2048, 1)]) ) assert_structural_equal(layout, layout_expected) case3() -def test_tma_shared_layout_alias(): - shape = (3, 64, 256) - layout = mma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, shape) - alias_layout = tma_shared_layout("float16", SwizzleMode.SWIZZLE_128B_ATOM, shape) - assert_structural_equal(alias_layout, layout) - - def test_pool_allocator_alloc_mma(): def alloc_layout(shape, dtype, swizzle_mode="auto"): with IRBuilder(): with Tx_builder.prim_func(): pool = T.SMEMPool(Var("smem_ptr", PointerType(PrimType("uint8")))) - buf = pool.alloc_mma(shape, dtype, swizzle_mode=swizzle_mode) + buf = pool.alloc_tcgen05_mma_AB(shape, dtype, swizzle_mode=swizzle_mode) return buf.layout cases = [ @@ -1450,6 +1518,37 @@ def alloc_layout(shape, dtype, swizzle_mode="auto"): assert_structural_equal(layout_none, expected_none) +def test_tmem_mma_operand_layout_grouped_d(): + head64_o = tmem_mma_operand_layout( + "D", (64, 512), "float32", M=64, cta_group=1, ws=True, group=(2, 2, 128) + ) + expected_head64_o = TileLayout( + S[(64, 2, 2, 128) : (1 @ TLane, 128 @ TCol, 64 @ TLane, 1 @ TCol)] + ).canonicalize() + assert_structural_equal(head64_o, expected_head64_o) + + head128_o = tmem_mma_operand_layout( + "D", (64, 512), "float32", M=128, cta_group=2, group=(2, 2, 128) + ) + assert_structural_equal(head128_o, expected_head64_o) + + with pytest.raises(ValueError, match="ws=True only valid"): + tmem_mma_operand_layout("D", (128, 256), "float32", M=128, cta_group=1, ws=True) + + with pytest.raises(ValueError, match="group must be"): + tmem_mma_operand_layout( + "D", (64, 512), "float32", M=64, cta_group=1, ws=True, group=(4, 128) + ) + + with pytest.raises(ValueError, match="identity needs 2D"): + tmem_mma_operand_layout("A", (1, 128, 64), "bfloat16", M=128, cta_group=1) + + # M=64 non-.ws A occupies lanes 0..63 identically (Layout F datapath is the + # accumulator's scatter, not the A operand's). + a_f = tmem_mma_operand_layout("A", (64, 64), "float32", M=64, cta_group=1) + assert_structural_equal(a_f, TileLayout(S[(64, 64) : (1 @ TLane, 1 @ TCol)]).canonicalize()) + + def test_storage(): def case1(): layout = TileLayout(S[(8, 8) : (8, 1)]) @@ -1465,7 +1564,7 @@ def case2(): case2() def case3(): - layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert_structural_equal(layout.storage(), layout) case3() @@ -1496,21 +1595,15 @@ def case1(): case1() def case2(): - layout = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) - layout_expected = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) + layout_expected = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) case2() def case3(): - layout = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) - layout_expected = ComposeLayout( - SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 128) : (128, 1)]), - ) + layout = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) + layout_expected = ComposeLayout(4, 3, 3, TileLayout(S[(8, 128) : (128, 1)])) assert_structural_equal(layout.unpack(2).canonicalize(), layout_expected.canonicalize()) case3() @@ -1525,21 +1618,15 @@ def case1(): case1() def case2(): - layout = SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3) - layout_expected = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + layout = ComposeLayout(4, 3, 3, TileLayout(S[(1024,)])) + layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) case2() def case3(): - layout = ComposeLayout( - SwizzleLayout(per_element=4, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 128) : (128, 1)]), - ) - layout_expected = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) + layout = ComposeLayout(4, 3, 3, TileLayout(S[(8, 128) : (128, 1)])) + layout_expected = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) assert_structural_equal(layout.pack(2).canonicalize(), layout_expected.canonicalize()) case3() @@ -1615,9 +1702,20 @@ def case4(): case4() + def case_adds_region_offset_to_base(): + warp_offset = Var("warp_offset", "int32") + layout = TileLayout(S[(4, 64)] + warp_offset * 64) + shape = [4, 64] + region = [(2, 4), (0, 64)] + sliced = layout.slice(shape, region).canonicalize() + assert sliced is not None + assert Analyzer().can_prove_equal(sliced.offset[m], warp_offset * 64 + 128) + + case_adds_region_offset_to_base() + def case_swizzle_slice(): - # SwizzleLayout slice - delegates to ComposeLayout - swizzle = SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + # bare-swizzle ComposeLayout slice - delegates to ComposeLayout + swizzle = ComposeLayout(3, 3, 3, TileLayout(S[(512,)])) shape = [512] region = [(64, 128)] sliced = swizzle.slice(shape, region) @@ -1628,10 +1726,7 @@ def case_swizzle_slice(): def case_compose_slice(): # ComposeLayout slice - compose = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) + compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) shape = [512] region = [(64, 128)] sliced = compose.slice(shape, region) @@ -1642,10 +1737,7 @@ def case_compose_slice(): def case_compose_slice_2d(): # ComposeLayout slice with 2D shape - compose = ComposeLayout( - SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - TileLayout(S[(8, 64) : (64, 1)]), - ) + compose = ComposeLayout(3, 3, 3, TileLayout(S[(8, 64) : (64, 1)])) shape = [8, 64] region = [(2, 4), (0, 64)] sliced = compose.slice(shape, region) @@ -1661,7 +1753,7 @@ def test_cuda_copy_extract_swizzle_tile_simplifies_constant_region_extents(): zero = tvm.tirx.Mul(tvm.tirx.IntImm("int32", 0), tvm.tirx.IntImm("int32", 64)) end = tvm.tirx.Add(zero, tvm.tirx.IntImm("int32", 64)) - tile = _extract_tile(SwizzleLayout(3, 3, 3), [(zero, end)]) + tile = _extract_tile(ComposeLayout(3, 3, 3, TileLayout(S[(512,)])), [(zero, end)]) assert [int(it.extent) for it in tile.shard] == [64] diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py index efd245ddbf67..8d6e6326e189 100644 --- a/tests/python/tirx/test_op.py +++ b/tests/python/tirx/test_op.py @@ -21,7 +21,7 @@ from tvm.ir import Op, assert_structural_equal from tvm.tirx.buffer import decl_buffer from tvm.tirx.exec_scope import ExecScope -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def _test(op: str, *args): diff --git a/tests/python/tirx/test_op_namespace_cleanup.py b/tests/python/tirx/test_op_namespace_cleanup.py index 26871df5211c..6cdf447c7420 100644 --- a/tests/python/tirx/test_op_namespace_cleanup.py +++ b/tests/python/tirx/test_op_namespace_cleanup.py @@ -26,7 +26,7 @@ from tvm.ir import Op, assert_structural_equal from tvm.script import tirx as T from tvm.script.tirx import tile as Tx -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def _tile_calls(func): diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 12f29c0d4c83..d99f70924747 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -14,6 +14,8 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import math + import pytest import tvm @@ -23,7 +25,7 @@ from tvm.script import ir as I from tvm.script import tirx as T from tvm.script.tirx import tile as Tx -from tvm.tirx.layout import laneid, warpid +from tvm.tirx.layout import TCol, TLane, laneid, warpid def from_source(code): @@ -140,13 +142,10 @@ def get_layout3(): return T.TileLayout(T.S[(8, 16, 8, 16) : (1024, 16, 128, 1)]) def get_layout4(): - return T.SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3) + return T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(512,)])) def get_layout5(): - return T.ComposeLayout( - T.SwizzleLayout(per_element=3, swizzle_len=3, atom_len=3), - T.TileLayout(T.S[(64, 64, 4) : (64, 1, 64 * 64)]), - ) + return T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(64, 64, 4) : (64, 1, 64 * 64)])) # fmt: off @T.prim_func @@ -1024,7 +1023,7 @@ def test_kwargs_op_call(): @T.prim_func(private=True) def test(A: T.Buffer((10, 10), "float32"), B: T.Buffer((10, 10), "float32")): T.device_entry() - kwargs = T.meta_var({"dispatch": "tma", "cta_group": 2}) + kwargs = T.meta_var({"dispatch": "tma_auto", "cta_group": 2}) Tx.copy_async(A[:, :], B[:, :], **kwargs) # fmt: on code = test.script() @@ -1472,6 +1471,459 @@ def func() -> None: assert from_source(code).script() == code +def test_buffer_rearrange_allows_arbitrary_axis_names(): + @T.prim_func + def ordinary_axis() -> None: + T.device_entry() + A = T.alloc_buffer( + (8, 4), + "float16", + scope="local", + layout=T.TileLayout(T.S[(8, 4) : (4, 1)]), + ) + B = A.rearrange("(outer inner) tail -> outer tail inner", outer=2) + B[0, 0, 0] = T.float16(0) + + @T.prim_func + def buf_axis() -> None: + T.device_entry() + A = T.alloc_buffer( + (8, 4), + "float16", + scope="local", + layout=T.TileLayout(T.S[(8, 4) : (4, 1)]), + ) + B = A.rearrange("(buf inner) tail -> buf tail inner", buf=2) + B[0, 0, 0] = T.float16(0) + + @T.prim_func + def self_axis() -> None: + T.device_entry() + A = T.alloc_buffer( + (8, 4), + "float16", + scope="local", + layout=T.TileLayout(T.S[(8, 4) : (4, 1)]), + ) + B = A.rearrange("(self inner) tail -> self tail inner", self=2) + B[0, 0, 0] = T.float16(0) + + @T.prim_func + def pattern_axis() -> None: + T.device_entry() + A = T.alloc_buffer( + (8, 4), + "float16", + scope="local", + layout=T.TileLayout(T.S[(8, 4) : (4, 1)]), + ) + B = A.rearrange("(pattern inner) tail -> pattern tail inner", pattern=2) + B[0, 0, 0] = T.float16(0) + + @T.prim_func + def keyword_pattern() -> None: + T.device_entry() + A = T.alloc_buffer( + (8, 4), + "float16", + scope="local", + layout=T.TileLayout(T.S[(8, 4) : (4, 1)]), + ) + B = A.rearrange(pattern="(outer inner) tail -> outer tail inner", outer=2) + B[0, 0, 0] = T.float16(0) + + expected = _collect_buffers(ordinary_axis)["B"] + for func in (buf_axis, self_axis, pattern_axis, keyword_pattern): + actual = _collect_buffers(func)["B"] + assert_structural_equal(actual.shape, expected.shape) + assert_structural_equal(actual.layout, expected.layout) + + +def test_buffer_permute_compose_layout_ir(): + """Verify .permute on a swizzle-composed layout: the swizzle is preserved + and the inner tile layout's dim groups are permuted (the reshape-permute- + reshape idiom used to refactor gather views without restating strides).""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer( + [4, 4, 4, 64], dtype="bfloat16", scope="shared.dyn", + layout=T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(4, 4, 4, 64) : (1024, 256, 64, 1)])), + ) + B = A.permute(1, 0, 2, 3) + B[0, 0, 0, 0] = T.bfloat16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf = bufs["A"] + b_buf = bufs["B"] + + assert b_buf.data.same_as(a_buf.data) + assert [int(s) for s in b_buf.shape] == [4, 4, 4, 64] + expected = tvm.tirx.layout.ComposeLayout( + a_buf.layout.per_element, + a_buf.layout.swizzle_len, + a_buf.layout.atom_len, + a_buf.layout.tile_layout.permute_dims([1, 0, 2, 3]), + a_buf.layout.swizzle_inner, + ) + assert_structural_equal(b_buf.layout, expected) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_sub_multi_iter_dim_ir(): + """sub with an int index on a dim carried by several layout iters + decomposes the index mixed-radix across the iters' strides.""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([8, 16], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(2, 4, 16) : (1024, 64, 1)])) + B = A.sub[5] + B[0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf, b_buf = bufs["A"], bufs["B"] + # 5 -> (5 // 4, 5 % 4) = (1, 1) -> 1 * 1024 + 1 * 64 + assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 1088 + assert [int(s) for s in b_buf.shape] == [16] + assert_structural_equal(b_buf.layout, tvm.tirx.layout.TileLayout(T.S[(16,) : (1,)])) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_sub_multi_iter_misaligned_rejected(): + buf = tvm.tirx.decl_buffer( + (8, 16), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(2, 4, 16) : (1024, 64, 1)]) + ) + # sub[2:6] narrows the multi-iter dim 0 at a misaligned offset. + with pytest.raises(ValueError, match="multiples of the inner iter block"): + buf.sub[2:6] + + +def test_buffer_sub_ir(): + """buf.sub follows numpy basic indexing as a view constructor: int drops + the dim, a:b narrows, a::s strides. Offsets fold into elem_offset through + the dim's layout iter strides; the derived layout carries the survivors.""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([4, 8, 16], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(4, 8, 16) : (256, 16, 1)])) + B = A.sub[1, 2:6] + B[0, 0] = T.float16(0) + C = A.sub[:, 1::2] + C[0, 0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf, b_buf, c_buf = bufs["A"], bufs["B"], bufs["C"] + # sub[1, 2:6]: drop dim 0 at 1 (1 * 256) then narrow dim 1 to [2, 6) (2 * 16) + assert [int(s) for s in b_buf.shape] == [4, 16] + assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 288 + assert_structural_equal(b_buf.layout, tvm.tirx.layout.TileLayout(T.S[(4, 16) : (16, 1)])) + # sub[:, 1::2]: keep dim 0, split dim 1 into (4, 2) and fix the remainder at 1 + assert [int(s) for s in c_buf.shape] == [4, 4, 16] + assert int(tvm.arith.Analyzer().simplify(c_buf.elem_offset - a_buf.elem_offset)) == 16 + assert_structural_equal( + c_buf.layout, tvm.tirx.layout.TileLayout(T.S[(4, 4, 16) : (256, 32, 1)]) + ) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_view_surgery_static_bounds_rejected(): + """Statically-known out-of-range sub arguments must be rejected loudly + (review finding: OOB offsets were silent).""" + buf = tvm.tirx.decl_buffer( + (10,), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(10,) : (1,)]) + ) + grid = tvm.tirx.decl_buffer( + (4, 8), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(4, 8) : (8, 1)]) + ) + # int index: static bounds + with pytest.raises(ValueError, match="out of range"): + buf.sub[10] + with pytest.raises(ValueError, match="out of range"): + buf.sub[-1] + # slice narrow: static range bounds + with pytest.raises(ValueError, match="exceeds"): + buf.sub[8:12] + with pytest.raises(ValueError, match="must be non-negative"): + buf.sub[-2:2] + with pytest.raises(ValueError, match="must be positive"): + buf.sub[5:3] + # grid.sub: out-of-range int, exceeding narrow, stepped-start out of range + with pytest.raises(ValueError, match="out of range"): + grid.sub[10, :] + with pytest.raises(ValueError, match="exceeds"): + grid.sub[:, 4:12] + with pytest.raises(ValueError, match=r"in \[0, 2\)"): + grid.sub[:, -1::2] + + +def test_buffer_sub_swizzle_commutation(): + """A folded view offset moves into elem_offset only when it commutes + with the swizzle, i.e. is a multiple of the swizzle period + 2^(per_element + atom_len + swizzle_len). Sub-period offsets stay inside + the derived tile layout's offset so the swizzle keeps applying to them + (review finding: folding them outside produced wrong addresses). Both + placements must be address-equivalent to the parent layout.""" + + def addr(buf, base, *coords): + analyzer = tvm.arith.Analyzer() + if len(coords) == 1: + rel = buf.layout.apply(coords[0])["m"] + else: + rel = buf.layout.apply(*coords, shape=[int(s) for s in buf.shape])["m"] + return int(analyzer.simplify((buf.elem_offset - base) + rel)) + + analyzer = tvm.arith.Analyzer() + compose = T.ComposeLayout( + 3, 3, 3, T.TileLayout(T.S[(4, 1024) : (1024, 1)]) + ) # period = 2^(3+3+3) = 512 elements + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([4, 1024], dtype="bfloat16", scope="shared.dyn", layout=compose) + B = A.sub[1] # offset 1024 = 2 * period: folds into elem_offset + B[0] = T.bfloat16(0) + C = A.sub[:, 512:1024] # offset 512 = period: folds into elem_offset + C[0, 0] = T.bfloat16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf, b_buf, c_buf = bufs["A"], bufs["B"], bufs["C"] + base = a_buf.elem_offset + assert int(analyzer.simplify(b_buf.elem_offset - base)) == 1024 + for j in (0, 1, 63, 511, 1023): + assert addr(a_buf, base, 1024 + j) == addr(b_buf, base, j) + for j in (0, 1, 255, 511): + assert addr(a_buf, base, 512 + j) == addr(c_buf, base, 0, j) + + code = func.script() + assert from_source(code).script() == code + + # Sub-period offsets do not commute: they stay inside the tile layout's + # offset (elem_offset unchanged) and every address matches the parent. + compose2 = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(2, 16, 8) : (128, 8, 1)])) + + # fmt: off + @T.prim_func + def func2() -> None: + T.device_entry() + A = T.alloc_buffer([2, 16, 8], dtype="float16", scope="shared.dyn", layout=compose2) + B = A.sub[:, 1] # offset 8 + B[0, 0] = T.float16(0) + C = A.sub[:, :, 1] # offset 1 + C[0, 0] = T.float16(0) + D = A.sub[:, 1:3] # offset 8 + D[0, 0, 0] = T.float16(0) + E = A.sub[:, 1:3] # narrow via sub + E[0, 0, 0] = T.float16(0) + for w in T.serial(16): + F = A.sub[:, w] # dynamic sub-period offset + F[0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func2) + a2, base2 = bufs["A"], bufs["A"].elem_offset + shape2 = [2, 16, 8] + for name, to_parent in { + "B": lambda c: (c[0], 1, c[1]), + "C": lambda c: (c[0], c[1], 1), + "D": lambda c: (c[0], 1 + c[1], c[2]), + "E": lambda c: (c[0], 1 + c[1], c[2]), + }.items(): + child = bufs[name] + assert int(analyzer.simplify(child.elem_offset - base2)) == 0 + child_shape = [int(s) for s in child.shape] + for flat in range(math.prod(child_shape)): + coords, rem = [], flat + for extent in reversed(child_shape): + coords.append(rem % extent) + rem //= extent + coords = tuple(reversed(coords)) + assert addr(a2, base2, *to_parent(coords)) == addr(child, base2, *coords), ( + name, + coords, + ) + + code = func2.script() + assert from_source(code).script() == code + + # fixed-point windows (all touched addresses below 2^(per_element + + # atom_len)) are correct through the same layout-offset placement + compose3 = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(64,) : (1,)])) + + # fmt: off + @T.prim_func + def func3() -> None: + T.device_entry() + A = T.alloc_buffer([64], dtype="bfloat16", scope="shared.dyn", layout=compose3) + B = A.sub[8:16] + B[0] = T.bfloat16(0) + # fmt: on + + bufs = _collect_buffers(func3) + a3, b3 = bufs["A"], bufs["B"] + for j in range(8): + assert addr(a3, a3.elem_offset, 8 + j) == addr(b3, a3.elem_offset, j) == 8 + j + + +def test_buffer_tile_ir(): + """buf.tile((dim, factors))[picks] splits dims into factors and picks + chunks in one call: int/Expr picks a factor, ':' keeps it, kept + factors merge back. Equivalent to the view (reshape) + sub chain.""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([3, 64, 512], dtype="float16", scope="shared", + layout=T.TileLayout(T.S[(3, 64, 512) : (64 * 512, 512, 1)])) + for w in T.serial(4): + B = A.tile((1, (-1, 4, 4)))[:, w, :] + B[0, 0, 0] = T.float16(0) + C = A.view(3, 4, 4, 4, 512).sub[:, :, w].view(3, 16, 512) + C[0, 0, 0] = T.float16(0) + D = A.tile((1, (-1, 4)))[:, 2] + D[0, 0, 0] = T.float16(0) + E = A.sub[:, 2::4] + E[0, 0, 0] = T.float16(0) + F = A.tile((1, (4, -1)))[2, :] + F[0, 0, 0] = T.float16(0) + G = A.sub[:, 32:48] + G[0, 0, 0] = T.float16(0) + + @T.prim_func + def func_multi() -> None: + T.device_entry() + A = T.alloc_buffer([64, 128], dtype="float16", scope="shared", + layout=T.TileLayout(T.S[(64, 128) : (128, 1)])) + for wx in T.serial(4): + for wy in T.serial(2): + H = A.tile((0, (-1, 4, 2)), (1, (-1, 2, 8)))[:, wx, :, :, wy, :] + H[0, 0] = T.float16(0) + J = (A.view(8, 4, 2, 128).sub[:, wx].view(16, 128) + .view(16, 8, 2, 8).sub[:, :, wy].view(16, 64)) + J[0, 0] = T.float16(0) + + @T.prim_func + def func_multipick() -> None: + T.device_entry() + A = T.alloc_buffer([128, 16], dtype="float16", scope="shared", + layout=T.TileLayout(T.S[(128, 16) : (16, 1)])) + for a in T.serial(2): + for b in T.serial(4): + K = A.tile((0, (2, 4, -1)))[a, b, :] + K[0, 0] = T.float16(0) + L = A.view(2, 4, 16, 16).sub[a, b] + L[0, 0] = T.float16(0) + # fmt: on + + b = _collect_buffers(func) + assert [int(s) for s in b["B"].shape] == [3, 16, 512] + assert_structural_equal(b["B"].layout, b["C"].layout) + assert_structural_equal(b["D"].layout, b["E"].layout) + assert_structural_equal(b["F"].layout, b["G"].layout) + m = _collect_buffers(func_multi) + assert [int(s) for s in m["H"].shape] == [16, 64] + assert_structural_equal(m["H"].layout, m["J"].layout) + mp = _collect_buffers(func_multipick) + assert [int(s) for s in mp["K"].shape] == [16, 16] + assert_structural_equal(mp["K"].layout, mp["L"].layout) + + code = func_multipick.script() + assert from_source(code).script() == code + + +def test_buffer_tile_rejected(): + buf = tvm.tirx.decl_buffer( + (3, 64, 512), + "float16", + layout=tvm.tirx.layout.TileLayout(T.S[(3, 64, 512) : (64 * 512, 512, 1)]), + ) + with pytest.raises(ValueError, match="takes a dim and a factors"): + buf.tile(1, 4, 4) # positional single-dim form takes exactly (dim, factors) + with pytest.raises(ValueError, match="picks no factor"): + buf.tile(1, (-1, 4))[:, :] # a chunk must pick at least one factor + with pytest.raises(ValueError, match="non-empty tuple"): + buf.tile((1, 4)) # factors must be a tuple + with pytest.raises(ValueError, match="non-empty tuple"): + buf.tile((1, ())) + with pytest.raises(ValueError, match="tiled more than once"): + buf.tile((1, (4, -1)), (1, (2, -1))) + with pytest.raises(ValueError, match="index"): + buf.tile((1, (-1, 4)))[2] # 2 factors, 1 index + with pytest.raises(ValueError, match="must be ':'"): + buf.tile((1, (-1, 4)))[:, 1:3] # sub-slice on a factor + + +def test_buffer_chunk_ir(): + """buf.chunk(spec)[picks] narrows each chunked dim to its picked chunk's + contiguous [c*k : (c+1)*k) range (k = E // n), rank-preserving: a per-dim + tuple where None passes the pick straight through and n divides that dim + into n equal chunks. chunk(spec)[picks] is the exact same BufferRegion as + the hand-written a*k:(a+1)*k slice — no reshape, no extra dim.""" + + from tvm.tirx.stmt import BufferRegion + + compose = T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(4, 512) : (512, 1)])) + A = tvm.tirx.decl_buffer( + (4, 8, 16), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(4, 8, 16) : (128, 16, 1)]) + ) + C = tvm.tirx.decl_buffer((4, 512), "bfloat16", layout=compose) + + # chunk((None, None, 2))[:, :, 1] narrows dim 2 (extent 16) to chunk 1 of 2 + # → [8:16] (k = 16 // 2 = 8); rank preserved, dims 0/1 pass through as ':'. + reg = A.chunk((None, None, 2))[:, :, 1] + assert isinstance(reg, BufferRegion) + assert len(reg.region) == 3 # rank-preserving: no extra extent-1 chunk dim + assert (int(reg.region[2].min), int(reg.region[2].extent)) == (8, 8) + assert_structural_equal(reg, A[:, :, 8:16]) + + # a None dim passes an int pick straight through (int → extent-1 region), + # while the chunked dim still narrows to its picked chunk. + reg2 = A.chunk((None, None, 2))[3, :, 0] + assert_structural_equal(reg2, A[3, :, 0:8]) + + # chunk((None, 4))[:, 2] on the swizzle-carrying compose layout: dim 1 + # (extent 512) → chunk 2 of 4 → [256:384] (k = 128), byte-identical slice. + reg_c = C.chunk((None, 4))[:, 2] + assert (int(reg_c.region[1].min), int(reg_c.region[1].extent)) == (256, 128) + assert_structural_equal(reg_c, C[:, 256:384]) + + # a symbolic (Expr) chunk index translates to c*k : (c+1)*k as well. + c = T.Var(name="c", ty="int32") + assert_structural_equal(A.chunk((None, None, 2))[:, :, c], A[:, :, c * 8 : c * 8 + 8]) + + # validation + with pytest.raises(ValueError, match="per-dim tuple"): + A.chunk(2) # spec must be a per-dim tuple, not a bare int + with pytest.raises(ValueError, match="spec length"): + A.chunk((None, 2)) # length 2 != rank 3 + with pytest.raises(ValueError, match="None or a positive int"): + A.chunk((None, None, 0)) # 0 is not a positive chunk count + with pytest.raises(ValueError, match="chunk index, not a slice"): + A.chunk((None, None, 2))[:, :, 0:1] # a chunked dim takes a chunk index + with pytest.raises(ValueError, match="rank-3 spec"): + A.chunk((None, None, 2))[0, 0, 0, 0] # too many indices + + def test_buffer_view_dtype_ir(): """Verify .view('float32') on float16: dtype correct, last dim halved, shared data.""" @@ -1794,6 +2246,67 @@ def func(): assert_structural_equal(func, from_source(code)) +def test_buffer_sub_tmem_offset_uses_physical_columns(): + """A tmem layout measures TCol in elements, but allocated_addr measures + physical 32-bit columns. Folding a sub-view offset must scale by dtype + width exactly once (the FlashMLA Q-tail view is the bf16 regression).""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + Q = T.decl_buffer( + (2, 64, 288), "bfloat16", scope="tmem", allocated_addr=256, + layout=T.TileLayout(T.S[(2, 64, 288) : (64 @ TLane, 1 @ TLane, 1 @ TCol)]), + ) + Q_tail = Q.sub[:, :, 256:288] + F8 = T.decl_buffer( + (64, 128), "float8_e4m3fn", scope="tmem", allocated_addr=32, + layout=T.TileLayout(T.S[(64, 128) : (1 @ TLane, 1 @ TCol)]), + ) + F8_tail = F8.sub[:, 64:96] + F32 = T.decl_buffer( + (64, 128), "float32", scope="tmem", allocated_addr=64, + layout=T.TileLayout(T.S[(64, 128) : (1 @ TLane, 1 @ TCol)]), + ) + F32_tail = F32.sub[:, 32:64] + T.evaluate(Q_tail[0, 0, 0]) + T.evaluate(F8_tail[0, 0]) + T.evaluate(F32_tail[0, 0]) + # fmt: on + + bufs = _collect_buffers(func) + assert int(bufs["Q_tail"].allocated_addr[0]) == 384 # 256 + 256 * 16 / 32 + assert int(bufs["F8_tail"].allocated_addr[0]) == 48 # 32 + 64 * 8 / 32 + assert int(bufs["F32_tail"].allocated_addr[0]) == 96 # 64 + 32 * 32 / 32 + for name in ("Q_tail", "F8_tail", "F32_tail"): + assert int(bufs[name].layout.offset.get(TCol, 0)) == 0 + + code = func.script() + assert from_source(code).script() == code + assert_structural_equal(func, from_source(code)) + + +def test_buffer_sub_tmem_rejects_partial_column_offset(): + buf_layout = tvm.tirx.layout.TileLayout(T.S[(64, 16) : (1 @ TLane, 1 @ TCol)]) + + def build(): + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.decl_buffer( + (64, 16), "bfloat16", scope="tmem", allocated_addr=0, layout=buf_layout, + ) + _ = A.sub[:, 1:3] + # fmt: on + + return func + + with pytest.raises(tvm.error.DiagnosticError, match="aligned to a physical 32-bit column"): + build() + + def test_vector_annotation_shorthand_aliases(): """Test shorthand aliases: T.f32, T.i32, T.f16, etc.""" @@ -1893,8 +2406,8 @@ def func(): assert_structural_equal(func, from_source(code)) -def test_roundtrip_cp_async_bulk_tensor_g2c(): - """cp.async.bulk.tensor.g2c must round-trip with *coords at end.""" +def test_roundtrip_cp_async_bulk_tensor_g2s_cluster(): + """cp.async.bulk.tensor.g2s_cluster must round-trip with *coords at end.""" # fmt: off @T.prim_func(check_well_formed=False) @@ -1904,7 +2417,7 @@ def func(A_ptr: T.handle): with T.launch_thread("blockIdx.x", 1): T.launch_thread("threadIdx.x", 128) A_smem = T.alloc_buffer((16, 16), "float32", scope="shared") - T.ptx.cp_async.bulk.tensor.g2c( + T.ptx.cp_async.bulk.tensor.g2s_cluster( 2, A_smem.data, 0, T.address_of(A_map), 0, 1, "", 0, 0 ) # fmt: on @@ -1935,8 +2448,8 @@ def func(A_ptr: T.handle): assert_structural_equal(func, from_source(code)) -def test_roundtrip_cp_async_bulk_tensor_g2c_prefetch(): - """cp.async.bulk.tensor.g2c_prefetch must round-trip with *coords at end.""" +def test_roundtrip_cp_async_bulk_tensor_prefetch(): + """cp.async.bulk.tensor.prefetch must round-trip with *coords at end.""" # fmt: off @T.prim_func(check_well_formed=False) @@ -1945,7 +2458,7 @@ def func(A_ptr: T.handle): A_map: T.let[T.handle("tensormap")] = T.tvm_stack_alloca("tensormap", 1) with T.launch_thread("blockIdx.x", 1): T.launch_thread("threadIdx.x", 128) - T.ptx.cp_async.bulk.tensor.g2c_prefetch( + T.ptx.cp_async.bulk.tensor.prefetch( 2, T.address_of(A_map), "", 0, 0 ) # fmt: on diff --git a/tests/python/tirx/test_printer_tir_namespaces.py b/tests/python/tirx/test_printer_tir_namespaces.py index 57c989bd4a32..f81307c8a883 100644 --- a/tests/python/tirx/test_printer_tir_namespaces.py +++ b/tests/python/tirx/test_printer_tir_namespaces.py @@ -16,6 +16,8 @@ # under the License. +import pytest + import tvm from tvm import tirx as tir from tvm.script import tirx as T @@ -82,7 +84,14 @@ def test_printer_ptx_more(): cuda_op.ptx_cp_async_bulk_wait_group(0, True), "T.ptx.cp_async.bulk.wait_group(0, T.bool(True))", ) - _assert_print(cuda_op.ptx_cp_async_mbarrier_arrive(0), "T.ptx.cp_async.mbarrier.arrive(0)") + _assert_print( + cuda_op.ptx_cp_async_mbarrier_arrive(r), + 'r = T.handle()\nT.ptx.cp_async.mbarrier.arrive(r, T.bool(False), "shared")', + ) + _assert_print( + cuda_op.ptx_cp_async_mbarrier_arrive_noinc(r), + 'r = T.handle()\nT.ptx.cp_async.mbarrier.arrive(r, T.bool(True), "shared::cta")', + ) _assert_print(cuda_op.ptx_fence("acq_rel", "gpu"), 'T.ptx.fence("acq_rel", "gpu")') _assert_print(cuda_op.ptx_fence("sc", "cta"), 'T.ptx.fence("sc", "cta")') _assert_print( @@ -100,11 +109,16 @@ def test_printer_ptx_more(): cuda_op.ptx_ld_global_acquire(r, s), "r = T.handle()\ns = T.handle()\nT.ptx.ld_global_acquire(r, s)", ) + _assert_print( + cuda_op.cuda_fdividef(1.0, 2.0), + "T.cuda.fdividef(T.float32(1.0), T.float32(2.0))", + ) _assert_print( cuda_op.ptx_map_shared_rank(r, 2), 'r = T.handle()\nT.ptx.mapa(r, 2, "", "u64", "uint64")' ) _assert_print(cuda_op.ptx_bar_arrive(0, 128), "T.ptx.bar.arrive(0, 128)") _assert_print(cuda_op.ptx_bar_sync(0, 128), "T.ptx.bar.sync(0, 128)") + _assert_print(cuda_op.ptx_barrier_sync(0, 128), "T.ptx.barrier.sync(0, 128)") _assert_print( cuda_op.ptx_tcgen05_alloc(s, 64, 1), "s = T.handle()\nT.ptx.tcgen05.alloc(s, 64, 1)" ) @@ -168,6 +182,12 @@ def test_printer_ptx_more(): "d = T.handle()\n" 'T.ptx.tcgen05.cp(a, d, "64x128b", 1, "warpx2::02_13", "", 0, 0)', ) + _assert_print( + cuda_op.ptx_tcgen05_cp(a, d, shape="128x128b", cta_group=2, decompress="b8x16.b6x16_p32"), + "a = T.handle()\n" + "d = T.handle()\n" + 'T.ptx.tcgen05.cp(a, d, "128x128b", 2, "", "b8x16.b6x16_p32", 0, 0)', + ) _assert_print(cuda_op.ptx_tcgen05_shift(a, 1), "a = T.handle()\nT.ptx.tcgen05.shift(a, 1)") _assert_print( cuda_op.ptx_tcgen05_ld(a, 0, shape="16x64b", num=1, row=0, col=0, pack=False), @@ -187,15 +207,46 @@ def test_printer_ptx_more(): ) +@pytest.mark.parametrize( + "kwargs,match", + [ + ({"shape": "bad"}, "invalid shape"), + ({"shape": "128x256b", "multicast": "warpx4"}, "requires multicast=''"), + ({"shape": "64x128b"}, "requires multicast in warpx2"), + ({"shape": "64x128b", "multicast": "warpx4"}, "requires multicast in warpx2"), + ({"shape": "32x128b"}, "requires multicast='warpx4'"), + ({"shape": "32x128b", "multicast": "warpx2::02_13"}, "requires multicast='warpx4'"), + ({"shape": "128x128b", "decompress": "bad"}, "invalid decompress"), + ], +) +def test_ptx_tcgen05_cp_validation(kwargs, match): + a = tir.Var("a", "handle") + d = tir.Var("d", "handle") + with pytest.raises(ValueError, match=match): + cuda_op.ptx_tcgen05_cp(a, d, cta_group=1, **kwargs) + + def test_printer_ptx_mbarrier(): bar = tir.Var("bar", "handle") _assert_print( cuda_op.ptx_mbarrier_init(bar, 32), "bar = T.handle()\nT.ptx.mbarrier.init(bar, 32)" ) - _assert_print(cuda_op.ptx_mbarrier_arrive(bar), "bar = T.handle()\nT.ptx.mbarrier.arrive(bar)") + _assert_print( + cuda_op.ptx_mbarrier_arrive(bar), + 'bar = T.handle()\nT.ptx.mbarrier.arrive(bar, "", "", "shared", 0, 0, 0)', + ) _assert_print( cuda_op.ptx_mbarrier_arrive_expect_tx(bar, 128), - "bar = T.handle()\nT.ptx.mbarrier.arrive.expect_tx(bar, 128)", + 'bar = T.handle()\nT.ptx.mbarrier.arrive.expect_tx(bar, 128, "", "", "shared", 0, 0)', + ) + _assert_print( + cuda_op.ptx_mbarrier_arrive_no_complete(bar, 2), + 'bar = T.handle()\nT.ptx.mbarrier.arrive.no_complete(bar, 2, "shared", 0)', + ) + _assert_print( + cuda_op.ptx_mbarrier_complete_tx(bar, 128), + 'bar = T.handle()\nT.ptx.mbarrier.complete_tx(bar, 128, "relaxed", ' + '"cluster", "shared::cluster", 0, 0)', ) _assert_print( cuda_op.ptx_mbarrier_try_wait(bar, 1), "bar = T.handle()\nT.ptx.mbarrier.try_wait(bar, 1)" @@ -449,21 +500,22 @@ def test_printer_ptx_mma_and_wgmma(): def test_printer_ptx_cp_async_tensor(): tmap = tir.Var("tm", "handle") _assert_print( - cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster(2, tmap, 0, tmap, 0, 1, "", 0, 1, ""), + cuda_op.ptx_cp_async_bulk_tensor_g2s_cluster(2, tmap, 0, tmap, 0, 1, "", 0, 1, ""), "tm = T.handle()\n" - 'T.ptx.cp_async.bulk.tensor.g2c(2, tm, 0, tm, 0, 1, T.uint64(0), 0, 0, 1, "")', + "T.ptx.cp_async.bulk.tensor.g2s_cluster" + '(2, tm, 0, tm, 0, 1, T.uint64(0), 0, "tile", 0, 0, 0, 1, "")', ) _assert_print( - cuda_op.ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster( - 2, tmap, 0, tmap, 0, 1, "", 0, 1, "" + cuda_op.ptx_cp_async_bulk_tensor_g2s_cta( + 2, tmap, 0, tmap, 1, "", 0, 1, 2, 3, 4, load_mode="tile_gather4" ), "tm = T.handle()\n" - "T.ptx.cp_async.bulk.tensor.g2c_tile_gather4" - '(2, tm, 0, tm, 0, 1, T.uint64(0), 0, 0, 1, "")', + "T.ptx.cp_async.bulk.tensor.g2s_cta" + '(2, tm, 0, tm, 1, T.uint64(0), 0, "tile_gather4", 0, 0, 1, 2, 3, 4)', ) _assert_print( - cuda_op.ptx_cp_async_bulk_tensor_global_to_cluster_prefetch(2, tmap, "", 0, 0, ""), - 'tm = T.handle()\nT.ptx.cp_async.bulk.tensor.g2c_prefetch(2, tm, T.uint64(0), 0, 0, 0, "")', + cuda_op.ptx_cp_async_bulk_tensor_prefetch(2, tmap, "", 0, 0, ""), + 'tm = T.handle()\nT.ptx.cp_async.bulk.tensor.prefetch(2, tm, T.uint64(0), 0, 0, 0, "")', ) _assert_print( cuda_op.ptx_cp_async_bulk_tensor_shared_to_global(2, 0, tmap, "", 0, 0, ""), diff --git a/tests/python/tirx/test_tirx_kernels_registry_correctness.py b/tests/python/tirx/test_tirx_kernels_registry_correctness.py new file mode 100644 index 000000000000..eda5e9268276 --- /dev/null +++ b/tests/python/tirx/test_tirx_kernels_registry_correctness.py @@ -0,0 +1,164 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Correctness coverage for workloads maintained by tirx-kernels.""" + +from __future__ import annotations + +import os +import sys +from contextlib import contextmanager +from pathlib import Path + +import pytest + +_WORKSPACE_TIRX_KERNELS = Path(__file__).resolve().parents[4] / "tirx-kernels" +if _WORKSPACE_TIRX_KERNELS.exists(): + sys.path.insert(0, str(_WORKSPACE_TIRX_KERNELS)) + +kernel_registry = pytest.importorskip("tirx_kernels.registry") +kernel_runner = pytest.importorskip("tirx_kernels.runner") +bench_suite_run = pytest.importorskip("tirx_kernels.bench_suite.run") + +_WORKLOADS = bench_suite_run.load_workloads(bench_suite_run.DEFAULT_WORKLOADS) +_KERNELS = { + kernel_name: kernel_registry.load_kernel(kernel_name, strict=True) + for kernel_name in sorted({workload["kernel"] for workload in _WORKLOADS}) +} +_DISTRIBUTED_KERNELS = frozenset( + {"allgather_gemm", "deepgemm_fp8_fp4_mega_moe", "gemm_reduce_scatter"} +) +_XDIST_CUDA_DEVICE = None + + +def _manifest_kernel_config_cases(): + cases = [] + for workload in _WORKLOADS: + kernel_name = workload["kernel"] + label = workload["config"] + mod = _KERNELS[kernel_name] + configs = getattr(mod, "BENCH_CONFIGS", getattr(mod, "CONFIGS", [])) + matches = [config for config in configs if config.get("label") == label] + if len(matches) != 1: + raise ValueError( + f"{kernel_name}::{label} must resolve to exactly one benchmark config, " + f"found {len(matches)}" + ) + config = matches[0] + required_devices = int(config.get("num_processes", config.get("world_size", 1))) + if workload["num_gpus"] != required_devices: + raise ValueError( + f"{kernel_name}::{label} declares num_gpus={workload['num_gpus']}, " + f"but its config requires {required_devices}" + ) + marks = ( + pytest.mark.xdist_group(name="distributed_device_zero") + if kernel_name in _DISTRIBUTED_KERNELS + else () + ) + cases.append(pytest.param(kernel_name, config, id=f"{kernel_name}::{label}", marks=marks)) + return cases + + +def _set_cuda_device_for_xdist_worker(): + global _XDIST_CUDA_DEVICE + + try: + import torch + except ImportError: + return False + + if not torch.cuda.is_available(): + return False + + if _XDIST_CUDA_DEVICE is None: + worker = os.environ.get("PYTEST_XDIST_WORKER", "gw0") + worker_index = int(worker[2:]) if worker.startswith("gw") and worker[2:].isdigit() else 0 + _XDIST_CUDA_DEVICE = worker_index % torch.cuda.device_count() + torch.cuda.set_device(_XDIST_CUDA_DEVICE) + return True + + +def _visible_cuda_device_count(): + try: + import torch + except ImportError: + return 0 + + if not torch.cuda.is_available(): + return 0 + return torch.cuda.device_count() + + +def _required_cuda_device_count(config): + return int(config.get("num_processes", config.get("world_size", 1))) + + +@contextmanager +def _registry_gpu_lock(kernel_name, config): + try: + import fcntl + + import torch + except ImportError: + yield + return + + if not torch.cuda.is_available(): + yield + return + + if kernel_name in _DISTRIBUTED_KERNELS: + device_ids = range(_required_cuda_device_count(config)) + else: + device_ids = (torch.cuda.current_device(),) + + lock_files = [] + try: + for device_id in device_ids: + lock_path = Path("/tmp") / f"tirx-kernels-registry-correctness-device{device_id}.lock" + lock_file = lock_path.open("w") + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + lock_files.append(lock_file) + + torch.cuda.empty_cache() + min_free_gb = float(os.environ.get("TIRX_KERNEL_TEST_MIN_FREE_GB", "32")) + min_free_bytes = int(min_free_gb * 1024**3) + for device_id in device_ids: + free_bytes, _ = torch.cuda.mem_get_info(device_id) + if free_bytes < min_free_bytes: + pytest.skip( + f"CUDA device {device_id} has less than {min_free_gb:g} GiB free memory" + ) + yield + finally: + torch.cuda.empty_cache() + for lock_file in reversed(lock_files): + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + lock_file.close() + + +@pytest.mark.parametrize(("kernel_name", "config"), _manifest_kernel_config_cases()) +def test_manifest_tirx_kernel_correctness(kernel_name, config): + _set_cuda_device_for_xdist_worker() + required_devices = _required_cuda_device_count(config) + visible_devices = _visible_cuda_device_count() + if required_devices > visible_devices: + pytest.skip( + f"requires {required_devices} CUDA devices, but only {visible_devices} are visible" + ) + with _registry_gpu_lock(kernel_name, config): + kernel_runner.run_kernel_test(kernel_name, config, registry=_KERNELS) diff --git a/tests/python/tirx/test_verifier.py b/tests/python/tirx/test_verifier.py index 5ed20e7162fe..a4abb9b298e9 100644 --- a/tests/python/tirx/test_verifier.py +++ b/tests/python/tirx/test_verifier.py @@ -177,7 +177,7 @@ def test1(): # fmt: on verify(test1) - ### SwizzleLayout + ### ComposeLayout (bare swizzle) # fmt: off @T.prim_func(check_well_formed=False) def test2(): @@ -185,7 +185,9 @@ def test2(): T.cta_id([32]) T.warp_id([4]) T.lane_id([32]) - A = T.alloc_buffer((512,), scope="shared", layout=T.SwizzleLayout(3, 3, 3)) + A = T.alloc_buffer( + (512,), scope="shared", layout=T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(512,)])) + ) A[0] = 0 # fmt: on @@ -212,7 +214,7 @@ def test1(A_ptr: T.handle): if threadIdx == 0: T.ptx.mbarrier.init(bar.data, 1) T.ptx.fence.proxy_async("shared::cta") - T.ptx.cp_async.bulk.tensor.g2c(2, A_smem.data, bar.data, T.address_of(A_map), 0, 1, "", 0, 0) # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(2, A_smem.data, bar.data, T.address_of(A_map), 0, 1, "", 0, 0) # noqa: E501 T.ptx.mbarrier.arrive.expect_tx(bar.data, 16*16*4) T.ptx.mbarrier.try_wait(bar.data, phase[0]) phase[0] = phase[0] ^ 1 diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index 764431e1ae13..273f51544455 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -1126,13 +1126,13 @@ def op_call_with_config(A: T.Buffer((10,), "int32"), B: T.Buffer((10,), "int32") Tx.add(A, B, 1.0) op_call_stmt = op_call_with_config.body.body - assert isinstance(op_call_stmt, tir.stmt.TilePrimitiveCall) + assert isinstance(op_call_stmt, tir.TilePrimitiveCall) # Manually construct an TilePrimitiveCall with a Expr in config config_var = Var("config_val", "int32") new_config = dict(op_call_stmt.config) new_config["cta_mask"] = config_var + tir.IntImm("int32", 5) - op_call_with_var = tir.stmt.TilePrimitiveCall( + op_call_with_var = tir.TilePrimitiveCall( *op_call_stmt.args, op=op_call_stmt.op, config=new_config ) @@ -1158,20 +1158,20 @@ def op_call_with_config(A: T.Buffer((10,), "int32"), B: T.Buffer((10,), "int32") Tx.add(A, B, 1.0) op_call_stmt = op_call_with_config.body.body - assert isinstance(op_call_stmt, tir.stmt.TilePrimitiveCall) + assert isinstance(op_call_stmt, tir.TilePrimitiveCall) # Create TilePrimitiveCall with a Var in the config old_var = Var("old_scope_id", "int32") new_var = Var("new_let_var", "int32") new_config = dict(op_call_stmt.config) new_config["cta_mask"] = old_var + tir.IntImm("int32", 5) - op_call_with_var = tir.stmt.TilePrimitiveCall( + op_call_with_var = tir.TilePrimitiveCall( *op_call_stmt.args, op=op_call_stmt.op, config=new_config ) # Substitute old_var -> new_var result = substitute(op_call_with_var, {old_var: new_var}) - assert isinstance(result, tir.stmt.TilePrimitiveCall) + assert isinstance(result, tir.TilePrimitiveCall) # The config value should now reference new_var, not old_var cta_mask_expr = result.config["cta_mask"] @@ -1183,5 +1183,41 @@ def op_call_with_config(A: T.Buffer((10,), "int32"), B: T.Buffer((10,), "int32") ) +def test_op_call_nested_config_visited_and_substituted(): + """Nested selector arrays participate in the core visitor and mutator.""" + from tvm.tirx.stmt_functor import post_order_visit, substitute + + @T.prim_func + def selector( + A: T.Buffer((8,), "float16"), + B: T.Buffer((8,), "float16"), + C: T.Buffer((8,), "float16"), + flag: T.int32, + ): + Tx.copy_async( + C[:], + A[:], + dispatch="tma_explicit", + mbar=C.data, + src_selector=[(flag != 0, B)], + ) + + op_call = selector.body + assert isinstance(op_call, tir.TilePrimitiveCall) + original_b = selector.buffer_map[selector.params[1]] + seen = [] + post_order_visit(selector.body, seen.append) + assert any(isinstance(node, Var) and node.same_as(selector.params[-1]) for node in seen) + undefined = tvm.tirx.analysis.undefined_vars(op_call) + assert any(var.same_as(original_b.data) for var in undefined) + + replacement = Var("replacement", "int32") + updated = substitute(op_call, {selector.params[-1]: replacement}) + condition, candidate = updated.config["src_selector"][0] + assert isinstance(condition, tir.NE) + assert condition.a.same_as(replacement) + assert candidate.same_as(original_b) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/transform/test_transform_lower_tirx.py b/tests/python/tirx/transform/test_transform_lower_tirx.py index 272b4598d432..74b0dee1ca53 100644 --- a/tests/python/tirx/transform/test_transform_lower_tirx.py +++ b/tests/python/tirx/transform/test_transform_lower_tirx.py @@ -510,7 +510,10 @@ def before(A: T.Buffer((128, 32), "float16")) -> None: T.lane_id([32]) tid = T.thread_id([128]) A_smem = T.alloc_buffer( - [128, 32], dtype="float16", scope="shared", layout=T.SwizzleLayout(3, 3, 3) + [128, 32], + dtype="float16", + scope="shared", + layout=T.ComposeLayout(3, 3, 3, T.TileLayout(T.S[(512,)])), ) thread_col = T.meta_var(4) thread_row = T.meta_var(32) @@ -974,6 +977,27 @@ def before(A_ptr: T.handle, B_ptr: T.handle): assert any("T.selector(lane_id, T.ptx.elect_sync() != T.uint32(0))" in item for item in seen) +def test_lower_cleanup_accepts_bool_elect_sync_else_path(): + @T.prim_func(private=True) + def before(A_ptr: T.handle): + A = T.match_buffer(A_ptr, (32,), "int32", scope="global") + T.device_entry() + T.cta_id([1]) + T.warp_id([1]) + lane_id = T.lane_id([32]) + if T.ptx.elect_sync() != T.uint32(0): + A[lane_id] = 1 + else: + A[lane_id] = 0 + + with tvm.target.Target("cuda"): + lowered = LowerTIRx()(tvm.IRModule({"main": before})) + + script = lowered.script(extra_config={"tirx.prefix": "T"}) + assert "T.ptx.elect_sync() != T.uint32(0)" in script + assert "else:" in script + + def test_lower_exec_context_scope_guard_mixes_structural_and_selector(): import tvm.tirx.operator.tile_primitive as _ # noqa: F401 from tvm.tirx.operator.tile_primitive.dispatcher import register_dispatch diff --git a/tests/scripts/setup-pytest-env.sh b/tests/scripts/setup-pytest-env.sh new file mode 100755 index 000000000000..171ddbc2d0d6 --- /dev/null +++ b/tests/scripts/setup-pytest-env.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# NOTE: allow unbound variable here +set +u + +if [[ ! -z $CI_PYTEST_ADD_OPTIONS ]]; then + export PYTEST_ADDOPTS="-s -vv $CI_PYTEST_ADD_OPTIONS $PYTEST_ADDOPTS" +else + export PYTEST_ADDOPTS="-s -vv $PYTEST_ADDOPTS" +fi +set -ux + +export TVM_PATH=`pwd` +export PYTHONPATH="${TVM_PATH}/python" + +export TVM_PYTEST_RESULT_DIR="${TVM_PATH}/build/pytest-results" +mkdir -p "${TVM_PYTEST_RESULT_DIR}" +pytest_errors=() + +# This ensures that all pytest invocations that are run through run_pytest will +# complete and errors will be reported once Bash is done executing all scripts. +function cleanup() { + set +x + if [ "${#pytest_errors[@]}" -gt 0 ]; then + echo "These pytest invocations failed, the results can be found in the Jenkins 'Tests' tab or by scrolling up through the raw logs here." + python3 ci/scripts/jenkins/pytest_wrapper.py "${pytest_errors[@]}" + exit 1 + fi + set -x +} +trap cleanup 0 + +function run_pytest() { + ffi_type="cython" + set -e + local test_suite_name="$1" + shift + extra_args=( "$@" ) + if [ -z "${ffi_type}" -o -z "${test_suite_name}" ]; then + echo "error: run_pytest called incorrectly: run_pytest ${test_suite_name}" "${extra_args[@]}" + echo "usage: run_pytest [pytest args...]" + exit 2 + fi + + # Allow unbound variable here. + set +u + if [[ -z "${TVM_SHARD_INDEX}" ]]; then + current_shard="no-shard" + else + current_shard="shard-${TVM_SHARD_INDEX}" + fi + set -u + + has_reruns=$(python3 -m pytest --help 2>&1 | grep 'reruns=' || true) + if [ -n "$has_reruns" ]; then + if [[ ! "${extra_args[*]}" == *"--reruns"* ]]; then + extra_args+=('--reruns=3') + fi + fi + + suite_name="${test_suite_name}-${current_shard}-${ffi_type}" + + DEFAULT_PARALLELISM=1 + + if [[ ! "${extra_args[*]}" == *" -n"* ]] && [[ ! "${extra_args[*]}" == *" -dist"* ]]; then + extra_args+=("-n=$DEFAULT_PARALLELISM") + fi + + exit_code=0 + set +e + python3 -m pytest \ + -o "junit_suite_name=${suite_name}" \ + "--junit-xml=${TVM_PYTEST_RESULT_DIR}/${suite_name}.xml" \ + "--junit-prefix=${ffi_type}" \ + "${extra_args[@]}" || exit_code=$? + # Pytest will return error code -5 if no test is collected. + if [ "$exit_code" -ne "0" ] && [ "$exit_code" -ne "5" ]; then + pytest_errors+=("${suite_name}: $@") + fi + # To avoid overwriting. + set -e +} From dfd98ea6b8a9530db5453980c6855a2bd96187d6 Mon Sep 17 00:00:00 2001 From: spectrometerHBH Date: Sat, 1 Aug 2026 01:43:09 -0400 Subject: [PATCH 2/2] fix(lower-tirx): keep buffer identity coherent across buffer rebuilds MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit After the typed-buffer-variable migration, a buffer's identity is the variable itself. A pass that rewrites the program while rebuilding buffers is applying a substitution over buffer identities, and that substitution must reach every expression position — including the ones tucked inside buffer types and pass-internal folds. Several spots missed this and broke kernels using shared-memory pools, views sized/offset by local scalars, or host-side tensormap encodes of device-declared views. FlattenBuffer is restructured around its invariant instead of patched: each n-d buffer flattens to a 1-d storage husk buf' (same data origin, dtype, alignment and scope; no layout, no elem_offset), and every access buf[x] rewrites to buf'[f(x)] with f(x) = layout.apply(x, shape) + elem_offset. Since the folded indices live in the rewritten program, f's coefficients (runtime elem_offset, symbolic shapes/strides, layout iters) must be rewritten too. The pass now walks the AST top-down and derives, exactly once at each buffer's definition point (AllocBuffer/DeclBuffer/params), the pair {fold view = rewritten geometry, buf'}; use sites only look the pair up. The substitution is applied exactly once per subexpression by construction, and a use before its definition is a hard error at the fault instead of a stale reference that surfaces later in MakePackedAPI. A buffer_data projection is only resolvable where the buffer's definition is visible, and TMA host-init statements violated that: they are emitted inside the kernel region but hoisted to host scope, carrying projections of device-local views. TilePrimitiveDispatch now tracks each buffer's storage root as it walks definitions and rewrites those projections onto the root — a PrimFunc parameter, visible on the host — at the moment the statements are hoisted. With host statements self-contained, LowerTIRxCleanup's alias lookup no longer has to tolerate forward references, so a projection with no visible definition is an error there instead of silently resolving to itself. Two more substitution obligations were unmet: - LowerTIRxCleanup rebuilt buffers without rewriting the buffer-type shape/stride fields, so a view sized by a local scalar kept pointing at the pre-rebuild variable. - LowerTIRxOpaque's unit-loop Var visitor returned non-unit-loop variables verbatim, shadowing the base visitor's buffer remapping. The shared-memory pool no longer needs buffer-referencing metadata at all. Its allocation is an extern placeholder whose extent never reaches the generated code, so the dynamic-shared-memory launch size is declared directly instead of being patched into the allocation: - SMEMPool.commit() appends a leaf AttrStmt (node=0, key="tirx.dyn_smem_bytes", IntImm value) to the kernel region; hand-written kernels declare theirs with the existing T.attr({"tirx.dyn_smem_bytes": n}) sugar. - SplitHostDevice reads the declaration as the dynamic-shared-memory launch argument and strips it when finalizing the kernel; allocation extents are no longer consulted, and a shared.dyn allocation without a declaration is a hard error. - LowerTIRxOpaque's pool-size collection and extent patching are removed. Fork tests updated for the typed-buffer surface; test_tcgen05_mma_ss_no_tma builds its smem descriptors with first-class ptr_to() instead of a reinterpreting access_ptr over the raw arena, satisfying tvm_access_ptr's element-type check; and test_transform_flatten_buffer.py pins the FlattenBuffer invariant (rewritten references in view shapes and folded elem_offsets; identity preservation for already-flat buffers); all three of its tests fail on the pre-fix pass. --- python/tvm/backend/cuda/lang/alloc_pool.py | 20 +-- python/tvm/tirx/layout.py | 2 +- python/tvm/tirx/stmt.py | 6 + src/backend/cuda/codegen/codegen_cuda.cc | 6 + src/backend/cuda/codegen/codegen_cuda.h | 4 + src/tirx/transform/flatten_buffer.cc | 165 +++++++++++++----- src/tirx/transform/lower_tirx_cleanup.cc | 15 +- src/tirx/transform/lower_tirx_opaque.cc | 54 +----- src/tirx/transform/split_host_device.cc | 56 +++++- src/tirx/transform/tile_primitive_dispatch.cc | 73 +++++++- .../tirx/codegen/test_codegen_blackwell.py | 13 +- .../tirx/codegen/test_codegen_hopper.py | 8 +- .../tirx/iket/test_iket_orchestration.py | 1 + tests/python/tirx/iket/test_iket_profiler.py | 2 + .../cuda/copy/test_gmem_smem.py | 4 +- .../cuda/copy/test_ld_stmatrix.py | 8 +- .../tile_primitive/cuda/copy/test_reg.py | 2 +- .../cuda/copy_async/test_tcgen05_cp.py | 8 +- .../cuda/copy_async/test_tma.py | 8 +- tests/python/tirx/test_parser_printer.py | 2 +- .../tirx/transform/test_stmt_functor.py | 2 +- .../test_transform_flatten_buffer.py | 159 +++++++++++++++++ 22 files changed, 473 insertions(+), 145 deletions(-) create mode 100644 tests/python/tirx/transform/test_transform_flatten_buffer.py diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py b/python/tvm/backend/cuda/lang/alloc_pool.py index 75cf3034bbb2..766713f5a541 100644 --- a/python/tvm/backend/cuda/lang/alloc_pool.py +++ b/python/tvm/backend/cuda/lang/alloc_pool.py @@ -40,12 +40,6 @@ def _get_ir(): return _ir -def _get_frame(): - from tvm.tirx.script.builder import frame - - return frame - - # --------------------------------------------------------------------------- # Shared utilities # --------------------------------------------------------------------------- @@ -490,16 +484,16 @@ def commit(self, size=None): """ if not self._owns_buffer: return - ir = _get_ir() - frame_mod = _get_frame() resolved = size if size is not None else self.max_offset assert resolved >= self.max_offset, ( f"Specified smem size ({resolved}) is smaller than " f"the pool high-water mark ({self.max_offset})" ) - attr_frame = ir.attr(self.ptr, "tirx.pool_max_bytes", resolved) - if isinstance(attr_frame, frame_mod.AttrFrame): - from functools import partial + import tvm.tirx - attr_frame.add_callback(partial(attr_frame.__exit__, None, None, None)) - attr_frame.__enter__() + ir = _get_ir() + ir.add_to_parent( + tvm.tirx.AttrStmt( + 0, "tirx.dyn_smem_bytes", tvm.tirx.IntImm("int64", resolved), tvm.tirx.Evaluate(0) + ) + ) diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py index 49b6ebbe6dbd..92f4bb9044b6 100644 --- a/python/tvm/tirx/layout.py +++ b/python/tvm/tirx/layout.py @@ -712,7 +712,7 @@ def tmem_datapath_layout(datapath: str, rows: int, cols: int, sub_slab: int = 0) ) if cols % 2 != 0: raise ValueError( - f"tmem_datapath_layout: datapath={datapath!r} requires even cols, got {cols}" + f"tmem_datapath_layout: datapath={datapath!r} expects even cols, got {cols}" ) return TileLayout(S[(rows, 2, cols // 2) : (1 @ tlane, 64 @ tlane, 1 @ tcol)]) # Layouts C (M=128 cta_group::2 sparse) and F (M=64 non-.ws): rows-per-CTA diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index e68046774423..fe4c98a4025a 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -938,3 +938,9 @@ def stmt_list(stmt: Stmt) -> list[Stmt]: res += stmt_list(x) return res return [stmt] + + +# Source-compatibility re-export: TilePrimitiveCall lives in tile_primitive.py +# after the tile-primitive module merge. Imported last to avoid a cycle with +# tile_primitive's own ``from .stmt import Stmt``. +from .tile_primitive import TilePrimitiveCall # noqa: E402,F401 isort: skip diff --git a/src/backend/cuda/codegen/codegen_cuda.cc b/src/backend/cuda/codegen/codegen_cuda.cc index 8db7f3d690e2..e1c1811bbc6b 100644 --- a/src/backend/cuda/codegen/codegen_cuda.cc +++ b/src/backend/cuda/codegen/codegen_cuda.cc @@ -182,6 +182,7 @@ void CodeGenCUDA::PrintFunctionSignature(const ffi::String& function_name, const std::ostream& os) { CallingConv calling_conv = func->GetAttr(tvm::attr::kCallingConv, CallingConv::kDefault).value(); + in_kernel_launch_ = (calling_conv == CallingConv::kDeviceKernelLaunch); if (calling_conv == CallingConv::kDeviceKernelLaunch) { os << "extern \"C\" __global__ "; } else if (calling_conv == CallingConv::kDefault) { @@ -267,6 +268,11 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f, std::ostream& os) { } void CodeGenCUDA::VisitStmt_(const ReturnNode* op) { + if (!in_kernel_launch_) { + // __device__ subroutines return real values. + CodeGenC::VisitStmt_(op); + return; + } const auto* value = op->value.as(); TVM_FFI_ICHECK(value && value->value == 0) << "CUDA device kernel may only contain a successful early return, return 0"; diff --git a/src/backend/cuda/codegen/codegen_cuda.h b/src/backend/cuda/codegen/codegen_cuda.h index 6b43cd890905..26a657a2a415 100644 --- a/src/backend/cuda/codegen/codegen_cuda.h +++ b/src/backend/cuda/codegen/codegen_cuda.h @@ -109,6 +109,10 @@ class CodeGenCUDA final : public CodeGenC { // y/z cluster-CTA extents are both one. bool cluster_cta_x_is_linear_rank_{false}; + // Whether the function being generated is a __global__ kernel entry + // (kDeviceKernelLaunch); __device__ subroutines return real values. + bool in_kernel_launch_{false}; + // Codegen tags std::unordered_set codegen_tags_; diff --git a/src/tirx/transform/flatten_buffer.cc b/src/tirx/transform/flatten_buffer.cc index 8650099a2142..3dbb703f994c 100644 --- a/src/tirx/transform/flatten_buffer.cc +++ b/src/tirx/transform/flatten_buffer.cc @@ -39,8 +39,23 @@ namespace tvm { namespace tirx { /*! - * \brief Transform multi-dimension BufferLoad/BufferStore into device-supported dimension - * for the TIR not contains opaque block. + * \brief Flatten each n-d buffer ``buf`` into a 1-d storage view ``buf'``, + * rewriting every access ``buf[x]`` into ``buf'[f(x)]``. + * + * The invariant: ``f(x) = layout.apply(x, shape) + elem_offset`` is fully + * determined by ``buf``'s geometry, and ``buf'`` is only a storage husk — + * same data origin, dtype, alignment and scope; no layout, no elem_offset. + * + * The pass walks the AST top-down. At each buffer definition point + * (AllocBuffer/DeclBuffer; PrimFunc params are seeded up front) it derives, + * exactly once: + * - the fold view: the original geometry with its expression fields + * (runtime elem_offset, symbolic shapes/strides, layout iters) rewritten + * by the pass — the folded indices live in the rewritten program, so + * ``f``'s coefficients must reference rebuilt buffers; and + * - ``buf'``, the flattened storage husk. + * Every use site then only looks the pair up; a use before its definition is + * a hard error instead of a silently stale reference. */ class BufferFlattener : public arith::IRMutatorWithAnalyzer { public: @@ -50,6 +65,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { pass.MarkBufferMapShapes(func); for (const auto& [param, buffer] : func->buffer_map) { pass.extern_buffers_.insert(buffer); + pass.Define(buffer); } auto body = pass.VisitStmt(func->body); @@ -62,7 +78,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { if (auto opt = func->buffer_map.Get(handle)) { auto old_buf = opt.value(); if (pass.buffers_used_.count(old_buf)) { - auto new_buf = pass.GetFlattenedBuffer(old_buf); + auto new_buf = pass.Lookup(old_buf).flattened; if (!old_buf.same_as(new_buf)) { body = SeqStmt::Flatten(DeclBuffer(new_buf, old_buf.data()), std::move(body)); } @@ -84,6 +100,83 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { explicit BufferFlattener(const arith::Analyzer& ana) : IRMutatorWithAnalyzer(ana) {} + struct FlatInfo { + /*! \brief Original geometry with rewritten expression fields; the source + * of ``f``. Only used to fold indices, never emitted into the IR. */ + BufferVar fold_view; + /*! \brief The 1-d storage husk ``buf'``. */ + BufferVar flattened; + }; + + /*! \brief Derive {fold view, flattened husk} for ``buf`` at its definition + * point. Idempotent so params can be seeded up front. */ + const FlatInfo& Define(const BufferVar& buf) { + if (auto it = flat_map_.find(buf.var()); it != flat_map_.end()) { + return it->second; + } + + // Fold view: rewrite the geometry's expression leaves. + auto view_type = CopyBufferType(buf); + for (size_t i = 0; i < view_type->shape.size(); ++i) { + view_type->shape.Set(i, this->VisitPrimExpr(view_type->shape[i])); + } + for (size_t i = 0; i < view_type->strides.size(); ++i) { + view_type->strides.Set(i, this->VisitPrimExpr(view_type->strides[i])); + } + if (view_type->elem_offset.defined()) { + view_type->elem_offset = this->VisitPrimExpr(view_type->elem_offset); + } + if (auto tile = view_type->layout.as()) { + auto remap_iter = [this](const Iter& iter) { + PrimExpr extent = this->VisitPrimExpr(iter->extent); + PrimExpr stride = this->VisitPrimExpr(iter->stride); + if (extent.same_as(iter->extent) && stride.same_as(iter->stride)) { + return iter; + } + return Iter(extent, stride, iter->axis); + }; + auto shard = tile->shard.Map(remap_iter); + auto replica = tile->replica.Map(remap_iter); + if (!shard.same_as(tile->shard) || !replica.same_as(tile->replica)) { + view_type->layout = TileLayout(shard, replica, tile->offset); + } + } + BufferVar fold_view = RebuildBufferVar(buf, std::move(view_type)); + + // buf': the storage husk. The linearized indices carry layout and + // elem_offset, so the husk keeps neither. + auto flat = fold_view.GetFlattenedBuffer(); + auto type = CopyBufferType(flat); + for (size_t i = 0; i < type->shape.size(); ++i) { + type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i])); + } + type->layout = std::nullopt; + if (type->elem_offset.defined() && !is_zero(type->elem_offset)) { + type->elem_offset = IntImm(type->elem_offset.ty().as_or_throw(), 0); + } + // Body-local buffers keep their identity when flattening changes nothing. + // PrimFunc-parameter buffers always rebuild: the epilogue aliases the + // rebuilt view onto the argument buffer with an explicit DeclBuffer, and + // downstream s_tir passes pin that shape. + BufferVar flattened = + (!extern_buffers_.count(buf) && ffi::StructuralEqual()(BufferType(type), buf.type())) + ? buf + : RebuildBufferVar(buf, std::move(type)); + + // Feed the base mutator's remap so stray buffer-var expressions follow. + buffer_remap_.Set(buf, flattened); + auto [it, inserted] = flat_map_.emplace(buf.var(), FlatInfo{fold_view, flattened}); + return it->second; + } + + const FlatInfo& Lookup(const BufferVar& buf) { + auto it = flat_map_.find(buf.var()); + TVM_FFI_ICHECK(it != flat_map_.end()) + << "Buffer " << buf.name() + << " is used before its definition (AllocBuffer/DeclBuffer/PrimFunc param)"; + return it->second; + } + Stmt VisitStmt_(const SBlockNode* op) final { TVM_FFI_ICHECK_EQ(op->match_buffers.size(), 0) << "Unexpected MatchBufferRegion found during tirx.transform.FlattenBuffer. " @@ -92,7 +185,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { SBlock block = ffi::GetRef(op); ffi::Array alloc_buffers = op->alloc_buffers; - alloc_buffers.MutateByApply([this](BufferVar buf) { return GetFlattenedBuffer(buf); }); + alloc_buffers.MutateByApply([this](BufferVar buf) { return Define(buf).flattened; }); if (!alloc_buffers.same_as(op->alloc_buffers)) { block.CopyOnWrite()->alloc_buffers = alloc_buffers; } @@ -113,14 +206,13 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { } Stmt VisitStmt_(const AllocBufferNode* op) final { - auto node = StmtExprMutator::VisitStmt_(op).as_or_throw(); - - auto new_buf = GetFlattenedBuffer(node->buffer); - if (!node->buffer.same_as(new_buf)) { - node.CopyOnWrite()->buffer = new_buf; + const FlatInfo& info = Define(op->buffer); + if (info.flattened.same_as(op->buffer)) { + return ffi::GetRef(op); } - - return std::move(node); + auto n = CopyOnWrite(op); + n->buffer = info.flattened; + return Stmt(n); } Stmt VisitStmt_(const DeclBufferNode* op) final { @@ -135,29 +227,11 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { if (!is_extern_buffer_source) { data = VisitExpr(op->data); } - BufferVar flattened = GetFlattenedBuffer(op->buffer); - if (flattened.same_as(op->buffer) && data.same_as(op->data)) { + const FlatInfo& info = Define(op->buffer); + if (info.flattened.same_as(op->buffer) && data.same_as(op->data)) { return ffi::GetRef(op); } - return DeclBuffer(flattened, std::move(data), op->span); - } - - BufferVar GetFlattenedBuffer(BufferVar buf) { - if (auto remapped = buffer_remap_.Get(buf)) { - return remapped.value(); - } - auto flattened = buf.GetFlattenedBuffer(); - ffi::ObjectPtr type = CopyBufferType(flattened); - - // canonicalize shape - for (size_t i = 0; i < flattened->shape.size(); ++i) { - type->shape.Set(i, analyzer_->canonical_simplify(flattened->shape[i])); - } - type->layout = std::nullopt; - flattened = RebuildBufferVar(flattened, std::move(type)); - - buffer_remap_.Set(buf, flattened); - return flattened; + return DeclBuffer(info.flattened, std::move(data), op->span); } Stmt VisitStmt_(const BufferStoreNode* op) final { @@ -180,16 +254,15 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { if (var.value()->ty.as()) { BufferVar original(var.value()); buffers_used_.insert(original); - return GetFlattenedBuffer(original).data(); + return Lookup(original).flattened.data(); } } } return IRMutatorWithAnalyzer::VisitExpr_(op); } - ffi::Array GetSimplifiedElemOffset(const BufferVar& buffer, - const ffi::Array& indices) { - auto flattened_indices = buffer->ElemOffset(indices); + ffi::Array FoldIndices(const FlatInfo& info, const ffi::Array& indices) { + auto flattened_indices = info.fold_view->ElemOffset(indices); return this->IterMapSimplifyWithContext(flattened_indices, false); } @@ -197,19 +270,18 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { Node VisitBufferAccess(Node node, const BufferVar& original_buffer) { TVM_FFI_ICHECK(node->buffer.defined()); buffers_used_.insert(original_buffer); - auto flattened_indices = GetSimplifiedElemOffset(original_buffer, node->indices); - BufferVar flattened_buffer = GetFlattenedBuffer(original_buffer); + const FlatInfo& info = Lookup(original_buffer); + auto flattened_indices = FoldIndices(info, node->indices); auto writer = node.CopyOnWrite(); - writer->buffer = flattened_buffer; + writer->buffer = info.flattened; writer->indices = flattened_indices; return node; } BufferRegion MutateBufferRegion(BufferRegion region) { - BufferVar orig_buf = region->buffer; - BufferVar flattened_buf = GetFlattenedBuffer(orig_buf); - if (flattened_buf.same_as(orig_buf)) { + const FlatInfo& info = Lookup(region->buffer); + if (info.flattened.same_as(region->buffer)) { return region; } @@ -220,8 +292,8 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { max_values.push_back(range->min + range->extent - 1); } - ffi::Array flattened_min = GetSimplifiedElemOffset(orig_buf, min_values); - ffi::Array flattened_max = GetSimplifiedElemOffset(orig_buf, max_values); + ffi::Array flattened_min = FoldIndices(info, min_values); + ffi::Array flattened_max = FoldIndices(info, max_values); ffi::Array flattened_ranges; TVM_FFI_ICHECK_EQ(flattened_min.size(), flattened_max.size()); @@ -229,7 +301,7 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { flattened_ranges.push_back(Range(flattened_min[i], flattened_max[i] + 1)); } - return BufferRegion(flattened_buf, flattened_ranges); + return BufferRegion(info.flattened, flattened_ranges); } /*! \brief Set of buffers accessed during visitation (used to emit DeclBuffer for param buffers). @@ -238,6 +310,9 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer { /*! \brief Buffers whose storage is supplied by a PrimFunc parameter. */ std::unordered_set extern_buffers_; + + /*! \brief Per-buffer {fold view, flattened husk}, derived at definition points. */ + std::unordered_map flat_map_; }; PrimFunc FlattenBuffer(PrimFunc f) { return BufferFlattener::Flatten(f); } diff --git a/src/tirx/transform/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index b2e1ecd16ae2..4bf6176495bc 100644 --- a/src/tirx/transform/lower_tirx_cleanup.cc +++ b/src/tirx/transform/lower_tirx_cleanup.cc @@ -111,7 +111,11 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { if (auto var = op->args[0].as(); var.has_value() && var.value()->ty.as()) { - Var root = buffer_aliases_.Get(var.value()).value_or(var.value()); + auto root_opt = buffer_aliases_.Get(var.value()); + TVM_FFI_ICHECK(root_opt.has_value()) + << "buffer_data projects " << var.value()->name << ", which has no visible definition " + << "(AllocBuffer/DeclBuffer/PrimFunc parameter) at this point"; + Var root = root_opt.value(); if (auto it = var_remap_.find(root); it != var_remap_.end()) { root = it->second; } @@ -200,9 +204,14 @@ class LayoutApplier : public arith::IRMutatorWithAnalyzer { flattened = buf.GetFlattenedBuffer(); type = CopyBufferType(flattened); } - // canonicalize shape + // Remap variables the pass has already rebuilt (a shape may load from + // another local buffer), then canonicalize. for (size_t i = 0; i < type->shape.size(); ++i) { - type->shape.Set(i, analyzer_->canonical_simplify(type->shape[i])); + type->shape.Set( + i, analyzer_->canonical_simplify(StmtExprMutator::VisitPrimExpr(type->shape[i]))); + } + for (size_t i = 0; i < type->strides.size(); ++i) { + type->strides.Set(i, StmtExprMutator::VisitPrimExpr(type->strides[i])); } type->layout = std::nullopt; type->elem_offset = StmtExprMutator::VisitPrimExpr(buf->elem_offset); diff --git a/src/tirx/transform/lower_tirx_opaque.cc b/src/tirx/transform/lower_tirx_opaque.cc index 12763457f231..71fed942ec6c 100644 --- a/src/tirx/transform/lower_tirx_opaque.cc +++ b/src/tirx/transform/lower_tirx_opaque.cc @@ -44,55 +44,9 @@ namespace tirx { */ class TIRxOpaqueLower : public StmtExprMutator { public: - static Stmt Rewrite(Stmt body) { - TIRxOpaqueLower lower; - lower.pool_sizes_ = CollectPoolSizes(body); - return lower(std::move(body)); - } + static Stmt Rewrite(Stmt body) { return TIRxOpaqueLower()(std::move(body)); } private: - static std::unordered_map CollectPoolSizes( - const Stmt& body) { - class Collector : public StmtVisitor { - public: - void VisitStmt_(const AttrStmtNode* op) final { - if (op->attr_key == "tirx.pool_max_bytes") { - if (auto var = op->node.try_cast()) { - const auto* n = op->value.as(); - TVM_FFI_ICHECK(n) << "TIRxError: tirx.pool_max_bytes must be IntImm"; - pool_sizes_[var.value()] = n->value; - } - } - StmtVisitor::VisitStmt_(op); - } - - std::unordered_map pool_sizes_; - }; - - Collector collector; - collector(body); - return std::move(collector.pool_sizes_); - } - - Stmt VisitStmt_(const AttrStmtNode* op) final { - if (op->attr_key == "tirx.pool_max_bytes") { - // Strip the pool size AttrStmt after pre-collection in Rewrite(). - return VisitStmt(op->body); - } - return StmtExprMutator::VisitStmt_(op); - } - - Stmt VisitStmt_(const AllocBufferNode* op) final { - auto it = pool_sizes_.find(op->buffer.var()); - if (it != pool_sizes_.end()) { - auto type = CopyBufferType(op->buffer); - type->shape = {IntImm::Int64(it->second)}; - BufferVar alloc_buf = RebuildBufferVar(op->buffer, std::move(type)); - buffer_remap_.Set(op->buffer, alloc_buf); - } - return StmtExprMutator::VisitStmt_(op); - } - Stmt VisitStmt_(const ForNode* op) final { // Step 1. Update unit loop info. PrimExpr min = this->VisitPrimExpr(op->min); @@ -135,7 +89,9 @@ class TIRxOpaqueLower : public StmtExprMutator { Var var = ffi::GetRef(op); auto it = unit_loop_vars_.find(var); if (it == unit_loop_vars_.end()) { - return var; + // Fall through to the base visitor so buffer-variable remapping from + // any rebuild in this pass reaches remaining use sites. + return StmtExprMutator::VisitExpr_(op); } else { PrimExpr expr = it->second; PrimType var_ty = var->ty.as_or_throw(); @@ -205,8 +161,6 @@ class TIRxOpaqueLower : public StmtExprMutator { /*! \brief Record the loop_var and loop start value of unit loops, whose extent is one. */ std::unordered_map unit_loop_vars_; - /*! \brief Pool size annotations: buffer data var → size in bytes. */ - std::unordered_map pool_sizes_; }; namespace transform { diff --git a/src/tirx/transform/split_host_device.cc b/src/tirx/transform/split_host_device.cc index 0251afd360fc..cd87fc608700 100644 --- a/src/tirx/transform/split_host_device.cc +++ b/src/tirx/transform/split_host_device.cc @@ -352,8 +352,20 @@ class DeviceInfoCollector : public StmtVisitor { if (collector.use_cooperative_launch_) { collector.info_.launch_params.push_back(tvm::runtime::launch_param::kUseCooperativeLaunch); } - // The dynamic shared memory is required to be the last of the - // kernel launch parameters. + // The dynamic shared memory is required to be the last of the kernel + // launch parameters. An explicit tirx.dyn_smem_bytes declaration wins; + // otherwise fall back to the size inferred from the allocation extent. + // A zero-extent allocation is a pool-style extern placeholder, so having + // neither a declaration nor a usable extent is an authoring error. + if (!collector.dyn_shmem_size.has_value() && collector.inferred_shmem_size_.has_value()) { + const auto* inferred = collector.inferred_shmem_size_.value().as(); + TVM_FFI_ICHECK(!(inferred && inferred->value == 0)) + << "PrimFunc " << gvar->name_hint + << " allocates dynamic shared memory with a placeholder extent but does not declare " + "its size; annotate the kernel with tirx.dyn_smem_bytes (SMEMPool.commit() emits " + "it)."; + collector.dyn_shmem_size = collector.inferred_shmem_size_; + } if (collector.dyn_shmem_size) { collector.info_.launch_params.push_back( tvm::runtime::launch_param::kUseDynamicSharedMemoryTag); @@ -378,7 +390,7 @@ class DeviceInfoCollector : public StmtVisitor { if (launch_param == tvm::runtime::launch_param::kUseDynamicSharedMemoryTag) { TVM_FFI_ICHECK(dyn_shmem_size.has_value()) << "Compute kernel requires launch parameter \"" << launch_param - << "\", but PrimFunc did not contain AllocBuffer node with shared dynamic scope."; + << "\", but PrimFunc did not declare tirx.dyn_smem_bytes."; return dyn_shmem_size.value(); } @@ -407,6 +419,15 @@ class DeviceInfoCollector : public StmtVisitor { } void VisitStmt_(const AttrStmtNode* op) final { + if (op->attr_key == "tirx.dyn_smem_bytes") { + // Kernel-level declaration of the dynamic shared memory launch size. + // The backing shared.dyn allocation is an extern placeholder; this + // attribute is the single source of truth for the launch parameter. + TVM_FFI_ICHECK(!dyn_shmem_size.has_value()) + << "Only one tirx.dyn_smem_bytes declaration is allowed per kernel."; + TVM_FFI_ICHECK(op->value.as()) << "tirx.dyn_smem_bytes must be an IntImm"; + dyn_shmem_size = op->value; + } if (op->attr_key == attr::thread_extent) { ffi::String thread_tag; if (auto iv = op->node.as()) { @@ -437,21 +458,24 @@ class DeviceInfoCollector : public StmtVisitor { void VisitStmt_(const AllocBufferNode* op) final { auto storage_scope = runtime::StorageScope::Create(op->buffer.scope()); if (storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == ".dyn") { - TVM_FFI_ICHECK(!dyn_shmem_size.has_value()) + TVM_FFI_ICHECK(!saw_dyn_shared_alloc_) << "Only one dynamic shared memory allocation is allowed."; - TVM_FFI_ICHECK_GT(op->buffer->shape.size(), 0); + saw_dyn_shared_alloc_ = true; + // Fallback launch size inferred from the allocation extent, used when + // no tirx.dyn_smem_bytes declaration is present (e.g. s_tir schedules + // allocate shared.dyn with a concrete extent). A zero extent is a + // pool-style extern placeholder and carries no size information. + TVM_FFI_ICHECK_GT(op->buffer->shape.size(), 0); PrimExpr dyn_size = IntImm::Int32(1); for (const auto& extent : op->buffer->shape) { dyn_size *= extent; } dyn_size *= IntImm::Int64(static_cast(op->buffer->dtype.StorageBytes())); - - // Inline any locally-bound variables (e.g. from CSE). if (bind_map_.size()) { dyn_size = Substitute(dyn_size, bind_map_); } - dyn_shmem_size = dyn_size; + inferred_shmem_size_ = dyn_size; } StmtVisitor::VisitStmt_(op); } @@ -464,6 +488,11 @@ class DeviceInfoCollector : public StmtVisitor { ffi::Map thread_extent; // The amount of dynamic shared memory used. ffi::Optional dyn_shmem_size{std::nullopt}; + // Whether a shared.dyn allocation was seen. + bool saw_dyn_shared_alloc_{false}; + // Launch size inferred from the allocation extent (fallback when no + // tirx.dyn_smem_bytes declaration is present). + ffi::Optional inferred_shmem_size_{std::nullopt}; // Flag-only launch attributes requested by the original PrimFunc. bool use_programmatic_dependent_launch_{false}; bool use_cooperative_launch_{false}; @@ -578,6 +607,17 @@ class DeviceKernelMutator : public StmtExprMutator { Target target = func->GetAttr(tvm::attr::kTarget).value(); bool preserve_early_returns = target->kind->name == "cuda"; write_ptr->body = ReturnRemover::Apply(write_ptr->body, !preserve_early_returns); + // The dyn-smem size declaration was consumed by DeviceInfoCollector; + // it has no meaning inside the kernel body. + class StripDynSmemAttr : public StmtMutator { + Stmt VisitStmt_(const AttrStmtNode* op) final { + if (op->attr_key == "tirx.dyn_smem_bytes") { + return VisitStmt(op->body); + } + return StmtMutator::VisitStmt_(op); + } + }; + write_ptr->body = StripDynSmemAttr()(std::move(write_ptr->body)); } func = WithAttrs(std::move(func), diff --git a/src/tirx/transform/tile_primitive_dispatch.cc b/src/tirx/transform/tile_primitive_dispatch.cc index dfd00c39e638..b6b2491d268b 100644 --- a/src/tirx/transform/tile_primitive_dispatch.cc +++ b/src/tirx/transform/tile_primitive_dispatch.cc @@ -423,7 +423,12 @@ class TilePrimitiveDispatcher : public StmtExprMutator { // Insert host init stmts outside the outermost thread binding or block. if (is_first_thread_attr_) { for (const auto& stmt : host_init_stmts_) { - res = KernelReplacePointSearcher::Seek(stmt, std::move(res)); + // These statements leave the kernel region for host scope, where a + // ``buffer_data`` projection of a device-local view cannot be + // resolved. Rewrite each projection onto its storage root, which is + // a PrimFunc parameter and therefore visible on the host. + res = KernelReplacePointSearcher::Seek(StorageRootResolver::Apply(stmt, buffer_root_), + std::move(res)); } host_init_stmts_.clear(); } @@ -491,11 +496,74 @@ class TilePrimitiveDispatcher : public StmtExprMutator { return StmtExprMutator::VisitStmt_(op); } + /*! + * \brief Track the storage root of a buffer variable. + * + * A ``DeclBuffer`` whose data is ``buffer_data(src)`` is a view over + * ``src``'s storage, so it inherits ``src``'s root; anything else owns its + * storage. Buffers with no definition in the body (PrimFunc parameters) + * are absent from the map and are their own root. + */ + void RegisterStorageRoot(const Var& old_var, const Var& new_var, + const ffi::Optional& data) { + Var root = new_var; + if (data.has_value()) { + if (const auto* call = data.value().as(); + call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) { + if (auto src = call->args[0].as(); + src.has_value() && src.value()->ty.as()) { + root = StorageRootOf(src.value()); + } + } + } + buffer_root_[new_var] = root; + if (!old_var.same_as(new_var)) { + buffer_root_[old_var] = root; + } + } + + Var StorageRootOf(const Var& var) const { + auto it = buffer_root_.find(var); + return it == buffer_root_.end() ? var : it->second; + } + + /*! \brief Rewrite ``buffer_data(view)`` onto ``buffer_data(storage root)``. */ + class StorageRootResolver : public StmtExprMutator { + public: + static Stmt Apply( + Stmt stmt, + const std::unordered_map& buffer_root) { + StorageRootResolver resolver(buffer_root); + return resolver(std::move(stmt)); + } + + private: + explicit StorageRootResolver( + const std::unordered_map& buffer_root) + : buffer_root_(buffer_root) {} + + Expr VisitExpr_(const CallNode* op) final { + if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) { + if (auto var = op->args[0].as(); + var.has_value() && var.value()->ty.as()) { + auto it = buffer_root_.find(var.value()); + if (it != buffer_root_.end() && !it->second.same_as(var.value())) { + return BufferVar(it->second).data(); + } + } + } + return StmtExprMutator::VisitExpr_(op); + } + + const std::unordered_map& buffer_root_; + }; + Stmt VisitStmt_(const AllocBufferNode* op) final { BufferVar old_buffer = op->buffer; Stmt stmt = StmtExprMutator::VisitStmt_(op); op = stmt.as(); TVM_FFI_ICHECK(op); + RegisterStorageRoot(old_buffer.var(), op->buffer.var(), std::nullopt); std::vector seq{stmt}; AppendPostBufferDefStmts(&seq, old_buffer, op->buffer); @@ -507,6 +575,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator { Stmt stmt = StmtExprMutator::VisitStmt_(op); op = stmt.as(); TVM_FFI_ICHECK(op); + RegisterStorageRoot(old_buffer.var(), op->buffer.var(), op->data); std::vector seq{stmt}; AppendPostBufferDefStmts(&seq, old_buffer, op->buffer); @@ -1448,6 +1517,8 @@ class TilePrimitiveDispatcher : public StmtExprMutator { std::vector alloc_buffers_; std::vector device_init_stmts_; std::vector host_init_stmts_; + /*! \brief Storage root of each buffer variable defined in the body. */ + std::unordered_map buffer_root_; std::unordered_map, ffi::ObjectPtrHash, ffi::ObjectPtrEqual> post_buffer_def_stmts_; ffi::Map shared_state_; diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py index 2d2991261f0d..4340c9e0b3c7 100644 --- a/tests/python/tirx/codegen/test_codegen_blackwell.py +++ b/tests/python/tirx/codegen/test_codegen_blackwell.py @@ -60,8 +60,11 @@ def _assert_remote_mbarrier_ir(func, arrive_op_name): def visit(node): if isinstance(node, tvm.tirx.Bind) and node.var.name == "remote_mbar_ptr": bindings.append(node) - if isinstance(node, tvm.tirx.DeclBuffer) and node.buffer.data.name == "remote_mbar_ptr": - buffers.append(node.buffer) + if ( + isinstance(node, tvm.tirx.DeclBuffer) + and getattr(node.data, "name", None) == "remote_mbar_ptr" + ): + buffers.append(node) if isinstance(node, tvm.ir.Call) and node.op.name == "tirx.ptx.mapa": mapa_calls.append(node) if isinstance(node, tvm.ir.Call) and node.op.name == arrive_op_name: @@ -77,7 +80,7 @@ def visit(node): assert bindings[0].value.ty.storage_scope == "shared" assert buffers[0].data.same_as(bindings[0].var) assert buffers[0].data.ty.storage_scope == "shared" - assert buffers[0].scope() == "shared" + assert buffers[0].buffer.scope() == "shared" for arrive in arrive_calls: tvm.ir.assert_structural_equal(arrive.args[0], mapa_calls[0].args[0]) tvm.ir.assert_structural_equal( @@ -516,8 +519,8 @@ def test_mma_ss_no_tma(A: T.Buffer((M, K), a_type, layout=T.TileLayout(T.S[M, K] T.ptx.mbarrier.init(bar.data, 1) T.ptx.tcgen05.encode_instr_descriptor(descI.data, d_dtype=d_type, a_dtype=a_type, b_dtype=b_type, M=M, N=N, K=MMA_K, trans_a=False, trans_b=False, n_cta_groups=cta_group) # noqa: E501 for k in range(K // MMA_K): - T.ptx.tcgen05.encode_matrix_descriptor(descA.data, A_smem.access_ptr("r", offset=A_smem.elem_offset_of([0, k * MMA_K])), ldo=ldo, sdo=sdo, swizzle=SWIZZLE) # noqa: E501 - T.ptx.tcgen05.encode_matrix_descriptor(descB.data, B_smem.access_ptr("r", offset=B_smem.elem_offset_of([0, k * MMA_K])), ldo=ldo, sdo=sdo, swizzle=SWIZZLE) # noqa: E501 + T.ptx.tcgen05.encode_matrix_descriptor(descA.data, A_smem.ptr_to([0, k * MMA_K]), ldo=ldo, sdo=sdo, swizzle=SWIZZLE) # noqa: E501 + T.ptx.tcgen05.encode_matrix_descriptor(descB.data, B_smem.ptr_to([0, k * MMA_K]), ldo=ldo, sdo=sdo, swizzle=SWIZZLE) # noqa: E501 if k == 0: T.ptx.tcgen05.mma(tmem_addr, descA[0], descB[0], descI[0], d_dtype=d_type, a_dtype=a_type, b_dtype=b_type, use_a_tmem=False, cta_group=cta_group, enable_input_d=0) # noqa: E501 else: diff --git a/tests/python/tirx/codegen/test_codegen_hopper.py b/tests/python/tirx/codegen/test_codegen_hopper.py index ac5edb872afb..c8fee4061660 100644 --- a/tests/python/tirx/codegen/test_codegen_hopper.py +++ b/tests/python/tirx/codegen/test_codegen_hopper.py @@ -819,16 +819,16 @@ def main(A_ptr: T.handle, B_ptr: T.handle): T.ptx.fence.proxy_async("shared::cta") T.ptx.mbarrier.arrive.expect_tx(T.address_of(bar), total_bytes) if clusterCtaIdx == 0: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord0[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord0[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord0) # noqa: E501 if clusterCtaIdx == 1: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord1[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord1[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord1) # noqa: E501 if clusterCtaIdx == 2: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord2[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord2[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord2) # noqa: E501 if clusterCtaIdx == 3: - T.ptx.cp_async.bulk.tensor.g2c(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord3[::-1])), # noqa: E501 + T.ptx.cp_async.bulk.tensor.g2s_cluster(len(shape), A_smem.access_ptr(BufferAccessKind.WRITE, offset=A_smem.elem_offset_of(coord3[::-1])), # noqa: E501 T.address_of(bar), T.address_of(A_map), int("1111", 2), 1, "", *coord3) # noqa: E501 # wait for the copy to finish T.ptx.mbarrier.try_wait(T.address_of(bar), phase) diff --git a/tests/python/tirx/iket/test_iket_orchestration.py b/tests/python/tirx/iket/test_iket_orchestration.py index 5e66c6b91479..d72446716f4b 100644 --- a/tests/python/tirx/iket/test_iket_orchestration.py +++ b/tests/python/tirx/iket/test_iket_orchestration.py @@ -54,6 +54,7 @@ def run_process(argv, *, cwd, env, timeout): def test_public_namespace_and_old_bench_path_is_absent(): + pytest.importorskip("triton") # tvm.tirx.bench pulls in the triton harness import tvm.tirx.bench as bench assert iket.__all__ == [ diff --git a/tests/python/tirx/iket/test_iket_profiler.py b/tests/python/tirx/iket/test_iket_profiler.py index 2771b1c8bbbf..956f7f34bc58 100644 --- a/tests/python/tirx/iket/test_iket_profiler.py +++ b/tests/python/tirx/iket/test_iket_profiler.py @@ -32,6 +32,7 @@ import tvm from tvm.backend.cuda import transforms as cuda_transforms from tvm.script import tirx as T +from tvm.testing import env from tvm.tirx.cuda.iket import IketProfiler TARGET = tvm.target.Target({"kind": "cuda", "arch": "sm_100a"}) @@ -474,6 +475,7 @@ def test_declaration_module_and_architecture_boundaries(): _compile(serial_a, target=tvm.target.Target({"kind": "cuda", "arch": "sm_80"})) +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") def test_multi_kernel_module_has_no_tvm_control_plane(): executable = IketProfiler().compile( tvm.IRModule( diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py index f83486bd4c76..a80e34e8ac28 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py @@ -555,8 +555,8 @@ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: ex = tvm.compile(mod, target=target, tir_pipeline="tirx") src = ex.mod.imports[0].inspect_source() - bitsel = re.findall(r"& 1\) \* v_\d+\[", src) - v_decls = re.findall(r"alignas\(\d+\) int v_\d+\[(\d+)\]", src) + bitsel = re.findall(r"& 1\) \* \w+\[", src) + v_decls = re.findall(r"alignas\(\d+\) int \w+\[(\d+)\]", src) assert bitsel, ( "expected fast-path ``(bit & 1) * v_[i]`` adds; if missing, " "var_bounds wiring may have regressed" diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py index 64e66cd5e089..314753d6c126 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py @@ -453,10 +453,10 @@ def test_ldstmatrix_swizzle_multi_iter_pow2(): assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src # Fast-path fingerprint: 3-slot signed_strides + bit-select uses. - assert re.search(r"alignas\(\d+\) int v_\d+\[3\]", src), ( + assert re.search(r"alignas\(\d+\) int \w+\[3\]", src), ( "expected 3-slot signed_strides buffer for bjs [7, 6, 2]" ) - bitsel = re.findall(r"& 1\) \* v_\d+\[", src) + bitsel = re.findall(r"& 1\) \* \w+\[", src) assert bitsel, "fast-path bit-select pattern '& 1) * v_[' missing" n_elem = 1 @@ -529,10 +529,10 @@ def test_ldstmatrix_swizzle_multi_iter_linear(): assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src # Fast-path fingerprint: 1-slot signed_strides (just the inner BitIter). - assert re.search(r"alignas\(\d+\) int v_\d+\[1\]", src), ( + assert re.search(r"alignas\(\d+\) int \w+\[1\]", src), ( "expected 1-slot signed_strides buffer (only the inner Case-1.A bj=2)" ) - bitsel = re.findall(r"& 1\) \* v_\d+\[", src) + bitsel = re.findall(r"& 1\) \* \w+\[", src) assert bitsel, "fast-path bit-select pattern missing" n_elem = 1 diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py index 4b7ecc26d2db..e07dab5b441c 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py @@ -471,7 +471,7 @@ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: # such bit-select * signed-stride pattern. import re - bitsel_pattern = re.findall(r"& 1\) \* v_\d+\[", src) + bitsel_pattern = re.findall(r"& 1\) \* \w+\[", src) assert bitsel_pattern, ( "fast-path bit-select pattern '& 1) * v_[' not found; " "looks like emit_iter_offset's fast path didn't fire." diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py index de914da54367..657f3a0e31fc 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py @@ -378,7 +378,7 @@ def test_cp_shape_config_routes_to_generic_planner(): # Generic-planner signature: descriptor template encoded at address 0... assert "reinterpret_cast((uint64_t)0)" in src # ...and patched per cp via the 0x3FFF address-field mask. - assert src.count("cp_desc[0] &") == 4, f"expected 4 patched cps; src=\n{src}" + assert src.count("cp_desc_ptr[0] &") == 4, f"expected 4 patched cps; src=\n{src}" @pytest.mark.gpu @@ -677,7 +677,7 @@ def kernel(A_ptr: T.handle): assert t_tok in cp_lines[i], f"cp[{i}] tmem col: want {t_tok!r} in {cp_lines[i]!r}" s_tok = f"+ {s_off}))" assert s_tok in cp_lines[i], f"cp[{i}] smem byte off: want {s_tok!r} in {cp_lines[i]!r}" - assert "cp_desc[0] &" in cp_lines[i] and "16383" in cp_lines[i], cp_lines[i] + assert "cp_desc_ptr[0] &" in cp_lines[i] and "16383" in cp_lines[i], cp_lines[i] # Negative tests: readable ValueErrors from the planner @@ -1296,9 +1296,9 @@ def test_multi_cp_encodes_descriptor_once_and_patches_addr(): # Descriptor encoded once (single matrix-descriptor encode call), then # reused with a per-cp 14-bit SMEM address patch (0x3FFF == 16383 mask). assert "16383" in src, "expected 14-bit SMEM address-field patch (0x3FFF mask)" - assert src.count("cp_desc[0] &") == 4, ( + assert src.count("cp_desc_ptr[0] &") == 4, ( f"expected 4 address-patched cp's reusing one cp_desc; got " - f"{src.count('cp_desc[0] &')}\nsrc=\n{src}" + f"{src.count('cp_desc_ptr[0] &')}\nsrc=\n{src}" ) diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index f2431ab15fdf..318ffba03230 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py @@ -1496,6 +1496,7 @@ def rank_change(A_ptr: T.handle): T.cta_id([1]) tid = T.thread_id([1]) dyn = T.alloc_buffer((65,), "uint64", scope="shared.dyn") + T.attr({"tirx.dyn_smem_bytes": 65 * 8}) A_smem = T.decl_buffer((64,), "float16", dyn.data, layout=T.TileLayout(T.S[64])) mbar = T.decl_buffer((1,), "uint64", dyn.data, elem_offset=16) if tid == 0: @@ -1522,6 +1523,7 @@ def selector_gather( T.cta_id([1]) tid = T.thread_id([128]) dyn = T.alloc_buffer((520,), "uint64", scope="shared.dyn") + T.attr({"tirx.dyn_smem_bytes": 520 * 8}) A_smem = T.decl_buffer( (4, 64), "bfloat16", dyn.data, layout=T.TileLayout(T.S[4, 64]) ) @@ -1636,7 +1638,7 @@ def test_prefetch_only_main_selector_descriptor(): lowered = _lower_module(func) collector = _PrefetchCollector() collector.visit_stmt(lowered["main"].body) - assert collector.names == ["A_ptr_tensormap"] + assert collector.names == ["A_tensormap"] def _build_sparse_decode_qo_tma_regression(): @@ -1684,6 +1686,7 @@ def kernel( T.cta_id([1]) tid = T.thread_id([128]) dyn = T.alloc_buffer((shared_bytes + 8,), "uint8", scope="shared.dyn") + T.attr({"tirx.dyn_smem_bytes": shared_bytes + 8}) q_smem = T.decl_buffer( (64, 512), "bfloat16", dyn.data, scope="shared.dyn", layout=q_layout ) @@ -1825,7 +1828,7 @@ def test_sparse_decode_absent_extra_map_has_no_dummy_descriptor_or_select(): assert selects.nodes == [] prefetches = _PrefetchCollector() prefetches.visit_stmt(lowered["main"].body) - assert prefetches.names == ["A_ptr_tensormap"] + assert prefetches.names == ["A_tensormap"] @pytest.mark.parametrize( @@ -1959,6 +1962,7 @@ def kernel( T.cta_id([1]) tid = T.thread_id([128]) dyn = T.alloc_buffer((shared_bytes + 64,), "uint8", scope="shared.dyn") + T.attr({"tirx.dyn_smem_bytes": shared_bytes + 64}) A_smem = T.decl_buffer( (4, cols), dtype, dyn.data, layout=T.TileLayout(T.S[4, cols]) ) diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index d99f70924747..9a8a811294c2 100644 --- a/tests/python/tirx/test_parser_printer.py +++ b/tests/python/tirx/test_parser_printer.py @@ -1560,7 +1560,7 @@ def func() -> None: a_buf = bufs["A"] b_buf = bufs["B"] - assert b_buf.data.same_as(a_buf.data) + assert_structural_equal(_collect_buffer_sources(func)["B"], a_buf.data) assert [int(s) for s in b_buf.shape] == [4, 4, 4, 64] expected = tvm.tirx.layout.ComposeLayout( a_buf.layout.per_element, diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index 273f51544455..8d602cce3560 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -1209,7 +1209,7 @@ def selector( post_order_visit(selector.body, seen.append) assert any(isinstance(node, Var) and node.same_as(selector.params[-1]) for node in seen) undefined = tvm.tirx.analysis.undefined_vars(op_call) - assert any(var.same_as(original_b.data) for var in undefined) + assert any(var.same_as(original_b) for var in undefined) replacement = Var("replacement", "int32") updated = substitute(op_call, {selector.params[-1]: replacement}) diff --git a/tests/python/tirx/transform/test_transform_flatten_buffer.py b/tests/python/tirx/transform/test_transform_flatten_buffer.py new file mode 100644 index 000000000000..548640cdad00 --- /dev/null +++ b/tests/python/tirx/transform/test_transform_flatten_buffer.py @@ -0,0 +1,159 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""FlattenBuffer must keep buffer identity coherent across rebuilds. + +Rebuilding a buffer mints a fresh typed variable, so every remaining +reference — including loads embedded in another buffer's type fields and +loads spliced into indices by the elem_offset fold — must be remapped to +the rebuilt identity, or SplitHostDevice later sees them as undefined +and hoists dead variables into the kernel ABI. +""" + +import tvm +import tvm.testing +from tvm.script import tirx as T +from tvm.tirx.transform import FlattenBuffer + + +def _collect_defined_buffers(func): + defined = set() + + def visit(node): + if isinstance(node, tvm.tirx.AllocBuffer | tvm.tirx.DeclBuffer): + defined.add(node.buffer) + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + return defined + + +def _assert_loads_reference_defined_buffers(func): + """Every BufferLoad — direct, or embedded in a buffer's type fields — + must reference a buffer defined by an AllocBuffer/DeclBuffer in the + function.""" + defined = _collect_defined_buffers(func) + + def is_defined(buf): + return any(buf.same_as(d) for d in defined) + + stale = [] + + def check_expr(expr, where): + def visit(node): + if isinstance(node, tvm.tirx.BufferLoad) and not is_defined(node.buffer): + stale.append(f"{where}: load of {node.buffer.name}") + + tvm.tirx.stmt_functor.post_order_visit(expr, visit) + + def visit(node): + if isinstance(node, tvm.tirx.BufferLoad | tvm.tirx.BufferStore): + if not is_defined(node.buffer): + stale.append(f"access of {node.buffer.name}") + for index in node.indices: + check_expr(index, f"index of {node.buffer.name}") + if isinstance(node, tvm.tirx.AllocBuffer | tvm.tirx.DeclBuffer): + for extent in node.buffer.shape: + check_expr(extent, f"shape of {node.buffer.name}") + if node.buffer.elem_offset is not None: + check_expr(node.buffer.elem_offset, f"elem_offset of {node.buffer.name}") + + tvm.tirx.stmt_functor.post_order_visit(func.body, visit) + assert not stale, f"stale buffer references after FlattenBuffer: {stale}" + + +def _flatten(func): + mod = tvm.IRModule({"main": func}) + with tvm.target.Target("cuda"): + return next(iter(FlattenBuffer()(mod).functions_items()))[1] + + +def test_flatten_remaps_loads_in_view_shape(): + """A view sized by a local scalar: the scalar's rebuild must reach the + load embedded in the view's shape.""" + + @T.prim_func(private=True) + def before(): + n = T.alloc_local([1], "int32") + n[0] = 8 + data = T.alloc_buffer([64], "float16", scope="shared") + view = T.decl_buffer((n[0],), "float16", data.data, scope="shared") + view[0] = T.float16(0) + + _assert_loads_reference_defined_buffers(_flatten(before)) + + +def test_flatten_remaps_loads_in_folded_elem_offset(): + """A view offset by a local scalar: the elem_offset fold splices the + scalar load into every access index; those spliced loads must follow + the scalar's rebuild.""" + + @T.prim_func(private=True) + def before(): + n = T.alloc_local([1], "int32") + n[0] = 4 + base = T.alloc_buffer([128], "uint64", scope="shared") + mbar = T.decl_buffer((1,), "uint64", base.data, elem_offset=n[0], scope="shared") + mbar[0] = T.uint64(1) + + after = _flatten(before) + _assert_loads_reference_defined_buffers(after) + + # The fold must actually have spliced the offset into the store index. + found = [] + + def visit(node): + if isinstance(node, tvm.tirx.BufferStore) and node.buffer.name.startswith("mbar"): + + def inner(sub): + if isinstance(sub, tvm.tirx.BufferLoad): + found.append(sub) + + tvm.tirx.stmt_functor.post_order_visit(node.indices[0], inner) + + tvm.tirx.stmt_functor.post_order_visit(after.body, visit) + assert found, "expected the folded elem_offset load in the mbar store index" + + +def test_flatten_keeps_identity_of_already_flat_buffers(): + """A flat buffer whose type is unchanged by flattening must keep its + identity (no gratuitous rebuild).""" + + @T.prim_func(private=True) + def before(): + flat = T.alloc_buffer([32], "float32", scope="shared", layout=None) + flat[0] = T.float32(0) + + before_allocs = {} + + def collect_before(node): + if isinstance(node, tvm.tirx.AllocBuffer): + before_allocs[node.buffer.name] = node.buffer + + tvm.tirx.stmt_functor.post_order_visit(before.body, collect_before) + + after = _flatten(before) + preserved = [] + + def visit(node): + if isinstance(node, tvm.tirx.AllocBuffer) and node.buffer.name in before_allocs: + preserved.append(node.buffer.same_as(before_allocs[node.buffer.name])) + + tvm.tirx.stmt_functor.post_order_visit(after.body, visit) + assert preserved and all(preserved), "already-flat buffer identity was not preserved" + + +if __name__ == "__main__": + tvm.testing.main()