tlopex commented on code in PR #20236:
URL: https://github.com/apache/tvm/pull/20236#discussion_r3901394480
##########
src/backend/metal/codegen/codegen_metal.cc:
##########
@@ -357,14 +359,25 @@ void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
std::string vid = AllocVarID(op->buffer.get());
this->PrintIndent();
- // Compute constant_size from buffer shape
+ // 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 IntImmNode* dim_imm = dim.as<IntImmNode>();
- TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation
for now";
- constant_size *= dim_imm->value;
+ const auto* dim_imm = dim.as<IntImmNode>();
+ 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.
+ TVM_FFI_ICHECK(dtype_max && dim_size < dtype_max->value)
+ << "Metal allocation extent requires a finite compile-time upper
bound, but got " << dim;
Review Comment:
This incorrectly rejects bounded symbolic extents with dtype `uint64`.
`max_value(PrimType::UInt(64))` is represented by `LargeUIntImm`, because
`UINT64_MAX` cannot fit in `IntImmNode::value`. Therefore, `as<IntImmNode>()`
returns null and this check fails even when the analyzer has proven a small
finite upper bound.
For example:
def main(n: T.uint64):
scratch = T.alloc_buffer(
(T.min(n, T.uint64(64)),),
"float32",
scope="local",
)
`const_int_bound` reports `[0, 64]`, but Metal codegen rejects it with the
"requires a finite compile-time upper bound" error. The equivalent `int32`,
`int64`, and `uint32` cases all emit `thread float scratch[64]`.
Could we check `dim_size != arith::ConstIntBound::kPosInf` independently,
and only compare against the dtype maximum when it is representable as an
`IntImmNode`? That would continue rejecting an unbounded `uint64` extent while
accepting genuinely bounded expressions. Please also add a `uint64` regression
test.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
