This is an automated email from the ASF dual-hosted git repository.
tqchen 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 9b02ce500f [FIX][Codegen] Bind tensor views with their element pointer
type (#20657)
9b02ce500f is described below
commit 9b02ce500f0801f7427c94392e9a9b4d58382f88
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 11 11:07:00 2026 -0400
[FIX][Codegen] Bind tensor views with their element pointer type (#20657)
Tensor aliases backed by untyped or differently typed pointers currently
repeat casts at every access. Declare a typed view once and reuse it for
matching accesses, while preserving backing allocations, address spaces,
volatile accesses and required vector/packed casts. Reuse packed
argument decoding casts and keep WebGPU aliases on their existing WGSL
array representation.
---
src/backend/webgpu/codegen/codegen_webgpu.cc | 15 ++++++++++++
src/backend/webgpu/codegen/codegen_webgpu.h | 1 +
src/target/source/codegen_c.cc | 29 +++++++++++++++-------
tests/python/tirx/codegen/test_codegen_cuda.py | 33 ++++++++++++++++++++++++++
4 files changed, 69 insertions(+), 9 deletions(-)
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index ee2c92d2f3..20f26360aa 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -769,6 +769,21 @@ void CodeGenWebGPU::Dispatch_(const TensorStoreNode* op) {
}
}
+void CodeGenWebGPU::DispatchDeclTensor(const BindNode* op, const CallNode*
buffer_call) {
+ Expr data = buffer_call->args[0];
+ const VarNode* source = data.as<VarNode>();
+ if (const auto* call = data.as<CallNode>();
+ call && call->op.same_as(tirx::tensor_data_ptr_op()) &&
call->args.size() == 1) {
+ source = call->args[0].as<VarNode>();
+ }
+ TVM_FFI_CHECK(source && var_idmap_.count(source), ValueError)
+ << "WebGPU tensor views require a named storage array";
+ // WGSL storage is an array, not a C pointer. Preserve its physical element
+ // representation and access mode by sharing the existing array name.
+ TVM_FFI_ICHECK(!var_idmap_.count(op->var.get()));
+ var_idmap_[op->var.get()] = GetVarID(source);
+}
+
void CodeGenWebGPU::DispatchAllocTensor(const BindNode* op, const CallNode*
buffer_call) {
tvm::Tuple shape = buffer_call->args[0].as_or_throw<tvm::Tuple>();
DLDataType dtype = buffer_call->args[1].as_or_throw<DataTypeImm>()->value;
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h
b/src/backend/webgpu/codegen/codegen_webgpu.h
index 92f967ec17..90ccc28014 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.h
+++ b/src/backend/webgpu/codegen/codegen_webgpu.h
@@ -84,6 +84,7 @@ class CodeGenWebGPU final : public CodeGenC {
void Dispatch_(const TensorStoreNode* op) final;
void Dispatch_(const ForNode* op) final;
void DispatchAllocTensor(const BindNode* op, const CallNode* buffer_call);
+ void DispatchDeclTensor(const BindNode* op, const CallNode* buffer_call);
void Dispatch_(const AssertStmtNode* op) final;
void Dispatch_(const WhileNode* op) final;
void Dispatch_(const BreakNode* op) final;
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 6eef5cf4c9..ae2271e41d 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -23,6 +23,7 @@
#include "codegen_c.h"
#include <tvm/ffi/cast.h>
+#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/function.h>
#include <tvm/ir/prim/op.h>
#include <tvm/ir/unique_name_supply.h>
@@ -946,18 +947,20 @@ void CodeGenC::DispatchDeclTensor(const BindNode* op,
const CallNode* buffer_cal
}
if (source && var_idmap_.count(source)) {
TVM_FFI_ICHECK(!var_idmap_.count(buffer.get()));
- var_idmap_[buffer.get()] = GetVarID(source);
if (auto it = alloc_storage_scope_.find(source); it !=
alloc_storage_scope_.end()) {
- alloc_storage_scope_[buffer.get()] = it->second;
- } else {
- alloc_storage_scope_[buffer.get()] = scope;
+ scope = it->second;
}
if (IsVolatile(source)) {
MarkVolatile(buffer.get());
}
- auto it = handle_data_type_.find(source);
- RegisterHandleType(buffer.get(), it == handle_data_type_.end() ?
PrimType(dtype) : it->second);
- return;
+ // Reuse the name only when it already has the view's element type.
+ // Otherwise declare a typed view without changing the backing allocation.
+ if (HandleTypeMatch(source, PrimType(dtype))) {
+ var_idmap_[buffer.get()] = GetVarID(source);
+ alloc_storage_scope_[buffer.get()] = scope;
+ RegisterHandleType(buffer.get(), PrimType(dtype));
+ return;
+ }
}
alloc_storage_scope_[buffer.get()] = scope;
@@ -965,9 +968,17 @@ void CodeGenC::DispatchDeclTensor(const BindNode* op,
const CallNode* buffer_cal
if (IsScopePartOfType()) {
PrintStorageScope(scope, stream);
}
- PrintType(PtrType(PrimType(dtype), scope), stream);
+ PtrType pointer_type(PrimType(dtype), scope);
+ PrintType(pointer_type, stream);
stream << ' ' << AllocVarID(buffer.get()) << " = ";
- PrintExpr(Call(PtrType(PrimType(dtype), scope), tirx::reinterpret_op(),
{data}), stream);
+ // Packed argument decoding may already have established this typed view.
+ const auto* cast = data.as<CallNode>();
+ if (cast && cast->op.same_as(tirx::reinterpret_op()) &&
+ ffi::StructuralEqual()(cast->ty, pointer_type)) {
+ PrintExpr(data, stream);
+ } else {
+ PrintExpr(Call(pointer_type, tirx::reinterpret_op(), {data}), stream);
+ }
stream << ";\n";
RegisterHandleType(buffer.get(), PrimType(dtype));
}
diff --git a/tests/python/tirx/codegen/test_codegen_cuda.py
b/tests/python/tirx/codegen/test_codegen_cuda.py
index a1fc56db32..31e72be560 100644
--- a/tests/python/tirx/codegen/test_codegen_cuda.py
+++ b/tests/python/tirx/codegen/test_codegen_cuda.py
@@ -152,6 +152,39 @@ def test_cuda_module_destructor_preserves_current_device():
torch.cuda.set_device(original_device)
+def test_typed_buffer_bindings(monkeypatch):
+ @T.function
+ def main(data: T.handle("float32"), output: T.handle):
+ T.func_attr({"global_symbol": "main", "calling_conv": 2,
"tirx.kernel_launch_params": []})
+ A = T.decl_tensor((1,), "float32", data=data)
+ B = T.decl_tensor((1,), "float32", data=output)
+ storage = T.alloc_tensor((1,), "uint32", scope="shared")
+ S = T.decl_tensor((1,), "float32", data=storage.data, scope="shared")
+ pool = T.alloc_tensor((4,), "uint8", scope="shared.dyn")
+ D = T.decl_tensor((1,), "float32", data=pool.data, scope="shared.dyn")
+ alias = T.decl_tensor((1,), "float32", data=D.data, scope="shared.dyn")
+ S[0] = A[0]
+ D[0] = S[0]
+ B[0] = alias[0]
+ storage[0] = T.uint32(0)
+
+ monkeypatch.setenv("TVM_COMPILE_FORCE_FALLBACK", "1")
+ target = tvm.target.Target({"kind": "cuda", "arch": "sm_80"})
+ source = tvm.get_global_func("target.build.cuda")(
+ tvm.IRModule.from_expr(main), target
+ ).inspect_source()
+ assert "float* data, void* output" in source
+ assert "float* B = ((float*)output);" in source
+ assert "__shared__ alignas(64) uint storage_ptr[1];" in source
+ assert "float* S = ((float*)storage_ptr);" in source
+ assert "extern __shared__ __align__(64) uchar pool_ptr[];" in source
+ assert "float* D = ((float*)pool_ptr);" in source
+ assert "S[0] = data[0];" in source
+ assert "D[0] = S[0];" in source
+ assert "B[0] = D[0];" in source
+ assert "storage_ptr[0] = (uint)0;" in source
+
+
def test_vector_pointer_preserves_packed_offset(monkeypatch):
tensor = tvm.tirx.decl_tensor((8,), "int4x4", name="A")
data = tvm.tirx.Var("A_data", tensor.data.ty)
