`tests/python/codegen/test_target_codegen_aarch64.py` cross-compiles
AArch64 SVE kernels and regex-matches the generated assembly for
specific instruction forms and counts. Several of those assertions
encode the exact code shape produced by the LLVM versions used in CI
(15-17). On a TVM built against LLVM 20 the tests fail, even though the
emitted IR is correct (+sve target-features and vscale_range(1,16) are
present) -- the difference is entirely inside LLVM's loop vectorizer /
cost model, not in TVM's codegen.
Since LLVM 19, SVE/SVE2 are optional extensions of Armv9.0-A
(llvm/llvm-project#96007), so "+v9a" no longer implies "+sve" and
CodeGenAArch64 does not add the vscale_range() function attribute. Gate
the expectation in test_vscale_range_function_attribute on
llvm_version_major() < 19 so the test passes on both older LLVM (15-17,
as used in CI) and LLVM 19+.
## 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 PR brings up the tirx namespace. We have been spliting out the
original tir namespace to include high-level component s_tir and this PR
updates the remaining low-level part as tirx namespace
This PR phases out legacy target string format in favor of the json
style format that is more well formed. It also simplfies our overall
code in handling multiple formats.
This PR migrates all the codegen tests to explicitly using tvmscript
instead of indirectly via s_tir.Schedule. They makes the test surface
more unit, contains less dep and more maintainable.
This PR formalizes original runtime::Module into ffi
as ffi.Module and cleans the APIs around it.
The goal is to stablize the Module API as extra API that can benefit the overall
ffi interactions. We also refactors the c++ code that depends on the Module.
* [REFACTOR] Phase out te.schedule python components
This PR phases out te.schedule python components.
te.compute is kept around for future usages.
tir.Schedule is a more modern version of the scheduling that we can use onwards.
Doing so also helps us to cleanup the testcases that relies on
explicit full build and execution. As we move future unit testcases
towards structural equality based unit tests.
* Simplify CI to focus on UT
The main rationale is that we should only have very few target
dependent UT in tests/python/codegen and possible
a new category in future for op-level integration if needed.
* Re-enable wasm
* fix lint
* remove hybrid,sparse autodoc and remove tests
---------
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
This PR starts the step 0 to phase out relay from the current
development main branch. This PR focuses on the python
components of relay, autotvm, auto_scheduler. To make the change
manageable, we will also do followup steps on te.Schedule and
c++ components in followup PRs.
To continue support community members who depends on
legacy flows, the [v0.19.0](https://github.com/apache/tvm/tree/v0.19.0)
branch will continue contain these components.
As noted in [discussion on phasing out legacy components](https://discuss.tvm.apache.org/t/phasing-out-legacy-components/17703/30),
this would help us to do two purposes:
- By removing outdated or redundant elements, we can significantly
reduce complexity and improve maintainability.
- Unify our focus: Concentrating our efforts on the new unity flow
will allow for more efficient development and innovation.
It is also a good opportunity for us to revisit and reduce CI time.
The past relay legacy flow contains a lot of end to end tests that
requires hardware resources to run and causing long CI time.
Moving onwards, we can focus more on unit-tests that focuses
on structural equality and runs within seconds, while be mindful
about tests that requires hardware resources (by restricting them
to specific folders and CI nightly in some cases).
---
Co-authored-by: Siyuan Feng <hzfengsy@sjtu.edu.cn>
Prior to this commit, the matmul and conv2d schedules required padding
of the inputs to some multiple of vscale and a final "unpadding" stage.
Instead, we can leverage predicated operations to avoid the
the requirement for padding. Both the transpose interleave and outer
product fp32 intrinsics are updated to use predication. The
`get_active_lane_mask` intrinsic is utilized to generate a variably
sized mask of active lanes depending on the global position the tensor
intrinsic is operating on.
For now this relies on using `offset_of` and `stride` information from
the tensor we're predicating an access on. Likely we will want to
build on this in the future with a more intuitive API for determining
the current tile location.
Support for batched conv2d was removed since this causes numerical
issues which is suspected to be due to how the current tile is
determined (paragraph above).
[SME] Extract gemm block correctly when fused with bias/activation
Prior to this commit, the scheduling assumed the gemm block would
be the second to last block in the function ("unpadding" step is the
final block). However, when dense is fused with a bias or activation
the gemm block is no longer the second to last block. This commit
instead searches a single reduction block to use as the gemm block.
This commit adds a scalable `arm_cpu` conv2d NHWC schedule for fp32 which generates SME instructions by using the tensor intrinsics introduced in #16921.
Alongside the SME schedule, the logic of the TE schedule `schedule_conv2d_gemm_native()` for both non-scalable and scalable vector implementations has also been translated into the new TIR schedule. This means that the TE compute definition `compute_conv2d_NHWC_hybrid()` is now compatible with both the original TE schedules (e.g. `schedule_conv2d_NHWC_hybrid()`) and the newly introduced TIR schedule `schedule_conv2d_NHWC_hybrid_TIR()`. The corresponding TOPI test has been extended to reflect that.
This commit extends the functionality of the SME dense and matmul
schedules to support operations with fp16 inputs and an fp32 output,
where `transpose_a=False` and `transpose_b=True`.
For convenience, it also adds a utility called `get_vscale_factor`
which created the correct multiplier for `vscale` given a data type,
reflecting ideas from an early design of the
[SVE](https://github.com/apache/tvm-rfcs/pull/104) RFC.
* [SVE] Add support for representing and creating buffer-level predicates
Representation
--------------
This commit extends `BufferLoad` and `BufferStore` to accept a predicate
mask argument indicating which lanes in a vectorized buffer load/store
should be read/written.
As a simple example, we can load all lanes:
```
tir.BufferLoad(buf, [tir.Ramp(0, 1, 8)], predicate=tir.Broadcast(1, 8))
```
Or disable loading all lanes:
```
tir.BufferLoad(buf, [tir.Ramp(0, 1, 8)], predicate=tir.Broadcast(0, 8))
```
In TVMScript, buffer loads and stores are currently displayed using a
"short-hand" notation e.g. `A[0:4]`, but there was no clear path for
extending this notation to support predicates. Therefore, a "long-hand"
notation is introduced e.g. `A.load([T.Ramp(0, 1, 4)], predicate=...)`.
The TVMScript printer falls back to the long-hand notation whenever
predicates are specified.
Creation
--------
Buffer-level predication becomes more motivating when combined with the
`tir.get_active_lane_mask` intrinsic. It can be used to mask off lanes
when the vectorized axis is not divisible by the vector length. A
detailed example and rationale can be found in the
[RFC](https://github.com/apache/tvm-rfcs/blob/main/rfcs/0104-scalable-vectors-in-tir.md#predication).
Predicated buffer load/stores are created in the `VectorizeLoop` pass
via `TryPredicateBufferAccesses`. This pass aims to convert block-level
predicates e.g.
```
for i_0 in T.serial(4):
for i_1 in T.vectorized(4):
if i_0 * 4 + i_1 < 14:
B[i_0 * 4 + i_1] = A[i_0 * 4 + i_1] + 1.0
```
to buffer-level predicates, e.g.
```
for i_0 in T.serial(4):
predicate = T.get_active_lane_mask("int1x4", i_0 * 4, 14)
A_load = T.meta_var(A.load([T.Ramp(i_0 * 4, 1, 4)], predicate=predicate))
B.store(A_load, [T.Ramp(i_0 * 4, 1, 4)], predicate=predicate)
```
It takes a conservative approach for now, focussing only on expressions
produced by the split scheduling primitive, but more complex expressions
could be supported in the future.
`TryPredicateBufferAccesses` can be explicitly enabled/disabled with the
`tir.enable_buffer_level_predication` pass context option. By default it
will be disabled, unless the target supports SVE, in which case it will
be enabled by default.
Co-authored-by: Elen Kalda <elen.kalda@arm.com>
Co-authored-by: Neil Hickey <neil.hickey@arm.com>
Change-Id: Idde259a7d7e4536f00ed3a1dafedd0a5d24a1593
* Fix lint and correct test config option name
Change-Id: I864475c3d03e9b426ce5ef987989216d57f3e019
* Address review comments
This includes:
* Taking into account possibility of target being overridden in
the vectorize pass.
* Predicate PrimExpr -> Optional<PrimExpr>
* Checking that predicate is not used for any target that doesn't
support it.
* Use vload/vstore API as opposed to load/store
* int1 mask -> uint1 mask for boolean representation. This is converted
to int1 in the LLVM backend.
Change-Id: I4da0705352e321f6be6333a5bb777caa6a6ca9ef
* Fix lint
Change-Id: Idd3f3593fe524f3444487c520d947dfd53386db0
* Fix some failing tests
* vload/vstore updates that were missed previously
* int1 -> bool updates
* fix gpu target tests
Fixes a test and updates comments referencing old load/store api
Change-Id: I26a0c480d2dedee442ca0116909a7751d1dfa9ac
* Address comments
- Correct doc strings
- Correct typo in error message
- Add some additional checks for BufferLoad
Change-Id: Ie25563d569c0ed729ac915a6ba3a724a9e191014
* Account for buffer lanes in predicate lane check
Change-Id: I821210665e36c26bfa37fc9ed380b5d03c9e816e
This commit adds a new scalable fp32 dense schedule that calls SME intrinsics according to the SME RFC: https://github.com/apache/tvm-rfcs/pull/107.
Currently the schedule does not make use of predication, meaning the output from the matmul compute must be copied in a subsequent compute stage. This will be removed once support for predication is added.
This commit adds support for the `vscale_range()` LLVM function attribute to be generated for SVE and SME targets.
Some LLVM optimisation passes make use of the `vscale_range()` function attribute when scalable vectors are present (e.g. BasicAA llvm/llvm-project/pull/80445), so we include it alongside the "target_cpu" and "target-features" attributes.
This commit adds an `arm_cpu` conv2d NHWC schedule which generates SVE instructions by extending the hybrid GeMM approach implemented in #16106 to use scalable expressions as splitting factors.
Various vscale-related fixes needed to implement the schedule are also included, such as:
- adding vscale bounds in the `ConstIntBoundAnalyzer` and `IntervalSetEvaluator`
- simplifying `MinNode` and `MaxNode` that have scalable expression operands in `RewriteSimplifier`, which would appear when defining the shape of a buffer padded to be a multiple of vscale and in its respective buffer access indices (e.g. `C_1 = T.Buffer((1024 * (T.vscale() * 16 + 256 - 16 % T.vscale() * 16),), data=C)` instead of `C_1 = T.Buffer((1024 * (T.max(255, T.vscale() * 16 + 255 - 16 % T.vscale() * 16) + 1),), data=C)`)
The correctness of the new schedule is checked using a TOPI test, while the presence of generated SVE instructions is verified by a codegen_aarch64 test. The new rewrite_simplify rules are also covered by additional test cases.
Execution of SME instructions requires the processor be in a certain
state. This functionality can be can be controlled using LLVM function
level annotations such as "aarch64_pstate_sm_enabled" or
"aarch64_pstate_za_new" (see arm_utils.py for more information).
This commit exposes this functionality for AArch64 schedules where SME
intrinsics will be called. The attributes are intended to be added
at the block-level around the compute definition. They are prepended
with "pragma" to ensure they remain in the lowering.
In order to detect these attributes and convert them to the relevant
LLVM function attributes, a new AArch64 LLVM codegen backend is added.
This backend extends the functionality of `codegen_llvm` for AArch64
specific compilation.
Tests to check these attributes propagate correctly have been added.
This commit adds support for generating code for scalable loads and
stores. It also adds support for the creation of scalable broadcast
operations.
Co-authored-by: Elen Kalda <elen.kalda@arm.com>
Co-authored-by: Neil Hickey <neil.hickey@arm.com>
* [SVE] Add vscale builtin
Add a vscale builtin and lowering to `llvm.vscale`. This will be used in
subsequent patches for expressing scalable vectors in TIR.
Co-authored-by: Luke Hutton <luke.hutton@arm.com>
Co-authored-by: Neil Hickey <neil.hickey@arm.com>
* Improve documentation and fix LLVM versioning in a test
Change-Id: I5b364a8b50c8622d21d3d6d30b9d44a56e0418db
---------
Co-authored-by: Luke Hutton <luke.hutton@arm.com>
Co-authored-by: Neil Hickey <neil.hickey@arm.com>
The current unittest folder is too large and contains too many files and
too many components. This PR refactors the unittest folder by moving the
files to the corresponding folders.
Shushi Hong