akaashrp commented on code in PR #20236:
URL: https://github.com/apache/tvm/pull/20236#discussion_r3909455950
##########
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:
Thanks—this is correct. Updated in c13fab94fe: `kPosInf` is rejected
independently, and the intrinsic dtype maximum is compared only when it is
representable as an `IntImm`. Bounded and unbounded uint64 regressions cover
both paths. After rebasing onto current main, all 19 applicable Metal codegen
tests and pre-commit checks pass.
--
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]
