This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 23fc37c4cf [Fix][TIRx] Handle vector access pointer addresses in C
codegen (#20058)
23fc37c4cf is described below
commit 23fc37c4cf95569d417e8e657dc67d2d0f171f89
Author: Shushi Hong <[email protected]>
AuthorDate: Tue Jul 28 19:29:13 2026 -0400
[Fix][TIRx] Handle vector access pointer addresses in C codegen (#20058)
This PR fixes invalid pointer arithmetic emitted by C-family codegen for
vector-typed `tvm_access_ptr`.
A vector access pointer is lowered to `address_of(BufferLoad(...))` with
a `Ramp` index describing its lane indices. For example, `Ramp(4, 1, 2)`
represents scalar elements `[4, 5]`, so its address should be the
address of the first lane, `&A[4]`.
LLVM codegen already extracts `Ramp::base` in this case. However,
`CodeGenC` previously passed the complete ramp to pointer arithmetic,
which could generate invalid CUDA code such as:
```cpp
(float*)A + make_int2(4, 5)
```
This PR makes `CodeGenC` use `Ramp::base` when generating the address of
a vector `BufferLoad`. The normalized index is applied to both the
direct pointer-offset path and the general `GetBufferRef` path.
The existing scalar-buffer plus `Ramp` lowering is preserved. This
avoids regressions for padded vector types such as `float32x3` and
packed vector types such as `int4x4`, while making C-family codegen
consistent with LLVM codegen.
Regression tests cover:
- `float32x2` C codegen.
- Padded `float32x3` LLVM codegen.
- Packed `int4x4` CUDA codegen.
---
src/target/source/codegen_c.cc | 11 ++++--
tests/python/codegen/test_target_codegen_c_host.py | 13 +++++++
.../test_tir_transform_lower_intrin.py | 40 ++++++++++++++++++++++
tests/python/tirx/codegen/test_codegen_cuda.py | 20 +++++++++++
4 files changed, 81 insertions(+), 3 deletions(-)
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 6f9c25c592..8d0dadb55e 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -758,15 +758,20 @@ void CodeGenC::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLINT(*)
if (load) {
TVM_FFI_ICHECK_EQ(load->indices.size(), 1)
<< "CodeGenC only supports flat memory allocations.";
+ PrimExpr index = load->indices[0];
+ // A vector BufferLoad uses a Ramp to describe its lane indices. The
+ // address of that load is the address of its first lane.
+ if (const RampNode* ramp = index.as<RampNode>()) {
+ index = ramp->base;
+ }
const VarNode* data = load->buffer->data.get();
if (pointer_offset_vars_.count(data) && HandleTypeMatch(data,
load->buffer->dtype) &&
!IsVolatile(data)) {
os << "(" << GetVarID(data) << " + ";
- this->PrintExpr(load->indices[0], os);
+ this->PrintExpr(index, os);
os << ")";
} else {
- os << "(&("
- << GetBufferRef(load->ty.as_or_throw<PrimType>(),
load->buffer.get(), load->indices[0])
+ os << "(&(" << GetBufferRef(load->ty.as_or_throw<PrimType>(),
load->buffer.get(), index)
<< "))";
}
} else {
diff --git a/tests/python/codegen/test_target_codegen_c_host.py
b/tests/python/codegen/test_target_codegen_c_host.py
index 8af6be0b79..f33eb92210 100644
--- a/tests/python/codegen/test_target_codegen_c_host.py
+++ b/tests/python/codegen/test_target_codegen_c_host.py
@@ -245,5 +245,18 @@ def test_workspace_allocation_cast():
built.export_library(temp.relpath("workspace.so"))
+def test_vector_access_ptr_address_uses_ramp_base():
+ buffer = tvm.tirx.decl_buffer((8,), "float32x2", name="A")
+ access_ptr = buffer.access_ptr(access_mask=3, offset=2, extent=4)
+ body = tvm.tirx.Evaluate(tvm.tirx.call_extern("void", "consume",
access_ptr))
+ func = tvm.tirx.PrimFunc([buffer], body).with_attr("global_symbol", "main")
+
+ source = tvm.tirx.build(tvm.IRModule.from_expr(func),
target="c").inspect_source()
+ call = next(line.strip() for line in source.splitlines() if
line.strip().startswith("consume("))
+ assert "int32_t2" not in call
+ assert "float2*" in call
+ assert " + 4" in call
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
index bde693fbce..c5378cfcdb 100644
--- a/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
+++ b/tests/python/tirx-transform/test_tir_transform_lower_intrin.py
@@ -114,6 +114,46 @@ def test_lower_nested_access_ptr():
assert int(tvm.arith.Analyzer().simplify(load.indices[0])) == 5
+def test_lower_vector_access_ptr():
+ buffer = tvm.tirx.decl_buffer((8,), "float32x2", name="A")
+ access_ptr = buffer.access_ptr(access_mask=3, offset=2, extent=4)
+
+ assert access_ptr.op.name == "tirx.tvm_access_ptr"
+ assert int(access_ptr.args[2]) == 2
+ assert int(access_ptr.args[3]) == 4
+ assert int(access_ptr.args[4]) == 3
+
+ mod = tvm.IRModule.from_expr(
+ tvm.tirx.PrimFunc([buffer], tvm.tirx.Evaluate(access_ptr)).with_attr(
+ "target", tvm.target.Target("llvm")
+ )
+ )
+ lowered = tvm.tirx.transform.LowerIntrin()(mod)["main"].body.value
+ assert lowered.op.name == "tirx.address_of"
+ assert lowered.ty == access_ptr.ty
+
+ load = lowered.args[0]
+ assert isinstance(load, tvm.tirx.BufferLoad)
+ assert load.buffer.data.same_as(buffer.data)
+ assert load.buffer.dtype == tvm.ir.PrimType("float32")
+ assert len(load.indices) == 1
+ ramp = load.indices[0]
+ assert isinstance(ramp, tvm.tirx.Ramp)
+ assert int(ramp.base) == 4
+ assert int(ramp.stride) == 1
+ assert ramp.lanes == 2
+
+
[email protected](not env.has_llvm(), reason="need llvm")
+def test_lower_vector_access_ptr_with_padded_vector_dtype():
+ buffer = tvm.tirx.decl_buffer((8,), "float32x3", name="A")
+ access_ptr = buffer.access_ptr(access_mask=1, offset=2, extent=4)
+ body = tvm.tirx.Evaluate(tvm.tirx.call_extern("void", "consume",
access_ptr))
+ func = tvm.tirx.PrimFunc([buffer], body).with_attr("global_symbol", "main")
+
+ tvm.tirx.build(tvm.IRModule.from_expr(func), target="llvm")
+
+
def get_ref_data():
"""Get reference data for every pairs"""
import itertools
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index f5fc621c69..745d10d928 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -43,6 +43,26 @@ def _helper_source(src: str, helper_name: str) -> str:
return src[start:next_helper]
+def test_vector_access_ptr_preserves_packed_offset(monkeypatch):
+ buffer = tvm.tirx.decl_buffer((8,), "int4x4", name="A")
+ access_ptr = buffer.access_ptr(access_mask=3, offset=2, extent=4)
+ body = tvm.tirx.Evaluate(tvm.tirx.call_extern("void", "consume",
access_ptr))
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+ func = (
+ tvm.tirx.PrimFunc([buffer.data], body)
+ .with_attr("global_symbol", "main")
+ .with_attr("target", target)
+ )
+ lowered = tvm.tirx.transform.LowerIntrin()(tvm.IRModule.from_expr(func))
+
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+ source = tvm.get_global_func("target.build.cuda")(lowered,
target).inspect_source()
+ call = next(line.strip() for line in source.splitlines() if
line.strip().startswith("consume("))
+
+ assert "make_int4" not in call
+ assert " + 8 / 4" in call
+
+
def test_tirx_launch_bounds_omits_min_blocks_without_persistent_schedule():
@T.prim_func
def main(A: T.Buffer((4,), "int32")):
