Skip to content

[Fix][Metal][WebGPU] Preserve variable bindings when resolving allocation bounds #20337

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akaashrp wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[Fix][Metal][WebGPU] Preserve variable bindings when resolving allocation bounds #20337

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
730fcb8
[FIX][Metal] Retain bound expressions for symbolic allocations
akaashrp Sep 13, 2026
0d7f7fc
[FIX][WebGPU] Retain bound expressions for symbolic allocations
akaashrp Sep 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions src/backend/metal/codegen/codegen_metal.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ Var GetSimdgroupBufferVar(const Expr& data) {

void CodeGenMetal::InitFuncState(const PrimFunc& f) {
CodeGenC::InitFuncState(f);
analyzer_ = arith::Analyzer();
// analyze the data;
for (Var arg : f->params) {
if (arg->ty.as<PointerTypeNode>()) {
Expand Down Expand Up @@ -330,6 +331,9 @@ void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os)
}

void CodeGenMetal::VisitStmt_(const BindNode* op) {
if (auto prim_value = op->value.as<PrimExpr>()) {
analyzer_->Bind(op->var, prim_value.value());
}
Comment on lines +334 to +336

@tlopex tlopex Sep 16, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please only bind pure expressions here, as in stmt_simplify.cc. Buffer loads are also PrimExpr: after snapshot = state[0] and an intervening store, the analyzer can incorrectly simplify state[0] - snapshot to zero. I reproduced both Metal and WebGPU allocating only one element while retaining a write to index 31. Guarding Bind with SideEffect(value) <= CallEffectKind::kPure would prevent this.

const auto* pointer_type = op->var->ty.as<PointerTypeNode>();
if (pointer_type == nullptr || pointer_type->storage_scope.empty()) {
return CodeGenC::VisitStmt_(op);
Expand Down Expand Up @@ -361,10 +365,9 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
this->PrintIndent();
// Compute a compile-time upper bound on the number of buffer elements.
size_t constant_size = 1;
arith::Analyzer analyzer;
for (const auto& dim : op->buffer->shape) {
const auto* dim_imm = dim.as<IntImmNode>();
int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value;
int64_t dim_size = dim_imm ? dim_imm->value : analyzer_->const_int_bound(dim)->max_value;
if (dim_imm == nullptr) {
// An integer dtype's intrinsic maximum is not a program-derived allocation bound.
TVM_FFI_ICHECK(dim_size != arith::ConstIntBound::kPosInf)
Expand Down
2 changes: 2 additions & 0 deletions src/backend/metal/codegen/codegen_metal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#ifndef TVM_TARGET_METAL_CODEGEN_METAL_H_
#define TVM_TARGET_METAL_CODEGEN_METAL_H_

#include <tvm/arith/analyzer.h>
#include <tvm/target/codegen.h>

#include <string>
Expand Down Expand Up @@ -63,6 +64,7 @@ class CodeGenMetal final : public CodeGenC {
using CodeGenC::PrintType;

private:
arith::Analyzer analyzer_;
std::unordered_map<const VarNode*, std::string> simdgroup_dtype_;
int thread_index_bits_{32};
int thread_work_dim_{0};
Expand Down
7 changes: 5 additions & 2 deletions src/backend/webgpu/codegen/codegen_webgpu.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ std::string CodeGenWebGPU::Finish() {

void CodeGenWebGPU::InitFuncState(const PrimFunc& f) {
CodeGenC::InitFuncState(f);
analyzer_ = arith::Analyzer();
workgroup_memory_bytes_ = 0;
// analyze the data;
for (Var arg : f->params) {
Expand Down Expand Up @@ -643,6 +644,9 @@ void CodeGenWebGPU::VisitExpr_(const TensorLoadNode* op, std::ostream& os) { //
}

void CodeGenWebGPU::VisitStmt_(const BindNode* op) {
if (auto prim_value = op->value.as<PrimExpr>()) {
analyzer_->Bind(op->var, prim_value.value());
}
// use ssa form.
if (print_ssa_form_) {
std::string value = PrintExpr(op->value);
Expand Down Expand Up @@ -721,10 +725,9 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get());
size_t constant_size = 1;
arith::Analyzer analyzer;
for (const auto& dim : op->buffer->shape) {
const auto* dim_imm = dim.as<IntImmNode>();
int64_t dim_size = dim_imm ? dim_imm->value : analyzer->const_int_bound(dim)->max_value;
int64_t dim_size = dim_imm ? dim_imm->value : analyzer_->const_int_bound(dim)->max_value;
if (dim_imm == nullptr) {
const auto* dtype_max = max_value(dim.ty()).as<IntImmNode>();
// An integer dtype's intrinsic maximum is not a program-derived allocation bound.
Expand Down
3 changes: 3 additions & 0 deletions src/backend/webgpu/codegen/codegen_webgpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
#ifndef TVM_TARGET_WEBGPU_CODEGEN_WEBGPU_H_
#define TVM_TARGET_WEBGPU_CODEGEN_WEBGPU_H_

#include <tvm/arith/analyzer.h>
#include <tvm/target/codegen.h>

#include <cstddef>
Expand Down Expand Up @@ -86,6 +87,8 @@ class CodeGenWebGPU final : public CodeGenC {
void VisitStmt_(const ContinueNode* op) final;

private:
arith::Analyzer analyzer_;

/*!
* \brief Enforce value to be U32.
*/
Expand Down
36 changes: 35 additions & 1 deletion tests/python/codegen/test_target_codegen_metal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def main(A: T.Buffer((1, 2), "int32")):
for i in T.thread_binding(1, thread="threadIdx.x"):
with T.sblock("block"):
tx = T.axis.spatial(1, i)
r = T.ramp(tx, 3, 2)
r: T.let = T.ramp(tx, 3, 2)
A[0, T.ramp(0, 1, 2)] = r

f = tvm.compile(IRModule, target=target)
Expand Down Expand Up @@ -410,6 +410,40 @@ def main(n: T.int32):
assert "thread float scratch[128]" in source


@pytest.mark.parametrize("bounded", [True, False])
def test_bound_symbolic_stack_allocation(bounded):
limit = 64 if bounded else 2147483647

@I.ir_module
class Module:
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr(
{
"calling_conv": 2,
"global_symbol": "main",
"target": T.target("metal"),
"tirx.kernel_launch_params": [],
"tirx.is_global_func": True,
}
)
# Common subexpression elimination can hoist the bounded extent.
extent: T.let[T.int32] = T.min(n, limit)
elements: T.let[T.int32] = extent * 2
scratch = T.alloc_buffer((elements,), "float32", scope="local")
T.evaluate(scratch.data)

if bounded:
source = _build_metal(Module).inspect_source()
assert "thread float scratch[128]" in source
else:
with pytest.raises(
tvm.error.InternalError,
match="Metal allocation extent requires a finite compile-time upper bound",
):
_build_metal(Module)


def test_bounded_uint64_symbolic_stack_allocation():
@I.ir_module
class Module:
Expand Down
68 changes: 68 additions & 0 deletions tests/python/codegen/test_target_codegen_webgpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,74 @@ def main(n: T.int32):
assert re.search(r"\bvar\s+\w+\s*:\s*array<f32,\s*128>;", source)


@pytest.mark.parametrize("scope", ["local", "shared"])
@pytest.mark.parametrize("bounded", [True, False])
def test_bound_symbolic_allocation(scope, bounded):
limit = 64 if bounded else 2147483647

@I.ir_module
class Module:
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr(
{
"calling_conv": 2,
"global_symbol": "main",
"target": T.target("webgpu"),
"tirx.is_global_func": True,
}
)
# Common subexpression elimination can hoist the bounded extent.
extent: T.let[T.int32] = T.min(n, limit)
first = T.alloc_buffer((extent * 2,), "float32", scope=scope)
elements: T.let[T.int32] = extent * 2
second = T.alloc_buffer((elements,), "float32", scope=scope)
first[0] = 1.0
second[0] = first[0]

if bounded:
source = _build_webgpu(Module).inspect_source()
declaration = r"var<workgroup>" if scope == "shared" else r"\bvar"
assert len(re.findall(declaration + r"\s+\w+\s*:\s*array<f32,\s*128>;", source)) == 2
else:
with pytest.raises(
tvm.error.InternalError,
match="WebGPU allocation extent requires a finite compile-time upper bound",
):
_build_webgpu(Module)


@pytest.mark.parametrize("target_limit", [512, 496])
def test_bound_symbolic_workgroup_allocation_respects_target_limit(target_limit):
@I.ir_module
class Module:
@T.prim_func(s_tir=True)
def main(n: T.int32):
T.func_attr(
{
"calling_conv": 2,
"global_symbol": "main",
"target": T.target("webgpu"),
"tirx.is_global_func": True,
}
)
extent: T.let[T.int32] = T.min(n, 64)
elements: T.let[T.int32] = extent * 2
scratch = T.alloc_buffer((elements,), "float32", scope="shared")
scratch[0] = 1.0

target = {"kind": "webgpu", "max_shared_memory_per_block": target_limit}
if target_limit == 512:
source = _build_webgpu(Module, target).inspect_source()
assert re.search(r"var<workgroup>\s+\w+\s*:\s*array<f32,\s*128>;", source)
else:
with pytest.raises(
tvm.error.InternalError,
match=r"WebGPU workgroup allocations use 512 bytes, .* supports only 496 bytes",
):
_build_webgpu(Module, target)


def test_unbounded_symbolic_stack_allocation_rejected():
@I.ir_module
class Module:
Expand Down
Loading