Bohan Hou
57c638fc7c
## Summary Follow-up work on top of the TIRx infrastructure bring-up (#19581). It extends the TIRx operator-dispatch, codegen, and TVMScript surfaces with the next batch of low-level programming features for Blackwell-class GPUs, while keeping `s_tir` script support intact. ## Main Changes - **op-dispatch**: warp `ldmatrix`/`stmatrix` copy dispatch; split CUDA copy into register / gmem-smem / `ldgsts` paths; `tcgen05.ld/st` `.16x{64,128,256}b` dispatch with a factory and M=128 layout; element-wise broadcast at the layout level with a copy vec-alignment fix. - **gemm**: CUDA synchronous `mma.sync` tensor-core dispatch; accept a Layout F C operand for M=64 MMAs. - **op**: add the `permute_layout` primitive (replaces `permute_dims`). - **tvmscript**: add the `Tx.jit` decorator, `Tx.constexpr` compile-time params, and `Tx.wg_reg_tile`. - **lower-tirx**: introduce the `Tx.device_entry()` marker (replacing `ScopeKind::kKernel`); canonical thread filters that drop the `Tx.filter` wrapper. - **codegen**: add a typed-pointer byte-offset intrinsic; remove the `entry_cluster_sync` codegen attribute. ## Validation - `pre-commit run` (changed files) — clean - `ninja -C build -j$(nproc)` — builds - `pytest tests/python/tirx/ -n 16` - `1997 passed, 39 skipped, 3 xpassed` - `python -m pytest tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm pytest tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `630 passed, 25 skipped, 8 xfailed, 1 xpassed` ## Local CI Notes Several full CI-equivalent jobs are not locally reproducible because this machine is missing parts of the Apache TVM CI environment (e.g., specific `llvm-config` versions, Vulkan, ROCm, ARM/QEMU cross-toolchain, and web/wasm components). The Blackwell/Trainium kernel tests are maintained downstream and are intentionally not part of this PR.
44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
# 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 pytest
|
|
|
|
from tvm.tirx.exec_scope import ExecScope
|
|
|
|
|
|
def test_exec_scope_create():
|
|
def is_trivial_scope(scope, name):
|
|
return isinstance(scope, ExecScope) and scope.name == name
|
|
|
|
thread = ExecScope("thread")
|
|
warp = ExecScope("warp")
|
|
wg = ExecScope("warpgroup")
|
|
cta = ExecScope("cta")
|
|
cluster = ExecScope("cluster")
|
|
|
|
assert is_trivial_scope(thread, "thread")
|
|
assert is_trivial_scope(warp, "warp")
|
|
assert is_trivial_scope(wg, "warpgroup")
|
|
assert is_trivial_scope(cta, "cta")
|
|
assert is_trivial_scope(cluster, "cluster")
|
|
|
|
with pytest.raises(Exception, match="Unknown scope kind name"):
|
|
ExecScope("aaa")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_exec_scope_create()
|