[Fix][TIR] Check allocation size arithmetic instead of wrapping around

src/s_tir/analysis/verify_gpu_code.cc

@@ -33,6 +33,9 @@
 #include <tvm/tirx/analysis.h>
 #include <tvm/tirx/stmt.h>
 
+#include <algorithm>
+#include <limits>
+
 #include "../../runtime/thread_storage_scope.h"
 #include "../../tirx/transform/ir_utils.h"
 
@@ -71,22 +74,22 @@ class GPUCodeVerifier : public StmtExprVisitor {
     TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(StmtExprVisitor::Visit_(op));
     auto scope = op->buffer.scope();
     runtime::StorageScope storage_scope = runtime::StorageScope::Create(scope);
-    int64_t const_size = 1;
-    for (const PrimExpr& e : op->buffer->shape) {
-      if (auto* imm = e.as<IntImmNode>()) {
-        const_size = static_cast<int64_t>(const_size * imm->value);
-      } else {
-        const_size = 0;
-        break;
-      }
-    }
     PrimType dtype_ty = op->buffer->dtype;
     TVM_FFI_ICHECK(!dtype_ty.IsScalableVector())
         << "Cannot verify GPU memory usage for scalable vector dtype " << dtype_ty;
-    if (storage_scope.rank == runtime::StorageRank::kLocal) {
-      local_memory_per_block_ += static_cast<size_t>(const_size) * ElementBytes(dtype_ty);
+    // A dynamic shape leaves const_nbytes at zero and is not accounted for, as before.
+    uint64_t const_nbytes = 0;
+    ConstantSizeKind size_kind =
+        GetConstantAllocationSize(op->buffer->shape, ElementBytes(dtype_ty), &const_nbytes);
+    if (size_kind == ConstantSizeKind::kUnrepresentable) {
+      std::stringstream s;
+      s << "Size of buffer " << op->buffer.name() << " with shape " << op->buffer->shape
+        << " and dtype " << dtype_ty << " does not fit in a 64-bit byte count";
+      errors_.push_back(s.str());
+    } else if (storage_scope.rank == runtime::StorageRank::kLocal) {
+      AccumulateBytes(&local_memory_per_block_, const_nbytes);
     } else if (storage_scope.rank == runtime::StorageRank::kShared) {
-      shared_memory_per_block_ += static_cast<size_t>(const_size) * ElementBytes(dtype_ty);
+      AccumulateBytes(&shared_memory_per_block_, const_nbytes);
     }
     if (dtype_ty.IsFixedLengthVector()) {
       if (ElementBytes(dtype_ty) > max_vector_bytes_) {
@@ -288,6 +291,18 @@ class GPUCodeVerifier : public StmtExprVisitor {
 
   static size_t ElementBytes(const PrimType& ty) { return ty.StorageBytes(); }
 
+  /*!
+   * \brief Add to a memory total, saturating instead of wrapping around.
+   *
+   * The totals are only ever compared against a maximum, so saturating keeps an
+   * oversized kernel detectable.
+   */
+  static void AccumulateBytes(size_t* total, uint64_t nbytes) {
+    constexpr size_t kMaxTotal = std::numeric_limits<size_t>::max();
+    size_t value = static_cast<size_t>(std::min<uint64_t>(nbytes, kMaxTotal));
+    *total = *total > kMaxTotal - value ? kMaxTotal : *total + value;
+  }
+
   void Reset_() {
     local_memory_per_block_ = 0;
     shared_memory_per_block_ = 0;

src/s_tir/transform/merge_shared_memory_allocations.cc

@@ -85,24 +85,6 @@ bool IsStaticSharedMemory(const BufferVar& buffer) {
   return storage_scope.rank == runtime::StorageRank::kShared && storage_scope.tag == "";
 }
 
-/*!
- * \brief Compute constant allocation size from buffer's allocation shape.
- * \return Product of extents if all constant, 0 otherwise.
- */
-static int64_t ConstantAllocationSize(const ffi::Array<PrimExpr>& extents) {
-  int64_t result = 1;
-  for (size_t i = 0; i < extents.size(); ++i) {
-    if (const IntImmNode* int_size = extents[i].as<IntImmNode>()) {
-      auto product = (result * int_size->value).as<int64_t>();
-      if (!product.has_value()) return 0;
-      result = *product;
-    } else {
-      return 0;
-    }
-  }
-  return result;
-}
-
 /*!
  * \brief collect the mapping from the buffer var to its BufferVar within a subtree
  */
@@ -861,8 +843,16 @@ class SharedMemoryRewriter : public StmtExprMutator {
     ffi::Array<PrimExpr> alloc_shape = GetBufferAllocationShape(buf);
     DLDataType dtype = buf->dtype->dtype;
     uint64_t op_elem_bits = static_cast<uint64_t>(dtype.bits) * dtype.lanes;
-    uint64_t const_nbits =
-        static_cast<uint64_t>(ConstantAllocationSize(alloc_shape) * op_elem_bits);
+    // A size that is not known at compile time leaves const_nbits at zero. A
+    // constant size that does not fit in uint64_t is rejected rather than
+    // planned with the wrapped value.
+    uint64_t const_nbits = 0;
+    if (tirx::GetConstantAllocationSize(alloc_shape, op_elem_bits, &const_nbits) ==
+        tirx::ConstantSizeKind::kUnrepresentable) {
+      TVM_FFI_THROW(ValueError) << "Cannot plan shared memory for buffer " << buf.name()
+                                << " with shape " << alloc_shape << " and dtype " << buf->dtype
+                                << ": its size in bits does not fit in 64 bits";
+    }
     // disable reuse of small arrays, they will be lowered to registers in LLVM
     // This rules only apply if we are using non special memory
     if (const_nbits > 0 && const_nbits <= 32) {
@@ -873,7 +863,7 @@ class SharedMemoryRewriter : public StmtExprMutator {
       // constant allocation.
       auto begin = scope.const_free_map.lower_bound(0);
       auto mid = scope.const_free_map.lower_bound(const_nbits);
-      auto end = scope.const_free_map.upper_bound(const_nbits * match_range);
+      auto end = scope.const_free_map.upper_bound(tirx::SaturatingMul(const_nbits, match_range));
       // Start looking at the buffer that is bigger than the required size first.
       // If we find one, directly allocate the buffer in its location and remove its entry in the
       // free list

src/tirx/transform/ir_utils.h

@@ -24,6 +24,7 @@
 #ifndef TVM_TIR_TRANSFORM_IR_UTILS_H_
 #define TVM_TIR_TRANSFORM_IR_UTILS_H_
 
+#include <tvm/ffi/big_int.h>
 #include <tvm/ir/prim/builtin.h>
 #include <tvm/ir/prim/expr.h>
 #include <tvm/ir/scope_stack.h>
@@ -198,6 +199,55 @@ inline int GetTempAllocaAlignment(const PrimType& type, int64_t const_size) {
   return align;
 }
 
+/*! \brief Outcome of computing the constant size of an allocation. */
+enum class ConstantSizeKind : int {
+  /*! \brief At least one extent is not known at compile time. */
+  kDynamic = 0,
+  /*! \brief The size is known at compile time and fits in uint64_t. */
+  kConstant = 1,
+  /*! \brief The extents are constant, but the size is not representable. */
+  kUnrepresentable = 2,
+};
+
+/*!
+ * \brief Compute the constant size of an allocation, in units of \p unit_size.
+ *
+ * The product of the extents and \p unit_size is formed exactly, so a shape
+ * whose size does not fit in uint64_t is reported instead of wrapping around.
+ * A negative extent is reported the same way, so that it is never converted
+ * into a very large unsigned size.
+ *
+ * \param shape The allocation shape.
+ * \param unit_size The size of one element, in the unit the caller works in.
+ * \param out Set to the size when the result is kConstant, left untouched otherwise.
+ * \return Whether the size is dynamic, constant, or not representable.
+ */
+inline ConstantSizeKind GetConstantAllocationSize(const ffi::Array<PrimExpr>& shape,
+                                                  uint64_t unit_size, uint64_t* out) {
+  ffi::BigInt size(unit_size);
+  for (const PrimExpr& extent : shape) {
+    const auto* imm = extent.as<IntImmNode>();
+    if (imm == nullptr) return ConstantSizeKind::kDynamic;
+    size *= imm->value;
+  }
+  std::optional<uint64_t> fits = size.as<uint64_t>();
+  if (!fits.has_value()) return ConstantSizeKind::kUnrepresentable;
+  *out = *fits;
+  return ConstantSizeKind::kConstant;
+}
+
+/*!
+ * \brief Multiply two sizes, saturating instead of wrapping around.
+ * \param a The left operand.
+ * \param b The right operand.
+ * \return The product, or the largest uint64_t if the product does not fit.
+ */
+inline uint64_t SaturatingMul(uint64_t a, uint64_t b) {
+  constexpr uint64_t kMaxSize = std::numeric_limits<uint64_t>::max();
+  if (a == 0 || b == 0) return 0;
+  return a > kMaxSize / b ? kMaxSize : a * b;
+}
+
 /*!
  * \brief Create an int32 constant
  * \param index the value of the constant

src/tirx/transform/lower_tvm_builtin.cc

@@ -265,13 +265,23 @@ class BuiltinLower : public StmtExprMutator {
       return stmt;
     }
     int64_t nbytes = GetVectorBytes(op->buffer->dtype);
+    uint64_t const_nbytes = 0;
+    ConstantSizeKind size_kind =
+        GetConstantAllocationSize(op->buffer->shape, static_cast<uint64_t>(nbytes), &const_nbytes);
+    if (size_kind == ConstantSizeKind::kUnrepresentable) {
+      // total_bytes below would fold to a wrapped constant, so the workspace
+      // request would be smaller than the buffer it is meant to hold.
+      TVM_FFI_THROW(ValueError) << "Cannot allocate buffer " << op->buffer.name() << " with shape "
+                                << op->buffer->shape << " and dtype " << op->buffer->dtype
+                                << ": its size in bytes does not fit in the 64-bit byte count "
+                                   "passed to TVMBackendAllocWorkspace";
+    }
     if (const auto* dev_type = device_type_.as<IntImmNode>();
         dev_type && dev_type->value == kDLCPU) {
       auto storage_scope = op->buffer->storage_scope;
       if (storage_scope == "global") {
-        auto constant_size = stmt.as_or_throw<AllocBuffer>().ConstantAllocationSize();
-        if (constant_size.has_value() && constant_size.value() > 0 &&
-            static_cast<size_t>(constant_size.value()) * nbytes < runtime::kMaxStackAlloca) {
+        if (size_kind == ConstantSizeKind::kConstant && const_nbytes > 0 &&
+            const_nbytes < runtime::kMaxStackAlloca) {
           return stmt;
         }
       }

src/tirx/transform/storage_rewrite.cc

@@ -1058,12 +1058,15 @@ class StoragePlanRewriter : public StmtExprMutator {
                     !alloc->buffer->dtype.IsScalableVector() &&
                     src_entry->elem_type == alloc->buffer->dtype.WithLanes(1) &&
                     visitor->Check(s.stmt, var, src)) {
-                  int64_t const_size = AllocBuffer(ffi::GetRef<AllocBuffer>(alloc))
-                                           .ConstantAllocationSize()
-                                           .value_or(0);
-                  uint64_t const_nbits = static_cast<uint64_t>(const_size) *
-                                         alloc->buffer->dtype.bits() * alloc->buffer->dtype.lanes();
-                  if (src_entry->const_nbits == const_nbits && !inplace_found) {
+                  uint64_t elem_bits = static_cast<uint64_t>(alloc->buffer->dtype.bits()) *
+                                       alloc->buffer->dtype.lanes();
+                  uint64_t const_nbits = 0;
+                  ConstantSizeKind size_kind =
+                      GetConstantAllocationSize(alloc->buffer->shape, elem_bits, &const_nbits);
+                  // A size that does not fit in uint64_t cannot be compared against
+                  // the source entry; FindAlloc below rejects such a buffer.
+                  if (size_kind != ConstantSizeKind::kUnrepresentable &&
+                      src_entry->const_nbits == const_nbits && !inplace_found) {
                     // successfully inplace
                     dst_entry = src_entry;
                     inplace_flag.insert(src);
@@ -1141,10 +1144,20 @@ class StoragePlanRewriter : public StmtExprMutator {
     bool is_scalable_vector = op->buffer->dtype.IsScalableVector();
     uint64_t op_elem_bits =
         is_scalable_vector ? 0 : op->buffer->dtype.bits() * op->buffer->dtype.lanes();
-    int64_t const_size =
-        AllocBuffer(ffi::GetRef<AllocBuffer>(op)).ConstantAllocationSize().value_or(0);
-    uint64_t const_nbits =
-        is_scalable_vector ? 0 : static_cast<uint64_t>(const_size * op_elem_bits);
+    // A size that is not known at compile time leaves const_nbits at zero, so the
+    // buffer gets an allocation of its own below. A constant size that does not
+    // fit in uint64_t is rejected rather than planned with the wrapped value.
+    uint64_t const_nbits = 0;
+    if (!is_scalable_vector) {
+      ConstantSizeKind size_kind =
+          GetConstantAllocationSize(op->buffer->shape, op_elem_bits, &const_nbits);
+      if (size_kind == ConstantSizeKind::kUnrepresentable) {
+        TVM_FFI_THROW(ValueError) << "Cannot plan storage for buffer " << op->buffer.name()
+                                  << " with shape " << op->buffer->shape << " and dtype "
+                                  << op->buffer->dtype
+                                  << ": its size in bits does not fit in 64 bits";
+      }
+    }
 
     // If the size of the array isn't known at compile-time, it must
     // have its own allocation with size determined at runtime.
@@ -1168,7 +1181,7 @@ class StoragePlanRewriter : public StmtExprMutator {
       // constant allocation.
       auto begin = const_free_map_.lower_bound(const_nbits / match_range);
       auto mid = const_free_map_.lower_bound(const_nbits);
-      auto end = const_free_map_.upper_bound(const_nbits * match_range);
+      auto end = const_free_map_.upper_bound(SaturatingMul(const_nbits, match_range));
       // start looking at the buffer that is bigger than the required size first
       for (auto it = mid; it != end; ++it) {
         StorageEntry* e = it->second;

tests/python/s_tir/analysis/test_s_tir_analysis_verify_gpu_code.py

@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=missing-module-docstring,missing-function-docstring,missing-class-docstring
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+
+CONSTRAINTS = {
+    "max_shared_memory_per_block": 49152,
+    "max_local_memory_per_block": 2147483647,
+    "max_threads_per_block": 1024,
+}
+
+
+def _shared_memory_kernel(shape, dtype):
+    @T.prim_func(s_tir=True)
+    def main(a: T.handle) -> None:
+        T.func_attr({"global_symbol": "main", "T.noalias": True})
+        threadIdx_x = T.env_thread("threadIdx.x")
+        blockIdx_x = T.env_thread("blockIdx.x")
+        A = T.match_buffer(a, [64], dtype=dtype)
+        T.launch_thread(blockIdx_x, 1)
+        T.launch_thread(threadIdx_x, 64)
+        shared = T.decl_buffer(shape, dtype, scope="shared")
+        shared[0, 0] = A[threadIdx_x]
+        A[threadIdx_x] = shared[0, 0]
+
+    return main
+
+
+def test_shared_memory_within_limit():
+    assert tvm.s_tir.analysis.verify_gpu_code(_shared_memory_kernel((8, 8), "float32"), CONSTRAINTS)
+
+
+def test_shared_memory_over_limit():
+    assert not tvm.s_tir.analysis.verify_gpu_code(
+        _shared_memory_kernel((256, 256), "float32"), CONSTRAINTS
+    )
+
+
+@pytest.mark.parametrize(
+    "shape,dtype",
+    [
+        # 2**32 * 2**32 elements: the element count alone does not fit.
+        ((2**32, 2**32), "int8"),
+        # 2**62 elements of 4 bytes: the element count fits, the byte count does not.
+        ((2**31, 2**31), "float32"),
+    ],
+)
+def test_shared_memory_size_that_does_not_fit(shape, dtype):
+    """A size that wraps around must not be counted as zero shared memory."""
+    assert not tvm.s_tir.analysis.verify_gpu_code(_shared_memory_kernel(shape, dtype), CONSTRAINTS)
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py

@@ -319,5 +319,36 @@ def main():
     tvm.ir.assert_structural_equal(After, Expected)
 
 
+def _global_alloc_func(shape, dtype="int8"):
+    """A CPU function holding a single global allocation of the given shape."""
+    buf = tvm.tirx.decl_buffer(shape, dtype=dtype, scope="global")
+    body = tvm.tirx.AttrStmt(
+        tvm.tirx.Var("dev", "int32"),
+        "device_id",
+        tvm.tirx.IntImm("int32", 0),
+        tvm.tirx.AllocBuffer(buf),
+    )
+    attrs = tvm.ir.DictAttrs({"target": tvm.target.Target("llvm")})
+    return tvm.tirx.PrimFunc([], body, attrs=attrs)
+
+
+@pytest.mark.parametrize("shape", [(2**32, 2**32), (2**62, 4), (2**62, 5)])
+def test_workspace_size_that_does_not_fit_is_rejected(shape):
+    """An int8 allocation of 2**64 bytes or more must not wrap around.
+
+    Folding the byte count in uint64 turns the first two shapes into a request
+    for 0 bytes and the third into a request for 2**62 instead of 5 * 2**62.
+    """
+    mod = tvm.IRModule({"main": _global_alloc_func(shape)})
+    with pytest.raises(ValueError, match="does not fit"):
+        tvm.tirx.transform.LowerTVMBuiltin()(mod)
+
+
+def test_workspace_size_that_fits_is_unchanged():
+    mod = tvm.IRModule({"main": _global_alloc_func((2**31, 4))})
+    lowered = tvm.tirx.transform.LowerTVMBuiltin()(mod)
+    assert str(2**31 * 4) in str(lowered["main"])
+
+
 if __name__ == "__main__":
     tvm.testing.main()