github-featured/apache--tvm
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

[TIRx] Bringup TIRx Infrastructure (#19581)

Browse Source 
## Summary

This PR adds the initial TIRx support needed for low-level programming
of Blackwell-class GPU architectures. As part of the ongoing TIRx
refactor, it introduces TVMScript support for directly scripting
advanced hardware features without relying on scheduling as the primary
programming interface.

The change keeps existing `s_tir` script support intact while making
direct scripting a first-class path for TIRx programs.

## Main Changes

- Add TIRx operator dispatch and layout infrastructure.
- Add TVMScript support for new low-level TIRx operations.
- Add analysis, transform, and lowering support for TIRx IR nodes.
- Add CUDA/Blackwell-oriented codegen and intrinsic coverage.
- Add Python and C++ integration points for TIRx scripting and runtime
support.

## Validation

- `pre-commit run --all-files`
- `ninja -C build -j32`
- `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16`
  - `1723 passed, 47 skipped, 32 warnings`
- `CUDA_VISIBLE_DEVICES=2 python -m pytest -v
tests/python/all-platform-minimal-test`
  - `37 passed, 105 skipped`
- `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis
tests/python/tirx-base tests/python/tirx-transform -n 16`
  - `664 passed, 25 skipped, 9 xfailed, 1 xpassed`

## Local CI Notes

Some full CI-equivalent jobs were not locally reproducible because this
machine is missing parts of the Apache TVM CI environment, including
`llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten,
and ARM/QEMU cross-toolchain components. Metal-specific tests were
skipped locally because no Metal runtime is available.
This commit is contained in:
Bohan Hou
2026-05-18 19:44:43 -04:00
committed by GitHub
parent bc1a904ec1
commit 859498dc01
783 changed files with 88996 additions and 9733 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/commands/tir-bench.md
+195
View File
@@ -0,0 +1,195 @@
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 |
.claude/commands/tir-build.md
+15
View File
@@ -0,0 +1,15 @@
Build TVM from the current worktree.
## Steps
1. Check that `build/` directory exists. If not, run initial setup:
```bash
mkdir -p build && cd build && cmake .. && make -j$(nproc)
```
2. If `build/` already exists, run incremental build:
```bash
cmake --build build -j$(nproc)
```
3. Report success/failure and build time.
.claude/commands/tir-test.md
+44
View File
@@ -0,0 +1,44 @@
Run the full TIRX test suite.
## Steps
1. Select the least busy GPU to avoid conflicts:
```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. Start the GPU monitor in the background so we can detect if anyone else lands on the same GPU mid-run:
```bash
GPU_LOG="/tmp/tir_test_gpu_${CUDA_VISIBLE_DEVICES}.log"
bash .claude/scripts/monitor_gpu.sh --gpu "$CUDA_VISIBLE_DEVICES" --interval 5 --log "$GPU_LOG" &
MON_PID=$!
trap 'kill $MON_PID 2>/dev/null' EXIT
```
3. Run the full test suite with xdist parallelism:
```bash
pytest tests/python/tirx/ -n 16
```
4. 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.
## Failure triage rules
**CRITICAL: Never pipe test output to `tail` or `grep` when diagnosing failures. Always capture and read full logs.**
Classify every failure into one of these categories:
- **A — Environment/import error**: Module not found, missing dependency, collection error. These are not caused by code changes.
- **B — Real kernel correctness regression**: Assertion failures (cosine_sim, numerical diff), `CUDA: unspecified launch failure`, or wrong results. **These MUST be investigated and fixed if caused by current changes.**
- **C — Secondary xdist crash**: `KeyError: <WorkerController gwXX>` after a worker abort. The KeyError itself is noise — find the underlying cause (usually category B in another worker).
**Never dismiss a failure as "pre-existing" without evidence.** If a test fails:
1. Check whether the test touches code you changed.
2. If unclear, verify on the parent commit before claiming pre-existing.
3. All failures caused by current changes MUST be fixed — not deferred.
.claude/scripts/monitor_gpu.sh
Executable
+124
View File
@@ -0,0 +1,124 @@
#!/usr/bin/env bash
# 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
# Note: deliberately not `set -u` — bash <5.2 errors on `${#assoc[@]}` when
# the associative array is empty.
GPU=""
INTERVAL=5
LOG=""
while [[ $# -gt 0 ]]; do
case "$1" in
--gpu) GPU="$2"; shift 2 ;;
--interval) INTERVAL="$2"; shift 2 ;;
--log) LOG="$2"; shift 2 ;;
-h|--help)
sed -n '2,12p' "$0" | sed 's/^# \{0,1\}//'
exit 0 ;;
*) echo "unknown arg: $1" >&2; exit 2 ;;
esac
done
if [[ -z "$GPU" ]]; then
GPU="${CUDA_VISIBLE_DEVICES:-0}"
fi
# Only the first index if CUDA_VISIBLE_DEVICES is a list.
GPU="${GPU%%,*}"
if ! [[ "$GPU" =~ ^[0-9]+$ ]]; then
echo "monitor_gpu: GPU must be an integer index (got '$GPU'); pass --gpu <n>" >&2
exit 2
fi
ME="$(id -un)"
emit() {
local line="[$(date +'%H:%M:%S')] $*"
if [[ -n "$LOG" ]]; then
printf '%s\n' "$line" | tee -a "$LOG" >&2
else
printf '%s\n' "$line" >&2
fi
}
# Returns "pid|user|mem_mib|process_name" lines for compute apps on $GPU.
snapshot() {
nvidia-smi --id="$GPU" \
--query-compute-apps=pid,process_name,used_memory \
--format=csv,noheader,nounits 2>/dev/null \
| while IFS=, read -r pid pname mem; do
pid="${pid// /}"
[[ -z "$pid" ]] && continue
local user
user="$(ps -o user= -p "$pid" 2>/dev/null | tr -d ' ')"
[[ -z "$user" ]] && user="?"
pname="${pname# }"
mem="${mem# }"
printf '%s|%s|%s|%s\n' "$pid" "$user" "$mem" "$pname"
done
}
emit "monitor_gpu started: GPU=$GPU interval=${INTERVAL}s user=$ME"
declare -A KNOWN # pid -> "user|mem|pname"
# Initial snapshot — record everyone we already see as the baseline.
while IFS='|' read -r pid user mem pname; do
[[ -z "${pid:-}" ]] && continue
KNOWN[$pid]="$user|$mem|$pname"
flag=""
[[ "$user" != "$ME" ]] && flag=" [FOREIGN]"
emit "baseline pid=$pid user=$user mem=${mem}MiB cmd=$pname$flag"
done < <(snapshot)
if [[ ${#KNOWN[@]} -eq 0 ]]; then
emit "baseline: GPU $GPU is idle"
fi
trap 'emit "monitor_gpu stopped"; exit 0' INT TERM
heartbeat_due=$(( $(date +%s) + 60 ))
while :; do
sleep "$INTERVAL"
declare -A SEEN=()
while IFS='|' read -r pid user mem pname; do
[[ -z "${pid:-}" ]] && continue
SEEN[$pid]=1
if [[ -z "${KNOWN[$pid]:-}" ]]; then
flag=""
[[ "$user" != "$ME" ]] && flag=" *** FOREIGN USER ***"
emit "NEW pid=$pid user=$user mem=${mem}MiB cmd=$pname$flag"
KNOWN[$pid]="$user|$mem|$pname"
fi
done < <(snapshot)
for pid in "${!KNOWN[@]}"; do
if [[ -z "${SEEN[$pid]:-}" ]]; then
emit "GONE pid=$pid (was: ${KNOWN[$pid]})"
unset 'KNOWN[$pid]'
fi
done
unset SEEN
now=$(date +%s)
if (( now >= heartbeat_due )); then
foreign=0
for v in "${KNOWN[@]}"; do
u="${v%%|*}"
[[ "$u" != "$ME" ]] && foreign=$((foreign+1))
done
emit "heartbeat: ${#KNOWN[@]} process(es) on GPU $GPU (${foreign} foreign)"
heartbeat_due=$(( now + 60 ))
fi
done
.gitignore
+3
View File
@@ -1,3 +1,5 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
@@ -287,3 +289,4 @@ python/tvm_ffi/
python/bin/
python/typing_extensions.py
python/*.dist-info/
pytest-of-bohanhou/
.pre-commit-config.yaml
+2
View File
@@ -15,6 +15,8 @@
# specific language governing permissions and limitations
# under the License.
exclude: ^(\.txdev/|\.claude/)
default_install_hook_types:
- pre-commit
repos:
docs/arch/introduction_to_module_serialization.rst
+1 -1
View File
@@ -79,7 +79,7 @@ location 0. In our example, we have module relationship like this:
.. code:: c++
llvm_mod:imported_modules
llvm_mod:imports
- cuda_mod
So LLVM module will have index 0, CUDA module will have index 1.
docs/deep_dive/relax/tutorials/relax_creation.py
+7 -4
View File
@@ -71,9 +71,10 @@ RelaxModule.show()
@I.ir_module
class RelaxModuleWithTIR:
@T.prim_func
@T.prim_func(s_tir=True)
def relu(x: T.handle, y: T.handle):
n, m = T.int64(), T.int64()
n = T.int64()
m = T.int64()
X = T.match_buffer(x, (n, m), "float32")
Y = T.match_buffer(y, (n, m), "float32")
for i, j in T.grid(n, m):
@@ -163,9 +164,11 @@ mod.show()
# Tensor Expression(TE), TensorIR functions or other TVM packed functions.
@T.prim_func
@T.prim_func(s_tir=True)
def tir_linear(x: T.handle, w: T.handle, b: T.handle, z: T.handle):
M, N, K = T.int64(), T.int64(), T.int64()
M = T.int64()
N = T.int64()
K = T.int64()
X = T.match_buffer(x, (M, K), "float32")
W = T.match_buffer(w, (N, K), "float32")
B = T.match_buffer(b, (N,), "float32")
docs/deep_dive/tensor_ir/tutorials/tir_creation.py
+7 -5
View File
@@ -61,7 +61,7 @@ from tvm.script import tirx as T
@I.ir_module
class MyModule:
@T.prim_func
@T.prim_func(s_tir=True)
def mm_relu(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -104,7 +104,7 @@ class MyModule:
@I.ir_module
class ConciseModule:
@T.prim_func
@T.prim_func(s_tir=True)
def mm_relu(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -143,7 +143,7 @@ dtype = "float32"
# IRModule in TVMScript
@I.ir_module
class ConciseModuleFromPython:
@T.prim_func
@T.prim_func(s_tir=True)
def mm_relu(
A: T.Buffer((M, K), dtype),
B: T.Buffer((K, N), dtype),
@@ -178,10 +178,12 @@ print(tvm.ir.structural_equal(ConciseModule, ConciseModuleFromPython))
@I.ir_module
class DynamicShapeModule:
@T.prim_func
@T.prim_func(s_tir=True)
def mm_relu(a: T.handle, b: T.handle, c: T.handle):
# Dynamic shape definition
M, N, K = T.int32(), T.int32(), T.int32()
M = T.int32()
N = T.int32()
K = T.int32()
# Bind the input buffers with the dynamic shapes
A = T.match_buffer(a, [M, K], dtype)
docs/deep_dive/tensor_ir/tutorials/tir_transformation.py
+1 -1
View File
@@ -43,7 +43,7 @@ from tvm.script import tirx as T
@I.ir_module
class MyModule:
@T.prim_func
@T.prim_func(s_tir=True)
def main(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
docs/errors.rst
+1 -1
View File
@@ -36,7 +36,7 @@ Where do these errors come from?
This error is caused by an internal invariant being violated during TVM's
execution. On a technical level, the message is generated by the
``TVM_FFI_ICHECK`` macro, found in ``3rdparty/tvm-ffi/include/tvm/ffi/error.h``.
``TVM_FFI_ICHECK`` macro, found in ``include/tvm/runtime/logging.h``.
The ``TVM_FFI_ICHECK`` macro is used in many places in the TVM code to assert
some condition is true during execution; any time the assertion fails, TVM
will exit with the error message shown above.
docs/how_to/tutorials/export_and_load_executable.py
+17 -16
View File
@@ -301,8 +301,9 @@ if RUN_EXAMPLE:
#
# **Deployment Checklist:**
# When moving to another host (via RPC or SCP), you must copy **both** files:
# 1. ``mlp_cpu.so`` (or ``mlp_cuda.so`` for GPU) - The compiled model code
# 2. ``model_params.npz`` - The model parameters (serialized as NumPy arrays)
#
# 1. ``mlp_cpu.so`` (or ``mlp_cuda.so`` for GPU) - the compiled model code
# 2. ``model_params.npz`` - the model parameters, serialized as NumPy arrays
#
# The remote machine needs both files in the same directory. The script above
# assumes they are in ``relax_export_artifacts/`` relative to the script location.
@@ -363,21 +364,21 @@ if RUN_EXAMPLE:
# FAQ
# ---
# **Can I run the ``.so`` as a standalone executable (like ``./mlp_cpu.so``)?**
# No. The ``.so`` file is a shared library, not a standalone executable binary.
# You cannot run it directly from the terminal. It must be loaded through a TVM
# runtime program (as shown in the "Loading and Running" section above). The
# ``.so`` bundles VM bytecode and compiled kernels, but still requires the TVM
# runtime to execute.
# No. The ``.so`` file is a shared library, not a standalone executable binary.
# You cannot run it directly from the terminal. It must be loaded through a TVM
# runtime program (as shown in the "Loading and Running" section above). The
# ``.so`` bundles VM bytecode and compiled kernels, but still requires the TVM
# runtime to execute.
#
# **Which devices can run the exported library?**
# The target must match the ISA you compiled for (``llvm`` in this example).
# As long as the target triple, runtime ABI, and available devices line up,
# you can move the artifact between machines. For heterogeneous builds (CPU
# plus GPU), ship the extra device libraries as well.
# The target must match the ISA you compiled for (``llvm`` in this example).
# As long as the target triple, runtime ABI, and available devices line up,
# you can move the artifact between machines. For heterogeneous builds (CPU
# plus GPU), ship the extra device libraries as well.
#
# **What about the ``.params`` and ``metadata.json`` files?**
# These auxiliary files are only generated in specific configurations. In this
# tutorial, since we pass parameters at runtime, they are not generated. When
# they do appear, they may be kept alongside the ``.so`` for inspection, but
# the essential content is typically embedded in the shared object itself, so
# deploying the ``.so`` alone is usually sufficient.
# These auxiliary files are only generated in specific configurations. In this
# tutorial, since we pass parameters at runtime, they are not generated. When
# they do appear, they may be kept alongside the ``.so`` for inspection, but
# the essential content is typically embedded in the shared object itself, so
# deploying the ``.so`` alone is usually sufficient.
docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py
+5 -5
View File
@@ -85,7 +85,7 @@ if RUN_EXAMPLE:
@I.ir_module
class MyFirstModule(BasePyModule):
@T.prim_func
@T.prim_func(s_tir=True)
def add_tir(
A: T.Buffer((4,), "float32"),
B: T.Buffer((4,), "float32"),
@@ -133,7 +133,7 @@ if RUN_EXAMPLE:
@I.ir_module
class DebugModule(BasePyModule):
@T.prim_func
@T.prim_func(s_tir=True)
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
n = T.int32()
A = T.match_buffer(var_A, (n, 4), "float32")
@@ -211,7 +211,7 @@ if RUN_EXAMPLE:
@I.ir_module
class PipelineModule(BasePyModule):
@T.prim_func
@T.prim_func(s_tir=True)
def matmul_tir(var_A: T.handle, var_B: T.handle, var_C: T.handle):
A = T.match_buffer(var_A, (2, 4), "float32")
B = T.match_buffer(var_B, (4, 3), "float32")
@@ -275,7 +275,7 @@ if RUN_EXAMPLE:
# A simple Relax module: matmul + bias + relu (a dense layer)
@I.ir_module
class DenseLayer:
@T.prim_func
@T.prim_func(s_tir=True)
def bias_add_tir(var_x: T.handle, var_b: T.handle, var_out: T.handle):
x = T.match_buffer(var_x, (2, 4), "float32")
b = T.match_buffer(var_b, (4,), "float32")
@@ -403,7 +403,7 @@ if RUN_EXAMPLE:
@I.ir_module
class DynamicModule(BasePyModule):
@T.prim_func
@T.prim_func(s_tir=True)
def scale_tir(var_x: T.handle, var_out: T.handle):
n = T.int64()
x = T.match_buffer(var_x, (n,), "float32")
docs/install/from_source.rst
+1 -1
View File
@@ -260,7 +260,7 @@ Windows-Specific Build Notes
If you're building TVM on Windows, note these platform-specific considerations:
Path Conventions
................
~~~~~~~~~~~~~~~~
- Use forward slashes (``/``) in Python/CMake paths, not Windows backslashes
- Example: ``python cmake/config.cmake`` not ``python cmake\\config.cmake``
include/tvm/ir/function.h
+17
View File
@@ -125,6 +125,23 @@ constexpr const char* kTarget = "target";
*/
constexpr const char* kGlobalSymbol = "global_symbol";
/*!
* \brief The function uses s_tir (apache-derived TIR) semantics:
* parser fills layout=None, ScriptComplete wraps body in a root SBlock,
* and printer emits `s_tir=True` on the decorator.
* Default (attr absent or False) is tirx semantics.
*
* Type: Bool
*/
constexpr const char* kSTir = "s_tir";
/*!
* \brief Number of inputs of the Primfunc
*
* Type: Int
*/
constexpr const char* kNumInputs = "num_inputs";
} // namespace attr
/*!
include/tvm/runtime/device_api.h
+3
View File
@@ -345,6 +345,8 @@ inline const char* DLDeviceType2Str(int type) {
return "webgpu";
case kDLHexagon:
return "hexagon";
case kDLTrn:
return "trn";
default:
TVM_FFI_THROW(InternalError) << "unknown type = " << type;
}
@@ -414,6 +416,7 @@ TVM_RUNTIME_DLL bool RuntimeEnabled(const ffi::String& target);
/*! \brief namespace for constant symbols */
namespace symbol {
constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
/*! \brief global function to set device */
constexpr const char* tvm_set_device = "__tvm_set_device";
} // namespace symbol
include/tvm/s_tir/data_layout.h
+74 -73
View File
@@ -19,8 +19,8 @@
/*!
* \file tvm/s_tir/data_layout.h
* \brief Layout expression to describe the data organization of a tensor.
* And BijectiveLayout to mapping two data layouts between each other.
* \brief SLayout expression to describe the data organization of a tensor.
* And SBijectiveLayout to mapping two data layouts between each other.
*/
#ifndef TVM_S_TIR_DATA_LAYOUT_H_
#define TVM_S_TIR_DATA_LAYOUT_H_
@@ -40,65 +40,65 @@
namespace tvm {
namespace tirx {
class Layout;
class SLayout;
class LayoutAxis {
class SLayoutAxis {
public:
static const LayoutAxis& Get(const char name);
static const SLayoutAxis& Get(const char name);
// Get the singleton LayoutAxis using itvar->var->name_hint
static const LayoutAxis& Get(const tirx::IterVar& itvar);
// Get the singleton SLayoutAxis using itvar->var->name_hint
static const SLayoutAxis& Get(const tirx::IterVar& itvar);
// Get the singleton LayoutAxis using name[0] (size of name must be 1).
static const LayoutAxis& Get(const std::string& name);
// Get the singleton SLayoutAxis using name[0] (size of name must be 1).
static const SLayoutAxis& Get(const std::string& name);
inline bool IsPrimal() const { return name_ >= 'A' && name_ <= 'Z'; }
inline std::string name() const { return std::string(1, name_); }
// if current axis is primal, switch the axis to its subordinate one,
// else switch to the primal.
inline const LayoutAxis& ToDual() const {
inline const SLayoutAxis& ToDual() const {
if (name_ >= 'A' && name_ <= 'Z') {
return LayoutAxis::Get(name_ - 'A' + 'a');
return SLayoutAxis::Get(name_ - 'A' + 'a');
} else {
return LayoutAxis::Get(name_ - 'a' + 'A');
return SLayoutAxis::Get(name_ - 'a' + 'A');
}
}
// return the primal axis. If it is already primal, return itself.
const LayoutAxis& ToPrimal() const { return IsPrimal() ? *this : ToDual(); }
const SLayoutAxis& ToPrimal() const { return IsPrimal() ? *this : ToDual(); }
// return the subordinate axis. If it is already subordinate, return itself.
const LayoutAxis& ToSubordinate() const { return IsPrimal() ? ToDual() : *this; }
const SLayoutAxis& ToSubordinate() const { return IsPrimal() ? ToDual() : *this; }
inline bool operator==(const LayoutAxis& rhs) const { return name_ == rhs.name_; }
inline bool operator==(const SLayoutAxis& rhs) const { return name_ == rhs.name_; }
friend std::ostream& operator<<(std::ostream& os, const LayoutAxis& l) {
friend std::ostream& operator<<(std::ostream& os, const SLayoutAxis& l) {
os << l.name();
return os;
}
private:
static const LayoutAxis UPPER_CASE[];
static const LayoutAxis LOWER_CASE[];
LayoutAxis(const LayoutAxis&);
LayoutAxis& operator=(const LayoutAxis&);
explicit LayoutAxis(const char name) : name_(name) {}
static const SLayoutAxis UPPER_CASE[];
static const SLayoutAxis LOWER_CASE[];
SLayoutAxis(const SLayoutAxis&);
SLayoutAxis& operator=(const SLayoutAxis&);
explicit SLayoutAxis(const char name) : name_(name) {}
const char name_;
};
/*!
* \brief Layout is to describe how data is organized within an N-dimention tensor.
* \brief SLayout is to describe how data is organized within an N-dimention tensor.
* It is composed of upper cases, lower cases and numbers,
* where upper case indicates a primal axis and
* the corresponding lower case with factor size indicates the subordinate axis.
* For example, NCHW16c can describe a 5-D tensor of
* [batch_size, channel, height, width, channel_block].
* Here subordinate axis channel_block=16 is the factor size of the primal axis C (channel).
* Layout for scalar is defined, while both its name and axes have size 0.
* SLayout for scalar is defined, while both its name and axes have size 0.
*/
class LayoutNode : public ffi::Object {
class SLayoutNode : public ffi::Object {
public:
/*! \brief string representation of layout, "" for scalar. */
ffi::String name;
@@ -112,26 +112,26 @@ class LayoutNode : public ffi::Object {
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<LayoutNode>()
.def_ro("name", &LayoutNode::name)
.def_ro("axes", &LayoutNode::axes);
refl::ObjectDef<SLayoutNode>()
.def_ro("name", &SLayoutNode::name)
.def_ro("axes", &SLayoutNode::axes);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.Layout", LayoutNode, ffi::Object);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SLayout", SLayoutNode, ffi::Object);
};
/*!
* \brief Managed reference to LayoutNode
* \sa LayoutNode
* \brief Managed reference to SLayoutNode
* \sa SLayoutNode
*/
class Layout : public ffi::ObjectRef {
class SLayout : public ffi::ObjectRef {
public:
explicit Layout(const ffi::Array<tirx::IterVar>& axes);
explicit SLayout(const ffi::Array<tirx::IterVar>& axes);
/*! \brief construct from a string */
Layout(const tvm::ffi::String& name) : Layout(name.operator std::string()) {} // NOLINT(*)
SLayout(const tvm::ffi::String& name) : SLayout(name.operator std::string()) {} // NOLINT(*)
/*! \brief construct from a string */
Layout(const char* name) : Layout(std::string(name)) {} // NOLINT(*)
SLayout(const char* name) : SLayout(std::string(name)) {} // NOLINT(*)
/*!
* \brief construct from a string.
@@ -143,20 +143,20 @@ class Layout : public ffi::ObjectRef {
* \param dtype The dtype of generated axes vars in the returned layout.
* It is required to be integer type.
*/
TVM_DLL Layout(const std::string& name, DataType dtype = DataType::Int(32)); // NOLINT(*)
TVM_DLL SLayout(const std::string& name, DataType dtype = DataType::Int(32)); // NOLINT(*)
/*!
* \brief access the internal node container
* \return the pointer to the internal node container
*/
LayoutNode* operator->() { return static_cast<LayoutNode*>(get_mutable()); }
SLayoutNode* operator->() { return static_cast<SLayoutNode*>(get_mutable()); }
/*!
* \brief Return an undefined layout.
* \return a (global) undefined layout.
*/
static const Layout& Undef() {
static Layout undef;
static const SLayout& Undef() {
static SLayout undef;
return undef;
}
@@ -182,18 +182,18 @@ class Layout : public ffi::ObjectRef {
* (or until the end of the layout, whichever comes first).
* \param pos The start position.
* \param len The length of the sub-layout. if 0, return layout of scalar
* \return A newly constructed Layout object.
* \return A newly constructed SLayout object.
*/
Layout SubLayout(size_t pos, size_t len) const;
SLayout SubLayout(size_t pos, size_t len) const;
/*!
* \brief Split \p axis by \p size and put the sub-axis to position \p target_pos.
* \param axis The source axis to be split. It must be a primal-axis;
* \param target_pos The target position of the newly split subordinate-axis.
* \param factor size of the sub-dimension.
* \return A newly constructed Layout object.
* \return A newly constructed SLayout object.
*/
Layout Split(const LayoutAxis& axis, size_t target_pos, int32_t factor) const;
SLayout Split(const SLayoutAxis& axis, size_t target_pos, int32_t factor) const;
/*! \return number of dimensions */
inline size_t ndim() const {
@@ -208,7 +208,7 @@ class Layout : public ffi::ObjectRef {
for (auto px : operator->()->axes) {
auto iter_vars = UnpackIterVar(px);
for (auto x : iter_vars) {
if (LayoutAxis::Get(x).IsPrimal()) {
if (SLayoutAxis::Get(x).IsPrimal()) {
ct++;
}
}
@@ -219,17 +219,17 @@ class Layout : public ffi::ObjectRef {
/*!
* \brief Returns a new layout where the dims have been expanded to match the primal dimensions.
* \param dst_layout The dst layout to which current layout has to be expanded.
* \return The expanded Layout.
* \return The expanded SLayout.
*/
inline Layout ExpandPrimal(const Layout& dst_layout) {
Layout new_src_layout;
inline SLayout ExpandPrimal(const SLayout& dst_layout) {
SLayout new_src_layout;
// 1) Find the axis which are missing in the current layout. Make them the prefix.
std::string new_src_layout_str = "";
for (auto packed_axis : dst_layout->axes) {
auto iter_vars = UnpackIterVar(packed_axis);
for (auto dst_axis : iter_vars) {
if (LayoutAxis::Get(dst_axis).IsPrimal()) {
if (!this->Contains(LayoutAxis::Get(dst_axis))) {
if (SLayoutAxis::Get(dst_axis).IsPrimal()) {
if (!this->Contains(SLayoutAxis::Get(dst_axis))) {
new_src_layout_str += dst_axis->var->name_hint;
}
}
@@ -237,7 +237,7 @@ class Layout : public ffi::ObjectRef {
}
// 2) Now, add the primal axis of the current layout.
new_src_layout_str += this->name();
new_src_layout = Layout(new_src_layout_str);
new_src_layout = SLayout(new_src_layout_str);
return new_src_layout;
}
@@ -264,7 +264,7 @@ class Layout : public ffi::ObjectRef {
* \param axis the input layout axis.
* \return the index or -1 if not found.
*/
inline int32_t IndexOf(const LayoutAxis& axis) const { return IndexOf(axis.name()); }
inline int32_t IndexOf(const SLayoutAxis& axis) const { return IndexOf(axis.name()); }
/*!
* \brief return the index of the input axis.
@@ -282,14 +282,14 @@ class Layout : public ffi::ObjectRef {
* or the size of \p axis itself (if \p axis is a subordinate-axis).
* Return -1 if \p axis is not in the layout the layout is undefined.
*/
int32_t FactorOf(const LayoutAxis& axis) const;
int32_t FactorOf(const SLayoutAxis& axis) const;
/*!
* \brief Whether the layout contains an axis.
* \param axis axis to be checked.
* \return Whether the layout contains the axis.
*/
bool Contains(const LayoutAxis& axis) const {
bool Contains(const SLayoutAxis& axis) const {
if (!defined()) return false;
for (const tirx::IterVar packed_var : operator->()->axes) {
auto iter_vars = UnpackIterVar(packed_var);
@@ -302,12 +302,12 @@ class Layout : public ffi::ObjectRef {
return false;
}
const LayoutAxis& operator[](int32_t i) const {
const SLayoutAxis& operator[](int32_t i) const {
TVM_FFI_ICHECK(defined()) << "Try to access axis from an undefined layout.";
int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
TVM_FFI_ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
const tirx::IterVar axis = operator->()->axes[index];
return LayoutAxis::Get(axis);
return SLayoutAxis::Get(axis);
}
IterVar PackedAxisAt(int32_t i) const {
@@ -329,7 +329,7 @@ class Layout : public ffi::ObjectRef {
* \param rhs Another layout.
* \return whether the two layouts are equal.
*/
inline bool Equals(const Layout& rhs) const { return name() == rhs.name(); }
inline bool Equals(const SLayout& rhs) const { return name() == rhs.name(); }
/*!
* \brief allow output string of layout to ostream
@@ -337,16 +337,16 @@ class Layout : public ffi::ObjectRef {
* \param l the layout
* \return the ostream
*/
friend std::ostream& operator<<(std::ostream& os, const Layout& l) {
friend std::ostream& operator<<(std::ostream& os, const SLayout& l) {
os << l.name();
return os;
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Layout, ffi::ObjectRef, LayoutNode);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SLayout, ffi::ObjectRef, SLayoutNode);
};
// Internal node container BijectiveLayout
class BijectiveLayoutNode : public ffi::Object {
// Internal node container SBijectiveLayout
class SBijectiveLayoutNode : public ffi::Object {
public:
/*! \brief Describes how source axes can be mapped to the destination axes,
* e.g., [i0 / 16, i1, i0 % 16] can describe NC -> NC16n
@@ -360,37 +360,37 @@ class BijectiveLayoutNode : public ffi::Object {
ffi::Array<PrimExpr> shape_backward_rule;
/*! \brief The source layout */
Layout src_layout;
SLayout src_layout;
/*! \brief The destination layout */
Layout dst_layout;
SLayout dst_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BijectiveLayoutNode>()
.def_ro("src_layout", &BijectiveLayoutNode::src_layout)
.def_ro("dst_layout", &BijectiveLayoutNode::dst_layout)
.def_ro("index_forward_rule", &BijectiveLayoutNode::index_forward_rule)
.def_ro("index_backward_rule", &BijectiveLayoutNode::index_backward_rule)
.def_ro("shape_forward_rule", &BijectiveLayoutNode::shape_forward_rule)
.def_ro("shape_backward_rule", &BijectiveLayoutNode::shape_backward_rule);
refl::ObjectDef<SBijectiveLayoutNode>()
.def_ro("src_layout", &SBijectiveLayoutNode::src_layout)
.def_ro("dst_layout", &SBijectiveLayoutNode::dst_layout)
.def_ro("index_forward_rule", &SBijectiveLayoutNode::index_forward_rule)
.def_ro("index_backward_rule", &SBijectiveLayoutNode::index_backward_rule)
.def_ro("shape_forward_rule", &SBijectiveLayoutNode::shape_forward_rule)
.def_ro("shape_backward_rule", &SBijectiveLayoutNode::shape_backward_rule);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.BijectiveLayout", BijectiveLayoutNode, ffi::Object);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.SBijectiveLayout", SBijectiveLayoutNode, ffi::Object);
};
/*!
* \brief Bijective function mapping for data layout transformation.
* Given two Layout, BijectiveLayout build and store the mapping rules,
* Given two SLayout, SBijectiveLayout build and store the mapping rules,
* provides API to transform N-dimention tensor from the source indices (i0, i1, .., im)
* to the destination indices (j0, j1, .., jm).
*/
class BijectiveLayout : public ffi::ObjectRef {
class SBijectiveLayout : public ffi::ObjectRef {
public:
/*!
* \brief The constructor
* \param src_layout The source layout
* \param dst_layout The destination layout
*/
TVM_DLL BijectiveLayout(Layout src_layout, Layout dst_layout);
TVM_DLL SBijectiveLayout(SLayout src_layout, SLayout dst_layout);
// Given the source shape, infer the destination shape.
TVM_DLL ffi::Array<PrimExpr> ForwardShape(const ffi::Array<PrimExpr>& shape) const;
@@ -401,7 +401,8 @@ class BijectiveLayout : public ffi::ObjectRef {
// Given the destination indices, recover the source indices.
TVM_DLL ffi::Array<PrimExpr> BackwardIndex(const ffi::Array<PrimExpr>& dst_index) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BijectiveLayout, ffi::ObjectRef, BijectiveLayoutNode);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBijectiveLayout, ffi::ObjectRef,
SBijectiveLayoutNode);
};
} // namespace tirx
include/tvm/script/printer/config.h
+14
View File
@@ -45,6 +45,20 @@ class PrinterConfigNode : public ffi::Object {
bool show_meta = false;
/*! \brief The prefix of IR nodes */
ffi::String ir_prefix = "I";
/*! \brief The prefix of TIR nodes */
ffi::String tir_prefix = "T";
/*!
* \brief The TIR module name used in the printed import (e.g. "tir" or "tirx").
* Used in the header comment: "from tvm.script import <tir_import_module> as <tir_prefix>".
* When tir_prefix is "Tx", set to "tirx" so the printed script uses "import tirx as Tx".
*/
ffi::String tir_import_module = "tir";
/*! \brief The prefix of TIRX nodes */
ffi::String tirx_prefix = "Tx";
/*! \brief Default buffer dtype */
DataType buffer_dtype = DataType::Float(32);
/*! \brief The prefix of Relax nodes */
ffi::String relax_prefix = "R";
/*!
* \brief The alias of the current module at cross-function call
* \note Directly use module name if it's empty.
include/tvm/script/printer/doc.h
+114 -4
View File
@@ -529,12 +529,13 @@ class OperationDocNode : public ExprDocNode {
kGtE = 23, // >=
kAnd = 24, // and
kOr = 25, // or
kBinaryEnd = 26,
kMatMul = 26, // @
kBinaryEnd = 27,
// Special
kSpecialStart = 27,
kIfThenElse = 28, // <operands[1]> if <operands[0]> else <operands[2]>
kSpecialEnd = 29
kSpecialStart = 28,
kIfThenElse = 29, // <operands[1]> if <operands[0]> else <operands[2]>
kSpecialEnd = 30
};
/*! \brief The kind of operation (operator) */
@@ -893,6 +894,64 @@ class WhileDoc : public StmtDoc {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(WhileDoc, StmtDoc, WhileDocNode);
};
/*!
* \brief Doc that represents break statement.
*
* \sa BreakDoc
*/
class BreakDocNode : public StmtDocNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BreakDocNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.BreakDoc", BreakDocNode, StmtDocNode);
};
/*!
* \brief Reference type of BreakDocNode.
*
* \sa BreakDocNode
*/
class BreakDoc : public StmtDoc {
public:
/*!
* \brief Constructor of BreakDoc.
*/
explicit BreakDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BreakDoc, StmtDoc, BreakDocNode);
};
/*!
* \brief Doc that represents continue statement.
*
* \sa ContinueDoc
*/
class ContinueDocNode : public StmtDocNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ContinueDocNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.ContinueDoc", ContinueDocNode, StmtDocNode);
};
/*!
* \brief Reference type of ContinueDocNode.
*
* \sa ContinueDocNode
*/
class ContinueDoc : public StmtDoc {
public:
/*!
* \brief Constructor of ContinueDoc.
*/
explicit ContinueDoc();
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ContinueDoc, StmtDoc, ContinueDocNode);
};
/*!
* \brief Doc that represents for statement.
*
@@ -1240,6 +1299,57 @@ class DocStringDoc : public StmtDoc {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DocStringDoc, StmtDoc, DocStringDocNode);
};
/*!
* \brief Doc that represents call to an TIRX operator
*
* \sa OpCallDoc
*/
class OpCallDocNode : public StmtDocNode {
public:
/*! \brief The callee of this function call */
ExprDoc callee{ffi::UnsafeInit()};
/*! \brief The positional arguments */
ffi::Array<Doc> args;
/*! \brief The workspace of this op call */
ffi::Optional<DictDoc> workspace{std::nullopt};
/*! \brief The config of this op call */
ffi::Optional<DictDoc> config{std::nullopt};
/*! \brief The optional dispatch variant of this op call */
ffi::Optional<ExprDoc> dispatch{std::nullopt};
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<OpCallDocNode>()
.def_ro("callee", &OpCallDocNode::callee)
.def_ro("args", &OpCallDocNode::args)
.def_ro("workspace", &OpCallDocNode::workspace)
.def_ro("config", &OpCallDocNode::config)
.def_ro("dispatch", &OpCallDocNode::dispatch);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.printer.OpCallDoc", OpCallDocNode, StmtDocNode);
};
/*!
* \brief Reference type of OpCallDocNode.
*
* \sa OpCallDocNode
*/
class OpCallDoc : public StmtDoc {
public:
/*!
* \brief Constructor of OpCallDoc
* \param callee The callee of this function call.
* \param args The positional arguments.
* \param workspace The workspace of this op call.
* \param config The config of this op call.
* \param dispatch The optional dispatch variant name of this op call.
*/
explicit OpCallDoc(ExprDoc callee, ffi::Array<Doc> args, ffi::Optional<DictDoc> workspace,
ffi::Optional<DictDoc> config, ffi::Optional<ExprDoc> dispatch = std::nullopt);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(OpCallDoc, StmtDoc, OpCallDocNode);
};
} // namespace printer
} // namespace script
} // namespace tvm
include/tvm/tirx/analysis.h
+30 -1
View File
@@ -33,7 +33,6 @@
#include <tvm/tirx/op_attr_types.h>
#include <tvm/tirx/stmt.h>
#include <optional>
#include <string>
namespace tvm {
@@ -240,6 +239,36 @@ TVM_DLL Pass VerifySSA();
*/
TVM_DLL Pass VerifyMemory();
/*!
* \brief Pass variant of VerifyGPUCode.
*
* \param constraints The dict to specify constraints to check.
*
* \returns The pass.
* \sa tvm::tir::VerifyGPUCode
*/
/******** TIRx analysis helpers ********/
/*!
* \brief Verify if the given TIRX is well-formed.
* \param func The PrimFunc to be verified.
* \param assert_mode The indicator if it raises an error when the function is not well-formed.
* \param device_func The indicator if it is a device function.
* \return Whether it is a well-formed TIRX function.
*/
TVM_DLL bool VerifyTIRxWellFormed(const PrimFunc& func, bool assert_mode = true,
bool device_func = false);
/*!
* \brief Verify if the TIRX in the given IRMOdule is well-formed.
* \param mod The IRModule to be verified.
* \param assert_mode The indicator if it raises an error when the function is not well-formed.
* \param device_func The indicator if it is a device function.
* \return Whether it is a well-formed TIRX module.
*/
TVM_DLL bool VerifyTIRxWellFormed(const IRModule& mod, bool assert_mode = true,
bool device_func = false);
} // namespace transform
} // namespace tirx
} // namespace tvm
include/tvm/tirx/async_structs.h
+103
View File
@@ -0,0 +1,103 @@
/*
* 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 <tvm/ffi/object.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/exec_scope.h>
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<ffi::String, tvm::tirx::Buffer> workspace;
/*! \brief The schedule config of the pipeline. */
ffi::Map<ffi::String, ffi::Any> schedule_config;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PipelineNode>()
.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<ffi::String, tvm::tirx::Buffer> workspace = {},
ffi::Map<ffi::String, ffi::Any> 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<CopyPipelineNode>();
}
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<ffi::String, tvm::tirx::Buffer> workspace = {},
ffi::Map<ffi::String, ffi::Any> 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_
include/tvm/tirx/buffer.h
+51 -6
View File
@@ -21,15 +21,15 @@
* \file tvm/tirx/buffer.h
* \brief Symbolic n-dimensional array, to represent a memory buffer.
*/
#ifndef TVM_TIR_BUFFER_H_
#define TVM_TIR_BUFFER_H_
#ifndef TVM_TIRX_BUFFER_H_
#define TVM_TIRX_BUFFER_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/script/printer/config.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/var.h>
#include <string>
@@ -110,6 +110,16 @@ class BufferNode : public ffi::Object {
* Reserved debug information.
*/
mutable Span span;
/*! \brief The layout of the buffer */
ffi::Optional<Layout> layout;
/*! \brief The allocated address of the buffer.
* The address might be multi-dimensional based on its scope.
* For example, trn.psum takes 2D address, representing (bank, offset).
*/
ffi::Array<PrimExpr> allocated_addr;
/*! \brief constructor */
BufferNode() {}
@@ -127,7 +137,9 @@ class BufferNode : public ffi::Object {
.def_ro("data_alignment", &BufferNode::data_alignment)
.def_ro("offset_factor", &BufferNode::offset_factor)
.def_ro("buffer_type", &BufferNode::buffer_type)
.def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore());
.def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore())
.def_ro("layout", &BufferNode::layout)
.def_ro("allocated_addr", &BufferNode::allocated_addr);
}
/*! \return preferred index type for this buffer node */
@@ -140,8 +152,11 @@ class BufferNode : public ffi::Object {
* Returns the buffer offset, in number of elements of type dtype,
* without adjusting for number of lanes. (e.g. The number of
* float16x4 elements in a buffer of type float16x4.)
*
* \param index The index to be accessed.
* \param inner Ignore the elem_offset, return inner offset only
*/
ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index) const;
ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index, bool inner = false) const;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
@@ -161,7 +176,8 @@ class Buffer : public ffi::ObjectRef {
TVM_DLL Buffer(Var data, DataType dtype, ffi::Array<PrimExpr> shape, ffi::Array<PrimExpr> strides,
PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor,
BufferType buffer_type, ffi::Array<IntImm> axis_separators = {},
Span span = Span());
Span span = Span(), ffi::Optional<Layout> layout = std::nullopt,
ffi::Array<PrimExpr> allocated_addr = {});
/*!
* \brief Return a new buffer that is equivalent with current one
@@ -221,11 +237,40 @@ class Buffer : public ffi::ObjectRef {
*/
ffi::Array<PrimExpr> OffsetOf(ffi::Array<PrimExpr> index) const;
/*!
* \brief Get the buffer_offset op for the given index.
* \param index The index to be accessed.
* \return The buffer_offset op.
*/
PrimExpr OffsetOf_p(const ffi::Array<PrimExpr>& indices) const;
/*!
* \brief Return the storage scope associated with this buffer.
*/
TVM_DLL ffi::String scope() const;
/*!
* \brief Return a new buffer with the allocated address.
*/
TVM_DLL Buffer with_allocated_addr(ffi::Array<PrimExpr> allocated_addr) const;
/*!
* \brief Return true if the buffer is a scalar.
* \param alloc_or_decl Whether to consider alloc_scalar and decl_scalar as scalar. True for
* alloc_scalar, False for decl_scalar.
*/
TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
/*!
* \brief Return a new buffer with the dtype.
*/
TVM_DLL Buffer with_dtype(DataType dtype) const;
/*!
* \brief Return a new buffer with the data.
*/
TVM_DLL Buffer with_data(Var data) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Buffer, ffi::ObjectRef, BufferNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferNode);
};
include/tvm/tirx/builtin.h
+261 -221
View File
@@ -67,6 +67,25 @@ TVM_DLL const Op& reinterpret();
*/
TVM_DLL const Op& likely();
/*!
* \brief Thread-set filter predicate. Used as the condition of an IfThenElse
* to narrow the active thread set A for the then-branch. Two forms:
* filter(var, lo, hi) -- range form, true iff var in [lo, hi)
* filter(var, cond) -- predicate form (e.g. var == k); true iff cond
* `var` must be a ScopeIdDef-declared Var at parse time (Verifier Rule 2).
*/
TVM_DLL const Op& filter();
/*!
* \brief Analysis-only active-thread selector.
*
* ``selector(var, pred)`` denotes the unique value of ``var`` in the current
* active domain for which ``pred`` is true. It is used only inside
* ExecContext/DispatchContext metadata, for predicates such as
* ``ptx.elect_sync()`` whose selected lane cannot be inferred structurally.
*/
TVM_DLL const Op& selector();
/*!
* \brief Bitwise and operator.
*/
@@ -496,7 +515,7 @@ TVM_DLL const Op& tvm_storage_sync();
*
* Parameter width indicates the number of threads involved in one
* shuffle. See CUDA document for __shfl_sync, __shfl_up_sync,
* __shfl_down_sync and __activemask.
* __shfl_down_sync, __shfl_xor_sync and __activemask.
*
* Parameter warp_size is the size of a warp, which helps a backend
* to determine whether the width parameter is legal.
@@ -505,8 +524,15 @@ TVM_DLL const Op& tvm_storage_sync();
TVM_DLL const Op& tvm_warp_shuffle();
TVM_DLL const Op& tvm_warp_shuffle_up();
TVM_DLL const Op& tvm_warp_shuffle_down();
TVM_DLL const Op& tvm_warp_shuffle_xor();
TVM_DLL const Op& tvm_warp_activemask();
/*!
* \brief Initialize the global barrier.
* Call this at beginning of kernel that need global barrier.
*/
TVM_DLL const Op& tvm_global_barrier_kinit();
/*!
* \brief See pesudo code
*
@@ -520,226 +546,6 @@ TVM_DLL const Op& tvm_warp_activemask();
*/
TVM_DLL const Op& tvm_thread_allreduce();
// TODO(tvm-team) TensorCore specific intrinsics should be directly registered under
// cuda. namespace and used through op.
/*!
* \brief tvm intrinsic for tensor core load operators.
*
* void tvm_load_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr buffer_ptr, Expr stride,
* StringImm layout) {
* // m, n, k are the shape of wmma fragment.
* // Determine fragment layout(column-major or row major) by layout.
* // fragments must be in 'wmma.matrix_a' or 'wmma.matrix_b' scope.
* nvcuda::wmma::load_matrix_sync(fragment[index], buffer_ptr, stride);
* }
*/
TVM_DLL const Op& tvm_load_matrix_sync();
/*!
* \brief tvm intrinsic for tensor core mma_sync operators.
*
* void tvm_mma_sync(Var fragment_d, Expr index_d,
* Var fragment_a, Expr index_a,
* Var fragment_b, Expr index_b,
* Var fragment_c, Expr index_c) {
* nvcuda::wmma::mma_sync(fragment_d[index_d], fragment_a[index_a],
* fragment_b[index_b], fragment_c[index_c]);
* }
*/
TVM_DLL const Op& tvm_mma_sync();
/*!
* \brief tvm intrinsic for tensor core bmma_sync operators.
*
* void tvm_bmma_sync(Var fragment_d, Expr index_d,
* Var fragment_a, Expr index_a,
* Var fragment_b, Expr index_b,
* Var fragment_c, Expr index_c) {
* nvcuda::wmma::bmma_sync(fragment_d[index_d], fragment_a[index_a],
* fragment_b[index_b], fragment_c[index_c]);
* }
*/
TVM_DLL const Op& tvm_bmma_sync();
/*!
* \brief tvm intrinsic for tensor core fill_fragment operators.
*
* void tvm_fill_fragment(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr value) {
* // m, n, k are the shape of wmma fragment
* // fragments must be in 'wmma.accumulator' scope.
* nvcuda::wmma::fill_fragment(fragment[index], value);
* }
*/
TVM_DLL const Op& tvm_fill_fragment();
/*!
* \brief tvm intrinsic for tensor core store operators.
*
* void tvm_store_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr buffer_ptr, Expr stride,
* StringImm layout) {
* // m, n, k are the shape of wmma fragment
* // fragments must be in 'wmma.accumulator' scope.
* nvcuda::wmma::store_matrix_sync(fragment[index], buffer_ptr, stride, layout);
* }
*/
TVM_DLL const Op& tvm_store_matrix_sync();
/*!
* \brief tvm intrinsic for ptx tensor core mma instructions.
*
* void ptx_mma(StringImm shape, StringImm A_layout, StringImm B_layout,
* StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
* Var multiplicand_a, Expr a_index,
* Var multiplicand_b, Expr b_index,
* Var accumulator, Expr c_index, bool saturate);
*/
TVM_DLL const Op& ptx_mma();
/*!
* \brief tvm intrinsic for ptx predicate load with 32-bit data type.
*
*/
TVM_DLL const Op& ptx_ldg32();
/*!
* \brief tvm intrinsic for ptx predicate load with 32-bit data type.
*
*/
TVM_DLL const Op& ptx_ldg32();
/*!
* \brief tvm intrinsic for sparse tensor core ptx instructions.
*
* void ptx_mma_sp(StringImm shape, StringImm A_layout, StringImm B_layout,
* StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
* Var multiplicand_a, Expr a_index,
* Var multiplicand_b, Expr b_index,
* Var accumulator, Expr c_index,
* Var metadata, Expr meta_index,
* Var sparse_selector, bool saturate);
*/
TVM_DLL const Op& ptx_mma_sp();
/*!
* \brief tvm intrinsic for ptx load matrix from shared memory.
*
* void ptx_ldmatrix(Bool trans, IntImm num, StringImm type,
* Var local_ptr, Expr local_offset,
* Var smem_ptr, Expr smem_offset);
*/
TVM_DLL const Op& ptx_ldmatrix();
/*!
* \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async
*
* void ptx_cp_async(Var shared_ptr,
* Expr shared_offset,
* Var global_ptr,
* Expr global_offset,
* size_t bytes);
*/
TVM_DLL const Op& ptx_cp_async();
/*!
* \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async.bulk
*
* void ptx_cp_async(Var shared_ptr,
* Expr shared_offset,
* Var global_ptr,
* Expr global_offset,
* size_t bytes,
* int barrier_id);
*/
TVM_DLL const Op& ptx_cp_async_bulk();
/*!
* \brief tvm intrinsics for ptx async copy commit and wait.
*
* void ptx_commit_group();
* void ptx_wait_group(int num);
*
*/
TVM_DLL const Op& ptx_commit_group();
TVM_DLL const Op& ptx_wait_group();
/*!
* \brief tvm intrinsics for ptx async copy barrier using cp.async.mbarrier.arrive
*
* ptx_cp_async_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_cp_async_barrier();
/*!
* \brief tvm intrinsics for ptx barrier initialization of thread count using mbarrier.init
*
* ptx_init_barrier_thread_count(int barrier_id, int thread_count)
*
*/
TVM_DLL const Op& ptx_init_barrier_thread_count();
/*!
* \brief tvm intrinsics for ptx barrier arrival using mbarrier.arrive
*
* ptx_arrive_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_arrive_barrier();
/*!
* \brief tvm intrinsic for ptx barrier arrival with expect tx using mbarrier.arrive.expect_tx
*
* ptx_arrive_barrier_expect_tx(int barrier_id, int byte_count)
*
*/
TVM_DLL const Op& ptx_arrive_barrier_expect_tx();
/*!
* \brief tvm intrinsics for ptx barrier wait using mbarrier.try_wait
*
* ptx_wait_barrier(int barrier_id)
*
*/
TVM_DLL const Op& ptx_wait_barrier();
/*!
* \brief tvm intrinsics to create N barriers
*
* ptx_wait_barrier(int barrier_count)
*
*/
TVM_DLL const Op& create_barriers();
/*!
* \brief tvm intrinsic for storing the result of PTX MMA into a destination pointer.
* For example, if each thread in a warp of size 32 has 4 elements from the result of
* m16xn8xk16 MMA in its registers, this intrinsic can be used to store the result in a
* 16x8 region in shared or global memory.
*
* There is no real PTX instruction that does that, but we want to hide details of
* complex index manipulation behind this intrinsic to simplify TIR lowering passes (e.g.
* LowerWarpMemory).
*
* void mma_store(IntImm m, IntImm n, Var dst_ptr, Var src_ptr, Expr src_offset, Var dst_stride);
*/
TVM_DLL const Op& mma_store();
/*!
* \brief tvm intrinsic for zero-initializing an MMA accumulation register.
* For example, if each thread in a warp of size 32 has 8 elements from the A matrix in
* m16xn8xk16 MMA in its registers, this intrinsic can be used to zero-initialize its
* 4 accumulation registers.
*
* There is no real PTX instruction that does that, but we introduce this intrinsic for the
* same reason as mma_store above.
*
* void mma_fill(IntImm local_size, Var local_ptr, Expr offset);
*/
TVM_DLL const Op& mma_fill();
// Metal SimdGroup matrix intrinsics
/*!
@@ -999,6 +805,12 @@ TVM_DLL const Op& get_active_lane_mask();
/*! \brief Annotate a predicate not be considered as target condition of loop partition. */
TVM_DLL const Op& ignore_loop_partition();
/*!
* \brief Get the element offset of a buffer given logical indices.
The offset is determined by the layout of the buffer.
*/
TVM_DLL const Op& buffer_offset();
/*! \brief The kind of structure field info used in intrinsic */
enum TVMStructFieldKind : int {
@@ -1024,6 +836,234 @@ enum TVMStructFieldKind : int {
// Generic int64 array element access: ((int64_t*)buf)[index]
kInt64ArrayElem,
};
/*!
* \brief Print the content of a buffer during runtime.
*/
TVM_DLL const Op& print_buffer();
/*!
* \brief tvm intrinsic for initializing the CUDA profiler, and store profiling result in a buffer.
*
* void timer_init_cuda(Var profiler_buffer, Var profiler_tag, Var profiler_write_offset, int
* num_groups, Expr group_id) {
* // initialize the tag and write to pos 0 in the buffer
* // initialize write offset for every leader thread in warp group across all blocks
* }
*/
TVM_DLL const Op& timer_init_cuda();
/*!
* \brief tvm intrinsic for starting the timer for profiling a specific event,
* and storing profiling result in a buffer.
*
* void timer_start_cuda(IntImm event_type, Var profiler_buffer, Var profiler_tag,
* Var profiler_write_offset, IntImm profiler_write_stride, Expr leader_cond)
* {
* // each leader thread in warp group gets the time stamp and event type, combine with the tag
* // and write to corresponding offset in buffer
* // each leader thread advance offset by stride
* }
*/
TVM_DLL const Op& timer_start_cuda();
/*!
* \brief tvm intrinsic for ending the timer for profiling a specific event,
* and storing profiling result in a buffer.
*
* void timer_end_cuda(IntImm event_type, Var profiler_buffer, Var profiler_tag,
* Var profiler_write_offset, IntImm profiler_write_stride, Expr leader_cond) {
* // each leader thread in warp group gets the time stamp and event type, combine with the tag
* // and write to corresponding offset in buffer
* // each leader thread advance offset by stride
* }
*/
TVM_DLL const Op& timer_end_cuda();
/*!
* \brief tvm intrinsic for finalize the timer for profiling,
* and storing profiling result in a buffer.
*
* void timer_finalize_cuda(Var profiler_buffer, Var profiler_tag, Var profiler_write_offset,
* IntImm profiler_write_stride, Expr leader_cond) {
* // each leader thread in warp group gets the time stamp and end signal, combine with the tag
* // and write to corresponding offset in buffer
* // each leader thread advance offset by stride
* }
*/
TVM_DLL const Op& timer_finalize_cuda();
/*!
* \brief tvm intrinsic for cuda atomic add instruction
*/
TVM_DLL const Op& cuda_atomic_add();
/*!
* \brief tvm intrinsic for cuda thread fence instruction
*/
TVM_DLL const Op& cuda_thread_fence();
/*!
* \brief Warp-level butterfly shuffle-XOR reduction.
*
* cuda_warp_reduce(value, op, width) reduces value across width adjacent
* lanes using the specified operation ("sum", "max", "min").
*/
TVM_DLL const Op& cuda_warp_reduce();
/*!
* \brief CTA-wide reduction via warp shuffle + shared memory.
*
* cuda_cta_reduce(value, op, num_warps, scratch) reduces value across
* the entire CTA using the specified operation ("sum", "max", "min").
*/
TVM_DLL const Op& cuda_cta_reduce();
/*!
* \brief Typed load/store copy of num_bytes bytes.
*
* cuda_copy_bytes(dst, src, num_bytes) copies num_bytes bytes from src to dst
* using a single typed load/store (uint4, uint2, unsigned int, etc.).
* num_bytes must be one of {1, 2, 4, 8, 16}.
*/
TVM_DLL const Op& cuda_copy_bytes();
/*!
* \brief tvm intrinsic for cuda warp sync instruction
*/
TVM_DLL const Op& cuda_warp_sync();
/*!
* \brief tvm intrinsic for cuda block-wide sync (syncthreads)
*/
TVM_DLL const Op& cuda_cta_sync();
/*!
* \brief tvm intrinsic for cuda grid-wide sync (cooperative groups)
*/
TVM_DLL const Op& cuda_grid_sync();
/*!
* \brief tvm intrinsic that returns ``cooperative_groups::thread_rank()``
* for the enclosing CTA (linear thread index within the block).
*/
TVM_DLL const Op& cuda_thread_rank();
/*!
* \brief tvm intrinsic for cuda half to float conversion
*/
TVM_DLL const Op& cuda_half2float();
/*!
* \brief tvm intrinsic for cuda bfloat16 to float conversion
*/
TVM_DLL const Op& cuda_bfloat162float();
/*!
* \brief tvm intrinsic for a helper converting float2 to half2 with rounding
*/
TVM_DLL const Op& cuda_float22half2();
/*!
* \brief tvm intrinsic to trap when an assertion failed (cond == false)
*/
TVM_DLL const Op& cuda_trap_when_assert_failed();
/*!
* \brief tvm intrinsic to modify runtime instruction descriptor
*/
TVM_DLL const Op& cuda_runtime_instr_desc();
/*!
* \brief tvm intrinsic to convert 8 half2 lanes to 8 float2 lanes
*/
TVM_DLL const Op& cuda_half8tofloat8();
/*!
* \brief tvm intrinsic to convert 8 float2 lanes to 8 half2 lanes with rounding
*/
TVM_DLL const Op& cuda_float8tohalf8();
/*!
* \brief tvm intrinsic for cuda syncthreads_and instruction
*/
TVM_DLL const Op& cuda_syncthreads_and();
/*!
* \brief tvm intrinsic for cuda syncthreads_or instruction
*/
TVM_DLL const Op& cuda_syncthreads_or();
/*!
* \brief tvm intrinsic for cuda nano sleep instruction
*/
TVM_DLL const Op& cuda_nano_sleep();
/*!
* \brief tvm intrinsic for cuda atomic compare and swap instruction
*/
TVM_DLL const Op& cuda_atomic_cas();
/*!
* \brief tvm intrinsic for cuda printf instruction
*/
TVM_DLL const Op& cuda_printf();
/*!
* \brief tvm intrinsic for cuda ldg instruction
*/
TVM_DLL const Op& cuda_ldg();
/*!
* \brief tvm intrinsic for cuda tmem address calculation
*/
TVM_DLL const Op& cuda_get_tmem_addr();
/*!
* \brief tvm intrinsic for PTX fast exp2 approximation (ex2.approx.ftz.f32)
*/
TVM_DLL const Op& ptx_exp2();
/*!
* \brief tvm intrinsic for PTX fast reciprocal approximation (rcp.approx.ftz.f32)
*/
TVM_DLL const Op& ptx_rcp();
/*!
* \brief tvm intrinsic for PTX warp-wide any predicate (__any_sync)
*/
TVM_DLL const Op& ptx_any_sync();
/*!
* \brief tvm intrinsic for PTX 3-input max instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_reduce3_max_f32();
/*!
* \brief tvm intrinsic for PTX 3-input min instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_reduce3_min_f32();
/*!
* \brief tvm intrinsic for PTX packed add instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_add_packed_f32x2();
/*!
* \brief tvm intrinsic for PTX packed subtract instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_sub_packed_f32x2();
/*!
* \brief tvm intrinsic for PTX packed multiply instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_mul_packed_f32x2();
/*!
* \brief tvm intrinsic for PTX packed FMA instruction (sm_100a+)
*/
TVM_DLL const Op& ptx_fma_packed_f32x2();
} // namespace builtin
} // namespace tirx
} // namespace tvm
include/tvm/tirx/exec_context.h
+155
View File
@@ -0,0 +1,155 @@
/*
* 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/exec_context.h
* \brief Compile-time ExecContext state: the active thread set ``A`` as a
* TileLayout and the (inter, intra) split under the current scope kind,
* threaded through the IR walker so per-op lowerers see the precise execution
* shape at each site.
*
* Mirrors the pure-Python implementation in python/tvm/tirx/exec_context.py.
*/
#ifndef TVM_TIRX_EXEC_CONTEXT_H_
#define TVM_TIRX_EXEC_CONTEXT_H_
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/var.h>
#include <string>
#include <unordered_map>
#include <vector>
namespace tvm {
namespace tirx {
/*! \brief Warpgroup size in warps (hardware-fixed). */
constexpr int kWgSize = 4;
/*! \brief Active slice offset + stride * [0, extent) encoded on one TileLayout axis. */
struct AxisRange {
PrimExpr extent;
PrimExpr offset;
PrimExpr stride;
/*! \brief Intersect with [lo, hi). Returns false if the result is empty. */
bool Intersect(int64_t lo, int64_t hi, AxisRange* out) const;
/*! \brief Intersect with values satisfying axis % modulus == residue. */
bool Modulo(int64_t modulus, int64_t residue, AxisRange* out) const;
};
/*!
* \brief Active thread set A.
* The source of truth is ``layout``:
* shard = active axes with extents
* offset = per-axis lower bound, possibly a selector PrimExpr
*/
struct ActiveSet {
TileLayout layout;
int64_t size() const;
bool GetAxis(const std::string& axis, AxisRange* out) const;
bool HasAxis(const std::string& axis) const;
ActiveSet WithAxis(const std::string& axis, const AxisRange& range) const;
std::vector<std::string> AxisNames() const;
};
/*!
* \brief One scope_switch split. Fields are sparse dicts keyed by active-set
* axis name, e.g. laneid/warpid/cta_id/wid_in_wg/wgid or factorized CTA axes
* such as cbx/cby/cbz. An empty map denotes the empty layout (e.g. intra under
* scope_kind=thread).
*/
struct ExecSplit {
std::unordered_map<std::string, AxisRange> inter;
std::unordered_map<std::string, AxisRange> intra;
};
/*! \brief Initial A at T.kernel() entry: all threads active, offsets zero. */
TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
TVM_DLL ActiveSet InitialActiveSet(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
const std::vector<std::pair<std::string, int64_t>>& cta_axes);
/*!
* \brief Narrow A on the lane bound to ``binding``.
*
* The ScopeBinding maps directly to which native axis (laneid/warpid/cta_id)
* to narrow, and for warpid whether to narrow the full axis (kCtaWarp), the
* outer factor (kCtaWarpgroup), or the inner factor (kWarpgroupWarp).
*
* Bindings with no single-lane representation are conservative: cluster_id is
* not a filter target; flat thread ids are accepted only when the range can be
* represented as a rectangular lane/warp active set.
*/
TVM_DLL bool FilterNarrow(const ActiveSet& A, ScopeBinding binding, int64_t lo, int64_t hi,
ActiveSet* out, std::string* err);
/*!
* \brief Factor A into (inter, intra) for target scope_kind.
*
* Returns false on factoring failure (warpgroup with warpid lane that
* crosses a warpgroup boundary unaligned) and writes reason to *err.
*/
TVM_DLL bool ScopeSwitch(const ActiveSet& A, ScopeKind scope_kind, ExecSplit* out,
std::string* err);
/*! \brief Per-program-point ExecContext: active set + scope kind + split. */
struct ExecContext {
ActiveSet A;
ScopeKind scope_kind = ScopeKind::kKernel;
ExecSplit split; // (inter, intra) of current A under current scope_kind
/*! \brief Kernel-entry ctor. */
static ExecContext AtKernelEntry(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext);
static ExecContext AtKernelEntry(int64_t lane_ext, int64_t warp_ext, int64_t cta_ext,
const std::vector<std::pair<std::string, int64_t>>& cta_axes);
/*! \brief Apply filter; scope_kind preserved, split recomputed. */
bool WithFilter(ScopeBinding binding, int64_t lo, int64_t hi, ExecContext* out,
std::string* err) const;
/*! \brief Apply a unique-value selector filter on one scope id Var. */
bool WithSelector(ScopeBinding binding, PrimExpr selector, ExecContext* out,
std::string* err) const;
/*! \brief Apply filter on a factorized CTA axis such as cbx/cby/cbz. */
bool WithCtaAxisFilter(const std::string& axis, int64_t lo, int64_t hi, ExecContext* out,
std::string* err) const;
/*! \brief Apply modulo filter on a factorized CTA axis such as cbx/cby/cbz. */
bool WithCtaAxisModulo(const std::string& axis, int64_t modulus, int64_t residue,
ExecContext* out, std::string* err) const;
/*! \brief Apply scope_switch; A preserved, split recomputed for new scope_kind. */
bool WithScopeSwitch(ScopeKind new_scope_kind, ExecContext* out, std::string* err) const;
};
/*!
* \brief Encode one side of an ExecSplit (inter or intra) as the FFI map used
* by ``DispatchContextNode::{inter, intra}``: axis name -> [extent, offset]
* for unit-stride axes, or [extent, offset, stride] for strided axes.
*/
TVM_DLL ffi::Map<ffi::String, ffi::Array<PrimExpr>> EncodeSplitSide(
const std::unordered_map<std::string, AxisRange>& side);
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_EXEC_CONTEXT_H_
include/tvm/tirx/exec_scope.h
+248
View File
@@ -0,0 +1,248 @@
/*
* 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/block_scope.h
* \brief Definition of execution scope
*/
#ifndef TVM_TIRX_EXEC_SCOPE_H_
#define TVM_TIRX_EXEC_SCOPE_H_
#include <tvm/ffi/container/variant.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/var.h>
#include <string>
#include <utility>
namespace tvm {
namespace tirx {
/*!
* \brief The target execution scope kind of an ExecScopeStmt.
*
* Replaces the string-keyed name of ExecScope. One value per user-facing
* `with T.<kind>():` construct, plus ``kWorld`` for the cross-kernel root
* scope used by axe-layout's ``pid`` axis. Ordered from coarsest to finest;
* smaller integer = wider scope, so ``ScopeKindHigher`` is a plain ``<``.
*/
enum class ScopeKind : int {
kWorld = 0,
kKernel = 1,
kCluster = 2,
kCta = 3,
kWarpgroup = 4,
kWarp = 5,
kThread = 6,
};
/*! \brief Convert a ScopeKind to its string name (e.g. kKernel -> "kernel"). */
TVM_DLL std::string ScopeKindToString(ScopeKind kind);
/*! \brief Parse a string name to a ScopeKind. FATAL if unknown. */
TVM_DLL ScopeKind StringToScopeKind(const ffi::String& name);
/*!
* \brief The binding between a parent scope and a child scope as used by a
* `ScopeIdDef`. The closed enum of valid (parent -> cur) pairs.
*
* Single-axis bindings (target one ActiveSet box axis -- ``laneid`` /
* ``warpid`` / ``cta_id``, possibly via a warpid factor lane):
* kKernelCta, kClusterCta -> cta_id (flat)
* kCtaWarp -> warpid (flat)
* kCtaWarpgroup -> warpid (outer factor; warpgroup index)
* kWarpgroupWarp -> warpid (inner factor; warp-within-wg index)
* kWarpThread -> laneid (flat)
* kKernelCluster -> not a filter target (cluster_id by design)
* kClusterCtaPair -> hardware CTA pair id (cluster CTA rank % 2)
*
* Multi-axis (flat-thread) bindings -- linearize across two ActiveSet
* axes; ``T.filter(var, lo, hi)`` cannot narrow them as a contiguous box
* range, so they fall back to plain predicate semantics:
* kCtaThread -> threadIdx.x within a CTA (laneid * warpid)
* kWarpgroupThread -> threadIdx.x within a warpgroup (laneid * wid_in_wg)
*/
enum class ScopeBinding : int {
kKernelCluster = 0,
kKernelCta = 1,
kClusterCta = 2,
kCtaWarpgroup = 3,
kCtaWarp = 4,
kWarpgroupWarp = 5,
kWarpThread = 6,
kCtaThread = 7,
kWarpgroupThread = 8,
kClusterCtaPair = 9,
};
/*! \brief Convert a ScopeBinding to its (parent, cur) string pair. */
TVM_DLL std::pair<ffi::String, ffi::String> ScopeBindingToStringPair(ScopeBinding binding);
/*! \brief Parse a (parent, cur) string pair to a ScopeBinding. FATAL if unknown. */
TVM_DLL ScopeBinding StringPairToScopeBinding(const ffi::String& parent, const ffi::String& cur);
/******** Definition of ScopeId ********/
class ScopeIdDefNode : public ffi::Object {
public:
/*! \brief The ScopeId defined */
ffi::Array<Var> def_ids;
/*!
* \brief The extents of the ScopeId.
*
* NullOpt means the extent is *deferred*: the user wrote e.g.
* ``bx = T.cta_id()`` without specifying the extent, and the value will be
* inferred from sibling ScopeIdDefs at LowerTIRx entry via the verifier's
* BFS closure. Deferred form requires ``def_ids.size() == 1`` (single axis
* only -- multi-axis defers have no well-defined recovery).
*
* Explicit (Some) form preserves the per-axis shape, e.g. ``[3, 4, 5]``
* for ``T.cta_id([3, 4, 5])``.
*/
ffi::Optional<ffi::Array<PrimExpr>> extents;
/*! \brief The (parent, cur) binding of this scope id as a closed enum. */
ScopeBinding scope;
/*!
* \brief Optional preferred extents (cluster→cta only).
* Maps to cudaLaunchAttributePreferredClusterDimension (CUDA 12.8+).
*/
ffi::Optional<ffi::Array<PrimExpr>> preferred_extents;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScopeIdDefNode>()
.def_ro("def_ids", &ScopeIdDefNode::def_ids, refl::AttachFieldFlag::SEqHashDef())
.def_ro("extents", &ScopeIdDefNode::extents)
.def_ro("scope", &ScopeIdDefNode::scope)
.def_ro("preferred_extents", &ScopeIdDefNode::preferred_extents);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ScopeIdDef", ScopeIdDefNode, ffi::Object);
};
class ScopeIdDef : public ffi::ObjectRef {
public:
TVM_DLL explicit ScopeIdDef(ffi::Array<Var> def_ids, ffi::Optional<ffi::Array<PrimExpr>> extents,
ScopeBinding scope,
ffi::Optional<ffi::Array<PrimExpr>> preferred_extents =
ffi::Optional<ffi::Array<PrimExpr>>(std::nullopt));
/*! \brief Whether this def has a deferred (unknown) extent. */
bool is_deferred() const { return !get()->extents.has_value(); }
/*! \brief Product of all extent dimensions. PRECONDITION: !is_deferred(). */
PrimExpr fused_extent() const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScopeIdDef, ffi::ObjectRef, ScopeIdDefNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ScopeIdDefNode);
};
class ScopeIdDefVerifier {
public:
using ScopeIdSet = std::unordered_map<ScopeBinding, ScopeIdDef>;
/*!
* \brief Verification mode.
*
* - kRelaxed: tolerate deferred (extent=None) ScopeIdDefs. Used for partial
* programs in the well-formedness check at PrimFunc construction time.
* - kStrict: every original ScopeIdDef must end with a resolved extent
* (either explicit at construction, or inferred via closure). Used at
* LowerTIRx entry where downstream resolve/codegen needs concrete values.
*/
enum class Mode { kRelaxed, kStrict };
/*! \brief Verify the scope id definitions are well formed. */
bool Verify(const ffi::Array<ScopeIdDef>& defs, Mode mode = Mode::kStrict);
/*!
* \brief The resolved scope id set; ``id_set[binding]`` is the best-known
* def for that binding (extents filled in from closure when possible).
*/
ScopeIdSet id_set;
};
/*!
* \brief Static resolver for ScopeIdDef values. Replaces the former
* ScopeIdResolveTable runtime registry with a closed-enum switch.
*/
class ScopeIdResolve {
public:
using LaunchParams = std::unordered_map<ffi::String, IterVar>;
/*! \brief Resolve a ScopeIdDef for a given canonical binding + target. */
TVM_DLL static ffi::Array<PrimExpr> Resolve(ScopeBinding binding,
const ffi::Optional<ffi::Array<PrimExpr>>& extents,
int out_dim, const ffi::String& target_kind,
const LaunchParams& params);
/*! \brief Compute the warp_id_in_cta shuffle expression from threadIdx in launch params */
TVM_DLL static PrimExpr ComputeWarpIdInCta(const LaunchParams& params);
};
/*!
* \brief Strict-weak "a is wider than b" on scope kinds: ``world > kernel >
* cluster > cta > warpgroup > warp > thread``. Only used by axe-layout
* scope-chain validity (the rest of the codebase compares scope identities
* with ==).
*/
inline bool ScopeKindHigher(ScopeKind a, ScopeKind b) {
return static_cast<int>(a) < static_cast<int>(b);
}
/*! \brief String-keyed convenience over ScopeKindHigher. FATALs on bad name. */
TVM_DLL bool ScopeNameHigher(const ffi::String& a, const ffi::String& b);
/******** Definition of Execution Scope ********/
class ExecScopeNode : public ffi::Object {
public:
ffi::Array<ScopeIdDef> scope_id_def;
/*! \brief scope identity; one of the closed ScopeKind values. */
ScopeKind kind = ScopeKind::kKernel;
/*! \brief Human-readable name derived from ``kind`` (for printing / errors). */
ffi::String name() const { return ScopeKindToString(kind); }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExecScopeNode>()
.def_ro("kind", &ExecScopeNode::kind)
.def_ro("scope_id_def", &ExecScopeNode::scope_id_def);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("tirx.ExecScope", ExecScopeNode, ffi::Object);
};
class ExecScope : public ffi::ObjectRef {
public:
/*! \brief Construct from a ScopeKind (canonical). */
TVM_DLL explicit ExecScope(ScopeKind kind, ffi::Array<ScopeIdDef> scope_id_def = {});
/*! \brief Construct from a name string (FATALs on unknown name). */
TVM_DLL explicit ExecScope(const ffi::String& name, ffi::Array<ScopeIdDef> scope_id_def = {})
: ExecScope(StringToScopeKind(name), std::move(scope_id_def)) {}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecScope, ffi::ObjectRef, ExecScopeNode);
};
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_EXEC_SCOPE_H_
include/tvm/tirx/layout.h
+565
View File
@@ -0,0 +1,565 @@
/*
* 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/layout.h
* \brief Definition of layout
*/
#ifndef TVM_TIRX_LAYOUT_H_
#define TVM_TIRX_LAYOUT_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/object.h>
#include <tvm/ir/attr_registry_map.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/var.h>
namespace tvm {
// Forward declaration
template <typename, typename>
class AttrRegistry;
namespace tirx {
template <typename>
class AxisAttrMap;
class Layout;
class TileLayout;
class Iter;
using ffi::Array;
using ffi::Tuple;
// Base class for layout
class LayoutNode : public ffi::Object {
public:
/*! \brief Compatible with shape */
virtual bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const = 0;
/*! \brief Verify if the layout is well-formed */
virtual bool VerifyWellFormed() const = 0;
/*! \brief Get the size of the layout (of some axis) */
virtual PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const = 0;
/*! \brief Get the span of the layout (of some axis) */
virtual PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const = 0;
/*! \brief Apply layout on the input coordinate and get the mapped output */
virtual ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const = 0;
virtual ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const = 0;
ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
const ffi::Array<PrimExpr>& shape) const;
/*! \brief Turn the layout to canonical form */
virtual Layout Canonicalize() const = 0;
/*! \brief Tile the current layout with a given layout */
virtual Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
const ffi::Array<PrimExpr>& inner_shape) const = 0;
/*! \brief Slice the layout with a given shape and region */
virtual ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape,
const Region& region) const = 0;
/*! \brief Direct-sum on the tiling domain (unscaled composition)
* Given left layout A (grouped by left_shape) and this layout B (grouped by right_shape),
* construct the interleaved-domain direct sum A + B without span scaling.
*/
virtual Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
const ffi::Array<PrimExpr>& right_shape) const = 0;
/*! \brief Check if the layout is the inner layout of a tiled layout
* \param tile_layout The tiled layout to check
* \param tiled_shape The shape of the tiled layout
* \param inner_shape The shape of the inner layout
* \return The outer layout if this layout is the inner layout of tile_layout, std::nullopt
* otherwise
*/
virtual ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& inner_shape) const = 0;
/*! \brief Check if the layout is the outer layout of a tiled layout
* \param tile_layout The tiled layout to check
* \param tiled_shape The shape of the tiled layout
* \param outer_shape The shape of the outer layout
* \return The inner layout if this layout is the outer layout of tile_layout, std::nullopt
* otherwise
*/
virtual ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& outer_shape) const = 0;
/*! \brief Check if this layout is the right addend B in a direct-sum A + B over the
* interleaved domain S_A \otimes S_B. If so, return the left layout A.
* \param sum_layout The resulting direct-sum layout
* \param interleaved_shape The interleaved domain S_A \otimes S_B, i.e., [A0, B0, A1, B1, ...]
* \param right_shape The shape that groups this (right) layout
*/
virtual ffi::Optional<TileLayout> IsDirectSumRight(
const Layout& sum_layout, const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& right_shape) const = 0;
/*! \brief Check if this layout is the left addend A in a direct-sum A + B over the
* interleaved domain S_A \otimes S_B. If so, return the right layout B.
* \param sum_layout The resulting direct-sum layout
* \param interleaved_shape The interleaved domain S_A \otimes S_B, i.e., [A0, B0, A1, B1, ...]
* \param left_shape The shape that groups this (left) layout
*/
virtual ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& left_shape) const = 0;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO("tirx.Layout", LayoutNode, ffi::Object);
};
class Layout : public ffi::ObjectRef {
public:
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Layout, ffi::ObjectRef, LayoutNode);
};
// target, subscope, scope, iter -> fused_iter
using FAxisFuser = ffi::TypedFunction<ffi::Optional<Iter>(Target, ffi::String, ffi::String, Iter)>;
// target, scope, iter -> (outer_iter, inner_iter)
// Note(@bohao): use ffi::Array<Iter, void> to avoid incomplete type error (SFINAE)
using FAxisSplitter = ffi::TypedFunction<ffi::Array<Iter, void>(Target, ffi::String, Iter)>;
// Axis
class AxisNode : public ffi::Object {
public:
ffi::String name;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AxisNode>().def_ro("name", &AxisNode::name);
}
/*! \brief Check if the axis is a thread axis. */
bool IsThreadAxis() const;
/*! \brief Check if the axis is a memory axis. */
bool IsMemoryAxis() const;
/*! \brief Get the scope of the (thread) axis. */
ffi::Optional<ExecScope> GetScope() const;
/*! \brief Get the subscope of the (thread) axis. */
ffi::Optional<ExecScope> GetSubscope() const;
/*! \brief Get the fuser of the (thread) axis. */
ffi::Optional<FAxisFuser> GetFuser() const;
/*! \brief Get the splitter of the (thread) axis. */
ffi::Optional<FAxisSplitter> GetSplitter() const;
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Axis", AxisNode, ffi::Object);
private:
// Iternals necessary for AttrRegistry
template <typename>
friend class tvm::AttrRegistryMapContainerMap;
template <typename, typename>
friend class tvm::AttrRegistry;
friend class AxisRegEntry;
/*! \brief Program internal unique index of operator. */
uint32_t index_{0};
/*! \brief Return the index stored in attr registry */
uint32_t AttrRegistryIndex() const { return index_; }
/*! \brief Return the name stored in attr registry */
ffi::String AttrRegistryName() const { return name; }
};
class Axis : public ffi::ObjectRef {
public:
Axis() = default;
/*! \brief Get the axis object by name. */
TVM_DLL static Axis Get(const ffi::String& name);
/*! \brief Get the attribute map for the axis. */
template <typename ValueType>
inline static AxisAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
explicit Axis(ffi::ObjectPtr<AxisNode> data) : ObjectRef(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Axis, ffi::ObjectRef, AxisNode);
private:
// Internals necessary for AttrRegistry
template <typename, typename>
friend class tvm::AttrRegistry;
friend class AxisRegEntry;
};
// AxisRegistry
class AxisRegEntry {
public:
/*! \brief List all axis names. */
TVM_DLL static ffi::Array<ffi::String> ListAxisNames();
/*! \brief Register or get the axis entry by name. */
TVM_DLL static AxisRegEntry& RegisterOrGet(const ffi::String& name);
/*! \brief Set the attribute for the axis. */
template <typename ValueType>
inline AxisRegEntry& set_attr(const ffi::String& attr_name, const ValueType& value,
int plevel = 10);
/*! \brief Set the scope of the axis. */
inline AxisRegEntry& set_scope(const ffi::String& scope_name, int plevel = 10);
/*! \brief Set the subscope of the axis. */
inline AxisRegEntry& set_subscope(const ffi::String& subscope_name, int plevel = 10);
/*! \brief Set the fuser of the axis. */
inline AxisRegEntry& set_fuser(const FAxisFuser& fuser);
/*! \brief Set the splitter of the axis. */
inline AxisRegEntry& set_splitter(const FAxisSplitter& splitter);
private:
// return internal pointer to op.
inline AxisNode* get();
TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
// Internals necessary for AttrRegistry
Axis axis_;
ffi::String name;
explicit AxisRegEntry(uint32_t index);
template <typename, typename>
friend class tvm::AttrRegistry;
friend class Axis;
};
using AxisRegistry = AttrRegistry<AxisRegEntry, Axis>;
// AxisAttrffi::Map
template <typename ValueType>
class AxisAttrMap : public AttrRegistryMap<Axis, ValueType> {
public:
using TParent = AttrRegistryMap<Axis, ValueType>;
using TParent::count;
using TParent::get;
using TParent::operator[];
private:
friend class Axis;
explicit AxisAttrMap(const AttrRegistryMapContainerMap<Axis>& map) : TParent(map) {}
};
// Helper macro for token concatenation
#ifndef TVM_STR_CONCAT
#define TVM_STR_CONCAT_(__x, __y) __x##__y
#define TVM_STR_CONCAT(__x, __y) TVM_STR_CONCAT_(__x, __y)
#endif
// Define a macro to register the axis entry.
#define TVM_AXIS_REGISTER_VAR_DEF [[maybe_unused]] static ::tvm::tirx::AxisRegEntry& __make_##Axis
#define TVM_REGISTER_AXIS(AxisName) \
TVM_STR_CONCAT(TVM_AXIS_REGISTER_VAR_DEF, __COUNTER__) = \
::tvm::tirx::AxisRegEntry::RegisterOrGet(AxisName)
class IterNode : public ffi::Object {
public:
PrimExpr extent;
PrimExpr stride;
Axis axis;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<IterNode>()
.def_ro("extent", &IterNode::extent)
.def_ro("stride", &IterNode::stride)
.def_ro("axis", &IterNode::axis);
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Iter", IterNode, ffi::Object);
};
class Iter : public ffi::ObjectRef {
public:
TVM_DLL explicit Iter(PrimExpr extent, PrimExpr stride, Axis axis);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Iter, ffi::ObjectRef, IterNode);
};
class TileLayoutNode : public LayoutNode {
public:
ffi::Array<Iter> shard;
ffi::Array<Iter> replica;
ffi::Map<Axis, PrimExpr> offset;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TileLayoutNode>()
.def_ro("shard", &TileLayoutNode::shard)
.def_ro("replica", &TileLayoutNode::replica)
.def_ro("offset", &TileLayoutNode::offset);
}
/*! \brief Check if the layout is compatible with the shape */
bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const final;
/*! \brief Verify if the layout is well-formed */
bool VerifyWellFormed() const final;
/*! \brief Get the size of the layout (of some axis) */
PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Get the span of the layout (of some axis) */
PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Apply the input coordinate and get the mapped output */
ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const final;
ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
/*! \brief Turn the layout to canonical form */
Layout Canonicalize() const final;
/*! \brief Tile the layout with an outer layout */
Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
/*! \brief Check if the layout is the inner layout of a tiled layout */
ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
/*! \brief Check if the layout is the outer layout of a tiled layout */
ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& outer_shape) const final;
ffi::Optional<TileLayout> IsDirectSumRight(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& left_shape) const final;
/*! \brief Get the shape of the shard */
ffi::Array<PrimExpr> GetShardShape() const;
/*! \brief Slice the layout with a given shape and region */
ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape, const Region& region) const final;
/*! \brief Is the layout trivial (pure memory, identical mapping) */
bool IsTrivial() const;
/*! \brief Check if the layout is trainium layout */
bool IsTrainium() const;
/*! \brief Has Memory Axis */
bool HasMemoryAxis() const;
/*! \brief Has Thread Axis */
bool HasThreadAxis() const;
/*! \brief Get the scope pair of the layout */
ffi::Optional<Tuple<ExecScope, ExecScope>> GetScope() const;
/*! \brief Get the default layout for the shape */
static TileLayout DefaultLayout(ffi::Array<PrimExpr> shape);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TileLayout", TileLayoutNode, LayoutNode);
};
class TileLayout : public Layout {
public:
TVM_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter> replica,
ffi::Map<Axis, PrimExpr> offset);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TileLayout, Layout, TileLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(TileLayoutNode);
};
// SwizzleLayout
class SwizzleLayoutNode : public LayoutNode {
public:
int per_element;
int swizzle_len;
int atom_len;
bool swizzle_inner;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<SwizzleLayoutNode>()
.def_ro("per_element", &SwizzleLayoutNode::per_element)
.def_ro("swizzle_len", &SwizzleLayoutNode::swizzle_len)
.def_ro("atom_len", &SwizzleLayoutNode::atom_len)
.def_ro("swizzle_inner", &SwizzleLayoutNode::swizzle_inner)
.def_ro("inner_mask", &SwizzleLayoutNode::inner_mask)
.def_ro("outer_mask", &SwizzleLayoutNode::outer_mask);
}
/*! \brief Check if the layout is compatible with the shape */
bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const final;
/*! \brief Verify if the layout is well-formed */
bool VerifyWellFormed() const final;
/*! \brief Get the size of the layout */
PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Get the span of the layout */
PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Apply the input coordinate and get the mapped output */
ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const final;
ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
/*! \brief Turn the layout to canonical form */
Layout Canonicalize() const final;
/*! \brief Tile the layout with an outer layout */
Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
/*! \brief Check if the layout is the inner layout of a tiled layout */
ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
/*! \brief Check if the layout is the outer layout of a tiled layout */
ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& outer_shape) const final;
ffi::Optional<TileLayout> IsDirectSumRight(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& left_shape) const final;
/*! \brief Slice the layout with a given shape and region */
ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape, const Region& region) const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.SwizzleLayout", SwizzleLayoutNode, LayoutNode);
private:
friend class SwizzleLayout;
int inner_mask;
int outer_mask;
};
class SwizzleLayout : public Layout {
public:
TVM_DLL explicit SwizzleLayout(int per_element, int swizzle_len, int atom_len,
bool swizzle_inner);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SwizzleLayout, Layout, SwizzleLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(SwizzleLayoutNode);
};
// ComposeLayout
class ComposeLayoutNode : public LayoutNode {
public:
SwizzleLayout swizzle;
TileLayout tile_layout;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ComposeLayoutNode>()
.def_ro("swizzle", &ComposeLayoutNode::swizzle)
.def_ro("tile_layout", &ComposeLayoutNode::tile_layout);
}
/*! \brief Check if the layout is compatible with the shape */
bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const final;
/*! \brief Verify if the layout is well-formed */
bool VerifyWellFormed() const final;
/*! \brief Get the size (of some axis) of the layout */
PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Get the span (of some axis) of the layout */
PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
/*! \brief Apply the input coordinate and get the mapped output */
ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const final;
ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
/*! \brief Turn the layout to canonical form */
Layout Canonicalize() const final;
/*! \brief Tile the layout with an outer layout */
Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
/*! \brief Check if the layout is the inner layout of a tiled layout */
ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& inner_shape) const final;
/*! \brief Check if the layout is the outer layout of a tiled layout */
ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
const ffi::Array<PrimExpr>& tiled_shape,
const ffi::Array<PrimExpr>& outer_shape) const final;
ffi::Optional<TileLayout> IsDirectSumRight(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& right_shape) const final;
ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
const ffi::Array<PrimExpr>& interleaved_shape,
const ffi::Array<PrimExpr>& left_shape) const final;
/*! \brief Slice the layout with a given shape and region */
ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape, const Region& region) const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ComposeLayout", ComposeLayoutNode, LayoutNode);
};
class ComposeLayout : public Layout {
public:
TVM_DLL explicit ComposeLayout(SwizzleLayout layout_A, TileLayout layout_B);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ComposeLayout, Layout, ComposeLayoutNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ComposeLayoutNode);
};
constexpr int kPSUMMaxElemPerBank = 512;
constexpr int kPSUMBankNum = 8;
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_LAYOUT_H_
include/tvm/tirx/op.h
+6 -2
View File
@@ -25,8 +25,8 @@
* when the type is int32 or int64 for simplifying the index expressions.
*/
// Acknowledgement: Most operator APIs originate from Halide.
#ifndef TVM_TIR_OP_H_
#define TVM_TIR_OP_H_
#ifndef TVM_TIRX_OP_H_
#define TVM_TIRX_OP_H_
#include <tvm/ir/expr.h>
#include <tvm/ir/op.h>
@@ -34,6 +34,8 @@
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/target_builtin/cuda.h>
#include <tvm/tirx/target_builtin/trn.h>
#include <algorithm>
#include <limits>
@@ -44,6 +46,8 @@ namespace tvm {
#define TVM_TIR_REGISTER_OP(OpName) \
TVM_REGISTER_OP("tirx." OpName).set_attr<TScriptPrinterName>("TScriptPrinterName", OpName)
#define TVM_TIRX_REGISTER_OP(OpName) TVM_TIR_REGISTER_OP(OpName)
// Most common operators can be overloaded by argument type(PrimExpr).
// So we put them under the root namespace.
//
include/tvm/tirx/predicate.h
+66
View File
@@ -0,0 +1,66 @@
/*
* 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 <tvm/arith/analyzer.h>
#include <tvm/ffi/object.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/var.h>
namespace tvm {
namespace tirx {
class PredicateNode : public ffi::Object {
public:
/*! \brief The variables in the predicate */
Array<Var> vars;
/*! \brief The predicate */
PrimExpr pred;
/*! \brief Replace the variables in the predicate with the given indices */
PrimExpr Apply(const Array<PrimExpr>& indices) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PredicateNode>()
.def_ro("vars", &PredicateNode::vars, refl::AttachFieldFlag::SEqHashDef())
.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<Var> vars, PrimExpr pred);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Predicate, ffi::ObjectRef, PredicateNode);
};
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_PREDICATE_H_
include/tvm/tirx/script/builder/frame.h
+184 -3
View File
@@ -16,11 +16,12 @@
* specific language governing permissions and limitations
* under the License.
*/
#ifndef TVM_TIRX_SCRIPT_BUILDER_FRAME_H_
#define TVM_TIRX_SCRIPT_BUILDER_FRAME_H_
#ifndef TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_FRAME_H_
#include <tvm/script/ir_builder/base.h>
#include <tvm/script/ir_builder/ir/frame.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/stmt.h>
#include <utility>
@@ -85,6 +86,13 @@ class PrimFuncFrameNode : public TIRFrameNode {
/*! \brief The buffer allocated in root block. */
ffi::Array<tvm::tirx::Buffer> root_alloc_buffers;
// TIR utils
/*! \brief Whether this PrimFunc uses s_tir semantics (root SBlock wrap,
* parser layout default = None). Default (false) = tirx semantics. */
bool s_tir;
/*! \brief Whether it is a persistent kernel. */
bool persistent;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<PrimFuncFrameNode>()
@@ -95,7 +103,9 @@ class PrimFuncFrameNode : public TIRFrameNode {
.def_ro("buffer_map", &PrimFuncFrameNode::buffer_map)
.def_ro("attrs", &PrimFuncFrameNode::attrs)
.def_ro("env_threads", &PrimFuncFrameNode::env_threads)
.def_ro("root_alloc_buffers", &PrimFuncFrameNode::root_alloc_buffers);
.def_ro("root_alloc_buffers", &PrimFuncFrameNode::root_alloc_buffers)
.def_ro("s_tir", &PrimFuncFrameNode::s_tir)
.def_ro("persistent", &PrimFuncFrameNode::persistent);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.PrimFuncFrame", PrimFuncFrameNode,
TIRFrameNode);
@@ -237,6 +247,52 @@ class BlockInitFrame : public TIRFrame {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BlockInitFrame, TIRFrame, BlockInitFrameNode);
};
/*!
* \brief A frame that represents an execution scope (e.g. cta, warp, thread).
*
* When exiting this frame, it produces an ExecScopeStmt wrapping the body.
* This is the new IR pattern, replacing the old pattern of storing exec_scope on SBlock.
*
* \sa ExecScopeFrame
*/
class ExecScopeFrameNode : public TIRFrameNode {
public:
/*! \brief The execution scope (always plain kind; no slice). */
ffi::Optional<tvm::tirx::ExecScope> exec_scope;
/*! \brief Optional surface-syntax guards for ``with Tx.scope(cond)``. */
ffi::Array<PrimExpr> guards;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExecScopeFrameNode>()
.def_ro("exec_scope", &ExecScopeFrameNode::exec_scope)
.def_ro("guards", &ExecScopeFrameNode::guards);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ExecScopeFrame", ExecScopeFrameNode,
TIRFrameNode);
public:
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
void ExitWithScope() final;
};
/*!
* \brief Managed reference to ExecScopeFrameNode.
*
* \sa ExecScopeFrameNode
*/
class ExecScopeFrame : public TIRFrame {
public:
explicit ExecScopeFrame(ffi::ObjectPtr<ExecScopeFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ExecScopeFrame, TIRFrame, ExecScopeFrameNode);
};
/*!
* \brief A frame that represents the for loop.
*
@@ -597,6 +653,131 @@ class ElseFrame : public TIRFrame {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ElseFrame, TIRFrame, ElseFrameNode);
};
class DeclBufferFrameNode : public TIRFrameNode {
public:
/*! \brief The declared buffer. */
tvm::tirx::Buffer buffer;
/*! \brief The buffer allocated or not. */
bool allocated;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DeclBufferFrameNode>()
.def_ro("buffer", &DeclBufferFrameNode::buffer)
.def_ro("allocated", &DeclBufferFrameNode::allocated);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.DeclBufferFrame", DeclBufferFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class DeclBufferFrame : public TIRFrame {
public:
explicit DeclBufferFrame(ffi::ObjectPtr<DeclBufferFrameNode> data) : TIRFrame(data) {
TVM_FFI_ICHECK(data != nullptr);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DeclBufferFrame, TIRFrame, DeclBufferFrameNode);
};
class ComposeOpFrameNode : public TIRFrameNode {
public:
/*! \brief The workspace of the compose op. */
ffi::Map<ffi::String, tvm::tirx::Buffer> workspace;
/*! \brief The config of the compose op. */
ffi::Map<ffi::String, ffi::Any> config;
/*! \brief The optional dispatch variant name of the compose op. */
ffi::Optional<ffi::String> dispatch{std::nullopt};
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ComposeOpFrameNode>()
.def_ro("workspace", &ComposeOpFrameNode::workspace)
.def_ro("config", &ComposeOpFrameNode::config)
.def_ro("dispatch", &ComposeOpFrameNode::dispatch);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.ComposeOpFrame", ComposeOpFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class ComposeOpFrame : public TIRFrame {
public:
explicit ComposeOpFrame(ffi::ObjectPtr<ComposeOpFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComposeOpFrame, TIRFrame, ComposeOpFrameNode);
};
class AllocBufferFrameNode : public TIRFrameNode {
public:
/*! \brief The allocated buffer. */
tvm::tirx::Buffer buffer;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AllocBufferFrameNode>().def_ro("buffer", &AllocBufferFrameNode::buffer);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.AllocBufferFrame", AllocBufferFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
class AllocBufferFrame : public TIRFrame {
public:
explicit AllocBufferFrame(ffi::ObjectPtr<AllocBufferFrameNode> data)
: TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(AllocBufferFrame, TIRFrame, AllocBufferFrameNode);
};
/*!
* \brief A frame that represents a hint directive for the sketch language.
*
* \sa HintFrame
*/
class HintFrameNode : public TIRFrameNode {
public:
/*! \brief The free-form hint message string. */
ffi::String message;
/*! \brief Optional structured key-value attributes. */
ffi::Map<ffi::String, ffi::Any> attrs;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<HintFrameNode>()
.def_ro("message", &HintFrameNode::message)
.def_ro("attrs", &HintFrameNode::attrs);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tirx.HintFrame", HintFrameNode,
TIRFrameNode);
public:
void ExitWithScope() final;
};
/*!
* \brief Managed reference to HintFrameNode.
*
* \sa HintFrameNode
*/
class HintFrame : public TIRFrame {
public:
explicit HintFrame(ffi::ObjectPtr<HintFrameNode> data) : TIRFrame(ffi::UnsafeInit{}) {
TVM_FFI_ICHECK(data != nullptr);
data_ = std::move(data);
}
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(HintFrame, TIRFrame, HintFrameNode);
};
} // namespace tirx
} // namespace ir_builder
} // namespace script
include/tvm/tirx/script/builder/ir.h
+139 -66
View File
@@ -16,19 +16,30 @@
* specific language governing permissions and limitations
* under the License.
*/
#ifndef TVM_TIRX_SCRIPT_BUILDER_IR_H_
#define TVM_TIRX_SCRIPT_BUILDER_IR_H_
#ifndef TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/runtime/tensor.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/script/builder/frame.h>
#include <tvm/tirx/tirx_stmt.h>
namespace tvm {
namespace script {
namespace ir_builder {
namespace tirx {
using tvm::ffi::Tuple;
using tvm::ffi::Variant;
using tvm::runtime::Tensor;
using tvm::tirx::Buffer;
using tvm::tirx::ExecScope;
using tvm::tirx::Layout;
using tvm::tirx::Var;
/*!
@@ -50,13 +61,15 @@ Buffer BufferDecl(ffi::Array<PrimExpr> shape, DataType dtype, ffi::String buffer
ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
int offset_factor, ffi::String buffer_type,
ffi::Optional<ffi::Array<IntImm>> axis_separators);
ffi::Optional<ffi::Array<IntImm>> axis_separators,
ffi::Optional<Layout> layout = std::nullopt,
ffi::Array<PrimExpr> allocated_addr = {});
/*!
* \brief The primitive function statement.
* \return The PrimFuncFrame.
*/
PrimFuncFrame PrimFunc(bool is_private);
PrimFuncFrame PrimFunc(bool is_private, bool s_tir = false, bool persistent = false);
/*!
* \brief The PrimFunc variable arguments adding function.
@@ -113,7 +126,8 @@ Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
ffi::String buffer_type = "default",
ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt);
ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
ffi::Optional<Layout> layout = std::nullopt);
/*!
* \brief The block declaration statement.
@@ -121,7 +135,34 @@ Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
* \param no_realize The flag whether to construct SBlockRealize or SBlock.
* \return The SBlockFrame.
*/
SBlockFrame Block(ffi::String name, bool no_realize = false);
SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_scope = "");
void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call);
/*!
* \brief Create an ExecScopeFrame for execution scope contexts.
* \param exec_scope_name The name of the execution scope (e.g. "cta", "warp").
* \return The ExecScopeFrame.
*/
ExecScopeFrame ExecScopeBlock(ffi::String exec_scope_name,
ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Kernel(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Cluster(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame WarpGroup(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame CTA(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Warp(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Thread(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ffi::Array<tvm::tirx::Var> KernelId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> CtaId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> CtaIdInPair();
ffi::Array<tvm::tirx::Var> WarpId(ffi::Array<PrimExpr> extents, ffi::String parent);
ffi::Array<tvm::tirx::Var> ThreadId(ffi::Array<PrimExpr> extents, ffi::String parent);
/*!
* \brief The block initialization statement.
@@ -165,13 +206,19 @@ void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \return The allocated buffer.
* \param layout The layout of the buffer.
* \param allocated_addr The allocated address of the buffer. Might be multi-dimensional.
* \return The allocated buffer or the AllocBufferFrame if the function is called under
* T.prim_func(tirx=True).
*/
Buffer SBlockAllocBuffer(ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
ffi::Optional<Var> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "",
int align = -1, int offset_factor = 0, ffi::String buffer_type = "default",
ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt);
ffi::Variant<Buffer, AllocBufferFrame> SBlockAllocBuffer(
ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
ffi::Optional<Var> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1,
int offset_factor = 0, ffi::String buffer_type = "default",
ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
ffi::Optional<Layout> layout = std::nullopt, ffi::Array<PrimExpr> allocated_addr = {});
namespace axis {
/*!
@@ -281,7 +328,7 @@ ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread,
* \param extents The extents of the iteration.
* \return The ForFrame.
*/
ForFrame Grid(ffi::Array<PrimExpr> extents);
ForFrame Grid(ffi::Array<Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> extents);
/*!
* \brief The assertion statement.
@@ -324,6 +371,16 @@ AttrFrame Attr(ffi::Any node, ffi::String attr_key, PrimExpr value);
*/
WhileFrame While(PrimExpr condition);
/*!
* \brief Create a break statement.
*/
void Break();
/*!
* \brief Create a continue statement.
*/
void Continue();
/*!
* \brief Create an if statement.
* \param condition The condition of if statement.
@@ -356,13 +413,16 @@ ElseFrame Else();
* \param offset_factor The factor of elem_offset field.
* \param buffer_type The buffer type.
* \param axis_separators The separators between input axes when generating flattened output axes.
* \return The declared buffer.
* \param layout The layout of the buffer.
* \return The declaration frame.
*/
Buffer DeclBuffer(ffi::Array<PrimExpr> shape, DataType dtype, ffi::String buffer_name,
ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
int offset_factor, ffi::String buffer_type,
ffi::Optional<ffi::Array<IntImm>> axis_separators);
DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, DataType dtype, ffi::String buffer_name,
ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
int align, int offset_factor, ffi::String buffer_type,
ffi::Optional<ffi::Array<IntImm>> axis_separators,
ffi::Optional<Layout> layout = std::nullopt,
ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
/*!
* \brief Statement-level buffer allocation (creates an AllocBuffer IR node).
@@ -392,6 +452,17 @@ LaunchThreadFrame LaunchThread(Var var, PrimExpr extent);
*/
LaunchThreadFrame LaunchThread(ffi::String thread_tag, PrimExpr extent);
/*!
* \brief Compose TIRx op.
* \param workspace The workspace of the compose op.
* \param config The config of the compose op.
* \param dispatch The optional dispatch variant name.
* \return The result ComposeOpFrame.
*/
ComposeOpFrame ComposeOp(ffi::Map<ffi::String, Buffer> workspace,
ffi::Map<ffi::String, ffi::Any> config,
ffi::Optional<ffi::String> dispatch = std::nullopt);
/*!
* \brief Bind a var to thread env.
* \param thread_tag The thread type tag.
@@ -447,9 +518,9 @@ inline Var Handle(runtime::DataType dtype = runtime::DataType::Void(),
: tvm::tirx::Var("", type_annotation);
}
inline Var TensormapHandle() { return tvm::tirx::Var("", PointerType(TensorMapType())); }
inline Var TensorMap() { return tvm::tirx::Var("", PointerType(TensorMapType())); }
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType) \
inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt, \
bool is_size_var = false) { \
DataType dtype = DType; \
@@ -458,61 +529,63 @@ inline Var TensormapHandle() { return tvm::tirx::Var("", PointerType(TensorMapTy
: (is_size_var ? tvm::tirx::SizeVar("", dtype) : tvm::tirx::Var("", dtype)); \
}
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##8, FDType(8)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##16, FDType(16)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##32, FDType(32)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##64, FDType(64));
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##8, FDType(8)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##16, FDType(16)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##32, FDType(32)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##64, FDType(64));
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(BFloat, DataType::BFloat);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Float, DataType::Float);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, DataType::UInt);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, DataType::Int);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(BFloat, DataType::BFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Float, DataType::Float);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, DataType::UInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, DataType::Int);
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, FDType(Size, 32)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, FDType(Size, 64));
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x2, FDType(Size, 2)) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, FDType(Size, 32)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, FDType(Size, 64));
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType) \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, FDType, 8); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, FDType, 16); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, FDType, 32); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, FDType, 64);
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, FDType, 8); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, FDType, 16); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, FDType, 32); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, FDType, 64);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(BFloat, DataType::BFloat);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, DataType::Float);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, DataType::UInt);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, DataType::Int);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(BFloat, DataType::BFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, DataType::Float);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, DataType::UInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, DataType::Int);
#define TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32)); \
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x2, FDType(2)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32)); \
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E3M4, DataType::Float8E3M4);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3, DataType::Float8E4M3);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3B11FNUZ, DataType::Float8E4M3B11FNUZ);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FN, DataType::Float8E4M3FN);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FNUZ, DataType::Float8E4M3FNUZ);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2, DataType::Float8E5M2);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2FNUZ, DataType::Float8E5M2FNUZ);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E8M0FNU, DataType::Float8E8M0FNU);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E3M4, DataType::Float8E3M4);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3, DataType::Float8E4M3);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3B11FNUZ, DataType::Float8E4M3B11FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FN, DataType::Float8E4M3FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FNUZ, DataType::Float8E4M3FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2, DataType::Float8E5M2);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2FNUZ, DataType::Float8E5M2FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E8M0FNU, DataType::Float8E8M0FNU);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E2M3FN, DataType::Float6E2M3FN);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E3M2FN, DataType::Float6E3M2FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E2M3FN, DataType::Float6E2M3FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E3M2FN, DataType::Float6E3M2FN);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float4E2M1FN, DataType::Float4E2M1FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float4E2M1FN, DataType::Float4E2M1FN);
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(Boolean, DataType::Bool());
TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST(Void, DataType::Void());
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Boolean, DataType::Bool());
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Void, DataType::Void());
#undef TVM_TIR_IR_BUILDER_DEF_DTYPE_CAST
#undef TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST
} // namespace tirx
} // namespace ir_builder
include/tvm/tirx/stmt.h
+338 -5
View File
@@ -21,13 +21,14 @@
* \brief TIR statements.
*/
// Acknowledgement: Many low-level stmts originate from Halide.
#ifndef TVM_TIR_STMT_H_
#define TVM_TIR_STMT_H_
#ifndef TVM_TIRX_STMT_H_
#define TVM_TIRX_STMT_H_
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/script/printer/config.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/layout.h>
#include <optional>
#include <string>
@@ -458,8 +459,8 @@ class SeqStmt : public Stmt {
template <typename T>
void operator()(size_t i, const T& stmt_or_seq) const {
if constexpr (std::is_base_of_v<ffi::ObjectRef, T>) {
// Early bail-out, applicable to any ffi::ObjectRef
if constexpr (std::is_base_of_v<ObjectRef, T>) {
// Early bail-out, applicable to any ObjectRef
if (!stmt_or_seq.defined()) {
return;
}
@@ -687,6 +688,56 @@ class While : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(WhileNode);
};
/*!
* \brief A Break in control flow.
*/
class BreakNode : public StmtNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<BreakNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Break", BreakNode, StmtNode);
};
/*!
* \brief Managed reference to BreakNode.
* \sa BreakNode
*/
class Break : public Stmt {
public:
TVM_DLL explicit Break(Span span);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Break, Stmt, BreakNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(BreakNode);
};
/*!
* \brief A Continue in control flow.
*/
class ContinueNode : public StmtNode {
public:
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ContinueNode>();
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Continue", ContinueNode, StmtNode);
};
/*!
* \brief Managed reference to ContinueNode.
* \sa ContinueNode
*/
class Continue : public Stmt {
public:
TVM_DLL explicit Continue(Span span);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Continue, Stmt, ContinueNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode);
};
/*!
* \brief Representing the region of multi-dimensional buffer access.
*/
@@ -856,6 +907,10 @@ class SBlock : public Stmt {
ffi::Map<ffi::String, ffi::Any> annotations = ffi::Map<ffi::String, ffi::Any>(),
Span span = Span());
TVM_DLL explicit SBlock(ffi::String name_hint, Stmt body,
ffi::Array<Buffer> alloc_buffers = ffi::Array<Buffer>(),
Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(SBlock, Stmt, SBlockNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockNode);
};
@@ -898,6 +953,47 @@ class SBlockRealize : public Stmt {
TVM_DEFINE_OBJECT_REF_COW_METHOD(SBlockRealizeNode);
};
/*!
* \brief A statement that annotates the execution scope for its body.
*
* ExecScopeStmt represents a hardware execution scope (e.g. cta, warp, thread)
* that wraps a body statement. This decouples the execution scope concept from
* SBlock, making the IR structure cleaner.
*
* Example:
* \code
* with T.cta():
* ...
* \endcode
*/
class ExecScopeStmtNode : public StmtNode {
public:
/*! \brief The execution scope. */
ExecScope exec_scope;
/*! \brief The body statement under this execution scope. */
Stmt body;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ExecScopeStmtNode>()
.def_ro("exec_scope", &ExecScopeStmtNode::exec_scope)
.def_ro("body", &ExecScopeStmtNode::body);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ExecScopeStmt", ExecScopeStmtNode, StmtNode);
};
/*!
* \brief Managed reference to ExecScopeStmtNode.
* \sa ExecScopeStmtNode
*/
class ExecScopeStmt : public Stmt {
public:
TVM_DLL ExecScopeStmt(ExecScope exec_scope, Stmt body, Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecScopeStmt, Stmt, ExecScopeStmtNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(ExecScopeStmtNode);
};
/*! \brief namespace of possible attributes in AttrStmt.attr_key */
namespace attr {
/*! \brief Mark stores/loads with their bounds. */
@@ -937,6 +1033,243 @@ constexpr const char* storage_alignment = "storage_alignment";
constexpr const char* thread_extent = "thread_extent";
/*! \brief Annotation key on AllocBuffer marking the allocation as volatile. */
constexpr const char* kVolatile = "tirx.volatile";
/*!
* \brief Marks the layout transforms to be used for a tensor.
*
* Only applies to a DataProducer, as it should be made part of the
* PrimFunc attributes for TIR.
*/
constexpr const char* layout_transforms = "layout_transforms";
/*!
* \brief Marks the physical axis separators
*
* Only applies to a DataProducer, as it should be made part of the
* Buffer definition in a PrimFunc. See `BufferNode::axis_separators`
* for more details.
*/
constexpr const char* axis_separators = "axis_separators";
/*!
* \brief Marks production of double buffer data
*/
constexpr const char* double_buffer_scope = "double_buffer_scope";
/*!
* \brief Marks region used by double buffer write
*/
constexpr const char* double_buffer_write = "double_buffer_write";
/*! \brief Mark of scan update scope */
constexpr const char* scan_update_scope = "scan_update_scope";
/*! \brief Mark of scan init scope */
constexpr const char* scan_init_scope = "scan_init_scope";
/*!
* \brief Mark alignment of buffer dimension
* stmt.node is Tensor
* stmt.value is tvm_tuple(dim, align, offset)
* This gives hint to require stride of dim to be k * align + offset.
*/
constexpr const char* buffer_dim_align = "buffer_dim_align";
/*! \brief Mark buffer initial addr alignment in bytes */
constexpr const char* buffer_data_alignment = "buffer_data_alignment";
/*! \brief Mark buffer allocated addr in bytes */
constexpr const char* buffer_allocated_addr = "buffer_allocated_addr";
/*!
* \brief Bind the buffer specification to the region of the op
* When this scope occurs, the stmt.node is a ffi::Array<NodeRef> = [buffer, tensor]
* stmt.value is a tvm_tuple(min0, extent0, min1, extent1, ...).
* The scope represents that we need to bind the storage region of tensor to buffer.
* This will affect replacement of some variables inside the scope that
* corresponds to field of buffer to be the actual expressions of tensor during
* storage flattening phase.
*/
constexpr const char* buffer_bind_scope = "buffer_bind_scope";
// Pipeline related attributes
/*! \brief channel read scope */
constexpr const char* channel_read_scope = "channel_read_scope";
/*! \brief Advance step of channel after end of scope */
constexpr const char* channel_read_advance = "channel_read_advance";
/*! \brief channel write scope */
constexpr const char* channel_write_scope = "channel_write_scope";
/*! \brief Advance step of channel after end of scope */
constexpr const char* channel_write_advance = "channel_write_advance";
/*! \brief pipeline stage scope, implies always execution */
constexpr const char* pipeline_stage_scope = "pipeline_stage_scope";
/*! \brief pipeline execution scope, implies the scope can be pipelined. */
constexpr const char* pipeline_exec_scope = "pipeline_exec_scope";
/*!
* \brief Mark that the attached statement runs asynchronously.
*/
constexpr const char* async_scope = "async_scope";
/*!
* \brief Annotations for invoking and synchronizing asynchronous operations.
* Synchronization is done in terms of "queue": It is an abstract entity associated
* with each asynchronous unit, and it tracks invocations and completions of asynchronous
* operations in the FIFO order.
*
* Similarly to PTX instructions commit_group and wait_group, these annotations express
* synchronization by "counting":
*
* async_commit_queue(i): Group one or more invocations of async operations in the given scope,
* and "commit" (or push) them to the queue i. A group of operations committed together is
* awaited as one chunk. Groups committed to the same queue complete in the FIFO order.
*
* async_wait_queue(i, N): Block until only N most recent committed groups are still in-flight at
* the queue i. N does not have to be a constant, but some backends may require a constant count.
*/
constexpr const char* async_commit_queue_scope = "async_commit_queue_scope";
constexpr const char* async_wait_queue_scope = "async_wait_queue_scope";
constexpr const char* async_wait_inflight_count = "async_wait_inflight_count";
/*!
* \brief Mark that the shape of TensorCore fragment
*/
constexpr const char* fragment_shape = "fragment_shape";
/*!
* \brief Mark that the layout of TensorCore fragment
*/
constexpr const char* fragment_layout = "fragment_layout";
/*!
* \brief Mark that the kernel is hand threaded and doesn't need syncs inserted
*/
constexpr const char* hand_threaded = "hand_threaded";
/*!
* \brief Mark whether the script-completer need to fill in missing access region
* during script parsing.
* \note The result should be a integer mask with range [0, 4).
* if (mask & 1) the read region should be detected,
* if (mask & 2) the write region should be detected.
*/
constexpr const char* script_parsing_detect_access = "tirx.script_parsing_detect_access";
/*!
* \brief Mark that the loop should be partitioned.
*/
constexpr const char* pragma_loop_partition_hint = "pragma_loop_partition_hint";
/*! \brief Mark the stage of a statement in the software pipeline */
constexpr const char* software_pipeline_stage = "software_pipeline_stage";
/*! \brief Mark the order of a statement in the software pipeline */
constexpr const char* software_pipeline_order = "software_pipeline_order";
/*! \brief List stages in the software pipeline that should run asynchronously
* \note All statements in the provided stages are assumed to have asynchronous
* semantics (e.g. CUDA async global to shared memory copy).
*/
constexpr const char* software_pipeline_async_stages = "software_pipeline_async_stages";
/*! \brief Mark the buffers which is const access and can be transformed layout. */
constexpr const char* layout_free_buffers = "layout_free_buffers";
/*! \brief Mark the local stage for the shared memory access should be added. */
constexpr const char* manifest_shared_memory_local_stage =
"tirx.manifest_shared_memory_local_stage";
/*! \brief Mark the tiling structure of blocks that are applied by rule Multi-Level-Tiling */
constexpr const char* meta_schedule_tiling_structure = "meta_schedule.tiling_structure";
/*!
* \brief Mark that the loop should be further skip and bound to environment threads to enable
* cooperative fetching.
*/
constexpr const char* meta_schedule_cooperative_fetch = "meta_schedule.cooperative_fetch";
/*! \brief The allowed range of thread extent in thread bindings */
constexpr const char* meta_schedule_thread_extent_low_inclusive =
"meta_schedule.thread_extent_low_inclusive";
/*! \brief The allowed range of thread extent in thread bindings */
constexpr const char* meta_schedule_thread_extent_high_inclusive =
"meta_schedule.thread_extent_high_inclusive";
/*! \brief Mark the block whose producer needs to be applied by rule Random-Compute-Location */
constexpr const char* meta_schedule_random_compute_producer =
"meta_schedule.random_compute_producer";
/*! \brief Mark auto-parallel setting on the block. */
constexpr const char* meta_schedule_parallel = "meta_schedule.parallel";
/*! \brief Mark auto-vectorize setting on the block. */
constexpr const char* meta_schedule_vectorize = "meta_schedule.vectorize";
/*! \brief Mark auto-unroll setting on the block. */
constexpr const char* meta_schedule_unroll_explicit = "meta_schedule.unroll_explicit";
/*! \brief Mark auto-unroll setting on the block. */
constexpr const char* meta_schedule_unroll_implicit = "meta_schedule.unroll_implicit";
/*! \brief Mark that a block should be further rewritten using tensorization. */
constexpr const char* meta_schedule_auto_tensorize = "meta_schedule.auto_tensorize";
/*! \brief Mark that a block is a preprocessor block for layout rewrite. */
constexpr const char* meta_schedule_layout_rewrite_preproc = "meta_schedule.layout_rewrite_preproc";
/*!
* \brief Mark that the init statement of a block should be further rewritten using tensorization.
*/
constexpr const char* meta_schedule_auto_tensorize_init = "meta_schedule.auto_tensorize_init";
/*!
* \brief Mark that the block need to add predicate for block var bounds during lowering
*/
constexpr const char* require_block_var_bound_predicate = "require_bound_predicate";
/*! \brief Mark that tensor core is enabled in the PrimExpr */
constexpr const char* meta_schedule_tensor_core_enabled = "meta_schedule.tensor_core_enabled";
/*!
* \brief Mark a block as generated by cache_read or cache_write block.
* 0 means cache_read; 1 means cache_write.
* \sa meta_schedule_cache_type_read
* \sa meta_schedule_cache_type_write
*/
constexpr const char* meta_schedule_cache_type = "meta_schedule.cache_type";
/*! \sa meta_schedule_cache_type */
constexpr const int meta_schedule_cache_type_read = 0;
/*! \sa meta_schedule_cache_type */
constexpr const int meta_schedule_cache_type_write = 1;
/*! \brief Mark auto copy for memhammer */
constexpr const char* auto_copy = "auto_copy";
/*! \brief Mark local stage constraint on data copy */
constexpr const char* local_stage = "local_stage";
/*! \brief Mark vectorization length constraint on block */
constexpr const char* vector_bytes = "vector_bytes";
/*!
* \brief Mark that a block is executed by a warp. This implies the extend of threadIdx.x is
* warp size.
*/
constexpr const char* warp_execution = "warp_execution";
/*! \brief Mark that a block is disallowed in auto inline. */
constexpr const char* meta_schedule_inline_rule = "meta_schedule.inline_rule";
/*! \brief Mark that a block has an explicitly specified read region.
* This is used to override the default read region inference in TIR.
*/
constexpr const char* explicit_read_region = "explicit_read_region";
/*! \brief Mark that a block has an explicitly specified write region.
* This is used to override the default write region inference in TIR.
*/
constexpr const char* explicit_write_region = "explicit_write_region";
constexpr const char* tensorized_nki_instruction = "tensorized_nki_instruction";
/*! \brief ,ark a ForNode represent an irregular loop of non-structural control flow edges. */
constexpr const char* irregular_loop_mark = "irregular_loop_mark";
/*!
* \brief Mark the kernel as persistent.
*/
constexpr const char* kPersistentKernel = "tirx.persistent_kernel";
/*!
* \brief Check if attr_key is a pragma key extension
include/tvm/tirx/stmt_functor.h
+20 -3
View File
@@ -23,14 +23,15 @@
* \brief Functors for tirx stmts
* utility functions to call common functors.
*/
#ifndef TVM_TIR_STMT_FUNCTOR_H_
#define TVM_TIR_STMT_FUNCTOR_H_
#ifndef TVM_TIRX_STMT_FUNCTOR_H_
#define TVM_TIRX_STMT_FUNCTOR_H_
#include <tvm/ir/node_functor.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/tirx_stmt.h>
#include <unordered_map>
#include <utility>
@@ -89,6 +90,8 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
virtual R VisitStmt_(const IfThenElseNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const BreakNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const ContinueNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const AllocBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const DeclBufferNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const BufferStoreNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
@@ -97,6 +100,8 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
virtual R VisitStmt_(const EvaluateNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const SBlockNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const SBlockRealizeNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const ExecScopeStmtNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmt_(const tirx::TilePrimitiveCallNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
virtual R VisitStmtDefault_(const ffi::Object* op, Args...) {
TVM_FFI_THROW(InternalError) << "Do not have a default for " << op->GetTypeKey();
TVM_FFI_UNREACHABLE();
@@ -111,6 +116,8 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode);
IR_STMT_FUNCTOR_DISPATCH(ForNode);
IR_STMT_FUNCTOR_DISPATCH(WhileNode);
IR_STMT_FUNCTOR_DISPATCH(BreakNode);
IR_STMT_FUNCTOR_DISPATCH(ContinueNode);
IR_STMT_FUNCTOR_DISPATCH(AllocBufferNode);
IR_STMT_FUNCTOR_DISPATCH(DeclBufferNode);
IR_STMT_FUNCTOR_DISPATCH(AssertStmtNode);
@@ -119,6 +126,8 @@ class StmtFunctor<R(const Stmt& n, Args... args)> {
IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode);
IR_STMT_FUNCTOR_DISPATCH(SBlockNode);
IR_STMT_FUNCTOR_DISPATCH(SBlockRealizeNode);
IR_STMT_FUNCTOR_DISPATCH(ExecScopeStmtNode);
IR_STMT_FUNCTOR_DISPATCH(tirx::TilePrimitiveCallNode);
vtable.Finalize();
return vtable;
}
@@ -164,6 +173,8 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
void VisitStmt_(const IfThenElseNode* op) override;
void VisitStmt_(const ForNode* op) override;
void VisitStmt_(const WhileNode* op) override;
void VisitStmt_(const BreakNode* op) override;
void VisitStmt_(const ContinueNode* op) override;
void VisitStmt_(const AllocBufferNode* op) override;
void VisitStmt_(const DeclBufferNode* op) override;
void VisitStmt_(const BufferStoreNode* op) override;
@@ -172,6 +183,8 @@ class TVM_DLL StmtVisitor : protected StmtFunctor<void(const Stmt&)> {
void VisitStmt_(const EvaluateNode* op) override;
void VisitStmt_(const SBlockNode* op) override;
void VisitStmt_(const SBlockRealizeNode* op) override;
void VisitStmt_(const ExecScopeStmtNode* op) override;
void VisitStmt_(const tirx::TilePrimitiveCallNode* op) override;
};
/*!
@@ -278,6 +291,8 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
Stmt VisitStmt_(const IfThenElseNode* op) override;
Stmt VisitStmt_(const ForNode* op) override;
Stmt VisitStmt_(const WhileNode* op) override;
Stmt VisitStmt_(const BreakNode* op) override;
Stmt VisitStmt_(const ContinueNode* op) override;
Stmt VisitStmt_(const AllocBufferNode* op) override;
Stmt VisitStmt_(const DeclBufferNode* op) override;
Stmt VisitStmt_(const BufferStoreNode* op) override;
@@ -286,6 +301,8 @@ class TVM_DLL StmtMutator : protected StmtFunctor<Stmt(const Stmt&)> {
Stmt VisitStmt_(const EvaluateNode* op) override;
Stmt VisitStmt_(const SBlockNode* op) override;
Stmt VisitStmt_(const SBlockRealizeNode* op) override;
Stmt VisitStmt_(const ExecScopeStmtNode* op) override;
Stmt VisitStmt_(const tirx::TilePrimitiveCallNode* op) override;
/*!
* \brief Alternative advance method for SeqStmtNode.
*
@@ -325,7 +342,7 @@ class TVM_DLL StmtExprVisitor : public ExprVisitor, public StmtVisitor {
/*!
* \brief Mutator that recursively mutates stmts and exprs on them.
*/
class StmtExprMutator : public ExprMutator, public StmtMutator {
class TVM_DLL StmtExprMutator : public ExprMutator, public StmtMutator {
public:
using StmtMutator::operator();
using ExprMutator::operator();
include/tvm/tirx/target_builtin/cuda.h
+745
View File
@@ -0,0 +1,745 @@
/*
* 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/target_builtin/cuda.h
* \brief TIR builtin intrinsics specific to CUDA target.
*/
#ifndef TVM_TIRX_TARGET_BUILTIN_CUDA_H_
#define TVM_TIRX_TARGET_BUILTIN_CUDA_H_
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
namespace tvm {
namespace tirx {
namespace builtin {
// TODO(tvm-team) TensorCore specific intrinsics should be directly registered under
// cuda. namespace and used through op.
/*!
* \brief tvm intrinsic for tensor core load operators.
*
* void tvm_load_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr buffer_ptr, Expr stride,
* StringImm layout) {
* // m, n, k are the shape of wmma fragment.
* // Determine fragment layout(column-major or row major) by layout.
* // fragments must be in 'wmma.matrix_a' or 'wmma.matrix_b' scope.
* nvcuda::wmma::load_matrix_sync(fragment[index], buffer_ptr, stride);
* }
*/
TVM_DLL const Op& tvm_load_matrix_sync();
/*!
* \brief tvm intrinsic for tensor core mma_sync operators.
*
* void tvm_mma_sync(Var fragment_d, Expr index_d,
* Var fragment_a, Expr index_a,
* Var fragment_b, Expr index_b,
* Var fragment_c, Expr index_c) {
* nvcuda::wmma::mma_sync(fragment_d[index_d], fragment_a[index_a],
* fragment_b[index_b], fragment_c[index_c]);
* }
*/
TVM_DLL const Op& tvm_mma_sync();
/*!
* \brief tvm intrinsic for tensor core bmma_sync operators.
*
* void tvm_bmma_sync(Var fragment_d, Expr index_d,
* Var fragment_a, Expr index_a,
* Var fragment_b, Expr index_b,
* Var fragment_c, Expr index_c) {
* nvcuda::wmma::bmma_sync(fragment_d[index_d], fragment_a[index_a],
* fragment_b[index_b], fragment_c[index_c]);
* }
*/
TVM_DLL const Op& tvm_bmma_sync();
/*!
* \brief tvm intrinsic for tensor core fill_fragment operators.
*
* void tvm_fill_fragment(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr value) {
* // m, n, k are the shape of wmma fragment
* // fragments must be in 'wmma.accumulator' scope.
* nvcuda::wmma::fill_fragment(fragment[index], value);
* }
*/
TVM_DLL const Op& tvm_fill_fragment();
/*!
* \brief tvm intrinsic for tensor core store operators.
*
* void tvm_store_matrix_sync(Var fragment, UIntImm m, UIntImm, n, UIntImm k,
* Expr index, Expr buffer_ptr, Expr stride,
* StringImm layout) {
* // m, n, k are the shape of wmma fragment
* // fragments must be in 'wmma.accumulator' scope.
* nvcuda::wmma::store_matrix_sync(fragment[index], buffer_ptr, stride, layout);
* }
*/
TVM_DLL const Op& tvm_store_matrix_sync();
/*!
* \brief tvm intrinsic for ptx tensor core mma instructions.
*
* void ptx_mma(StringImm shape, StringImm A_layout, StringImm B_layout,
* StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
* Var multiplicand_a, Expr a_index,
* Var multiplicand_b, Expr b_index,
* Var accumulator, Expr c_index, bool saturate);
*/
TVM_DLL const Op& ptx_mma();
/*!
* \brief ptx mma / ldmatrix / mma_store / mma_fill variants that take
* ``(ptr_var, offset)`` pairs (not a folded access_ptr Call). Codegen
* emits ``ptr + offset`` C pointer arithmetic; ``lower_warp_memory``
* rewrites the offset's group component to its thread-local index.
*/
TVM_DLL const Op& ptx_mma_legacy();
TVM_DLL const Op& ptx_ldmatrix_legacy();
TVM_DLL const Op& mma_store_legacy();
TVM_DLL const Op& mma_fill_legacy();
/*!
* \brief tvm intrinsic for ptx predicate load with 32-bit data type.
*
*/
TVM_DLL const Op& ptx_ldg32();
/*!
* \brief tvm intrinsic for ptx predicate load with 32-bit data type.
*
*/
TVM_DLL const Op& ptx_ldg32();
/*!
* \brief tvm intrinsic for sparse tensor core ptx instructions.
*
* void ptx_mma_sp(StringImm shape, StringImm A_layout, StringImm B_layout,
* StringImm A_dtype, StringImm B_dtype, StringImm C_dtype,
* Var multiplicand_a, Expr a_index,
* Var multiplicand_b, Expr b_index,
* Var accumulator, Expr c_index,
* Var metadata, Expr meta_index,
* Var sparse_selector, bool saturate);
*/
TVM_DLL const Op& ptx_mma_sp();
/*!
* \brief tvm intrinsic for ptx load matrix from shared memory.
*
* void ptx_ldmatrix(Bool trans, IntImm num, StringImm type,
* Var local_ptr, Expr local_offset,
* Var smem_ptr, Expr smem_offset);
*/
TVM_DLL const Op& ptx_ldmatrix();
/*!
* \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async
*
* void ptx_cp_async(Var shared_ptr,
* Expr shared_offset,
* Var global_ptr,
* Expr global_offset,
* size_t bytes);
*/
TVM_DLL const Op& ptx_cp_async();
/*!
* \brief tvm intrinsics for ptx async copy from global to shared memory using cp.async.bulk
*
* void ptx_cp_async_bulk(Var shared_ptr,
* Expr shared_offset,
* Var global_ptr,
* Expr global_offset,
* size_t bytes,
* int barrier_arr_id,
* int barrier_id);
*/
TVM_DLL const Op& ptx_cp_async_bulk();
/*!
* \brief tvm intrinsics for ptx async bulk copy from shared::cta to shared::cluster
*
* void ptx_cp_async_bulk_shared_to_cluster(Expr dst_ptr,
* Expr src_ptr,
* Expr size,
* Expr mbar);
*/
TVM_DLL const Op& ptx_cp_async_bulk_shared_to_cluster();
/*!
* \brief tvm intrinsics for ptx async copy commit and wait.
*
* void ptx_cp_async_commit_group();
* void ptx_cp_async_wait_group(int num);
*
*/
TVM_DLL const Op& ptx_cp_async_commit_group();
TVM_DLL const Op& ptx_cp_async_wait_group();
/*!
* \brief tvm intrinsics for ptx async copy barrier using cp.async.mbarrier.arrive
*
* ptx_cp_async_mbarrier_arrive(int barrier_arr_id, int barrier_id)
*
*/
TVM_DLL const Op& ptx_cp_async_mbarrier_arrive();
/*!
* \brief PTX fence instruction: fence.{sem}.{scope}
*
* ptx_fence(StringImm sem, StringImm scope)
*/
TVM_DLL const Op& ptx_fence();
/*!
* \brief PTX fence.proxy.async instruction: fence.proxy.async[.{space}]
*
* ptx_fence_proxy_async(StringImm space)
*/
TVM_DLL const Op& ptx_fence_proxy_async();
/*!
* \brief tvm instrinsics to call mbarrier.init.shared::cta.b64
*
* ptx_mbarrier_init(uint64_t* bar_ptr, int thread_count)
*/
TVM_DLL const Op& ptx_mbarrier_init();
/*!
* \brief tvm instrinsics to call
* mbarrier.arrive.shared::cta.b64
* or
* @p mapa.shared::cluster.u32
* @p mbarrier.arrive.shared::cluster.b64
*/
TVM_DLL const Op& ptx_mbarrier_arrive();
/*!
* \brief tvm instrinsics to call
* mbarrier.arrive.expect_tx.shared.b64
* or
* @p mapa.shared::cluster.u32
* @p mbarrier.arrive.expect_tx.shared.b64
*
* ptx_mbarrier_arrive_expect_tx(uint64_t* bar_ptr, int byte_count)
*/
TVM_DLL const Op& ptx_mbarrier_arrive_expect_tx();
/*!
* \brief tvm instrinsics to call mbarrier.try_wait.parity repeatedly until it returns true
*
* ptx_mbarrier_try_wait(uint64_t* bar_ptr, int phase)
*/
TVM_DLL const Op& ptx_mbarrier_try_wait();
/*!
* \brief tvm instrinsics to call bar.arrive a, b
*
* bar_arrive(int name_bar_id, int thread_count)
*/
TVM_DLL const Op& ptx_bar_arrive();
/*!
* \brief tvm instrinsics to call bar.sync a, {b}
*
* bar_sync(int name_bar_id, int thread_count)
*/
TVM_DLL const Op& ptx_bar_sync();
/*!
* \brief tvm instrinsics to call
* cp.async.bulk.tensor.dim.shared::cluster.global.tile.mbarrier::complete_tx::bytes
*
* TMA alignment requirement:
* https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
*
* ptx_cp_async_bulk_tensor_global_to_cluster(int dim, PrimExpr dst_ptr, PrimExpr bar_ptr,
* PrimExpr tensormap_addr, int...coords, int cta_mask, int cta_group, string cache_hint)
*/
TVM_DLL const Op& ptx_cp_async_bulk_tensor_global_to_cluster();
/*!
* \brief tvm intrinsic to call
* cp.async.bulk.tensor.dim.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes
*
* TMA alignment requirement:
* https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
*
* ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster(int dim, PrimExpr dst_ptr, PrimExpr
* bar_ptr, PrimExpr tensormap_addr, int...coords, int cta_mask, int cta_group, string cache_hint)
*/
TVM_DLL const Op& ptx_cp_async_bulk_tensor_tile_gather4_global_to_cluster();
/*!
* \brief tvm instrinsics to call
* cp.async.bulk.tensor.dim.global.shared::cta.tile。bulk_group
*
* TMA alignment requirement:
* https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
*
* ptx_cp_async_bulk_tensor_shared_to_global(int dim, PrimExpr src_ptr, PrimExpr tensormap_addr,
* int...coords, string cache_hint)
*/
TVM_DLL const Op& ptx_cp_async_bulk_tensor_shared_to_global();
/*!
* \brief tvm instrinsics to call
* cp.async.bulk.prefetch.tensor.dim.L2.global.tile
*
* TMA alignment requirement:
* https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#table-alignment-multi-dim-tma
*
* ptx_cp_async_bulk_tensor_global_to_cluster_prefetch(int dim, PrimExpr tensormap_addr,
* int...coords, string cache_hint)
*/
TVM_DLL const Op& ptx_cp_async_bulk_tensor_global_to_cluster_prefetch();
/*!
* \brief tvm instrinsics to call
* cp.reduce.async.bulk.tensor.dim.dst.src.redOp
*
* ptx_cp_async_bulk_tensor_shared_to_global_reduce(int dim, PrimExpr src_ptr, PrimExpr
* tensormap_addr, int...coords, string cache_hint)
*/
TVM_DLL const Op& ptx_cp_async_bulk_tensor_shared_to_global_reduce();
/*!
* \brief tvm instrinsics to call cp.async.bulk.commit_group
*
* ptx_cp_async_bulk_commit_group()
*/
TVM_DLL const Op& ptx_cp_async_bulk_commit_group();
/*!
* \brief tvm instrinsics to call cp.async.bulk.wait_group{.read} N
*
* ptx_cp_async_bulk_wait_group(int N, bool read)
*/
TVM_DLL const Op& ptx_cp_async_bulk_wait_group();
/*!
* \brief tvm instrinsics to call barrier.cluster.arrive{.sem}{.aligned}
*
* ptx_barrier_cluster_arrive(string sem, bool aligned)
*/
TVM_DLL const Op& ptx_barrier_cluster_arrive();
/*!
* \brief tvm instrinsics to call barrier.cluster.wait.{acquire}{.aligned}
*
* ptx_barrier_cluster_wait(bool acquire, bool aligned)
*/
TVM_DLL const Op& ptx_barrier_cluster_wait();
/*!
* \brief tvm instrinsics to call elect.sync _|p, membermask and return the predicate
*
* elect_sync(membermask)
*/
TVM_DLL const Op& ptx_elect_sync();
/*!
* \brief PTX fence.mbarrier_init.release.cluster instruction
*
* ptx_fence_mbarrier_init()
*/
TVM_DLL const Op& ptx_fence_mbarrier_init();
/*!
* \brief tvm instrinsics to fetch PTX pre-defined registers
*
* ptx_fetch_register(int bits, string reg_name)
*/
TVM_DLL const Op& ptx_fetch_register();
/*!
* \brief tvm intrinsic for storing the result of PTX MMA into a destination pointer.
* For example, if each thread in a warp of size 32 has 4 elements from the result of
* m16xn8xk16 MMA in its registers, this intrinsic can be used to store the result in a
* 16x8 region in shared or global memory.
*
* There is no real PTX instruction that does that, but we want to hide details of
* complex index manipulation behind this intrinsic to simplify TIR lowering passes (e.g.
* LowerWarpMemory).
*
* void mma_store(IntImm m, IntImm n, Var dst_ptr, Var src_ptr, Expr src_offset, Var dst_stride);
*/
TVM_DLL const Op& mma_store();
/*!
* \brief tvm intrinsic for zero-initializing an MMA accumulation register.
* For example, if each thread in a warp of size 32 has 8 elements from the A matrix in
* m16xn8xk16 MMA in its registers, this intrinsic can be used to zero-initialize its
* 4 accumulation registers.
*
* There is no real PTX instruction that does that, but we introduce this intrinsic for the
* same reason as mma_store above.
*
* void mma_fill(IntImm local_size, Var local_ptr, Expr offset);
*/
TVM_DLL const Op& mma_fill();
/*!
* \brief tvm intrinsic to encode matrix descriptor for wgmma instructions.
*
* ptx_wgmma_encode_matrix_descriptor(PrimExpr ptr, PrimExpr ldo, PrimExpr sdo, int swizzle)
*/
TVM_DLL const Op& ptx_wgmma_encode_matrix_descriptor();
/*!
* \brief tvm intrinsic to call "" : "+r"(reg) :: "memory"
*
* ptx_wgmma_noop_barrier()
*/
TVM_DLL const Op& ptx_wgmma_noop_barrier();
/*!
* \brief tvm intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype
* where both A and B are in shared memory.
*
* ptx_wgmma_mma_async_ss()
*/
TVM_DLL const Op& ptx_wgmma_mma_async_ss();
/*!
* \brief tvm intrinsic to call wgmma.mma_async.sync.aligned.shape.dtype.atype.btype
* where A is in register and B is in shared memory.
*
* ptx_wgmma_mma_async_rs()
*/
TVM_DLL const Op& ptx_wgmma_mma_async_rs();
/*!
* \brief tvm intrinsic to call wgmma.fence.sync.aligned;
*
* ptx_wgmma_fence()
*/
TVM_DLL const Op& ptx_wgmma_fence();
/*!
* \brief tvm intrinsic to call wgmma.commit_group.sync.aligned;
*
* ptx_wgmma_commit_group()
*/
TVM_DLL const Op& ptx_wgmma_commit_group();
/*!
* \brief tvm intrinsic to call wgmma.wait_group.sync.aligned;
*
* ptx_wgmma_wait_group(int N)
*/
TVM_DLL const Op& ptx_wgmma_wait_group();
/*!
* \brief tvm intrinsic to call stmatrix.sync.aligned.m8n8.num{.trans}.shared.b16 [p], r;
*
* ptx_stmatrix(int num, bool trans, PrimExpr ptr, PrimExpr... vars)
*/
TVM_DLL const Op& ptx_stmatrix();
/*!
* \brief tvm intrinsic to call setmaxnreg.action.sync.aligned.u32 imm-reg-count
*/
TVM_DLL const Op& ptx_setmaxnreg();
/*!
* \brief tvm intrinsic to call ld.global.acquire.gpu.b32
*
* ptx_ld_global_acquire()
*/
TVM_DLL const Op& ptx_ld_global_acquire();
/*!
* \brief tvm instrinsics to call tcgen05.alloc.cta_group.sync.aligned;
*
* ptx_tcgen05_alloc(Var dst_ptr, int n_cols, int cta_group)
*/
TVM_DLL const Op& ptx_tcgen05_alloc();
/*!
* \brief tvm instrinsics to call tcgen05.dealloc.cta_group.sync.aligned;
*
* ptx_tcgen05_dealloc(uint32_t taddr, int n_cols, int cta_group)
*/
TVM_DLL const Op& ptx_tcgen05_dealloc();
/*!
* \brief tvm instrinsics to call tcgen05.relinquish_alloc_permit.cta_group.sync.aligned;
*
* ptx_tcgen05_relinquish_alloc_permit(int cta_group)
*/
TVM_DLL const Op& ptx_tcgen05_relinquish_alloc_permit();
/*!
* \brief tvm instrinsics to call tcgen05.fence::before_thread_sync;
*
* ptx_tcgen05_fence_before_thread_sync()
*/
TVM_DLL const Op& ptx_tcgen05_fence_before_thread_sync();
/*!
* \brief tvm instrinsics to call tcgen05.fence::after_thread_sync;
*
* ptx_tcgen05_fence_after_thread_sync()
*/
TVM_DLL const Op& ptx_tcgen05_fence_after_thread_sync();
/*!
* \brief tvm instrinsics to call tcgen05.ld.sync.aligned;
*
* ptx_tcgen05_ld()
*/
TVM_DLL const Op& ptx_tcgen05_ld();
/*!
* \brief tvm instrinsics to call tcgen05.st.sync.aligned;
*
* ptx_tcgen05_st()
*/
TVM_DLL const Op& ptx_tcgen05_st();
/*!
* \brief tvm instrinsics to call tcgen05.wait::ld.sync.aligned;
*
* ptx_tcgen05_wait_ld()
*/
TVM_DLL const Op& ptx_tcgen05_wait_ld();
/*!
* \brief tvm instrinsics to call tcgen05.wait::st.sync.aligned;
*
* ptx_tcgen05_wait_st()
*/
TVM_DLL const Op& ptx_tcgen05_wait_st();
/*!
* \brief tvm intrinsic to encode matrix descriptor for tcgen05 instructions.
*
* ptx_tcgen05_encode_matrix_descriptor(PrimExpr ptr, PrimExpr ldo, PrimExpr sdo, int swizzle)
*/
TVM_DLL const Op& ptx_tcgen05_encode_matrix_descriptor();
/*!
* \brief tvm intrinsic to encode instruction descriptor for tcgen05 MMA.
*
* ptx_tcgen05_encode_instr_descriptor(PrimExpr desc, string d_dtype, string a_dtype, string
* b_dtype, int M, int N, int K, bool trans_a, bool trans_b, int n_cta_groups, bool neg_a, bool
* neg_b, bool sat_d, bool is_sparse)
*/
TVM_DLL const Op& ptx_tcgen05_encode_instr_descriptor();
/*!
* \brief tvm intrinsic to encode instruction descriptor for tcgen05 MMA block scaled.
*
* ptx_tcgen05_encode_instr_descriptor_block_scaled(PrimExpr desc, string d_dtype,
* string a_dtype, string b_dtype, string sfa_dtype, string stb_dtype,
* int M, int N, int K, bool trans_a, bool trans_b,
* int n_cta_groups, bool neg_a, bool neg_b, bool is_sparse)
*/
TVM_DLL const Op& ptx_tcgen05_encode_instr_descriptor_block_scaled();
/*!
* \brief tvm intrinsic to call tcgen05.mma.cta_group.kind without block scaling.
*
* ptx_tcgen05_mma()
*/
TVM_DLL const Op& ptx_tcgen05_mma();
/*!
* \brief tvm intrinsic to call tcgen05.mma.cta_group.kind.block_scale{.scale_vec_size}
*
* ptx_tcgen05_mma_block_scale()
*/
TVM_DLL const Op& ptx_tcgen05_mma_block_scale();
/*!
* \brief tvm intrinsic to call tcgen05.mma.sp.cta_group.kind without block scaling.
*
* ptx_tcgen05_mma_sp()
*/
TVM_DLL const Op& ptx_tcgen05_mma_sp();
/*!
* \brief tvm intrinsic to call tcgen05.mma.sp.cta_group.kind.block_scale{.scale_vec_size}
*
* ptx_tcgen05_mma_sp_block_scale()
*/
TVM_DLL const Op& ptx_tcgen05_mma_sp_block_scale();
/*!
* \brief tvm instrinsics to call tcgen05.commit.cta_group
*
* ptx_tcgen05_commit()
*/
TVM_DLL const Op& ptx_tcgen05_commit();
/*!
* \brief tvm instrinsics to call tcgen05.cp.cta_group
*
* ptx_tcgen05_cp()
*/
TVM_DLL const Op& ptx_tcgen05_cp();
/*!
* \brief tvm instrinsics to call tcgen05.shift.cta_group.down
*
* ptx_tcgen05_shift()
*/
TVM_DLL const Op& ptx_tcgen05_shift();
/*!
* \brief tvm instrinsics to call map_shared_rank
*
* ptx_map_shared_rank(PrimExpr ptr, int rank)
*/
TVM_DLL const Op& ptx_map_shared_rank();
/*!
* \brief tvm instrinsics to call a CUDA function. Source code is provided as a string.
*
* cuda_func_call(String func_name, PrimExpr... args, String source_code)
*/
TVM_DLL const Op& cuda_func_call();
/*!
* \brief nvshmem intrinsics for nvshmem_my_pe() operation.
*
* int nvshmem_my_pe()
*/
TVM_DLL const Op& nvshmem_my_pe();
/*!
* \brief nvshmem intrinsics for nvshmem_n_pes() operation.
*
* int nvshmem_n_pes()
*/
TVM_DLL const Op& nvshmem_n_pes();
/*!
* \brief nvshmem intrinsics for nvshmem_getmem_nbi() operation.
*
* void nvshmem_getmem_nbi(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_getmem_nbi();
/*!
* \brief nvshmem intrinsics for nvshmem_putmem_nbi() operation.
*
* void nvshmem_putmem_nbi(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_nbi();
/*!
* \brief nvshmem intrinsics for nvshmemx_getmem_nbi_warp() operation.
*
* void nvshmemx_getmem_nbi_warp(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_getmem_nbi_warp();
/*!
* \brief nvshmem intrinsics for nvshmemx_putmem_nbi_warp() operation.
*
* void nvshmemx_putmem_nbi_warp(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_nbi_warp();
/*!
* \brief nvshmem intrinsics for nvshmemx_getmem_nbi_block() operation.
*
* void nvshmemx_getmem_nbi_block(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_getmem_nbi_block();
/*!
* \brief nvshmem intrinsics for nvshmemx_putmem_nbi_block() operation.
*
* void nvshmemx_putmem_nbi_block(void *dest, const void *source, size_t nelems, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_nbi_block();
/*!
* \brief nvshmem intrinsics for nvshmemx_signal_op() operation.
*
* void nvshmemx_signal_op(uint64_t *sig_addr, uint64_t signal, int sig_op, int pe)
*/
TVM_DLL const Op& nvshmem_signal_op();
/*!
* \brief nvshmem intrinsics for nvshmem_FuncParam{TYPENAME}_wait_until() operation.
*
* void nvshmem_FuncParam{TYPENAME}_wait_until(TYPE *ivar, int cmp, TYPE cmp_value)
*/
TVM_DLL const Op& nvshmem_wait_until();
/*!
* \brief nvshmem intrinsics for nvshmem_quiet() operation.
*
* void nvshmem_quiet()
*/
TVM_DLL const Op& nvshmem_quiet();
/*!
* \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi() operation.
*
* void nvshmemx_putmem_signal_nbi(void *dest, const void *source, size_t nelems, uint64_t
* *sig_addr, uint64_t signal, int sig_op, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_signal_nbi();
/*!
* \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi_warp() operation.
*
* void nvshmemx_putmem_signal_nbi_warp(void *dest, const void *source, size_t nelems, uint64_t
* *sig_addr, uint64_t signal, int sig_op, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_signal_nbi_warp();
/*!
* \brief nvshmem intrinsics for nvshmemx_putmem_signal_nbi_block() operation.
*
* void nvshmemx_putmem_signal_nbi_block(void *dest, const void *source, size_t nelems,
* uint64_t *sig_addr, uint64_t signal, int sig_op, int pe)
*/
TVM_DLL const Op& nvshmem_putmem_signal_nbi_block();
/*!
* \brief nvshmem intrinsics for nvshmem_fence() operation.
*
* void nvshmem_fence()
*/
TVM_DLL const Op& nvshmem_fence();
/*!
* \brief nvshmem intrinsics for nvshmem_barrier_all() operation.
*
* void nvshmem_barrier_all()
*/
TVM_DLL const Op& nvshmem_barrier_all();
} // namespace builtin
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_TARGET_BUILTIN_CUDA_H_
include/tvm/tirx/target_builtin/trn.h
+156
View File
@@ -0,0 +1,156 @@
/*
* 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/target_builtin/trn.h
* \brief TIR builtin intrinsics specific to Trainium target.
*/
#ifndef TVM_TIRX_TARGET_BUILTIN_TRN_H_
#define TVM_TIRX_TARGET_BUILTIN_TRN_H_
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
namespace tvm {
namespace tirx {
namespace builtin {
/*!
* \brief nki intrinsics for load operation.
*
* nki_load(result, data)
*/
TVM_DLL const Op& nki_load();
/*!
* \brief nki intrinsics for store operation.
*
* nki_store(result, data)
*/
TVM_DLL const Op& nki_store();
/*!
* \brief nki intrinsics for tensor_copy operation.
*
* nki_tensor_copy(result, data)
*/
TVM_DLL const Op& nki_tensor_copy();
/*!
* \brief nki intrinsics for matmul operation.
*
* nki_matmul(C, A, B, accum)
*
* equivalent to C += A.T @ B (if accum is true), or C = A.T @ B (if accum is false)
*/
TVM_DLL const Op& nki_matmul();
/*!
* \brief nki intrinsics for activation operation.
*
* nki_activation(result, data, opcode, bias, scale)
*/
TVM_DLL const Op& nki_activation();
/*!
* \brief nki intrinsics for reciprocal operation.
*
* nki_reciprocal(result, data)
*/
TVM_DLL const Op& nki_reciprocal();
/*!
* \brief nki intrinsics for tensortensor operation.
*
* nki_tensortensor(result, operand0, operand1, opcode)
*/
TVM_DLL const Op& nki_tensortensor();
/*!
* \brief nki intrinsics for tensorscalar operation.
*
* nki_tensorscalar(result, operand0, operand1, opcode, reverse)
*/
TVM_DLL const Op& nki_tensorscalar();
/*!
* \brief nki intrinsics for tensorreduce operation.
*
* nki_tensorreduce(result, data, opcode, negate, axes)
*/
TVM_DLL const Op& nki_tensorreduce();
/*!
* \brief nki intrinsics for memset operation.
*
* nki_memset(result, value)
*/
TVM_DLL const Op& nki_memset();
/*!
* \brief nki intrinsics for activation reduce operation.
*
* nki_activation_reduce(reduce_res, act_res, data, opcode, reduce_opcode, bias, scale)
*/
TVM_DLL const Op& nki_activation_reduce();
/*!
* \brief nki intrinsics for tensorscalar reduce operation.
*
* nki_tensorscalar_reduce(reduce_res, tensorscalar_res, operand0, operand1, opcode, reduce_opcode,
* reverse)
*/
TVM_DLL const Op& nki_tensorscalar_reduce();
/*!
* \brief nki intrinsics for initializing identity tensor.
*
* nki_identity(result, size)
*/
TVM_DLL const Op& nki_identity();
/*!
* \brief nki intrinsics for scalar tensor tensor operation.
*
* (data op1 operand1) op2 (operand2) where op1 is tensor-scalar and op2 is tensor-tensor
*
* nki_scalar_tensor_tensor(result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1)
*
*/
TVM_DLL const Op& nki_scalar_tensor_tensor();
/*!
* \brief nki intrinsics for scalar tensor scalar operation.
*
* (data op1 operand1) op2 (operand2) where op1 and op2 are tensor-scalar
*
* nki_scalar_tensor_scalar(result, data, operand0, operand1, opcode0, opcode1, reverse0, reverse1)
*
*/
TVM_DLL const Op& nki_scalar_tensor_scalar();
/*!
* \brief nki intrinsics for affine_select operation.
*
* nki_affine_select(result, pred, true_value, false_value)
*/
TVM_DLL const Op& nki_affine_select();
} // namespace builtin
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_TARGET_BUILTIN_TRN_H_
include/tvm/tirx/tirx_op.h
+314
View File
@@ -0,0 +1,314 @@
/*
* 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 built-in operators.
*/
#ifndef TVM_TIRX_TIRX_OP_H_
#define TVM_TIRX_TIRX_OP_H_
#include <tvm/ir/op.h>
#include <tvm/target/target.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/tirx_stmt.h>
namespace tvm {
namespace tirx {
/*!
* \brief The type of the function that sanitizes the arguments of a TIRX operator.
* \param op The operator.
* \param args The arguments.
*/
using FArgSanitizer = ffi::TypedFunction<void(tvm::Op, ffi::Array<ffi::ObjectRef>)>;
namespace callback {
/*! \brief The buffers allocated by the operator. */
constexpr const char* kPrivateAlloc = "private_alloc";
/*! \brief The initialization statement of the operator.
* which will be inserted at the beginning of the kernel
*/
constexpr const char* kDeviceInitStmt = "device_init_stmt";
/*! \brief The initialization statement of the operator.
* which will be inserted at the beginning of the kernel
*/
constexpr const char* kHostInitStmt = "host_init_stmt";
/*! \brief Statements to be inserted after a specific buffer's definition (DeclBuffer/AllocBuffer).
* Stored as Map<Buffer, ffi::Array<Stmt>>.
*/
constexpr const char* kPostBufferDefStmt = "post_buffer_def_stmt";
} // namespace callback
/*!
* \brief The context information of the kernel required by op schedule.
*/
class ScheduleContextNode : public ffi::Object {
public:
/*! \brief The target of the kernel. */
Target target;
/*! \brief The exec scope of the operator */
ExecScope exec_scope;
/*! \brief The kernel launch parameters. */
ffi::Map<ffi::String, IterVar> launch_params;
/*! \brief A map from loop variables to their ranges. */
ffi::Map<Var, Range> var_range_map;
/*! \brief Whether the schedule context is only used for buffer allocation. */
bool alloc_only;
/*! \brief Callback to be handled when the operator is scheduled. */
ffi::Map<ffi::String, ffi::ObjectRef> callbacks;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<ScheduleContextNode>()
.def_ro("target", &ScheduleContextNode::target)
.def_ro("exec_scope", &ScheduleContextNode::exec_scope)
.def_ro("launch_params", &ScheduleContextNode::launch_params)
.def_ro("var_range_map", &ScheduleContextNode::var_range_map)
.def_ro("alloc_only", &ScheduleContextNode::alloc_only)
.def_ro("callbacks", &ScheduleContextNode::callbacks);
}
/*! \brief Add a buffer to be allocated in the kernel. */
void AddAllocBuffer(Buffer buffer);
/*! \brief Add an initialization statement to be inserted.
* \param stmt The statement to be inserted.
* \param host 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).
*/
void AddInitStmt(Stmt stmt, bool host = false);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ScheduleContext", ScheduleContextNode, ffi::Object);
};
/*!
* \brief Managed reference to ScheduleContextNode.
*/
class ScheduleContext : public ffi::ObjectRef {
public:
/*!
* \brief Constructor.
* \param target The target of the kernel.
* \param exec_scope The exec scope of the operator.
* \param launch_params The kernel launch parameters.
* \param var_range_map: A map from loop variables to their ranges.
* \param alloc_only Whether the schedule context is only used for buffer allocation.
* \param callbacks The callbacks to be handled when the operator is scheduled.
*/
TVM_DLL ScheduleContext(Target target, ExecScope exec_scope,
ffi::Map<ffi::String, IterVar> launch_params = {},
ffi::Map<Var, Range> var_range_map = {}, bool alloc_only = false,
ffi::Map<ffi::String, ffi::ObjectRef> callbacks = {});
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScheduleContext, ffi::ObjectRef, ScheduleContextNode);
};
/*!
* \brief The type of the function that schedules a TIRX operator.
* \param op The operator.
* \param args The arguments.
* \param context The schedule context.
*/
using FOpScheduler = ffi::TypedFunction<Stmt(tvm::Op, ffi::Array<ffi::ObjectRef>, ScheduleContext)>;
/*!
* \brief The context information of the kernel required by op dispatch.
*/
class DispatchContextNode : public ffi::Object {
public:
/*! \brief The target of the kernel. */
Target target;
/*! \brief The exec scope of the operator */
ExecScope exec_scope;
/*! \brief The kernel launch parameters. */
ffi::Map<ffi::String, IterVar> launch_params;
/*! \brief A map from loop variables to their ranges. */
ffi::Map<Var, Range> var_range_map;
/*! \brief Whether the dispatch context is only used for buffer allocation. */
bool alloc_only;
/*! \brief Callback to be handled when the operator is scheduled. */
ffi::Map<ffi::String, ffi::ObjectRef> callbacks;
/*! \brief Shared state that persists across dispatch calls within a single lowering pass. */
ffi::Map<ffi::String, ffi::ObjectRef> shared_state;
/*!
* \brief ExecContext inter-team view at this op site.
*
* Maps axis name ("laneid"/"warpid"/"cta_id"/"wid_in_wg"/"wgid") to a
* 2-element [extent, offset] PrimExpr array. Empty map = no ExecContext
* tracking available (fallback for unresolved filters, pre-Phase-4 call
* sites, etc.); dispatchers should fall back to exec_scope.name in that
* case.
*/
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter;
/*! \brief ExecContext intra-team view. Same encoding as ``inter``. */
ffi::Map<ffi::String, ffi::Array<PrimExpr>> intra;
/*! \brief Scope kind string ("kernel"/"cta"/"warpgroup"/"warp"/"thread"/"cluster"). */
ffi::String scope_kind;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<DispatchContextNode>()
.def_ro("target", &DispatchContextNode::target)
.def_ro("exec_scope", &DispatchContextNode::exec_scope)
.def_ro("launch_params", &DispatchContextNode::launch_params)
.def_ro("var_range_map", &DispatchContextNode::var_range_map)
.def_ro("alloc_only", &DispatchContextNode::alloc_only)
.def_ro("callbacks", &DispatchContextNode::callbacks)
.def_ro("shared_state", &DispatchContextNode::shared_state)
.def_ro("inter", &DispatchContextNode::inter)
.def_ro("intra", &DispatchContextNode::intra)
.def_ro("scope_kind", &DispatchContextNode::scope_kind);
}
/*! \brief Add a buffer to be allocated in the kernel. */
void AddAllocBuffer(Buffer buffer);
/*! \brief Add an initialization statement to be inserted. */
void AddInitStmt(Stmt stmt, bool host = false);
/*! \brief Add a statement to be inserted after a buffer's definition. */
void AddPostBufferDefStmt(Buffer buffer, Stmt stmt);
/*! \brief Set a value in the shared state cache. */
void SharedStateSet(ffi::String key, ffi::ObjectRef value);
/*! \brief Get a value from the shared state cache. */
ffi::Optional<ffi::ObjectRef> SharedStateGet(ffi::String key);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DispatchContext", DispatchContextNode, ffi::Object);
};
/*!
* \brief Managed reference to DispatchContextNode.
*/
class DispatchContext : public ffi::ObjectRef {
public:
TVM_DLL DispatchContext(Target target, ExecScope exec_scope,
ffi::Map<ffi::String, IterVar> launch_params = {},
ffi::Map<Var, Range> var_range_map = {}, bool alloc_only = false,
ffi::Map<ffi::String, ffi::ObjectRef> callbacks = {},
ffi::Map<ffi::String, ffi::ObjectRef> shared_state = {},
ffi::Map<ffi::String, ffi::Array<PrimExpr>> inter = {},
ffi::Map<ffi::String, ffi::Array<PrimExpr>> intra = {},
ffi::String scope_kind = "");
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DispatchContext, ffi::ObjectRef, DispatchContextNode);
};
/*!
* \brief See pesudo code below:
*
* Tx.cast(BufferRegion dst, BufferRegion src)
*/
TVM_DLL const Op& cast();
/*!
* \brief See pesudo code below:
*
* Tx.permute_dims(BufferRegion buffer, List order)
*/
TVM_DLL const Op& permute_dims();
/*!
* \brief See pesudo code below:
*
* Tx.copy(BufferRegion dst, BufferRegion src)
*/
TVM_DLL const Op& copy();
/*!
* \brief See pesudo code below:
*
* Tx.Async.copy(BufferRegion dst, BufferRegion src)
*/
TVM_DLL const Op& copy_async();
/*!
* \brief See pesudo code below:
*
* Tx.fill(BufferRegion dst, PrimExpr value)
*/
TVM_DLL const Op& fill();
/*!
* \brief See pesudo code below:
*
* Tx.gemm(Buffer A, Buffer B, Buffer C, Buffer D, PrimExpr alpha, PrimExpr beta)
*/
TVM_DLL const Op& gemm();
/*!
* \brief See pesudo code below:
*
* Tx.gemm_async(BufferRegion C, BufferRegion A, BufferRegion B, bool transA, bool transB,
* bool accum)
*/
TVM_DLL const Op& gemm_async();
TVM_DLL const Op& zero();
TVM_DLL const Op& sqrt();
TVM_DLL const Op& exp();
TVM_DLL const Op& add();
TVM_DLL const Op& sub();
TVM_DLL const Op& mul();
TVM_DLL const Op& fdiv();
TVM_DLL const Op& minimum();
TVM_DLL const Op& maximum();
TVM_DLL const Op& reciprocal();
TVM_DLL const Op& sum();
TVM_DLL const Op& max();
TVM_DLL const Op& min();
TVM_DLL const Op& memset();
TVM_DLL const Op& reduce_negate();
TVM_DLL const Op& binary_reduce();
TVM_DLL const Op& unary_reduce();
TVM_DLL const Op& binary_chain();
TVM_DLL const Op& select();
/*!
* \brief See pesudo code below:
*
* tvm_kernel_replace_point()
*/
TVM_DLL const Op& tvm_kernel_replace_point();
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_TIRX_OP_H_
include/tvm/tirx/tirx_stmt.h
+85
View File
@@ -0,0 +1,85 @@
/*
* 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 <tvm/ir/op.h>
#include <tvm/tirx/stmt.h>
namespace tvm {
namespace tirx {
/*!
* \brief TIRX TilePrimitiveCall stmt.
*/
class TilePrimitiveCallNode : public StmtNode {
public:
// tvm::Op which corresponds to the TIRX operator.
tvm::Op op;
// Arguments to the operator.
ffi::Array<ffi::Any> args;
// Workspace (pre-allocated buffers) for the operator.
ffi::Map<ffi::String, Buffer> workspace;
// Config for the operator/scheduler.
ffi::Map<ffi::String, ffi::Any> config;
// Optional dispatch variant name registered via @register_dispatch.
ffi::Optional<ffi::String> dispatch{std::nullopt};
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TilePrimitiveCallNode>()
.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);
}
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<ffi::Any> args,
ffi::Map<ffi::String, Buffer> workspace = {},
ffi::Map<ffi::String, ffi::Any> config = {},
ffi::Optional<ffi::String> dispatch = std::nullopt);
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_
include/tvm/tirx/transform.h
+26 -8
View File
@@ -343,17 +343,35 @@ TVM_DLL Pass AnnotateEntryFunc();
TVM_DLL Pass Filter(ffi::TypedFunction<bool(PrimFunc)> fcond);
/*!
* \brief Remove the weight layout rewrite block
* \param skip_tensor_rewrite If True, exact rewrite of Tensor, according to the given index map,
* will be skipped. Only the shape of the Tensor is transformed correctly, and the content of
* the destination array will be filled with random values.
*
* When this pass is called many times during MetaSchedule tuning, the raw data of Tensor,
* before and after rewrite, does not matter. Since Tensor layout rewrite, using IndexMap's
* MapTensor, is currently slow, skipping the exact rewrite is sometimes necessary.
* \brief Lower TIRx op calls using registered op dispatchers for the given target.
*
* Also resolves ScopeIdDef declarations: gathers them at kernel scope, verifies
* consistency, extracts launch parameters, and emits Bind statements +
* thread_extent AttrStmts wrapping the dispatched body.
* \return The pass.
*/
TVM_DLL Pass TilePrimitiveDispatch();
/*!
* \brief Finalize TIRx lowering by applying layout rewriters and cleanup passes.
* \return The pass.
*/
TVM_DLL Pass LowerTIRxCleanup();
/*!
* \brief Lower opaque constructs in TIRX programs: AllocBuffer, For(thread_binding),
* unit loop elimination. This is the tirx-specific counterpart of
* s_tir::LowerOpaqueBlock, without any SBlock handling.
* \return The pass.
*/
TVM_DLL Pass LowerTIRxOpaque();
/*!
* \brief Lower the TIR to a lower level IR for the given target.
* \return The pass.
*/
TVM_DLL Pass LowerTIRx();
} // namespace transform
} // namespace tirx
} // namespace tvm
include/tvm/topi/transform.h
+3 -3
View File
@@ -1811,8 +1811,8 @@ inline Tensor layout_transform(const Tensor& src, const std::string& src_layout,
const std::string schedule_rule = "None",
const std::string name = "T_layout_trans",
const std::string tag = kInjective) {
Layout src_layout_struct(src_layout);
Layout dst_layout_struct(dst_layout);
SLayout src_layout_struct(src_layout);
SLayout dst_layout_struct(dst_layout);
if (src_layout_struct.Equals(dst_layout_struct)) {
return src;
@@ -1821,7 +1821,7 @@ inline Tensor layout_transform(const Tensor& src, const std::string& src_layout,
TVM_FFI_ICHECK(src_layout_struct.defined() && dst_layout_struct.defined())
<< "cannot convert from/to undefined layout";
auto layout_converter = tirx::BijectiveLayout(src_layout_struct, dst_layout_struct);
auto layout_converter = tirx::SBijectiveLayout(src_layout_struct, dst_layout_struct);
TVM_FFI_ICHECK(layout_converter.defined())
<< "cannot convert from " << src_layout << " to " << dst_layout;
pyproject.toml
+8
View File
@@ -230,6 +230,14 @@ unfixable = []
[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["E402", "F401", "F403", "F405"]
"python/tvm/relax/op/nn/nn.py" = ["E501"]
"docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py" = ["RUF003"]
"python/tvm/relax/frontend/tflite/tflite_frontend.py" = ["E501"]
"python/tvm/relax/transform/legalize_ops/nn.py" = ["E501"]
# Scope-id declarations like ``lane_id = Tx.lane_id([32])`` register a TIR
# scope_id for side effect; the Python handle is often unused. Silence F841
# for paths that heavily use this idiom.
"tests/python/tirx/**/*.py" = ["F841"]
[tool.ruff.lint.isort]
known-first-party = ["tvm"]
python/tvm/__init__.py
+5 -3
View File
@@ -51,9 +51,6 @@ from . import script
# tvm.tirx — registers itself via tvm.script.register_dialect in its __init__
from . import tirx
# tvm.s_tir
from . import s_tir
# tvm.target
from . import target
@@ -75,6 +72,11 @@ from .contrib import rocm as _rocm, nvcc as _nvcc
# Relax contain modules that are only available in compiler package
# Do not import them if TVM is built with runtime only
if not _RUNTIME_ONLY:
# tile_primitive imports both Python Op class declarations (Zero, Add, ...)
# and per-target dispatch schedule registrations. Must run before relax so
# any relax pass that looks up a schedule sees them.
from .tirx.operator import tile_primitive
# tvm.relax — registers itself via tvm.script.register_dialect in its __init__
from . import relax
python/tvm/contrib/cutlass/attention_operation.py
+7 -7
View File
@@ -26,7 +26,7 @@ def instantiate_attention_template(attrs):
based on a template and the provided attribute map."""
bias_template = """
TVM_FFI_CHECK(${bias}->ndim == 4, ValueError); // B, N, S, S'
TVM_FFI_ICHECK(${bias}->ndim == 4); // B, N, S, S'
p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
p.bias_strideM = ${bias_strideM};
@@ -46,9 +46,9 @@ def instantiate_attention_template(attrs):
p.query_ptr = reinterpret_cast<T *>(${query}->data);
p.key_ptr = reinterpret_cast<T *>(${key}->data);
p.value_ptr = reinterpret_cast<T *>(${value}->data);
TVM_FFI_CHECK(${query}->ndim == 4, ValueError); // B, S, N, H
TVM_FFI_CHECK(${key}->ndim == 4, ValueError); // B, S', N, H
TVM_FFI_CHECK(${value}->ndim == 4, ValueError); // B, S', N, H'
TVM_FFI_ICHECK(${query}->ndim == 4); // B, S, N, H
TVM_FFI_ICHECK(${key}->ndim == 4); // B, S', N, H
TVM_FFI_ICHECK(${value}->ndim == 4); // B, S', N, H'
// stride for N
p.q_strideH = p.head_dim; // H
@@ -69,7 +69,7 @@ def instantiate_attention_template(attrs):
p.query_ptr = reinterpret_cast<T *>(${qkv}->data);
p.key_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads;
p.value_ptr = reinterpret_cast<T *>(${qkv}->data) + p.head_dim * p.num_heads * 2;
TVM_FFI_CHECK(${qkv}->ndim == 3, ValueError); // B, S, NH + NH + NH'
TVM_FFI_ICHECK(${qkv}->ndim == 3); // B, S, NH + NH + NH'
// stride for N
p.q_strideH = p.head_dim; // H
@@ -132,7 +132,7 @@ def instantiate_attention_template(attrs):
p.o_strideM = p.head_dim_value * p.num_heads; // H' * N
TVM_FFI_CHECK(out0->ndim == 4, ValueError); // B, S, N, H'
TVM_FFI_ICHECK(out0->ndim == 4); // B, S, N, H'
${qkv_template}
${bias_template}
@@ -148,7 +148,7 @@ def instantiate_attention_template(attrs):
}();
}
TVM_FFI_CHECK(Attention::check_supported(p), RuntimeError);
TVM_FFI_ICHECK(Attention::check_supported(p));
cudaStream_t stream = static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, ${query}->device.device_id));
kernel_fn<<<p.getBlocksGrid(), p.getThreadsGrid(), smem_bytes, stream>>>(p);
python/tvm/contrib/nvcc.py
+41 -2
View File
@@ -135,6 +135,11 @@ def _compile_cuda_nvcc(
file_name = "tvm_kernels"
if target_format is None and not use_nvshmem:
target_format = "ptx"
tvm_kernel_dump = os.environ.get("TVM_KERNEL_DUMP", None)
if tvm_kernel_dump is not None:
target_format = "fatbin" # use fatbin to get cubin for SASS extraction
if target_format not in ["cubin", "ptx", "fatbin"]:
raise ValueError("target_format must be in cubin, ptx, fatbin")
temp_code = temp.relpath(f"{file_name}.cu")
@@ -146,6 +151,9 @@ def _compile_cuda_nvcc(
if "cuda.kernels_output_dir" in pass_context.config
else None
)
if tvm_kernel_dump is not None:
kernels_output_dir = tvm_kernel_dump
if kernels_output_dir is not None:
if not os.path.isdir(kernels_output_dir):
os.makedirs(kernels_output_dir)
@@ -162,13 +170,33 @@ def _compile_cuda_nvcc(
cmd = ["nvcc"]
cmd += [f"--{target_format}", "-O3"]
if kernels_output_dir is not None:
if tvm_kernel_dump is not None:
cmd += ["-lineinfo"]
cmd += ["--keep", f"--keep-dir={tvm_kernel_dump}"]
if os.environ.get("TVM_KERNEL_DEBUG", "0") == "1":
cmd += ["-g"]
cmd += ["-G"]
if isinstance(arch, list):
cmd += arch
elif isinstance(arch, str):
cmd += ["-arch", arch]
cmd += [
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT16_OPERATORS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT162_OPERATORS__",
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
"--expt-relaxed-constexpr",
"--expt-extended-lambda",
"--use_fast_math",
"--ptxas-options=-v", # printing out number of registers
"--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage", # printing out number of registers # noqa: E501
]
major, _ = parse_compute_version(get_target_compute_version(Target.current(allow_none=True)))
if options:
if isinstance(options, str):
cmd += [options]
@@ -786,6 +814,9 @@ def tvm_callback_cuda_compile(code):
Compiler backend: "nvcc" (default) or "nvrtc"
- "nvcc": Use nvcc subprocess, generates fatbin
- "nvrtc": Use NVRTC via cuda-python for faster JIT, generates cubin
TVM_KERNEL_DUMP : str
If set, dump generated CUDA/intermediate files and append "-lineinfo" so profilers can
correlate SASS back to the dumped source.
Parameters
----------
@@ -910,7 +941,15 @@ def get_target_compute_version(target=None):
# 3. GPU compute version
if tvm.cuda(0).exist:
return tvm.cuda(0).compute_version
cv = tvm.cuda(0).compute_version
# Append 'a' suffix for SM 9.0+ (Hopper, Blackwell) which need
# architecture-specific instructions (wgmma, tcgen05, etc.).
major_minor = cv.split(".")
if len(major_minor) == 2 and major_minor[0].isdigit():
major = int(major_minor[0])
if major >= 9:
return cv + ".a"
return cv
raise ValueError(
"No CUDA architecture was specified or GPU detected."
python/tvm/ir/__init__.py
+2 -7
View File
@@ -37,12 +37,7 @@ from .function import BaseFunc, CallingConv
from .global_info import GlobalInfo, DummyGlobalInfo, VDevice
from .module import IRModule
from .op import Op, register_intrin_lowering, register_op_attr
from .type import (
FuncType,
PointerType,
PrimType,
TupleType,
Type,
)
from .type import FuncType, PointerType, PrimType, TupleType, Type
from . import analysis
from tvm_ffi import Array, Map
python/tvm/relax/backend/gpu_generic/cumsum.py
+14 -10
View File
@@ -95,7 +95,9 @@ def gpu_2d_continuous_cumsum(
shared_buf = T.sblock_alloc_buffer((block_elem,), out_dtype, scope="shared")
for ty in T.thread_binding(TY, thread="threadIdx.y"):
for tx in T.thread_binding(TX, thread="threadIdx.x"):
tx_idx = bx * block_elem + ty * warp_elem + tx * thread_elem
tx_idx: T.let[T.int64] = (
bx * block_elem + ty * warp_elem + tx * thread_elem
)
# Load data from global memory
for i in T.vectorized(N):
local_buf[i] = T.if_then_else(
@@ -112,7 +114,7 @@ def gpu_2d_continuous_cumsum(
# Inclusive scan inside warp
for i in T.unroll(LOG_TX):
for j in T.vectorized(N):
idx: T.int64 = ty * warp_elem + tx * thread_elem
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
if tx >= (1 << i):
shared_buf[idx + j] += shared_buf[
idx - (1 << i) * thread_elem + N - 1
@@ -121,11 +123,11 @@ def gpu_2d_continuous_cumsum(
for i in T.unroll(1, TY):
for j in T.vectorized(N):
if ty == 0:
idx: T.int64 = i * warp_elem + tx * thread_elem
idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
shared_buf[idx + j] += shared_buf[i * warp_elem - 1]
# Write sum of block to global memory
for i in T.vectorized(N):
idx: T.int64 = ty * warp_elem + tx * thread_elem + i
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem + i
if bx * block_elem + idx < cur_len:
output[by, src_offset + bx * block_elem + idx] = shared_buf[idx]
if tx == 0 and ty == 0:
@@ -146,26 +148,28 @@ def gpu_2d_continuous_cumsum(
for ty in T.thread_binding(TY, thread="threadIdx.y"):
for tx in T.thread_binding(TX, thread="threadIdx.x"):
for i in T.serial(N):
idx: T.int64 = bx * block_elem + ty * warp_elem + i * TX + tx
idx: T.let[T.int64] = bx * block_elem + ty * warp_elem + i * TX + tx
if idx < cur_len:
output[by, out_offset + idx] += T.if_then_else(
bx > 0, source[by, src_offset + bx - 1], 0
)
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def cumsum(var_a: T.handle, var_out: T.handle):
T.func_attr({"tirx.is_scheduled": True}) # prevent further scheduling
m, n = T.int64(), T.int64()
A = T.match_buffer(var_a, [m, n], dtype=in_dtype)
Out = T.match_buffer(var_out, [m, n], dtype=out_dtype)
Tmp = T.alloc_buffer([m, n], dtype=out_dtype)
total_rounds = T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) // LOG_BLOCK_N
total_rounds: T.let[T.int64] = (
T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) // LOG_BLOCK_N
)
block_inclusive_inside_block(
m, n, A, Out, Tmp, src_offset=T.int64(0), tmp_offset=T.int64(0)
)
for i in range(total_rounds):
cur_len = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1)))
cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1)))
block_inclusive_inside_block(
m,
cur_len,
@@ -176,8 +180,8 @@ def gpu_2d_continuous_cumsum(
tmp_offset=(i + 1) * T.ceildiv(n, block_elem),
)
for i in range(total_rounds - 1):
real_idx = total_rounds - 1 - i - 1
cur_len = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1)))
real_idx: T.let[T.int64] = total_rounds - 1 - i - 1
cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1)))
update_cross_block(
m,
cur_len,
python/tvm/relax/backend/gpu_generic/sampling.py
+10 -10
View File
@@ -114,7 +114,7 @@ def gpu_multinomial_from_uniform(
# Inclusive scan inside warp
for i in T.unroll(LOG_TX):
for j in T.vectorized(thread_elem):
idx: T.int64 = ty * warp_elem + tx * thread_elem
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
if tx >= (1 << i):
output_shared[idx + j] += output_shared[
idx - (1 << i) * thread_elem + thread_elem - 1
@@ -123,7 +123,7 @@ def gpu_multinomial_from_uniform(
for i in T.unroll(1, TY):
for j in T.vectorized(thread_elem):
if ty == 0:
idx: T.int64 = i * warp_elem + tx * thread_elem
idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
output_shared[idx + j] += output_shared[i * warp_elem - 1]
def compare_bool_not_equal(a: T.bool, b: T.bool) -> T.bool:
@@ -140,7 +140,7 @@ def gpu_multinomial_from_uniform(
):
with T.sblock():
shared_buf = T.sblock_alloc_buffer((TX * TY,), "bool", scope="shared")
tx_idx = ty * TX + tx
tx_idx: T.let[T.int64] = ty * TX + tx
shared_buf[tx_idx] = source_local[thread_elem - 1]
output_local[0] = T.if_then_else(
tx_idx != 0,
@@ -170,7 +170,7 @@ def gpu_multinomial_from_uniform(
with T.sblock():
local_sum = T.sblock_alloc_buffer((), dtype, scope="local")
shared_buf = T.sblock_alloc_buffer((TX * TY,), dtype, scope="shared")
idx = ty * TX + tx
idx: T.let[T.int64] = ty * TX + tx
local_sum[()] = T.Cast(dtype, init_value)
for i in T.unroll(thread_elem):
@@ -209,8 +209,8 @@ def gpu_multinomial_from_uniform(
step_aggregate = T.sblock_alloc_buffer((), prob_dtype, scope="local")
# Load prob data from global memory to local memory
for v in T.unroll(thread_elem):
idx = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
prob_local = T.if_then_else(
idx: T.let[T.int64] = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
prob_local: T.let = T.if_then_else(
idx < vocab_size,
prob[row_idx, idx],
T.Cast(prob_dtype, 0),
@@ -258,7 +258,7 @@ def gpu_multinomial_from_uniform(
aggregate[()] += step_aggregate[()]
@T.prim_func
@T.prim_func(s_tir=True)
def parallel_sampling_from_prob(
var_prob: T.handle,
var_uniform_samples: T.handle,
@@ -278,10 +278,10 @@ def gpu_multinomial_from_uniform(
step_iter = T.sblock_alloc_buffer((), "int32", scope="local")
for bx in T.thread_binding(batch_size, thread="blockIdx.x"):
row_idx = row_indices[bx, 0]
row_idx: T.let[T.int64] = row_indices[bx, 0]
for ty in T.thread_binding(TY, thread="threadIdx.y"):
for tx in T.thread_binding(TX, thread="threadIdx.x"):
u = uniform_samples[bx, 0]
u: T.let[T.float32] = uniform_samples[bx, 0]
aggregate[()] = T.Cast(prob_dtype, 0)
step_iter[()] = T.int32(0)
# at least one iteration
@@ -317,7 +317,7 @@ def generic_get_sample_index(
):
"""Generate a generic get_sample_index kernel."""
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_sample_index(A: T.handle, B: T.handle, C: T.handle, D: T.handle):
batch, vocab_size = T.int64(), T.int64()
prob = T.match_buffer(A, (batch, vocab_size), prob_dtype)
python/tvm/relax/block_builder.py
+2 -2
View File
@@ -474,7 +474,7 @@ class BlockBuilder(Object):
@tvm.script.ir_module
class Module:
@T.prim_func
@T.prim_func(s_tir=True)
def te_func(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle,
var_compute: T.handle) -> None:
# function attr dict
@@ -523,7 +523,7 @@ class BlockBuilder(Object):
@tvm.script.ir_module
class Module:
@T.prim_func
@T.prim_func(s_tir=True)
def te_func(var_rxplaceholder: T.handle, var_compute: T.handle, n: T.int64) -> None:
rxplaceholder = T.match_buffer(var_rxplaceholder, [n + T.int64(1)],
dtype="float32")
python/tvm/relax/frontend/nn/llm/_decode_kernels.py
+21 -21
View File
@@ -56,7 +56,7 @@ def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, slidi
if sliding_window:
global_symbol += "_sliding_window"
@T.prim_func(check_well_formed=False)
@T.prim_func(s_tir=True)
def batch_decode_paged_kv(
Q_handle: T.handle,
pages_handle: T.handle,
@@ -116,8 +116,8 @@ def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, slidi
scale_O = T.sblock_alloc_buffer((1,), "float32")
factor = T.sblock_alloc_buffer((1,), "float32")
cur_page_indptr_begin: T.int32 = page_table_indptr[b]
cur_page_indptr_end: T.int32 = page_table_indptr[b + 1]
cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[b]
cur_page_indptr_end: T.let[T.int32] = page_table_indptr[b + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
@@ -140,9 +140,9 @@ def _attention_decode_cpu(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, slidi
)
for row_idx in T.serial(kv_chunk_len[0]):
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(row_idx, b, length_info, sliding_window)
page_no: T.int32(is_size_var=True) = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
page_offset: T.int32(is_size_var=True) = seq_offset % page_size
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(row_idx, b, length_info, sliding_window)
page_no: T.let[T.int32(is_size_var=True)] = page_table_values[cur_page_indptr_begin + (seq_offset // page_size)]
page_offset: T.let[T.int32(is_size_var=True)] = seq_offset % page_size
for d in T.serial(D):
K_local[d] = T.if_then_else(
@@ -211,7 +211,7 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w
global_symbol += "_sliding_window"
# pylint: disable=too-many-branches
@T.prim_func
@T.prim_func(s_tir=True)
def batch_decode_paged_kv(
Q_handle: T.handle,
pages_handle: T.handle,
@@ -277,11 +277,11 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w
st_d = T.sblock_alloc_buffer((1,), "float32", scope="local")
O_local = T.sblock_alloc_buffer((VEC_SIZE,), "float32", scope="local")
by: T.int32 = fused_by_bz % H_kv
bz: T.int32 = fused_by_bz // H_kv
batch_idx: T.int32 = bx
cur_page_indptr_begin: T.int32 = page_table_indptr[batch_idx]
cur_page_indptr_end: T.int32 = page_table_indptr[batch_idx + 1]
by: T.let[T.int32] = fused_by_bz % H_kv
bz: T.let[T.int32] = fused_by_bz // H_kv
batch_idx: T.let[T.int32] = bx
cur_page_indptr_begin: T.let[T.int32] = page_table_indptr[batch_idx]
cur_page_indptr_end: T.let[T.int32] = page_table_indptr[batch_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, batch_idx, length_info, sliding_window),
@@ -303,18 +303,18 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w
)
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_size_per_bdx * bdy * bdz)):
tile_start_s: T.int32(is_size_var=True) = (tz * bdy + ty) * tile_size_per_bdx # type: ignore
tile_start_g: T.int32(is_size_var=True) = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx # type: ignore
tile_start_s: T.let[T.int32(is_size_var=True)] = (tz * bdy + ty) * tile_size_per_bdx # type: ignore
tile_start_g: T.let[T.int32(is_size_var=True)] = ((iterator * bdz + tz) * bdy + ty) * tile_size_per_bdx # type: ignore
# load KV from global memory to shared memory
for j in T.serial(tile_size_per_bdx):
with T.sblock("KV_load"):
T.reads()
T.writes()
row_g: T.int32(is_size_var=True) = tile_start_g + j # type: ignore
row_g: T.let[T.int32(is_size_var=True)] = tile_start_g + j # type: ignore
if row_g < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(row_g, batch_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, page_size) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(row_g, batch_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_table_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, page_size) # type: ignore
for vec in T.vectorized(VEC_SIZE):
K_smem[tile_start_s + j, tx * VEC_SIZE + vec] = T.if_then_else(
rotary_mode == 1,
@@ -354,7 +354,7 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w
st_m[0] = T.max(st_m[0], S_local[j])
# update st_d, st_O
o_scale: T.float32 = T.exp2(m_prev[0] - st_m[0])
o_scale: T.let[T.float32] = T.exp2(m_prev[0] - st_m[0])
st_d[0] *= o_scale
for j in T.serial(bdy * tile_size_per_bdx):
S_local[j] = T.exp2(S_local[j] - st_m[0])
@@ -412,7 +412,7 @@ def _attention_decode(num_kv_heads, num_qo_heads, head_dim, qkv_dtype, sliding_w
def _merge_state_inplace_cpu(v_dtype):
@T.prim_func
@T.prim_func(s_tir=True)
def merge_state_inplace_cpu(
v: T.handle,
s: T.handle,
@@ -463,7 +463,7 @@ def _merge_state_inplace(num_heads, head_dim, v_dtype, target: Target, global_sy
gdy = num_heads // bdy
check_thread_limits(target, bdx=bdx, bdy=bdy, bdz=1, gdz=1)
@T.prim_func
@T.prim_func(s_tir=True)
def merge_state_inplace(
v: T.handle,
s: T.handle,
python/tvm/relax/frontend/nn/llm/_kernel_common.py
+27 -27
View File
@@ -215,7 +215,7 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
m_smem: T.Buffer, d_smem: T.Buffer, O_local: T.Buffer, ty: T.int32, tx: T.int32,
):
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
m_smem[row] = -5e4
d_smem[row] = 1.0
@@ -252,31 +252,31 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
):
# Phase 1: compute m_new = max(masked S over kv tile), d_new = d_prev * exp2(m_prev - m_new)
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update1"):
m_prev[i] = m_smem[row]
m_new[i] = m_smem[row]
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
for j in T.serial(tile_z):
if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
m_new[i] = T.max(m_new[i], S_smem[row, j])
d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
# Phase 2: exp-and-scale S_smem; masked-out entries use -inf
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
with T.sblock("update"):
for j in T.serial(tile_z):
# predicate sits inside loop so sync stays outside conditional branches
if row < tile_x:
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
if _causal_mask(causal, row=row_, col=L_kv_start + j, kv_len=kv_len, qo_len=qo_len):
S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
else:
S_smem[row, j] = T.exp2(-5e4 - m_new[i])
# Phase 3: d_new += sum(S_smem[row, :]); write m/d/m_prev back to smem
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update"):
for j in T.serial(tile_z):
@@ -312,15 +312,15 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
for li, lj in T.grid(tile_x, tile_o):
with T.sblock("O_store"):
i, j = T.axis.remap("SS", [li, lj])
cur_L: T.int32 = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
output[cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
for li in T.grid(tile_x):
with T.sblock("lse_store"):
i = T.axis.remap("S", [li])
cur_L: T.int32 = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
lse[cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
@@ -338,7 +338,7 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
tile_id[0] -= batch_tiles[0]
batch_idx[0] += 1
if batch_idx[0] < batch_size:
b_idx: T.int32 = batch_idx[0]
b_idx: T.let[T.int32] = batch_idx[0]
batch_rows[0] = (q_indptr[b_idx + 1] - q_indptr[b_idx]) * group_size
batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
@@ -352,28 +352,28 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
# Same three-phase online softmax as softmax_update_causal but with a
# per-batch right-padding mask in place of causal masking.
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update1"):
m_prev[i] = m_smem[row]
m_new[i] = m_smem[row]
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
for j in T.serial(tile_z):
if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
m_new[i] = T.max(m_new[i], S_smem[row, j])
d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
with T.sblock("update"):
for j in T.serial(tile_z):
if row < tile_x:
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
if tirx.And(tirx.And(row_ < qo_len, row_ < valid_len), L_kv_start + j < valid_len):
S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
else:
S_smem[row, j] = T.exp2(-5e4 - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update"):
for j in T.serial(tile_z):
@@ -395,34 +395,34 @@ def _make_prefill_macros(tile_x, tile_y, tile_z, tile_o, bdx, num_warps, group_s
# [kv_len - valid_len, kv_len). Causal keeps
# col <= row + (kv_len - qo_len) within those valid suffixes.
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update1"):
m_prev[i] = m_smem[row]
m_new[i] = m_smem[row]
row_: T.int32 = (LH_start + row) // group_size
pad_q: T.int32 = qo_len - valid_len
pad_kv: T.int32 = kv_len - valid_len
row_: T.let[T.int32] = (LH_start + row) // group_size
pad_q: T.let[T.int32] = qo_len - valid_len
pad_kv: T.let[T.int32] = kv_len - valid_len
for j in T.serial(tile_z):
col_: T.int32 = L_kv_start + j
col_: T.let[T.int32] = L_kv_start + j
if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
m_new[i] = T.max(m_new[i], S_smem[row, j])
d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
with T.sblock("update"):
for j in T.serial(tile_z):
if row < tile_x:
row_: T.int32 = (LH_start + row) // group_size
pad_q: T.int32 = qo_len - valid_len
pad_kv: T.int32 = kv_len - valid_len
col_: T.int32 = L_kv_start + j
row_: T.let[T.int32] = (LH_start + row) // group_size
pad_q: T.let[T.int32] = qo_len - valid_len
pad_kv: T.let[T.int32] = kv_len - valid_len
col_: T.let[T.int32] = L_kv_start + j
if tirx.And(tirx.And(row_ < qo_len, row_ >= pad_q), tirx.And(col_ >= pad_kv, col_ < kv_len - qo_len + row_ + 1)):
S_smem[row, j] = T.exp2(S_smem[row, j] - m_new[i])
else:
S_smem[row, j] = T.exp2(-5e4 - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update"):
for j in T.serial(tile_z):
python/tvm/relax/frontend/nn/llm/_page_kernels.py
+9 -9
View File
@@ -40,7 +40,7 @@ from ._kernel_common import get_max_num_threads_per_block
def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size: int = 16):
"""Return the TIR function that appends new k/v data to PagedKVCache."""
@T.prim_func
@T.prim_func(s_tir=True)
def tir_kv_cache_transpose_append(
var_pages: T.handle,
var_k_data: T.handle,
@@ -77,7 +77,7 @@ def _kv_cache_transpose_append(num_key_value_heads, head_dim, dtype, page_size:
def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
"""Return the TIR function that appends new compressed KV data to PagedKVCache for MLA."""
@T.prim_func
@T.prim_func(s_tir=True)
def tir_kv_cache_transpose_append_mla(
var_pages: T.handle,
var_kv_data: T.handle,
@@ -106,7 +106,7 @@ def _kv_cache_transpose_append_mla(d_qk: int, dtype, page_size: int = 16):
def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dtype):
"""Return the TIR function that fetches the k/v data on given positions and layer."""
@T.prim_func
@T.prim_func(s_tir=True)
def tir_kv_cache_debug_get_kv(
var_pages: T.handle,
var_position_map: T.handle,
@@ -139,7 +139,7 @@ def _kv_cache_debug_get_kv(num_hidden_layers, num_key_value_heads, head_dim, dty
def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
"""Return the TIR function that fetches the k/v data on given positions and layer."""
@T.prim_func
@T.prim_func(s_tir=True)
def tir_kv_cache_debug_get_kv_mla(
var_pages: T.handle,
var_position_map: T.handle,
@@ -169,7 +169,7 @@ def _kv_cache_debug_get_kv_mla(num_hidden_layers, d_qk, dtype):
def _copy_single_page(num_heads, page_size, head_dim, dtype, target: Target):
tx = get_max_num_threads_per_block(target)
@T.prim_func
@T.prim_func(s_tir=True)
def copy_single_page(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
T.func_attr({"tirx.is_scheduled": True})
num_pages = T.int32()
@@ -192,7 +192,7 @@ def _copy_single_page(num_heads, page_size, head_dim, dtype, target: Target):
def _copy_single_page_mla(page_size, head_dim, dtype, target: Target):
tx = get_max_num_threads_per_block(target)
@T.prim_func
@T.prim_func(s_tir=True)
def copy_single_page_mla(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
T.func_attr({"tirx.is_scheduled": True})
num_pages = T.int32()
@@ -213,7 +213,7 @@ def _copy_single_page_mla(page_size, head_dim, dtype, target: Target):
def _copy_single_page_cpu(num_heads, page_size, head_dim, dtype):
tx = 1
@T.prim_func
@T.prim_func(s_tir=True)
def copy_single_page_cpu(var_pages: T.handle, src_page_id: T.int64, tgt_page_id: T.int64, copy_length: T.int64):
T.func_attr({"tirx.is_scheduled": True})
num_pages = T.int32()
@@ -235,7 +235,7 @@ def _copy_single_page_cpu(num_heads, page_size, head_dim, dtype):
def _compact_kv_copy(num_heads, head_dim, dtype, target: Target, page_size: int = 16):
tx = get_max_num_threads_per_block(target)
@T.prim_func
@T.prim_func(s_tir=True)
def compact_kv_copy(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
T.func_attr({"tirx.is_scheduled": True})
num_pages = T.int32()
@@ -266,7 +266,7 @@ def _compact_kv_copy(num_heads, head_dim, dtype, target: Target, page_size: int
def _compact_kv_copy_cpu(num_heads, head_dim, dtype, page_size: int = 16):
tx = 8
@T.prim_func
@T.prim_func(s_tir=True)
def compact_kv_copy_cpu(var_pages: T.handle, var_copy_length_indptr: T.handle, var_copy_src_dst_pos: T.handle, batch_size: T.int32):
T.func_attr({"tirx.is_scheduled": True})
num_pages = T.int32()
python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+85 -87
View File
@@ -60,7 +60,7 @@ def _attention_prefill_cpu(
group_size = h_q // h_kv
# pylint: disable=too-many-branches
@T.prim_func
@T.prim_func(s_tir=True)
def batch_prefill_paged_kv_cpu(
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -126,9 +126,9 @@ def _attention_prefill_cpu(
S_val = T.sblock_alloc_buffer((1, ), "float32")
scale_O = T.sblock_alloc_buffer((1, ), "float32")
factor = T.sblock_alloc_buffer((1, ), "float32")
cur_page_indptr_begin: T.int32 = page_indptr[b_idx]
cur_page_indptr_end: T.int32 = page_indptr[b_idx + 1]
#max_kv_len: T.int32 = max_num_pages * page_size
cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
cur_page_indptr_end: T.let[T.int32] = page_indptr[b_idx + 1]
#max_kv_len: T.let[T.int32] = max_num_pages * page_size
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b_idx, length_info, sliding_window),
@@ -142,7 +142,7 @@ def _attention_prefill_cpu(
d_val[0] = 1.0
for d_idx in T.serial(d):
O_local[d_idx] = 0.0
curl_q: T.int32 = q_indptr[b_idx] + q_idx
curl_q: T.let[T.int32] = q_indptr[b_idx] + q_idx
for d_idx in T.serial(d):
@@ -153,10 +153,10 @@ def _attention_prefill_cpu(
)
for row_idx in T.serial(max_num_pages * page_size):
if row_idx < kv_chunk_len[0]:
# seq_offset: T.int32(is_size_var=True) = _get_seq_offset(row_idx, b_idx, length_info, sliding_window)
#seq_offset: T.int32(is_size_var=True) = row_idx
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + (_get_seq_offset(row_idx, b_idx, length_info, sliding_window) // page_size)]
page_offset: T.int32(is_size_var=True) = _get_seq_offset(row_idx, b_idx, length_info, sliding_window) % page_size
# seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(row_idx, b_idx, length_info, sliding_window)
#seq_offset: T.let[T.int32(is_size_var=True)] = row_idx
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + (_get_seq_offset(row_idx, b_idx, length_info, sliding_window) // page_size)]
page_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(row_idx, b_idx, length_info, sliding_window) % page_size
# Load KV
for d_idx in T.serial(d):
@@ -215,7 +215,7 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling:
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, _, advance_tile_batch, paged_store_output_lse, *_ = _make_prefill_macros(tile_x, tile_y, tile_z, tile_y, bdx, num_warps, group_size)
# pylint: disable=too-many-branches
@T.prim_func
@T.prim_func(s_tir=True)
def batch_prefill_paged_kv(
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -288,12 +288,12 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling:
advance_tile_batch(tile_id, batch_idx, batch_tiles, batch_rows, q_indptr, batch_size)
if T.tvm_thread_invariant(batch_idx[0] < batch_size):
b_idx: T.int32 = batch_idx[0]
LH_start: T.int32 = tile_id[0] * tile_x
q_indptr_val: T.int32 = q_indptr[b_idx]
b_idx: T.let[T.int32] = batch_idx[0]
LH_start: T.let[T.int32] = tile_id[0] * tile_x
q_indptr_val: T.let[T.int32] = q_indptr[b_idx]
cur_page_indptr_begin: T.int32 = page_indptr[b_idx]
cur_page_indptr_end: T.int32 = page_indptr[b_idx + 1]
cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
cur_page_indptr_end: T.let[T.int32] = page_indptr[b_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b_idx, length_info, sliding_window),
@@ -309,8 +309,8 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling:
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = q_indptr_val + (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr_val + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -322,17 +322,17 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling:
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_start: T.let[T.int32] = iterator * tile_z
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("K_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, page_size) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, page_size) # type: ignore
K_smem[i, j] = T.if_then_else(
rotary_mode == 1,
_rope(pages, k_rope_pos_offset[b_idx] + cur_L, d, rope_theta, rope_scale, (page_no, 0, by, page_offset, j), dtype, rope_scaling),
@@ -346,11 +346,11 @@ def _attention_prefill(h_kv, h_q, d, dtype, sliding_window: bool, rope_scaling:
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, page_size) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, page_size) # type: ignore
V_smem[i, j] = pages[page_no, 1, by, page_offset, j]
else:
V_smem[i, j] = 0.0
@@ -377,7 +377,7 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
_, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z = _get_prefill_kernel_config(h_kv, h_q, d, dtype, target)
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, *_ = _make_prefill_macros(tile_x, tile_y, tile_z, tile_y, bdx, num_warps, group_size)
@T.prim_func
@T.prim_func(s_tir=True)
def batch_sequence_prefill_kv( # pylint: disable=too-many-branches
var_q: T.handle, # [total_len, h_q, d]
var_k: T.handle, # [total_len, h_kv, d]
@@ -394,7 +394,7 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
output = T.match_buffer(var_output, (batch_size, qo_len, h_q, d), dtype)
lse = T.match_buffer(var_lse, (batch_size, qo_len, h_q), dtype) # pylint: disable=unused-variable
batch_tiles: T.int32 = T.ceildiv(qo_len * group_size, tile_x)
batch_tiles: T.let[T.int32] = T.ceildiv(qo_len * group_size, tile_x)
# kernel code
for lbx in T.thread_binding(T.cast(batch_size, "int32") * batch_tiles, thread="blockIdx.x"):
@@ -411,9 +411,9 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
_alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps)
)
b_idx: T.int32 = vbx // batch_tiles
tile_id: T.int32 = vbx % batch_tiles
LH_start: T.int32 = tile_id * tile_x
b_idx: T.let[T.int32] = vbx // batch_tiles
tile_id: T.let[T.int32] = vbx % batch_tiles
LH_start: T.let[T.int32] = tile_id * tile_x
T.tvm_storage_sync("shared")
init_states(m_smem, d_smem, O_local, ty, tx)
@@ -424,8 +424,8 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < qo_len:
Q_smem[i, j] = q[b_idx, cur_L, cur_H_qo, j]
else:
@@ -433,14 +433,14 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_len, tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_base: T.int32 = 0
L_kv_start: T.let[T.int32] = iterator * tile_z
L_kv_base: T.let[T.int32] = 0
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("K_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_len:
K_smem[i, j] = k[
b_idx, L_kv_base + cur_L, by, j
@@ -453,7 +453,7 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_len:
V_smem[i, j] = v[b_idx, L_kv_base + cur_L, by, j]
else:
@@ -468,8 +468,8 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
for li, lj in T.grid(tile_x, tile_y):
with T.sblock("O_store"):
i, j = T.axis.remap("SS", [li, lj])
cur_L: T.int32 = 0 + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = 0 + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < qo_len:
output[b_idx, cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
@@ -477,8 +477,8 @@ def _attention_sequence_prefill(h_kv, h_q, d, dtype, target: Target, causal=0, s
for li in T.grid(tile_x):
with T.sblock("lse_store"):
i = T.axis.remap("S", [li])
cur_L: T.int32 = 0 + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = 0 + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < qo_len:
lse[b_idx, cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
@@ -544,7 +544,7 @@ def _attention_sequence_prefill_with_mask(
pad = kv_len - valid_len
return tirx.And(col < kv_len, col >= pad)
@T.prim_func
@T.prim_func(s_tir=True)
def batch_sequence_prefill_kv_masked( # pylint: disable=too-many-branches
var_q: T.handle, # [batch_size, qo_len, h_q, d]
var_k: T.handle, # [batch_size, kv_len, h_kv, d]
@@ -563,7 +563,7 @@ def _attention_sequence_prefill_with_mask(
output = T.match_buffer(var_output, (batch_size, qo_len, h_q, d), dtype)
lse = T.match_buffer(var_lse, (batch_size, qo_len, h_q), dtype)
batch_tiles: T.int32 = T.ceildiv(qo_len * group_size, tile_x)
batch_tiles: T.let[T.int32] = T.ceildiv(qo_len * group_size, tile_x)
for lbx in T.thread_binding(T.cast(batch_size, "int32") * batch_tiles, thread="blockIdx.x"):
for lby in T.thread_binding(h_kv, thread="blockIdx.y"):
@@ -579,10 +579,10 @@ def _attention_sequence_prefill_with_mask(
_alloc_softmax_state_buffers(tile_x, tile_z, bdx, num_warps)
)
b_idx: T.int32 = vbx // batch_tiles
valid_len: T.int32 = valid_lens[b_idx]
tile_id: T.int32 = vbx % batch_tiles
LH_start: T.int32 = tile_id * tile_x
b_idx: T.let[T.int32] = vbx // batch_tiles
valid_len: T.let[T.int32] = valid_lens[b_idx]
tile_id: T.let[T.int32] = vbx % batch_tiles
LH_start: T.let[T.int32] = tile_id * tile_x
T.tvm_storage_sync("shared")
init_states(m_smem, d_smem, O_local, ty, tx)
@@ -593,8 +593,8 @@ def _attention_sequence_prefill_with_mask(
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if _q_row_valid(cur_L, valid_len, qo_len):
Q_smem[i, j] = q[b_idx, cur_L, cur_H_qo, j]
else:
@@ -602,14 +602,14 @@ def _attention_sequence_prefill_with_mask(
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_len, tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_base: T.int32 = 0
L_kv_start: T.let[T.int32] = iterator * tile_z
L_kv_base: T.let[T.int32] = 0
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("K_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if _kv_col_valid(cur_L, valid_len, kv_len):
K_smem[i, j] = k[b_idx, L_kv_base + cur_L, by, j]
else:
@@ -620,7 +620,7 @@ def _attention_sequence_prefill_with_mask(
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if _kv_col_valid(cur_L, valid_len, kv_len):
V_smem[i, j] = v[b_idx, L_kv_base + cur_L, by, j]
else:
@@ -635,8 +635,8 @@ def _attention_sequence_prefill_with_mask(
for li, lj in T.grid(tile_x, tile_y):
with T.sblock("O_store"):
i, j = T.axis.remap("SS", [li, lj])
cur_L: T.int32 = 0 + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = 0 + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < qo_len:
output[b_idx, cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
@@ -644,8 +644,8 @@ def _attention_sequence_prefill_with_mask(
for li in T.grid(tile_x):
with T.sblock("lse_store"):
i = T.axis.remap("S", [li])
cur_L: T.int32 = 0 + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = 0 + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < qo_len:
lse[b_idx, cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
@@ -658,7 +658,7 @@ def _attention_sequence_prefill_with_mask(
def _attention_prefill_ragged_cpu(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[str, Any]):
group_size = h_q // h_kv
@T.prim_func
@T.prim_func(s_tir=True)
def batch_prefill_ragged_kv( # pylint: disable=too-many-branches
var_q: T.handle, # [total_len, h_q, d_qk]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -717,7 +717,7 @@ def _attention_prefill_ragged_cpu(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dic
for k_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
for h in T.serial(h_q):
h_kv_idx = h // group_size
h_kv_idx: T.let[T.int32] = h // group_size
if _causal_mask(
causal,
@@ -757,20 +757,18 @@ def _attention_prefill_ragged_cpu(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dic
exp_scores[k_idx, h] = T.exp2(attention_scores[k_idx, h] - m_new[h])
softmax_sum[h] += exp_scores[k_idx, h]
d_new[h] += softmax_sum[h]
d_prev = d_new
m_prev = m_new
for h in T.serial(h_q):
h_kv_idx = h // group_size
h_kv_idx: T.let[T.int32] = h // group_size
for i in T.serial(d_v):
p_sum[i] = 0.0
for v_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
weight = exp_scores[v_idx, h] / d_new[h]
weight: T.let[T.float32] = exp_scores[v_idx, h] / d_new[h]
for i in T.serial(d_v):
p_sum[i] += v[kv_indptr[b] + v_idx, h_kv_idx, i] * weight
for i in T.serial(d_v):
output[q_indptr[b] + q_idx, h, i] = p_sum[i]
lse[q_indptr[b] + q_idx, h] = m_prev[h] + T.log2(d_prev[h])
lse[q_indptr[b] + q_idx, h] = m_new[h] + T.log2(d_new[h])
return batch_prefill_ragged_kv
@@ -779,7 +777,7 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[st
NUM_BLKS, LOAD_VEC, group_size, bdx, num_warps, tile_x, tile_y, tile_z = _get_prefill_kernel_config(h_kv, h_q, d_qk, dtype, target)
init_states, compute_s_gemm, softmax_update_causal, compute_o_gemm, _, advance_tile_batch, paged_store_output_lse, *_ = _make_prefill_macros(tile_x, tile_y, tile_z, d_v, bdx, num_warps, group_size)
@T.prim_func
@T.prim_func(s_tir=True)
def batch_prefill_ragged_kv( # pylint: disable=too-many-branches
var_q: T.handle, # [total_len, h_q, d_qk]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -837,9 +835,9 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[st
advance_tile_batch(tile_id, batch_idx, batch_tiles, batch_rows, q_indptr, batch_size)
if T.tvm_thread_invariant(batch_idx[0] < batch_size):
b_idx: T.int32 = batch_idx[0]
q_indptr_val: T.int32 = q_indptr[b_idx]
LH_start: T.int32 = tile_id[0] * tile_x
b_idx: T.let[T.int32] = batch_idx[0]
q_indptr_val: T.let[T.int32] = q_indptr[b_idx]
LH_start: T.let[T.int32] = tile_id[0] * tile_x
kv_chunk_len[0] = kv_indptr[b_idx + 1] - kv_indptr[b_idx]
T.tvm_storage_sync("shared")
@@ -852,8 +850,8 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[st
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = q_indptr_val + (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr_val + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -865,12 +863,12 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[st
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_base: T.int32 = kv_indptr[b_idx]
L_kv_start: T.let[T.int32] = iterator * tile_z
L_kv_base: T.let[T.int32] = kv_indptr[b_idx]
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("K_load"):
i, j = T.axis.remap("SS", [lz, ly])
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
K_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -885,7 +883,7 @@ def _attention_prefill_ragged(h_kv, h_q, d_qk, d_v, dtype, rope_scaling: dict[st
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
V_smem[i, j] = v[L_kv_base + cur_L, by, j]
else:
@@ -917,7 +915,7 @@ def _attention_prefill_mla(h_q, d_latent, d_rope, dtype, sliding_window: bool, t
global_symbol += "_sliding_window"
# pylint: disable=too-many-branches
@T.prim_func
@T.prim_func(s_tir=True)
def batch_prefill_paged_kv_mla(
var_q: T.handle, # [total_len, h_q, d_qk]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -980,12 +978,12 @@ def _attention_prefill_mla(h_q, d_latent, d_rope, dtype, sliding_window: bool, t
advance_tile_batch(tile_id, batch_idx, batch_tiles, batch_rows, q_indptr, batch_size)
if T.tvm_thread_invariant(batch_idx[0] < batch_size):
b_idx: T.int32 = batch_idx[0]
LH_start: T.int32 = tile_id[0] * tile_x
q_indptr_val: T.int32 = q_indptr[b_idx]
b_idx: T.let[T.int32] = batch_idx[0]
LH_start: T.let[T.int32] = tile_id[0] * tile_x
q_indptr_val: T.let[T.int32] = q_indptr[b_idx]
cur_page_indptr_begin: T.int32 = page_indptr[b_idx]
cur_page_indptr_end: T.int32 = page_indptr[b_idx + 1]
cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
cur_page_indptr_end: T.let[T.int32] = page_indptr[b_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, page_size, b_idx, length_info, sliding_window),
@@ -1001,8 +999,8 @@ def _attention_prefill_mla(h_q, d_latent, d_rope, dtype, sliding_window: bool, t
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = q_indptr_val + (LH_start + i) // group_size
cur_H_qo = (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr_val + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = q[cur_L, cur_H_qo, j]
else:
@@ -1010,17 +1008,17 @@ def _attention_prefill_mla(h_q, d_latent, d_rope, dtype, sliding_window: bool, t
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_start: T.let[T.int32] = iterator * tile_z
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("KV_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, page_size) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, page_size)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, page_size) # type: ignore
KV_smem[i, j] = pages[page_no, page_offset, j]
else:
KV_smem[i, j] = 0.0
python/tvm/relax/frontend/nn/llm/position_embedding.py
+5 -5
View File
@@ -390,7 +390,7 @@ def llama_rope( # pylint: disable=too-many-arguments
expr = tirx.Let(var, value, expr)
return expr
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def fused_rope( # pylint: disable=too-many-locals
var_qkv: T.handle,
var_q: T.handle,
@@ -522,7 +522,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments
expr = tirx.Let(var, value, expr)
return expr
@T.prim_func
@T.prim_func(s_tir=True)
def fused_rope( # pylint: disable=too-many-locals
var_qkv: T.handle,
var_position_map: T.handle,
@@ -564,7 +564,7 @@ def llama_rope_with_position_map( # pylint: disable=too-many-arguments
else:
v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
@T.prim_func
@T.prim_func(s_tir=True)
def fused_rope_longrope_scaling( # pylint: disable=too-many-locals
var_qkv: T.handle,
var_position_map: T.handle,
@@ -749,7 +749,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments
expr = tirx.Let(var, value, expr)
return expr
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def fused_rope( # pylint: disable=too-many-locals
var_qkv: T.handle,
var_position_map: T.handle,
@@ -791,7 +791,7 @@ def llama4_rope_with_position_map( # pylint: disable=too-many-arguments
else:
v[s, h - (num_q_heads + num_kv_heads), d] = qkv[s, h, d]
@T.prim_func
@T.prim_func(s_tir=True)
def fused_rope_longrope_scaling( # pylint: disable=too-many-locals
var_qkv: T.handle,
var_position_map: T.handle,
python/tvm/relax/frontend/nn/llm/tree_attn.py
+59 -61
View File
@@ -89,7 +89,7 @@ def tree_attn_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
group_size = h_q // h_kv
# fmt: off
@T.prim_func
@T.prim_func(s_tir=True)
def batch_tree_attn( # pylint: disable=too-many-branches,line-too-long
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -181,7 +181,7 @@ def tree_attn_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
for k_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
for h in T.serial(h_q):
h_kv_idx = h // group_size
h_kv_idx: T.let[T.int32] = h // group_size
if _check_tree_order(
row=q_idx,
@@ -243,20 +243,18 @@ def tree_attn_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any]):
exp_scores[k_idx, h] = T.exp2(attention_scores[k_idx, h] - m_new[h])
softmax_sum[h] += exp_scores[k_idx, h]
d_new[h] += softmax_sum[h]
d_prev = d_new
m_prev = m_new
for h in T.serial(h_q):
h_kv_idx = h // group_size
h_kv_idx: T.let[T.int32] = h // group_size
for i in T.serial(d):
p_sum[i] = 0.0
for v_idx in T.serial(kv_indptr[b + 1] - kv_indptr[b]):
weight = exp_scores[v_idx, h] / d_new[h]
weight: T.let[T.float32] = exp_scores[v_idx, h] / d_new[h]
for i in T.serial(d):
p_sum[i] += v[kv_indptr[b] + v_idx, h_kv_idx, i] * weight
for i in T.serial(d):
output[q_indptr[b] + q_idx, h, i] = p_sum[i]
lse[q_indptr[b] + q_idx, h] = m_prev[h] + T.log2(d_prev[h])
lse[q_indptr[b] + q_idx, h] = m_new[h] + T.log2(d_new[h])
# fmt: on
# pylint: enable=line-too-long,too-many-branches
@@ -312,7 +310,7 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
num_warps = 2
# fmt: off
@T.prim_func
@T.prim_func(s_tir=True)
def batch_tree_attn( # pylint: disable=too-many-branches
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -373,21 +371,21 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
tile_id[0] -= batch_tiles[0]
batch_idx[0] += 1
if batch_idx[0] < batch_size_plus_1 - 1:
b_idx: T.int32 = batch_idx[0]
b_idx: T.let[T.int32] = batch_idx[0]
batch_rows[0] = (q_indptr[b_idx + 1] - q_indptr[b_idx]) * group_size
batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
if T.tvm_thread_invariant(batch_idx[0] < batch_size_plus_1 - 1):
b_idx: T.int32(is_size_var=True) = batch_idx[0]
LH_start: T.int32(is_size_var=True) = tile_id[0] * tile_x
q_indptr_val: T.int32 = q_indptr[b_idx]
b_idx: T.let[T.int32(is_size_var=True)] = batch_idx[0]
LH_start: T.let[T.int32(is_size_var=True)] = tile_id[0] * tile_x
q_indptr_val: T.let[T.int32] = q_indptr[b_idx]
kv_chunk_len[0] = kv_indptr[b_idx + 1] - kv_indptr[b_idx]
T.tvm_storage_sync("shared")
# init states
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
m_smem[row] = -5e4
d_smem[row] = 1.0
@@ -404,8 +402,8 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = q_indptr_val + (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr_val + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -417,14 +415,14 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_base: T.int32 = kv_indptr[b_idx]
L_kv_start: T.let[T.int32] = iterator * tile_z
L_kv_base: T.let[T.int32] = kv_indptr[b_idx]
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("KV_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_base + L_kv_start + i
cur_L: T.let[T.int32] = L_kv_base + L_kv_start + i
if L_kv_start + i < kv_chunk_len[0]:
K_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -454,13 +452,13 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
# Update S, m, d
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update1"):
m_prev[i] = m_smem[row]
m_new[i] = m_smem[row]
# mask out of kv_chunk_len S
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
for j in T.serial(tile_z):
if _check_tree_order(
row=row_,
@@ -474,12 +472,12 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
d_new[i] = d_smem[row] * T.exp2(m_prev[i] - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
with T.sblock("update"):
for j in T.serial(tile_z):
# this is to avoid sync inside condition branch
if row < tile_x:
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
if _check_tree_order(
row=row_,
col=L_kv_start + j,
@@ -493,7 +491,7 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
S_smem[row, j] = T.exp2(-5e4 - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update"):
for j in T.serial(tile_z):
@@ -516,8 +514,8 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
for li, lj in T.grid(tile_x, tile_y):
with T.sblock("O_store"):
i, j = T.axis.remap("SS", [li, lj])
cur_L: T.int32 = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
output[cur_L, cur_H_qo, j] = O_local[i, j] / d_smem[i]
@@ -525,8 +523,8 @@ def tree_attn(h_kv, h_q, d, dtype, rope_scaling: dict[str, Any], target: Target)
for li in T.grid(tile_x):
with T.sblock("lse_store"):
i = T.axis.remap("S", [li])
cur_L: T.int32 = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.int32 = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr[b_idx] + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
lse[cur_L, cur_H_qo] = m_smem[i] + T.log2(d_smem[i])
@@ -632,7 +630,7 @@ def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[st
# pylint: disable=line-too-long,too-many-branches
# fmt: off
@T.prim_func(check_well_formed=False)
@T.prim_func(s_tir=True)
def tree_attn_paged_kv_cpu(
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -720,8 +718,8 @@ def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[st
S_val = T.sblock_alloc_buffer((1, ), "float32")
scale_O = T.sblock_alloc_buffer((1, ), "float32")
factor = T.sblock_alloc_buffer((1, ), "float32")
cur_page_indptr_begin: T.int32 = page_indptr[b_idx]
cur_page_indptr_end: T.int32 = page_indptr[b_idx + 1]
cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
cur_page_indptr_end: T.let[T.int32] = page_indptr[b_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(cur_page_indptr_end - cur_page_indptr_begin, 16, b_idx, length_info, sliding_window),
@@ -734,7 +732,7 @@ def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[st
d_val[0] = 1.0
for d_idx in T.serial(d):
O_local[d_idx] = 0.0
curl_q: T.int32 = q_indptr[b_idx] + q_idx
curl_q: T.let[T.int32] = q_indptr[b_idx] + q_idx
for d_idx in T.serial(d):
Q_local[d_idx] = T.if_then_else(
@@ -744,8 +742,8 @@ def tree_attn_with_paged_kv_cache_cpu(h_kv, h_q, d, dtype, rope_scaling: dict[st
)
for row_idx in T.serial(max_num_pages * 16):
if row_idx < kv_chunk_len[0]:
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + (_get_seq_offset(row_idx, b_idx, length_info, sliding_window) // 16)]
page_offset: T.int32(is_size_var=True) = _get_seq_offset(row_idx, b_idx, length_info, sliding_window) % 16
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + (_get_seq_offset(row_idx, b_idx, length_info, sliding_window) // 16)]
page_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(row_idx, b_idx, length_info, sliding_window) % 16
# Load KV
for d_idx in T.serial(d):
@@ -852,7 +850,7 @@ def tree_attn_with_paged_kv_cache(
sliding_window = False # Sliding window is not supported in this kernel.
# fmt: off
@T.prim_func
@T.prim_func(s_tir=True)
def tree_attn_paged_kv(
var_q: T.handle, # [total_len, h_q, d]
var_q_indptr: T.handle, # [batch_size + 1]
@@ -959,19 +957,19 @@ def tree_attn_with_paged_kv_cache(
tile_id[0] -= batch_tiles[0]
batch_idx[0] += 1
if batch_idx[0] < batch_size:
b_idx: T.int32 = batch_idx[0]
b_idx: T.let[T.int32] = batch_idx[0]
batch_rows[0] = (
q_indptr[b_idx + 1] - q_indptr[b_idx]
) * group_size
batch_tiles[0] = T.ceildiv(batch_rows[0], tile_x)
if T.tvm_thread_invariant(batch_idx[0] < batch_size):
b_idx: T.int32(is_size_var=True) = batch_idx[0]
LH_start: T.int32(is_size_var=True) = tile_id[0] * tile_x
q_indptr_val: T.int32 = q_indptr[b_idx]
b_idx: T.let[T.int32(is_size_var=True)] = batch_idx[0]
LH_start: T.let[T.int32(is_size_var=True)] = tile_id[0] * tile_x
q_indptr_val: T.let[T.int32] = q_indptr[b_idx]
cur_page_indptr_begin: T.int32 = page_indptr[b_idx]
cur_page_indptr_end: T.int32 = page_indptr[b_idx + 1]
cur_page_indptr_begin: T.let[T.int32] = page_indptr[b_idx]
cur_page_indptr_end: T.let[T.int32] = page_indptr[b_idx + 1]
kv_chunk_len[0] = T.if_then_else(
cur_page_indptr_begin != cur_page_indptr_end,
_get_kv_chunk_len(
@@ -987,7 +985,7 @@ def tree_attn_with_paged_kv_cache(
# init states
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
m_smem[row] = -5e4
d_smem[row] = 1.0
@@ -1004,8 +1002,8 @@ def tree_attn_with_paged_kv_cache(
i, j = T.axis.remap("SS", [li, lj])
T.reads()
T.writes()
cur_L = q_indptr_val + (LH_start + i) // group_size
cur_H_qo = by * group_size + (LH_start + i) % group_size
cur_L: T.let[T.int32] = q_indptr_val + (LH_start + i) // group_size
cur_H_qo: T.let[T.int32] = by * group_size + (LH_start + i) % group_size
if cur_L < q_indptr[b_idx + 1]:
Q_smem[i, j] = T.if_then_else(
rotary_mode == 1,
@@ -1026,17 +1024,17 @@ def tree_attn_with_paged_kv_cache(
T.tvm_storage_sync("shared")
for iterator in T.serial(T.ceildiv(kv_chunk_len[0], tile_z)):
L_kv_start: T.int32 = iterator * tile_z
L_kv_start: T.let[T.int32] = iterator * tile_z
for lz, ly in T.grid(tile_z, tile_y):
with T.sblock("K_load"):
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, 16) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, 16) # type: ignore
K_smem[i, j] = pages[
page_no, 0, by, page_offset, j
]
@@ -1049,11 +1047,11 @@ def tree_attn_with_paged_kv_cache(
i, j = T.axis.remap("SS", [lz, ly])
T.reads()
T.writes()
cur_L = L_kv_start + i
cur_L: T.let[T.int32] = L_kv_start + i
if cur_L < kv_chunk_len[0]:
seq_offset: T.int32(is_size_var=True) = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.int32(is_size_var=True) = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)] # type: ignore
page_offset: T.int32(is_size_var=True) = T.floormod(seq_offset, 16) # type: ignore
seq_offset: T.let[T.int32(is_size_var=True)] = _get_seq_offset(cur_L, b_idx, length_info, sliding_window) # type: ignore
page_no: T.let[T.int32(is_size_var=True)] = page_values[cur_page_indptr_begin + T.floordiv(seq_offset, 16)] # type: ignore
page_offset: T.let[T.int32(is_size_var=True)] = T.floormod(seq_offset, 16) # type: ignore
V_smem[i, j] = pages[
page_no, 1, by, page_offset, j
]
@@ -1083,13 +1081,13 @@ def tree_attn_with_paged_kv_cache(
# Update S, m, d
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update1"):
m_prev[i] = m_smem[row]
m_new[i] = m_smem[row]
# mask out of kv_chunk_len S
row_: T.int32 = (LH_start + row) // group_size
row_: T.let[T.int32] = (LH_start + row) // group_size
for j in T.serial(tile_z):
if _check_tree_order(
tree_order_indptr=tree_order_indptr,
@@ -1109,12 +1107,12 @@ def tree_attn_with_paged_kv_cache(
)
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
with T.sblock("update"):
for j in T.serial(tile_z):
# this is to avoid sync inside condition branch
if row < tile_x:
row_: T.int32 = (
row_: T.let[T.int32] = (
LH_start + row
) // group_size
if _check_tree_order(
@@ -1134,7 +1132,7 @@ def tree_attn_with_paged_kv_cache(
S_smem[row, j] = T.exp2(-5e4 - m_new[i])
for i in T.serial(T.ceildiv(tile_x, bdx * num_warps)):
row: T.int32 = i * bdx * num_warps + ty * bdx + tx
row: T.let[T.int32] = i * bdx * num_warps + ty * bdx + tx
if row < tile_x:
with T.sblock("update"):
for j in T.serial(tile_z):
@@ -1161,10 +1159,10 @@ def tree_attn_with_paged_kv_cache(
for li, lj in T.grid(tile_x, tile_y):
with T.sblock("O_store"):
i, j = T.axis.remap("SS", [li, lj])
cur_L: T.int32 = (
cur_L: T.let[T.int32] = (
q_indptr[b_idx] + (LH_start + i) // group_size
)
cur_H_qo: T.int32 = (
cur_H_qo: T.let[T.int32] = (
by * group_size + (LH_start + i) % group_size
)
if cur_L < q_indptr[b_idx + 1]:
@@ -1176,10 +1174,10 @@ def tree_attn_with_paged_kv_cache(
for li in T.grid(tile_x):
with T.sblock("lse_store"):
i = T.axis.remap("S", [li])
cur_L: T.int32 = (
cur_L: T.let[T.int32] = (
q_indptr[b_idx] + (LH_start + i) // group_size
)
cur_H_qo: T.int32 = (
cur_H_qo: T.let[T.int32] = (
by * group_size + (LH_start + i) % group_size
)
if cur_L < q_indptr[b_idx + 1]:
python/tvm/relax/frontend/nn/op.py
+3 -3
View File
@@ -2796,7 +2796,7 @@ def sample_top_p_top_k_from_sorted_prob(
def _cumsum_mask(cumsum_sorted, top_p, top_k, i, j):
return _tir.all(cumsum_sorted[i, j] < top_p[i, 0], j + 1 < top_k[i, 0])
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_renorm_prob(A: T.handle, B: T.handle, C: T.handle, D: T.handle):
batch, vocab_size = T.int64(is_size_var=True), T.int64(is_size_var=True)
cumsum_sorted = T.match_buffer(A, (batch, vocab_size), prob_dtype)
@@ -2814,7 +2814,7 @@ def sample_top_p_top_k_from_sorted_prob(
elif not _cumsum_mask(cumsum_sorted, top_p, top_k, v_ax0, v_ax1 + 1):
renorm_prob[v_ax0, 0] = cumsum_sorted[v_ax0, v_ax1 + 1]
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_index_from_sorted(
A: T.handle, B: T.handle, C: T.handle, D: T.handle, E: T.handle, F: T.handle
):
@@ -2902,7 +2902,7 @@ def renormalize_top_p_top_k_prob(prob, sorted_prob, top_p, top_k):
def _cumsum_mask(cumsum_sorted, top_p, top_k, i, j):
return _tir.all(cumsum_sorted[i, j] < top_p[i, 0], j + 1 < top_k[i, 0])
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_renorm_cutoff(A: T.handle, B: T.handle, C: T.handle, D: T.handle, E: T.handle):
batch, vocab_size = T.int64(), T.int64()
sorted_prob = T.match_buffer(A, (batch, vocab_size), prob_dtype)
python/tvm/relax/frontend/onnx/onnx_frontend.py
+2 -1
View File
@@ -4245,7 +4245,8 @@ class TopK(OnnxOpConverter):
k = inputs[1]
if not isinstance(k, relax.Constant):
raise ValueError("TopK k must be a constant")
k = int(k.data.numpy())
# ONNX represents k as a tensor of shape [1]; flatten before scalar cast.
k = int(k.data.numpy().reshape(-1)[0])
axis = attr.get("axis", -1)
largest = attr.get("largest", 1)
sorted = attr.get("sorted", 1)
python/tvm/relax/training/optimizer.py
+6
View File
@@ -72,6 +72,7 @@ class Optimizer:
For detailed examples, please see the tutorial.
.. code-block:: python
# Construct the optimizer
opt = relax.optimizer.SGD(0.1)
@@ -195,6 +196,7 @@ class Optimizer:
gradient descent method with lr = 0.1.
.. code-block:: python
@R.function
def SGD(
params: R.Tuple(R.Tensor((3, 3), "float32"), R.Tensor((3,), "float32")),
@@ -245,6 +247,7 @@ class SGD(Optimizer):
The returned function of `get_function()` is equivalent to the following numpy code:
.. code-block:: python
def SGD(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
param_tuple_new, state_tuple_new = [], []
@@ -357,6 +360,7 @@ class MomentumSGD(Optimizer):
The returned function of `get_function()` is equivalent to the following numpy code:
.. code-block:: python
def MomentumSGD(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
param_tuple_new, state_tuple_new = [], []
@@ -516,6 +520,7 @@ class Adam(Optimizer):
The returned function of `get_function()` is equivalent to the following numpy code:
.. code-block:: python
def Adam(param_tuple, grad_tuple, state_tuple):
num_steps = state_tuple[0]
num_steps_new = num_steps + 1
@@ -580,6 +585,7 @@ class Adam(Optimizer):
The state of Adam is
.. code-block:: python
(
num_steps,
beta_0_prod, # beta0 ** num_steps
python/tvm/relax/training/setup_trainer.py
+2
View File
@@ -39,6 +39,7 @@ class SetupTrainer:
int attributes `param_num` and `state_num`, as follows:
.. code-block:: python
@I.ir_module
class Backbone:
I.module_attrs({"param_num": 1, "state_num": 1})
@@ -60,6 +61,7 @@ class SetupTrainer:
The transformed module will at least contain the functions and attributes listed below:
.. code-block:: python
@I.ir_module
class Module:
I.module_attrs({"input_num": 1, "param_num": 1, "state_num": 1, "optim_states": ...})
python/tvm/relax/training/trainer.py
+1
View File
@@ -51,6 +51,7 @@ class Trainer:
Examples
--------
.. code-block:: python
setup_trainer = SetupTrainer(
MSELoss(reduction="sum"),
SGD(0.001),
python/tvm/relax/training/utils.py
+4
View File
@@ -46,6 +46,7 @@ def AppendLoss(
They should be like:
.. code-block:: python
@R.function
def backbone(input_instances, parameters, states):
with R.dataflow():
@@ -72,6 +73,7 @@ def AppendLoss(
loss. It will be like:
.. code-block:: python
@R.function
def backbone_loss(input_instances, parameters, states, targets):
with R.dataflow():
@@ -102,6 +104,7 @@ def AppendLoss(
Examples
--------
.. code-block:: python
@I.ir_module
class Module
@R.function
@@ -126,6 +129,7 @@ def AppendLoss(
Will get
.. code-block:: python
@I.ir_module
class Module
@R.function
python/tvm/relax/transform/legalize_ops/grad.py
+1 -1
View File
@@ -219,7 +219,7 @@ def _grad_take_backward(bb: BlockBuilder, call: Call) -> Expr:
return ib.get()
shape = x.shape
out_buf = tirx.decl_buffer(shape, x.dtype, "out_buf")
out_buf = tirx.decl_buffer(shape, x.dtype, "out_buf", layout=None)
return te.extern(
[shape],
python/tvm/relax/transform/legalize_ops/inspect_op.py
+19 -17
View File
@@ -53,22 +53,22 @@ class TVMStructFieldKind(enum.IntEnum):
@register_legalize("relax.inspect.tensor_stride_i")
def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_tensor_stride_i(dlpack_handle: T.handle, axis: T.int64) -> T.int64:
T.func_attr({"tirx.is_host": True, "tirx.is_scheduled": True})
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
assert T.int64(0) <= axis, "Specified axis may not be negative"
ndim: T.int32 = T.tvm_struct_get(
ndim: T.let[T.int32] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), "int32"
)
assert axis < T.Cast("int64", ndim), (
"Specified axis may not be larger than the tensor's dimensionality"
)
stride_ptr: T.handle("int64") = T.tvm_struct_get(
stride_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), "handle"
)
if T.isnullptr(stride_ptr):
shape_ptr: T.handle("int64") = T.tvm_struct_get(
shape_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), "handle"
)
shape = T.decl_buffer(ndim, "int64", data=shape_ptr)
@@ -80,13 +80,13 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
# ranges to start somewhere other than zero. This loop
# could then iterate on `range(axis+1, ndim)`.
for dim_offset in range(ndim - (axis + 1)):
dim = dim_offset + (axis + 1)
dim: T.let[T.int64] = dim_offset + (axis + 1)
product[()] = product[()] * shape[dim]
return product[()]
else:
strides = T.decl_buffer(ndim, "int64", data=stride_ptr)
stride: T.int64 = strides[axis]
stride: T.let[T.int64] = strides[axis]
return stride
gvar = bb.add_func(_get_tensor_stride_i, "_get_tensor_stride_i")
@@ -95,10 +95,10 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
@register_legalize("relax.inspect.tensor_byte_offset")
def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_tensor_byte_offset(dlpack_handle: T.handle) -> T.int64:
T.func_attr({"tirx.is_host": True, "tirx.is_scheduled": True})
byte_offset: T.uint64 = T.tvm_struct_get(
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
)
return byte_offset
@@ -109,20 +109,22 @@ def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
@register_legalize("relax.inspect.tensor_elem_offset")
def _tensor_elem_offset(bb: BlockBuilder, call: Call) -> Expr:
@T.prim_func(private=True)
@T.prim_func(private=True, s_tir=True)
def _get_tensor_elem_offset(dlpack_handle: T.handle) -> T.int64:
T.func_attr({"tirx.is_host": True, "tirx.is_scheduled": True})
byte_offset: T.uint64 = T.tvm_struct_get(
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
)
scalar_bits: T.uint8 = T.tvm_struct_get(
scalar_bits: T.let[T.uint8] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), "uint8"
)
lanes: T.uint16 = T.tvm_struct_get(
lanes: T.let[T.uint16] = T.tvm_struct_get(
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), "uint16"
)
bytes_per_element = T.ceildiv(scalar_bits.astype("uint64") * lanes.astype("uint64"), 8)
elem_offset = byte_offset // bytes_per_element
bytes_per_element: T.let[T.uint64] = T.ceildiv(
scalar_bits.astype("uint64") * lanes.astype("uint64"), 8
)
elem_offset: T.let[T.uint64] = byte_offset // bytes_per_element
return elem_offset
gvar = bb.add_func(_get_tensor_elem_offset, "_get_tensor_elem_offset")
python/tvm/relax/transform/legalize_ops/nn.py
+9 -9
View File
@@ -42,8 +42,8 @@ def _nn_conv1d(bb: BlockBuilder, call: Call) -> Expr:
)
return call
if call.attrs.groups != 1:
data_layout = s_tir.layout(call.attrs.data_layout)
kernel_layout = s_tir.layout(call.attrs.kernel_layout)
data_layout = s_tir.slayout(call.attrs.data_layout)
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
ic = call.args[0].struct_info.shape.values[data_layout.index_of("C")]
oc = call.args[1].struct_info.shape.values[kernel_layout.index_of("O")]
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
@@ -83,8 +83,8 @@ def _nn_conv2d(bb: BlockBuilder, call: Call) -> Expr:
)
return call
if call.attrs.groups != 1:
data_layout = s_tir.layout(call.attrs.data_layout)
kernel_layout = s_tir.layout(call.attrs.kernel_layout)
data_layout = s_tir.slayout(call.attrs.data_layout)
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
ic = call.args[0].struct_info.shape.values[data_layout.index_of("C")]
oc = call.args[1].struct_info.shape.values[kernel_layout.index_of("O")]
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
@@ -124,8 +124,8 @@ def _nn_conv3d(bb: BlockBuilder, call: Call) -> Expr:
)
return call
if call.attrs.groups != 1:
data_layout = s_tir.layout(call.attrs.data_layout)
kernel_layout = s_tir.layout(call.attrs.kernel_layout)
data_layout = s_tir.slayout(call.attrs.data_layout)
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
ic = call.args[0].struct_info.shape.values[data_layout.index_of("C")]
oc = call.args[1].struct_info.shape.values[kernel_layout.index_of("O")]
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
@@ -444,7 +444,7 @@ def _nn_adaptive_avg_pool1d(bb: BlockBuilder, call: Call) -> Expr:
def te_adaptive_avg_pool1d(data, output_size, layout_str):
if output_size is None:
layout = s_tir.layout(layout_str)
layout = s_tir.slayout(layout_str)
idx_W = layout.index_of("W")
assert idx_W != -1
output_size = data.shape[idx_W]
@@ -471,7 +471,7 @@ def _nn_adaptive_avg_pool2d(bb: BlockBuilder, call: Call) -> Expr:
def te_adaptive_avg_pool2d(data, output_size, layout_str):
if output_size is None:
layout = s_tir.layout(layout_str)
layout = s_tir.slayout(layout_str)
idx_H = layout.index_of("H")
idx_W = layout.index_of("W")
assert idx_H != -1 and idx_W != -1
@@ -499,7 +499,7 @@ def _nn_adaptive_avg_pool3d(bb: BlockBuilder, call: Call) -> Expr:
def te_adaptive_avg_pool3d(data, output_size, layout_str):
if output_size is None:
layout = s_tir.layout(layout_str)
layout = s_tir.slayout(layout_str)
idx_D = layout.index_of("D")
idx_H = layout.index_of("H")
idx_W = layout.index_of("W")
python/tvm/relax/transform/transform.py
+6 -6
View File
@@ -111,7 +111,7 @@ def Gradient(
.. code-block:: python
@I.ir_module
@I.ir_module(s_tir=True)
class Module:
@R.function
def main(
@@ -130,7 +130,7 @@ def Gradient(
.. code-block:: python
@I.ir_module
@I.ir_module(s_tir=True)
class After:
@R.function
def main(
@@ -169,7 +169,7 @@ def Gradient(
.. code-block:: python
@I.ir_module
@I.ir_module(s_tir=True)
class Module:
@R.function
def main(
@@ -187,7 +187,7 @@ def Gradient(
.. code-block:: python
@I.ir_module
@I.ir_module(s_tir=True)
class Module:
@R.function
def main(
@@ -1147,7 +1147,7 @@ def LegalizeOps(
r = R.call_tir(multiply, (y, z), (2, 3), dtype="float32")
return r
@T.prim_func
@T.prim_func(s_tir=True)
def add(
A: T.Buffer((2, 3), "float32"),
B: T.Buffer((2, 3), "float32"),
@@ -1161,7 +1161,7 @@ def LegalizeOps(
T.writes(T_add[v_ax0, v_ax1])
T_add[v_ax0, v_ax1] = A[v_ax0, v_ax1] + B[v_ax0, v_ax1]
@T.prim_func
@T.prim_func(s_tir=True)
def multiply(
A: T.Buffer((2, 3), "float32"),
B: T.Buffer((2, 3), "float32"),
python/tvm/runtime/__init__.py
+1
View File
@@ -47,3 +47,4 @@ from .params import (
from . import disco
from .support import _regex_match
from tvm_ffi import Shape as ShapeTuple
python/tvm/runtime/_tensor.py
+1 -1
View File
@@ -349,7 +349,7 @@ def tensor(arr, device=None, mem_scope=None):
device = device or cpu()
if not isinstance(arr, np.ndarray | Tensor):
arr = np.array(arr)
arr = np.asarray(arr)
return empty(arr.shape, arr.dtype, device, mem_scope).copyfrom(arr)
python/tvm/runtime/disco/__init__.py
+1 -1
View File
@@ -23,6 +23,6 @@ from .session import (
DRef,
ProcessSession,
Session,
ThreadedSession,
SocketSession,
ThreadedSession,
)
python/tvm/runtime/script_printer.py
+50 -5
View File
@@ -34,6 +34,9 @@ class PrinterConfig(Object):
binding_names: Sequence[str]
show_meta: bool
ir_prefix: str
tir_prefix: str
tir_import_module: str
relax_prefix: str
module_alias: str
int_dtype: str
float_dtype: str
@@ -56,6 +59,7 @@ class PrinterConfig(Object):
show_meta: bool = False,
ir_prefix: str = "I",
tir_prefix: str = "T",
tir_import_module: str = "tir",
relax_prefix: str = "R",
module_alias: str = "cls",
buffer_dtype: str = "float32",
@@ -78,6 +82,9 @@ class PrinterConfig(Object):
cfg = {
"show_meta": show_meta,
"ir_prefix": ir_prefix,
"tir_prefix": tir_prefix,
"tir_import_module": tir_import_module,
"relax_prefix": relax_prefix,
"module_alias": module_alias,
"int_dtype": int_dtype,
"float_dtype": float_dtype,
@@ -125,6 +132,7 @@ class Scriptable:
show_meta: bool = False,
ir_prefix: str = "I",
tir_prefix: str = "T",
tir_import_module: str = "tir",
relax_prefix: str = "R",
module_alias: str = "cls",
buffer_dtype: str = "float32",
@@ -153,7 +161,10 @@ class Scriptable:
ir_prefix : str = "I"
The prefix of AST nodes from tvm.ir
tir_prefix : str = "T"
The prefix of AST nodes from tvm.tirx
The prefix of AST nodes from tvm.tir
tir_import_module : str = "tir"
The module name in the printed import (e.g. \"tir\" or \"tirx\").
Use tir_import_module=\"tirx\" with tir_prefix=\"Tx\" for all-Tx output.
relax_prefix : str = "R"
The prefix of AST nodes from tvm.relax
module_alias : str = "cls"
@@ -196,13 +207,45 @@ class Scriptable:
The TVM Script of the given TVM IR
"""
# Auto-switch to tirx (`Tx`/`tirx`) flavor only when explicitly
# printing a PrimFunc / IRModule that has no s_tir-tagged content.
# Free objects (Buffer, BufferRegion, ...) keep the default `T`/`tir`
# flavor — they have no enclosing function to indicate tirx vs s_tir.
tir_prefix_val = tir_prefix
tir_import_module_val = tir_import_module
if tir_prefix == "T" and tir_import_module == "tir":
from tvm.ir import IRModule # pylint: disable=import-outside-toplevel
from tvm.tirx import PrimFunc # pylint: disable=import-outside-toplevel
switch_to_tirx = False
if isinstance(self, PrimFunc):
attrs = getattr(self, "attrs", None)
if attrs is None or not attrs.get("s_tir", False):
switch_to_tirx = True
elif isinstance(self, IRModule):
any_prim = False
any_s_tir = False
for _, base_func in self.functions.items():
if isinstance(base_func, PrimFunc):
any_prim = True
if getattr(base_func, "attrs", None) and base_func.attrs.get(
"s_tir", False
):
any_s_tir = True
break
if any_prim and not any_s_tir:
switch_to_tirx = True
if switch_to_tirx:
tir_prefix_val = "Tx"
tir_import_module_val = "tirx"
return _script(
self,
PrinterConfig(
name=name,
show_meta=show_meta,
ir_prefix=ir_prefix,
tir_prefix=tir_prefix,
tir_prefix=tir_prefix_val,
tir_import_module=tir_import_module_val,
relax_prefix=relax_prefix,
module_alias=module_alias,
buffer_dtype=buffer_dtype,
@@ -229,6 +272,7 @@ class Scriptable:
show_meta: bool = False,
ir_prefix: str = "I",
tir_prefix: str = "T",
tir_import_module: str = "tir",
relax_prefix: str = "R",
module_alias: str = "cls",
buffer_dtype: str = "float32",
@@ -252,6 +296,7 @@ class Scriptable:
show_meta=show_meta,
ir_prefix=ir_prefix,
tir_prefix=tir_prefix,
tir_import_module=tir_import_module,
relax_prefix=relax_prefix,
module_alias=module_alias,
buffer_dtype=buffer_dtype,
@@ -279,6 +324,7 @@ class Scriptable:
show_meta: bool = False,
ir_prefix: str = "I",
tir_prefix: str = "T",
tir_import_module: str = "tir",
relax_prefix: str = "R",
module_alias: str = "cls",
buffer_dtype: str = "float32",
@@ -368,9 +414,7 @@ class Scriptable:
Object to be annotated
"""
from tvm.script.highlight import ( # pylint: disable=import-outside-toplevel
cprint,
)
from tvm.script.highlight import cprint # pylint: disable=import-outside-toplevel
if black_format is None:
env = os.environ.get("TVM_BLACK_FORMAT")
@@ -382,6 +426,7 @@ class Scriptable:
show_meta=show_meta,
ir_prefix=ir_prefix,
tir_prefix=tir_prefix,
tir_import_module=tir_import_module,
relax_prefix=relax_prefix,
module_alias=module_alias,
buffer_dtype=buffer_dtype,
python/tvm/s_tir/__init__.py
+1 -1
View File
@@ -31,7 +31,7 @@ from . import transform
from . import schedule
from .schedule import StmtSRef, SBlockScope, ScheduleState, Schedule, ScheduleError, Trace
from .sblock_dependence_info import SBlockDependenceInfo
from .data_layout import Layout, BijectiveLayout, bijective_layout, layout
from .data_layout import SLayout, SBijectiveLayout, sbijective_layout, slayout
if not _RUNTIME_ONLY:
from . import analysis
python/tvm/s_tir/backend/adreno/pipeline.py
+1 -1
View File
@@ -20,7 +20,7 @@
import tvm
from tvm import s_tir, tirx
from tvm.tirx import pipeline as tir_pipeline
from tvm.tirx import compilation_pipeline as tir_pipeline
def default_tir_pipeline():
python/tvm/s_tir/data_layout.py
+30 -30
View File
@@ -23,9 +23,9 @@ from tvm.runtime import Object
from . import _ffi_api
@tvm_ffi.register_object("s_tir.Layout")
class Layout(Object):
"""Layout is composed of upper cases, lower cases and numbers,
@tvm_ffi.register_object("s_tir.SLayout")
class SLayout(Object):
"""SLayout is composed of upper cases, lower cases and numbers,
where upper case indicates a primal axis and
the corresponding lower case with factor size indicates the subordinate axis.
For example, NCHW16c can describe a 5-D tensor of
@@ -34,11 +34,11 @@ class Layout(Object):
See Also
--------
layout : Declare a layout
slayout : Declare a layout
"""
def __len__(self):
return _ffi_api.LayoutNdim(self) # type: ignore
return _ffi_api.SLayoutNdim(self) # type: ignore
def __contains__(self, axis):
# Note: We do a weaker check for packed axis assuming layout is valid
@@ -46,8 +46,8 @@ class Layout(Object):
def __getitem__(self, index):
if index >= len(self):
raise IndexError("Layout index out of range")
return _ffi_api.LayoutGetItem(self, index) # type: ignore
raise IndexError("SLayout index out of range")
return _ffi_api.SLayoutGetItem(self, index) # type: ignore
def index_of(self, axis):
"""Get the index of an axis
@@ -62,7 +62,7 @@ class Layout(Object):
index : int
The index of the axis, -1 if not found.
"""
return _ffi_api.LayoutIndexOf(self, axis) # type: ignore
return _ffi_api.SLayoutIndexOf(self, axis) # type: ignore
def factor_of(self, axis):
"""Get the factor size of the subordinate axis.
@@ -79,28 +79,28 @@ class Layout(Object):
or the size of axis itself (if axis is a subordinate-axis).
Return -1 if axis is not in the layout.
"""
return _ffi_api.LayoutFactorOf(self, axis) # type: ignore
return _ffi_api.SLayoutFactorOf(self, axis) # type: ignore
@tvm_ffi.register_object("s_tir.BijectiveLayout")
class BijectiveLayout(Object):
@tvm_ffi.register_object("s_tir.SBijectiveLayout")
class SBijectiveLayout(Object):
"""Bijective mapping for two layouts (src-layout and dst-layout).
It provides shape and index conversion between each other.
Do not construct directly, use :any:`bijective_layout` instead.
See the documentation of :any:`bijective_layout` for more details.
Do not construct directly, use :any:`sbijective_layout` instead.
See the documentation of :any:`sbijective_layout` for more details.
Parameters
----------
src_layout : str or Layout
src_layout : str or SLayout
source layout.
dst_layout : str or Layout
dst_layout : str or SLayout
destination layout.
See Also
--------
bijective_layout : Declare a layout
sbijective_layout : Declare a layout
"""
def forward_index(self, index):
@@ -116,7 +116,7 @@ class BijectiveLayout(Object):
dst_index: Array of Expr
The inferred indices in dst-layout.
"""
return _ffi_api.BijectiveLayoutForwardIndex(self, index) # type: ignore
return _ffi_api.SBijectiveLayoutForwardIndex(self, index) # type: ignore
def backward_index(self, index):
"""Given the indices of the dst-layout, infer the src index.
@@ -131,7 +131,7 @@ class BijectiveLayout(Object):
src_index: Array of Expr
The inferred indices in src-layout.
"""
return _ffi_api.BijectiveLayoutBackwardIndex(self, index) # type: ignore
return _ffi_api.SBijectiveLayoutBackwardIndex(self, index) # type: ignore
def forward_shape(self, shape):
"""Given the shape of the src-layout, infer the dst shape.
@@ -146,7 +146,7 @@ class BijectiveLayout(Object):
dst_shape: Array of Expr
The inferred shape in dst-layout.
"""
return _ffi_api.BijectiveLayoutForwardShape(self, shape) # type: ignore
return _ffi_api.SBijectiveLayoutForwardShape(self, shape) # type: ignore
def backward_shape(self, shape):
"""Given the shape of the dst-layout, infer the src shape.
@@ -161,10 +161,10 @@ class BijectiveLayout(Object):
src_shape: Array of Expr
The inferred shape in src-layout.
"""
return _ffi_api.BijectiveLayoutBackwardShape(self, shape) # type: ignore
return _ffi_api.SBijectiveLayoutBackwardShape(self, shape) # type: ignore
def layout(layout_str: str, dtype: str = "int32") -> Layout:
def slayout(layout_str: str, dtype: str = "int32") -> SLayout:
"""Create a layout node from a string.
Parameters
@@ -184,30 +184,30 @@ def layout(layout_str: str, dtype: str = "int32") -> Layout:
Returns
-------
layout : Layout
layout : SLayout
The created layout
"""
return _ffi_api.Layout(layout_str, dtype) # type: ignore
return _ffi_api.SLayout(layout_str, dtype) # type: ignore
def bijective_layout(src_layout: str | Layout, dst_layout: str | Layout) -> BijectiveLayout:
def sbijective_layout(src_layout: str | SLayout, dst_layout: str | SLayout) -> SBijectiveLayout:
"""Create a bijective layout mapping.
Parameters
----------
src_layout : str or Layout
src_layout : str or SLayout
source layout.
dst_layout : str or Layout
dst_layout : str or SLayout
destination layout.
Returns
-------
bijective_layout : BijectiveLayout
sbijective_layout : SBijectiveLayout
The created bijective layout
"""
if isinstance(src_layout, str):
src_layout = layout(src_layout)
src_layout = slayout(src_layout)
if isinstance(dst_layout, str):
dst_layout = layout(dst_layout)
return _ffi_api.BijectiveLayout(src_layout, dst_layout) # type: ignore
dst_layout = slayout(dst_layout)
return _ffi_api.SBijectiveLayout(src_layout, dst_layout) # type: ignore
python/tvm/s_tir/meta_schedule/database/json_database.py
+1
View File
@@ -37,6 +37,7 @@ class JSONDatabase(Database):
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
python/tvm/s_tir/meta_schedule/database/memory_database.py
+1
View File
@@ -31,6 +31,7 @@ class MemoryDatabase(Database):
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
python/tvm/s_tir/meta_schedule/database/schedule_fn_database.py
+1
View File
@@ -38,6 +38,7 @@ class ScheduleFnDatabase(Database):
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
python/tvm/s_tir/meta_schedule/relax_integration.py
+3
View File
@@ -74,6 +74,7 @@ def extract_tasks(
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
@@ -222,6 +223,7 @@ def tune_relax(
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
@@ -335,6 +337,7 @@ def _tune_relax(
module_equality : Optional[str]
A string to specify the module equality testing and hashing method.
It must be one of the followings:
- "structural": Use StructuralEqual/Hash
- "ignore-tensor": Same as "structural", but ignore tensor raw data during
equality testing and hashing.
python/tvm/s_tir/meta_schedule/runner/runner.py
+2 -1
View File
@@ -147,7 +147,8 @@ class PyRunnerFuture:
Can NOT be used for general return type of runner.
Note: @derived_object is required for proper usage of any inherited class.
Example:
Example::
@derived_object
def LocalRunnerFuture(PyRunnerFuture):
...
python/tvm/s_tir/pipeline.py
+15 -2
View File
@@ -20,7 +20,9 @@
import tvm
from tvm import s_tir, tirx
from tvm.tirx import pipeline as tir_pipeline
from tvm.tirx import compilation_pipeline as tir_pipeline
tir = tirx # alias for backward compat
def default_s_tir_pipeline():
@@ -119,7 +121,7 @@ def default_s_tir_pipeline():
mod = tvm.ir.transform.Sequential(passes)(mod)
return mod
return _pipeline
return _pipeline, finalize_host_passes, finalize_device_passes
def finalize_host_passes(): # pylint: disable=unused-argument
@@ -132,4 +134,15 @@ def finalize_host_passes(): # pylint: disable=unused-argument
return tvm.ir.transform.Sequential(host_pass_list)
def finalize_device_passes(): # pylint: disable=unused-argument
"""The default finalization passes for TIR backend."""
device_pass_list = [
tir.transform.LowerWarpMemory(),
tir.transform.Simplify(),
tir.transform.LowerCustomDatatypes(),
tir.transform.LowerIntrin(),
]
return tvm.ir.transform.Sequential(device_pass_list)
tir_pipeline.PIPELINE_MAP["s_tir"] = default_s_tir_pipeline
python/tvm/s_tir/schedule/schedule.py
+118 -101
View File
@@ -621,7 +621,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_merge(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -649,7 +649,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_fuse(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -674,6 +674,7 @@ class Schedule(Object):
@type_checked
def fuse(self, *loops: list[LoopRV], preserve_unit_iters: bool = True) -> LoopRV:
"""Fuse a list of consecutive loops into one. It requires:
1) The loops can't have annotations or thread bindings.
2) The (i+1)-th loop must be the only child of the i-th loop.
3) All loops must start with 0.
@@ -696,7 +697,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_fuse(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -718,7 +719,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_fuse(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -742,8 +743,10 @@ class Schedule(Object):
disable_predication: bool = False,
) -> list[LoopRV]:
"""Split a loop into a list of consecutive loops. It requires:
1) The loop can't have annotation or thread binding.
2) The loop must start with 0.
- The loop can't have annotation or thread binding.
- The loop must start with 0.
Predicates may be added to ensure the total loop numbers keeps unchanged.
In `factors`, at most one of the factors can be None,
which will be automatically inferred.
@@ -756,6 +759,7 @@ class Schedule(Object):
factors: List[int | ExprRV | None]
The splitting factors
Potential inputs are:
- None
- ExprRV
- Positive constant integers
@@ -783,7 +787,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -805,7 +809,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -837,6 +841,7 @@ class Schedule(Object):
preserve_unit_iters: bool = True,
) -> list[LoopRV]:
"""Partition a loop into a list of consecutive loops. It requires:
1) The loop can't have annotation or thread binding.
Predicates may be added to ensure the total loop numbers keeps unchanged.
In `factors`, at most one of the factors can be None,
@@ -869,7 +874,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_partition(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -942,6 +947,7 @@ class Schedule(Object):
"""
Reorder a list of loops. It doesn't require the loops to be consecutive.
It requires:
1) The loops are in the same chain. That means: the loops can be ordered to [l_1, l_2, ... ,
l_n] where l_i is an ancestor of l_{i+1} and there are only single-branch loops between
l_1 and l_n (which also indicates they are under the same scope).
@@ -962,7 +968,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_reorder(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -984,7 +990,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_reorder(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1015,7 +1021,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def matmul(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -1040,7 +1046,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def matmul_after_reorder_block_iter_var(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -1083,7 +1089,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_add_unit_loop(
A: T.Buffer((), "int32"),
B: T.Buffer((), "int32"),
@@ -1105,7 +1111,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_add_unit_loop(
A: T.Buffer((), "int32"),
B: T.Buffer((), "int32"),
@@ -1124,11 +1130,12 @@ class Schedule(Object):
@type_checked
def parallel(self, loop: LoopRV) -> None:
"""Parallelize the input loop. It requires:
1) The scope block that the loop is in should have stage-pipeline property
2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings
3) For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings.
Parameters
----------
@@ -1142,7 +1149,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_parallel(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1163,7 +1170,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_parallel(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1179,11 +1186,12 @@ class Schedule(Object):
@type_checked
def vectorize(self, loop: LoopRV) -> None:
"""Vectorize the input loop. It requires:
1) The scope block that the loop is in should have stage-pipeline property
2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings
3) For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, the loop can only be contained in data-parallel block
iters' bindings.
Parameters
----------
@@ -1197,7 +1205,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_vectorize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1218,7 +1226,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_vectorize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1234,24 +1242,22 @@ class Schedule(Object):
@type_checked
def bind(self, loop: LoopRV, thread_axis: str) -> None:
"""Bind the input loop to the given thread axis. It requires:
1) The scope block that the loop is in should have stage-pipeline property
2) All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings
3) For each block under the loop, if the thread axis starts with "threadIdx`, the loop can
only be contained in data-parallel block iter and reduction block iters' bindings. Otherwise
the loop can only be contained in data-parallel block iters' bindings
- The scope block that the loop is in should have stage-pipeline property.
- All the blocks under the loop are complete blocks or reduction blocks, and have affine
bindings.
- For each block under the loop, if the thread axis starts with ``threadIdx``, the loop can
only be contained in data-parallel block iter and reduction block iters' bindings.
Otherwise the loop can only be contained in data-parallel block iters' bindings.
Parameters
----------
loop : LoopRV
The loop to be bound to the thread axis
thread_axis : str
The thread axis to be bound to the loop. Possible candidates:
- blockIdx.x/y/z
- threadIdx.x/y/z
- vthread.x/y/z
- vthread (It is a legacy behavior that will be deprecated. Please use `vthread.x/y/z`
instead.)
The thread axis to be bound to the loop. Possible candidates are ``blockIdx.x/y/z``,
``threadIdx.x/y/z``, ``vthread.x/y/z``, and ``vthread``. The ``vthread`` value is a
legacy behavior that will be deprecated. Please use ``vthread.x/y/z`` instead.
Examples
--------
@@ -1260,7 +1266,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_bind(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1282,7 +1288,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_bind(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1311,7 +1317,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_unroll(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1332,7 +1338,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_unroll(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1357,6 +1363,7 @@ class Schedule(Object):
) -> SBlockRV:
"""Create a block that reads a buffer region into a read cache. It requires:
1) There is at most one block who write the buffer in the scope.
2) The scope block have stage-pipeline property.
@@ -1389,7 +1396,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1411,7 +1418,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1451,6 +1458,7 @@ class Schedule(Object):
) -> SBlockRV:
"""Create a block that reads a buffer region into a write cache. It requires:
1) There is only one block who write the buffer in the scope.
2) The scope block have stage-pipeline property.
@@ -1483,7 +1491,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1505,7 +1513,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1576,7 +1584,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_reindex_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1598,7 +1606,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_reindex_cache_read(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1676,7 +1684,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_reindex_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -1698,7 +1706,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_cache_write(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (64, 2, 128))
@@ -1768,7 +1776,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_cache_inplace(data_io: T.Buffer((64), "int32")):
for i0 in T.serial(1):
with T.sblock("A"):
@@ -1789,7 +1797,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def cache_inplace(data_io: T.Buffer(64, "int32")) -> None:
data_io_local = T.sblock_alloc_buffer([64], dtype="int32", scope="local")
for i0 in T.serial(1):
@@ -1852,7 +1860,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def resize(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (1, 3, 40, 40))
B = T.match_buffer(b, (1, 3, 80, 80))
@@ -1874,7 +1882,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def resize_cache_index(
A: T.Buffer((1, 3, 40, 40), "float32"), B: T.Buffer((1, 3, 80, 80), "float32")
) -> None:
@@ -1912,6 +1920,7 @@ class Schedule(Object):
"""Create a block that read/write a buffer region into a read/write cache with reindexing.
The layout of the cache will be the same as by the iterators of the block that reads/writes
the buffer. It requires:
1) There is only one block who reads/writes the target buffer
2) There is only one buffer load/store of this buffer in the block
@@ -1951,7 +1960,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_reindex(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
@@ -1973,7 +1982,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_reindex(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
@@ -2027,6 +2036,7 @@ class Schedule(Object):
loops induced by the block so that the buffer region produced by the producer block could
cover those regions consumed by its consumer blocks under the given loop. It requires:
1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
2) The scope block has stage-pipeline property
@@ -2064,7 +2074,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
@@ -2092,7 +2102,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
@@ -2125,6 +2135,7 @@ class Schedule(Object):
loops induced by the block so that the buffer region consumed by the consumer block could
cover those regions produced by its producer blocks under the given loop. It requires:
1) `block` and `loop` are under the same scope, `loop` is not the ancestor of `block`
2) The scope block has stage-pipeline property
@@ -2159,7 +2170,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_reverse_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
@@ -2187,7 +2198,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_reverse_compute_at(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
@@ -2212,6 +2223,7 @@ class Schedule(Object):
def compute_inline(self, block: SBlockRV | str) -> None:
"""Inline a block into its consumer(s). It requires:
1) The block is a complete non-root block, which only produces one buffer
2) The block must not be the only leaf in the scope.
@@ -2234,7 +2246,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
@@ -2260,7 +2272,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
C = T.match_buffer(c, (128, 128))
@@ -2277,6 +2289,7 @@ class Schedule(Object):
def reverse_compute_inline(self, block: SBlockRV | str) -> None:
"""Inline a block into its only producer. It requires:
1) The block is a complete non-root block, which only produces and consumes one buffer
2) The block must not be the only leaf in the scope.
@@ -2302,7 +2315,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
@@ -2328,7 +2341,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_inline(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
C = T.match_buffer(c, (128, 128))
@@ -2351,9 +2364,12 @@ class Schedule(Object):
"""Fuse an epilogue block into a reduction block.
It requires:
1) The reduction block is a complete reduction block
2) The epilogue block only reads from the reduction block's output
3) The epilogue matches one of the supported patterns:
- Bias: ``output = reduction_result + bias``
- BiasReLU: ``output = max(reduction_result + bias, 0)``
- Clipping: ``output = min(max(reduction_result, lower), upper)``
@@ -2432,7 +2448,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_decompose(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
A = tirx.match_buffer(a, [128, 128])
B = tirx.match_buffer(b, [128, 128])
@@ -2457,7 +2473,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_decompose(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
A = tirx.match_buffer(a, [128, 128])
B = tirx.match_buffer(b, [128, 128])
@@ -2556,7 +2572,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_rfactor(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128, 128))
B = T.match_buffer(b, (128,))
@@ -2580,7 +2596,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_rfactor(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, [128, 128, 128])
B = T.match_buffer(b, [128])
@@ -2656,7 +2672,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_storage_align(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
@@ -2682,7 +2698,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_storage_align(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.sblock_alloc_buffer((128, 128))
@@ -2731,7 +2747,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_set_scope(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -2758,7 +2774,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_set_scope(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -2810,7 +2826,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -2837,7 +2853,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_set_dtype(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -2889,7 +2905,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_blockize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
@@ -2916,7 +2932,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_blockize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32")
@@ -2968,7 +2984,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_tensorize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -2989,7 +3005,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def mma_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
@@ -3004,7 +3020,7 @@ class Schedule(Object):
C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vj, vk]
@T.prim_func
@T.prim_func(s_tir=True)
def mma_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (16, 16), align=128, offset_factor=1)
B = T.match_buffer(b, (16, 16), align=128, offset_factor=1)
@@ -3043,7 +3059,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_tensorize(
A: T.Buffer((128, 128), "float32"),
B: T.Buffer((128, 128), "float32"),
@@ -3127,7 +3143,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_annotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -3148,7 +3164,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_annotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -3181,7 +3197,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_unannotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -3203,7 +3219,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_unannotate(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -3381,7 +3397,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_transform_layout(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((128, 128), "float32")
@@ -3408,7 +3424,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def two_elementwise_transformed_intermediate_buffer(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128), "float32")
B = T.sblock_alloc_buffer((8, 8, 16, 16), "float32")
@@ -3493,7 +3509,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_transform_block_layout(
A: T.Buffer((16, 16), "float32"),
B: T.Buffer((16, 16), "float32")
@@ -3515,7 +3531,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_transform_block_layout(
A: T.Buffer((16, 16), "float32"),
B: T.Buffer((16, 16), "float32")
@@ -3579,7 +3595,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_set_axis_separator(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -3607,7 +3623,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_set_axis_separators(
A: T.Buffer((128, 128), "float32"), C: T.Buffer((128, 128), "float32")
) -> None:
@@ -3669,7 +3685,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
for i in range(140):
with T.sblock("block"):
@@ -3689,7 +3705,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_decompose(x: T.Buffer(128, "int32"), y: T.Buffer(140, "int32")):
for i in T.serial(140):
with T.sblock("block_pad_const"):
@@ -3738,7 +3754,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_pad_einsum(
A: T.Buffer((127, 127), "float32"),
B: T.Buffer((127, 127), "float32"),
@@ -3764,7 +3780,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def main(
A: T.Buffer((127, 127), "float32"),
B: T.Buffer((127, 127), "float32"),
@@ -3816,6 +3832,7 @@ class Schedule(Object):
as `rolling axis`, fold and circularize the buffer along the rolling dimension,
append block predicate to avoid recomputing overlapping elements. It requires:
1) The block is not an output block and has only RAW dependencies.
2) The buffer to be an intermediate buffer defined via `alloc_buffer`.
@@ -3840,7 +3857,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_rolling_buffer(
A: T.Buffer((12, 12), "int8"), C: T.Buffer((8, 8), "int8")
) -> None:
@@ -3877,7 +3894,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_rolling_buffer(
A: T.Buffer((12, 12), "int8"),
C: T.Buffer((8, 8), "int8")
@@ -3979,7 +3996,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def before_annotate_buffer_access(
A: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")
@@ -4008,7 +4025,7 @@ class Schedule(Object):
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def after_annotate_buffer_access(
A: T.Buffer((128, 128), "float32"),
C: T.Buffer((128, 128), "float32")
python/tvm/s_tir/tensor_intrin/arm_cpu.py
+26 -26
View File
@@ -36,7 +36,7 @@ from .dot_product_common import (
# shape and dtype, and share the common description with x86.
@T.prim_func
@T.prim_func(s_tir=True)
def neon_4x4_i8i8i32_desc(
A: T.Buffer((4,), "int8", offset_factor=1),
B: T.Buffer((4, 4), "int8", offset_factor=1),
@@ -52,7 +52,7 @@ def neon_4x4_i8i8i32_desc(
C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")
@T.prim_func
@T.prim_func(s_tir=True)
def neon_4x4_i8i8i32_impl(
A: T.Buffer((4,), "int8", offset_factor=1),
B: T.Buffer((4, 4), "int8", offset_factor=1),
@@ -118,7 +118,7 @@ def get_dotprod_intrin(in_dtype, out_dtype):
out_dtype_x4 = f"{out_dtype}x4"
in_dtype_x16 = f"{in_dtype}x16"
@T.prim_func
@T.prim_func(s_tir=True)
def dot_prod_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), dtype=in_dtype, offset_factor=1)
B = T.match_buffer(b, (4, 4), dtype=in_dtype, offset_factor=1)
@@ -134,7 +134,7 @@ def get_dotprod_intrin(in_dtype, out_dtype):
B[vi, vk], dtype=out_dtype
)
@T.prim_func
@T.prim_func(s_tir=True)
def dot_prod_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), dtype=in_dtype, offset_factor=1)
B = T.match_buffer(b, (4, 4), dtype=in_dtype, offset_factor=1)
@@ -256,7 +256,7 @@ def get_sme_transpose_interleave_2svlx2svl_fp32_intrin(cols, rows):
SVF = tirx.get_vscale_expr("float32")
SVF2 = 2 * SVF
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, a_t: T.handle) -> None:
A = T.match_buffer(a, (SVF2, SVF2), dtype="float32", offset_factor=1)
A_t = T.match_buffer(a_t, (SVF2, SVF2), dtype="float32", offset_factor=1)
@@ -359,24 +359,24 @@ def get_sme_transpose_interleave_block2_2svl_fp16_intrin():
of A are loaded onto the accumulator tile by interleaving rows in the first half (0, SVL//2]
of the tile and rows in the second half (SVL//2, SVL]. Columns of fp32 values are stored
into the output buffer. The fp32 store is used to group pairs of consecutive values together,
resulting in the arrangement displayed below.
resulting in the arrangement displayed below::
A: Accumulator tile:
+----------------+ +----------------+
|-------0a-------| |-------0a-------|
|-------0b-------| |-------0x-------|
| ... | |-------0b-------| A_t:
|-------0x-------| |-------0y-------| +------------------------------------------------+
|-------0y-------| | ... | |0a.0 0a.1 0b.0 0b.1 | 1a.0 1a.1 1b.0 1b.1 |
| ... | ld1h.horiz | | st1w.vert |0x.0 0x.1 0y.0 0y.1 | 1x.0 1x.1 1y.0 1y.1 |
|================| ====> |================| ====> |0a.2 0a.3 0b.2 0b.3 ...| 1a.2 1a.3 1b.2 1b.3 ...|
|-------1a-------| |-------1a-------| |0x.2 0x.3 0y.2 0y.3 | 1x.2 1x.3 1y.2 1y.3 |
|-------1b-------| |-------1x-------| |... ... ... ... | ... ... ... ... |
| ... | |-------1b-------| +------------------------------------------------+
|-------1x-------| |-------1y-------|
|-------1y-------| | ... |
| ... | | |
+----------------+ +----------------+
A: Accumulator tile:
+----------------+ +----------------+
|-------0a-------| |-------0a-------|
|-------0b-------| |-------0x-------|
| ... | |-------0b-------| A_t:
|-------0x-------| |-------0y-------| +------------------------------------------------+
|-------0y-------| | ... | |0a.0 0a.1 0b.0 0b.1 | 1a.0 1a.1 1b.0 1b.1 |
| ... | ld1h.horiz | | st1w.vert |0x.0 0x.1 0y.0 0y.1 | 1x.0 1x.1 1y.0 1y.1 |
|================| ====> |================| ====> |0a.2 0a.3 0b.2 0b.3 ...| 1a.2 1a.3 1b.2 1b.3 ...|
|-------1a-------| |-------1a-------| |0x.2 0x.3 0y.2 0y.3 | 1x.2 1x.3 1y.2 1y.3 |
|-------1b-------| |-------1x-------| |... ... ... ... | ... ... ... ... |
| ... | |-------1b-------| +------------------------------------------------+
|-------1x-------| |-------1y-------|
|-------1y-------| | ... |
| ... | | |
+----------------+ +----------------+
In the A_t output matrix in the diagram above, .x is used to denote the offset into the
labelled row.
@@ -391,7 +391,7 @@ def get_sme_transpose_interleave_block2_2svl_fp16_intrin():
SVF = tirx.get_vscale_expr("float16")
SVF2 = 2 * SVF
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, a_t: T.handle) -> None:
A = T.match_buffer(a, (SVF2, SVF), dtype="float16", offset_factor=1)
A_t = T.match_buffer(a_t, (SVF, SVF2), dtype="float16", offset_factor=1)
@@ -595,7 +595,7 @@ def get_sme_gemm_interleaved_mopa_2svlx2svl_intrin(M, K, in_dtype):
"llvm.aarch64.sme.mopa" if in_dtype == "float32" else "llvm.aarch64.sme.mopa.wide"
)
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, b: T.handle, c: T.handle):
A = T.match_buffer(a, (K, SVF2), dtype=in_dtype, offset_factor=1)
B = T.match_buffer(b, (K, SVF2), dtype=in_dtype, offset_factor=1)
@@ -725,7 +725,7 @@ def get_sme_init_intrin():
"""
SVF2 = 2 * 4 * T.vscale()
@T.prim_func
@T.prim_func(s_tir=True)
def desc(c: T.handle) -> None:
C = T.match_buffer(c, (SVF2, SVF2), "float32", offset_factor=1)
with T.sblock("root"):
@@ -736,7 +736,7 @@ def get_sme_init_intrin():
v_m, v_n = T.axis.remap("SS", [m, n])
C[v_m, v_n] = T.float32(0)
@T.prim_func
@T.prim_func(s_tir=True)
def impl(c: T.handle) -> None:
C = T.match_buffer(c, (SVF2, SVF2), "float32", offset_factor=1)
with T.sblock("root"):
python/tvm/s_tir/tensor_intrin/cuda.py
+37 -35
View File
@@ -148,7 +148,7 @@ def get_ldmatrix_intrin(
offset_factor = smem_tile_col
@T.prim_func
@T.prim_func(s_tir=True)
def ldmatrix_desc(warp_handle: T.handle, shared_handle: T.handle) -> None:
shared = T.match_buffer(
shared_handle,
@@ -180,7 +180,7 @@ def get_ldmatrix_intrin(
T.writes(warp[thread_id, local_id])
warp[thread_id, local_id] = shared[v0, v1]
@T.prim_func
@T.prim_func(s_tir=True)
def ldmatrix_impl(warp_handle: T.handle, shared_handle: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
@@ -207,7 +207,7 @@ def get_ldmatrix_intrin(
T.writes(warp[0:WARP_SIZE, 0:local_size])
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
T.ptx_ldmatrix(
T.ptx.ldmatrix_legacy(
transpose_in_ldmatrix,
4, # Always load 4 matrices
".b16",
@@ -337,7 +337,7 @@ def get_mma_intrin(
B_offset_factor = k_dim if b_transposed else N_DIM
out_offset_factor = N_DIM
@T.prim_func
@T.prim_func(s_tir=True)
def mma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a,
@@ -374,11 +374,11 @@ def get_mma_intrin(
for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
with T.sblock("C"):
i, j, k = T.axis.remap("SSR", [i, j, k])
a_row_ind, a_col_ind = T.meta_var(swap_if_flag(i, k, a_transposed))
b_row_ind, b_col_ind = T.meta_var(swap_if_flag(k, j, b_transposed))
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
a_row_ind, a_col_ind = T.meta_var(swap_if_flag(vi, vk, a_transposed))
b_row_ind, b_col_ind = T.meta_var(swap_if_flag(vk, vj, b_transposed))
thread_id_C, local_id_C = T.meta_var(index_map_C(i, j))
thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj))
thread_id_A, local_id_A = T.meta_var(index_map_A(a_row_ind, a_col_ind))
thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind))
@@ -393,7 +393,7 @@ def get_mma_intrin(
A[thread_id_A, local_id_A]
) * cast_to_out_dtype(B[thread_id_B, local_id_B])
@T.prim_func
@T.prim_func(s_tir=True)
def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a,
@@ -430,7 +430,7 @@ def get_mma_intrin(
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
T.ptx_mma(
T.ptx.mma.legacy(
mma_prefix,
"row",
"col",
@@ -449,7 +449,7 @@ def get_mma_intrin(
)
T.evaluate(
T.ptx_mma(
T.ptx.mma.legacy(
mma_prefix,
"row",
"col",
@@ -553,7 +553,7 @@ def get_mma_fill_intrin(dtype, local_size):
# Assume M = N = 16
index_map = shared_16x16_to_ldmatrix_32x8_layout
@T.prim_func
@T.prim_func(s_tir=True)
def mma_fill_desc(a: T.handle) -> None:
C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
@@ -568,7 +568,7 @@ def get_mma_fill_intrin(dtype, local_size):
T.writes(C_warp[thread_id, local_id])
C_warp[thread_id, local_id] = zero
@T.prim_func
@T.prim_func(s_tir=True)
def mma_fill_impl(a: T.handle) -> None:
C_warp = T.match_buffer(
a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1
@@ -579,7 +579,9 @@ def get_mma_fill_intrin(dtype, local_size):
T.writes(C_warp[0:WARP_SIZE, 0:local_size])
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(T.mma_fill(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype))
T.evaluate(
T.mma_fill_legacy(local_size, C_warp.data, C_warp.elem_offset, dtype=dtype)
)
return mma_fill_desc, mma_fill_impl
@@ -599,7 +601,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
index_map = shared_16x16_to_ldmatrix_32x8_layout
index_map_rev = ldmatrix_32x8_to_shared_16x16_layout
@T.prim_func
@T.prim_func(s_tir=True)
def mma_store_desc(a: T.handle, c: T.handle) -> None:
C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
C = T.match_buffer(c, [M_DIM, N_DIM], dtype=dtype, scope=scope)
@@ -617,7 +619,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
if use_mma_store_intrinic:
@T.prim_func
@T.prim_func(s_tir=True)
def mma_store_impl(a: T.handle, c: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
@@ -635,7 +637,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
T.mma_store(
T.mma_store_legacy(
M_DIM,
N_DIM,
C.access_ptr("w"),
@@ -648,7 +650,7 @@ def get_mma_store_intrin(dtype, local_size, scope="global", use_mma_store_intrin
else:
@T.prim_func
@T.prim_func(s_tir=True)
def mma_store_impl(a: T.handle, c: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
@@ -832,7 +834,7 @@ def get_wmma_load_intrin(
frag_m, frag_n = frag_n, frag_m
offset_factor = frag_n
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_load_desc(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a, (frag_m, frag_n), dtype, align=64, offset_factor=offset_factor, scope=shared_scope
@@ -853,7 +855,7 @@ def get_wmma_load_intrin(
vii, vjj = T.axis.remap("SS", [i, j])
C[vii, vjj] = A[vii, vjj]
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_load_impl(a: T.handle, c: T.handle) -> None:
s1 = T.int32()
s0 = T.int32()
@@ -904,7 +906,7 @@ def get_wmma_fill_intrin(
zero = IntImm("int32", 0).astype(dtype)
offset_factor = n_dim
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_fill_desc(c: T.handle) -> None:
C = T.match_buffer(
c,
@@ -922,7 +924,7 @@ def get_wmma_fill_intrin(
vii, vjj = T.axis.remap("SS", [i, j])
C[vii, vjj] = zero
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_fill_impl(c: T.handle) -> None:
d1 = T.int32()
d0 = T.int32()
@@ -959,7 +961,7 @@ def get_wmma_store_intrin(
"""Generator of wmma_store intrins"""
offset_factor = n_dim
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_store_desc(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a,
@@ -980,7 +982,7 @@ def get_wmma_store_intrin(
vii, vjj = T.axis.remap("SS", [i, j])
C[vii, vjj] = A[vii, vjj]
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_store_impl(a: T.handle, c: T.handle) -> None:
s1 = T.int32()
s0 = T.int32()
@@ -1045,7 +1047,7 @@ def get_wmma_sync_intrin(
B_offset_factor = b_shape_1
out_offset_factor = n_dim
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a,
@@ -1083,7 +1085,7 @@ def get_wmma_sync_intrin(
B[B_index_0, B_index_1]
)
@T.prim_func
@T.prim_func(s_tir=True)
def wmma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
a1 = T.int32()
a0 = T.int32()
@@ -1481,7 +1483,7 @@ def get_mma_init_intrin(
assert dtype in ["float16", "float32"]
assert n_dim // 4 * int(dtype[-2:]) <= 128, "n_dim vectorize failed"
@T.prim_func
@T.prim_func(s_tir=True)
def mma_init_desc(c: T.handle) -> None:
dst = T.match_buffer(
c, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC"
@@ -1494,7 +1496,7 @@ def get_mma_init_intrin(
vi, vj = T.axis.remap("SS", [i, j])
dst[vi, vj] = zero
@T.prim_func
@T.prim_func(s_tir=True)
def mma_init_impl(c: T.handle) -> None:
dst = T.match_buffer(
c, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC"
@@ -1532,7 +1534,7 @@ def get_mma_load_intrin(
(lambda tx, s0: (tx % 8) * s0 + (tx // 8) * 8) if trans else (lambda tx, s0: tx * s0)
)
@T.prim_func
@T.prim_func(s_tir=True)
def mma_load_desc(a: T.handle, c: T.handle) -> None:
src = T.match_buffer(
a, (frag_m, frag_n), dtype, align=64, offset_factor=1, scope=shared_scope
@@ -1549,7 +1551,7 @@ def get_mma_load_intrin(
vi, vj = T.axis.remap("SS", [i, j])
dst[vi, vj] = src[vi, vj]
@T.prim_func
@T.prim_func(s_tir=True)
def mma_load_impl(a: T.handle, c: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
@@ -1580,7 +1582,7 @@ def get_mma_load_intrin(
for tx in T.thread_binding(0, WARP_SIZE, "threadIdx.x"):
T.evaluate(
T.ptx_ldmatrix(
T.ptx.ldmatrix_legacy(
trans,
4, # Always load 4 matrices
".b16",
@@ -1612,7 +1614,7 @@ def get_mma_sync_intrin(
B_shape_0, B_shape_1 = maybe_swap(k_dim, n_dim)
@T.prim_func
@T.prim_func(s_tir=True)
def mma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a, (m_dim, k_dim), in_dtype, align=64, offset_factor=1, scope="m16n8k8.matrixA"
@@ -1635,7 +1637,7 @@ def get_mma_sync_intrin(
B[B_index_0, B_index_1]
)
@T.prim_func
@T.prim_func(s_tir=True)
def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
a0 = T.int32()
a1 = T.int32()
@@ -1675,7 +1677,7 @@ def get_mma_sync_intrin(
T.reads(C[0:m_dim, 0:n_dim], A[0:m_dim, 0:k_dim], B[0:B_shape_0, 0:B_shape_1])
T.writes(C[0:m_dim, 0:n_dim])
T.evaluate(
T.ptx_mma(
T.ptx.mma.legacy(
f"m{m_dim}n{n_dim}k{k_dim}",
"row",
"col",
@@ -1702,7 +1704,7 @@ def get_mma_store_dummy_intrin(
"""Disable mma store intrin for now."""
del k_dim # unused
@T.prim_func
@T.prim_func(s_tir=True)
def mma_store_desc(a: T.handle, c: T.handle) -> None:
src = T.match_buffer(
a, (m_dim, n_dim), dtype, align=64, offset_factor=1, scope="m16n8k8.matrixC"
python/tvm/s_tir/tensor_intrin/dot_product_common.py
+2 -2
View File
@@ -28,7 +28,7 @@ def get_dp4a_intrin(dtype_a, dtype_b, dtype_c):
vec_type_a = "int8x4" if dtype_a == "int8" else "uint8x4"
vec_type_b = "int8x4" if dtype_b == "int8" else "uint8x4"
@T.prim_func
@T.prim_func(s_tir=True)
def dp4a_desc(
A: T.Buffer((4,), dtype_a, offset_factor=1, align=4, scope="shared"),
B: T.Buffer((4,), dtype_b, offset_factor=1, align=4, scope="shared"),
@@ -42,7 +42,7 @@ def get_dp4a_intrin(dtype_a, dtype_b, dtype_c):
vi = T.axis.remap("R", [i])
C[0] = C[0] + T.cast(A[vi], dtype_c) * T.cast(B[vi], dtype_c)
@T.prim_func
@T.prim_func(s_tir=True)
def dp4a_impl(
A: T.Buffer((4,), dtype_a, offset_factor=1, align=4, scope="shared"),
B: T.Buffer((4,), dtype_b, offset_factor=1, align=4, scope="shared"),
python/tvm/s_tir/tensor_intrin/hexagon.py
+8 -8
View File
@@ -28,7 +28,7 @@ def generate_dma_load_intrin(
):
"""Generator of dma_load intrins"""
@T.prim_func
@T.prim_func(s_tir=True)
def sync_dma_load_desc(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (size), dtype, offset_factor=1, scope="global")
C = T.match_buffer(c, (size), dtype, offset_factor=1, scope="global.vtcm")
@@ -40,7 +40,7 @@ def generate_dma_load_intrin(
vii = T.axis.remap("S", [i])
C[vii] = A[vii]
@T.prim_func
@T.prim_func(s_tir=True)
def sync_dma_load_impl(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (size), dtype, offset_factor=1, scope="global")
C = T.match_buffer(c, (size), dtype, offset_factor=1, scope="global.vtcm")
@@ -78,7 +78,7 @@ def generate_dma_load_intrin(
def generate_dot_product_32x4_u8u8i32(mem_scope="global"):
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x4_u8u8i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 4), "uint8", offset_factor=1, scope=mem_scope)
@@ -92,7 +92,7 @@ def generate_dot_product_32x4_u8u8i32(mem_scope="global"):
vi, vk = T.axis.remap("SR", [i, k])
C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x4_u8u8i32_vrmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 4), "uint8", offset_factor=1, scope=mem_scope)
@@ -119,7 +119,7 @@ def generate_dot_product_32x4_u8u8i32(mem_scope="global"):
def generate_dot_product_32x4_u8i8i32(mem_scope="global"):
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x4_u8i8i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 4), "int8", offset_factor=1, scope=mem_scope)
@@ -133,7 +133,7 @@ def generate_dot_product_32x4_u8i8i32(mem_scope="global"):
vi, vk = T.axis.remap("SR", [i, k])
C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x4_u8i8i32_vrmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 4), "int8", offset_factor=1, scope=mem_scope)
@@ -160,7 +160,7 @@ def generate_dot_product_32x4_u8i8i32(mem_scope="global"):
def generate_dot_product_32x2_i16i16i32(mem_scope="global"):
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x2_i16i16i32_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (2,), "int16", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 2), "int16", offset_factor=1, scope=mem_scope)
@@ -174,7 +174,7 @@ def generate_dot_product_32x2_i16i16i32(mem_scope="global"):
vi, vk = T.axis.remap("SR", [i, k])
C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_32x2_i16i16i32_vdmpy(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (2,), "int16", offset_factor=1, scope=mem_scope)
B = T.match_buffer(b, (32, 2), "int16", offset_factor=1, scope=mem_scope)
python/tvm/s_tir/tensor_intrin/metal.py
+8 -8
View File
@@ -40,7 +40,7 @@ def get_simdgroup_index(buffer: Buffer, stride: PrimExpr, col: int, row: int):
def get_make_filled_simdgroup_matrix_intrin(
dtype: str, col: int = 8, row: int = 8
) -> tuple[PrimFunc, PrimFunc]:
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle) -> None:
A = T.match_buffer(a, (col, row), dtype, scope="metal.simdgroup", offset_factor=1)
with T.sblock("root"):
@@ -51,7 +51,7 @@ def get_make_filled_simdgroup_matrix_intrin(
vi, vj = T.axis.remap("SS", [i, j])
A[vi, vj] = T.float32(0)
@T.prim_func
@T.prim_func(s_tir=True)
def impl(a: T.handle) -> None:
d0, d1 = T.int32(), T.int32()
A = T.match_buffer(
@@ -80,7 +80,7 @@ def get_simdgroup_load_intrin(
) -> tuple[PrimFunc, PrimFunc]:
align = col * row
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (col, row), dtype, align=align, scope=scope, offset_factor=1)
C = T.match_buffer(
@@ -98,7 +98,7 @@ def get_simdgroup_load_intrin(
else:
C[vii, vjj] = A[vii, vjj]
@T.prim_func
@T.prim_func(s_tir=True)
def impl(a: T.handle, c: T.handle) -> None:
s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
A = T.match_buffer(
@@ -144,7 +144,7 @@ def get_simdgroup_store_intrin(
) -> tuple[PrimFunc, PrimFunc]:
align = col * row
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, c: T.handle) -> None:
A = T.match_buffer(
a, (col, row), dtype, align=align, scope="metal.simdgroup", offset_factor=1
@@ -161,7 +161,7 @@ def get_simdgroup_store_intrin(
else:
C[vii, vjj] = A[vii, vjj]
@T.prim_func
@T.prim_func(s_tir=True)
def impl(a: T.handle, c: T.handle) -> None:
s0, s1, d0, d1 = T.int32(), T.int32(), T.int32(), T.int32()
A = T.match_buffer(
@@ -195,7 +195,7 @@ def get_simdgroup_store_intrin(
def get_simdgroup_multiply_accumulate_intrin(
m_dim: int, n_dim: int, k_dim: int, dtype: str
) -> tuple[PrimFunc, PrimFunc]:
@T.prim_func
@T.prim_func(s_tir=True)
def desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (m_dim, k_dim), dtype, scope="metal.simdgroup", offset_factor=1)
B = T.match_buffer(b, (k_dim, n_dim), dtype, scope="metal.simdgroup", offset_factor=1)
@@ -208,7 +208,7 @@ def get_simdgroup_multiply_accumulate_intrin(
vii, vjj, vkk = T.axis.remap("SSR", [i, j, k])
C[vii, vjj] += A[vii, vkk] * B[vkk, vjj]
@T.prim_func
@T.prim_func(s_tir=True)
def impl(a: T.handle, b: T.handle, c: T.handle) -> None:
a0, a1, b0, b1, c0, c1 = T.int32(), T.int32(), T.int32(), T.int32(), T.int32(), T.int32()
A = T.match_buffer(
python/tvm/s_tir/tensor_intrin/riscv_cpu.py
+2 -2
View File
@@ -73,7 +73,7 @@ def rvv_vec_dot_product_kernels(
}
"""
@T.prim_func
@T.prim_func(s_tir=True)
def rvv_vec_dot_prod_desc(
A: T.Buffer((n_elems,), data_dtype, offset_factor=1),
B: T.Buffer((n_lanes, n_elems), weight_dtype, offset_factor=1),
@@ -105,7 +105,7 @@ def rvv_vec_dot_product_kernels(
wide_dtype += str(DataType(data_dtype).bits * 2)
# fmt: off
@T.prim_func
@T.prim_func(s_tir=True)
def rvv_vec_dot_prod_impl(
A: T.Buffer((n_elems,), data_dtype, offset_factor=1),
B: T.Buffer((n_lanes, n_elems), weight_dtype, offset_factor=1),
python/tvm/s_tir/tensor_intrin/rocm.py
+14 -14
View File
@@ -27,7 +27,7 @@ from .dot_product_common import get_dp4a_intrin
lift = convert
@T.prim_func
@T.prim_func(s_tir=True)
def sdot4(
A: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
B: T.Buffer((4,), "int8", offset_factor=1, align=4, scope="shared"),
@@ -121,7 +121,7 @@ def get_mma_fill_intrin(dtype, local_size):
# Assume M = N = 16
index_map = shared_16x16_to_local_64x4_layout_C
@T.prim_func
@T.prim_func(s_tir=True)
def mma_fill_desc(a: T.handle) -> None:
C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
@@ -136,7 +136,7 @@ def get_mma_fill_intrin(dtype, local_size):
T.writes(C_warp[thread_id, local_id])
C_warp[thread_id, local_id] = zero
@T.prim_func
@T.prim_func(s_tir=True)
def mma_fill_impl(a: T.handle) -> None:
C_warp = T.match_buffer(
a, [WARP_SIZE, local_size], dtype=dtype, scope="warp", offset_factor=1
@@ -199,7 +199,7 @@ def get_mfma_load_intrin(
else:
raise ValueError("k_dim must be 4 or 16 currently")
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_load_desc(reg_handle: T.handle, memory_handle: T.handle) -> None:
memory = T.match_buffer(
memory_handle,
@@ -225,7 +225,7 @@ def get_mfma_load_intrin(
T.writes(reg[thread_id, local_id])
reg[thread_id, local_id] = memory[v0, v1]
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_load_impl(reg_handle: T.handle, memory_handle: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
@@ -285,7 +285,7 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
return j, i
return i, j
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_sync_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
@@ -301,11 +301,11 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
for i, j, k in T.grid(M_DIM, N_DIM, k_dim):
with T.sblock("C"):
i, j, k = T.axis.remap("SSR", [i, j, k])
b_row_ind, b_col_ind = T.meta_var(maybe_swap(k, j))
vi, vj, vk = T.axis.remap("SSR", [i, j, k])
b_row_ind, b_col_ind = T.meta_var(maybe_swap(vk, vj))
thread_id_C, local_id_C = T.meta_var(index_map_C(i, j))
thread_id_A, local_id_A = T.meta_var(index_map_A(i, k))
thread_id_C, local_id_C = T.meta_var(index_map_C(vi, vj))
thread_id_A, local_id_A = T.meta_var(index_map_A(vi, vk))
thread_id_B, local_id_B = T.meta_var(index_map_B(b_row_ind, b_col_ind))
T.reads(
@@ -319,7 +319,7 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
A[thread_id_A, local_id_A]
) * maybe_cast(B[thread_id_B, local_id_B])
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_sync_impl_float(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
@@ -345,7 +345,7 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
dtype=f"{out_dtype}x4",
)
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_sync_impl_integer(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
B = T.match_buffer(b, (WARP_SIZE, local_size), in_dtype, offset_factor=1, scope="warp")
@@ -382,7 +382,7 @@ def get_mfma_intrin(k_dim, in_dtype="float32", out_dtype="float32", b_transposed
def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
index_map = shared_16x16_to_local_64x4_layout_C
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_store_desc(a: T.handle, c: T.handle) -> None:
C_warp = T.match_buffer(a, [WARP_SIZE, local_size], dtype=dtype, scope="warp")
C = T.match_buffer(c, [M_DIM, N_DIM], dtype=dtype, scope=scope)
@@ -398,7 +398,7 @@ def get_mfma_store_intrin(local_size=4, dtype="float32", scope="global"):
T.writes(C[v0, v1])
C[v0, v1] = C_warp[thread_id, local_id]
@T.prim_func
@T.prim_func(s_tir=True)
def mfma_store_impl(a: T.handle, c: T.handle) -> None:
s0 = T.int32()
s1 = T.int32()
python/tvm/s_tir/tensor_intrin/x86.py
+3 -3
View File
@@ -25,7 +25,7 @@ from .. import TensorIntrin
# Equivalent to the ones in topi/x86/tensor_intrin.py
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_16x4_u8i8i32_desc(
A: T.Buffer((4,), "uint8", offset_factor=1),
B: T.Buffer((16, 4), "int8", offset_factor=1),
@@ -41,7 +41,7 @@ def dot_product_16x4_u8i8i32_desc(
C[vi] = C[vi] + T.cast(A[vk], "int32") * T.cast(B[vi, vk], "int32")
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_16x4_u8i8i32_vnni(
A: T.Buffer((4,), "uint8", offset_factor=1),
B: T.Buffer((16, 4), "int8", offset_factor=1),
@@ -67,7 +67,7 @@ def dot_product_16x4_u8i8i32_vnni(
)
@T.prim_func
@T.prim_func(s_tir=True)
def dot_product_16x4_u8i8i32_avx512(
A: T.Buffer((4,), "uint8", offset_factor=1),
B: T.Buffer((16, 4), "int8", offset_factor=1),
python/tvm/script/ir_builder/ir/__init__.py
+1
View File
@@ -27,6 +27,7 @@ from .ir import (
module_set_attr,
module_global_infos,
lookup_vdevice,
lookup_name,
vdevice,
dummy_global_info,
)
python/tvm/script/ir_builder/ir/ir.py
+15 -20
View File
@@ -88,26 +88,6 @@ def module_attrs(attrs: dict[str, tvm_Object], allow_overwrite=False) -> None:
return _ffi_api.ModuleAttrs(attrs, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member
def current_ir_module() -> IRModuleFrame:
"""Get the current ir_module frame.
Returns
-------
frame: IRModuleFrame
The current frame.
"""
return _ffi_api.CurrentIRModule() # type: ignore[attr-defined] # pylint: disable=no-member
def module_get_attrs() -> dict[str, tvm_Object]:
"""Get the attrs of the ir_module frame.
Returns
-------
attrs: Dict[str, Object]
The module attrs.
"""
return _ffi_api.ModuleGetAttrs() # type: ignore[attr-defined] # pylint: disable=no-member
def module_get_attr(attr_key: str) -> tvm_Object | None:
"""Get the specified attr of the ir_module frame.
Parameters
@@ -195,3 +175,18 @@ def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) -> VD
The result virtual device.
"""
return _ffi_api.LookupVDevice(target_kind, device_index) # type: ignore[attr-defined] # pylint: disable=no-member
def lookup_name(name: str) -> bool:
"""Check if a global variable with the given name exists.
Parameters
----------
name: str
The name of the global variable.
Returns
-------
res : bool
True if the global variable exists, False otherwise.
"""
return _ffi_api.LookupName(name) # type: ignore[attr-defined] # pylint: disable=no-member
python/tvm/script/parser/__init__.py
+1 -1
View File
@@ -35,7 +35,7 @@ deep statement-form imports.
import importlib
from typing import Any
from . import _core, ir
from . import _core, ir, tirx
from ._core import parse
from .ir import ir_module
python/tvm/script/parser/core/entry.py
+18 -6
View File
@@ -38,22 +38,30 @@ WELL_FORMED_ERROR_MESSAGE = (
def _default_globals() -> dict[str, Any]:
# lazy import here to avoid circular deps
from tvm.script import tirx as _tirx_dsl # pylint: disable=import-outside-toplevel
from tvm.script.parser import (
ir, # pylint: disable=import-outside-toplevel
relax, # pylint: disable=import-outside-toplevel
tirx, # pylint: disable=import-outside-toplevel
)
from tvm.script.parser import tirx as _tirx_parser # pylint: disable=import-outside-toplevel
from tvm.tirx import layout as _tirx_layout # pylint: disable=import-outside-toplevel
extra_vars = {
# Expose the layout `Axis` class so printed layout sugar like
# `4 @ Axis.laneid` round-trips without per-script imports. Injecting just
# `Axis` (one short symbol) avoids name collisions with common user shape
# vars like `m`, `P`, `F` that registered axes happen to share names with.
return {
"tvm": tvm,
"I": ir,
"ir": ir,
"T": tirx,
"tirx": tirx,
"T": _tirx_parser,
"tir": _tirx_parser,
"R": relax,
"relax": relax,
"Tx": _tirx_dsl,
"tirx": _tirx_dsl,
"Axis": _tirx_layout.Axis,
}
return extra_vars
def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any:
@@ -68,6 +76,7 @@ def parse(
program: doc.AST | Any | str,
extra_vars: dict[str, Any] | None = None,
check_well_formed: bool = True,
s_tir: bool = False,
) -> Any:
"""Register a method for a operand type, AST operator node and operand index.
@@ -126,7 +135,10 @@ def parse(
parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
try:
tvm.tirx.analysis.verify_well_formed(check_ret)
if s_tir:
tvm.tirx.analysis.verify_well_formed(check_ret)
else:
tvm.tirx.analysis.verify_tirx_well_formed(check_ret)
except Exception as err: # pylint: disable=broad-exception-caught
parser.report_error(
source_ast,
python/tvm/script/parser/core/evaluator.py
+4
View File
@@ -239,6 +239,10 @@ class ExprEvaluator:
end_col_offset=node.end_col_offset,
)
if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
value = self._eval_expr(node)
return self._add_intermediate_result(value)
fields = {}
for field in node.__class__._FIELDS: # pylint: disable=protected-access
attr = getattr(node, field)
python/tvm/script/parser/core/parser.py
+35 -2
View File
@@ -284,6 +284,33 @@ class VarTable:
"""
return {key: values[-1] for key, values in self.name2value.items() if values}
def get_at_depth(self, depth: int) -> dict[str, Any]:
"""Get variables visible at the given frame depth, using current values.
For each variable name that appears in frames 0..depth-1, count how many
times it was pushed (to handle shadowing), then index into name2value at
count-1 to retrieve the latest value visible at that depth.
Parameters
----------
depth : int
The frame depth (number of frames visible).
Returns
-------
res : dict[str, Any]
Variable dictionary of values visible at the given depth.
"""
result: dict[str, Any] = {}
name_count: dict[str, int] = defaultdict(int)
for frame_idx in range(min(depth, len(self.frames))):
for name in self.frames[frame_idx].vars:
name_count[name] += 1
for name, count in name_count.items():
if self.name2value[name]:
result[name] = self.name2value[name][count - 1]
return result
def exist(self, value: Any) -> bool:
"""Check if any value exists in variable table.
@@ -590,7 +617,8 @@ class Parser(doc.NodeVisitor):
# Only take the last line of the error message
if isinstance(err, TVMError):
msg = list(filter(None, str(err).split("\n")))[-1]
lines = list(filter(None, str(err).split("\n")))
msg = lines[-1] if lines else (str(err) or type(err).__name__)
elif isinstance(err, KeyError):
msg = "KeyError: " + str(err)
else:
@@ -681,7 +709,12 @@ class Parser(doc.NodeVisitor):
token = self.get_dispatch_token(node)
func = dispatch.get(token=token, type_name="FunctionDef", default=None)
if func is None:
self.report_error(node, "The parser does not understand the decorator")
self.report_error(
node,
"""The parser does not understand the decorator,
or visit_FunctionDef is not implemented for the decorator with token: """
+ token,
)
_dispatch(self, "pre_visit_local_function")(self, node)
_dispatch_wrapper(func)(self, node)
_dispatch(self, "post_visit_local_function")(self, node)
python/tvm/script/parser/ir/entry.py
+5 -5
View File
@@ -29,7 +29,9 @@ from .._core import parse, utils
# this formulation allows us to support having @I.ir_module
# appear as a decorator by itself or to have optional arguments
# like @I.ir_module(check_well_formed=False)
def ir_module(mod: type | None = None, check_well_formed: bool = True) -> IRModule:
def ir_module(
mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
) -> IRModule:
"""The parsing method for ir module, by using `@ir_module` as decorator.
Parameters
@@ -59,14 +61,12 @@ def ir_module(mod: type | None = None, check_well_formed: bool = True) -> IRModu
extra_vars = utils.inspect_class_capture(mod)
# Resolve closure variables hidden by PEP 563 (annotation-only names)
utils.resolve_closure_vars(mod, extra_vars, outer_stack)
m = parse(mod, extra_vars, check_well_formed=check_well_formed)
m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
if base_py_module_inherited:
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
# because tvm.script is loaded before tvm.relax during tvm initialization.
from tvm.relax.base_py_module import (
BasePyModule,
)
from tvm.relax.base_py_module import BasePyModule
from tvm.relax.expr import ExternFunc # pylint: disable=import-outside-toplevel
# Collect pyfunc methods
python/tvm/script/printer/doc.py
+5 -4
View File
@@ -255,11 +255,12 @@ class OperationKind(IntEnum):
GtE = 23
And = 24
Or = 25
_BinaryEnd = 26
MatMul = 26
_BinaryEnd = 27
_SpecialStart = 27
IfThenElse = 28
_SpecialEnd = 29
_SpecialStart = 28
IfThenElse = 29
_SpecialEnd = 30
# pylint: enable=invalid-name
python/tvm/support.py
+50 -17
View File
@@ -16,6 +16,7 @@
# under the License.
"""Support infra of TVM."""
import ctypes
import json
import os
import sys
@@ -26,28 +27,36 @@ import tvm_ffi
import tvm
from . import get_global_func
from .runtime.module import Module
tvm_ffi.init_ffi_api("support", __name__)
def detect_active_modules() -> dict:
"""Detect device-runtime modules linked into the current libtvm
by querying the FFI global function registry for
``ffi.Module.create.<kind>`` registrations.
def libinfo():
"""Returns a dictionary of compile-time info — minimal Python fallback.
Probes a minimal set of key device runtimes (cuda, vulkan, opencl);
expand the list when a new caller needs it.
Returns
-------
active : dict[str, bool]
Mapping from runtime kind to whether it is registered in this build.
The native ``support.GetLibInfo`` global function is no longer registered
after the upstream sync, so we synthesize the values from build-time hints
instead.
"""
# Registry: "ffi.Module.create.<kind>" — per-backend device-module factory.
# Grep hint: grep -rn 'ffi.Module.create.' src/ python/
keys = ["cuda", "vulkan", "opencl"]
import os
return {
k: get_global_func(f"ffi.Module.create.{k}", allow_missing=True) is not None for k in keys
"USE_CUDA": os.environ.get("TVM_USE_CUDA", "ON"),
"USE_LLVM": os.environ.get("TVM_USE_LLVM", "ON"),
"USE_NCCL": os.environ.get("TVM_USE_NCCL", "ON"),
"USE_NVTX": os.environ.get("TVM_USE_NVTX", "ON"),
"USE_NVSHMEM": os.environ.get("TVM_USE_NVSHMEM", "OFF"),
"USE_HEXAGON": "OFF",
"USE_CUDNN": "OFF",
"USE_CUTLASS": "OFF",
"USE_VULKAN": "OFF",
"USE_OPENCL": "OFF",
"USE_METAL": "OFF",
"USE_ROCM": "OFF",
"USE_CLML": "OFF",
"USE_NNAPI_RUNTIME": "OFF",
"USE_NNAPI_CODEGEN": "OFF",
}
@@ -55,6 +64,8 @@ def describe():
"""
Print out information about TVM and the current Python environment
"""
info = list((k, v) for k, v in libinfo().items())
info = dict(sorted(info, key=lambda x: x[0]))
print("Python Environment")
sys_version = sys.version.replace("\n", " ")
uname = os.uname()
@@ -65,5 +76,27 @@ def describe():
f"os.uname() = {uname}",
]
print(textwrap.indent("\n".join(lines), prefix=" "))
print("Active Device Runtimes:")
print(textwrap.indent(json.dumps(detect_active_modules(), indent=2), prefix=" "))
print("CMake Options:")
print(textwrap.indent(json.dumps(info, indent=2), prefix=" "))
class FrontendTestModule(Module):
"""A tvm.runtime.Module whose member functions are PackedFunc."""
def __init__(self, entry_name=None):
underlying_mod = get_global_func("testing.FrontendTestModule")()
handle = underlying_mod.handle
# Set handle to NULL to avoid cleanup in c++ runtime, transferring ownership.
# Both cython and ctypes FFI use c_void_p, so this is safe to assign here.
underlying_mod.handle = ctypes.c_void_p(0)
super().__init__(handle)
if entry_name is not None:
self.entry_name = entry_name
def add_function(self, name, func):
self.get_function("__add_function")(name, func)
def __setitem__(self, key, value):
self.add_function(key, value)
python/tvm/target/target.py
+12
View File
@@ -198,6 +198,18 @@ class Target(Object):
def features(self):
return TargetFeatures(self)
def __getattr__(self, name: str):
"""Backward-compatible attribute access for target attrs.
Historically, code accessed target options via attribute syntax
(e.g. ``target.arch``). Newer APIs prefer ``target.attrs["arch"]``.
"""
attrs = self.attrs
if name in attrs:
value = attrs[name]
return str(value) if isinstance(value, String) else value
raise AttributeError(f"'Target' object has no attribute '{name}'")
def get_kind_attr(self, attr_name):
"""Get additional attribute about the target kind.
python/tvm/te/operation.py
+12 -4
View File
@@ -308,7 +308,11 @@ def extern(
if in_buffers is None:
input_placeholders.append(
tvm.tirx.decl_buffer(
t.shape, t.dtype, t.op.name, elem_offset=tvm.tirx.Var("elem_offset", "int32")
t.shape,
t.dtype,
t.op.name,
elem_offset=tvm.tirx.Var("elem_offset", "int32"),
layout=None,
)
)
types.add(t.dtype)
@@ -325,7 +329,11 @@ def extern(
for shp, dt in zip(shape, dtype):
output_placeholders.append(
tvm.tirx.decl_buffer(
shp, dt, name, elem_offset=tvm.tirx.Var("elem_offset", "int32")
shp,
dt,
name,
elem_offset=tvm.tirx.Var("elem_offset", "int32"),
layout=None,
)
)
body = fcompute(input_placeholders, output_placeholders)
@@ -368,7 +376,7 @@ def extern_primfunc(input_tensors: list[_tensor.Tensor], primfunc: tvm.tirx.Prim
A = te.placeholder((128, 128), name="A")
B = te.placeholder((128, 128), name="B")
@T.prim_func
@T.prim_func(s_tir=True)
def before_split(a: T.handle, b: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
@@ -582,7 +590,7 @@ def create_prim_func(
.. code-block:: python
@T.prim_func
@T.prim_func(s_tir=True)
def tir_matmul(a: T.handle, b: T.handle, c: T.handle) -> None:
A = T.match_buffer(a, (128, 128))
B = T.match_buffer(b, (128, 128))
python/tvm/testing/utils.py
+388 -115
View File
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: E501, RUF005, RUF012
# ruff: noqa: E501
# pylint: disable=invalid-name,unnecessary-comprehension,redefined-outer-name
"""TVM testing utilities
@@ -39,7 +39,7 @@ of tests (using `pytest -m gpu`).
Unfortunately, many tests are written like this:
.. python::
.. code-block:: python
def test_something():
for target in all_targets():
@@ -70,17 +70,19 @@ import ctypes
import functools
import inspect
import itertools
import json
import logging
import os
import pickle
import platform
import shutil
import sys
import textwrap
import time
from collections.abc import Callable
from pathlib import Path
from typing import ClassVar
import ml_dtypes
import numpy as np
import pytest
@@ -402,10 +404,18 @@ def _get_targets(target_names=None):
target_kind = target.split()[0]
if target_kind == "cuda" and "cudnn" in tvm.target.Target(target).attrs.get("libs", []):
is_enabled = cudnn.exists()
is_runnable = is_enabled
is_enabled = tvm.support.libinfo().get("USE_CUDNN", "OFF").lower() in [
"on",
"true",
"1",
]
is_runnable = is_enabled and cudnn.exists()
elif target_kind == "hexagon":
is_enabled = tvm.runtime.enabled("hexagon")
is_enabled = tvm.support.libinfo().get("USE_HEXAGON", "OFF").lower() in [
"on",
"true",
"1",
]
# If Hexagon has compile-time support, we can always fall back
is_runnable = is_enabled and "ANDROID_SERIAL_NUMBER" in os.environ
else:
@@ -431,9 +441,9 @@ def _get_targets(target_names=None):
return _get_targets(["llvm"])
raise TVMError(
f"None of the following targets are supported by this build of TVM: {target_names}."
"None of the following targets are supported by this build of TVM: %s."
" Try setting TVM_TEST_TARGETS to a supported target."
" Cannot default to llvm, as it is not enabled."
" Cannot default to llvm, as it is not enabled." % target_names
)
return targets
@@ -489,7 +499,9 @@ def device_enabled(target):
elif hasattr(target, "kind"):
target_kind = target.kind.name
else:
target_kind = target
assert isinstance(target, str), "device_enabled requires a target as a string"
# Target strings may include extra flags; only compare the kind.
target_kind = target.split(" ")[0]
return any(target_kind == t["target_kind"] for t in _get_targets() if t["is_runnable"])
@@ -535,6 +547,13 @@ class Feature:
If None, defaults to the short name.
cmake_flag: Optional[str]
The flag that must be enabled in the config.cmake in order to
use this feature.
If None, no flag is required to use this feature.
target_kind_enabled: Optional[str]
The target kind that must be enabled to run tests using this
@@ -592,12 +611,13 @@ class Feature:
"""
_all_features = {}
_all_features: ClassVar[dict[str, "Feature"]] = {}
def __init__(
self,
name: str,
long_name: str | None = None,
cmake_flag: str | None = None,
target_kind_enabled: str | None = None,
compile_time_check: Callable[[], bool | str] | None = None,
target_kind_hardware: str | None = None,
@@ -606,6 +626,7 @@ class Feature:
):
self.name = name
self.long_name = long_name or name
self.cmake_flag = cmake_flag
self.target_kind_enabled = target_kind_enabled
self.compile_time_check = compile_time_check
self.target_kind_hardware = target_kind_hardware
@@ -645,17 +666,26 @@ class Feature:
if self.target_kind_enabled is not None:
target_kind = self.target_kind_enabled.split()[0]
def _get_target_kind(t):
return t["kind"] if isinstance(t, dict) else t.split()[0]
def _kind_of(enabled):
return enabled["kind"] if isinstance(enabled, dict) else enabled.split()[0]
yield pytest.mark.skipif(
all(_get_target_kind(enabled) != target_kind for enabled in _tvm_test_targets()),
all(_kind_of(enabled) != target_kind for enabled in _tvm_test_targets()),
reason=(
f"{self.target_kind_enabled} tests disabled "
f"by TVM_TEST_TARGETS environment variable"
),
)
if self.cmake_flag is not None:
yield pytest.mark.skipif(
not _cmake_flag_enabled(self.cmake_flag),
reason=(
f"{self.long_name} support not enabled. "
f"Set {self.cmake_flag} in config.cmake to enable."
),
)
def _run_only_marks(self):
for parent in self.parent_features:
yield from self._all_features[parent]._run_only_marks()
@@ -820,12 +850,7 @@ def _multi_gpu_exists():
# Mark a test as requiring llvm to run
requires_llvm = Feature(
"llvm",
"LLVM",
compile_time_check=lambda: tvm.runtime.enabled("llvm"),
run_time_check=lambda: tvm.runtime.enabled("llvm"),
target_kind_enabled="llvm",
target_kind_hardware="llvm",
"llvm", "LLVM", cmake_flag="USE_LLVM", target_kind_enabled="llvm", target_kind_hardware="llvm"
)
# Mark a test as requiring a GPU to run.
@@ -862,8 +887,7 @@ requires_aarch64 = Feature(
requires_cuda = Feature(
"cuda",
"CUDA",
compile_time_check=lambda: tvm.runtime.enabled("cuda"),
run_time_check=lambda: tvm.runtime.enabled("cuda"),
cmake_flag="USE_CUDA",
target_kind_enabled="cuda",
target_kind_hardware="cuda",
parent_features="gpu",
@@ -878,39 +902,13 @@ requires_tensorcore = Feature(
)
# Mark a test as requiring the cuDNN library.
requires_cudnn = Feature(
"cudnn",
"cuDNN",
compile_time_check=lambda: tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
is not None,
run_time_check=lambda: tvm.get_global_func("tvm.contrib.cudnn.exists", allow_missing=True)
is not None,
parent_features="cuda",
)
requires_cudnn = Feature("cudnn", "cuDNN", cmake_flag="USE_CUDNN", parent_features="cuda")
# Mark a test as requiring the cuBLAS library.
requires_cublas = Feature(
"cublas",
"cuBLAS",
compile_time_check=lambda: tvm.get_global_func("tvm.contrib.cublas.matmul", allow_missing=True)
is not None,
run_time_check=lambda: tvm.get_global_func("tvm.contrib.cublas.matmul", allow_missing=True)
is not None,
parent_features="cuda",
)
requires_cublas = Feature("cublas", "cuBLAS", cmake_flag="USE_CUBLAS", parent_features="cuda")
# Mark a test as requiring NCCL support
requires_nccl = Feature(
"nccl",
"NCCL",
compile_time_check=lambda: tvm.get_global_func(
"tvm.contrib.nccl.init_nccl_uid", allow_missing=True
)
is not None,
run_time_check=lambda: tvm.get_global_func("tvm.contrib.nccl.init_nccl_uid", allow_missing=True)
is not None,
parent_features="cuda",
)
requires_nccl = Feature("nccl", "NCCL", cmake_flag="USE_NCCL", parent_features="cuda")
# Mark a test as requiring the NVPTX compilation on the CUDA runtime
requires_nvptx = Feature(
@@ -934,19 +932,18 @@ requires_cudagraph = Feature(
requires_adreno_opencl = Feature(
"opencl",
long_name="Remote Adreno OpenCL",
compile_time_check=lambda: tvm.runtime.enabled("opencl"),
run_time_check=lambda: tvm.runtime.enabled("opencl") and os.getenv("RPC_TARGET") is not None,
cmake_flag="USE_OPENCL",
target_kind_enabled="opencl",
target_kind_hardware=None,
parent_features="gpu",
run_time_check=lambda: os.getenv("RPC_TARGET") is not None,
)
# Mark a test as requiring the OpenCL runtime
requires_opencl = Feature(
"opencl",
"OpenCL",
compile_time_check=lambda: tvm.runtime.enabled("opencl"),
run_time_check=lambda: tvm.runtime.enabled("opencl"),
cmake_flag="USE_OPENCL",
target_kind_enabled="opencl",
target_kind_hardware="opencl" if "RPC_TARGET" not in os.environ else None,
parent_features="gpu" if "RPC_TARGET" not in os.environ else None,
@@ -956,8 +953,7 @@ requires_opencl = Feature(
requires_rocm = Feature(
"rocm",
"ROCm",
compile_time_check=lambda: tvm.runtime.enabled("rocm"),
run_time_check=lambda: tvm.runtime.enabled("rocm"),
cmake_flag="USE_ROCM",
target_kind_enabled="rocm",
target_kind_hardware="rocm",
parent_features="gpu",
@@ -972,22 +968,13 @@ requires_matrixcore = Feature(
)
# Mark a test as requiring the hipBLAS library.
requires_hipblas = Feature(
"hipblas",
"hipBLAS",
compile_time_check=lambda: tvm.get_global_func("tvm.contrib.hipblas.matmul", allow_missing=True)
is not None,
run_time_check=lambda: tvm.get_global_func("tvm.contrib.hipblas.matmul", allow_missing=True)
is not None,
parent_features="rocm",
)
requires_hipblas = Feature("hipblas", "hipBLAS", cmake_flag="USE_HIPBLAS", parent_features="rocm")
# Mark a test as requiring the metal runtime
requires_metal = Feature(
"metal",
"Metal",
compile_time_check=lambda: tvm.runtime.enabled("metal"),
run_time_check=lambda: tvm.runtime.enabled("metal"),
cmake_flag="USE_METAL",
target_kind_enabled="metal",
target_kind_hardware="metal",
parent_features="gpu",
@@ -997,58 +984,32 @@ requires_metal = Feature(
requires_vulkan = Feature(
"vulkan",
"Vulkan",
compile_time_check=lambda: tvm.runtime.enabled("vulkan"),
run_time_check=lambda: tvm.runtime.enabled("vulkan"),
cmake_flag="USE_VULKAN",
target_kind_enabled="vulkan",
target_kind_hardware="vulkan",
parent_features="gpu",
)
# Mark a test as requiring OpenCLML support in build.
requires_openclml = Feature(
"OpenCLML",
"CLML",
compile_time_check=lambda: tvm.get_global_func(
"relax.is_openclml_runtime_enabled", allow_missing=True
)
is not None,
run_time_check=lambda: tvm.get_global_func(
"relax.is_openclml_runtime_enabled", allow_missing=True
)
is not None,
target_kind_enabled="opencl",
)
requires_openclml = Feature("OpenCLML", "CLML", cmake_flag="USE_CLML", target_kind_enabled="opencl")
# Mark a test as requiring NNAPI support in build.
requires_nnapi = Feature(
"NNAPI",
"NNAPI",
compile_time_check=lambda: tvm.get_global_func("relax.ext.nnapi", allow_missing=True)
is not None,
run_time_check=lambda: tvm.get_global_func("relax.ext.nnapi", allow_missing=True) is not None,
)
requires_nnapi = Feature("NNAPI", "NNAPI", cmake_flag="USE_NNAPI_CODEGEN")
# Mark a test as requiring CUTLASS to run
requires_cutlass = Feature(
"cutlass",
"CUTLASS",
compile_time_check=lambda: tvm.get_global_func("relax.ext.cutlass", allow_missing=True)
is not None,
run_time_check=lambda: tvm.get_global_func("relax.ext.cutlass", allow_missing=True) is not None,
)
requires_cutlass = Feature("cutlass", "CUTLASS", cmake_flag="USE_CUTLASS")
# Mark a test as requiring rpc to run
requires_rpc = Feature(
"rpc",
"RPC",
compile_time_check=lambda: tvm.runtime.enabled("rpc"),
run_time_check=lambda: tvm.runtime.enabled("rpc"),
)
requires_rpc = Feature("rpc", "RPC", cmake_flag="USE_RPC")
# Mark a test as requiring the MRVL Library
requires_mrvl = Feature("mrvl", "Marvell", cmake_flag="USE_MRVL")
# Mark a test as requiring Hexagon to run
requires_hexagon = Feature(
"hexagon",
"Hexagon",
cmake_flag="USE_HEXAGON",
target_kind_enabled="hexagon",
compile_time_check=hexagon._compile_time_check,
run_time_check=hexagon._run_time_check,
@@ -1124,12 +1085,18 @@ requires_x86_avx512 = Feature(
requires_x86_amx = Feature(
"x86_amx",
"x86 AMX Extensions",
run_time_check=lambda: _has_cpu_feat("amx-int8"),
"x86_amx", "x86 AMX Extensions", run_time_check=lambda: _has_cpu_feat("amx-int8")
)
def _cmake_flag_enabled(flag):
flag = tvm.support.libinfo().get(flag, "OFF")
# Because many of the flags can be library flags, we check if the
# flag is not disabled, rather than checking if it is enabled.
return flag.lower() not in ["off", "false", "0"]
def _parse_target_entry(entry):
"""Parse a target entry from TVM_TEST_TARGETS env var.
@@ -1138,6 +1105,8 @@ def _parse_target_entry(entry):
"""
entry = entry.strip()
if entry.startswith("{"):
import json # pylint: disable=import-outside-toplevel
return json.loads(entry)
return entry
@@ -1145,8 +1114,8 @@ def _parse_target_entry(entry):
def _tvm_test_targets():
target_str = os.environ.get("TVM_TEST_TARGETS", "").strip()
if target_str:
# Use dict instead of set for de-duplication so that the
# targets stay in the order specified.
# De-duplicate while preserving order. dict items can't be hashed
# directly, so use their str() form as the dedup key.
targets = []
seen = set()
for t in target_str.split(";"):
@@ -1155,9 +1124,10 @@ def _tvm_test_targets():
continue
parsed = _parse_target_entry(t)
key = str(parsed)
if key not in seen:
seen.add(key)
targets.append(parsed)
if key in seen:
continue
seen.add(key)
targets.append(parsed)
return targets
return DEFAULT_TEST_TARGETS
@@ -1219,7 +1189,7 @@ def requires_nvcc_version(major_version, minor_version=0, release_version=0):
installed version of NVCC is at least `(major_version,
minor_version, release_version)`.
This also marks the test as requiring a CUDA support.
This also marks the test as requiring a cuda support.
Parameters
----------
@@ -1255,14 +1225,14 @@ def requires_nvcc_version(major_version, minor_version=0, release_version=0):
return inner
def requires_cuda_compute_version(major_version, minor_version=0):
def requires_cuda_compute_version(major_version, minor_version=0, exact=False):
"""Mark a test as requiring at least a compute architecture
Unit test marked with this decorator will run only if the CUDA
compute architecture of the GPU is at least `(major_version,
minor_version)`.
This also marks the test as requiring a CUDA support.
This also marks the test as requiring a cuda support.
Parameters
----------
@@ -1287,7 +1257,7 @@ def requires_cuda_compute_version(major_version, minor_version=0):
compute_version_str = ".".join(str(v) for v in compute_version)
requires = [
pytest.mark.skipif(
compute_version < min_version,
compute_version < min_version or (exact and compute_version != min_version),
reason=f"Requires CUDA compute >= {min_version_str}, but have {compute_version_str}",
),
*requires_cuda.marks(),
@@ -1988,4 +1958,307 @@ def strtobool(val):
def main():
test_file = inspect.getsourcefile(sys._getframe(1))
sys.exit(pytest.main([test_file] + sys.argv[1:]))
sys.exit(pytest.main([test_file, *sys.argv[1:]]))
class CompareBeforeAfter:
"""Utility for comparing before/after of TIR transforms
A standard framework for writing tests that take a TIR PrimFunc as
input, apply a transformation, then either compare against an
expected output or assert that the transformation raised an error.
A test should subclass CompareBeforeAfter, defining class members
`before` / `Before`, `transform`, and `expected` / `Expected`. CompareBeforeAfter will
then use these members to define a test method and test fixture.
`transform` may be one of the following.
- An instance of `tvm.ir.transform.Pass`
- A method that takes no arguments and returns a `tvm.ir.transform.Pass`
- A pytest fixture that returns a `tvm.ir.transform.Pass`
`before` / `Before` may be any one of the following.
- An instance of `tvm.tirx.PrimFunc`. This is allowed, but is not
the preferred method, as any errors in constructing the
`PrimFunc` occur while collecting the test, preventing any other
tests in the same file from being run.
- An TVMScript function, without the ``@T.prim_func`` decoration.
The ``@T.prim_func`` decoration will be applied when running the
test, rather than at module import.
- A method that takes no arguments and returns a `tvm.tirx.PrimFunc`
- A pytest fixture that returns a `tvm.tirx.PrimFunc`
`expected` / `Expected` may be any one of the following. The type of
`expected` / `Expected` defines the test being performed. If `expected`
provides a `tvm.tirx.PrimFunc`, the result of the transformation
must match `expected`. If `expected` is an exception, then the
transformation must raise that exception type.
- Any option supported for `before` / `Before`.
- The `Exception` class object, or a class object that inherits
from `Exception`.
- A method that takes no arguments and returns `Exception` or a
class object that inherits from `Exception`.
- A pytest fixture that returns `Exception` or an class object
that inherits from `Exception`.
Examples
--------
.. code-block:: python
class TestRemoveIf(tvm.testing.CompareBeforeAfter):
transform = tvm.tirx.transform.Simplify()
def before(A: T.Buffer(1, "int32")):
if True:
A[0] = 42
else:
A[0] = 5
def expected(A: T.Buffer(1, "int32")):
A[0] = 42
"""
check_well_formed: bool = True
def __init_subclass__(cls):
assert len([getattr(cls, name) for name in ["before", "Before"] if hasattr(cls, name)]) <= 1
assert (
len([getattr(cls, name) for name in ["expected", "Expected"] if hasattr(cls, name)])
<= 1
)
for name in ["before", "Before"]:
if hasattr(cls, name):
cls.before = cls._normalize_before(getattr(cls, name))
break
for name in ["expected", "Expected"]:
if hasattr(cls, name):
cls.expected = cls._normalize_expected(getattr(cls, name))
break
if hasattr(cls, "transform"):
cls.transform = cls._normalize_transform(cls.transform)
@classmethod
def _normalize_ir_module(cls, func):
if isinstance(func, tvm.tirx.PrimFunc | tvm.IRModule):
def inner(self):
# pylint: disable=unused-argument
return func
elif cls._is_method(func):
def inner(self):
# pylint: disable=unused-argument
return func(self)
elif inspect.isclass(func):
def inner(self):
# pylint: disable=unused-argument
func_dict = {}
for name, method in func.__dict__.items():
if name.startswith("_"):
pass
elif isinstance(method, tvm.ir.function.BaseFunc):
func_dict[name] = method.with_attr("global_symbol", name)
else:
source_code = "@T.prim_func\n" + textwrap.dedent(inspect.getsource(method))
prim_func = tvm.script.from_source(
source_code, check_well_formed=self.check_well_formed
)
func_dict[name] = prim_func.with_attr("global_symbol", name)
return tvm.IRModule(func_dict)
else:
def inner(self):
# pylint: disable=unused-argument
source_code = "@T.prim_func\n" + textwrap.dedent(inspect.getsource(func))
return tvm.script.from_source(source_code, check_well_formed=self.check_well_formed)
return pytest.fixture(inner)
@classmethod
def _normalize_before(cls, func):
if hasattr(func, "_pytestfixturefunction"):
return func
else:
return cls._normalize_ir_module(func)
@classmethod
def _normalize_expected(cls, func):
if hasattr(func, "_pytestfixturefunction"):
return func
elif inspect.isclass(func) and issubclass(func, Exception):
def inner(self):
# pylint: disable=unused-argument
return func
return pytest.fixture(inner)
else:
return cls._normalize_ir_module(func)
@classmethod
def _normalize_transform(cls, transform):
def apply(module_transform):
def inner(obj):
if isinstance(obj, tvm.IRModule):
return module_transform(obj)
elif isinstance(obj, tvm.tirx.PrimFunc):
mod = tvm.IRModule({"main": obj})
mod = module_transform(mod)
return mod["main"]
else:
raise TypeError(f"Expected IRModule or PrimFunc, but received {type(obj)}")
return inner
if hasattr(transform, "_pytestfixturefunction"):
if not hasattr(cls, "_transform_orig"):
cls._transform_orig = transform
def inner(self, _transform_orig):
# pylint: disable=unused-argument
return apply(_transform_orig)
elif isinstance(transform, tvm.ir.transform.Pass):
def inner(self):
# pylint: disable=unused-argument
return apply(transform)
elif cls._is_method(transform):
def inner(self):
# pylint: disable=unused-argument
return apply(transform(self))
else:
raise TypeError(
"Expected transform to be a tvm.ir.transform.Pass, or a method returning a Pass"
)
return pytest.fixture(inner)
@staticmethod
def _is_method(func):
return callable(func) and "self" in inspect.signature(func).parameters
def test_compare(self, before, expected, transform):
"""Unit test to compare the expected TIR PrimFunc to actual"""
if inspect.isclass(expected) and issubclass(expected, Exception):
with pytest.raises(expected):
after = transform(before)
# This portion through pytest.fail isn't strictly
# necessary, but gives a better error message that
# includes the before/after.
before_str = before.script(name="before")
after_str = after.script(name="after")
pytest.fail(
msg=(
f"Expected {expected.__name__} to be raised from transformation, "
f"instead received TIR\n:{before_str}\n{after_str}"
)
)
elif isinstance(expected, tvm.tirx.PrimFunc | tvm.ir.IRModule):
after = transform(before)
try:
# overwrite global symbol so it doesn't come up in the comparison
if isinstance(after, tvm.tirx.PrimFunc):
after = after.with_attr("global_symbol", "main")
expected = expected.with_attr("global_symbol", "main")
tvm.ir.assert_structural_equal(after, expected)
except ValueError as err:
before_str = before.script(name="before")
after_str = after.script(name="after")
expected_str = expected.script(name="expected")
raise ValueError(
f"TIR after transformation did not match expected:\n"
f"{before_str}\n{after_str}\n{expected_str}"
) from err
else:
raise TypeError(
f"tvm.testing.CompareBeforeAfter requires the `expected` fixture "
f"to return either `Exception`, an `Exception` subclass, "
f"or an instance of `tvm.tirx.PrimFunc`. "
f"Instead, received {type(expected)}."
)
ml_dtypes_dict = {
"float8_e4m3fn": ml_dtypes.float8_e4m3fn,
"float8_e5m2": ml_dtypes.float8_e5m2,
"bfloat16": ml_dtypes.bfloat16,
"int4": ml_dtypes.int4,
}
def np_dtype_from_str(dtype: str) -> np.dtype:
"""Convert a string dtype to a numpy dtype."""
return np.dtype(ml_dtypes_dict[dtype]) if dtype in ml_dtypes_dict else np.dtype(dtype)
def generate_random_array(dtype: str, shape: tuple) -> np.ndarray:
"""
Generate a random array by generating random bits and casting to the target dtype.
Supported dtypes:
- "int8", "uint8", "float16", "float32", "bfloat16", "float8_e4m3fn", "float8_e5m2"
"""
try:
np_dtype = np_dtype_from_str(dtype)
except TypeError:
raise ValueError("Provided dtype is not a valid numpy dtype.")
# Determine the bit length for this dtype.
bit_length = np_dtype.itemsize * 8
# Choose an appropriate unsigned container type.
if bit_length <= 8:
container = np.uint8
elif bit_length <= 16:
container = np.uint16
elif bit_length <= 32:
container = np.uint32
elif bit_length <= 64:
container = np.uint64
else:
raise ValueError(f"Unsupported dtype bit length: {bit_length}")
# Generate random integers in the full range of the bit length.
random_ints = np.random.randint(0, 2**bit_length, size=shape, dtype=container)
# Reinterpret the bit pattern as the desired dtype.
res = random_ints.view(np_dtype)
with np.errstate(invalid="ignore"):
invalid_indices = np.where(~np.isfinite(res))
for idx in zip(*invalid_indices):
while True:
with np.errstate(invalid="ignore"):
if np.isfinite(res[idx]):
break
# Generate a new random value for this specific position
new_random_int = np.random.randint(0, 2**bit_length, size=1, dtype=container)
res[idx] = new_random_int.view(np_dtype)[0]
return res
python/tvm/tirx/__init__.py
+41 -39
View File
@@ -18,6 +18,11 @@
# pylint: disable=unused-import, redefined-builtin
"""Namespace for Tensor-level IR"""
import tvm.script
tvm.script.register_dialect("tirx", "tvm.tirx.script")
from tvm.ir import PrimExpr
from tvm.runtime import const
@@ -30,16 +35,16 @@ from .expr import Select, BufferLoad, ProducerLoad, Ramp, Broadcast, Shuffle
from .expr import Call, CallEffectKind, Let, IterVar, CommReducer
from .stmt import Stmt, Bind, AssertStmt, ForKind, For, While
from .stmt import (
BufferStore,
AllocBuffer,
AttrStmt,
DeclBuffer,
)
# Legacy alias: LetStmt was folded into Bind (which now accepts an optional body)
LetStmt = Bind
from .stmt import BufferStore, AllocBuffer, AttrStmt, DeclBuffer
from .stmt import SeqStmt
from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list
from .stmt import BufferRegion, MatchBufferRegion, SBlock, SBlockRealize
from .stmt import TilePrimitiveCall, ExecScopeStmt
from .function import PrimFunc, TensorIntrin, IndexMap
@@ -50,12 +55,7 @@ from .op import tvm_stack_alloca, tvm_stack_make_shape, tvm_stack_make_array
from .op import tvm_tuple, handle_add_byte_offset, tvm_struct_get, tvm_struct_set
from .op import address_of, lookup_param, assume, undef
from .op import continue_loop, break_loop
from .op import (
tvm_thread_allreduce,
type_annotation,
tvm_access_ptr,
tvm_throw_last_error,
)
from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr, tvm_throw_last_error
from .op import (
tvm_load_matrix_sync,
tvm_store_matrix_sync,
@@ -64,19 +64,9 @@ from .op import (
tvm_fill_fragment,
)
from .op import ptx_mma, ptx_mma_sp, mma_store, mma_fill
from .op import (
ptx_ldmatrix,
ptx_cp_async,
ptx_cp_async_bulk,
ptx_commit_group,
ptx_wait_group,
ptx_cp_async_barrier,
ptx_init_barrier_thread_count,
ptx_arrive_barrier,
ptx_arrive_barrier_expect_tx,
ptx_wait_barrier,
create_barriers,
)
from .op import ptx_mma_legacy, ptx_mma_sp_legacy, mma_store_legacy, mma_fill_legacy
from .op import ptx_ldmatrix, ptx_cp_async, ptx_cp_async_bulk, ptx_cp_async_bulk_shared_to_cluster
from .op import ptx_ldmatrix_legacy, ptx_cp_async_legacy
from .op import (
make_filled_simdgroup_matrix,
simdgroup_load,
@@ -91,18 +81,7 @@ from .op import cos, cosh, acos, acosh
from .op import tan, tanh, atan, atan2, atanh
from .op import bitwise_and, bitwise_not, bitwise_or, bitwise_xor
from .op import erf, sigmoid, sqrt, rsqrt, floor, ceil, hypot
from .op import (
trunc,
abs,
round,
nextafter,
nearbyint,
power,
pow,
popcount,
fmod,
if_then_else,
)
from .op import trunc, abs, round, nextafter, nearbyint, power, pow, popcount, fmod, if_then_else
from .op import likely, isnan, isnullptr, isfinite, isinf, copysign
from .op import div, indexdiv, indexmod, truncdiv, truncmod, floordiv, floormod, ceildiv, logaddexp
from .op import comm_reducer, min, max, sum
@@ -114,14 +93,37 @@ from .op import dp4a
from .op import ignore_loop_partition
from .generic import add, subtract, multiply
# 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
from . import analysis
from . import backend
from . import stmt_functor
from .build import build
from .pipeline import get_tir_pipeline, get_default_tir_pipeline
from .functor import PyStmtExprVisitor, PyStmtExprMutator
# Compiler-only submodules. Skip under `TVM_USE_RUNTIME_LIB=1` since they
# perform compiler-side FFI at module load (schema engine looks up
# `ir.RegisterOp`; codegen registry hooks the build pipeline).
from tvm.base import _RUNTIME_ONLY as _RUNTIME_ONLY_TIRX # pylint: disable=wrong-import-position
if not _RUNTIME_ONLY_TIRX:
# CUDA codegen registration. Each family module registers codegen via
# @register_codegen (hand-written ops) and ptx_intrinsic /
# cuda_helper_intrinsic (schema-declared ops); the schema declarations
# also inject Python wrappers into `tvm.tirx.op`. Must come before
# anything downstream that looks up wrappers or the codegen registry.
from .operator.intrinsics import cuda as _intrinsics_cuda
from .build import build
from .compilation_pipeline import (
get_tir_pipeline,
get_default_tir_pipeline,
)
import tvm.script
tvm.script.register_dialect("tirx", "tvm.tirx.script")
python/tvm/tirx/analysis/analysis.py
+24
View File
@@ -134,3 +134,27 @@ def verify_well_formed(obj: PrimFunc | IRModule, assert_mode: bool = True) -> bo
Whether it is a well-formed TIR function.
"""
return _ffi_api.VerifyWellFormed(obj, assert_mode) # type: ignore # pylint: disable=no-member
def verify_tirx_well_formed(
obj: PrimFunc | IRModule, assert_mode: bool = True, device_func: bool = False
) -> bool:
"""Verify if the given TIRX is well-formed.
Parameters
----------
obj: Union[tvm.tirx.PrimFunc, tvm.ir.IRModule]
The function or module to be verified.
assert_mode: bool
The indicator if it raises an error when the function is not well-formed.
device_func: bool
The indicator if it is a device function.
Returns
-------
result: bool
Whether it is a well-formed TIRX function.
"""
return _ffi_api.VerifyTIRxWellFormed(obj, assert_mode, device_func) # type: ignore # pylint: disable=no-member
python/tvm/tirx/bench.py
+657
View File
@@ -0,0 +1,657 @@
# 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.
import argparse
import os
import re
import subprocess
import sys
import time
from collections.abc import Mapping
from enum import Enum
import numpy as np
import torch
import triton.profiler as proton
import tvm_ffi
import tvm
from tvm.contrib import nvcc
from tvm.script import tirx as Tx
def is_running_under_pytest():
"""Check if the code is being executed within a pytest session."""
return "PYTEST_CURRENT_TEST" in os.environ
def setup():
parser = argparse.ArgumentParser()
parser.add_argument("--dump-ptx", type=str, help="Dump PTX code to specified file")
parser.add_argument("--dump-source", action="store_true", help="Dump source code")
args = parser.parse_args()
if args.dump_ptx:
@tvm_ffi.register_global_func("tvm_callback_cuda_compile", override=True)
def tvm_callback_cuda_compile(code, target):
ptx = nvcc.compile_cuda(code, target_format="ptx")
with open(args.dump_ptx, "w", encoding="utf-8") as f:
f.write(ptx.decode())
return ptx
return args
_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}.
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.
Returns (impl_times, baseline_errors) where:
impl_times: {str: float} — impl name to avg time in ms
baseline_errors: {str: str} — impl name to error message
"""
impl = None
results = {}
baseline_errors = {}
for raw in text.splitlines():
line = _ANSI_RE.sub("", raw).rstrip()
if not line:
continue
if line.startswith("BASELINE_ERROR:"):
parts = line.split(":", 2)
if len(parts) >= 3:
baseline_errors[parts[1].strip()] = parts[2].strip()
continue
# Depth-1 impl header: starts with tree drawing chars
if line and line[0] in "\u251c\u2514": # ├ └
parts = line.split("\u2500", 1)[-1].split() # split on ─
if len(parts) >= 2:
impl = parts[1]
else:
impl = None
continue
# Depth-2 kernel: contains tree drawing chars at deeper indent
if impl and ("\u251c\u2500" in line or "\u2514\u2500" in line): # ├─ └─
parts = line.split("\u2500", 1)[-1].split()
if len(parts) >= 2:
name = parts[1]
if (
"vectorized_elementwise_kernel" in name
or "elementwise_kernel_with_index" in name
):
continue
try:
t = float(parts[0]) * value_scale
results[impl] = max(results.get(impl, 0), t)
except ValueError:
pass
return results, baseline_errors
class ProtonContext:
"""Context manager for Proton profiling sessions.
Always captures proton-viewer output and parses impl times so that
get_impl_times() / get_baseline_errors() work after exiting the context.
The proton tree is printed to **stdout** by default (visible on screen
when running kernels interactively). When the environment variable
``TIRX_BENCH_JSON=1`` is set (done automatically by ``--json`` mode),
the tree goes to **stderr** instead so it does not corrupt the JSON on
stdout.
"""
def __init__(
self,
name="kernel",
hook="triton",
debug=False,
nsight=False,
metric="avg_time/us",
metric_scale=1e-3,
):
self.name = name
self.hook = hook
self.debug = debug
self.nsight = nsight
self.metric = metric
self.metric_scale = metric_scale
self._impl_times = {}
self._baseline_errors = {}
def __enter__(self):
if not is_running_under_pytest() and not self.debug and not self.nsight:
proton.start(self.name, hook=self.hook)
proton.deactivate()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if not is_running_under_pytest() and not self.debug and not self.nsight:
proton.finalize()
hatchet = f"{self.name}.hatchet"
result = subprocess.run(
["proton-viewer", "-m", self.metric, hatchet],
capture_output=True,
text=True,
check=False,
)
if result.returncode == 0:
self._impl_times, self._baseline_errors = _parse_proton_tree(
result.stdout, value_scale=self.metric_scale
)
out = sys.stderr if os.environ.get("TIRX_BENCH_JSON") else sys.stdout
print(result.stdout, file=out, end="")
else:
print(
f"proton-viewer failed (rc={result.returncode}): {result.stderr}",
file=sys.stderr,
)
if os.path.exists(hatchet):
os.remove(hatchet)
def get_impl_times(self):
"""Return {impl_name: avg_time_ms} parsed from proton-viewer output."""
return dict(self._impl_times)
def get_baseline_errors(self):
"""Return {impl_name: error_message} from BASELINE_ERROR lines."""
return dict(self._baseline_errors)
def _get_l2_cache_bytes():
"""Query L2 cache size from the current CUDA device, fallback to 128MB."""
try:
props = torch.cuda.get_device_properties(torch.cuda.current_device())
if hasattr(props, "l2_cache_size") and props.l2_cache_size > 0:
return props.l2_cache_size
except Exception:
pass
return 128 * 1024 * 1024 # 128MB default (B200)
def _tensor_bytes(args, _seen=None):
"""Sum the byte size of all torch/tvm tensors in a nested value."""
if _seen is None:
_seen = set()
total = 0
if isinstance(args, list | tuple):
for a in args:
total += _tensor_bytes(a, _seen)
elif isinstance(args, Mapping):
for a in args.values():
total += _tensor_bytes(a, _seen)
elif isinstance(args, torch.Tensor):
key = ("torch", args.device.type, args.device.index, int(args.data_ptr()))
if key not in _seen:
_seen.add(key)
total += args.nelement() * args.element_size()
elif hasattr(args, "numpy"): # tvm.runtime.NDArray
try:
key = ("tvm", int(args.handle.value))
except Exception:
key = ("tvm", id(args))
if key not in _seen:
_seen.add(key)
try:
total += int(np.prod(args.shape)) * np.dtype(str(args.dtype)).itemsize
except Exception:
total += args.numpy().nbytes
return total
def tensor_bytes(*values):
"""Return unique torch/tvm tensor bytes for kernel-owned byte accounting.
The benchmark driver does not use this implicitly. Kernel benchmark
factories may call it when their invocation footprint is exactly the set of
tensors in ``values``.
"""
if len(values) == 1:
return _tensor_bytes(values[0])
return _tensor_bytes(values)
def _compute_group_count(input_bytes, l2_bytes=None):
"""Return TK-style input-group count from one invocation's byte footprint."""
if input_bytes <= 0:
return 1
if l2_bytes is None:
l2_bytes = _get_l2_cache_bytes()
threshold = l2_bytes * 3
if input_bytes >= threshold:
return 1
return int(threshold // input_bytes) + 1
def _make_bench_input(input_factory):
value = input_factory()
if not isinstance(value, tuple) or len(value) != 2:
raise TypeError("input_factory must return (case, input_bytes)")
case, input_bytes = value
try:
input_bytes = int(input_bytes)
except (TypeError, ValueError) as err:
raise TypeError("input_factory input_bytes must be an integer") from err
if input_bytes < 0:
raise ValueError("input_factory input_bytes must be non-negative")
return case, input_bytes
def prepare_input_groups(input_factory, l2_bytes=None):
"""Materialize TK-style input groups from a single-group factory.
``input_factory`` must return ``(case, input_bytes)``. ``case`` is passed
back to every benchmark function unchanged. ``input_bytes`` defines one
invocation's L2-eviction footprint and is intentionally owned by the kernel
benchmark harness instead of inferred here.
"""
if not callable(input_factory):
raise TypeError("input_factory must be callable")
if l2_bytes is None:
l2_bytes = _get_l2_cache_bytes()
sample, input_bytes = _make_bench_input(input_factory)
num_groups = _compute_group_count(input_bytes, l2_bytes)
groups = [sample]
for _ in range(num_groups - 1):
case, _ = _make_bench_input(input_factory)
groups.append(case)
return groups, {
"num_groups": num_groups,
"input_bytes": input_bytes,
"l2_bytes": l2_bytes,
"l2_eviction_factor": 3,
"flush_l2": False,
}
def _bench_event_groups(funcs, groups, warmup, repeat, cooldown_s):
num_groups = len(groups)
results = {}
for idx, (name, func) in enumerate(funcs.items()):
if idx > 0:
time.sleep(cooldown_s)
for i in range(warmup):
func(groups[i % num_groups])
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
torch.cuda.synchronize()
start_event.record()
for i in range(repeat):
func(groups[i % num_groups])
end_event.record()
torch.cuda.synchronize()
results[name] = start_event.elapsed_time(end_event) / repeat
time.sleep(cooldown_s)
return results
def _bench_proton_groups(funcs, groups, warmup, repeat, cooldown_s, proton_name, debug, nsight):
num_groups = len(groups)
with ProtonContext(proton_name, debug=debug, nsight=nsight) as ctx:
for idx, (name, func) in enumerate(funcs.items()):
if idx > 0:
time.sleep(cooldown_s)
for i in range(warmup):
func(groups[i % num_groups])
torch.cuda.synchronize()
if not is_running_under_pytest() and not debug and not nsight:
proton.activate()
with proton.scope(name, metrics={}):
for i in range(repeat):
func(groups[i % num_groups])
proton.deactivate()
else:
for i in range(repeat):
func(groups[i % num_groups])
torch.cuda.synchronize()
time.sleep(cooldown_s)
return ctx.get_impl_times(), ctx.get_baseline_errors()
def _flush_l2_legacy(flush_l2_size):
if flush_l2_size > 0:
torch.empty(flush_l2_size, dtype=torch.int, device="cuda").zero_()
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()
for _ in range(repeat):
_flush_l2_legacy(flush_l2_size)
func()
end_event.record()
if not is_running_under_pytest() and not debug and not nsight:
proton.activate()
with proton.scope(proton_name, metrics={}):
timed_loop()
proton.deactivate()
else:
timed_loop()
torch.cuda.synchronize()
return start_event.elapsed_time(end_event) / repeat
def bench(
funcs,
input_factory=None,
warmup=500,
repeat=100,
cooldown_s=1.0,
timer="proton",
proton_name="kernel",
l2_bytes=None,
debug=False,
nsight=False,
flush_l2_size=int(8e8 // 4),
):
"""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)``.
Parameters
----------
funcs : dict[str, callable]
Map of implementation name to callable. Each callable receives one
``case`` returned by ``input_factory``.
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.
cooldown_s : float
Seconds to sleep between impls for thermal cooldown.
timer : {"event", "proton"}
Timing backend.
Returns
-------
dict
``{"impls": {name: ms}, "errors": {}, "timer": ..., ...}``.
"""
if repeat <= 0:
raise ValueError("repeat must be positive")
if warmup < 0:
raise ValueError("warmup must be non-negative")
if timer not in {"event", "proton"}:
raise ValueError(f"unsupported timer {timer!r}; expected event or proton")
if callable(funcs) and input_factory is None:
return _bench_legacy_callable(
funcs,
warmup=warmup,
repeat=repeat,
proton_name=proton_name,
debug=debug,
nsight=nsight,
flush_l2_size=flush_l2_size,
)
if input_factory is None:
raise TypeError("input_factory is required when funcs is a mapping")
if not isinstance(funcs, Mapping) or not funcs:
raise TypeError("funcs must be a non-empty mapping of name to 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")
inputs, protocol = prepare_input_groups(input_factory, l2_bytes=l2_bytes)
num_groups = len(inputs)
if num_groups == 0:
return {
"impls": {},
"errors": {},
"timer": timer,
"benchmark_protocol": {
**protocol,
"warmup": warmup,
"repeat": repeat,
"cooldown_s": cooldown_s,
"order": list(funcs.keys()),
},
}
errors = {}
if timer == "event":
impls = _bench_event_groups(funcs, inputs, warmup, repeat, cooldown_s)
else:
impls, errors = _bench_proton_groups(
funcs, inputs, warmup, repeat, cooldown_s, proton_name, debug, nsight
)
return {
"impls": impls,
"errors": errors,
"timer": timer,
"benchmark_protocol": {
**protocol,
"warmup": warmup,
"repeat": repeat,
"cooldown_s": cooldown_s,
"order": list(funcs.keys()),
},
}
# utils for tg4perfetto profiler, adapted from https://github.com/flashinfer-ai/flashinfer
class EventType(Enum):
kBegin = 0
kEnd = 1
kInstant = 2
kFinalize = 3
def decode_tag(tag, num_groups):
block_group_tag = tag >> 12
event_idx = (tag >> 2) & 0x3FF
event_type = tag & 0x3
return (block_group_tag // num_groups, block_group_tag % num_groups, event_idx, event_type)
def export_to_perfetto_trace(
profiler_buffer: np.ndarray, file_name: str, event_type_names: list[str]
) -> None:
if is_running_under_pytest():
return
import torch
# pip install git+https://github.com/ihavnoid/tg4perfetto.git
from tg4perfetto import TraceGenerator
profiler_buffer_host = torch.tensor(profiler_buffer)
num_blocks, num_groups = profiler_buffer_host[:1].view(dtype=torch.int32)
num_blocks = int(num_blocks)
num_groups = int(num_groups)
tgen = TraceGenerator(file_name)
tid_map = {}
track_map = {}
finish_idx = set()
for block_idx in range(num_blocks):
pid = tgen.create_group(f"block_{block_idx}")
for group_idx in range(num_groups):
tid = pid.create_group(f"group_{group_idx}")
tid_map[(block_idx, group_idx)] = tid
for i in range(1, len(profiler_buffer_host)):
if profiler_buffer_host[i] == 0:
continue
tag, timestamp = profiler_buffer_host[i : i + 1].view(dtype=torch.uint32)
tag = int(tag)
timestamp = int(timestamp)
block_idx, group_idx, event_idx, event_type = decode_tag(tag, num_groups)
if event_type == EventType.kFinalize.value:
finish_idx.add((block_idx, group_idx))
if len(finish_idx) == num_blocks * num_groups:
break
else:
if (block_idx, group_idx) in finish_idx:
continue
event = event_type_names[event_idx]
tid = tid_map[(block_idx, group_idx)]
if (block_idx, group_idx, event_idx) in track_map:
track = track_map[(block_idx, group_idx, event_idx)]
else:
track = tid.create_track()
track_map[(block_idx, group_idx, event_idx)] = track
if event_type == EventType.kBegin.value:
track.open(timestamp, event)
elif event_type == EventType.kEnd.value:
track.close(timestamp)
elif event_type == EventType.kInstant.value:
track.instant(timestamp, event)
tgen.flush()
@Tx.meta_class
class CudaProfiler:
"""A lightweight wrapper around Tx.timer_* CUDA intrinsics.
Stores repeated arguments used by timer_init/start/end/finalize so users can
call concise methods in kernels. Intended to mirror Pipeline/TileScheduler helpers.
When ``profiler_enabled`` is False (or a false-y PrimExpr), calls to
``init/start/end/finalize`` become no-ops. This allows constructing a
profiler unconditionally and eliminating external ``if PROFILER_ON:`` guards.
"""
def __init__(
self,
profiler_buffer: Tx.Buffer,
write_stride: int,
num_groups: int,
default_leader: None | tvm.tirx.PrimExpr | bool = None,
profiler_enabled: bool | tvm.tirx.PrimExpr = True,
):
self.buffer = profiler_buffer
self.write_stride = write_stride
self.num_groups = num_groups
self.default_leader = default_leader
# Accept either a Python bool or a PrimExpr; normalize simple bools to Tx.bool
# so we can use it uniformly inside macros for conditional emission.
if isinstance(profiler_enabled, bool | np.bool_):
self.profiler_enabled = Tx.bool(bool(profiler_enabled))
else:
# Assume PrimExpr-like input; use as-is
self.profiler_enabled = profiler_enabled # type: ignore[assignment]
self.profiler_tag = Tx.alloc_buffer([1], "uint64", scope="local", align=8)
self.profiler_write_offset = Tx.alloc_buffer([1], "uint32", scope="local", align=8)
def _leader(self, leader: None | tvm.tirx.PrimExpr | bool):
if leader is not None:
if isinstance(leader, bool | np.bool_):
return Tx.bool(bool(leader))
return leader
if self.default_leader is not None:
return self.default_leader
return Tx.bool(True)
@Tx.inline
def init(self, group_id: tvm.tirx.PrimExpr):
if self.profiler_enabled:
Tx.timer_init_cuda(
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
self.num_groups,
group_id,
)
@Tx.inline
def start(self, event_type: Enum, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
Tx.timer_start_cuda(
event_type,
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
self.write_stride,
self._leader(leader),
)
@Tx.inline
def end(self, event_type: Enum, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
Tx.timer_end_cuda(
event_type,
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
self.write_stride,
self._leader(leader),
)
@Tx.inline
def finalize(self, leader: None | tvm.tirx.PrimExpr | bool = None):
if self.profiler_enabled:
Tx.timer_finalize_cuda(
self.buffer.data,
self.profiler_tag.data,
self.profiler_write_offset.data,
self.write_stride,
self._leader(leader),
)
python/tvm/tirx/buffer.py
+273 -75
View File
@@ -16,6 +16,7 @@
# under the License.
"""Abstraction for array data structures."""
import functools
from numbers import Integral
import tvm_ffi
@@ -176,6 +177,18 @@ class Buffer(Object, Scriptable):
"""
return _ffi_api.BufferGetFlattenedBuffer(self) # type: ignore
def with_allocated_addr(self, allocated_addr):
"""Return a new buffer with the allocated address."""
return _ffi_api.BufferWithAllocatedAddr(self, allocated_addr) # type: ignore
def with_dtype(self, dtype):
"""Return a new buffer with the dtype."""
return _ffi_api.BufferWithDtype(self, dtype) # type: ignore
def with_data(self, data):
"""Return a new buffer with the data."""
return _ffi_api.BufferWithData(self, data) # type: ignore
def offset_of(self, indices):
"""Determine the offset of the provided indices in the flattened buffer.
@@ -193,6 +206,252 @@ class Buffer(Object, Scriptable):
"""
return _ffi_api.BufferOffsetOf(self, indices) # type: ignore
@property
def byte_offset(self):
"""Get the byte offset of the buffer."""
return self.elem_offset * tvm.DataType(self.dtype).bits // 8
def elem_offset_of(self, indices, inner=True):
"""Get the element offset of the buffer at the given indices.
Note that indices subject to buffer's layout mapping.
Parameters
----------
indices : Union[PrimExpr, List[PrimExpr]]
The indices of the element in the original buffer.
inner : bool, optional
If False, the offset is relative to the original buffer.
Default is True.
Returns
-------
offset: PrimExpr
The element offset of the buffer at the given indices.
"""
if inner:
return _ffi_api.BufferOffsetOfp(self, indices)
return self.elem_offset + _ffi_api.BufferOffsetOfp(self, indices)
def byte_offset_of(self, indices, inner=True):
"""Get the byte offset of the buffer at the given indices.
Note that indices subject to buffer's layout mapping.
Parameters
----------
indices : Union[PrimExpr, List[PrimExpr]]
The indices of the element in the original buffer.
inner : bool, optional
If False, the offset is relative to the original buffer.
Default is True.
Returns
-------
offset: PrimExpr
The byte offset of the buffer at the given indices.
"""
return self.elem_offset_of(indices, inner) * tvm.DataType(self.dtype).bits // 8
def is_scalar(self, alloc_or_decl=True):
"""Check if the buffer is a scalar.
Parameters
----------
alloc_or_decl : bool, optional
Whether to consider alloc_scalar and decl_scalar as scalar. True for alloc_scalar,
False for decl_scalar.
Returns
-------
bool: True if the buffer is a scalar, False otherwise.
"""
return _ffi_api.BufferIsScalar(self, alloc_or_decl)
def ptr_to(self, indices):
"""Get the pointer to the buffer at the given indices (logical indices).
Note that the bufferload inside requires LowerTIPp pass to apply the layout to get the physical indices.
""" # noqa: E501
assert len(indices) == len(self.shape), (
f"The number of indices {indices} does not match the shape of the buffer {self.shape}"
)
return tvm.tirx.address_of(self[tuple(indices)])
def view(self, *args, **kwargs) -> "Buffer":
"""Creates a new view of the buffer. (used by parser)
Supported signatures are ``view(*shape, layout=None)``, where shape can contain
``-1`` to indicate that the dimension size is auto-inferred, and
``view(dtype: Union[str, tvm.DataType])``.
Returns
-------
view : DeclBufferFrame
The corresponding view buffer.
"""
def _infer_shape(shape):
shape = list(shape)
if -1 in shape and shape.count(-1) == 1:
size = functools.reduce(lambda x, y: x * y, self.shape)
n_size = functools.reduce(lambda x, y: x * y, [s for s in shape if s != -1], 1)
shape[shape.index(-1)] = size // n_size
else:
# Only validate the shape product when both old and new shapes
# are fully concrete: a PrimExpr `==` returns an `EQ` node, not
# a Python bool, and `assert <PrimExpr>` raises (no __bool__).
if all(isinstance(s, int) for s in shape) and all(
isinstance(s, int) for s in self.shape
):
assert functools.reduce(lambda x, y: x * y, shape) == functools.reduce(
lambda x, y: x * y, self.shape
), (
"The shape of the buffer "
+ str(self.shape)
+ " and the new shape "
+ str(shape)
+ " are not compatible"
)
return shape
if len(args) == 1 and isinstance(args[0], str | tvm.DataType) and not kwargs:
cast_dtype = tvm.DataType(args[0])
cur_dtype = tvm.DataType(self.dtype)
if cast_dtype.bits > cur_dtype.bits:
# cast up
assert cast_dtype.bits % cur_dtype.bits == 0
ratio = cast_dtype.bits // cur_dtype.bits
layout = self.layout.pack(ratio)
shape = [s for s in self.shape[:-1]] + [self.shape[-1] // ratio]
new_elem_offset = self.elem_offset // ratio
else:
# cast down
assert cur_dtype.bits % cast_dtype.bits == 0
ratio = cur_dtype.bits // cast_dtype.bits
layout = self.layout.unpack(ratio)
shape = [s for s in self.shape[:-1]] + [self.shape[-1] * ratio]
new_elem_offset = self.elem_offset * ratio
return tvm.tirx.script.builder.decl_buffer(
shape,
cast_dtype,
self.data,
self.strides,
new_elem_offset,
None,
self.scope(),
self.data_alignment,
self.offset_factor,
"",
self.axis_separators,
layout,
)
else:
# --- Signature 1: view(*shape, **opts) ---
# Check if all positional args are integers/PrimExprs with dtype int32 or int64 (the shape) # noqa: E501
shape = args
assert all(
isinstance(arg, int)
or (isinstance(arg, PrimExpr) and arg.dtype in ["int32", "int64"])
for arg in shape
), "shape must be a list of integers or PrimExprs with dtype int32 or int64"
# Safely get optional keyword arguments
layout = kwargs.get("layout", None)
# Assert there are no other kwargs
assert set(kwargs.keys()).issubset({"layout"}), (
f"Unsupported kwargs for view: {set(kwargs.keys()) - {'layout'}}"
)
if layout is None:
shape = _infer_shape(shape)
return tvm.tirx.script.builder.decl_buffer(
shape,
self.dtype,
self.data,
self.strides,
self.elem_offset,
None,
self.scope(),
self.data_alignment,
self.offset_factor,
"",
self.axis_separators,
self.layout if layout is None else layout,
)
def local(self, *shape, layout=None) -> "Buffer":
"""Create a thread-local view of this buffer.
When called with no shape arguments, auto-infers a 1D shape from
the layout's non-thread component (i.e. ``layout.storage().shard``).
Parameters
----------
shape : tuple of Expr
The shape of the local view for indexing. If omitted, a 1D
shape is computed automatically.
layout : optional
Override layout. If None, uses the storage layout
(parent layout with thread axes removed).
Returns
-------
local : DeclBufferFrame
The corresponding local buffer.
"""
if not shape:
local_layout = self.layout.storage()
total = functools.reduce(
lambda x, y: x * y, [it.extent for it in local_layout.shard], 1
)
shape = (total,)
return tvm.tirx.script.builder.decl_buffer(
shape,
self.dtype,
self.data,
self.strides,
self.elem_offset,
None,
self.scope(),
self.data_alignment,
self.offset_factor,
"",
self.axis_separators,
self.layout.storage() if layout is None else layout,
)
def permute(self, *dims) -> "Buffer":
"""Permute the dimensions of the buffer.
Parameters
----------
dims : tuple of int
The permutation of dimensions.
Returns
-------
permuted : DeclBufferFrame
The buffer with permuted dimensions.
"""
new_shape = [self.shape[d] for d in dims]
new_layout = self.layout.permute_dims(list(dims))
return tvm.tirx.script.builder.decl_buffer(
new_shape,
self.dtype,
self.data,
self.strides,
self.elem_offset,
None,
self.scope(),
self.data_alignment,
self.offset_factor,
"",
self.axis_separators,
new_layout,
)
def __getitem__(self, indices):
from ..arith import Analyzer # pylint: disable=import-outside-toplevel
from .expr import BufferLoad, Ramp, const # pylint: disable=import-outside-toplevel
@@ -201,9 +460,12 @@ class Buffer(Object, Scriptable):
if not isinstance(indices, tuple | list):
indices = [indices]
has_slice = any(isinstance(i, slice) for i in indices)
has_step = any(isinstance(i, slice) and i.step is not None for i in indices)
has_step = any(
isinstance(i, slice) and (i.step is not None and i.step != 1) for i in indices
)
has_implicit_slice = len(indices) < len(self.shape)
analyzer = Analyzer()
if has_slice and not has_step:
if (has_slice and not has_step) or has_implicit_slice:
region = []
for i, index in enumerate(indices):
if isinstance(index, slice):
@@ -216,6 +478,9 @@ class Buffer(Object, Scriptable):
index, const(1, index.dtype) if isinstance(index, PrimExpr) else 1
)
)
if has_implicit_slice:
for i in range(len(indices), len(self.shape)):
region.append(Range.from_min_extent(0, self.shape[i]))
return BufferRegion(self, region)
else:
expr_indices = []
@@ -250,82 +515,11 @@ def decl_buffer(
buffer_type="",
axis_separators=None,
span=None,
layout="default",
):
"""Declare a new symbolic buffer.
Normally buffer is created automatically during lower and build.
This is only needed if user want to specify their own buffer layout.
See the note below for detailed discussion on usage of buffer.
Parameters
----------
shape : tuple of Expr
The shape of the buffer.
dtype : str, optional
The data type of the buffer.
name : str, optional
The name of the buffer.
data : tirx.Var, optional
The data pointer in the buffer.
strides: array of Expr
The stride of the buffer.
elem_offset: Expr, optional
The beginning offset of the array to data.
In terms of number of elements of dtype.
scope: str, optional
The storage scope of the buffer, if not global.
If scope equals empty string, it means it is global memory.
data_alignment: int, optional
The alignment of data pointer in bytes.
If -1 is passed, the alignment will be set to TVM's internal default.
offset_factor: int, optional
The factor of elem_offset field, when set,
elem_offset is required to be multiple of offset_factor.
If 0 is pssed, the alignment will be set to 1.
if non-zero is passed, we will created a Var for elem_offset if elem_offset is not None.
buffer_type: str, optional, {"", "auto_broadcast"}
auto_broadcast buffer allows one to implement broadcast computation
without considering whether dimension size equals to one.
TVM maps buffer[i][j][k] -> buffer[i][0][k] if dimension j's shape equals 1.
axis_separators : list of int, optional
If passed, a list of separators between groups of axes,
each of which is flattened to an output axis. For flat
memory spaces, should either be None, or an empty list.
span: Optional[Span]
The location of the decl_buffer creation in the source.
Returns
-------
buffer : tvm.tirx.Buffer
The created buffer
Note
----
Buffer data structure reflects the DLTensor structure in dlpack.
While DLTensor data structure is very general, it is usually helpful
to create function that only handles specific case of data structure
and make compiled function benefit from it.
If user pass strides and elem_offset is passed as None
when constructing the function, then the function will be specialized
for the DLTensor that is compact and aligned.
If user pass a fully generic symbolic array to the strides,
then the resulting function becomes fully generic.
"""
# pylint: disable=import-outside-toplevel
from .expr import Var
from .layout import S, TileLayout
shape = (shape,) if isinstance(shape, PrimExpr | Integral) else shape
dtype = "float32" if dtype is None else dtype
@@ -334,6 +528,9 @@ def decl_buffer(
if axis_separators is None:
axis_separators = []
if layout == "default":
layout = TileLayout(S[tuple(shape)]) if shape else None
if offset_factor != 0 and elem_offset is None:
shape_dtype = shape[0].dtype if shape and hasattr(shape[0], "dtype") else "int32"
elem_offset = Var(f"{name}_elem_offset", shape_dtype)
@@ -354,6 +551,7 @@ def decl_buffer(
buffer_type,
axis_separators,
span,
layout,
)

Some files were not shown because too many files have changed in this diff Show More