From 80819853121a083c25bb45331b2d99cb443ff8f1 Mon Sep 17 00:00:00 2001 From: Chengye YU Date: Thu, 19 Mar 2026 16:12:30 +0800 Subject: [PATCH 1/3] [Vulkan] Avoid explicit layout decoration on non-interface allocations --- src/target/spirv/ir_builder.cc | 19 ++++++++-------- .../codegen/test_target_codegen_vulkan.py | 22 +++++++++++++++++++ 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/src/target/spirv/ir_builder.cc b/src/target/spirv/ir_builder.cc index 135888c23d8b..f912e482761c 100644 --- a/src/target/spirv/ir_builder.cc +++ b/src/target/spirv/ir_builder.cc @@ -178,23 +178,24 @@ SType IRBuilder::GetStructArrayType(const SType& value_type, uint32_t num_elems, } else { ib_.Begin(spv::OpTypeRuntimeArray).AddSeq(arr_type, value_type).Commit(&global_); } - int nbits = value_type.type.bits() * value_type.type.lanes(); - TVM_FFI_ICHECK_EQ(nbits % 8, 0); - uint32_t nbytes = static_cast(nbits) / 8; - // decorate the array type. - this->Decorate(spv::OpDecorate, arr_type, spv::DecorationArrayStride, nbytes); + if (interface_block) { + int nbits = value_type.type.bits() * value_type.type.lanes(); + TVM_FFI_ICHECK_EQ(nbits % 8, 0); + uint32_t nbytes = static_cast(nbits) / 8; + // Explicit layout is required for descriptor-backed interface blocks. + this->Decorate(spv::OpDecorate, arr_type, spv::DecorationArrayStride, nbytes); + } // declare struct of array SType struct_type; struct_type.id = id_counter_++; struct_type.type = DataType::Handle(); struct_type.element_type_id = value_type.id; ib_.Begin(spv::OpTypeStruct).AddSeq(struct_type, arr_type).Commit(&global_); - // decorate the array type. - ib_.Begin(spv::OpMemberDecorate) - .AddSeq(struct_type, 0, spv::DecorationOffset, 0) - .Commit(&decorate_); if (interface_block) { + ib_.Begin(spv::OpMemberDecorate) + .AddSeq(struct_type, 0, spv::DecorationOffset, 0) + .Commit(&decorate_); // Runtime array are always decorated as Block or BufferBlock // (shader storage buffer) if (spirv_support_.supports_storage_buffer_storage_class) { diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py index 38830ae96f30..8f2aed7035de 100644 --- a/tests/python/codegen/test_target_codegen_vulkan.py +++ b/tests/python/codegen/test_target_codegen_vulkan.py @@ -515,6 +515,28 @@ def kernel(): vulkan_codegen(Module, target) +@tvm.testing.requires_vulkan(support_required="compile-only") +def test_codegen_static_shared_memory(): + """The codegen should accept static shared/workgroup allocations.""" + + A = te.placeholder((128,), name="A", dtype="float32") + B = te.compute((128,), lambda i: A[i], name="B") + + sch = tir.Schedule(te.create_prim_func([A, B])) + block = sch.get_block("B") + (loop,) = sch.get_loops(block) + bx, tx = sch.split(loop, factors=[None, 128]) + sch.bind(bx, "blockIdx.x") + sch.bind(tx, "threadIdx.x") + + block_read = sch.cache_read(block, 0, "shared") + sch.compute_at(block_read, bx) + _, fetch_tx = sch.get_loops(block_read) + sch.bind(fetch_tx, "threadIdx.x") + + tvm.compile(sch.mod, target="vulkan") + + @tvm.testing.requires_gpu @tvm.testing.requires_vulkan def test_unary(): From fe31b87f0fe754783b5d772683293461caac704f Mon Sep 17 00:00:00 2001 From: Chengye YU Date: Thu, 19 Mar 2026 19:05:56 +0800 Subject: [PATCH 2/3] [Vulkan] Fix missing imports in static shared-memory test --- tests/python/codegen/test_target_codegen_vulkan.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py index 8f2aed7035de..ba7736a915e4 100644 --- a/tests/python/codegen/test_target_codegen_vulkan.py +++ b/tests/python/codegen/test_target_codegen_vulkan.py @@ -23,6 +23,7 @@ import tvm import tvm.testing +from tvm import te, tir from tvm.script import ir as I from tvm.script import tir as T from tvm.script.ir_builder import IRBuilder From 9516ca4309653df2278a51d32f89b9adbef88667 Mon Sep 17 00:00:00 2001 From: Chengye YU Date: Fri, 20 Mar 2026 23:09:13 +0800 Subject: [PATCH 3/3] [Vulkan] Use tvm script in static shared-memory test --- .../codegen/test_target_codegen_vulkan.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py index ba7736a915e4..9af08c1a04bb 100644 --- a/tests/python/codegen/test_target_codegen_vulkan.py +++ b/tests/python/codegen/test_target_codegen_vulkan.py @@ -23,7 +23,6 @@ import tvm import tvm.testing -from tvm import te, tir from tvm.script import ir as I from tvm.script import tir as T from tvm.script.ir_builder import IRBuilder @@ -520,22 +519,18 @@ def kernel(): def test_codegen_static_shared_memory(): """The codegen should accept static shared/workgroup allocations.""" - A = te.placeholder((128,), name="A", dtype="float32") - B = te.compute((128,), lambda i: A[i], name="B") - - sch = tir.Schedule(te.create_prim_func([A, B])) - block = sch.get_block("B") - (loop,) = sch.get_loops(block) - bx, tx = sch.split(loop, factors=[None, 128]) - sch.bind(bx, "blockIdx.x") - sch.bind(tx, "threadIdx.x") + @I.ir_module + class Module: + @T.prim_func + def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")): + A_shared = T.alloc_buffer((128,), dtype="float32", scope="shared") - block_read = sch.cache_read(block, 0, "shared") - sch.compute_at(block_read, bx) - _, fetch_tx = sch.get_loops(block_read) - sch.bind(fetch_tx, "threadIdx.x") + for bx in T.thread_binding(1, thread="blockIdx.x"): + for tx in T.thread_binding(128, thread="threadIdx.x"): + A_shared[tx] = A[tx] + B[tx] = A_shared[tx] - tvm.compile(sch.mod, target="vulkan") + tvm.compile(Module, target="vulkan") @tvm.testing.requires_gpu