[Fix][TIRx] Handle vector access pointer addresses in C codegen

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 {

tests/python/codegen/test_target_codegen_c_host.py

@@ -245,5 +245,18 @@ def main(A: T.Buffer((256,), "float32")):
     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()

tests/python/tirx-transform/test_tir_transform_lower_intrin.py

@@ -114,6 +114,46 @@ def collect(node):
     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
+
+
+@pytest.mark.skipif(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

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")):