diff --git a/.agents/scripts/monitor_gpu.sh b/.agents/scripts/monitor_gpu.sh index e1d91ae3b49c..85963da93089 100755 --- a/.agents/scripts/monitor_gpu.sh +++ b/.agents/scripts/monitor_gpu.sh @@ -23,18 +23,7 @@ while [[ $# -gt 0 ]]; do --interval) INTERVAL="$2"; shift 2 ;; --log) LOG="$2"; shift 2 ;; -h|--help) - cat <<'EOF' -Watch a single GPU for foreign processes (anyone other than the current -user) appearing during a long-running test. Intended companion to -`/tir-test`: leave this running in a side terminal while pytest runs, and -it will alert if someone else lands on the same GPU. - -Usage: - monitor_gpu.sh # uses $CUDA_VISIBLE_DEVICES, defaults to 0 - monitor_gpu.sh --gpu 3 # watch GPU 3 - monitor_gpu.sh --gpu 3 --interval 2 # poll every 2 seconds - monitor_gpu.sh --log /tmp/gpu.log # also tee to a log file -EOF + sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//' exit 0 ;; *) echo "unknown arg: $1" >&2; exit 2 ;; esac diff --git a/.agents/skills/tir-bench/SKILL.md b/.agents/skills/tir-bench/SKILL.md deleted file mode 100644 index 515863829bd6..000000000000 --- a/.agents/skills/tir-bench/SKILL.md +++ /dev/null @@ -1,195 +0,0 @@ -Run kernel performance benchmarks to verify codegen changes. - -## Kernels to benchmark - -All commands use `--warmup 100 --repeat 30` for ~3-minute total runtime with reliable medians. Drop to defaults only when chasing a sub-2% regression. - -- **GEMM**: square GEMM at M=N=K in {1024, 2048, 4096, 8192, 16384} for three variants: - - fp16: `python -m tirx_kernels.bench --kernel fp16_bf16_gemm --warmup 100 --repeat 30` - - fp8: `python -m tirx_kernels.bench --kernel fp8_blockwise_gemm --warmup 100 --repeat 30` - - nvfp4: `python -m tirx_kernels.bench --kernel nvfp4_gemm --warmup 100 --repeat 30` -- **FA4** (flash_attention4): all registered configs - - `python -m tirx_kernels.bench --kernel flash_attention4 --warmup 100 --repeat 30` -- **MQA logits** (fp8 / fp4): all registered configs - - `python -m tirx_kernels.bench --kernel deepgemm_sm100_fp8_mqa_logits --warmup 100 --repeat 30` - - `python -m tirx_kernels.bench --kernel deepgemm_sm100_fp4_mqa_logits --warmup 100 --repeat 30` - -## Steps - -1. Select the least busy GPU: - ```bash - 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 ' ') - ``` - -2. Run benchmarks for each kernel using the commands above. - -3. Present results in a table: kernel x config, with times in ms. - -## When to use - -When modifying anything that affects code generation: kernels, op dispatches, lowering passes, codegen, device ops. - -## Reference baseline - -Captured 2026-05-17 on B200 (sm_100a), GPU 7, `warmup=100 repeat=30`, `timer=proton`. - -- `tir` @ `587f439c4c` (branch `scope-id`, with `feat(exec-scope): infer scope_id extent from sibling defs when omitted` on top of upstream tirx `c9ee147baf`) -- `tirx-kernels` @ `fdab8ac5` (branch `scope-id`, with `perf(kernel): hoist mqa_fp8 warpgroup index` on top of upstream `ae8673c9`) - -All times in us. `baseline/tirx` > 1 means TIRX faster. - -### `fp16_bf16_gemm` (baseline=`torch-cublas`) - - -| config | torch-cublas | tir | baseline/tirx | -|---|---:|---:|---:| -| `fp16_1024x1024x1024` | 5.73us | 16.54us | 0.347 | -| `fp16_2048x2048x2048` | 16.40us | 27.91us | 0.588 | -| `fp16_4096x4096x4096` | 95.19us | 94.34us | 1.009 | -| `fp16_8192x8192x8192` | 823.15us | 843.04us | 0.976 | -| `fp16_16384x16384x16384` | 6093.33us | 6128.95us | 0.994 | -| `bf16_1024x1024x1024` | 5.72us | 16.51us | 0.347 | -| `bf16_2048x2048x2048` | 16.13us | 27.77us | 0.581 | -| `bf16_4096x4096x4096` | 92.25us | 91.35us | 1.010 | -| `bf16_8192x8192x8192` | 756.17us | 781.91us | 0.967 | -| `bf16_16384x16384x16384` | 5823.27us | 5809.98us | 1.002 | - -### `fp8_blockwise_gemm` (baseline=`deepgemm`) - - -| config | deepgemm | tir | baseline/tirx | -|---|---:|---:|---:| -| `smoke_1024x1024x1024` | 6.07us | 5.91us | 1.026 | -| `deepgemm_m4096_n2112_k7168` | 49.86us | 48.96us | 1.018 | -| `deepgemm_m4096_n576_k7168` | 19.12us | 18.84us | 1.015 | -| `deepgemm_m4096_n24576_k1536` | 116.18us | 115.68us | 1.004 | -| `deepgemm_m4096_n32768_k512` | 75.54us | 71.28us | 1.060 | -| `deepgemm_m4096_n7168_k16384` | 320.22us | 329.80us | 0.971 | -| `deepgemm_m4096_n4096_k7168` | 83.19us | 82.69us | 1.006 | -| `deepgemm_m4096_n7168_k2048` | 44.04us | 43.59us | 1.010 | -| `stress_m8192_n7168_k4096` | 159.30us | 159.99us | 0.996 | - -### `nvfp4_gemm` (baseline=`flashinfer`) - - -| config | flashinfer | tir | baseline/tirx | -|---|---:|---:|---:| -| `1024x1024x1024` | 5.13us | 6.59us | 0.778 | -| `2048x2048x2048` | 8.39us | 8.84us | 0.950 | -| `4096x4096x4096` | 32.50us | 30.56us | 1.064 | -| `8192x8192x8192` | 199.24us | 186.39us | 1.069 | -| `16384x16384x16384` | 2128.05us | 1511.81us | 1.408 | - -### `flash_attention4` (baseline=`flashattn_sm100`) - - -| config | flashattn_sm100 | tir | baseline/tirx | -|---|---:|---:|---:| -| `s1024_h32kv4` | 20.34us | 20.80us | 0.978 | -| `s1024_h32kv4_causal` | 19.85us | 19.66us | 1.009 | -| `s1024_h32kv8` | 20.50us | 20.91us | 0.980 | -| `s1024_h32kv8_causal` | 19.85us | 19.75us | 1.005 | -| `s1024_h32kv16` | 20.51us | 21.05us | 0.974 | -| `s1024_h32kv16_causal` | 20.24us | 20.68us | 0.979 | -| `s1024_h32kv32` | 20.75us | 21.18us | 0.980 | -| `s1024_h32kv32_causal` | 21.07us | 22.24us | 0.947 | -| `s2048_h32kv4` | 59.47us | 60.85us | 0.977 | -| `s2048_h32kv4_causal` | 39.40us | 37.51us | 1.050 | -| `s2048_h32kv8` | 60.23us | 61.84us | 0.974 | -| `s2048_h32kv8_causal` | 39.49us | 37.76us | 1.046 | -| `s2048_h32kv16` | 60.60us | 62.83us | 0.965 | -| `s2048_h32kv16_causal` | 39.94us | 38.57us | 1.036 | -| `s2048_h32kv32` | 61.59us | 63.62us | 0.968 | -| `s2048_h32kv32_causal` | 40.29us | 42.38us | 0.951 | -| `s4096_h32kv4` | 203.59us | 204.89us | 0.994 | -| `s4096_h32kv4_causal` | 114.98us | 111.69us | 1.029 | -| `s4096_h32kv8` | 204.46us | 207.67us | 0.985 | -| `s4096_h32kv8_causal` | 116.24us | 112.45us | 1.034 | -| `s4096_h32kv16` | 208.31us | 211.63us | 0.984 | -| `s4096_h32kv16_causal` | 117.59us | 113.66us | 1.035 | -| `s4096_h32kv32` | 211.75us | 216.02us | 0.980 | -| `s4096_h32kv32_causal` | 118.98us | 122.09us | 0.975 | -| `s8192_h32kv4` | 816.39us | 818.33us | 0.998 | -| `s8192_h32kv4_causal` | 429.56us | 420.64us | 1.021 | -| `s8192_h32kv8` | 795.55us | 852.89us | 0.933 | -| `s8192_h32kv8_causal` | 411.97us | 440.47us | 0.935 | -| `s8192_h32kv16` | 779.83us | 841.29us | 0.927 | -| `s8192_h32kv16_causal` | 412.70us | 399.01us | 1.034 | -| `s8192_h32kv32` | 784.06us | 821.54us | 0.954 | -| `s8192_h32kv32_causal` | 459.55us | 420.57us | 1.093 | - -### `deepgemm_sm100_fp8_mqa_logits` (baseline=`deepgemm`) - - -| config | deepgemm | tirx | baseline/tirx | -|---|---:|---:|---:| -| `s2048_skv4096_h64_d128_f32_dense_cp` | 43.80us | 44.49us | 0.984 | -| `s2048_skv4096_h64_d128_f32_dense_nocp` | 58.50us | 58.59us | 0.999 | -| `s2048_skv8192_h64_d128_f32_dense_cp` | 77.25us | 78.07us | 0.990 | -| `s2048_skv8192_h64_d128_f32_dense_nocp` | 118.40us | 118.97us | 0.995 | -| `s4096_skv4096_h64_d128_f32_dense_cp` | 78.02us | 77.94us | 1.001 | -| `s4096_skv4096_h64_d128_f32_dense_nocp` | 77.89us | 78.37us | 0.994 | -| `s4096_skv8192_h64_d128_f32_dense_cp` | 136.98us | 136.12us | 1.006 | -| `s4096_skv8192_h64_d128_f32_dense_nocp` | 196.36us | 202.57us | 0.969 | -| `s2048_skv4096_h64_d128_f32_compressed_cp` | 46.60us | 44.88us | 1.038 | -| `s2048_skv4096_h64_d128_f32_compressed_nocp` | 61.46us | 59.54us | 1.032 | -| `s2048_skv8192_h64_d128_f32_compressed_cp` | 81.83us | 78.99us | 1.036 | -| `s2048_skv8192_h64_d128_f32_compressed_nocp` | 125.40us | 120.15us | 1.044 | -| `s4096_skv4096_h64_d128_f32_compressed_cp` | 83.89us | 78.42us | 1.070 | -| `s4096_skv4096_h64_d128_f32_compressed_nocp` | 83.94us | 78.89us | 1.064 | -| `s4096_skv8192_h64_d128_f32_compressed_cp` | 147.25us | 137.97us | 1.067 | -| `s4096_skv8192_h64_d128_f32_compressed_nocp` | 209.79us | 196.89us | 1.066 | -| `s2048_skv4096_h64_d128_bf16_dense_cp` | 44.73us | 44.81us | 0.998 | -| `s2048_skv4096_h64_d128_bf16_dense_nocp` | 58.90us | 59.29us | 0.993 | -| `s2048_skv8192_h64_d128_bf16_dense_cp` | 79.48us | 79.03us | 1.006 | -| `s2048_skv8192_h64_d128_bf16_dense_nocp` | 121.27us | 121.16us | 1.001 | -| `s4096_skv4096_h64_d128_bf16_dense_cp` | 78.87us | 78.84us | 1.000 | -| `s4096_skv4096_h64_d128_bf16_dense_nocp` | 79.02us | 78.66us | 1.005 | -| `s4096_skv8192_h64_d128_bf16_dense_cp` | 139.18us | 138.40us | 1.006 | -| `s4096_skv8192_h64_d128_bf16_dense_nocp` | 199.50us | 197.53us | 1.010 | -| `s2048_skv4096_h64_d128_bf16_compressed_cp` | 46.91us | 46.09us | 1.018 | -| `s2048_skv4096_h64_d128_bf16_compressed_nocp` | 61.15us | 60.29us | 1.014 | -| `s2048_skv8192_h64_d128_bf16_compressed_cp` | 82.17us | 80.09us | 1.026 | -| `s2048_skv8192_h64_d128_bf16_compressed_nocp` | 126.02us | 123.97us | 1.017 | -| `s4096_skv4096_h64_d128_bf16_compressed_cp` | 84.10us | 82.16us | 1.024 | -| `s4096_skv4096_h64_d128_bf16_compressed_nocp` | 83.94us | 82.05us | 1.023 | -| `s4096_skv8192_h64_d128_bf16_compressed_cp` | 147.98us | 144.28us | 1.026 | -| `s4096_skv8192_h64_d128_bf16_compressed_nocp` | 209.74us | 204.18us | 1.027 | - -### `deepgemm_sm100_fp4_mqa_logits` (baseline=`deepgemm`) - - -| config | deepgemm | tirx | baseline/tirx | -|---|---:|---:|---:| -| `s2048_skv4096_h64_d128_f32_dense_cp` | 41.25us | 41.52us | 0.994 | -| `s2048_skv4096_h64_d128_f32_dense_nocp` | 53.67us | 54.10us | 0.992 | -| `s2048_skv8192_h64_d128_f32_dense_cp` | 71.99us | 72.44us | 0.994 | -| `s2048_skv8192_h64_d128_f32_dense_nocp` | 111.41us | 111.13us | 1.003 | -| `s4096_skv4096_h64_d128_f32_dense_cp` | 73.25us | 73.47us | 0.997 | -| `s4096_skv4096_h64_d128_f32_dense_nocp` | 73.21us | 73.52us | 0.996 | -| `s4096_skv8192_h64_d128_f32_dense_cp` | 130.21us | 129.54us | 1.005 | -| `s4096_skv8192_h64_d128_f32_dense_nocp` | 186.20us | 184.96us | 1.007 | -| `s2048_skv4096_h64_d128_f32_compressed_cp` | 45.14us | 42.37us | 1.066 | -| `s2048_skv4096_h64_d128_f32_compressed_nocp` | 59.05us | 54.82us | 1.077 | -| `s2048_skv8192_h64_d128_f32_compressed_cp` | 79.09us | 73.69us | 1.073 | -| `s2048_skv8192_h64_d128_f32_compressed_nocp` | 122.95us | 113.08us | 1.087 | -| `s4096_skv4096_h64_d128_f32_compressed_cp` | 80.41us | 73.88us | 1.088 | -| `s4096_skv4096_h64_d128_f32_compressed_nocp` | 80.32us | 73.81us | 1.088 | -| `s4096_skv8192_h64_d128_f32_compressed_cp` | 144.14us | 131.25us | 1.098 | -| `s4096_skv8192_h64_d128_f32_compressed_nocp` | 206.26us | 187.68us | 1.099 | -| `s2048_skv4096_h64_d128_bf16_dense_cp` | 42.24us | 42.51us | 0.994 | -| `s2048_skv4096_h64_d128_bf16_dense_nocp` | 55.24us | 55.44us | 0.996 | -| `s2048_skv8192_h64_d128_bf16_dense_cp` | 74.32us | 74.16us | 1.002 | -| `s2048_skv8192_h64_d128_bf16_dense_nocp` | 114.28us | 113.84us | 1.004 | -| `s4096_skv4096_h64_d128_bf16_dense_cp` | 74.91us | 74.90us | 1.000 | -| `s4096_skv4096_h64_d128_bf16_dense_nocp` | 74.90us | 74.84us | 1.001 | -| `s4096_skv8192_h64_d128_bf16_dense_cp` | 133.11us | 132.55us | 1.004 | -| `s4096_skv8192_h64_d128_bf16_dense_nocp` | 190.79us | 189.49us | 1.007 | -| `s2048_skv4096_h64_d128_bf16_compressed_cp` | 44.99us | 45.73us | 0.984 | -| `s2048_skv4096_h64_d128_bf16_compressed_nocp` | 59.06us | 60.01us | 0.984 | -| `s2048_skv8192_h64_d128_bf16_compressed_cp` | 79.27us | 80.35us | 0.987 | -| `s2048_skv8192_h64_d128_bf16_compressed_nocp` | 122.57us | 123.86us | 0.990 | -| `s4096_skv4096_h64_d128_bf16_compressed_cp` | 79.93us | 81.00us | 0.987 | -| `s4096_skv4096_h64_d128_bf16_compressed_nocp` | 79.78us | 80.97us | 0.985 | -| `s4096_skv8192_h64_d128_bf16_compressed_cp` | 142.89us | 144.28us | 0.990 | -| `s4096_skv8192_h64_d128_bf16_compressed_nocp` | 204.95us | 206.88us | 0.991 | diff --git a/.agents/skills/tir-test/SKILL.md b/.agents/skills/tir-test/SKILL.md index b4c45069163b..e4a5fe6ca2d9 100644 --- a/.agents/skills/tir-test/SKILL.md +++ b/.agents/skills/tir-test/SKILL.md @@ -2,9 +2,13 @@ Run the full TIRX test suite. ## Steps -1. Select the least busy GPU to avoid conflicts: +1. Install the kernel package and select the least busy GPU: ```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" ``` 2. Start the GPU monitor in the background so we can detect if anyone else lands on the same GPU mid-run: @@ -15,18 +19,29 @@ Run the full TIRX test suite. trap 'kill $MON_PID 2>/dev/null' EXIT ``` -3. Run the full test suite with xdist parallelism: +3. Import gate — bench workloads: fail fast if any kernel listed in `workloads.yaml` fails to import: ```bash - pytest tests/python/tirx/ -n auto + python -m tirx_kernels.bench_suite --check-imports ``` + A non-zero exit means a pinned workload kernel failed to import — fix it before proceeding. -4. Stop the monitor and check for foreign GPU usage during the run: +4. Full kernel import gate (correctness test suite coverage): + ```bash + python -m tirx_kernels.registry --cc 10 --strict + ``` + +5. Run the full test suite with xdist parallelism: + ```bash + pytest tests/python/tirx/ -n 16 + ``` + +6. Stop the monitor and check for foreign GPU usage during the run: ```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" ``` -5. Report results: total passed, failed, skipped, errors. If any foreign-user events are present in step 4, mention them — flaky failures should be re-evaluated on a clean GPU before being attributed to code changes. +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. ## Failure triage rules diff --git a/.gitignore b/.gitignore index 180eea6c4ead..2a59857ea7b0 100644 --- a/.gitignore +++ b/.gitignore @@ -183,7 +183,6 @@ cscope* perf .bash_history *.json -!.claude/commands/tir-bench/baseline.json *.params *.ro *.onnx @@ -295,7 +294,3 @@ python/bin/ python/typing_extensions.py python/*.dist-info/ pytest-of-bohanhou/ - -# tir-bench run artifacts (regenerable; see .claude/commands/tir-bench.md) -.tir-bench/ -.tir-bench-*/ diff --git a/AGENTS.md b/AGENTS.md index 4d70e7c79c43..cd38b2adf3da 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -85,8 +85,8 @@ the branch instead of the whole repository: pre-commit run --files ... ``` -Use `.agents/skills/tir-build`, `.agents/skills/tir-test`, and -`.agents/skills/tir-bench` when their workflows apply. +Use `python -m tirx_kernels.bench_suite` in the **tirx-kernels** repo +(`tirx_kernels/bench_suite/`) when that workflow applies. ## Coding Conventions diff --git a/include/tvm/tirx/async_structs.h b/include/tvm/tirx/async_structs.h deleted file mode 100644 index eb140309cb17..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/predicate.h b/include/tvm/tirx/predicate.h deleted file mode 100644 index f9e7667cfe2f..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 4460b28e6ffa..cf174c86c11a 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 0262a167918d..15f3e2f355bb 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 63% rename from include/tvm/tirx/tirx_op.h rename to include/tvm/tirx/tile_primitive.h index ad2ec8e80fee..460365bc5357 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,72 @@ 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: + 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 +342,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 41618eb1d96c..000000000000 --- a/include/tvm/tirx/tirx_stmt.h +++ /dev/null @@ -1,101 +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: - 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/backend/cuda/lang/pipeline.py b/python/tvm/backend/cuda/lang/pipeline.py index 40fd40c3fac6..9aa44612d7b2 100644 --- a/python/tvm/backend/cuda/lang/pipeline.py +++ b/python/tvm/backend/cuda/lang/pipeline.py @@ -184,18 +184,20 @@ def arrive(self, stage, tx_count=None, cta_id=None, pred=None): class TCGen05Bar(MBarrier): """Barrier signaled by ``tcgen05`` commit. - The caller is responsible for ensuring only one thread issues the - commit, e.g. by wrapping the call in ``if T.ptx.elect_sync():``. + The caller is responsible for ensuring only one thread issues the commit, + either with a surrounding election or by passing ``pred=``. """ @T.inline - def arrive(self, stage, cta_group=1, cta_mask=None): + def arrive(self, stage, cta_group=1, cta_mask=None, pred=None): # NOTE: this arrive() kwarg set intentionally differs from # MBarrier.arrive (hardware necessity, LSP-incompatible by design). if cta_mask is None and cta_group == 1: - T.ptx.tcgen05.commit(self.buf.ptr_to([stage])) + T.ptx.tcgen05.commit(self.buf.ptr_to([stage]), pred=pred) else: - T.ptx.tcgen05.commit(self.buf.ptr_to([stage]), cta_group=cta_group, cta_mask=cta_mask) + T.ptx.tcgen05.commit( + self.buf.ptr_to([stage]), cta_group=cta_group, cta_mask=cta_mask, pred=pred + ) # Barrier-type tags accepted by Pipeline's ``full=`` / ``empty=`` arguments. diff --git a/python/tvm/backend/cuda/lang/tile_scheduler.py b/python/tvm/backend/cuda/lang/tile_scheduler.py index c6154f2462f6..d35bb9e5284e 100644 --- a/python/tvm/backend/cuda/lang/tile_scheduler.py +++ b/python/tvm/backend/cuda/lang/tile_scheduler.py @@ -211,6 +211,7 @@ def __init__( serpentine: bool = False, ): super().__init__(prefix) + serpentine = getattr(serpentine, "value", serpentine) self._num_m_tiles = num_m_tiles self._num_n_tiles = num_n_tiles self._num_clusters = num_clusters @@ -304,17 +305,13 @@ def _gm_emit_zero(self, set_tile_coords): @T.inline def _gm_emit_full_only(self, tile_linear, set_tile_coords): - FULL_GROUPS = T.meta_var(self._FULL_GROUPS) GROUP_SIZE = T.meta_var(self._l2_group_size) GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS) - if (FULL_GROUPS > 0) & (tile_linear < FULL_GROUPS * GROUP_SPAN): - group_id: T.let = tile_linear // GROUP_SPAN - within_group: T.let = tile_linear % GROUP_SPAN - tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE) - tile_col: T.let = within_group // GROUP_SIZE - set_tile_coords(tile_row, tile_col) - else: - set_tile_coords(0, 0) + group_id: T.let = tile_linear // GROUP_SPAN + within_group: T.let = tile_linear % GROUP_SPAN + tile_row: T.let = group_id * GROUP_SIZE + (within_group % GROUP_SIZE) + tile_col: T.let = within_group // GROUP_SIZE + set_tile_coords(tile_row, tile_col) @T.inline def _gm_emit_tail_only(self, tile_linear, set_tile_coords): @@ -322,13 +319,10 @@ def _gm_emit_tail_only(self, tile_linear, set_tile_coords): TAIL_ROWS = T.meta_var(self._TAIL_ROWS) GROUP_SIZE = T.meta_var(self._l2_group_size) GROUP_SPAN = T.meta_var(self._l2_group_size * self._N_TILE_COLS) - if TAIL_ROWS > 0: - rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN - tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS) - tile_col: T.let = rem // TAIL_ROWS - set_tile_coords(tile_row, tile_col) - else: - set_tile_coords(0, 0) + rem: T.let = tile_linear - FULL_GROUPS * GROUP_SPAN + tile_row: T.let = FULL_GROUPS * GROUP_SIZE + (rem % TAIL_ROWS) + tile_col: T.let = rem // TAIL_ROWS + set_tile_coords(tile_row, tile_col) @T.inline def _gm_emit_full_and_tail(self, tile_linear, set_tile_coords): @@ -370,7 +364,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) diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py index 5755644d64ee..351df04b247e 100644 --- a/python/tvm/backend/cuda/op.py +++ b/python/tvm/backend/cuda/op.py @@ -159,57 +159,6 @@ def cuda_cta_min(value, num_warps, scratch): return cuda_cta_reduce(value, "min", num_warps, scratch) -def cuda_copy_bytes(dst, src, num_bytes): - """Typed load/store copy of ``num_bytes`` bytes. - - Copies ``num_bytes`` bytes from ``src`` to ``dst`` using a single - typed load/store instruction. Codegen selects the appropriate C++ - vector type (``uint4``, ``uint2``, ``unsigned int``, etc.). - - Parameters - ---------- - dst : Var - Destination pointer. - - src : Var - Source pointer. - - num_bytes : int - Number of bytes to copy. Must be one of {1, 2, 4, 8, 16}. - - Returns - ------- - call : PrimExpr - A void call expression. - """ - return call_intrin("void", "tirx.cuda.copy_bytes", dst, src, num_bytes) - - -def cuda_copy_128b(dst, src): - """Convenience wrapper: ``cuda_copy_bytes(dst, src, 16)`` — copies 128 bits.""" - return cuda_copy_bytes(dst, src, 16) - - -def cuda_copy_64b(dst, src): - """Convenience wrapper: ``cuda_copy_bytes(dst, src, 8)`` — copies 64 bits.""" - return cuda_copy_bytes(dst, src, 8) - - -def cuda_copy_32b(dst, src): - """Convenience wrapper: ``cuda_copy_bytes(dst, src, 4)`` — copies 32 bits.""" - return cuda_copy_bytes(dst, src, 4) - - -def cuda_copy_16b(dst, src): - """Convenience wrapper: ``cuda_copy_bytes(dst, src, 2)`` — copies 16 bits.""" - return cuda_copy_bytes(dst, src, 2) - - -def cuda_copy_8b(dst, src): - """Convenience wrapper: ``cuda_copy_bytes(dst, src, 1)`` — copies 8 bits.""" - return cuda_copy_bytes(dst, src, 1) - - def cuda_warp_sync(): """TVM intrinsic to synchronize threads within the current warp. @@ -2427,6 +2376,7 @@ def ptx_tcgen05_mma_block_scale( use_a_tmem, cta_group, enable_input_d=1, + pred=None, ): """TVM intrinsic to call tcgen05.mma.cta_group.kind.block_scale Performs matrix multiplication with block scaling: @@ -2477,12 +2427,14 @@ def ptx_tcgen05_mma_block_scale( enable_input_d : PrimExpr Scale operand for the input accumulator C/D. Zero means D = A*B, non-zero means D = A*B + D. + + pred : Optional[PrimExpr] + Runtime ``uint32`` instruction-level predicate. When given, emit + ``@p_issue tcgen05.mma...`` instead of a divergent C branch. """ _choice("cta_group", cta_group, _TCGEN05_CTA_GROUP) - return call_intrin( - "", - "tirx.ptx.tcgen05_mma_block_scale", + args = [ d_dtype, a_dtype, b_dtype, @@ -2497,7 +2449,10 @@ def ptx_tcgen05_mma_block_scale( use_a_tmem, cta_group, enable_input_d, - ) + ] + if pred is not None: + args.append(pred) + return call_intrin("", "tirx.ptx.tcgen05_mma_block_scale", *args) def ptx_tcgen05_mma_sp( @@ -3492,11 +3447,37 @@ def ptx_griddepcontrol_launch_dependents(): _PTX_LD_SCOPE = {"cta", "cluster", "gpu", "sys"} _PTX_LD_SPACE = {"global", "shared", "shared::cta", "shared::cluster", "local"} _PTX_LD_VOLATILE_SPACE = _PTX_LD_SPACE | {"const"} -_PTX_LD_TYPE = {"b32", "u32", "u64", "s32", "f32"} +_PTX_SCALAR_TYPE = { + "b8", + "u8", + "s8", + "b16", + "u16", + "s16", + "b32", + "b64", + "u32", + "u64", + "s32", + "s64", + "f32", + "f64", +} +_PTX_LD_TYPE = set(_PTX_SCALAR_TYPE) _PTX_LD_COP = {"", "ca", "cg", "cs", "lu", "cv"} +_PTX_LD_VEC = {"", "v2", "v4", "v8"} +_PTX_L1_EVICT = { + "", + "L1::evict_normal", + "L1::evict_unchanged", + "L1::evict_first", + "L1::evict_last", + "L1::no_allocate", +} +_PTX_L2_EVICT = {"", "L2::evict_normal", "L2::evict_first", "L2::evict_last"} +_PTX_PREFETCH = {"", "L2::64B", "L2::128B", "L2::256B"} _PTX_MEM_SCOPE = {"", "cta", "cluster", "gpu", "sys"} _PTX_MEM_SPACE = {"global", "shared", "shared::cta", "shared::cluster"} -_PTX_SCALAR_TYPE = {"b32", "b64", "u32", "u64", "s32", "s64", "f32", "f64"} _PTX_RED_OP = {"and", "or", "xor", "add", "inc", "dec", "min", "max"} _PTX_ATOM_OP = {"and", "or", "xor", "exch", "add", "inc", "dec", "min", "max"} _PTX_ST_VEC = {"", "v2", "v4", "v8"} @@ -3532,41 +3513,62 @@ def _resolve_cache_policy(cache_hint, cache_policy, choices=_CP_ASYNC_BULK_CACHE return const(0, dtype="uint64"), False -def ptx_ld_acquire(addr, return_type, ptx_type, *, scope="gpu", space="global"): - """TVM intrinsic for scalar PTX ``ld.acquire.scope{.ss}.type`` loads. - - This wrapper covers the scalar no-cache-policy/no-vector instances of the - PTX ISA ``ld.acquire`` form. ``scope``, state ``space``, PTX ``type`` and - TVM ``return_type`` are explicit so callers can request either raw-bit or - typed loads. - - Parameters - ---------- - addr : PrimExpr - The memory address to load. - - return_type : str - TVM dtype returned by the load. - - ptx_type : str - PTX type suffix such as ``"b32"``, ``"u64"``, or ``"s32"``. - - scope : str - PTX memory scope: ``"cta"``, ``"cluster"``, ``"gpu"``, or ``"sys"``. - - space : str - PTX state space suffix. - - Returns - ------- - call : PrimExpr - The loaded value. - """ +def ptx_ld_acquire( + addr, + return_type, + ptx_type, + *, + scope="gpu", + space="global", + vec="", + dst=None, + cache_hint="", + cache_policy=None, + l1_evict="", + l2_evict="", + prefetch_size="", +): + """TVM intrinsic for PTX ``ld.acquire.scope{.ss}...`` loads.""" _choice("scope", scope, _PTX_LD_SCOPE) _choice("space", space, _PTX_LD_SPACE) _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") + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld_acquire", + dst, + addr, + cache_policy, + return_type, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) return call_intrin( - return_type, "tirx.ptx.ld_acquire", addr, return_type, ptx_type, scope, space + return_type, + "tirx.ptx.ld_acquire", + addr, + return_type, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + 0, + l1_evict, + l2_evict, + prefetch_size, ) @@ -3575,21 +3577,45 @@ def ptx_ld( return_type, ptx_type, *, + dst=None, weak=False, space="global", cop="", + vec="", cache_hint="", cache_policy=None, + l1_evict="", + l2_evict="", + prefetch_size="", ): - """TVM intrinsic for scalar PTX ``ld{.weak}{.ss}{.cop}{.level::cache_hint}.type``. - - This wrapper covers scalar no-prefetch/no-vector instances of the weak - generic load form. - """ + """TVM intrinsic for PTX ``ld{.weak}{.ss}{.cop}...`` loads.""" _choice("space", space, _PTX_LD_SPACE | {"const", "param::entry", "param::func"}) _choice("cop", cop, _PTX_LD_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) + to_dst = int(dst is not None) + if vec and dst is None: + raise ValueError("vec ld requires dst") + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld", + dst, + addr, + cache_policy, + return_type, + int(bool(weak)), + space, + cop, + vec, + ptx_type, + int(has_cache_policy), + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) return call_intrin( return_type, "tirx.ptx.ld", @@ -3599,19 +3625,147 @@ def ptx_ld( int(bool(weak)), space, cop, + "", ptx_type, int(has_cache_policy), + 0, + l1_evict, + l2_evict, + prefetch_size, ) -def ptx_ld_volatile(addr, return_type, ptx_type, *, space="global"): - """TVM intrinsic for scalar PTX ``ld.volatile{.ss}.type`` loads. +def ptx_ld_relaxed( + addr, + return_type, + ptx_type, + *, + scope="gpu", + space="global", + vec="", + dst=None, + cache_hint="", + cache_policy=None, + l1_evict="", + l2_evict="", + prefetch_size="", +): + _choice("scope", scope, _PTX_LD_SCOPE) + _choice("space", space, _PTX_LD_SPACE) + _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") + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld_relaxed", + dst, + addr, + cache_policy, + return_type, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) + return call_intrin( + return_type, + "tirx.ptx.ld_relaxed", + addr, + cache_policy, + return_type, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + 0, + l1_evict, + l2_evict, + prefetch_size, + ) - This wrapper covers scalar no-prefetch/no-vector instances. - """ + +def ptx_ld_volatile( + addr, + return_type, + ptx_type, + *, + space="global", + vec="", + dst=None, + prefetch_size="", +): + """TVM intrinsic for PTX ``ld.volatile{.ss}...`` loads.""" _choice("space", space, _PTX_LD_VOLATILE_SPACE) _choice("ptx_type", ptx_type, _PTX_LD_TYPE) - return call_intrin(return_type, "tirx.ptx.ld_volatile", addr, return_type, ptx_type, space) + _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") + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld_volatile", + dst, + addr, + return_type, + space, + vec, + ptx_type, + to_dst, + prefetch_size, + ) + return call_intrin( + return_type, + "tirx.ptx.ld_volatile", + addr, + return_type, + space, + vec, + ptx_type, + 0, + prefetch_size, + ) + + +def ptx_ld_mmio(addr, return_type, ptx_type, *, sem="acquire", dst=None): + if sem not in ("acquire", "relaxed"): + raise ValueError(f"Unsupported PTX ld.mmio sem {sem!r}") + _choice("ptx_type", ptx_type, _PTX_LD_TYPE) + to_dst = int(dst is not None) + if to_dst: + return call_intrin( + "", + "tirx.ptx.ld_mmio", + dst, + addr, + return_type, + sem, + "sys", + "global", + ptx_type, + 1, + ) + return call_intrin( + return_type, + "tirx.ptx.ld_mmio", + addr, + return_type, + sem, + "sys", + "global", + ptx_type, + 0, + ) def ptx_ld_global_acquire(res, addr): @@ -3706,6 +3860,7 @@ def ptx_atom_scalar( def ptx_st( address, *values, + src=None, weak=False, space="shared", cop="", @@ -3713,17 +3868,15 @@ def ptx_st( ptx_type, cache_hint="", cache_policy=None, + l1_evict="", + l2_evict="", ): _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) cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) - return call_intrin( - "", - "tirx.ptx.st", - address, - *values, + attrs = ( cache_policy, int(bool(weak)), space, @@ -3731,7 +3884,106 @@ def ptx_st( vec, ptx_type, int(has_cache_policy), + l1_evict, + l2_evict, + int(src is not None), ) + if src is not None: + if values: + raise ValueError("ptx_st expects values or src, not both") + return call_intrin("", "tirx.ptx.st", address, src, *attrs) + return call_intrin("", "tirx.ptx.st", address, *values, *attrs) + + +def ptx_st_relaxed( + address, + *values, + src=None, + scope="gpu", + space="global", + vec="", + ptx_type, + cache_hint="", + cache_policy=None, + l1_evict="", + l2_evict="", +): + _choice("scope", scope, _PTX_LD_SCOPE) + _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"}) + _choice("vec", vec, _PTX_ST_VEC) + _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE) + cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) + attrs = ( + cache_policy, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + l1_evict, + l2_evict, + int(src is not None), + ) + if src is not None: + if values: + raise ValueError("ptx_st_relaxed expects values or src, not both") + return call_intrin("", "tirx.ptx.st_relaxed", address, src, *attrs) + return call_intrin("", "tirx.ptx.st_relaxed", address, *values, *attrs) + + +def ptx_st_release( + address, + *values, + src=None, + scope="gpu", + space="global", + vec="", + ptx_type, + cache_hint="", + cache_policy=None, + l1_evict="", + l2_evict="", +): + _choice("scope", scope, _PTX_LD_SCOPE) + _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"}) + _choice("vec", vec, _PTX_ST_VEC) + _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE) + cache_policy, has_cache_policy = _resolve_cache_policy(cache_hint, cache_policy) + attrs = ( + cache_policy, + scope, + space, + vec, + ptx_type, + int(has_cache_policy), + l1_evict, + l2_evict, + int(src is not None), + ) + if src is not None: + if values: + raise ValueError("ptx_st_release expects values or src, not both") + return call_intrin("", "tirx.ptx.st_release", address, src, *attrs) + return call_intrin("", "tirx.ptx.st_release", address, *values, *attrs) + + +def ptx_st_volatile(address, *values, src=None, space="global", vec="", ptx_type): + _choice("space", space, _PTX_MEM_SPACE | {"local", "param::func"}) + _choice("vec", vec, _PTX_ST_VEC) + _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE) + attrs = (space, vec, ptx_type, int(src is not None)) + if src is not None: + if values: + raise ValueError("ptx_st_volatile expects values or src, not both") + return call_intrin("", "tirx.ptx.st_volatile", address, src, *attrs) + return call_intrin("", "tirx.ptx.st_volatile", address, *values, *attrs) + + +def ptx_st_mmio(address, value, *, sem="release", ptx_type): + if sem not in ("release", "relaxed"): + raise ValueError(f"Unsupported PTX st.mmio sem {sem!r}") + _choice("ptx_type", ptx_type, _PTX_SCALAR_TYPE) + return call_intrin("", "tirx.ptx.st_mmio", address, value, sem, "sys", "global", ptx_type) def ptx_st_bulk(ptr, num_bytes, *, weak=False, space="shared::cta"): diff --git a/python/tvm/backend/cuda/operator/intrinsics/memory.py b/python/tvm/backend/cuda/operator/intrinsics/memory.py index 30bb88efe2bc..4b66ed26a25f 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/memory.py +++ b/python/tvm/backend/cuda/operator/intrinsics/memory.py @@ -25,7 +25,6 @@ * ``mapa.u64`` — map a SMEM ptr to a peer CTA's SMEM in the cluster. CUDA side: -* Typed N-byte copy helpers (1/2/4/8/16 bytes via uint{2,4} / unsigned). * ``__ldg`` (cache-as-read-only load). * Templated ``atomicAdd`` / ``atomicCAS``. * half↔float type-punned conversions (single, packed, batch-of-8). @@ -39,39 +38,6 @@ from .registry import CODEGEN_REGISTRY, register_codegen from .utils import parse_str -# ============================================================================= -# Typed N-byte copies — one helper per (1, 2, 4, 8, 16)-byte width. -# Dispatcher picks by ``num_bytes``. -# ============================================================================= -_TYPE_MAP = {16: "uint4", 8: "uint2", 4: "unsigned int", 2: "unsigned short", 1: "unsigned char"} - - -for _num_bytes, _cpp_type in _TYPE_MAP.items(): - device_intrinsic( - f"_cuda_copy_bytes_{_num_bytes}_impl", - helper_name=f"tvm_builtin_copy_{_num_bytes * 8}b", - c_signature="(void* dst_ptr, void* src_ptr)", - body=( - f" {_cpp_type}* src_ = reinterpret_cast<{_cpp_type}*>(src_ptr);\n" - f" {_cpp_type}* dst_ = reinterpret_cast<{_cpp_type}*>(dst_ptr);\n" - " *dst_ = *src_;" - ), - ) -del _num_bytes, _cpp_type - - -@register_codegen("cuda_copy_bytes") -def codegen_cuda_copy_bytes(dst, src, num_bytes): - """Dispatch to the size-specific helper based on ``num_bytes``.""" - num_bytes_int = int(num_bytes) - if num_bytes_int not in _TYPE_MAP: - raise ValueError( - f"Unsupported cuda_copy_bytes num_bytes {num_bytes_int}, " - f"expected one of {sorted(_TYPE_MAP)}" - ) - result = CODEGEN_REGISTRY[f"tirx._cuda_copy_bytes_{num_bytes_int}_impl"]([dst, src]) - return result[0] if isinstance(result, tuple) else result - # ============================================================================= # __ldg — templated read-only cached load; ``T`` resolved at call time from @@ -91,28 +57,87 @@ def codegen_cuda_ldg(addr, dtype): return cuda_func_call(func_name, addr, source_code=source_code, return_type=dtype) +# Shared PTX scalar type metadata (ld/st/red/atom). +_PTX_SCALAR_TYPE_INFO = { + "b8": ("unsigned int", "r", "uint32"), + "u8": ("unsigned int", "r", "uint32"), + "s8": ("int", "r", "int32"), + "b16": ("unsigned short", "h", "uint16"), + "u16": ("unsigned short", "h", "uint16"), + "s16": ("short", "h", "int16"), + "b32": ("unsigned int", "r", "uint32"), + "u32": ("unsigned int", "r", "uint32"), + "s32": ("int", "r", "int32"), + "b64": ("unsigned long long", "l", "uint64"), + "u64": ("unsigned long long", "l", "uint64"), + "s64": ("long long", "l", "int64"), + "f32": ("float", "f", "float32"), + "f64": ("double", "d", "float64"), +} +_PTX_LD_TYPE_RETURNS = { + "b32": {"uint32": "unsigned int", "int32": "int"}, + "b64": {"uint64": "unsigned long long", "int64": "long long"}, +} +_PTX_VEC_STORE_TYPE = { + 16: "uint4", + 8: "uint2", + 4: "unsigned int", + 2: "unsigned short", + 1: "unsigned char", +} + + +def _safe_attr(value): + return parse_str(value).replace("::", "_").replace(".", "_") + + +def _dot(value): + value = parse_str(value) + return f".{value}" if value else "" + + +def _cache_suffix(cache): + return ".L2::cache_hint" if cache else "" + + +def _type_info(ptx_type): + ptx_type = parse_str(ptx_type) + if ptx_type not in _PTX_SCALAR_TYPE_INFO: + raise ValueError( + f"Unsupported PTX scalar type {ptx_type!r}; expected {sorted(_PTX_SCALAR_TYPE_INFO)}" + ) + return (ptx_type, *_PTX_SCALAR_TYPE_INFO[ptx_type]) + + # ============================================================================= -# PTX ld forms: +# 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.acquire.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.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.volatile{.ss}{.level::prefetch_size}{.vec}.type d, [a]; -# -# These are registered from the PTX ISA ld grammar. The current helpers cover -# the scalar no-cache-policy/no-vector instances currently registered. Scope, -# state space, PTX type, and TVM return dtype are explicit instead of being -# inferred from a generic "load" helper. +# 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}; +# ld.mmio.sem.sys{.global}.type d, [a]; # ============================================================================= _PTX_LD_SCOPES = {"cta", "cluster", "gpu", "sys"} _PTX_LD_SPACES = {"global", "shared", "shared::cta", "shared::cluster", "local"} _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_TYPES = { - "b32": {"constraint": "r", "returns": {"uint32": "unsigned int", "int32": "int"}}, - "u32": {"constraint": "r", "returns": {"uint32": "unsigned int"}}, - "u64": {"constraint": "l", "returns": {"uint64": "unsigned long long"}}, - "s32": {"constraint": "r", "returns": {"int32": "int"}}, - "f32": {"constraint": "f", "returns": {"float32": "float"}}, +_PTX_VEC = {"", "v2", "v4", "v8"} +_PTX_L1_EVICT = { + "", + "L1::evict_normal", + "L1::evict_unchanged", + "L1::evict_first", + "L1::evict_last", + "L1::no_allocate", } +_PTX_L2_EVICT = {"", "L2::evict_normal", "L2::evict_first", "L2::evict_last"} +_PTX_PREFETCH = {"", "L2::64B", "L2::128B", "L2::256B"} + + +def _bool_attr(value): + return bool(int(value)) if hasattr(value, "value") else bool(value) def _parse_ld_attrs(return_dtype, ptx_type, scope=None, space="global"): @@ -120,21 +145,27 @@ def _parse_ld_attrs(return_dtype, ptx_type, scope=None, space="global"): ptx_type = parse_str(ptx_type) scope = None if scope is None else parse_str(scope) space = parse_str(space) - if ptx_type not in _PTX_LD_TYPES: - raise ValueError( - f"Unsupported PTX ld type {ptx_type!r}; expected one of {sorted(_PTX_LD_TYPES)}" - ) - returns = _PTX_LD_TYPES[ptx_type]["returns"] - if return_dtype not in returns: - raise ValueError( - f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; " - f"expected one of {sorted(returns)}" - ) + ptx_type, _ptx, constraint, default_tvm = _type_info(ptx_type) + if ptx_type in _PTX_LD_TYPE_RETURNS: + returns = _PTX_LD_TYPE_RETURNS[ptx_type] + if return_dtype not in returns: + raise ValueError( + f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; " + f"expected one of {sorted(returns)}" + ) + c_type = returns[return_dtype] + else: + if return_dtype != default_tvm: + raise ValueError( + f"PTX ld type {ptx_type!r} cannot return TVM dtype {return_dtype!r}; " + f"expected {default_tvm!r}" + ) + c_type = _ptx if scope is not None and scope not in _PTX_LD_SCOPES: raise ValueError( f"Unsupported PTX ld scope {scope!r}; expected one of {sorted(_PTX_LD_SCOPES)}" ) - return return_dtype, ptx_type, scope, space, returns[return_dtype] + return return_dtype, ptx_type, scope, space, c_type, constraint def _validate_ld_space(space: str, allowed: set[str]) -> None: @@ -144,158 +175,248 @@ def _validate_ld_space(space: str, allowed: set[str]) -> None: ) -def _ptx_ld_helper_name(kind: str, return_dtype: str, ptx_type: str, scope: str | None, space: str): - parts = ["tvm_builtin_ptx_ld", kind] - if scope is not None: - parts.append(scope.replace("::", "_")) - parts.extend([space.replace("::", "_"), ptx_type, return_dtype]) - return "_".join(parts) +def _ptx_level_suffix(has_cache, l1_evict, l2_evict, prefetch_size): + suffix = _cache_suffix("cache" if has_cache else "") + 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}" + if prefetch_size: + suffix += f".{prefetch_size}" + return suffix + + +def _ptx_shared_addr(space, ptr_name="address"): + if parse_str(space).startswith("shared"): + return ( + f" unsigned int addr = (unsigned int)__cvta_generic_to_shared({ptr_name});\n", + '"r"(addr)', + ) + return "", f'"l"({ptr_name})' + + +def _ptx_ld_vec_store(num_bytes, vec_len, ptx_type): + 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)) + + "};" + ) + return f" *reinterpret_cast<{store_type}*>(dst_ptr) = r0;" -def _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache_hint): - return_dtype, ptx_type, _scope, space, c_type = _parse_ld_attrs( - return_dtype, ptx_type, None, space +def _ptx_ld_form_parts(form, attr_args): + if form == "weak": + ( + return_dtype, + weak, + space, + cop, + vec, + ptx_type, + has_cache_hint, + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) = attr_args + sem, scope = "", "" + elif form == "relaxed": + ( + return_dtype, + scope, + space, + vec, + ptx_type, + has_cache_hint, + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) = attr_args + sem, weak, cop = "", False, "" + elif form == "acquire": + ( + return_dtype, + scope, + space, + vec, + ptx_type, + has_cache_hint, + to_dst, + l1_evict, + l2_evict, + prefetch_size, + ) = attr_args + sem, weak, cop = "", False, "" + elif form == "volatile": + return_dtype, space, vec, ptx_type, to_dst, prefetch_size = attr_args + sem, scope, weak, cop = "", "", False, "" + has_cache_hint, l1_evict, l2_evict = False, "", "" + elif form == "mmio": + 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, "", "", "" + else: + raise ValueError(f"unknown ld form {form!r}") + + return_dtype, ptx_type, scope, space, c_type, constraint = _parse_ld_attrs( + return_dtype, ptx_type, scope if form in ("relaxed", "acquire") else None, space ) + sem = parse_str(sem) + scope = parse_str(scope) + space = parse_str(space) cop = parse_str(cop) - if cop not in _PTX_LD_COPS: + vec = parse_str(vec) + l1_evict = parse_str(l1_evict) + l2_evict = parse_str(l2_evict) + 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: raise ValueError(f"Unsupported PTX ld cache operation {cop!r}") - weak = bool(int(weak)) if hasattr(weak, "value") else bool(weak) - has_cache = ( - bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint) + if vec and vec not in _PTX_VEC: + raise ValueError(f"Unsupported PTX ld vector modifier {vec!r}") + if l1_evict and l1_evict not in _PTX_L1_EVICT: + raise ValueError(f"Unsupported PTX ld L1 eviction {l1_evict!r}") + if l2_evict and l2_evict not in _PTX_L2_EVICT: + raise ValueError(f"Unsupported PTX ld L2 eviction {l2_evict!r}") + if prefetch_size and prefetch_size not in _PTX_PREFETCH: + raise ValueError(f"Unsupported PTX ld prefetch size {prefetch_size!r}") + if form == "mmio": + 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 == "relaxed": + if not scope: + raise ValueError("ld.relaxed requires scope") + _validate_ld_space(space, _PTX_LD_SPACES) + prefix = f"ld.relaxed.{scope}{_dot(space)}" + elif form == "acquire": + if not scope: + raise ValueError("ld.acquire requires scope") + _validate_ld_space(space, _PTX_LD_SPACES) + prefix = f"ld.acquire.{scope}{_dot(space)}" + elif form == "volatile": + _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES) + prefix = f"ld.volatile{_dot(space)}" + 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) + vec_len = int(vec[1:]) if vec else 1 + if vec and not to_dst: + raise ValueError("vector ld requires to_dst") + elem_bytes = ( + 8 + if ptx_type.endswith("64") + else 2 + if ptx_type in ("u16", "s16", "b16") + else 1 + if ptx_type in ("u8", "s8", "b8") + else 4 ) - _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES | {"param::entry", "param::func"}) - spec = _PTX_LD_TYPES[ptx_type]["constraint"] - addr_decl = "" - addr_operand = '"l"(address)' - if space.startswith("shared"): - addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n" - addr_operand = '"r"(addr)' - modifiers = f"{'.weak' if weak else ''}.{space}{('.' + cop) if cop else ''}" - cache_inst = ".L2::cache_hint" if has_cache else "" - cache_slot = ", %2" if has_cache else "" - cache_operand = ', "l"(cache_policy)' if has_cache else "" - name = ( - "tvm_builtin_ptx_ld" - f"{'_weak' if weak else ''}_{space.replace('::', '_').replace('.', '_')}" - f"{('_' + cop) if cop else ''}_{ptx_type}_{return_dtype}" - f"{'_cache_hint' if has_cache else ''}" + 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), + _safe_attr(cop) if cop else "", + _safe_attr(vec) if vec else "", + ptx_type, + return_dtype if not to_dst else "to_dst", + ] ) + if has_cache: + name_parts.append("cache_hint") + if l1_evict: + name_parts.append(_safe_attr(l1_evict)) + if l2_evict: + name_parts.append(_safe_attr(l2_evict)) + if prefetch_size: + name_parts.append(_safe_attr(prefetch_size)) + 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") + if 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)) + "}" + out_constraints = ", ".join(f'"={constraint}"(r{i})' for i in range(vec_len)) + addr_idx = vec_len + else: + out_slot = "%0" + out_constraints = f'"={constraint}"(r0)' + addr_idx = 1 + cache_slot = f", %{addr_idx + 1}" if has_cache else "" + instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}" + body = ( + f"{addr_decl}{reg_decls}" + 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)}" + ) + return ( + name, + "(void* dst_ptr, void* src_ptr, unsigned long long cache_policy)", + "void", + "", + body, + ) + cache_slot = ", %2" if has_cache else "" + instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}" body = ( f" {c_type} ret;\n" f"{addr_decl}" - f' asm volatile("ld{modifiers}{cache_inst}.{ptx_type} %0, [%1]{cache_slot};" ' - f': "={spec}"(ret) : {addr_operand}{cache_operand});\n' + f' asm volatile("{instr} %0, [%1]{cache_slot};"\n' + f' : "={constraint}"(ret)\n' + f" : {addr_operand}{cache_operand});\n" " return ret;" ) - return name, c_type, return_dtype, body - - -device_intrinsic( - "ptx_ld", - n_attrs=6, - helper_name=lambda _addr, _cache_policy, return_dtype, weak, space, cop, ptx_type, has_cache: ( - _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[0] - ), - c_signature="(void* address, unsigned long long cache_policy)", - return_type=lambda _addr, _cache_policy, return_dtype, weak, space, cop, ptx_type, has_cache: ( - _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[1] - ), - tvm_return_type=lambda _addr, - _cache_policy, - return_dtype, - _weak, - _space, - _cop, - _ptx_type, - _has_cache: (parse_str(return_dtype)), - body=lambda _addr, _cache_policy, return_dtype, weak, space, cop, ptx_type, has_cache: ( - _ptx_ld_parts(return_dtype, ptx_type, weak, space, cop, has_cache)[3] - ), -) - - -def _ptx_ld_acquire_parts(return_dtype, ptx_type, scope, space): - return_dtype, ptx_type, scope, space, c_type = _parse_ld_attrs( - return_dtype, ptx_type, scope, space + sig = ( + "(void* address, unsigned long long cache_policy)" if form == "weak" else "(void* address)" ) - _validate_ld_space(space, _PTX_LD_SPACES) - spec = _PTX_LD_TYPES[ptx_type]["constraint"] - addr_decl = "" - addr_operand = '"l"(address)' - if space.startswith("shared"): - addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n" - addr_operand = '"r"(addr)' - return ( - _ptx_ld_helper_name("acquire", return_dtype, ptx_type, scope, space), - c_type, - ( - f" {c_type} ret;\n" - f"{addr_decl}" - f' asm volatile("ld.acquire.{scope}.{space}.{ptx_type} %0, [%1];" ' - f': "={spec}"(ret) : {addr_operand});\n' - " return ret;" - ), - return_dtype, - ) - + return name, sig, c_type, return_dtype, body -device_intrinsic( - "ptx_ld_acquire", - n_attrs=4, - helper_name=lambda _addr, return_dtype, ptx_type, scope, space: _ptx_ld_acquire_parts( - return_dtype, ptx_type, scope, space - )[0], - c_signature="(void* address)", - return_type=lambda _addr, return_dtype, ptx_type, scope, space: _ptx_ld_acquire_parts( - return_dtype, ptx_type, scope, space - )[1], - tvm_return_type=lambda _addr, return_dtype, _ptx_type, _scope, _space: parse_str(return_dtype), - body=lambda _addr, return_dtype, ptx_type, scope, space: _ptx_ld_acquire_parts( - return_dtype, ptx_type, scope, space - )[2], -) +def _register_ptx_ld(op_name, form, n_attrs): + def _parts(*args): + return _ptx_ld_form_parts(form, args[-n_attrs:]) -def _ptx_ld_volatile_parts(return_dtype, ptx_type, space): - return_dtype, ptx_type, _scope, space, c_type = _parse_ld_attrs( - return_dtype, ptx_type, None, space - ) - _validate_ld_space(space, _PTX_LD_VOLATILE_SPACES) - spec = _PTX_LD_TYPES[ptx_type]["constraint"] - addr_decl = "" - addr_operand = '"l"(address)' - if space.startswith("shared"): - addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n" - addr_operand = '"r"(addr)' - return ( - _ptx_ld_helper_name("volatile", return_dtype, ptx_type, None, space), - c_type, - ( - f" {c_type} ret;\n" - f"{addr_decl}" - f' asm volatile("ld.volatile.{space}.{ptx_type} %0, [%1];" ' - f': "={spec}"(ret) : {addr_operand});\n' - " return ret;" + device_intrinsic( + op_name, + n_attrs=n_attrs, + helper_name=lambda *a, _p=_parts: _p(*a)[0], + c_signature=lambda *a, _p=_parts: _p(*a)[1], + return_type=lambda *a, _p=_parts: _p(*a)[2], + tvm_return_type=lambda *a, _p=_parts: ( + None if _p(*a)[2] == "void" else parse_str(a[-n_attrs]) ), - return_dtype, + body=lambda *a, _p=_parts: _p(*a)[4], ) -device_intrinsic( - "ptx_ld_volatile", - n_attrs=3, - helper_name=lambda _addr, return_dtype, ptx_type, space: _ptx_ld_volatile_parts( - return_dtype, ptx_type, space - )[0], - c_signature="(void* address)", - return_type=lambda _addr, return_dtype, ptx_type, space: _ptx_ld_volatile_parts( - return_dtype, ptx_type, space - )[1], - tvm_return_type=lambda _addr, return_dtype, _ptx_type, _space: parse_str(return_dtype), - body=lambda _addr, return_dtype, ptx_type, space: _ptx_ld_volatile_parts( - return_dtype, ptx_type, space - )[2], -) +_register_ptx_ld("ptx_ld", "weak", 11) +_register_ptx_ld("ptx_ld_relaxed", "relaxed", 10) +_register_ptx_ld("ptx_ld_acquire", "acquire", 10) +_register_ptx_ld("ptx_ld_volatile", "volatile", 6) +_register_ptx_ld("ptx_ld_mmio", "mmio", 6) # ============================================================================= @@ -483,39 +604,6 @@ def _ptx_mapa_parts(_addr, _rank, space, ptx_type, return_dtype): # concrete sem/scope/space/op/type parameters for existing call sites. # ============================================================================= -_PTX_SCALAR_TYPE_INFO = { - "b32": ("unsigned int", "r", "uint32"), - "u32": ("unsigned int", "r", "uint32"), - "s32": ("int", "r", "int32"), - "b64": ("unsigned long long", "l", "uint64"), - "u64": ("unsigned long long", "l", "uint64"), - "s64": ("long long", "l", "int64"), - "f32": ("float", "f", "float32"), - "f64": ("double", "d", "float64"), -} - - -def _safe_attr(value): - return parse_str(value).replace("::", "_").replace(".", "_") - - -def _dot(value): - value = parse_str(value) - return f".{value}" if value else "" - - -def _cache_suffix(cache): - return ".L2::cache_hint" if cache else "" - - -def _type_info(ptx_type): - ptx_type = parse_str(ptx_type) - if ptx_type not in _PTX_SCALAR_TYPE_INFO: - raise ValueError( - f"Unsupported PTX scalar type {ptx_type!r}; expected {sorted(_PTX_SCALAR_TYPE_INFO)}" - ) - return (ptx_type, *_PTX_SCALAR_TYPE_INFO[ptx_type]) - # PTX red scalar form: # red{.sem}{.scope}{.space}.op{.level::cache_hint}.type [a], b{, cache-policy}; @@ -641,57 +729,185 @@ def _prefetch_tensormap_parts(_tensor_map, tensormap_space): ) -# PTX st weak scalar/vector form: +# PTX st forms (ISA table entries registered via ``ptx_st`` and siblings): # st{.weak}{.ss}{.cop}{.level::cache_hint}{.vec}.type [a], b{, cache-policy}; -def _ptx_st_parts(*args): - weak, space, cop, vec, ptx_type, has_cache_hint = args[-6:] - weak = bool(int(weak)) if hasattr(weak, "value") else bool(weak) +# st{.weak}{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy}; +# st.volatile{.ss}{.vec}.type [a], b; +# st.relaxed.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy}; +# st.release.scope{.ss}{.level1::eviction_priority}{.level2::eviction_priority}{.level::cache_hint}{.vec}.type [a], b{, cache-policy}; +# st.mmio.sem.sys{.global}.type [a], b; +_PTX_ST_COPS = {"", "wb", "cg", "cs", "wt"} +_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] + 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" + + +def _ptx_st_form_parts(form, attr_args, from_src): + if form == "weak": + weak, space, cop, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args + sem, scope = "", "" + elif form == "relaxed": + scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args + sem, weak, cop = "relaxed", False, "" + elif form == "release": + scope, space, vec, ptx_type, has_cache_hint, l1_evict, l2_evict = attr_args + sem, weak, cop = "release", False, "" + elif form == "volatile": + space, vec, ptx_type = attr_args + sem, scope, weak, cop = "", "", False, "" + has_cache_hint, l1_evict, l2_evict = False, "", "" + elif form == "mmio": + sem, scope, space, ptx_type = attr_args + weak, cop, vec = False, "", "" + has_cache_hint, l1_evict, l2_evict = False, "", "" + else: + raise ValueError(f"unknown st form {form!r}") + + sem = parse_str(sem) + scope = parse_str(scope) space = parse_str(space) cop = parse_str(cop) vec = parse_str(vec) + l1_evict = parse_str(l1_evict) + l2_evict = parse_str(l2_evict) + weak = _bool_attr(weak) + has_cache = _bool_attr(has_cache_hint) ptx_type, c_type, constraint, _tvm_dtype = _type_info(ptx_type) - has_cache = ( - bool(int(has_cache_hint)) if hasattr(has_cache_hint, "value") else bool(has_cache_hint) - ) + if cop and cop not in _PTX_ST_COPS: + raise ValueError(f"Unsupported PTX st cache operation {cop!r}") + if vec and vec not in _PTX_VEC: + raise ValueError(f"Unsupported PTX st vector modifier {vec!r}") + if space not in _PTX_ST_SPACES and not (form == "mmio" and space == "global"): + raise ValueError(f"Unsupported PTX st state space {space!r}") vec_len = int(vec[1:]) if vec else 1 - modifiers = f"{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}" - instr = f"st{modifiers}{_cache_suffix('cache' if has_cache else '')}{_dot(vec)}.{ptx_type}" - name = ( - "tvm_builtin_ptx_st" - f"{'_weak' if weak else ''}_{_safe_attr(space)}" - f"{('_' + _safe_attr(cop)) if cop else ''}" - f"{('_' + _safe_attr(vec)) if vec else ''}_{ptx_type}" - f"{'_cache_hint' if has_cache else ''}" + elem_bytes = ( + 8 + if ptx_type.endswith("64") + else 2 + if ptx_type in ("u16", "s16", "b16") + else 1 + if ptx_type in ("u8", "s8", "b8") + else 4 ) - value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len)) - c_signature = f"(void* address, {value_params}, unsigned long long cache_policy)" - values = f"{{{', '.join(f'%{i + 1}' for i in range(vec_len))}}}" if vec else "%1" + num_bytes = vec_len * elem_bytes if vec else elem_bytes + use_cache_policy = form in ("weak", "relaxed", "release") + if form == "mmio": + if sem not in ("acquire", "relaxed", "release") or scope != "sys" or space != "global": + raise ValueError("st.mmio requires sem, scope=sys, space=global") + prefix = f"st.mmio.{sem}.{scope}" + elif form == "relaxed": + if not scope: + raise ValueError("st.relaxed requires scope") + prefix = f"st.relaxed.{scope}{_dot(space)}" + elif form == "release": + if not scope: + raise ValueError("st.release requires scope") + prefix = f"st.release.{scope}{_dot(space)}" + elif form == "volatile": + prefix = f"st.volatile{_dot(space)}" + else: + prefix = f"st{'.weak' if weak else ''}{_dot(space)}{_dot(cop)}" + level = _ptx_level_suffix(has_cache, l1_evict, l2_evict, "") + instr = f"{prefix}{level}{_dot(vec)}.{ptx_type}" + name_parts = ["tvm_builtin_ptx_st", 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, + "from_src" if from_src else "values", + ] + ) + if has_cache: + name_parts.append("cache_hint") + 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 = '"l"(address)' - if space.startswith("shared"): - addr_decl = " unsigned int addr = (unsigned int)__cvta_generic_to_shared(address);\n" - addr_operand = '"r"(addr)' - body = ( - f"{addr_decl}" - f' asm volatile("{instr} [%0], {values}{cache_slot};"\n' - " :\n" - f" : {addr_operand}{value_constraints}" - f"{cache_operand}\n" - ' : "memory");' - ) - return name, c_signature, body - - -device_intrinsic( - "ptx_st", - n_attrs=6, - helper_name=lambda *a: _ptx_st_parts(*a)[0], - c_signature=lambda *a: _ptx_st_parts(*a)[1], - body=lambda *a: _ptx_st_parts(*a)[2], -) + addr_decl, addr_operand = _ptx_shared_addr(space, "address") + if from_src: + load_regs = _ptx_st_load_src(num_bytes, vec_len, ptx_type, c_type) + if vec_len > 1: + in_constraints = ", ".join(f'"{constraint}"(r{i})' for i in range(vec_len)) + value_args = f", {in_constraints}" + else: + value_args = f', "{constraint}"(r0)' + body = ( + f"{addr_decl}{load_regs}" + f' asm volatile("{instr} [%0], {values}{cache_slot};"\n' + " :\n" + f" : {addr_operand}{value_args}{cache_operand}\n" + ' : "memory");' + ) + if use_cache_policy: + sig = "(void* address, void* src_ptr, unsigned long long cache_policy)" + else: + sig = "(void* address, void* src_ptr)" + else: + body = ( + f"{addr_decl}" + f' asm volatile("{instr} [%0], {values}{cache_slot};"\n' + " :\n" + f" : {addr_operand}{value_constraints}{cache_operand}\n" + ' : "memory");' + ) + if form == "mmio": + sig = f"(void* 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)" + else: + value_params = ", ".join(f"{c_type} value{i}" for i in range(vec_len)) + sig = f"(void* address, {value_params})" + return name, sig, body + + +def _register_ptx_st(op_name, form, n_attrs, *, with_cache_policy=True): + 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)] + 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) + + codegen.__name__ = f"codegen_{op_name}" + register_codegen(op_name)(codegen) + + +def _register_ptx_st_mmio(op_name, form, n_attrs): + def codegen(*args): + parts = _ptx_st_form_parts(form, args[-n_attrs:], False) + forward = args[:-n_attrs] + 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) + + codegen.__name__ = f"codegen_{op_name}" + register_codegen(op_name)(codegen) + + +_register_ptx_st("ptx_st", "weak", 9) +_register_ptx_st("ptx_st_relaxed", "relaxed", 8) +_register_ptx_st("ptx_st_release", "release", 8) +_register_ptx_st("ptx_st_volatile", "volatile", 4) +_register_ptx_st_mmio("ptx_st_mmio", "mmio", 4) # PTX st.bulk form: diff --git a/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py index 029a6bf87ce8..a99f8f633274 100644 --- a/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py +++ b/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py @@ -995,8 +995,8 @@ def _get_tcgen05_mma_scale_vec_size(kind, scale_dtype): def _mma_block_scaled_parts(*args): """Args layout: (d_tmem_addr, a_operand, b_desc[, sp_tmem_addr], i_desc, enable_input_d, sfa_tmem_addr, sfb_tmem_addr, - kind, scale_vec_size, sparse, use_a_tmem, cta_group).""" - attrs = args[-5:] + [pred], kind, scale_vec_size, sparse, use_a_tmem, cta_group, has_pred).""" + attrs = args[-6:] kind = parse_str(attrs[0]) scale_vec_size = int(attrs[1]) sparse_raw = attrs[2] @@ -1006,6 +1006,7 @@ def _mma_block_scaled_parts(*args): bool(int(use_a_tmem_raw)) if hasattr(use_a_tmem_raw, "value") else bool(use_a_tmem_raw) ) cta_group = int(attrs[4]) + has_pred = bool(int(attrs[5])) a_type = "uint32_t" if use_a_tmem else "uint64_t" a_constraint = "r" if use_a_tmem else "l" @@ -1016,11 +1017,14 @@ def _mma_block_scaled_parts(*args): sig_parts.extend( ["uint32_t i_desc", "uint32_t scaleC", "uint32_t sfa_tmem_addr", "uint32_t sfb_tmem_addr"] ) + if has_pred: + sig_parts.append("uint32_t pred") sig = "(" + ", ".join(sig_parts) + ")" name = ( f"ptx_tcgen05_mma_block_scaled_cta_{cta_group}_kind_{kind}_scale_vec_{scale_vec_size}" f"{'_sp' if sparse else ''}{'_TS' if use_a_tmem else '_SS'}" + f"{'_pred' if has_pred else ''}" ) sparse_suffix = ".sp" if sparse else "" @@ -1036,12 +1040,19 @@ def _mma_block_scaled_parts(*args): f' "r"(i_desc), "r"(scaleC), "r"(sfa_tmem_addr), "r"(sfb_tmem_addr)' f"{sp_input}" ) + pred_idx = 8 if sparse else 7 + 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 has_pred: + asm_inputs += ', "r"(pred)' body = ( " asm volatile(\n" ' "{\\n"\n' - ' ".reg .pred p;\\n"\n' + f' ".reg .pred p{pred_reg};\\n"\n' ' "setp.ne.b32 p, %4, 0;\\n"\n' - f' "{instr} "\n' + f"{pred_setp}" + f' "{pred_prefix}{instr} "\n' f' "[%0], {a_str}, %2, {sparse_placeholder}%3, [%5], [%6], p;\\n"\n' ' "}\\n"\n' " :\n" @@ -1053,7 +1064,7 @@ def _mma_block_scaled_parts(*args): device_intrinsic( "_ptx_tcgen05_mma_block_scaled_form", - n_attrs=5, + n_attrs=6, helper_name=lambda *a: _mma_block_scaled_parts(*a)[0], c_signature=lambda *a: _mma_block_scaled_parts(*a)[1], body=lambda *a: _mma_block_scaled_parts(*a)[2], @@ -1075,6 +1086,7 @@ def _dispatch_tcgen05_mma_block_scaled( use_a_tmem, cta_group, enable_input_d, + pred=None, sparse=False, sp_tmem_addr=None, ): @@ -1100,8 +1112,11 @@ def _dispatch_tcgen05_mma_block_scaled( if sparse: operand_args.append(sp_tmem_addr) operand_args.extend([i_desc, enable_input_d, sfa_tmem_addr, sfb_tmem_addr]) + has_pred = pred is not None + if has_pred: + operand_args.append(pred) - attr_args = [kind, scale_vec_size, sparse, use_a_tmem_b, cta_group_i] + attr_args = [kind, scale_vec_size, sparse, use_a_tmem_b, cta_group_i, int(has_pred)] return CODEGEN_REGISTRY["tirx._ptx_tcgen05_mma_block_scaled_form"](operand_args + attr_args) @@ -1121,7 +1136,13 @@ def codegen_ptx_tcgen05_mma_block_scale( use_a_tmem, cta_group, enable_input_d=1, + *pred_args, ): + if len(pred_args) > 1: + raise ValueError( + f"tcgen05.mma.block_scale expected at most one predicate, got {len(pred_args)}" + ) + pred = pred_args[0] if pred_args else None return _dispatch_tcgen05_mma_block_scaled( d_dtype, a_dtype, @@ -1137,6 +1158,7 @@ def codegen_ptx_tcgen05_mma_block_scale( use_a_tmem, cta_group, enable_input_d, + pred=pred, ) 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/_common.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py index d8fad5f6ae56..b6399a16fe44 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py @@ -538,3 +538,19 @@ def _outer_offsets(outer_iters_s, outer_iters_g, flat_idx): ds = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_s)) dg = sum(c * int(it.stride) for c, it in zip(coords, outer_iters_g)) return ds, dg + + +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 { + 16: ("v4", "u32"), + 8: ("v2", "u32"), + 4: ("", "u32"), + 2: ("", "u16"), + 1: ("", "u8"), + }[num_bytes] + + +def copy_ptx_ld_return_type(ptx_type: str) -> str: + """TVM dtype string for ``T.ptx.ld``'s ``return_type`` argument.""" + return {"u32": "uint32", "u16": "uint16", "u8": "uint32"}[ptx_type] 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 ab69d3924450..9318026bc3b6 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py @@ -27,7 +27,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 from .reg import _axis_decl 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 727f7a18df77..f7a28f6f77d3 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,12 +36,14 @@ 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, _thread_cnt, align_layouts_gs, + copy_ptx_form, + copy_ptx_ld_return_type, ) from ._swizzle_iter import ( emit_init, @@ -136,7 +138,8 @@ def _emit_gmem_smem(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFu # [outer x thread x vec] coord scheme below. vec_bits = vec_len * elem_bits - copy_op = getattr(T.cuda, f"copy_{vec_bits}b") + num_bytes = vec_bits // 8 + vec, ptx_type = copy_ptx_form(num_bytes) # Partition guarantees ``prod(s_p.shard.extents) == prod(g_p.shard.extents) # == n_elements`` (the total transfer count). Express the per-thread @@ -267,6 +270,8 @@ def _s_off(f, s_lin): def impl(): tid = _decl_tid() _setup_swizzle(tid) + tmp = T.alloc_local((vec_len,), src.dtype) + tmp_ptr = tmp.ptr_to([0]) # NB: pass typed ptr_to(...) directly to _ptr_off; caching in a # local var turns it into void* + offset = byte arithmetic → # misaligned vector ops. @@ -285,9 +290,29 @@ def impl(): 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: - copy_op(s_ptr, g_ptr) + 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: - copy_op(g_ptr, s_ptr) + 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/ld_stmatrix.py b/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py index 96a4545c2b92..5b799bdafca6 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,7 +32,7 @@ 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, @@ -95,17 +95,13 @@ def _emit(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: s_shape = list(s_buf.shape) s_layout = s_buf.layout - # Step 2: canonicalize, then slice, then canonicalize. Push target so the - # scope-aware fusers run (e.g. laneid+wid_in_wg -> tid_in_wg for warpgroup). - # Canonicalize *before* slicing too: a frag carrying separate laneid + - # wid_in_wg thread axes (e.g. a permuted tcgen05-ld atom) only fuses to a - # single tid_in_wg axis on the *full* layout — slicing first leaves a - # sub-layout whose scope chain is ill-formed and GetScope rejects it. 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 = r_layout.canonicalize().slice(r_shape, r_region).canonicalize() - s = s_layout.canonicalize().slice(s_shape, s_region).canonicalize() + r_sliced = r_layout.slice(r_shape, r_region) + s_sliced = s_layout.slice(s_shape, s_region) + r = r_sliced.canonicalize() + s = s_sliced.canonicalize() # Step 2.5: peel any S-side swizzle wrapper to expose the underlying # TileLayout. The ComposeLayout doesn't have ``.replica`` / ``.shard``, 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 3f30a83c038f..17a163aec5cf 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py @@ -37,9 +37,9 @@ from tvm.tirx.layout import ComposeLayout, S, SwizzleLayout, 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 ._common import _alignment_ok +from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_return_type from ._swizzle_iter import ( emit_fallback_offset, emit_init, @@ -122,7 +122,6 @@ def _r_side_layout_valid( return False, f"R layout is {type(layout).__name__}, not TileLayout" scope_rank = _SCOPE_RANK[sctx.scope_kind] - seen_thread_axes: set[str] = set() for it in layout.shard: ax = it.axis if not ax.is_thread(): @@ -138,16 +137,6 @@ def _r_side_layout_valid( False, f"R thread axis {ax.name!r} scope={ax_scope.name!r} > exec {sctx.scope_kind!r}", ) - # TODO: lift these two; for now i = thread_value (stride=1, each axis appears once). - if int(it.stride) != 1: - return ( - False, - f"R thread axis {ax.name!r} stride={int(it.stride)} != 1 (not supported yet)", - ) - if ax.name in seen_thread_axes: - return False, f"R thread axis {ax.name!r} appears more than once (not supported yet)" - seen_thread_axes.add(ax.name) - 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) @@ -203,30 +192,39 @@ def key(p): return [i for i, _ in sorted(enumerate(r.shard), key=key)] -def align_layouts_raw(r_layout, r_shape, r_region, s_layout, s_shape, s_region): - """Returns (r_p, s_p, s_seps).""" - r = r_layout.slice(list(r_shape), r_region).canonicalize() - s = s_layout.slice(list(s_shape), s_region).canonicalize() +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_p = r.permute_dims(perm).canonicalize() + 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 + return r_p, s_p, s_seps, r_perm -def _split_thread_loop(r_p, s_p, s_seps): +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.""" + 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_p.shard): + for k, r_it in enumerate(r_perm.shard): if r_it.axis.is_thread(): continue r_iters.append(r_it) @@ -296,17 +294,10 @@ def _align_layouts(op_call: TilePrimitiveCall, sctx: DispatchContext): 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] - # Push the dispatch target so layout.canonicalize() runs scope-aware - # fusers (e.g. laneid+wid_in_wg -> tid_in_wg). with sctx.target: - return align_layouts_raw( - r_buf.layout, - r_buf.shape, - r_region, - s_buf.layout, - s_buf.shape, - s_region, - ) + 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]: @@ -490,8 +481,8 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: r_buf, s_buf, r_is_src = dst, src, False with sctx.target: - r_p, s_p, s_seps = _align_layouts(op_call, sctx) - r_iters, s_groups = _split_thread_loop(r_p, s_p, s_seps) + 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) @@ -515,7 +506,10 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc: for _ax, val in r_p.offset.items(): r_off_base = r_off_base + val - copy_op = getattr(T.cuda, f"copy_{vec_bits}b") + 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" total_outer = 1 for a in outer: @@ -572,9 +566,16 @@ def impl(): 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: - copy_op(s_ptr, r_ptr) + T.ptx.st(s_ptr, src=r_ptr, space=space, vec=vec, ptx_type=ptx_type) else: - copy_op(r_ptr, s_ptr) + T.ptx.ld( + s_ptr, + copy_ptx_ld_return_type(ptx_type), + ptx_type, + dst=r_ptr, + space=space, + vec=vec, + ) # fmt: on import os 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..f3adb61ac26d 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py @@ -20,7 +20,7 @@ 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 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..39d2e5f74fdb 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 @@ -43,7 +43,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, 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 5cccc307f573..7b4e790c06f9 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 @@ -65,7 +65,8 @@ from tvm.tirx.layout import ComposeLayout, SwizzleLayout, 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 _is_valid_smem_tmem_copy, _single_thread_exec @@ -372,14 +373,15 @@ def _canon_segment(iters): # add_post_buffer_def_stmt. # ----------------------------------------------------------------------------- def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle): - cache_key = f"smem_tmem_desc:{hash(s_buf)}:{int(ldo)}:{int(sdo)}:{int(swizzle)}" + # Cache descriptor template at SMEM 0; patch addr per cp. + cache_key = f"smem_tmem_desc:{int(ldo)}:{int(sdo)}:{int(swizzle)}" cached = sctx.cache_get(cache_key) if cached is not None: return cached desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name="cp_desc", scope="local") encode_call = T.ptx.tcgen05.encode_matrix_descriptor( - desc_buf.data, s_buf.ptr_to([0] * len(s_buf.shape)), ldo, sdo, swizzle + desc_buf.data, T.reinterpret("handle", T.uint64(0)), ldo, sdo, swizzle ) wrap = SeqStmt([AllocBuffer(desc_buf), Evaluate(encode_call)]) sctx.add_post_buffer_def_stmt(s_buf, wrap) @@ -387,6 +389,19 @@ def _get_or_create_desc(sctx, s_buf, ldo, sdo, swizzle): return desc_buf +def _desc_set_addr(desc_val, addr_ptr): + """Patch a SMEM matrix descriptor's 14-bit address field with cvta(addr)>>4 — + matches the hand-rolled ``replace_smem_desc_addr`` (descriptor encoded at 0).""" + start_addr = T.cast( + T.bitwise_and( + T.shift_right(T.cuda.cvta_generic_to_shared(addr_ptr), T.uint32(4)), + T.uint32(0x3FFF), + ), + "uint64", + ) + return T.bitwise_or(T.bitwise_and(desc_val, T.bitwise_not(T.uint64(0x3FFF))), start_addr) + + # ----------------------------------------------------------------------------- # 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 @@ -402,13 +417,18 @@ def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> Pr init_off_16B = plan["init_off_16B"] t_col0 = plan["t_col0"] - LDO_field = 16 # cp 32x128b ignores LDO; placeholder + # cp ignores LDO for data; non-zero LDO bloats the addr-patch codegen. + LDO_field = 0 cta_group = op_call.config.get("cta_group", 1) desc_buf = _get_or_create_desc(sctx, s_buf, LDO_field, SDO_field, sw) t_addr = t_buf.allocated_addr - from tvm.backend.cuda.operator.tile_primitive.common import smem_desc_add_16B_offset + s_rank = len(s_buf.shape) + + 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. @@ -422,7 +442,7 @@ def copy_smem_tmem_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> Pr def impl(): T.ptx.tcgen05.cp( t_addr[0] + t_col0, - smem_desc_add_16B_offset(desc_buf[0], init_off_16B), + _cp_desc(init_off_16B), shape="32x128b", cta_group=cta_group, multicast="warpx4", ) else: @@ -443,7 +463,7 @@ def impl(): 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), + _cp_desc(init_off_16B + s_off), shape="32x128b", cta_group=cta_group, multicast="warpx4", ) # fmt: on 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 081ea5a772d3..974aa597b188 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 @@ -37,7 +37,7 @@ 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 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 7fd773103f1e..c0384dce8e15 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 @@ -24,7 +24,7 @@ Pipeline: -L1 Canonicalize smem+gmem layouts; group gmem by buffer shape; split any +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". @@ -57,7 +57,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 @@ -72,8 +72,8 @@ class GmemIter: """One gmem logical dim after multi-iter group splitting. - ``shape`` and ``stride`` come from the canonicalized gmem layout for - this dim. ``copy_start`` / ``copy_ext`` carve out the user-requested + ``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). """ @@ -307,26 +307,29 @@ class L1Result: smem_groups: list # list[SmemGroup] -def _canonicalize_gmem(g_buf: Buffer) -> TileLayout: +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.canonicalize() + return layout def _canonicalize_smem(s_buf: Buffer) -> TileLayout: return _to_tile_layout(s_buf.layout, s_buf.shape).canonicalize() -def _group_gmem_by_buffer_shape(gmem_canon: TileLayout, buffer_shape: list): - """Group gmem canonicalized layout by the buffer shape. Returns - ``(grouped, separators)`` or raises on failure.""" +def _group_gmem_by_buffer_shape(gmem_raw: TileLayout, buffer_shape: list): + """Group raw gmem first; each group is canonicalized locally before splitting.""" try: - grouped, seps = gmem_canon.group(list(buffer_shape)) + return gmem_raw.group(list(buffer_shape)) except Exception as err: raise ValueError(f"Cannot group gmem layout by buffer shape: {err}") from err - return grouped, seps + + +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)] def _split_multi_iter_group( @@ -343,10 +346,7 @@ def _split_multi_iter_group( """ start = separators[group_idx] end = separators[group_idx + 1] - # Drop ext=1 padding iters (canonicalize may have inserted trivial ones). - raw_shards = [ - sh for sh in grouped.shard[start:end] if not analyzer.can_prove_equal(sh.extent, 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. @@ -397,8 +397,7 @@ def _slice_and_canonicalize_smem( def _regroup_smem_by_extgt1_shape(sliced_smem: TileLayout, extgt1_shape: list) -> tuple: - """Group the sliced smem layout by the ext>1 copy shape. Returns - ``(grouped, separators)`` or ``None`` on failure.""" + """Group the sliced smem layout by the ext>1 copy shape.""" try: return sliced_smem.group(list(extgt1_shape)) except Exception: @@ -419,11 +418,11 @@ def _build_l1_result( smem_canon = _canonicalize_smem(s_buf) _assert_memory_only(smem_canon, "shared") - gmem_canon = _canonicalize_gmem(g_buf) - _assert_memory_only(gmem_canon, "global") + 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_canon, g_buf.shape) + 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 @@ -1118,6 +1117,33 @@ def copy_tma_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc 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) + + # CUtensorMapL2promotion values from cuda.h. Keep the existing 128B + # default, while allowing kernels to match a reference implementation's + # tensor-map promotion policy explicitly. + _TMA_L2_PROMOTION_TO_CU = {0: 0, 64: 1, 128: 2, 256: 3} + l2_promotion = op_call.config.get("l2_promotion", 128) + l2_promotion = getattr(l2_promotion, "value", l2_promotion) + if l2_promotion not in _TMA_L2_PROMOTION_TO_CU: + fail( + f"Unsupported l2_promotion={l2_promotion!r}; " + f"expected one of {sorted(_TMA_L2_PROMOTION_TO_CU)}" + ) + l2_promotion_cu = _TMA_L2_PROMOTION_TO_CU[l2_promotion] + # 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: @@ -1158,7 +1184,8 @@ def val_key(value) -> str: 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}" + f":{val_key(plan.swizzle_mode.value)}:{oob_fill_kind}:{force_cu_dtype}" + f":{l2_promotion_cu}" ) cached_tensormap = sctx.cache_get(tensormap_cache_key) @@ -1233,8 +1260,12 @@ def create_tensor_map(): *element_strides, 0, # CU_TENSOR_MAP_INTERLEAVE_NONE plan.swizzle_mode.value, - 2, # CU_TENSOR_MAP_L2_PROMOTION_L2_128B + l2_promotion_cu, 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() # fmt: on 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 97a62a040a49..a7b0c3943fd4 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py @@ -55,14 +55,25 @@ def buffer_regions(plan) -> list[BufferRegion]: return out +def dtype_name(dtype) -> str: + dtype_obj = getattr(dtype, "dtype", None) + if dtype_obj is not None: + return str(dtype_obj) + return str(dtype) + + +def dtype_bits(dtype) -> int: + return DataType(dtype_name(dtype)).bits + + def compute_dtype_of(plan) -> str: """Widest dtype in bits across dst + buffer/scalar srcs (dst breaks ties).""" - candidates = [plan.dst.buffer.dtype] + candidates = [dtype_name(plan.dst.buffer.dtype)] for s in plan.srcs: if s.buf_region is not None: - candidates.append(s.buf_region.buffer.dtype) + candidates.append(dtype_name(s.buf_region.buffer.dtype)) elif s.scalar is not None: - candidates.append(s.scalar.dtype) + candidates.append(scalar_dtype(s.scalar)) widest = candidates[0] widest_bits = DataType(widest).bits for d in candidates[1:]: @@ -72,6 +83,19 @@ def compute_dtype_of(plan) -> str: return widest +def scalar_dtype(scalar) -> str: + dtype = getattr(scalar, "dtype", None) + if dtype is not None: + return str(dtype) + ty = getattr(scalar, "ty", None) + if ty is None and hasattr(scalar, "expr_ty"): + ty = scalar.expr_ty() + dtype = getattr(ty, "dtype", None) + if dtype is None: + raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing PrimType") + return str(dtype) + + def n_elements(buf_region: BufferRegion) -> int: _, ext = get_st_extent(buf_region) return functools.reduce(operator.mul, ext, 1) @@ -375,6 +399,8 @@ def sync(): "align_operands_to_anchor", "buffer_regions", "compute_dtype_of", + "dtype_bits", + "dtype_name", "emit_scope_sync", "fetch_src_value", "n_elements", diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py index 38a2a9a894c9..2932b121c59a 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py @@ -15,14 +15,16 @@ # specific language governing permissions and limitations # under the License. -"""Binary elementwise ops: add / sub / mul / fdiv. +"""Binary elementwise ops: add / sub / mul / fdiv / maximum. Includes constant-lhs commute logic. Broadcasting (extent=1 dims) is handled at the layout level in dispatch's ``_broadcast_lift``, not here — parser just records each src as-is. ``add``/``sub``/``mul`` attach a ``VecImpl`` for sm_100+ packed f32x2; -``fdiv`` has no packed PTX (uses scalar fallback only). +``fdiv``/``maximum`` have no packed PTX (scalar fallback only — ``max`` lowers +to a single ``FMNMX``/``max.f32``, which is exact, so there is no rounding/ftz +variant to pack). """ from __future__ import annotations @@ -31,12 +33,13 @@ import operator from typing import Any +from tvm.script import tirx as Tx from tvm.tirx import BufferRegion, TilePrimitiveCall from ..vec_emit.binary_f32x2 import BINARY_F32X2_IMPLS from . import OpSpec, Plan, SrcSpec -_COMMUTATIVE = frozenset({"add", "mul"}) +_COMMUTATIVE = frozenset({"add", "mul", "maximum"}) def _parse_binary_for(op_name: str): @@ -100,6 +103,10 @@ def _compute_fdiv(src_vals, extras, dt): return src_vals[0] / src_vals[1] +def _compute_maximum(src_vals, extras, dt): + return Tx.max(src_vals[0], src_vals[1]) + + BINARY_OPS: dict[str, OpSpec] = { "add": OpSpec( "add", @@ -124,4 +131,9 @@ def _compute_fdiv(src_vals, extras, dt): _parse_binary_for("fdiv"), _compute_fdiv, ), + "maximum": OpSpec( + "maximum", + _parse_binary_for("maximum"), + _compute_maximum, + ), } 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 bdeb6c0f6a14..9b38eccc418c 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 @@ -30,6 +30,7 @@ from tvm.tirx import BufferRegion, TilePrimitiveCall from tvm.tirx.expr import FloatImm +from .._common import scalar_dtype from . import OpSpec, Plan, SrcSpec @@ -62,11 +63,14 @@ def _parse_unary(op: TilePrimitiveCall) -> tuple[Plan | None, str | None]: def _check_unary_extras(extras: dict, compute_dtype: str) -> tuple[bool, str | None]: scale = extras.get("scale") - if scale is not None and scale.dtype != compute_dtype: - return False, f"scale dtype {scale.dtype} != compute dtype {compute_dtype}" + if scale is not None and scalar_dtype(scale) != compute_dtype: + return False, f"scale dtype {scalar_dtype(scale)} != compute dtype {compute_dtype}" bias_const = extras.get("bias_const") - if bias_const is not None and bias_const.dtype != compute_dtype: - return False, f"bias_const dtype {bias_const.dtype} != compute dtype {compute_dtype}" + if bias_const is not None and scalar_dtype(bias_const) != compute_dtype: + return ( + False, + f"bias_const dtype {scalar_dtype(bias_const)} != compute dtype {compute_dtype}", + ) return True, None 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 64d77a21cf69..2c276d71eecb 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py @@ -88,7 +88,7 @@ def _validate_scope_level_anchor(anchor_br, sctx: DispatchContext) -> tuple[bool st, ext = get_st_extent(anchor_br) sliced = get_sublayout_from_region(anchor_br.buffer.layout, anchor_br.buffer.shape, st, ext) with sctx.target: - canon = sliced.canonicalize() if hasattr(sliced, "canonicalize") else sliced + canon = sliced.canonicalize() shard = getattr(canon, "shard", None) if shard is None: return False, f"{scope}-scope op operand layout is not a TileLayout after slicing" @@ -131,7 +131,7 @@ def _validate_scope_level_anchor(anchor_br, sctx: DispatchContext) -> tuple[bool return True, None -def _check_layout_operands_agree(plan) -> tuple[bool, str | None]: +def _check_layout_operands_agree(plan, sctx) -> tuple[bool, str | None]: """Replica sigs must match across non-trivial-layout operands. ``align_operands_to_anchor`` normalizes thread + local parts via @@ -149,8 +149,9 @@ def _check_layout_operands_agree(plan) -> tuple[bool, str | None]: replica_sigs = [] for br in layout_brs: st, ext = get_st_extent(br) - sliced = get_sublayout_from_region(br.buffer.layout, br.buffer.shape, st, ext) - canon = sliced.canonicalize() if hasattr(sliced, "canonicalize") else sliced + with sctx.target: + sliced = get_sublayout_from_region(br.buffer.layout, br.buffer.shape, st, ext) + canon = sliced.canonicalize() sig = layout_signature(canon) if sig is None: return False, "layout has no signature (not a TileLayout?)" @@ -210,7 +211,7 @@ def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str ok_b, reason_b = shape_broadcast_compat(op_tshape, anchor_tshape) if not ok_b: return False, f"shape incompat: {reason_b}" - ok4, reason4 = _check_layout_operands_agree(plan) + ok4, reason4 = _check_layout_operands_agree(plan, sctx) if not ok4: return False, reason4 return True, None 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 d85193c06ac9..404d7d0c1415 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py @@ -44,6 +44,7 @@ _tensor_shape_of, buffer_regions, compute_dtype_of, + dtype_bits, emit_scope_sync, fetch_src_value, n_elements, @@ -99,10 +100,10 @@ def check(op_call: TilePrimitiveCall, sctx: DispatchContext) -> tuple[bool, str def _max_layout_vec(plan, total: int, thread_cnt: int) -> int: """Widest vec_chunk dividing all operands' innermost extents AND ``total / thread_cnt``, within dtype-bit candidates ``{128,64,32,16,8}``.""" - max_bits = plan.dst.buffer.dtype.dtype.bits + max_bits = dtype_bits(plan.dst.buffer.dtype) for s in plan.srcs: if s.buf_region is not None: - max_bits = max(max_bits, s.buf_region.buffer.dtype.dtype.bits) + max_bits = max(max_bits, dtype_bits(s.buf_region.buffer.dtype)) per_thread = total // thread_cnt if thread_cnt > 0 else total if total % thread_cnt != 0: return 1 diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py index d8e995b45a65..ea1a76ef797d 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py @@ -29,6 +29,7 @@ from tvm.ir.expr import PrimExpr from tvm.script import tirx as T +from .._common import dtype_name, scalar_dtype from ..ops import VecImpl @@ -49,7 +50,7 @@ def _f32x2_applies(op_name): def applies(op_call, sctx, plan): from ...common import sm_version_ok - if plan.dst.buffer.dtype != "float32": + if dtype_name(plan.dst.buffer.dtype) != "float32": return False, "dst dtype not f32" if not sm_version_ok(op_call, sctx, min_version=100)[0]: return False, "sm version < 100" @@ -57,10 +58,10 @@ def applies(op_call, sctx, plan): return False, "binary requires 2 srcs" for s in plan.srcs: if s.is_scalar: - if s.scalar.dtype != "float32": + if scalar_dtype(s.scalar) != "float32": return False, "scalar src dtype not f32" else: - if s.buf_region.buffer.dtype != "float32": + if dtype_name(s.buf_region.buffer.dtype) != "float32": return False, "buffer src dtype not f32" if s.index_fn is not None: return False, "broadcasting src not supported by f32x2 packed" diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py index 46292761b28f..759ebefa2713 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py @@ -28,6 +28,7 @@ from tvm.ir.expr import PrimExpr from tvm.script import tirx as T +from .._common import dtype_name from ..ops import VecImpl _VEC2_CAST_INTRINSICS = { @@ -60,8 +61,8 @@ def _cast_vec2_applies(op_call, sctx, plan): src = plan.srcs[0] if src.index_fn is not None: return False, "broadcasting src not supported by cast vec2" - src_dtype = src.buf_region.buffer.dtype - dst_dtype = plan.dst.buffer.dtype + src_dtype = dtype_name(src.buf_region.buffer.dtype) + dst_dtype = dtype_name(plan.dst.buffer.dtype) if (src_dtype, dst_dtype) not in _VEC2_CAST_INTRINSICS: return False, f"no vec2 intrinsic for {src_dtype}->{dst_dtype}" return True, None @@ -72,8 +73,10 @@ def _emit_cast_vec2(dst_buf, dst_lane_indices, src_args, extras) -> PrimExpr: # cast_vec2 requires buffer src (guarded by applies()). assert isinstance(src_arg, tuple), "cast vec2 src must be a buffer" src_buf, src_lane_indices = src_arg - func_name = _intrinsic_name(src_buf.dtype, dst_buf.dtype) - source_code = _intrinsic_source(src_buf.dtype, dst_buf.dtype) + src_dtype = dtype_name(src_buf.dtype) + dst_dtype = dtype_name(dst_buf.dtype) + func_name = _intrinsic_name(src_dtype, dst_dtype) + source_code = _intrinsic_source(src_dtype, dst_dtype) return T.cuda.func_call( func_name, T.address_of(dst_buf[tuple(dst_lane_indices[0])]), diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py index f47476d6a5ce..a438a17a6d49 100644 --- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py +++ b/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py @@ -26,6 +26,7 @@ from tvm.ir.expr import PrimExpr from tvm.script import tirx as T +from .._common import dtype_name, scalar_dtype from ..ops import VecImpl from .binary_f32x2 import _lane @@ -33,7 +34,7 @@ def _fma_f32x2_applies(op_call, sctx, plan): from ...common import sm_version_ok - if plan.dst.buffer.dtype != "float32": + if dtype_name(plan.dst.buffer.dtype) != "float32": return False, "dst dtype not f32" if not sm_version_ok(op_call, sctx, min_version=100)[0]: return False, "sm version < 100" @@ -42,16 +43,16 @@ def _fma_f32x2_applies(op_call, sctx, plan): a, b, c = plan.srcs if a.is_scalar: return False, "fma 'a' must be a buffer (no scalar-a packed FMA)" - if a.buf_region.buffer.dtype != "float32": + if dtype_name(a.buf_region.buffer.dtype) != "float32": return False, "src a dtype not f32" if a.index_fn is not None: return False, "broadcasting src a not supported" for s in (b, c): if s.is_scalar: - if s.scalar.dtype != "float32": + if scalar_dtype(s.scalar) != "float32": return False, "scalar b/c dtype not f32" else: - if s.buf_region.buffer.dtype != "float32": + if dtype_name(s.buf_region.buffer.dtype) != "float32": return False, "buffer b/c dtype not f32" if s.index_fn is not None: return False, "broadcasting src b/c not supported" 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 4eace5e3c911..6e14f774c5fa 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 @@ -42,7 +42,8 @@ tmem_datapath_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 ..common import get_st_extent, smem_desc_add_16B_offset from ..exec_scope_utils import single_thread @@ -74,6 +75,19 @@ "int32": 2, } +_INSTR_DESC_SCALE_FORMAT_MAP = { + "float8_e4m3fn": 0, + "float8_e4m3fnuz": 0, + "float8_e8m0fnu": 1, +} + + +def _dtype_name(dtype) -> str: + dtype_obj = getattr(dtype, "dtype", None) + if dtype_obj is not None: + return str(dtype_obj) + return str(dtype) + def _encode_instr_descriptor_dense_uint32( M, @@ -96,6 +110,9 @@ def _encode_instr_descriptor_dense_uint32( local descriptor on every gemm_async call (which forces an inline ``asm`` block that ptxas cannot hoist out of the i_kv loop body). """ + d_dtype = _dtype_name(d_dtype) + a_dtype = _dtype_name(a_dtype) + b_dtype = _dtype_name(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] @@ -114,6 +131,36 @@ def _encode_instr_descriptor_dense_uint32( return desc & 0xFFFFFFFF +def _encode_instr_descriptor_block_scaled_uint32( + M, + N, + a_dtype, + b_dtype, + sf_dtype, + trans_a, + trans_b, + neg_a=False, + neg_b=False, + is_sparse=False, +): + """Compile-time port of ``InstrDescriptorBlockScaled`` bitfield packing.""" + a_format = _INSTR_DESC_FORMAT_MAP[_dtype_name(a_dtype)] + b_format = _INSTR_DESC_FORMAT_MAP[_dtype_name(b_dtype)] + scale_format = _INSTR_DESC_SCALE_FORMAT_MAP[_dtype_name(sf_dtype)] + desc = 0 + desc |= (int(is_sparse) & 0x1) << 2 + desc |= (a_format & 0x7) << 7 + desc |= (b_format & 0x7) << 10 + desc |= (int(neg_a) & 0x1) << 13 + desc |= (int(neg_b) & 0x1) << 14 + desc |= (int(trans_a) & 0x1) << 15 + desc |= (int(trans_b) & 0x1) << 16 + desc |= ((N >> 3) & 0x3F) << 17 + desc |= (scale_format & 0x1) << 23 + desc |= ((M >> 4) & 0x1F) << 24 + return desc & 0xFFFFFFFF + + def sf_smem_layout(rows, SF_K, sf_per_mma, sf_reuse=1, pipe_depth=None): """SMEM-side layout for SF in tcgen05.cp scale-factor copy. @@ -240,8 +287,8 @@ def _compute_sf_mma_k(data_dtype, sf_dtype): - fp4 data + e8m0fnu SF: MMA_K=64, SF_VEC=32 → sf_mma_k=2 - fp4 data + e4m3fn SF (nvfp4): MMA_K=64, SF_VEC=16 → sf_mma_k=4 """ - data_dtype = str(data_dtype) - sf_dtype = str(sf_dtype) + data_dtype = _dtype_name(data_dtype) + sf_dtype = _dtype_name(sf_dtype) if data_dtype in ("float8_e4m3fn", "float8_e5m2"): return 1 # MMA_K=32, one SF per MMA elif data_dtype == "float4_e2m1fn": @@ -384,9 +431,18 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - analyzer = Analyzer() - C_type, A_type, B_type = C_buffer.dtype, A_buffer.dtype, B_buffer.dtype + C_type, A_type, B_type = ( + _dtype_name(C_buffer.dtype), + _dtype_name(A_buffer.dtype), + _dtype_name(B_buffer.dtype), + ) assert C_type == "float32", f"tcgen05 schedule expected C_type=float32, got {C_type}" + # fp32/bf16 storage may still use tf32 MMA semantics via is_AB_tf32. + is_AB_tf32 = op_call.config.get("is_AB_tf32", False) + A_sem = "tf32" if is_AB_tf32 else A_type + B_sem = "tf32" if is_AB_tf32 else B_type + # Valid A/B dtypes for block-scaled MMA (low-precision with per-block scale factors) _BLOCK_SCALED_DTYPES = ["float4_e2m1fn", "float8_e4m3fn"] @@ -400,14 +456,22 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - f"tcgen05 block-scaled schedule expected B_type in {_BLOCK_SCALED_DTYPES}, got {B_type}" ) else: - assert A_type in ["float16", "bfloat16"], ( - f"tcgen05 schedule expected A_type=float16 or bfloat16, got {A_type}" + _DENSE_DTYPES = [ + "float16", + "bfloat16", + "float8_e4m3fn", + "float8_e5m2", + "tensor_float32", + "tf32", + ] + assert A_sem in _DENSE_DTYPES, ( + f"tcgen05 schedule expected A dtype in {_DENSE_DTYPES}, got {A_sem}" ) - assert B_type in ["float16", "bfloat16"], ( - f"tcgen05 schedule expected B_type=float16 or bfloat16, got {B_type}" + assert B_sem in _DENSE_DTYPES, ( + f"tcgen05 schedule expected B dtype in {_DENSE_DTYPES}, got {B_sem}" ) - assert A_type == B_type, ( - f"tcgen05 schedule expect A_type and B_type to be the same, got A_type={A_type}, B_type={B_type}" # noqa: E501 + assert A_sem == B_sem, ( + f"tcgen05 schedule expect A and B MMA dtype to be the same, got A={A_sem}, B={B_sem}" ) # Parse SFA/SFB and transA/transB/accum based on arg layout @@ -422,7 +486,7 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall, sctx: DispatchContext) - f"tcgen05 block-scaled schedule expected SFA_scope=tmem, SFB_scope=tmem, " f"got SFA_scope={SFA_scope}, SFB_scope={SFB_scope}" ) - SFA_type, SFB_type = SFA_buffer.dtype, SFB_buffer.dtype + SFA_type, SFB_type = _dtype_name(SFA_buffer.dtype), _dtype_name(SFB_buffer.dtype) SFA_slice_layout = SFA_buffer.layout.slice(SFA_buffer.shape, SFA_buffer_region.region) SFB_slice_layout = SFB_buffer.layout.slice(SFB_buffer.shape, SFB_buffer_region.region) SFA_elem_per_col = 32 // DataType(SFA_type).bits @@ -538,8 +602,15 @@ def _try_atom(atom, atom_shape): tiler_shape = [s // a for s, a in zip(shape_2d, atom_shape)] tiler_grouped, seps = tiler.canonicalize().group(tiler_shape) elem_per_128b = 128 // tvm.DataType(dtype).bits - ldo = (tiler_grouped.shard[-1].stride * atom_size) // elem_per_128b - sdo = (tiler_grouped.shard[-2].stride * atom_size) // elem_per_128b + + # extent==1 leading dim -> unused LBO/SBO offset. + def _atom_off(dim): + if int(dim.extent) == 1: + return 0 + return (dim.stride * atom_size) // elem_per_128b + + ldo = _atom_off(tiler_grouped.shard[-1]) + sdo = _atom_off(tiler_grouped.shard[-2]) return mode, ldo, sdo for mode in ( @@ -686,12 +757,13 @@ def _try_atom(atom, atom_shape): # transA=False [M, K]: K = dim[-1]; transA=True [K, M]: K = dim[-2] K = A_dim2 if transA else A_dim1 - # tcgen05 MMA hardware constraints - # K dimension per MMA iteration depends on A/B dtype - if A_type == "float4_e2m1fn": + # tcgen05 MMA hardware constraints (MMA_K keyed on semantic dtype A_sem). + if A_sem == "float4_e2m1fn": MMA_K = 64 - elif A_type in ["float8_e4m3fn", "float8_e5m2"]: + elif A_sem in ["float8_e4m3fn", "float8_e5m2"]: MMA_K = 32 + elif A_sem in ["tensor_float32", "tf32"]: + MMA_K = 8 else: # float16, bfloat16 MMA_K = 16 MMA_N_MIN = 8 if cta_group == 1 else 16 # Minimum N dimension @@ -789,15 +861,15 @@ def _try_atom(atom, atom_shape): if not a_is_tmem: A_elem_per_16B = 128 // DataType(A_type).bits - # Allocate descriptor cells and encode once, right after A/B buffer defs. - # The callback is inserted as a flat SeqStmt after the target buffer def. - # Descriptors with identical construction parameters are cached and reused - # across dispatch calls via sctx.shared_state. - B_base = [0] * len(B_buffer.shape) - krp = Evaluate(tirx_op.tvm_kernel_replace_point()) + issue_pred = op_call.config.get("pred", None) + if issue_pred is None and warp_scope: + issue_pred = T.ptx.elect_sync() + elect_pred = True if issue_pred is None else issue_pred + + _SWIZZLE_TO_LAYOUT = {0: 0, 1: 6, 2: 4, 3: 2, 4: 1} + _krp = Evaluate(tirx_op.tvm_kernel_replace_point()) def _make_lo_uniform(desc): - """Shuffle the lower 32 bits of the descriptor to ensure warp-uniformity.""" func_name = "smem_desc_make_lo_uniform_" source_code = f""" __forceinline__ __device__ void {func_name}(uint64_t* desc) {{ @@ -809,56 +881,83 @@ def _make_lo_uniform(desc): func_name, T.address_of(desc), source_code=source_code, return_type="void" ) - def _make_desc_wrap(desc_buf, smem_buf, base, ldo, sdo, swizzle_val): - """Build: { AllocBuffer(desc); encode(desc, smem); krp }""" + def _make_desc(smem_buf, ldo, sdo, swizzle_val, name): + cache_key = f"gemm_smem_desc:{smem_buf.name}:{int(ldo)}:{int(sdo)}:{int(swizzle_val)}" + cached = sctx.cache_get(cache_key) + if cached is not None: + return cached + + desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name=name, scope="local") encode_call = tvm.tirx.call_intrin( "", "tirx.ptx.tcgen05_encode_matrix_descriptor", tvm.tirx.address_of(desc_buf[0]), - smem_buf.ptr_to(base), + smem_buf.ptr_to([0] * len(smem_buf.shape)), ldo, sdo, swizzle_val, ) - return SeqStmt( + wrap = SeqStmt( [ AllocBuffer(desc_buf), Evaluate(encode_call), Evaluate(_make_lo_uniform(desc_buf[0])), - krp, + _krp, ] ) - - # Per-dispatch-call descriptor (no kernel-scope cache). Each gemm_async - # call allocates + encodes its own ``alignas(64) uint64_t descX[1]`` - # right after the smem buffer definition. Without the previous cache the - # descriptor's lifetime is bounded by the surrounding loop scope rather - # than the entire kernel, which lets ptxas free the register sooner and - # reduces register pressure on the fa4 hot path. The descriptor base is - # the buffer origin (stage=0); the per-MMA operand still adds the - # stage-dependent offset via ``smem_desc_add_16B_offset``. - def _make_desc(smem_buf, base, ldo, sdo, swizzle_val, name): - desc_buf = tvm.tirx.decl_buffer((1,), "uint64", name=name, scope="local") - wrap = _make_desc_wrap(desc_buf, smem_buf, base, ldo, sdo, swizzle_val) sctx.add_post_buffer_def_stmt(smem_buf, wrap) + sctx.cache_set(cache_key, desc_buf) return desc_buf - B_base = [0] * len(B_buffer.shape) - descB_buf = _make_desc(B_buffer, B_base, B_ldo, B_sdo, B_swizzle_mode.value, "descB") - if not a_is_tmem: - A_base = [0] * len(A_buffer.shape) - descA_buf = _make_desc(A_buffer, A_base, A_ldo, A_sdo, A_swizzle_mode.value, "descA") - elect_pred = T.ptx.elect_sync() if warp_scope else True + def _uniform_desc(smem_buf, off16, ldo, sdo, swizzle): + layout = _SWIZZLE_TO_LAYOUT[int(swizzle)] + const_hi = (int(sdo) & 0x3FFF) | (1 << 14) | (layout << 29) + lo_const = (int(ldo) & 0x3FFF) << 16 + base_ptr = smem_buf.ptr_to([0] * len(smem_buf.shape)) + addr = T.ptr_byte_offset(base_ptr, off16 * 16, smem_buf.dtype) + sa = T.bitwise_and( + T.shift_right(T.cuda.cvta_generic_to_shared(addr), T.uint32(4)), + T.uint32(0x3FFF), + ) + 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")) - # Helper: compute B descriptor value for a given (ni, ki) tile - def _b_desc_val(descB_in, ni, ki): + def _b_offset(ni, ki): B_linear = ( ki * MMA_K * B_extent[-1] + ni * N_mma_per_cta if transB else ni * N_mma_per_cta * B_extent[-1] + ki * MMA_K ) - B_offset = tvm.tirx.floordiv(B_slice_tile.apply(B_linear)["m"], B_elem_per_16B) - return smem_desc_add_16B_offset(descB_in, B_offset) + return tvm.tirx.floordiv(B_slice_tile.apply(B_linear)["m"], B_elem_per_16B) + + def _a_offset(mi, ki): + A_linear = ( + ki * MMA_K * A_extent[-1] + mi * M_mma + if transA + else mi * M_mma * A_extent[-1] + ki * MMA_K + ) + 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. + smem_desc_mode = op_call.config.get("smem_desc", "hoist") + if smem_desc_mode not in ("hoist", "recompute"): + raise ValueError(f"Unsupported tcgen05 smem_desc mode: {smem_desc_mode}") + use_add = smem_desc_mode != "recompute" + descB_buf = ( + _make_desc(B_buffer, B_ldo, B_sdo, B_swizzle_mode.value, "descB") if use_add 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 + ) + B_use_add = use_add + + # Helper: compute B descriptor value for a given (ni, ki) tile + 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) + return _uniform_desc(B_buffer, B_offset, B_ldo, B_sdo, B_swizzle_mode.value) # Helper: compute A operand (TMEM address or SMEM descriptor) for a given (mi, ki) tile def _a_operand(mi, ki, descA_in=None): @@ -867,14 +966,10 @@ def _a_operand(mi, ki, descA_in=None): a_row = 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) - else: - A_linear = ( - ki * MMA_K * A_extent[-1] + mi * M_mma - if transA - else mi * M_mma * A_extent[-1] + ki * MMA_K - ) - A_offset = tvm.tirx.floordiv(A_slice_tile.apply(A_linear)["m"], A_elem_per_16B) - return smem_desc_add_16B_offset(descA_in, A_offset) + A_offset = _a_offset(mi, ki) + if A_use_add: + return smem_desc_add_16B_offset(descA_buf[0], A_offset) + return _uniform_desc(A_buffer, A_offset, A_ldo, A_sdo, A_swizzle_mode.value) if is_block_scaled: # Compute per-ki SF element steps from region extents @@ -884,21 +979,36 @@ 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) - ) - SFB_init_addr = analyzer.simplify( - sfb_base + tvm.tirx.floordiv(sfb_tcol_0, SFB_elem_per_col) - ) + sf_ids = set() + for mi in range(M_tiles): + for ki in range(K_iters): + sfa_linear = mi * M_mma * SFA_K_total + ki * sfa_elems_per_ki + sfa_tcol = analyzer.simplify(SFA_slice_layout.apply(sfa_linear).get("TCol", 0)) + sf_ids.add(int(analyzer.simplify(tvm.tirx.floormod(sfa_tcol, SFA_elem_per_col)))) + fixed_sf_id = next(iter(sf_ids)) if len(sf_ids) == 1 else None + + def _patch_instr_desc(desc, sf_id): + return analyzer.simplify( + T.bitwise_or( + T.bitwise_and(desc, T.uint32(0x9FFFFFCF)), + T.bitwise_or( + T.shift_left(T.cast(sf_id, "uint32"), T.uint32(29)), + T.shift_left(T.cast(sf_id, "uint32"), T.uint32(4)), + ), + ) + ) - # Determine if sf_id rotation is needed: - # sf_mma_k < epc means multiple ki's pack in one column, AND we need per-ki - # distinct SF (i.e. sfa_elems_per_ki > 0 so each ki advances to a new element) - needs_sf_id = sfa_sf_mma_k < SFA_elem_per_col and sfa_elems_per_ki > 0 and descI is None + def _instr_desc_value(desc, sfa_tcol): + if fixed_sf_id is not None: + return desc + sf_id = analyzer.simplify(tvm.tirx.floormod(sfa_tcol, SFA_elem_per_col)) + return _patch_instr_desc(desc, sf_id) + + else: + fixed_sf_id = None + + def _instr_desc_value(desc, sfa_tcol): + return desc # Physical TMEM columns per MMA N tile. # 2x2 layout (Layout B): each MMA tile spans N_mma/2 physical columns @@ -913,30 +1023,35 @@ 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 = _a_operand(mi, ki, descA_in) - descB_val = _b_desc_val(descB_in, ni, ki) - should_accum = tvm.tirx.any(ki != 0, accum_expr) + # meta_var inlines operands into mma.block_scale (avoids LMEM temps). + 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)) 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 - sfa_tcol = SFA_slice_layout.apply(sfa_linear).get("TCol", 0) - sfb_tcol = SFB_slice_layout.apply(sfb_linear).get("TCol", 0) - sfa_addr = sfa_base + tvm.tirx.floordiv(sfa_tcol, SFA_elem_per_col) - sfb_addr = sfb_base + 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 - T.cuda.runtime_instr_desc(T.address_of(descI_in), sf_id) - tmem_col = tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) - if elect_pred: - T.ptx.tcgen05.mma.block_scale( - T.cuda.get_tmem_addr(tmem_addr, mi * M_mma, tmem_col), - a_val, descB_val, - sfa_addr, sfb_addr, - descI_in, - d_dtype=C_type, a_dtype=A_type, b_dtype=B_type, - sfa_dtype=SFA_type, sfb_dtype=SFB_type, - use_a_tmem=a_is_tmem, cta_group=cta_group, - enable_input_d=should_accum, - ) + # 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 + sfa_addr = T.meta_var( + analyzer.simplify(sfa_base + 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)) + ) + instr_desc = T.meta_var(_instr_desc_value(descI_in, sfa_tcol)) + tmem_col = T.meta_var( + tmem_offset_32b + ni * (N_mma_phys_cols // C_elem_per_32b) + ) + T.ptx.tcgen05.mma.block_scale( + T.cuda.get_tmem_addr(tmem_addr, mi * M_mma, tmem_col), + a_val, descB_val, + sfa_addr, sfb_addr, + instr_desc, + d_dtype=C_type, a_dtype=A_type, b_dtype=B_type, + sfa_dtype=SFA_type, sfb_dtype=SFB_type, + use_a_tmem=a_is_tmem, cta_group=cta_group, + enable_input_d=should_accum, pred=issue_pred, + ) else: # Wrap each per-MMA operand in ``T.meta_var`` so the parser inlines # the value directly into the ``T.ptx.tcgen05.mma`` call instead of @@ -961,25 +1076,46 @@ def main_impl(descA_in, descB_in, descI_in): 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_type, b_dtype=B_type, + 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, ) - descA_val = None if a_is_tmem else descA_buf[0] + @T.inline + def invoke_main(descI_value): + main_impl(None, None, descI_value) + + if descI is not None and is_block_scaled: + block_desc = _patch_instr_desc(descI, fixed_sf_id) if fixed_sf_id is not None else descI - if descI is not None: @T.prim_func(check_well_formed=False) def impl(): - main_impl(descA_val, descB_buf[0], descI) + descI_local: T.uint32 + descI_local = block_desc + invoke_main(descI_local) + elif descI is not None: + @T.prim_func(check_well_formed=False) + def impl(): + invoke_main(descI) elif is_block_scaled: + descI_value = _encode_instr_descriptor_block_scaled_uint32( + M=M_mma * cta_group, + N=N_mma, + a_dtype=A_type, + b_dtype=B_type, + sf_dtype=SFA_type, + trans_a=a_mn_major, + trans_b=b_mn_major, + ) + if fixed_sf_id is not None: + descI_value = (descI_value & 0x9FFFFFCF) | (fixed_sf_id << 29) | (fixed_sf_id << 4) + descI_const = tvm.tirx.const(descI_value, "uint32") + @T.prim_func(check_well_formed=False) def impl(): descI_local: T.uint32 - 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, descB_buf[0], descI_local) # noqa: F821 + descI_local = descI_const + invoke_main(descI_local) 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 @@ -990,8 +1126,8 @@ def impl(): M=M_mma * cta_group, N=N_mma, d_dtype="float32", - a_dtype=A_type, - b_dtype=B_type, + a_dtype=A_sem, + b_dtype=B_sem, trans_a=a_mn_major, trans_b=b_mn_major, ) @@ -999,7 +1135,7 @@ def impl(): @T.prim_func(check_well_formed=False) def impl(): - main_impl(descA_val, descB_buf[0], descI_const) + invoke_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/permute_layout/warp_xor_swizzle.py b/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py index 8bee97aae546..2646d546606d 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 Buffer, BufferRegion, IntImm, PrimFunc 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 @@ -305,28 +305,70 @@ def _project(iter_idx, st_list): tid_x = sctx.launch_params["threadIdx.x"] dtype = src_buf.dtype + # Shared 32/64b: base ptr + stride offset avoids buf[] flatten IMAD path. + direct = ( + dtype_bytes in (4, 8) + and "shared" in str(src_buf.scope()) + and "shared" in str(dst_buf.scope()) + ) + ptx_t = f"b{dtype_bytes * 8}" + # ld/st only move bits, so use an unsigned container of the matching width: + # ``ptx.ld(..., bN)`` rejects a float return dtype, and the permute is a pure + # byte shuffle, so a float32/float64 tile loads/stores correctly as uint. + bits_dtype = f"uint{dtype_bytes * 8}" + + def _iter_off(iter_idx, strides): + return sum(iter_idx[d] * strides[d] for d in range(len(strides))) + # fmt: off - @T.prim_func - def impl(): - warp_size = T.meta_var(32) - lane_id = T.meta_var(tid_x % warp_size) - regs = T.alloc_buffer((P,), dtype, scope="local") - # Phase 1: read via L_src - for r in T.unroll(0, P): - j = T.meta_var(r ^ ((lane_id >> shift) & mask)) - flat = T.meta_var(lane_id + j * warp_size) - iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) - src_idx = T.meta_var(_project(iter_idx, src_st)) - regs[r] = src_buf[tuple(src_idx)] - T.cuda.warp_sync() - # Phase 2: write via L_dst - for r in T.unroll(0, P): - j = T.meta_var(r ^ ((lane_id >> shift) & mask)) - flat = T.meta_var(lane_id + j * warp_size) - iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) - dst_idx = T.meta_var(_project(iter_idx, dst_st)) - dst_buf[tuple(dst_idx)] = regs[r] - T.cuda.warp_sync() + if direct: + @T.prim_func + def impl(): + warp_size = T.meta_var(32) + lane_id = T.meta_var(tid_x % warp_size) + regs = T.alloc_buffer((P,), bits_dtype, scope="local") + base_src = T.meta_var(src_buf.ptr_to(list(src_st))) + base_dst = T.meta_var(dst_buf.ptr_to(list(dst_st))) + # Phase 1: read via L_src + for r in T.unroll(0, P): + j = T.meta_var(r ^ ((lane_id >> shift) & mask)) + flat = T.meta_var(lane_id + j * warp_size) + iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) + off = T.meta_var(_iter_off(iter_idx, src_str_)) + ptr = T.meta_var(T.ptr_byte_offset(base_src, off * dtype_bytes, dtype)) + regs[r] = T.ptx.ld(ptr, bits_dtype, ptx_t, space="shared") + T.cuda.warp_sync() + # Phase 2: write via L_dst + for r in T.unroll(0, P): + j = T.meta_var(r ^ ((lane_id >> shift) & mask)) + flat = T.meta_var(lane_id + j * warp_size) + iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) + off = T.meta_var(_iter_off(iter_idx, dst_str_)) + ptr = T.meta_var(T.ptr_byte_offset(base_dst, off * dtype_bytes, dtype)) + T.evaluate(T.ptx.st(ptr, regs[r], space="shared", ptx_type=ptx_t)) + T.cuda.warp_sync() + else: + @T.prim_func + def impl(): + warp_size = T.meta_var(32) + lane_id = T.meta_var(tid_x % warp_size) + regs = T.alloc_buffer((P,), dtype, scope="local") + # Phase 1: read via L_src + for r in T.unroll(0, P): + j = T.meta_var(r ^ ((lane_id >> shift) & mask)) + flat = T.meta_var(lane_id + j * warp_size) + iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) + src_idx = T.meta_var(_project(iter_idx, src_st)) + regs[r] = src_buf[tuple(src_idx)] + T.cuda.warp_sync() + # Phase 2: write via L_dst + for r in T.unroll(0, P): + j = T.meta_var(r ^ ((lane_id >> shift) & mask)) + flat = T.meta_var(lane_id + j * warp_size) + iter_idx = T.meta_var(get_indices(flat, [0] * len(extent), extent)) + dst_idx = T.meta_var(_project(iter_idx, dst_st)) + dst_buf[tuple(dst_idx)] = regs[r] + T.cuda.warp_sync() # fmt: on return impl 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/script.py b/python/tvm/backend/cuda/script.py index 5016b4d5794a..9eb724839ee7 100644 --- a/python/tvm/backend/cuda/script.py +++ b/python/tvm/backend/cuda/script.py @@ -58,7 +58,9 @@ def __init__(self): self.fetch_register: Callable[..., Any] = _op_wrapper(_cuda_op.ptx_fetch_register) self.ld = _op_wrapper(_cuda_op.ptx_ld) 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) + self.ld_mmio = _op_wrapper(_cuda_op.ptx_ld_mmio) self.ld_global_acquire = _op_wrapper(_cuda_op.ptx_ld_global_acquire) self.red_scalar = _op_wrapper(_cuda_op.ptx_red_scalar) self.atom_scalar = _op_wrapper(_cuda_op.ptx_atom_scalar) @@ -69,6 +71,10 @@ def __init__(self): self.cp_async_bulk_s2s_cluster = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2s_cluster) self.cp_async_bulk_s2g = _op_wrapper(_cuda_op.ptx_cp_async_bulk_s2g) self.st = _op_wrapper(_cuda_op.ptx_st) + self.st_relaxed = _op_wrapper(_cuda_op.ptx_st_relaxed) + self.st_release = _op_wrapper(_cuda_op.ptx_st_release) + self.st_volatile = _op_wrapper(_cuda_op.ptx_st_volatile) + self.st_mmio = _op_wrapper(_cuda_op.ptx_st_mmio) self.st_bulk = _op_wrapper(_cuda_op.ptx_st_bulk) self.fns_b32 = _op_wrapper(_cuda_op.ptx_fns_b32) self.add_rn_f32_bf16 = _op_wrapper(_cuda_op.ptx_add_rn_f32_bf16) @@ -421,12 +427,6 @@ def __init__(self): self.cta_sum = _op_wrapper(_cuda_op.cuda_cta_sum) self.cta_max = _op_wrapper(_cuda_op.cuda_cta_max) self.cta_min = _op_wrapper(_cuda_op.cuda_cta_min) - self.copy_bytes = _op_wrapper(_cuda_op.cuda_copy_bytes) - self.copy_128b = _op_wrapper(_cuda_op.cuda_copy_128b) - self.copy_64b = _op_wrapper(_cuda_op.cuda_copy_64b) - self.copy_32b = _op_wrapper(_cuda_op.cuda_copy_32b) - self.copy_16b = _op_wrapper(_cuda_op.cuda_copy_16b) - self.copy_8b = _op_wrapper(_cuda_op.cuda_copy_8b) self.cta_sync = _op_wrapper(_cuda_op.cuda_cta_sync) self.grid_sync = _op_wrapper(_cuda_op.cuda_grid_sync) self.cluster_sync = _op_wrapper(_cuda_op.cuda_cluster_sync) 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 9546c2f6e43f..036c6127802f 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 123698b62164..ddaa2c5239ee 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 eba660ab6d6a..324f419b8d25 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 830ab45f2d11..3507c6dba2c1 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,18 @@ 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: + dtype = getattr(scalar, "dtype", None) + if dtype is not None: + return str(dtype) + ty = getattr(scalar, "ty", None) + dtype = getattr(ty, "dtype", None) + if dtype is None: + raise AttributeError(f"{type(scalar).__name__} has no dtype-bearing PrimType") + return str(dtype) def alloc_const_bias_trn( @@ -53,7 +63,9 @@ def alloc_const_bias_trn( return {"const_bias": ("const_bias", bias.value)} else: new_shape = (par_size, max_inst_size) - new_buffer = T.buffer(new_shape, dtype=bias.dtype, scope="trn.sbuf", buffer_name="const_bias") + new_buffer = T.buffer( + new_shape, dtype=_scalar_dtype(bias), scope="trn.sbuf", buffer_name="const_bias" + ) @T.prim_func def const_bias_init(): 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 43bd92d56d75..9c0dc698528f 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/support/nvcc.py b/python/tvm/support/nvcc.py index b421042fb30b..84976dfb15b2 100644 --- a/python/tvm/support/nvcc.py +++ b/python/tvm/support/nvcc.py @@ -20,6 +20,7 @@ import glob import os import platform +import shlex import subprocess import warnings @@ -31,6 +32,29 @@ from . import utils +def _ptxas_option_flags(): + """Return ptxas flags forwarded via ``--ptxas-options`` (without the prefix). + + Environment Variables + --------------------- + TVM_CUDA_PTXAS_REG_LEVEL : str + ptxas ``--register-usage-level`` (default ``10``). + TVM_CUDA_PTXAS_EXTRA_OPTS : str + Extra ptxas flags, shell-tokenized (e.g. ``"-O1"`` or ``"-O2 --def-load-cache=ca"``). + Each token becomes its own ``--ptxas-options=`` entry for NVRTC, or is + comma-joined for nvcc. + """ + flags = [ + "-v", + f"--register-usage-level={os.environ.get('TVM_CUDA_PTXAS_REG_LEVEL', '10')}", + "--warn-on-local-memory-usage", + ] + extra = os.environ.get("TVM_CUDA_PTXAS_EXTRA_OPTS", "").strip() + if extra: + flags.extend(shlex.split(extra)) + return flags + + def compile_cuda( code, target_format=None, arch=None, options=None, path_target=None, compiler="nvrtc" ): @@ -190,8 +214,7 @@ def _compile_cuda_nvcc( "--expt-relaxed-constexpr", "--expt-extended-lambda", "--use_fast_math", - "--ptxas-options=-v", # printing out number of registers - f"--ptxas-options=--verbose,--register-usage-level={os.environ.get('TVM_CUDA_PTXAS_REG_LEVEL', '10')},--warn-on-local-memory-usage", # noqa: E501 + f"--ptxas-options={','.join(_ptxas_option_flags())}", ] major, _ = parse_compute_version(get_target_compute_version(Target.current(allow_none=True))) @@ -524,14 +547,8 @@ def _compile_cuda_nvrtc( # NVRTC does not comma-split --ptxas-options, so each ptxas flag must be its # own entry. The nvcc-only --expt-relaxed-constexpr / --expt-extended-lambda # have no NVRTC equivalent and are intentionally not mirrored. - reg_level = os.environ.get("TVM_CUDA_PTXAS_REG_LEVEL", "10") - compile_opts.extend( - [ - b"--ptxas-options=-v", - f"--ptxas-options=--register-usage-level={reg_level}".encode(), - b"--ptxas-options=--warn-on-local-memory-usage", - ] - ) + for flag in _ptxas_option_flags(): + compile_opts.append(f"--ptxas-options={flag}".encode()) # Add user-provided options, filtering out nvcc-specific flags that nvrtc doesn't support if options: @@ -868,6 +885,10 @@ def tvm_callback_cuda_compile(code): TVM_KERNEL_DUMP : str If set, dump generated CUDA/intermediate files and append "-lineinfo" so profilers can correlate SASS back to the dumped source. + TVM_CUDA_PTXAS_REG_LEVEL : str + Forwarded to ptxas ``--register-usage-level`` (default ``10``). + TVM_CUDA_PTXAS_EXTRA_OPTS : str + Extra ptxas flags (shell-tokenized), e.g. ``"-O1"`` or ``"-O2"``. Parameters ---------- diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py index 5e8a2184bdad..57a6aeb36ca6 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -44,7 +44,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 @@ -87,7 +88,6 @@ # 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 .expr_functor import ExprFunctor from . import transform diff --git a/python/tvm/tirx/bench.py b/python/tvm/tirx/bench.py index aa0676375905..6c4e7d123ebe 100644 --- a/python/tvm/tirx/bench.py +++ b/python/tvm/tirx/bench.py @@ -16,12 +16,20 @@ # under the License. import argparse +import gc +import inspect +import json +import math import os import re +import statistics import subprocess import sys +import tempfile import time +import uuid from collections.abc import Mapping +from contextlib import contextmanager from enum import Enum import numpy as np @@ -60,19 +68,24 @@ def tvm_callback_cuda_compile(code, target): _ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") -def _parse_proton_tree(text, value_scale=1.0): - """Parse proton-viewer tree output into {impl: time_ms}. +# 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" + + +def _parse_proton_tree(text, *, kernel: str = ""): + """Parse proton-viewer tree output into {impl: time_us}. Accepts ALL depth-1 nodes (no KNOWN_IMPLS filter). For each depth-1 impl, takes the slowest depth-2 child kernel time. - ``value_scale`` converts the displayed metric to milliseconds. For - example, use ``1e-3`` when parsing ``avg_time/us`` output. + Tree numbers are microseconds when ProtonContext uses avg_time/us. Returns (impl_times, baseline_errors) where: - impl_times: {str: float} — impl name to avg time in ms + 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 = {} @@ -104,7 +117,7 @@ def _parse_proton_tree(text, value_scale=1.0): ): continue try: - t = float(parts[0]) * value_scale + t = float(parts[0]) results[impl] = max(results.get(impl, 0), t) except ValueError: pass @@ -130,15 +143,15 @@ def __init__( hook="triton", debug=False, nsight=False, - metric="avg_time/us", - metric_scale=1e-3, + metric=PROTON_AVG_TIME_METRIC, + kernel="", ): self.name = name self.hook = hook self.debug = debug self.nsight = nsight self.metric = metric - self.metric_scale = metric_scale + self.kernel = kernel self._impl_times = {} self._baseline_errors = {} @@ -161,9 +174,10 @@ def __exit__(self, exc_type, exc_val, exc_tb): ) if result.returncode == 0: self._impl_times, self._baseline_errors = _parse_proton_tree( - result.stdout, value_scale=self.metric_scale + 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( @@ -175,7 +189,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): os.remove(hatchet) def get_impl_times(self): - """Return {impl_name: avg_time_ms} parsed from proton-viewer output.""" + """Return {impl_name: avg_time_us} parsed from proton-viewer output.""" return dict(self._impl_times) def get_baseline_errors(self): @@ -313,16 +327,18 @@ def _bench_event_groups(funcs, groups, warmup, repeat, cooldown_s): end_event.record() torch.cuda.synchronize() - results[name] = start_event.elapsed_time(end_event) / repeat + 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): +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) as ctx: + 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) @@ -353,13 +369,8 @@ def _flush_l2_legacy(flush_l2_size): def _bench_legacy_callable(func, warmup, repeat, proton_name, debug, nsight, flush_l2_size): - for _ in range(warmup): - _flush_l2_legacy(flush_l2_size) - func() - start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() def timed_loop(): start_event.record() @@ -368,6 +379,10 @@ def timed_loop(): 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={}): @@ -375,12 +390,12 @@ def timed_loop(): proton.deactivate() else: timed_loop() - torch.cuda.synchronize() - return start_event.elapsed_time(end_event) / repeat + return start_event.elapsed_time(end_event) / repeat * 1000.0 -def bench( + +def bench_tk( funcs, input_factory=None, warmup=500, @@ -392,40 +407,67 @@ def bench( debug=False, nsight=False, flush_l2_size=int(8e8 // 4), + references=None, + rounds=1, + round_cooldown_s=1.0, + validate_case=None, ): """Benchmark implementations with a factory-owned input footprint. - 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)``. + This is the ThunderKittens-style TIRx benchmark API. It follows the + multi-input protocol for L2 eviction (adapted from ThunderKittens, + https://github.com/HazyResearch/ThunderKittens) 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)``. + + For the Triton-standard, pure-launch benchmark path (``do_bench`` / + ``do_bench_cudagraph`` semantics, no group protocol), use ``bench`` instead. Parameters ---------- funcs : dict[str, callable] Map of implementation name to callable. Each callable receives one - ``case`` returned by ``input_factory``. + ``case`` returned by ``input_factory``. 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. 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. + Number of timed iterations per round. cooldown_s : float Seconds to sleep between impls for thermal cooldown. + 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 tir-bench, ``run_kernel_bench`` holds + the per-GPU lock for the whole ``run_bench()`` call. timer : {"event", "proton"} Timing backend. Returns ------- dict - ``{"impls": {name: ms}, "errors": {}, "timer": ..., ...}``. + ``{"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}``. + Times are stored in microseconds (same unit as 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: + 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") @@ -450,32 +492,79 @@ def bench( if not callable(func): raise TypeError(f"funcs[{name!r}] must be callable") + # ``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] = {} + 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} + inputs, protocol = prepare_input_groups(input_factory, l2_bytes=l2_bytes) num_groups = len(inputs) if num_groups == 0: return { "impls": {}, - "errors": {}, + "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()), }, } - errors = {} - if timer == "event": - impls = _bench_event_groups(funcs, inputs, warmup, repeat, cooldown_s) + 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: - impls, errors = _bench_proton_groups( - funcs, inputs, warmup, repeat, cooldown_s, proton_name, debug, nsight - ) + aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()} return { - "impls": impls, + "impls": aggregated, + "round_samples": round_samples, "errors": errors, "timer": timer, "benchmark_protocol": { @@ -483,11 +572,542 @@ def bench( "warmup": warmup, "repeat": repeat, "cooldown_s": cooldown_s, + "rounds": rounds, + "round_cooldown_s": round_cooldown_s, "order": list(funcs.keys()), }, } +# --------------------------------------------------------------------------- +# Triton-standard benchmark path. +# +# Faithful in-repo port of triton.testing.do_bench / do_bench_cudagraph +# (see https://github.com/triton-lang/triton, python/triton/testing.py). This is +# torch-only and does NOT import or call into triton at runtime: Triton driver +# calls (get_device_interface / get_empty_cache_for_benchmark / clear_cache) are +# replaced with direct torch.cuda + a torch L2-flush buffer. 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 MB buffer whose .zero_() evicts the L2 cache between measured iters. + return torch.empty(int(256e6 // 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"] + + fn() + torch.cuda.synchronize() + + cache = _empty_cache_for_benchmark() + + # 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. + 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)) + 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) + + +def _collect_proton_scope_times(database, prefix): + """Port of triton.testing._collect_proton_scope_times. + + 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 _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``. + + 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). + + Falls back to ``_do_bench_event`` under pytest or when a Proton session cannot be + created. + """ + assert return_mode in ["min", "max", "mean", "median", "all"] + + if is_running_under_pytest(): + return _do_bench_event( + fn, + warmup=warmup, + rep=rep, + grad_to_none=grad_to_none, + quantiles=quantiles, + return_mode=return_mode, + ) + + fn() + torch.cuda.synchronize() + + 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 + + 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)) + + # 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 = proton.start(profile_path, context="shadow", data="tree") + if session is None: + print( + "proton: Proton session unavailable; falling back to event timing", file=sys.stderr + ) + return _do_bench_event( + fn, + warmup=warmup, + rep=rep, + grad_to_none=grad_to_none, + quantiles=quantiles, + return_mode=return_mode, + ) + 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 _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. + + 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). + + Falls back to ``_do_bench_event`` (cold-cache CUDA-event timing) under pytest or + when a Proton session cannot be created -- staying on a cold-cache timer keeps it + consistent with the rest of the baseline. + """ + assert return_mode in ["min", "max", "mean", "median", "all"] + + if is_running_under_pytest(): + return _do_bench_event( + fn, grad_to_none=grad_to_none, quantiles=quantiles, return_mode=return_mode + ) + + 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) + start_event.record() + for _ in range(5): + fn() + end_event.record() + torch.cuda.synchronize() + 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 = proton.start(profile_path, context="shadow", data="tree") + if session is None: + print( + "cudagraph_proton: Proton session unavailable; falling back to event timing", + file=sys.stderr, + ) + return _do_bench_event( + fn, grad_to_none=grad_to_none, quantiles=quantiles, return_mode=return_mode + ) + + 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 bench( + funcs, + *, + warmup=None, + repeat=None, + cudagraph_rep=None, + timer=None, + references=None, + rounds=1, + round_cooldown_s=1.0, +): + """Benchmark pure-launch implementations using Triton-standard timing. + + 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_cudagraph_proton`` time a function. The + timing core is a faithful in-repo port (torch-only, no triton dependency); + see ``_do_bench_event`` / ``_do_bench_cudagraph_proton``. + + For the ThunderKittens-style multi-input group protocol (``input_factory``, + per-workload L2 eviction, Proton per-kernel attribution), use ``bench_tk``. + + Parameters + ---------- + funcs : dict[str, callable] + 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 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`` timer. ``None`` (default) defers + to ``_do_bench_event``'s own default; pass a value only to override. Ignored + by the graph timers (which have no warmup). + repeat : int, optional + Rep time budget (ms) for the ``event`` timer. ``None`` (default) defers to + ``_do_bench_event``'s own default; pass a value only to override. + cudagraph_rep : int, optional + ``rep`` (ms) for the graph timers (graph unroll length). ``None`` (default) + defers to ``_do_bench_cudagraph_proton``'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"} + ``None`` (default) resolves to ``proton`` -- see the resolution in the body; + the default is defined in exactly one place so kernels pass ``None`` to inherit. + All three are **cold-cache** (per-iter L2 flush) so results are comparable. + ``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`` -> ``do_bench_cudagraph_proton`` (graph replay + Proton + + L2 flush); like ``proton`` but also removes launch overhead via graph replay -- + best for tiny back-to-back kernels, but the graph capture fails/misattributes + for some references (CuTeDSL) so it 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. (Plain no-flush ``do_bench_cudagraph`` is intentionally NOT offered.) + rounds : int + 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.) + round_cooldown_s : float + Seconds to sleep between rounds (ignored when ``rounds == 1``). + + Returns + ------- + dict + ``{"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}``. + Times are stored in microseconds (same unit as ``bench_tk`` and the + pinned tir-bench baselines). + """ + # warmup/repeat/cudagraph_rep default to None = "use the timer function's own + # (Triton-aligned) default". They are only forwarded to the timer below when a + # caller explicitly overrides, so the defaults live in exactly one place: the + # _do_bench_* signatures. + 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 rounds < 1: + raise ValueError("rounds must be >= 1") + if round_cooldown_s < 0: + raise ValueError("round_cooldown_s must be non-negative") + # ``timer=None`` means "use the default timer". The default lives in exactly one + # place -- here -- so kernels forward ``timer=None`` to inherit it and a future + # change is a one-line edit. proton = pure per-kernel GPU time; it is honest for + # references with heavy host dispatch (flashinfer, CuTeDSL) where the ``event`` + # wall would over/under-credit us, and matches event within a few % on + # compute-bound kernels. It auto-falls back to ``event`` when Proton/CUPTI is + # unavailable (non-NVIDIA, or under pytest). + 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" + ) + if not isinstance(funcs, Mapping) or not funcs: + 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") + + # ``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] = {} + 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} + + # 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 -- 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, + "rounds": rounds, + "round_cooldown_s": round_cooldown_s, + "order": list(funcs.keys()), + } + + round_samples: dict[str, list[float]] = {} + for round_idx in range(rounds): + if round_idx > 0: + time.sleep(round_cooldown_s) + for name, func in funcs.items(): + ms = _timer_fn(func, **_timer_kwargs) + # ms -> microseconds (matches bench_tk unit and pinned baselines). + round_samples.setdefault(name, []).append(ms * 1000.0) + + aggregated = {impl: statistics.mean(samples) for impl, samples in round_samples.items()} + + return { + "impls": aggregated, + "round_samples": round_samples, + "errors": build_errors, + "timer": timer, + "benchmark_protocol": protocol, + } + + # utils for tg4perfetto profiler, adapted from https://github.com/flashinfer-ai/flashinfer diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py index 43023b4c3cb9..3c9d9c5ce353 100644 --- a/python/tvm/tirx/buffer.py +++ b/python/tvm/tirx/buffer.py @@ -17,6 +17,7 @@ """Abstraction for array data structures.""" import functools +import re from numbers import Integral import tvm_ffi @@ -443,8 +444,18 @@ def permute(self, *dims) -> "Buffer": # 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.layout.group(list(self.shape)) + layout = self.layout + swizzle = None + if isinstance(layout, tvm.tirx.layout.ComposeLayout): + # The swizzle permutes the flat offset, so it commutes with a + # permutation of the logical dims: permute the inner tile layout + # and re-compose. + swizzle = layout.swizzle + layout = layout.tile_layout + grouped, seps = layout.group(list(self.shape)) new_layout = grouped.permute_by_groups(seps, list(dims)) + if swizzle is not None: + new_layout = tvm.tirx.layout.ComposeLayout(swizzle, new_layout) return tvm.tirx.script.builder.decl_buffer( new_shape, self.dtype, @@ -460,6 +471,333 @@ def permute(self, *dims) -> "Buffer": new_layout, ) + # ------------------------------------------------------------------ + # Dimension-surgery views (torch-aligned): unflatten / flatten / + # select / narrow, the numpy-style ``sub`` indexer, and einops-style + # ``rearrange``. All return derived views sharing this buffer's data; + # the physical placement (layout iters + swizzle) is carried, never + # restated, and data is never moved: an operation either yields a + # valid view or raises. + # ------------------------------------------------------------------ + + def _normalized_dim(self, dim, name): + ndim = len(self.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 + + @staticmethod + 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 unflatten(self, dim, sizes) -> "Buffer": + """Split ``dim`` into the given factor sizes (row-major). + + Torch-aligned: ``x.unflatten(1, (4, 4, 4))``; one size may be ``-1`` + and is inferred from the dim's extent. Pure view: the layout is + carried and regrouped against the new shape. The split must be exact: + when the extent and factors are statically known, a non-bijective + factorization is rejected (symbolic values are the caller's + responsibility). + """ + dim = self._normalized_dim(dim, "unflatten") + sizes = list(sizes) + extent = self._concrete_int(self.shape[dim]) + negatives = [i for i, s in enumerate(sizes) if isinstance(s, int) and s == -1] + if len(negatives) > 1: + raise ValueError("unflatten: at most one -1 is allowed in sizes") + for i, size in enumerate(sizes): + if negatives and i == negatives[0]: + continue + size_c = self._concrete_int(size) + if size_c is not None and size_c <= 0: + raise ValueError( + f"unflatten: sizes must be positive (or a single -1); got {size_c}" + ) + if negatives: + known = functools.reduce( + lambda a, b: a * b, + [s for i, s in enumerate(sizes) if i != negatives[0]], + 1, + ) + known_c = self._concrete_int(known) + if extent is not None and known_c is not None: + if known_c <= 0 or extent % known_c != 0: + raise ValueError( + f"unflatten: dim {dim} extent {extent} is not divisible " + f"by the known factors (product {known_c})" + ) + sizes[negatives[0]] = extent // known_c + else: + raw_extent = self.shape[dim] + raw_extent = ( + int(raw_extent) if isinstance(raw_extent, tvm.tirx.IntImm) else raw_extent + ) + sizes[negatives[0]] = raw_extent // known + else: + product = functools.reduce(lambda a, b: a * b, sizes, 1) + product_c = self._concrete_int(product) + if extent is not None and product_c is not None and product_c != extent: + raise ValueError( + f"unflatten: sizes {sizes} multiply to {product_c}, " + f"but dim {dim} has extent {extent}" + ) + new_shape = list(self.shape[:dim]) + sizes + list(self.shape[dim + 1 :]) + return self.view(*new_shape) + + def flatten(self, start_dim=0, end_dim=-1) -> "Buffer": + """Merge dims ``start_dim..end_dim`` into one (torch-aligned). + + Unlike torch this never copies: the merge is expressed in the layout + (a logical axis may carry several strided iters), so it either yields + a view or the downstream consumer rejects the layout loudly. + """ + start = self._normalized_dim(start_dim, "flatten") + end = self._normalized_dim(end_dim, "flatten") + if end < start: + raise ValueError(f"flatten: end_dim {end} < start_dim {start}") + merged = functools.reduce(lambda a, b: a * b, list(self.shape[start : end + 1]), 1) + new_shape = [*self.shape[:start], merged, *self.shape[end + 1 :]] + return self.view(*new_shape) + + def _surgery_parts(self): + """Split ``self.layout`` into (inner tile layout, optional swizzle).""" + layout = self.layout + swizzle = None + if isinstance(layout, tvm.tirx.layout.ComposeLayout): + swizzle = layout.swizzle + layout = layout.tile_layout + return layout, swizzle + + @staticmethod + def _rewrap_swizzle(layout, swizzle): + if swizzle is not None: + return tvm.tirx.layout.ComposeLayout(swizzle, layout) + return layout + + @staticmethod + def _dim_group_offset(group, index): + """Offset of ``index`` along a dim made of the given layout iters. + + The dim's iters are outer→inner in row-major coordinate order: + decompose ``index`` mixed-radix over their extents and weight each + coordinate by its iter's stride. A single-iter dim reduces to + ``index * stride``. ``index`` may be a PrimExpr; multi-iter dims + need concrete inner extents. + """ + if len(group) == 1: + return index * group[0].stride + extents = [] + for it in group[1:]: + extent = it.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, it 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 * it.stride + offset = term if offset is None else offset + term + return offset + + def _swizzle_offset_commutes(self, swizzle, extra_offset): + """Whether ``extra_offset`` commutes with the swizzle permutation. + + Folding an offset into ``elem_offset`` moves it *outside* the layout, + so the address becomes ``offset + swizzle(rest)``. The swizzle XORs + bits within a period of ``2^(per_element + atom_len + swizzle_len)`` + elements, so this equals the true ``swizzle(offset + rest)`` only + when the offset is a multiple of that period. + """ + if swizzle is None or extra_offset is None: + return True + sw_len = int(swizzle.swizzle_len) + if sw_len == 0: + return True # identity permutation, everything commutes + period = 1 << (int(swizzle.per_element) + int(swizzle.atom_len) + sw_len) + offset_c = self._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(self, new_shape, new_shard, grouped, swizzle, extra_offset): + """Rebuild a derived view; ``extra_offset`` goes into ``elem_offset`` + when it commutes with the swizzle (provable period multiple), + otherwise it stays inside the tile layout's offset so the swizzle + keeps applying to it and the view addresses the same bytes.""" + offset_map = dict(grouped.offset.items()) + if extra_offset is not None and not self._swizzle_offset_commutes(swizzle, extra_offset): + m_axis = tvm.tirx.layout.Axis.get("m") + prev = offset_map.get(m_axis) + offset_map[m_axis] = extra_offset if prev is None else prev + extra_offset + extra_offset = None + new_layout = tvm.tirx.layout.TileLayout.from_iters( + new_shard, list(grouped.replica), offset_map + ) + new_layout = self._rewrap_swizzle(new_layout, swizzle) + elem_offset = self.elem_offset + if extra_offset is not None: + elem_offset = elem_offset + extra_offset + return tvm.tirx.script.builder.decl_buffer( + new_shape, + self.dtype, + self.data, + self.strides, + elem_offset, + None, + self.scope(), + self.data_alignment, + self.offset_factor, + "", + self.axis_separators, + new_layout, + ) + + def select(self, dim, index) -> "Buffer": + """Return a view with ``dim`` removed, fixed at ``index``. + + Torch-aligned (``Tensor.select``); ``index`` may be a dynamic + PrimExpr, folded into the view's ``elem_offset`` through the dim's + layout iters. Statically known indices are bounds-checked; dynamic + indices are the caller's responsibility. On a swizzled layout the + offset folds into ``elem_offset`` only when it provably commutes + with the swizzle (a swizzle-period multiple); otherwise it stays + inside the derived layout's offset, where the swizzle applies to it. + """ + dim = self._normalized_dim(dim, "select") + index_c = self._concrete_int(index) + extent_c = self._concrete_int(self.shape[dim]) + if index_c is not None: + if index_c < 0 or (extent_c is not None and index_c >= extent_c): + raise ValueError( + f"select: index {index_c} out of range for dim {dim} " + f"of extent {extent_c if extent_c is not None else self.shape[dim]}" + ) + layout, swizzle = self._surgery_parts() + grouped, seps = layout.group(list(self.shape)) + iters = list(grouped.shard) + lo, hi = seps[dim], seps[dim + 1] + offset = self._dim_group_offset(iters[lo:hi], index) + new_shape = list(self.shape[:dim]) + list(self.shape[dim + 1 :]) + new_shard = iters[:lo] + iters[hi:] + return self._rebuild_view(new_shape, new_shard, grouped, swizzle, offset) + + def narrow(self, dim, start, length) -> "Buffer": + """Return a view of ``dim`` narrowed to ``[start, start + length)``. + + Torch-aligned (``Tensor.narrow``); ``start`` may be a dynamic + PrimExpr when the dim maps to a single layout iter. For dims made of + several iters, ``start`` and ``length`` must be concrete multiples of + the inner iter block so the range stays a contiguous iter prefix. + Statically known bounds are checked (``start >= 0``, ``length >= 1``, + ``start + length <= extent``); dynamic values are the caller's + responsibility. On a swizzled layout the offset folds into + ``elem_offset`` only when it provably commutes with the swizzle + (a swizzle-period multiple); otherwise it stays inside the derived + layout's offset, where the swizzle applies to it. + """ + dim = self._normalized_dim(dim, "narrow") + start_c = self._concrete_int(start) + length_c = self._concrete_int(length) + extent_c = self._concrete_int(self.shape[dim]) + if start_c is not None and start_c < 0: + raise ValueError(f"narrow: start {start_c} must be non-negative") + if length_c is not None and length_c < 1: + raise ValueError(f"narrow: 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"narrow: range [{start_c}, {start_c + length_c}) exceeds " + f"dim {dim} extent {extent_c}" + ) + layout, swizzle = self._surgery_parts() + grouped, seps = layout.group(list(self.shape)) + iters = list(grouped.shard) + lo, hi = seps[dim], seps[dim + 1] + group = iters[lo:hi] + Iter = tvm.tirx.layout.Iter + if len(group) == 1: + it = group[0] + offset = start * it.stride + new_group = [Iter(length, it.stride, it.axis)] + else: + inner = 1 + for it in group[1:]: + extent = it.extent + if not isinstance(extent, tvm.tirx.IntImm): + raise ValueError( + "narrow: dim with multiple layout iters requires concrete iter extents" + ) + inner *= int(extent) + if not (isinstance(start, Integral) and isinstance(length, Integral)): + raise ValueError( + f"narrow: 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"narrow: dim {dim} start/length must be multiples of the " + f"inner iter block {inner}; got start={start} length={length}" + ) + outer = group[0] + offset = (start // inner) * outer.stride + new_group = [Iter(length // inner, outer.stride, outer.axis), *group[1:]] + new_shape = list(self.shape) + new_shape[dim] = length + new_shard = iters[:lo] + new_group + iters[hi:] + return self._rebuild_view(new_shape, new_shard, grouped, swizzle, offset) + + @property + def sub(self) -> "_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 _SubIndexer(self) + + def rearrange(self, pattern, **sizes) -> "Buffer": + """Einops-style dimension rearrangement, e.g. + ``buf.rearrange("b (s w r) c -> b w (s r) c", w=4, r=4)``. + + Pure bijective reshuffling (split / permute / merge) compiled onto + :meth:`unflatten`, :meth:`permute` and :meth:`flatten`; indexing is + out of scope (use :meth:`select` / ``sub``). Composite-axis sizes + come from keyword arguments; at most one factor per composite may be + omitted and is inferred from the dim extent. + """ + return _rearrange(self, pattern, **sizes) + def __getitem__(self, indices): from ..arith import Analyzer # pylint: disable=import-outside-toplevel from .expr import BufferLoad, Ramp # pylint: disable=import-outside-toplevel @@ -511,6 +849,196 @@ def __getitem__(self, indices): return BufferLoad(self, expr_indices) +_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") + + +def _rearrange_parse_side(text: str, pattern: str) -> list[list[str]]: + """Parse one side of a pattern into a list of axis groups. + + Each top-level entry corresponds to one buffer dim: either a single axis + name or a parenthesized composite of axis names. + """ + groups: list[list[str]] = [] + tokens = text.replace("(", " ( ").replace(")", " ) ").split() + current: list[str] | None = None + for token in tokens: + if token == "(": + if current is not None: + raise ValueError(f"rearrange: nested '(' in pattern {pattern!r}") + current = [] + elif token == ")": + if current is None: + raise ValueError(f"rearrange: unmatched ')' in pattern {pattern!r}") + if not current: + raise ValueError(f"rearrange: empty '()' group in pattern {pattern!r}") + groups.append(current) + current = None + else: + if not _NAME_RE.match(token): + raise ValueError(f"rearrange: bad axis name {token!r} in pattern {pattern!r}") + if current is not None: + current.append(token) + else: + groups.append([token]) + if current is not None: + raise ValueError(f"rearrange: unmatched '(' in pattern {pattern!r}") + if not groups: + raise ValueError(f"rearrange: empty side in pattern {pattern!r}") + return groups + + +def _rearrange_flat_names(groups: list[list[str]], pattern: str) -> list[str]: + names = [name for group in groups for name in group] + if len(set(names)) != len(names): + raise ValueError(f"rearrange: duplicate axis name in pattern {pattern!r}") + return names + + +def _rearrange(buffer, pattern, **sizes): + """Apply an einops-style rearrangement pattern to ``buffer`` as a view.""" + if "->" not in pattern: + raise ValueError(f"rearrange: pattern {pattern!r} must contain '->'") + lhs_text, rhs_text = pattern.split("->", 1) + lhs = _rearrange_parse_side(lhs_text, pattern) + rhs = _rearrange_parse_side(rhs_text, pattern) + lhs_names = _rearrange_flat_names(lhs, pattern) + rhs_names = _rearrange_flat_names(rhs, pattern) + if set(lhs_names) != set(rhs_names): + missing = set(lhs_names) ^ set(rhs_names) + raise ValueError( + f"rearrange: axes {sorted(missing)} appear on only one side of {pattern!r}" + ) + unknown = set(sizes) - set(lhs_names) + if unknown: + raise ValueError(f"rearrange: sizes for unknown axes {sorted(unknown)} in {pattern!r}") + if len(lhs) != len(buffer.shape): + raise ValueError( + f"rearrange: pattern {pattern!r} has {len(lhs)} input dims, " + f"buffer has rank {len(buffer.shape)}" + ) + + # Resolve every axis size. Non-composite axes take the dim extent; + # composite factors come from ``sizes`` with at most one inferred. + axis_size: dict[str, object] = dict(sizes) + for dim, group in enumerate(lhs): + extent = buffer.shape[dim] + extent = int(extent) if isinstance(extent, tvm.tirx.IntImm) else extent + if len(group) == 1: + name = group[0] + if name not in axis_size: + axis_size[name] = extent + else: + given_c = Buffer._concrete_int(axis_size[name]) + extent_c = Buffer._concrete_int(extent) + if given_c is not None and extent_c is not None and given_c != extent_c: + raise ValueError( + f"rearrange: size {name}={given_c} does not match dim {dim} " + f"extent {extent_c} in {pattern!r}" + ) + continue + unknown_factors = [name for name in group if name not in axis_size] + if len(unknown_factors) > 1: + raise ValueError( + f"rearrange: composite {'(' + ' '.join(group) + ')'} in {pattern!r} " + f"has multiple unknown factors {unknown_factors}; pass their sizes" + ) + if unknown_factors: + known = 1 + for name in group: + if name != unknown_factors[0]: + known = known * axis_size[name] + known_c = Buffer._concrete_int(known) + extent_c = Buffer._concrete_int(extent) + if ( + extent_c is not None + and known_c is not None + and (known_c <= 0 or extent_c % known_c != 0) + ): + raise ValueError( + f"rearrange: composite {'(' + ' '.join(group) + ')'} in " + f"{pattern!r} does not factor dim extent {extent_c} " + f"(known factors multiply to {known_c})" + ) + axis_size[unknown_factors[0]] = extent // known + + # 1. Split composites, right-to-left so dim indices stay valid. + result = buffer + for dim in range(len(lhs) - 1, -1, -1): + group = lhs[dim] + if len(group) > 1: + result = result.unflatten(dim, tuple(axis_size[name] for name in group)) + + # 2. Permute flat axes into the output order. + perm = [lhs_names.index(name) for name in rhs_names] + if perm != list(range(len(perm))): + result = result.permute(*perm) + + # 3. Merge output composites, right-to-left over flat positions. + position = len(rhs_names) + for group in reversed(rhs): + position -= len(group) + if len(group) > 1: + result = result.flatten(position, position + len(group) - 1) + return result + + +class _SubIndexer: + """Indexer object returned by :attr:`Buffer.sub`. + + Translates numpy basic indexing into a chain of + :meth:`Buffer.select` / :meth:`Buffer.narrow` / + :meth:`Buffer.unflatten` view operations. + """ + + 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 = buf.narrow(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 + # a::s over N = unflatten into (N//s, s) and fix the + # remainder coordinate at a (requires 0 <= a < s). + start_c = Buffer._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})" + ) + buf = buf.unflatten(dim, (extent // step, step)).select(dim + 1, start) + dim += 1 + else: + buf = buf.select(dim, item) + return buf + + def decl_buffer( shape, dtype=None, diff --git a/python/tvm/tirx/layout.py b/python/tvm/tirx/layout.py index 11d1e140ae16..7ca6c9e7c0ad 100644 --- a/python/tvm/tirx/layout.py +++ b/python/tvm/tirx/layout.py @@ -984,14 +984,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 | PrimExpr | 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 | PrimExpr | int): return _LayoutSpec( shard=self.shard, replica=self.replica, offset=_merge_offset(other, self.offset) ) diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 9a54e915bb0b..d35fb9680597 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -934,6 +934,8 @@ def print_buffer(buffer_var, dtype, is_string, is_scalar, dim_num, *shape): final_shape_args = list(shape[0]) else: final_shape_args = list(shape) + if isinstance(dtype, tvm.ir.PrimType): + dtype = dtype.dtype return _ffi_api.print_buffer( buffer_var, dtype, is_string, is_scalar, dim_num, *final_shape_args ) 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 b6bfad133329..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, PrimExpr] - 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 7455a1ae7456..ec3a07a0fd3f 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 PrimExpr -from tvm.tirx.stmt import TilePrimitiveCall +from tvm.tirx.tile_primitive import TilePrimitiveCall def get_tirx_op(op_name: str): 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 78d1c0c3b8ed..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 PrimExpr, Var - -from . import _ffi_api - - -@register_object("tirx.Predicate") -class Predicate(Object): - """A predicate object for TIRX""" - - vars: list[Var] - pred: PrimExpr - - def __init__(self, f_pred: Callable[..., PrimExpr]): - 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[PrimExpr]) -> PrimExpr: - """Apply the predicate to the given indices""" - return _ffi_api.PredicateApply(self, indices) diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 12db12aa99db..57c6401979c7 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -140,6 +140,23 @@ def _get_layout(layout: str | Layout | None, shape: list[PrimExpr], scope: str) return layout +def _normalize_prim_type(dtype) -> ir.PrimType: + if isinstance(dtype, ir.PrimType): + return dtype + dtype_str = getattr(dtype, "_dtype_str", None) + if dtype_str is not None: + return ir.PrimType(dtype_str) + if callable(dtype): + value = dtype() + 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) + + def _get_elem_offset(elem_offset, byte_offset, dtype: str): assert elem_offset is None or byte_offset is None, ( "elem_offset and byte_offset cannot be set at the same time" @@ -148,7 +165,7 @@ def _get_elem_offset(elem_offset, byte_offset, dtype: str): return elem_offset if byte_offset is None: return None - return byte_offset * 8 // (DataType(dtype).bits) + return byte_offset * 8 // (DataType(_normalize_prim_type(dtype).dtype).bits) _block_name_suffix = threading.local() @@ -1801,6 +1818,7 @@ def decl_buffer( strides = [Var(s, "int32") if isinstance(s, str) else s for s in strides] else: strides = [] + dtype = _normalize_prim_type(dtype) decl_frame = _ffi_api.DeclBuffer( # type: ignore[attr-defined] # pylint: disable=no-member shape, dtype, diff --git a/python/tvm/tirx/script/builder/tirx.py b/python/tvm/tirx/script/builder/tirx.py index f2d211d6485d..31b44ebd9e87 100644 --- a/python/tvm/tirx/script/builder/tirx.py +++ b/python/tvm/tirx/script/builder/tirx.py @@ -21,11 +21,10 @@ import tvm.tirx.operator as tirx_op from tvm.ir import Op -from tvm.tirx import Buffer, BufferRegion, PrimExpr +from tvm.tirx import Buffer, BufferRegion, LambdaExpr, PrimExpr 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 . import _ffi_api, frame from .ir import decl_buffer, meta_class @@ -1513,7 +1512,7 @@ def select( dst: BufferRegion | Buffer, true_value: BufferRegion | Buffer | FloatImm, false_value: BufferRegion | Buffer | FloatImm, - pred: Predicate | Callable[..., PrimExpr], + pred: LambdaExpr | Callable[..., PrimExpr], scope: ExecScope | None = None, ): """Select between two values based on a predicate. @@ -1529,7 +1528,7 @@ def select( false_value : Union[BufferRegion, Buffer, FloatImm] The value to select if the predicate is false. - pred : Union[Predicate, Callable[..., PrimExpr]] + pred : Union[LambdaExpr, Callable[..., PrimExpr]] 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 +1536,8 @@ def select( true_value = _to_region(true_value) if isinstance(false_value, Buffer): 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)) diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 543ff99fed66..ec561950d6a5 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -29,22 +29,18 @@ from collections.abc import Mapping from enum import IntEnum -from typing import TYPE_CHECKING, Any, ClassVar import tvm_ffi -from tvm.ir import Op, PrimExpr, Range, Span +from tvm.ir import PrimExpr, Range, Span 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): @@ -913,181 +909,4 @@ def stmt_list(stmt: Stmt) -> list[Stmt]: for x in stmt: res += stmt_list(x) return res - return [stmt] - - -def normalize_const_arg(arg) -> PrimExpr: - 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[PrimExpr] - 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[PrimExpr] - workspace: dict[str, Buffer] - config: dict[str, Any] - dispatch: str | None - scope: ExecScope - _registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {} - - def __init__( - self, - *args: list[PrimExpr], - 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[PrimExpr]: - raise NotImplementedError("Subclass must implement this method") - - @property - def dsts(self) -> list[PrimExpr]: - 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 + return [stmt] diff --git a/python/tvm/tirx/tile_primitive.py b/python/tvm/tirx/tile_primitive.py new file mode 100644 index 000000000000..b9d7e90b9feb --- /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 Op, PrimExpr, 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: PrimExpr + + def __init__(self, f_pred: Callable[..., PrimExpr]): + 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[PrimExpr]) -> PrimExpr: + """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, PrimExpr] + 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) -> PrimExpr: + 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[PrimExpr] + 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[PrimExpr] + workspace: dict[str, Buffer] + config: dict[str, Any] + dispatch: str | None + scope: ExecScope + _registry: ClassVar[dict[Op, type["TilePrimitiveCall"]]] = {} + + def __init__( + self, + *args: list[PrimExpr], + 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[PrimExpr]: + raise NotImplementedError("Subclass must implement this method") + + @property + def dsts(self) -> list[PrimExpr]: + 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/const_fold.h b/src/arith/const_fold.h index 94e1219f0851..c006cae63171 100644 --- a/src/arith/const_fold.h +++ b/src/arith/const_fold.h @@ -317,6 +317,22 @@ inline ffi::Optional TryConstFold(PrimExpr a, PrimExpr template <> inline ffi::Optional TryConstFold(PrimExpr a, PrimExpr b) { + const IntImmNode* ua = a.as(); + const IntImmNode* ub = b.as(); + PrimType utype = a.ty(); + if (utype.MatchesCode(DLDataTypeCode::kDLUInt) && utype == b.ty() && ua && ub) { + auto as_uint = [](int64_t value, const PrimType& dtype) -> uint64_t { + uint64_t result = static_cast(value); + if (dtype.bits() < 64) { + result &= (uint64_t{1} << dtype.bits()) - 1; + } + return result; + }; + uint64_t lhs = as_uint(ua->value, utype); + uint64_t rhs = as_uint(ub->value, b.ty()); + TVM_FFI_ICHECK_NE(rhs, 0U) << "Divide by zero"; + return IntImm(utype, static_cast(lhs % rhs)); + } TVM_INDEX_CONST_PROPAGATION({ PrimType result_ty = a.ty(); if (pa && pb) { diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 60305bfe5924..2374c64a005f 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -1187,6 +1187,15 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) { ContainsVscaleCall(y.Eval()) && CanProveGreaterEqual(x.Eval(), 0) && CanProveGreaterEqual(y.Eval(), 0) && CanProve(x.Eval() < y.Eval())); } + + // Unsigned (uint32/uint64): the signed IsIndexType block above is skipped for + // unsigned operands (see the note in the FloorMod handler). Only the + // OVERFLOW-FREE identities are valid here. + PrimType op_ty = op->ty(); + 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 + } return ret; } @@ -1304,6 +1313,21 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) { } } } + + // Unsigned (uint32/uint64) index arithmetic. The signed IsIndexType rewrite + // block above is skipped for unsigned operands: signed-overflow-is-UB lets + // those rules assume no wraparound, which is UNSOUND for unsigned (e.g. + // floormod(x*y, y) -> 0 fails when x*y wraps mod 2^bits). Only the + // OVERFLOW-FREE identities are valid here. + PrimType op_ty = op->ty(); + if (op_ty.MatchesCode(DLDataTypeCode::kDLUInt) && (op_ty.bits() == 32 || op_ty.bits() == 64)) { + 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). + TVM_TRY_REWRITE_IF(floormod(x * c1, c2), ZeroWithTypeLike(x), + c2.Eval()->value != 0 && c1.Eval()->value % c2.Eval()->value == 0); + } return ret; } diff --git a/src/backend/cuda/op/target_builtin.cc b/src/backend/cuda/op/target_builtin.cc index 5c5ad0b12d9e..852cab036a61 100644 --- a/src/backend/cuda/op/target_builtin.cc +++ b/src/backend/cuda/op/target_builtin.cc @@ -178,7 +178,6 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(cuda_bfloat162float, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_clock64, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_cluster_sync, cuda, kOpaque), - TIRX_DEVICE_INTRIN_ALIAS(cuda_copy_bytes, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_cta_reduce, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_cta_sync, cuda, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(cuda_cvta_generic_to_shared, cuda, kOpaque), @@ -278,6 +277,8 @@ 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_mmio, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_relaxed, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ld_volatile, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ldmatrix, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_ldmatrix_legacy, ptx, kOpaque), @@ -305,6 +306,10 @@ const DeviceIntrinsicRegistration kDeviceIntrinsics[] = { TIRX_DEVICE_INTRIN_ALIAS(ptx_setmaxnreg, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_st, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_st_bulk, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_st_mmio, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_st_relaxed, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_st_release, ptx, kOpaque), + TIRX_DEVICE_INTRIN_ALIAS(ptx_st_volatile, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_stmatrix, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_sub_f32, ptx, kOpaque), TIRX_DEVICE_INTRIN_ALIAS(ptx_sub_f32x2, ptx, kOpaque), diff --git a/src/backend/cuda/runtime/cuda_device_api.cc b/src/backend/cuda/runtime/cuda_device_api.cc index 44d1acff4937..08c38a331c56 100644 --- a/src/backend/cuda/runtime/cuda_device_api.cc +++ b/src/backend/cuda/runtime/cuda_device_api.cc @@ -434,12 +434,19 @@ TVM_FFI_STATIC_INIT_BLOCK() { uint32_t tensor_rank = static_cast(raw_tensor_rank); void* tensor_ptr = static_cast(args[arg_cnt++].cast()); - TVM_FFI_ICHECK_EQ(args.size(), 4 + tensor_rank * 4 + 3) - << "cuTensorMapEncodeTiled expects " << 4 + tensor_rank * 4 + 3 << " arguments" + // The base arg list ends with oob_fill_kind. An OPTIONAL trailing int + // ``force_cu_dtype`` (>= 0) overrides the dtype-derived CUtensorMapDataType + // (e.g. 11 == TFLOAT32 for an fp32 gmem buffer feeding a tf32 MMA, so the TMA + // hardware RN-truncates fp32->tf32 on load). -1 / absent = derive from + // tensor_dtype. Omitting it keeps older callers backward-compatible. + int base_arg_count = static_cast(4 + tensor_rank * 4 + 3); + TVM_FFI_ICHECK(args.size() == base_arg_count || args.size() == base_arg_count + 1) + << "cuTensorMapEncodeTiled expects " << base_arg_count + << " (or +1 force_cu_dtype) arguments" << "tensor_map, tensor_dtype, tensor_rank, tensor_ptr, global_shape(" << tensor_rank << "), global_strides(" << tensor_rank - 1 << "), shared_shape(" << tensor_rank << "), shared_strides(" << tensor_rank << "), interleaved_kind, swizzle_kind" - << ", l2_promotion_kind, oob_fill_kind"; + << ", l2_promotion_kind, oob_fill_kind[, force_cu_dtype]"; std::vector global_shape(tensor_rank); std::vector global_strides( @@ -477,6 +484,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { auto swizzle_kind = static_cast(args[arg_cnt++].cast()); auto l2_promotion_kind = static_cast(args[arg_cnt++].cast()); 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_EQ(tensor_dtype.lanes, 1) << "Expect tensor_dtype to have lanes=1, but get " << tensor_dtype; @@ -570,6 +579,12 @@ TVM_FFI_STATIC_INIT_BLOCK() { TVM_FFI_THROW(InternalError) << "Unsupported data type " << ffi::DLDataTypeToString(tensor_dtype); } + // Caller override (e.g. TFLOAT32 == 11 for an fp32 buffer in a tf32 GEMM): + // 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) { + cu_dtype = static_cast(force_cu_dtype); + } auto is_valid_interleave = interleaved_kind == CU_TENSOR_MAP_INTERLEAVE_NONE || interleaved_kind == CU_TENSOR_MAP_INTERLEAVE_16B || diff --git a/src/backend/trn/transform/lower_trainium_layout.cc b/src/backend/trn/transform/lower_trainium_layout.cc index 38a51e930b4d..ebb4e4f35ecd 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/script/printer/doc_printer/python_doc_printer.cc b/src/script/printer/doc_printer/python_doc_printer.cc index 55da056f407a..85b25947dcc1 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/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 95f821be698b..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 66% rename from src/tirx/ir/predicate.cc rename to src/tirx/ir/lambdaexpr.cc index 0e5b6f7dac89..7d2e0e733545 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), [&](const Var& var) { return vmap.Get(var); }); } -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,13 +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/tirx_stmt.cc b/src/tirx/ir/tirx_stmt.cc index e2838fe2efa1..e43841fe5bfd 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 3e383af4264a..01958e0dadae 100644 --- a/src/tirx/op/tirx.cc +++ b/src/tirx/op/tirx.cc @@ -24,7 +24,7 @@ #include #include -#include +#include namespace tvm { namespace tirx { diff --git a/src/tirx/script/builder/ir.cc b/src/tirx/script/builder/ir.cc index a3b024f0d5f5..e6e89227570e 100644 --- a/src/tirx/script/builder/ir.cc +++ b/src/tirx/script/builder/ir.cc @@ -29,7 +29,7 @@ #include #include #include -#include +#include #include "./utils.h" diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc index a4150f5a832e..d52f2cbdadb4 100644 --- a/src/tirx/script/printer/expr.cc +++ b/src/tirx/script/printer/expr.cc @@ -229,9 +229,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) { @@ -242,11 +242,11 @@ LambdaDoc PrintPredicate(const ffi::ObjectRef& pred, const ffi::Array } TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable) - .set_dispatch("", - [](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 { @@ -469,7 +469,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 94bade18d8b1..4ed983650248 100644 --- a/src/tirx/script/printer/stmt.cc +++ b/src/tirx/script/printer/stmt.cc @@ -282,6 +282,16 @@ std::vector FindParentBuffers(const tirx::Buffer& child, const IRD } } } + // obj2info iteration order is unspecified; order candidates by name so the + // printed parent choice (and hence the script) is deterministic across + // print/parse roundtrips. + tvm::ffi::StructuralHash shash; + std::stable_sort(results.begin(), results.end(), + [&shash](const tirx::Buffer& a, const tirx::Buffer& b) { + std::string an(a->name), bn(b->name); + if (an != bn) return an < bn; + return shash(a) < shash(b); + }); return results; } @@ -301,7 +311,8 @@ bool IsDefaultLayout(const ffi::Optional& layout, const ffi::Array */ ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, const AccessPath& p, const IRDocsifier& d, - const tirx::Buffer& parent) { + const tirx::Buffer& parent, + bool require_same_layout) { ffi::Optional parent_doc = d->GetVarDoc(parent); if (!parent_doc.defined()) return std::nullopt; ExprDoc pdoc = parent_doc.value(); @@ -326,67 +337,11 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, c 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.defined()) { @@ -660,6 +615,9 @@ ffi::Optional TryDeclBufferSugarWithParent(const tirx::Buffer& child, c } else if (!child->layout.defined() && !parent->layout.defined()) { 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.defined() && !child_is_default) { kwargs_keys.push_back("layout"); kwargs_values.push_back( @@ -678,7 +636,14 @@ ffi::Optional TryDeclBufferSugar(const tirx::Buffer& child, const Acces const IRDocsifier& d) { auto parents = FindParentBuffers(child, 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 4c787b95fdd8..c26fbe400651 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 @@ -325,6 +324,14 @@ ExprDoc BufferAttn(const tirx::Buffer& buffer, const AccessPath& p, const Frame& */ 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/lower_tirx_cleanup.cc b/src/tirx/transform/lower_tirx_cleanup.cc index 3637e2af42cb..991ad87e070a 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/tile_primitive_dispatch.cc b/src/tirx/transform/tile_primitive_dispatch.cc index 213264b1a2ae..45b0abeb1065 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 diff --git a/tests/python/arith/test_arith_rewrite_simplify.py b/tests/python/arith/test_arith_rewrite_simplify.py index a6887db8bee2..4c560713aa5d 100644 --- a/tests/python/arith/test_arith_rewrite_simplify.py +++ b/tests/python/arith/test_arith_rewrite_simplify.py @@ -751,6 +751,13 @@ class TestFloorModPadded(BaseCompare): ) +def test_uint_floormod_const_fold(): + analyzer = tvm.arith.Analyzer() + expr = flm(8192, T.uint32(128)) + tvm.ir.assert_structural_equal(expr, T.uint32(0)) + assert analyzer.can_prove_equal(expr, 0) + + 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( diff --git a/tests/python/arith/test_arith_simplify.py b/tests/python/arith/test_arith_simplify.py index 556ce32b5b74..8ddd889b251f 100644 --- a/tests/python/arith/test_arith_simplify.py +++ b/tests/python/arith/test_arith_simplify.py @@ -167,6 +167,16 @@ def test_simplify_floor_mod_with_linear_offset(): assert ana.can_prove_equal(tvm.tirx.floormod(expr1, divisor2), 0) +def test_simplify_uint_floormod_const_scale_divisible(): + """uint32 floormod(x * c1, c2) -> 0 when c1 % c2 == 0 (overflow-free).""" + ana = tvm.arith.Analyzer() + q = tirx.Var("q_stage_idx", "uint32") + expr = q * tirx.Cast("uint32", 128) + mod = expr % tirx.const(4, "uint32") + assert ana.can_prove_equal(mod, tirx.const(0, "uint32")) + tvm.ir.assert_structural_equal(ana.rewrite_simplify(mod), tirx.const(0, "uint32")) + + def test_simplify_float_division(): # Test for the discussion: # https://discuss.tvm.apache.org/t/discuss-is-constant-division-to-multiplication-rewrite-in-tvm-necessary/18615 diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py b/tests/python/tirx/codegen/test_codegen_blackwell.py index 61348ca48e61..00258b9b7bfc 100644 --- a/tests/python/tirx/codegen/test_codegen_blackwell.py +++ b/tests/python/tirx/codegen/test_codegen_blackwell.py @@ -89,6 +89,32 @@ def test_try_wait_once(A: T.Buffer((16, 16), "float16")): assert "selp.u32" in src +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_tcgen05_block_scaled_mma_predicate_codegen(): + # fmt: off + @T.prim_func + def test_predicated_mma(A: T.Buffer((1,), "uint32")): + T.device_entry() + T.cta_id([2]) + T.thread_id([256]) + issue: T.uint32 = T.ptx.elect_sync() + T.ptx.tcgen05.mma.block_scale( + T.uint32(0), T.uint64(0), T.uint64(0), + T.uint32(0), T.uint32(0), T.uint32(0), + d_dtype="float32", a_dtype="float8_e4m3fn", b_dtype="float8_e4m3fn", + sfa_dtype="float8_e8m0fnu", sfb_dtype="float8_e8m0fnu", + use_a_tmem=False, cta_group=2, enable_input_d=0, pred=issue, + ) + # fmt: on + + target = tvm.target.Target("cuda") + with target: + src, _ = _get_source(test_predicated_mma) + assert "setp.ne.b32 p_issue" in src + assert "@p_issue tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale" in src + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") def test_fence_before_after_thread_sync(): diff --git a/tests/python/tirx/codegen/test_cuda_copy.py b/tests/python/tirx/codegen/test_cuda_copy.py deleted file mode 100644 index 047eb1f12ca3..000000000000 --- a/tests/python/tirx/codegen/test_cuda_copy.py +++ /dev/null @@ -1,219 +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. -"""Tests for T.cuda.copy_128b / copy_64b / copy_32b / copy_16b / copy_8b intrinsics.""" - -import numpy as np -import pytest - -import tvm -from tvm.script import tirx as T -from tvm.testing import env - -DEV = tvm.cuda(0) -TARGET = tvm.target.Target("cuda") - - -def _build_and_run(func, *np_args): - mod = tvm.IRModule({"main": func}) - mod = tvm.compile(mod, target=TARGET, tir_pipeline="tirx") - rt_args = [tvm.runtime.tensor(a, device=DEV) for a in np_args] - mod(*rt_args) - return (*tuple(a.numpy() for a in rt_args), mod) - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_copy_128b(): - """copy_128b: copies 16 bytes (4 float32 elements) via uint4 load/store.""" - - # fmt: off - @T.prim_func - def func(out_ptr: T.handle): - out = T.match_buffer(out_ptr, (4,), "float32") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - src_buf = T.alloc_buffer((4,), "float32", scope="shared") - dst_buf = T.alloc_buffer((4,), "float32", scope="shared") - if lane < 4: - src_buf[lane] = T.float32(lane + 1) - T.cuda.cta_sync() - if lane == 0: - T.cuda.copy_128b(dst_buf.ptr_to([0]), src_buf.ptr_to([0])) - T.cuda.cta_sync() - if lane < 4: - out[lane] = dst_buf[lane] - # fmt: on - - out_np = np.zeros(4, dtype="float32") - result, mod = _build_and_run(func, out_np) - np.testing.assert_allclose(result, [1.0, 2.0, 3.0, 4.0]) - assert "tvm_builtin_copy_128b" in mod.mod.imports[0].inspect_source() - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_copy_64b(): - """copy_64b: copies 8 bytes (2 float32 elements) via uint2 load/store.""" - - # fmt: off - @T.prim_func - def func(out_ptr: T.handle): - out = T.match_buffer(out_ptr, (2,), "float32") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - src_buf = T.alloc_buffer((2,), "float32", scope="shared") - dst_buf = T.alloc_buffer((2,), "float32", scope="shared") - if lane < 2: - src_buf[lane] = T.float32(lane + 10) - T.cuda.cta_sync() - if lane == 0: - T.cuda.copy_64b(dst_buf.ptr_to([0]), src_buf.ptr_to([0])) - T.cuda.cta_sync() - if lane < 2: - out[lane] = dst_buf[lane] - # fmt: on - - out_np = np.zeros(2, dtype="float32") - result, mod = _build_and_run(func, out_np) - np.testing.assert_allclose(result, [10.0, 11.0]) - assert "tvm_builtin_copy_64b" in mod.mod.imports[0].inspect_source() - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_copy_32b(): - """copy_32b: copies 4 bytes (1 float32 element) via unsigned int load/store.""" - - # fmt: off - @T.prim_func - def func(out_ptr: T.handle): - out = T.match_buffer(out_ptr, (1,), "float32") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - src_buf = T.alloc_buffer((1,), "float32", scope="shared") - dst_buf = T.alloc_buffer((1,), "float32", scope="shared") - if lane == 0: - src_buf[0] = T.float32(42) - T.cuda.cta_sync() - if lane == 0: - T.cuda.copy_32b(dst_buf.ptr_to([0]), src_buf.ptr_to([0])) - T.cuda.cta_sync() - if lane == 0: - out[0] = dst_buf[0] - # fmt: on - - out_np = np.zeros(1, dtype="float32") - result, mod = _build_and_run(func, out_np) - np.testing.assert_allclose(result, [42.0]) - assert "tvm_builtin_copy_32b" in mod.mod.imports[0].inspect_source() - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_copy_16b(): - """copy_16b: copies 2 bytes (1 float16 element) via unsigned short load/store.""" - - # fmt: off - @T.prim_func - def func(out_ptr: T.handle): - out = T.match_buffer(out_ptr, (1,), "float16") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - src_buf = T.alloc_buffer((1,), "float16", scope="shared") - dst_buf = T.alloc_buffer((1,), "float16", scope="shared") - if lane == 0: - src_buf[0] = T.float16(7) - T.cuda.cta_sync() - if lane == 0: - T.cuda.copy_16b(dst_buf.ptr_to([0]), src_buf.ptr_to([0])) - T.cuda.cta_sync() - if lane == 0: - out[0] = dst_buf[0] - # fmt: on - - out_np = np.zeros(1, dtype="float16") - result, mod = _build_and_run(func, out_np) - np.testing.assert_allclose(result, [7.0]) - assert "tvm_builtin_copy_16b" in mod.mod.imports[0].inspect_source() - - -@pytest.mark.gpu -@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_copy_8b(): - """copy_8b: copies 1 byte (1 uint8 element) via unsigned char load/store.""" - - # fmt: off - @T.prim_func - def func(out_ptr: T.handle): - out = T.match_buffer(out_ptr, (1,), "uint8") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - src_buf = T.alloc_buffer((1,), "uint8", scope="shared") - dst_buf = T.alloc_buffer((1,), "uint8", scope="shared") - if lane == 0: - src_buf[0] = T.uint8(255) - T.cuda.cta_sync() - if lane == 0: - T.cuda.copy_8b(dst_buf.ptr_to([0]), src_buf.ptr_to([0])) - T.cuda.cta_sync() - if lane == 0: - out[0] = dst_buf[0] - # fmt: on - - out_np = np.zeros(1, dtype="uint8") - result, mod = _build_and_run(func, out_np) - np.testing.assert_equal(result, np.array([255], dtype="uint8")) - assert "tvm_builtin_copy_8b" in mod.mod.imports[0].inspect_source() - - -@pytest.mark.parametrize( - "num_bytes,func_suffix", [(16, "128b"), (8, "64b"), (4, "32b"), (2, "16b"), (1, "8b")] -) -def test_codegen_function_names(num_bytes, func_suffix): - """Verify each copy variant generates the expected C++ function name.""" - - copy_fn = getattr(T.cuda, f"copy_{func_suffix}") - - # fmt: off - @T.prim_func - def func(dummy_ptr: T.handle): - dummy = T.match_buffer(dummy_ptr, (16,), "uint8") - T.device_entry() - cta_id = T.cta_id([1]) - warp_id = T.warp_id([1]) - lane = T.lane_id([32]) - a = T.alloc_buffer((16,), "uint8", scope="shared") - b = T.alloc_buffer((16,), "uint8", scope="shared") - if lane == 0: - copy_fn(b.ptr_to([0]), a.ptr_to([0])) - dummy[0] = T.uint8(0) - # fmt: on - - mod = tvm.IRModule({"main": func}) - mod = tvm.compile(mod, target=TARGET, tir_pipeline="tirx") - source = mod.mod.imports[0].inspect_source() - assert f"tvm_builtin_copy_{func_suffix}" in source diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py new file mode 100644 index 000000000000..dd56ac825a56 --- /dev/null +++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py @@ -0,0 +1,202 @@ +# 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. +"""Unit tests for generic PTX ``T.ptx.ld`` / ``T.ptx.st`` vector copy ops.""" + +import numpy as np +import pytest + +import tvm +from tvm.ir import Op +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.copy._common import ( + copy_ptx_form, + copy_ptx_ld_return_type, +) + +DEV = tvm.cuda(0) +TARGET = tvm.target.Target("cuda") + +# num_bytes → kernel layout. ``fill_offset`` fills lane i with ``i + fill_offset``. +_SHARED_COPY_CASES = { + 16: {"nelems": 4, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_offset": 1}, + 8: {"nelems": 2, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_offset": 10}, + 4: {"nelems": 1, "smem_dtype": "uint32", "tmp_dtype": "uint32", "fill_value": 42}, + 2: {"nelems": 1, "smem_dtype": "float16", "tmp_dtype": "uint16", "fill_fp16": 7.0}, + 1: {"nelems": 1, "smem_dtype": "uint8", "tmp_dtype": "uint32", "fill_u8": 255}, +} + + +def _build_and_run(func, *np_args): + mod = tvm.compile(tvm.IRModule({"main": func}), target=TARGET, tir_pipeline="tirx") + rt_args = [tvm.runtime.tensor(a, device=DEV) for a in np_args] + mod(*rt_args) + return (*tuple(a.numpy() for a in rt_args), mod) + + +def _expected_values(num_bytes: int) -> np.ndarray: + spec = _SHARED_COPY_CASES[num_bytes] + if "fill_offset" in spec: + off, nelems = spec["fill_offset"], spec["nelems"] + return np.array([off + i for i in range(nelems)], dtype=np.uint32) + if "fill_fp16" in spec: + return np.array([spec["fill_fp16"]], dtype=np.float16) + if "fill_u8" in spec: + return np.array([spec["fill_u8"]], dtype=np.uint8) + return np.array([spec["fill_value"]], dtype=np.uint32) + + +def _shared_scratch_copy_kernel(num_bytes: int): + """Build shared → local scratch → shared copy kernel for ``num_bytes`` width.""" + spec = _SHARED_COPY_CASES[num_bytes] + smem_dtype = spec["smem_dtype"] + tmp_dtype = spec["tmp_dtype"] + nelems = spec["nelems"] + fill_offset = spec.get("fill_offset") + fill_value = spec.get("fill_value") + fill_fp16 = spec.get("fill_fp16") + fill_u8 = spec.get("fill_u8") + vec, ptx_type = copy_ptx_form(num_bytes) + return_type = copy_ptx_ld_return_type(ptx_type) + + @T.prim_func + def func(out_ptr: T.handle): + out = T.match_buffer(out_ptr, (nelems,), smem_dtype) + T.device_entry() + T.cta_id([1]) + T.warp_id([1]) + lane = T.lane_id([32]) + src_buf = T.alloc_buffer((nelems,), smem_dtype, scope="shared") + dst_buf = T.alloc_buffer((nelems,), smem_dtype, scope="shared") + tmp = T.alloc_local((nelems,), tmp_dtype) + if fill_offset is not None: + if lane < nelems: + src_buf[lane] = T.uint32(lane + fill_offset) + elif fill_fp16 is not None: + if lane == 0: + src_buf[0] = T.float16(fill_fp16) + elif fill_u8 is not None: + if lane == 0: + src_buf[0] = T.uint8(fill_u8) + elif lane == 0: + src_buf[0] = T.uint32(fill_value) + T.cuda.cta_sync() + if lane == 0: + T.ptx.ld( + src_buf.ptr_to([0]), + return_type, + ptx_type, + dst=tmp.ptr_to([0]), + space="shared", + vec=vec, + ) + T.ptx.st( + dst_buf.ptr_to([0]), + src=tmp.ptr_to([0]), + space="shared", + vec=vec, + ptx_type=ptx_type, + ) + T.cuda.cta_sync() + if lane < nelems: + out[lane] = dst_buf[lane] + + 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"): + Op.get(name) # raises if unregistered + + for attr in ( + "ld", + "st", + "ld_acquire", + "st_release", + "ld_volatile", + "st_volatile", + ): + assert hasattr(T.ptx, attr), attr + + +def test_ptx_ld_st_codegen_emits_shared_asm(): + """Shared ↔ register typed copies must codegen to ``ld.shared`` / ``st.shared``.""" + + # fmt: off + @T.prim_func + def copy_kernel(d_ptr: T.handle) -> None: + D = T.match_buffer(d_ptr, (4,), "uint32") + T.device_entry() + T.warp_id([4]) + T.cta_id([1]) + T.warpgroup_id([1]) + tid_in_wg = T.thread_id_in_wg([128]) + smem = T.alloc_buffer((4,), "uint32", scope="shared") + reg = T.alloc_local((4,), "uint32") + if tid_in_wg == 0: + T.ptx.st( + smem.ptr_to([0]), src=reg.ptr_to([0]), space="shared", vec="v4", ptx_type="u32" + ) + T.cuda.cta_sync() + if tid_in_wg == 0: + T.ptx.ld( + smem.ptr_to([0]), "uint32", "u32", dst=reg.ptr_to([0]), space="shared", vec="v4" + ) + Tx.copy(D[0:4], reg[:]) + # fmt: on + + 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.shared" in src, "PTX ld did not emit ld.shared" + assert "st.shared" in src, "PTX st did not emit st.shared" + assert "tvm_builtin_ptx_ld" in src + assert "tvm_builtin_ptx_st" in src + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +@pytest.mark.parametrize( + "num_bytes", + [16, 8, 4, 2, 1], + ids=["128b", "64b", "32b", "16b", "8b"], +) +def test_ptx_ld_st_shared_copy_gpu(num_bytes): + """GPU roundtrip for each supported PTX ld/st copy width (shared → scratch → shared).""" + expected = _expected_values(num_bytes) + kernel = _shared_scratch_copy_kernel(num_bytes) + out_np = np.zeros_like(expected) + result, mod = _build_and_run(kernel, out_np) + if expected.dtype == np.uint8: + np.testing.assert_array_equal(result, expected) + elif expected.dtype == np.float16: + np.testing.assert_allclose(result, expected) + else: + np.testing.assert_array_equal(result, expected) + src = mod.mod.imports[0].inspect_source("cuda") + assert "tvm_builtin_ptx_ld" in src + assert "tvm_builtin_ptx_st" in src + vec, _ptx_type = copy_ptx_form(num_bytes) + if vec == "v4": + assert "ld.shared.v4" in src + assert "st.shared.v4" in src + elif vec == "v2": + assert "ld.shared.v2" in src + assert "st.shared.v2" in src 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 4c51c9535e5b..c141f42449bb 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 @@ -457,6 +457,45 @@ def test_ldstmatrix_swizzle_multi_iter_pow2(): np.testing.assert_allclose(B.numpy(), A_np) +def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix(): + """Regression: a warpgroup ``.16x256b`` tcgen05 register atom loaded from a + 128B-swizzled SMEM tile must dispatch to ``ldmatrix.x4``. + + The atom's laneid is split across two shard iters (stride 4 and stride 1) + plus ``wid_in_wg`` — it only fuses to a single ``tid_in_wg`` thread axis + after the slice. With the redundant pre-slice ``canonicalize()`` removed, + Step 2 relies on ``TileLayout.Slice``'s built-in global pre-canonicalize, so + the slice no longer leaves an ill-formed split-laneid sub-layout that + ``GetScope`` would reject (which silently fell back to a scalar reg path). + Compile-only (no GPU): asserts the instruction appears in generated source. + """ + from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout + from tvm.tirx.layout import tcgen05_atom_layout + + m, k = 64, 64 + # bf16 payload carrying the *fp32* atom layout (mirrors the tf32-prenorm + # cast warp's ``a_bf16``: one element per 32-bit slot). + reg_layout = tcgen05_atom_layout("16x256b", (m, k), "float32") + smem_layout = mma_shared_layout("bfloat16", 3, (m, k)) + + @T.prim_func + def kernel(s_ptr: T.handle) -> None: + smem = T.match_buffer(s_ptr, (m, k), "bfloat16", scope="shared", layout=smem_layout) + T.device_entry() + T.cta_id([1]) + T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + T.thread_id_in_wg([128]) + a_reg = T.alloc_buffer((m, k), "bfloat16", scope="local", layout=reg_layout) + Tx.wg.copy(a_reg, smem) + + _, src = _compile_src(kernel) + assert "ldmatrix.sync.aligned.m8n8.x4.shared.b16" in src, ( + f"warpgroup tcgen05 atom did not emit ldmatrix.x4; src=\n{src}" + ) + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") def test_ldstmatrix_swizzle_multi_iter_linear(): 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 26c4d5de9b18..55058dfbf64e 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 @@ -333,7 +333,7 @@ def copy_sync(A_ptr: T.handle, B_ptr: T.handle) -> None: 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 - ``1 @ tid_in_wg``) must pick the widest vec ``copy_128b`` AND use the + ``1 @ tid_in_wg``) must pick the widest vec PTX ``st.shared.v4`` AND use the swizzle fast path (precomputed ``signed_strides`` + per-iter bit-select), not the per-iter ``swizzle.apply()`` fallback. @@ -388,12 +388,13 @@ 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() - # (1) Widest variant: 8 fp16 elements per call. - assert "tvm_builtin_copy_128b" in src, ( - "expected copy_128b in generated CUDA, alignment check fell back to a narrower variant" + # (1) Widest variant: 8 fp16 elements per call (16 bytes → v4.u32 st). + assert "tvm_builtin_ptx_st" in src, ( + "expected PTX st in generated CUDA, alignment check fell back to a narrower variant" ) - assert "tvm_builtin_copy_16b" not in src, ( - "scalar copy_16b appeared — vec=8 was wrongly rejected" + assert "st.shared.v4" in src, "expected 128b vector store (st.shared.v4.u32)" + assert "tvm_builtin_copy_" not in src, ( + "copy_xxb helpers appeared — reg dispatch should use PTX ld/st only" ) # (2) Swizzle fast path fingerprint: # * emit_init allocates a size-N int buffer of "signed strides". @@ -410,5 +411,208 @@ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: ) +# --- 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 +# into 128B-swizzled MMA SMEM for the subsequent TMA store to gmem D. +_TCGEN05_D_ATOM = "16x256b" +_TCGEN05_D_SHAPE = (64, 64) +_TCGEN05_D_DTYPE = "float32" +_TCGEN05_D_SWIZZLE = 3 # SwizzleMode 128B → mma_shared_layout(..., 3, shape) +_TCGEN05_D_SLICE = (slice(0, 64), slice(0, 64)) + + +def _tcgen05_d_epilogue_layouts(): + from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout + from tvm.tirx.layout import tcgen05_atom_layout + + m, n = _TCGEN05_D_SHAPE + reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE) + smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n)) + return m, n, reg_layout, smem_layout + + +def _build_tcgen05_d_epilogue_deposit(): + """``Tx.wg.copy(smem[slice], d_reg[slice])``: R (tcgen05 atom) → S (128B swizzle).""" + from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout + from tvm.tirx.layout import tcgen05_atom_layout + + m, n = _TCGEN05_D_SHAPE + smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n)) + reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE) + sl_m, sl_n = _TCGEN05_D_SLICE + + @T.prim_func + def deposit( + d_reg: T.Buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout), + ) -> None: + smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout) + T.device_entry() + T.cta_id([1]) + T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + T.thread_id_in_wg([128]) + Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n]) + + return deposit + + +def _compile_tcgen05_d_epilogue_deposit(): + target = tvm.target.Target("cuda") + with target: + mod = tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit()}) + return tvm.compile(mod, target=target, tir_pipeline="tirx") + + +def test_reg_copy_tcgen05_d_epilogue_deposit_layout_pairing(): + """Pre-fix bug: canonical ``r_p`` collapses atom ``m`` groups and drops S pairings. + + Copy: ``Tx.wg.copy(smem_cd_mma[0:64,0:64], d_reg[0:64,0:64])`` (R→S). + ``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 ( + _split_thread_loop, + align_layouts_raw, + ) + from tvm.tirx.exec_scope import ExecScope + from tvm.tirx.operator.tile_primitive import DispatchContext + + m, n, reg_layout, smem_layout = _tcgen05_d_epilogue_layouts() + region = [(0, m), (0, n)] + sctx = DispatchContext( + tvm.target.Target("cuda"), ExecScope("warpgroup"), {}, {}, scope_kind="warpgroup" + ) + with sctx.target: + r_sliced = reg_layout.slice([m, n], region) + s_sliced = smem_layout.slice([m, n], region) + r_p, s_p, s_seps, r_perm = align_layouts_raw(r_sliced, s_sliced, region) + + r_iters, s_groups = _split_thread_loop(r_perm, s_p, s_seps) + r_iters_bug, s_groups_bug = _split_thread_loop(r_p, s_p, s_seps) + mem_extents = [int(it.extent) for it in r_iters] + bug_extents = [int(it.extent) for it in r_iters_bug] + + # Fixed path: 3 register m-groups stay 1:1 with 3 S-side groups. + assert mem_extents == [8, 2, 2] + assert len(r_iters) == len(s_groups) == 3 + # Pre-fix path (what _split_thread_loop used to take): 1 fused m-iter, only + # the first S group is paired — the other two are silently dropped. + assert bug_extents == [32] + assert len(r_iters_bug) == len(s_groups_bug) == 1 + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_reg_copy_tcgen05_d_epilogue_deposit_codegen(): + """``reg`` dispatch lowers the D epilogue deposit; pre-fix loop only covered half. + + Without the ``r_perm`` fix the emitted outer loop ran ``f < 8`` (one fused + register group) instead of ``f < 16`` (three atom m-groups x vec tail), and + the swizzled SMEM stores landed in the wrong (row, col) slots. + """ + import re + + 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 "tvm_builtin_ptx_st" in src + assert "st.shared.v2.u32" in src, "fp32 vec=2 → 8B shared store per outer iter" + + loop = re.search(r"for \(int f = 0; f < (\d+)", src) + assert loop is not None, "expected reg copy outer loop in generated CUDA" + assert loop.group(1) == "16", ( + f"fixed pairing emits 16 outer stores (pre-fix bug collapsed to 8); got f < {loop.group(1)}" + ) + + +def _tcgen05_16x256b_row_col(tid_wg: T.int32, lane: T.int32, reg_idx: T.int32): + """Map ``(tid_in_wg, reg)`` → logical ``(row, col)`` for ``.16x256b`` fp32 atom.""" + t0 = lane & T.int32(3) + t1 = lane >> 2 + v0p = reg_idx & T.int32(1) + va = (reg_idx >> 1) & T.int32(1) + vb = reg_idx >> 2 + wid = tid_wg >> 5 + row = t1 + T.int32(8) * va + T.int32(16) * wid + col = v0p + T.int32(2) * t0 + T.int32(8) * vb + return row, col + + +def _build_tcgen05_d_epilogue_deposit_roundtrip(): + """Fill ``d_reg``, R→S deposit, S→R reload, dump via ``.local()`` to gmem.""" + from tvm.tirx.cuda.operator.tile_primitive.tma_utils import mma_shared_layout + from tvm.tirx.layout import tcgen05_atom_layout + + m, n = _TCGEN05_D_SHAPE + regs_per_thread = 32 # ``.16x256b.x8`` fp32: 4 regs/slot x rep=8 + smem_layout = mma_shared_layout(_TCGEN05_D_DTYPE, _TCGEN05_D_SWIZZLE, (m, n)) + reg_layout = tcgen05_atom_layout(_TCGEN05_D_ATOM, (m, n), _TCGEN05_D_DTYPE) + sl_m, sl_n = _TCGEN05_D_SLICE + + @T.prim_func + def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (m, n), _TCGEN05_D_DTYPE) + B = T.match_buffer(B_ptr, (m, n), _TCGEN05_D_DTYPE) + T.device_entry() + T.cta_id([1]) + T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + tid_wg = T.thread_id_in_wg([128]) + lane = T.lane_id([32]) + d_reg = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout) + d_reg_out = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="local", layout=reg_layout) + smem_cd_mma = T.alloc_buffer((m, n), _TCGEN05_D_DTYPE, scope="shared", layout=smem_layout) + reg_in = d_reg.local(regs_per_thread) + reg_out = d_reg_out.local(regs_per_thread) + for r in T.serial(regs_per_thread): + row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32")) + reg_in[r] = A[row, col] + Tx.wg.copy(smem_cd_mma[sl_m, sl_n], d_reg[sl_m, sl_n]) + T.cuda.cta_sync() + Tx.wg.copy(d_reg_out[sl_m, sl_n], smem_cd_mma[sl_m, sl_n]) + for r in T.serial(regs_per_thread): + row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32")) + B[row, col] = reg_out[r] + + return kernel + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(9), reason="need cuda compute >= 9.0") +def test_reg_copy_tcgen05_d_epilogue_deposit_gpu(): + """GPU: R→S deposit + S→R reload must preserve the Layout-F register tile. + + Host fills gmem ``A[row,col]=row*100+col``; each thread scatters into ``d_reg`` + via the ``.16x256b`` (tid,reg)→(row,col) map, runs production + ``Tx.wg.copy(smem_cd_mma, d_reg)`` then the inverse + ``Tx.wg.copy(d_reg_out, smem_cd_mma)``, and dumps ``d_reg_out`` back to gmem + ``B`` through ``.local()``. Pre-fix pairing dropped 2/3 of the S groups — + ``max|B-A|`` was hundreds, not 0. + """ + m, n = _TCGEN05_D_SHAPE + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _build_tcgen05_d_epilogue_deposit_roundtrip()}), + target=target, + tir_pipeline="tirx", + ) + + dev = tvm.cuda(0) + rows = np.arange(m, dtype=np.int32)[:, None] + cols = np.arange(n, dtype=np.int32)[None, :] + a_np = (rows * 100 + cols).astype(np.float32) + b_np = np.zeros((m, n), dtype=np.float32) + 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, rtol=0, atol=0) + + 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 27bf74ed4082..9c3dd81fa973 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_smem_tmem.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_smem_tmem.py index 84f23cf8eaa2..497c365e464c 100644 --- 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 @@ -442,5 +442,34 @@ def test_dispatch_rejects_bad_inputs(bad): tvm.compile(tvm.IRModule({"main": kernel}), target=target, tir_pipeline="tirx") +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_tma.py b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py index 3e9cb455b039..f70b52258356 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 @@ -35,10 +35,10 @@ ) 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_functor import StmtExprVisitor +from tvm.tirx.tile_primitive import DispatchContext # =========================================================================== # Helpers @@ -406,6 +406,16 @@ def _tma_case( 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-l2-promotion-256b", + 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)), + config={"l2_promotion": 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, 3, 0], + ), _tma_case( id="g2s-2d-8x256-swizzle2", g_shape=(8, 256), g_region=((0, 8), (0, 256)), @@ -976,6 +986,15 @@ def _tma_case( config={"oob": "bogus"}, raises=(Exception, "Unsupported TMA oob mode"), ), + _tma_case( + id="reject-unknown-l2-promotion", + 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)), + config={"l2_promotion": 192}, + raises=(Exception, "Unsupported l2_promotion"), + ), _tma_case( id="reject-g2s-nan-on-non-float", g_shape=(128, 64), g_region=((120, 136), (0, 64)), @@ -1568,5 +1587,71 @@ def copy_async_dynamic_mask(A_ptr: T.handle) -> None: assert "multicast" in src, "Expected multicast TMA instruction in generated code" +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[:]) + + target = tvm.target.Target("cuda") + with target: + tvm.compile(tvm.IRModule({"main": tma_load}), target=target, tir_pipeline="tirx") + + +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") + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py index 8d39ba355633..c19f53bef7ef 100644 --- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py +++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_binary.py @@ -781,6 +781,41 @@ def k(A_ptr: T.handle, B_ptr: T.handle) -> None: ), f"expected packed add_f32x2; got:\n{src[:2000]}" +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_binary_maximum_reg(): + N = 128 + + @T.prim_func + def relu_max(a_ptr: T.handle, b_ptr: T.handle, c_ptr: T.handle) -> None: + A = T.match_buffer(a_ptr, (N,), "float32") + B = T.match_buffer(b_ptr, (N,), "float32") + C = T.match_buffer(c_ptr, (N,), "float32") + T.device_entry() + T.warp_id([4]) + T.cta_id([1]) + T.warpgroup_id([1]) + tid = T.thread_id_in_wg([N]) + a_reg = T.alloc_local((1,), "float32") + b_reg = T.alloc_local((1,), "float32") + Tx.copy(a_reg[:], A[tid : tid + 1]) + Tx.copy(b_reg[:], B[tid : tid + 1]) + Tx.maximum(a_reg[:], a_reg[:], b_reg[:]) + Tx.copy(C[tid : tid + 1], a_reg[:]) + + dev = tvm.cuda(0) + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": relu_max}), target=target, tir_pipeline="tirx") + np.random.seed(0) + a = np.random.randn(N).astype("float32") + b = np.random.randn(N).astype("float32") + c = np.zeros(N, dtype="float32") + a_t, b_t, c_t = (tvm.runtime.tensor(x, dev) for x in (a, b, c)) + mod["main"](a_t, b_t, c_t) + np.testing.assert_allclose(c_t.numpy(), np.maximum(a, b), atol=0, rtol=0) + + def test_binary_add_f16_scalar_fallback_dispatch(): """add f16 has no packed VecImpl → reg.py scalar fallback (T.vectorized).""" shape = (64, 32) 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 97a1be256e0a..250639423b59 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 @@ -1505,5 +1505,89 @@ def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: _sl_compile(kernel) +# --- tcgen05 split-laneid atom (B00011 / tf32 A-cast path) ----------------- +_TCGEN05_M, _TCGEN05_K = 64, 64 +_TCGEN05_REGS_PER_THREAD = 32 # ``.16x256b`` fp32 atom on (64, 64) + + +def _tcgen05_16x256b_row_col(tid_wg: T.int32, lane: T.int32, reg_idx: T.int32): + """Map ``(tid_in_wg, reg)`` → logical ``(row, col)`` for ``.16x256b`` fp32 atom.""" + t0 = lane & T.int32(3) + t1 = lane >> 2 + v0p = reg_idx & T.int32(1) + va = (reg_idx >> 1) & T.int32(1) + vb = reg_idx >> 2 + wid = tid_wg >> 5 + row = t1 + T.int32(8) * va + T.int32(16) * wid + col = v0p + T.int32(2) * t0 + T.int32(8) * vb + return row, col + + +def _tcgen05_cast_atom_layout(): + from tvm.tirx.layout import tcgen05_atom_layout + + # Production tf32_hc_prenorm_gemm pattern: bf16 payload in fp32 atom slots. + return tcgen05_atom_layout("16x256b", (_TCGEN05_M, _TCGEN05_K), "float32") + + +def _tcgen05_cast_warpgroup_kernel(): + """Scatter gmem bf16 → reg, ``Tx.wg.cast``, gather fp32 → gmem.""" + atom_layout = _tcgen05_cast_atom_layout() + m, k = _TCGEN05_M, _TCGEN05_K + + @T.prim_func + def kernel(A_ptr: T.handle, B_ptr: T.handle) -> None: + A = T.match_buffer(A_ptr, (m, k), "bfloat16") + B = T.match_buffer(B_ptr, (m, k), "float32") + T.device_entry() + T.cta_id([1]) + T.warpgroup_id([1]) + T.warp_id_in_wg([4]) + T.lane_id([32]) + tid_wg = T.thread_id_in_wg([128]) + lane = T.lane_id([32]) + a_bf16 = T.alloc_buffer((m, k), "bfloat16", scope="local", layout=atom_layout) + a_fp32 = T.alloc_buffer((m, k), "float32", scope="local", layout=atom_layout) + reg_in = a_bf16.local(_TCGEN05_REGS_PER_THREAD) + reg_out = a_fp32.local(_TCGEN05_REGS_PER_THREAD) + for r in T.serial(_TCGEN05_REGS_PER_THREAD): + row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32")) + reg_in[r] = A[row, col] + Tx.wg.cast(a_fp32, a_bf16) + for r in T.serial(_TCGEN05_REGS_PER_THREAD): + row, col = _tcgen05_16x256b_row_col(tid_wg, lane, T.cast(r, "int32")) + B[row, col] = reg_out[r] + + return kernel + + +def test_cast_tcgen05_atom_warpgroup_reg_dispatch_compiles(): + """Regression: tf32 A-cast ``Tx.wg.cast(a_fp32, a_bf16)`` on tcgen05 atom compiles.""" + _sl_compile(_tcgen05_cast_warpgroup_kernel()) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_cast_tcgen05_atom_warpgroup_gpu(): + """GPU: bf16→fp32 cast on tcgen05 ``.16x256b`` atom must preserve per-slot layout.""" + m, k = _TCGEN05_M, _TCGEN05_K + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _tcgen05_cast_warpgroup_kernel()}), + target=target, + tir_pipeline="tirx", + ) + + dev = tvm.cuda(0) + np.random.seed(0) + a_np = (np.random.randn(m, k).astype("float32") * 4.0).astype("bfloat16") + b_ref = a_np.astype("float32") + a = tvm.runtime.tensor(a_np, dev) + b = tvm.runtime.tensor(np.zeros((m, k), dtype="float32"), dev) + mod(a, b) + tvm.testing.assert_allclose(b.numpy(), b_ref, rtol=0, atol=1e-2) + + if __name__ == "__main__": tvm.testing.main() 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 32ac00e39d5f..d7f765a9777f 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 @@ -2067,5 +2067,225 @@ def gemm_async(A_ptr: T.handle, B_ptr: T.handle, C_ptr: T.handle) -> None: np.testing.assert_allclose(C_t.numpy(), C_ref, atol=1e-2, rtol=1e-2) +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"} + 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", 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) + + +@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 for fp8") +def test_gemm_dense_fp8(): + _run_dense_gemm("float8_e4m3fn", "float8_e4m3fn", "float32", 128, atol=2.0, rtol=0.15) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >= 10.0") +def test_gemm_tf32_with_tfloat32_tma(): + _run_dense_gemm( + "float32", + "float32", + "float32", + 64, + is_AB_tf32=True, + tma_dtype_B="tf32", + atol=2e-2, + rtol=2e-2, + ) + + +def _build_smem_desc_kernel(smem_desc): + """Minimal cta_group=1 fp16 gemm_async kernel parametrized on ``smem_desc``.""" + 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 + width = C_region[1][1] - C_region[1][0] + A_layout = mma_shared_layout(A_dtype, A_sw, A_shape) + B_layout = mma_shared_layout(B_dtype, B_sw, B_shape) + r_gmem_A = [slice(0, A_shape[i]) for i in range(len(A_shape))] + r_gmem_B = [slice(0, B_shape[i]) for i in range(len(B_shape))] + total_bytes = ( + functools.reduce(operator.mul, A_shape, 1) * 2 + + functools.reduce(operator.mul, B_shape, 1) * 2 + ) + r_tmem_C = [slice(C_region[i][0], C_region[i][1]) for i in range(len(C_shape))] + r_smem_A = [slice(1, 2), slice(0, 128), slice(0, 64)] + r_smem_B = [slice(2, 3), slice(0, 128), slice(0, 64)] + + # fmt: off + @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]) + cta_id = 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=128, cta_group=1) + 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])}) + 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 + 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(width, dtype=C_dtype) + C_view = C_reg.view(128, width, layout=TileLayout(S[(128, width) : (1@axis_tid_in_wg, 1)])) + if wg_id == 0: + Tx.wg.copy_async(C_view[:, :], tmem[tuple(r_tmem_C)]) + T.ptx.tcgen05.wait.ld() + T.cuda.cta_sync() + Tx.copy(C[tid_in_wg, C_region[1][0]:C_region[1][1]], 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 + + return gemm_async + + +@pytest.mark.parametrize("smem_desc", ["hoist", "recompute"]) +def test_gemm_smem_desc_hoist_vs_recompute(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``. + - ``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. + + Both must emit the MMA; the descriptor-construction fingerprints differ. + """ + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile( + tvm.IRModule({"main": _build_smem_desc_kernel(smem_desc)}), + target=target, + tir_pipeline="tirx", + ) + src = mod.mod.imports[0].inspect_source() + 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 "smem_desc_add_16B_offset" in src, "hoist mode must add the per-MMA 16B offset" + 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" + + 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 0402719ba1e5..dac4d06a15ec 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 @@ -414,5 +414,47 @@ def f(A: T.handle, B: T.handle): assert "warp" in str(exc_info.value) +@pytest.mark.parametrize("dtype", ["uint32", "float32"]) +def test_shared_to_shared_uses_direct_ldst(dtype): + """Compile-only: a shared->shared 32b transpose must take the direct + base-ptr + byte-offset ``ld.shared`` / ``st.shared`` path. + + For a 4/8-byte dtype with both operands in shared memory, indexing through + ``buf[...]`` lowers the swizzled layout to a per-element IMAD flatten. The + direct path computes one base ptr (``ptr_to(stride_offset)``) and adds a + compile-time ``off * dtype_bytes`` per register slot, then issues + ``T.ptx.ld/st(..., space="shared")``. The bits move through a uint + container, so ``float32`` (whose ``ld.b32`` cannot return a float) lowers + the same way as the ``uint32`` SF case. + """ + shape = (4, 32) + pre = TileLayout(S[shape : (32, 1)]) # linear + post = TileLayout(S[shape : (1, 4)]) # transposed within the 128-block + + # fmt: off + @T.prim_func + def f(A: T.handle, B: T.handle): + A_buf = T.match_buffer(A, shape, dtype, layout=pre) + B_buf = T.match_buffer(B, shape, dtype, layout=post) + T.device_entry() + T.cta_id([1]) + tid = T.thread_id([32]) + sA = T.alloc_buffer(shape, dtype, scope="shared", layout=pre) + sB = T.alloc_buffer(shape, dtype, scope="shared", layout=post) + Tx.cta.copy(sA[:, :], A_buf[:, :]) + T.cuda.cta_sync() + Tx.warp.permute_layout(sB[:, :], sA[:, :]) + T.cuda.cta_sync() + Tx.cta.copy(B_buf[:, :], sB[:, :]) + # fmt: on + + target = tvm.target.Target("cuda") + with target: + mod = tvm.compile(tvm.IRModule({"main": f}), target=target, tir_pipeline="tirx") + src = mod.mod.imports[0].inspect_source() + assert "ld.shared" in src, f"expected direct ld.shared in permute; src=\n{src}" + assert "st.shared" in src, f"expected direct st.shared in permute; src=\n{src}" + + if __name__ == "__main__": tvm.testing.main() 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_bench_utils.py b/tests/python/tirx/test_bench_utils.py index e5bf13bc845e..be7ab2ce7755 100644 --- a/tests/python/tirx/test_bench_utils.py +++ b/tests/python/tirx/test_bench_utils.py @@ -22,7 +22,13 @@ 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 ( + _compute_group_count, + _parse_proton_tree, + bench, + bench_tk, + tensor_bytes, +) # ── _parse_proton_tree ────────────────────────────────────────────────────── @@ -88,13 +94,13 @@ def test_parse_proton_tree_empty(): assert errors == {} -# ── bench ─────────────────────────────────────────────────────────────────── +# ── bench_tk (ThunderKittens group-input protocol) ────────────────────────── @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_basic(): - """bench returns positive times for each impl.""" +def test_bench_tk_basic(): + """bench_tk returns positive times for each impl.""" M, N = 256, 256 funcs = {"matmul": lambda case: torch.mm(case[0], case[1])} @@ -104,14 +110,14 @@ def make_input(): B = torch.randn(M, N, device="cuda", dtype=torch.float16) return (A, B), tensor_bytes(A, B) - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") + results = bench_tk(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") assert "matmul" in results["impls"] assert results["impls"]["matmul"] > 0 @pytest.mark.gpu @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") -def test_bench_multiple_impls(): +def test_bench_tk_multiple_impls(): """Multiple impls each get their own timing.""" M, N = 128, 128 funcs = { @@ -126,14 +132,14 @@ def make_input(): B = torch.randn(M, N, device="cuda", dtype=torch.float16) return (A, B), tensor_bytes(A, B) - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") + results = bench_tk(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_multiple_input_groups(): +def test_bench_tk_multiple_input_groups(): """Multiple input groups cycle correctly (L2 eviction).""" M, N = 128, 128 call_count = [0] @@ -145,13 +151,97 @@ def make_input(): return (A, B), tensor_bytes(A, B) funcs = {"mm": lambda case: torch.mm(case[0], case[1])} - results = bench( + results = bench_tk( funcs, make_input, warmup=5, repeat=20, cooldown_s=0.0, timer="event", l2_bytes=64 * 1024 ) assert results["impls"]["mm"] > 0 assert call_count[0] > 1 +# ── bench (Triton-standard, pure-launch) ───────────────────────────────────── + + +@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" + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_bench_default_timer_is_proton(): + """Omitting timer resolves to the proton default (recorded as timer='proton'). + + Under pytest _do_bench_proton falls back to event timing, but bench() still + records the resolved timer name, so this pins the timer=None -> 'proton' default. + """ + 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) + assert results["timer"] == "proton" + assert results["impls"]["mm"] > 0 + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_bench_cudagraph_proton_pure_launch(): + """New Triton-standard bench(): cudagraph_proton timer. + + Under pytest, _do_bench_cudagraph_proton falls back to event-based cudagraph + timing (no CUPTI/Proton required in CI), so this exercises the mode wiring and + the fallback path rather than the real Proton attribution. + """ + 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="cudagraph_proton") + assert results["impls"]["mm"] > 0 + assert results["timer"] == "cudagraph_proton" + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +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) + + funcs = {"tir": lambda: torch.mm(A, B)} + + def _addmm(): + C = torch.zeros(M, N, device="cuda", dtype=torch.float16) + return lambda: torch.addmm(C, A, B) + + 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()) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_cuda(), reason="need cuda") +def test_bench_rejects_unknown_timer(): + """New bench() only accepts event / proton / cudagraph_proton (plain cudagraph + was removed).""" + A = torch.randn(8, 8, device="cuda", dtype=torch.float16) + with pytest.raises(ValueError): + bench({"mm": lambda: torch.mm(A, A)}, timer="cudagraph") + + # ── _compute_group_count ─────────────────────────────────────────────────── @@ -181,13 +271,13 @@ def test_compute_groups_moderate_tensors(): @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.""" +def test_bench_tk_legacy_callable_api(): + """bench_tk still accepts the existing single-callable API used by TIRx tests.""" M, N = 128, 128 A = torch.randn(M, N, device="cuda", dtype=torch.float16) B = torch.randn(M, N, device="cuda", dtype=torch.float16) - result = bench( + result = bench_tk( lambda: torch.mm(A, B), warmup=1, repeat=2, proton_name="legacy", flush_l2_size=1 ) assert result > 0 @@ -195,8 +285,8 @@ def test_bench_legacy_callable_api(): @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.""" +def test_bench_tk_callable_inputs(): + """bench_tk accepts a factory callable and auto-computes groups.""" M, N = 256, 256 call_count = [0] @@ -210,7 +300,7 @@ def make_input(): return case, tensor_bytes(*case) funcs = {"mm": lambda case: torch.mm(case[0], case[1])} - results = bench(funcs, make_input, warmup=5, repeat=10, cooldown_s=0.0, timer="event") + results = bench_tk(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 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_op.py b/tests/python/tirx/test_op.py index 4c417033d697..bf4d43f8e048 100644 --- a/tests/python/tirx/test_op.py +++ b/tests/python/tirx/test_op.py @@ -18,7 +18,7 @@ from tvm.ir import Op from tvm.tirx.buffer import decl_buffer -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 e6d71cabd45a..2f75474f3f70 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): @@ -95,7 +95,7 @@ def test_builtin_expression_ops_are_not_tile_primitives(): cast = T.cast(x, "float32") assert isinstance(cast, tvm.tirx.Cast) - assert cast.dtype == "float32" + assert cast.ty.dtype == "float32" sqrt = T.sqrt(y) assert sqrt.op.name == "tirx.sqrt" @@ -292,7 +292,7 @@ def test_device_intrinsic_printer_roundtrips_canonical_namespaces(): def device_namespaces(dst: T.handle, src: T.handle): A = T.match_buffer(src, (1,), "float32") R = T.alloc_buffer((1,), "float32", scope="local") - T.cuda.copy_bytes(dst, src, 16) + T.cuda.cta_sync() T.ptx.ldg32(R[0], 1, A[0], 0) T.metal.simd_shuffle(A[0], 0) T.metal.simd_shuffle_up(A[0], 1) @@ -300,14 +300,14 @@ def device_namespaces(dst: T.handle, src: T.handle): calls = _expr_calls(device_namespaces) assert [call.op.name for call in calls] == [ - "tirx.cuda.copy_bytes", + "tirx.cuda.cta_sync", "tirx.ptx.ldg32", "tirx.metal.simd_shuffle", "tirx.metal.simd_shuffle_up", "tirx.metal.simd_shuffle_down", ] for op_name, namespace in [ - ("tirx.cuda.copy_bytes", "cuda"), + ("tirx.cuda.cta_sync", "cuda"), ("tirx.ptx.ldg32", "ptx"), ("tirx.metal.simd_shuffle", "metal"), ("tirx.metal.simd_shuffle_up", "metal"), @@ -318,7 +318,7 @@ def device_namespaces(dst: T.handle, src: T.handle): assert _op_attr(op_name, "TCallEffectKind") in (1, 3) code = device_namespaces.script() - assert "T.cuda.copy_bytes(" in code + assert "T.cuda.cta_sync(" in code assert "T.ptx.ldg32(" in code assert "T.metal.simd_shuffle(" in code assert "T.metal.simd_shuffle_up(" in code diff --git a/tests/python/tirx/test_parser_printer.py b/tests/python/tirx/test_parser_printer.py index 561adfc602ed..60c700e4eedd 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 @@ -1385,6 +1387,432 @@ def func() -> None: assert from_source(code).script() == code +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( + T.SwizzleLayout(3, 3, 3, swizzle_inner=True), + 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.swizzle, + a_buf.layout.tile_layout.permute_dims([1, 0, 2, 3]), + ) + assert_structural_equal(b_buf.layout, expected) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_unflatten_flatten_ir(): + """unflatten splits a dim (with -1 inference); flatten merges dims back. + Both keep the original layout object (pure reshapes).""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([8, 64], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(8, 64) : (64, 1)])) + B = A.unflatten(1, (4, 16)) + B[0, 0, 0] = T.float16(0) + C = A.unflatten(1, (-1, 16)) + C[0, 0, 0] = T.float16(0) + D = B.flatten(1, 2) + D[0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf, b_buf, c_buf, d_buf = bufs["A"], bufs["B"], bufs["C"], bufs["D"] + assert b_buf.data.same_as(a_buf.data) + assert [int(s) for s in b_buf.shape] == [8, 4, 16] + assert [int(s) for s in c_buf.shape] == [8, 4, 16] + assert [int(s) for s in d_buf.shape] == [8, 64] + assert_structural_equal(b_buf.layout, a_buf.layout) + assert_structural_equal(d_buf.layout, a_buf.layout) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_select_ir(): + """select removes a dim; the index (static or dynamic) folds into + elem_offset through the dim's layout iter stride.""" + + # 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.select(0, 2) + B[0, 0] = T.float16(0) + for i in T.serial(4): + C = A.select(0, i) + C[0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + a_buf, b_buf, c_buf = bufs["A"], bufs["B"], bufs["C"] + assert b_buf.data.same_as(a_buf.data) + assert [int(s) for s in b_buf.shape] == [8, 16] + assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - a_buf.elem_offset)) == 512 + assert_structural_equal(b_buf.layout, tvm.tirx.layout.TileLayout(T.S[(8, 16) : (16, 1)])) + assert [int(s) for s in c_buf.shape] == [8, 16] + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_select_multi_iter_dim_ir(): + """select 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.select(0, 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_narrow_ir(): + """narrow keeps the dim at a reduced extent; multi-iter dims require + start/length on the inner iter block boundary.""" + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([4, 64], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(4, 64) : (64, 1)])) + for i in T.serial(2): + B = A.narrow(1, i * 32, 32) + B[0, 0] = T.float16(0) + M = T.alloc_buffer([8, 16], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(2, 4, 16) : (1024, 64, 1)])) + C = M.narrow(0, 4, 4) + C[0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + b_buf, m_buf, c_buf = bufs["B"], bufs["M"], bufs["C"] + assert [int(s) for s in b_buf.shape] == [4, 32] + assert [int(s) for s in c_buf.shape] == [4, 16] + assert int(tvm.arith.Analyzer().simplify(c_buf.elem_offset - m_buf.elem_offset)) == 1024 + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_narrow_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)]) + ) + with pytest.raises(ValueError, match="multiples of the inner iter block"): + buf.narrow(0, 2, 4) + + +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 (via unflatten + select).""" + + # 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) + D = A.select(0, 1).narrow(0, 2, 4) + D[0, 0] = T.float16(0) + E = A.unflatten(1, (4, 2)).select(2, 1) + E[0, 0, 0] = T.float16(0) + # fmt: on + + bufs = _collect_buffers(func) + b_buf, c_buf, d_buf, e_buf = bufs["B"], bufs["C"], bufs["D"], bufs["E"] + assert [int(s) for s in b_buf.shape] == [4, 16] + assert_structural_equal(b_buf.layout, d_buf.layout) + assert int(tvm.arith.Analyzer().simplify(b_buf.elem_offset - d_buf.elem_offset)) == 0 + assert [int(s) for s in c_buf.shape] == [4, 4, 16] + assert_structural_equal(c_buf.layout, e_buf.layout) + assert int(tvm.arith.Analyzer().simplify(c_buf.elem_offset - e_buf.elem_offset)) == 0 + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_rearrange_ir(): + """rearrange compiles to the unflatten/permute/flatten chain; the swizzle + of a composed layout is carried.""" + + compose = T.ComposeLayout( + T.SwizzleLayout(3, 3, 3, swizzle_inner=True), + T.TileLayout(T.S[(2, 16, 8) : (128, 8, 1)]), + ) + + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([2, 16, 8], dtype="bfloat16", scope="shared.dyn", layout=compose) + B = A.rearrange("b (s w r) c -> b w (s r) c", w=2, r=4) + B[0, 0, 0, 0] = T.bfloat16(0) + + @T.prim_func + def func_chain() -> None: + T.device_entry() + A = T.alloc_buffer([2, 16, 8], dtype="bfloat16", scope="shared.dyn", layout=compose) + C = A.unflatten(1, (2, 2, 4)).permute(0, 2, 1, 3, 4).flatten(2, 3) + C[0, 0, 0, 0] = T.bfloat16(0) + # fmt: on + + b_buf = _collect_buffers(func)["B"] + c_buf = _collect_buffers(func_chain)["C"] + assert [int(s) for s in b_buf.shape] == [2, 2, 8, 8] + assert_structural_equal(b_buf.layout, c_buf.layout) + assert isinstance(b_buf.layout, tvm.tirx.layout.ComposeLayout) + + code = func.script() + assert from_source(code).script() == code + + +def test_buffer_rearrange_errors(): + buf = tvm.tirx.decl_buffer( + (2, 16, 8), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(2, 16, 8) : (128, 8, 1)]) + ) + with pytest.raises(ValueError, match="must contain '->'"): + buf.rearrange("b h c") + with pytest.raises(ValueError, match="appear on only one side"): + buf.rearrange("b h c -> b h d") + with pytest.raises(ValueError, match="duplicate axis name"): + buf.rearrange("b b c -> b c b") + with pytest.raises(ValueError, match="multiple unknown factors"): + buf.rearrange("b (s w r) c -> b w (s r) c") + with pytest.raises(ValueError, match="input dims"): + buf.rearrange("b c -> c b") + + +def test_buffer_view_surgery_static_bounds_rejected(): + """Statically-known out-of-range arguments must be rejected loudly + (review finding: non-bijective reshapes and 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)]) + ) + # unflatten: non-bijective factorizations + with pytest.raises(ValueError, match="multiply to"): + buf.unflatten(0, (3, 4)) + with pytest.raises(ValueError, match="not divisible"): + buf.unflatten(0, (-1, 4)) + # negative factors whose product happens to match the extent + with pytest.raises(ValueError, match="must be positive"): + buf.unflatten(0, (-2, -5)) + with pytest.raises(ValueError, match="must be positive"): + buf.rearrange("(a b) -> a b", a=-2, b=-5) + # rearrange reuses the same validation + with pytest.raises(ValueError, match="multiply to"): + buf.rearrange("(a b) -> a b", a=3, b=4) + with pytest.raises(ValueError, match="does not factor"): + buf.rearrange("(a b) -> a b", b=4) + # select: static index bounds + with pytest.raises(ValueError, match="out of range"): + buf.select(0, 10) + with pytest.raises(ValueError, match="out of range"): + buf.select(0, -1) + # narrow: static range bounds + with pytest.raises(ValueError, match="exceeds"): + buf.narrow(0, 8, 4) + with pytest.raises(ValueError, match="must be positive"): + buf.narrow(0, 0, -1) + with pytest.raises(ValueError, match="must be non-negative"): + buf.narrow(0, -2, 4) + # sub: negative indices/starts + with pytest.raises(ValueError, match=r"in \[0, 2\)"): + grid.sub[:, -1::2] + with pytest.raises(ValueError, match="out of range"): + grid.sub[-1, :] + with pytest.raises(ValueError, match="exceeds"): + grid.sub[:, 4:12] + + +def test_buffer_select_narrow_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( + T.SwizzleLayout(3, 3, 3, swizzle_inner=True), + 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.select(0, 1) # offset 1024 = 2 * period: folds into elem_offset + B[0] = T.bfloat16(0) + C = A.narrow(1, 512, 512) # 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( + T.SwizzleLayout(3, 3, 3, swizzle_inner=True), + 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.select(1, 1) # offset 8 + B[0, 0] = T.float16(0) + C = A.select(2, 1) # offset 1 + C[0, 0] = T.float16(0) + D = A.narrow(1, 1, 2) # 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.select(1, 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( + T.SwizzleLayout(3, 3, 3, swizzle_inner=True), + 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.narrow(0, 8, 8) + 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_rearrange_singleton_size_mismatch_rejected(): + buf = tvm.tirx.decl_buffer( + (2, 3), "float16", layout=tvm.tirx.layout.TileLayout(T.S[(2, 3) : (3, 1)]) + ) + with pytest.raises(ValueError, match="does not match dim"): + buf.rearrange("a b -> b a", a=99) + + # a matching explicit size on a singleton axis is accepted + # fmt: off + @T.prim_func + def func() -> None: + T.device_entry() + A = T.alloc_buffer([2, 3], dtype="float16", scope="local", + layout=T.TileLayout(T.S[(2, 3) : (3, 1)])) + B = A.rearrange("a b -> b a", a=2) + B[0, 0] = T.float16(0) + # fmt: on + + assert [int(s) for s in _collect_buffers(func)["B"].shape] == [3, 2] + + def test_buffer_view_dtype_ir(): """Verify .view('float32') on float16: dtype correct, last dim halved, shared data.""" diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index af8605d841bf..189214290ce1 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -1006,7 +1006,7 @@ class NegateIntImmMutator(StmtExprMutator): def visit_int_imm_(self, op): # Create a new IntImm with negated value - return tir.IntImm(op.dtype, -op.value) + return tir.IntImm(op.ty, -op.value) def test_mutator_transformation(): @@ -1099,13 +1099,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 PrimExpr 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 ) @@ -1131,20 +1131,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"] diff --git a/tests/scripts/setup-pytest-env.sh b/tests/scripts/setup-pytest-env.sh index f511c7578127..171ddbc2d0d6 100755 --- a/tests/scripts/setup-pytest-env.sh +++ b/tests/scripts/setup-pytest-env.sh @@ -29,20 +29,6 @@ set -ux export TVM_PATH=`pwd` export PYTHONPATH="${TVM_PATH}/python" -# Prefer a valid sibling tirx-kernels worktree over stale editable installs. -# Some environments export TIRX_KERNELS_PATH that does not actually contain the -# tirx_kernels package (e.g. ".../tirx-kernels/kernels"), so validate before use. -tirx_kernels_path="" -if [[ -n "${TIRX_KERNELS_PATH:-}" ]] && [[ -f "${TIRX_KERNELS_PATH}/tirx_kernels/__init__.py" ]]; then - tirx_kernels_path="${TIRX_KERNELS_PATH}" -elif [[ -d "${TVM_PATH}/../tirx-kernels/tirx_kernels" ]]; then - tirx_kernels_path="${TVM_PATH}/../tirx-kernels" -fi -if [[ -n "${tirx_kernels_path}" ]]; then - export TIRX_KERNELS_PATH="${tirx_kernels_path}" - export PYTHONPATH="${tirx_kernels_path}:${PYTHONPATH}" -fi - export TVM_PYTEST_RESULT_DIR="${TVM_PATH}/build/pytest-results" mkdir -p "${TVM_PYTEST_RESULT_DIR}" pytest_errors=()