diff --git a/src/s_tir/analysis/verify_gpu_code.cc b/src/s_tir/analysis/verify_gpu_code.cc index b56756ad2380..e11411935e2e 100644 --- a/src/s_tir/analysis/verify_gpu_code.cc +++ b/src/s_tir/analysis/verify_gpu_code.cc @@ -33,6 +33,9 @@ #include #include +#include +#include + #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()) { - const_size = static_cast(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(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(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::max(); + size_t value = static_cast(std::min(nbytes, kMaxTotal)); + *total = *total > kMaxTotal - value ? kMaxTotal : *total + value; + } + void Reset_() { local_memory_per_block_ = 0; shared_memory_per_block_ = 0; diff --git a/src/s_tir/transform/merge_shared_memory_allocations.cc b/src/s_tir/transform/merge_shared_memory_allocations.cc index 321db8f39d74..8e3e732467e0 100644 --- a/src/s_tir/transform/merge_shared_memory_allocations.cc +++ b/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& extents) { - int64_t result = 1; - for (size_t i = 0; i < extents.size(); ++i) { - if (const IntImmNode* int_size = extents[i].as()) { - auto product = (result * int_size->value).as(); - 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 alloc_shape = GetBufferAllocationShape(buf); DLDataType dtype = buf->dtype->dtype; uint64_t op_elem_bits = static_cast(dtype.bits) * dtype.lanes; - uint64_t const_nbits = - static_cast(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 diff --git a/src/tirx/transform/ir_utils.h b/src/tirx/transform/ir_utils.h index 4182bba09349..094d1ab06827 100644 --- a/src/tirx/transform/ir_utils.h +++ b/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 #include #include #include @@ -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& shape, + uint64_t unit_size, uint64_t* out) { + ffi::BigInt size(unit_size); + for (const PrimExpr& extent : shape) { + const auto* imm = extent.as(); + if (imm == nullptr) return ConstantSizeKind::kDynamic; + size *= imm->value; + } + std::optional fits = size.as(); + 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::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 diff --git a/src/tirx/transform/lower_tvm_builtin.cc b/src/tirx/transform/lower_tvm_builtin.cc index 27762104b997..4ca5ee6670f5 100644 --- a/src/tirx/transform/lower_tvm_builtin.cc +++ b/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(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(); dev_type && dev_type->value == kDLCPU) { auto storage_scope = op->buffer->storage_scope; if (storage_scope == "global") { - auto constant_size = stmt.as_or_throw().ConstantAllocationSize(); - if (constant_size.has_value() && constant_size.value() > 0 && - static_cast(constant_size.value()) * nbytes < runtime::kMaxStackAlloca) { + if (size_kind == ConstantSizeKind::kConstant && const_nbytes > 0 && + const_nbytes < runtime::kMaxStackAlloca) { return stmt; } } diff --git a/src/tirx/transform/storage_rewrite.cc b/src/tirx/transform/storage_rewrite.cc index d6a6d7bdd275..235b91ac9f11 100644 --- a/src/tirx/transform/storage_rewrite.cc +++ b/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(alloc)) - .ConstantAllocationSize() - .value_or(0); - uint64_t const_nbits = static_cast(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(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(op)).ConstantAllocationSize().value_or(0); - uint64_t const_nbits = - is_scalable_vector ? 0 : static_cast(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; diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_verify_gpu_code.py b/tests/python/s_tir/analysis/test_s_tir_analysis_verify_gpu_code.py new file mode 100644 index 000000000000..537f810b8783 --- /dev/null +++ b/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() diff --git a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py b/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py index bacbe70d42f2..8a27df35d7d8 100644 --- a/tests/python/tirx-transform/test_tir_transform_lower_tvm_builtin.py +++ b/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() diff --git a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py index b3133cba1e91..0bdcccb19a8a 100644 --- a/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py +++ b/tests/python/tirx-transform/test_tir_transform_storage_rewrite.py @@ -528,5 +528,58 @@ def main(A: T.Buffer(16, "float32"), D: T.Buffer(16, "float32")): tvm.ir.assert_structural_equal(After, Expected) +def test_allocation_size_that_does_not_fit_is_rejected(): + """A buffer whose size in bits wraps around must not be planned as a small one. + + 8 bits * (2**61 + 128) is 2**64 + 1024, which wraps to 1024 bits: exactly the + size of A, so the planner used to place B inside A's 128-byte allocation. + """ + huge = 2**61 + 128 + + @T.prim_func(s_tir=True) + def func(): + for j in range(128): + A = T.alloc_buffer((128,), "int8") + A[j] = T.int8(1) + for j in range(128): + B = T.alloc_buffer((huge,), "int8") + B[j] = T.int8(2) + + mod = tvm.IRModule.from_expr(func) + with pytest.raises(ValueError, match="does not fit"): + tvm.tirx.transform.StorageRewrite()(mod) + + +def test_reuse_search_range_does_not_wrap(): + """The free-list search range of a large buffer must not wrap around. + + The upper bound of the search is const_nbits * 16. For B that product is + 2**65, which wraps to 0, so the range ended before it started and the search + walked past the end of the free list. + """ + big = 2**58 + + @T.prim_func(s_tir=True) + def func(): + for j in range(128): + A = T.alloc_buffer((128,), "int8") + A[j] = T.int8(1) + for j in range(128): + B = T.alloc_buffer((big,), "int8") + B[j] = T.int8(2) + + mod = tvm.IRModule.from_expr(func) + body = tvm.tirx.transform.StorageRewrite()(mod)["func"].body + + sizes = [] + + def collect(n): + if isinstance(n, tvm.tirx.AllocBuffer): + sizes.append(n.buffer.ty.shape[0].value) + + tvm_ffi.structural_walk(body, collect) + assert sorted(sizes) == [128, big] + + if __name__ == "__main__": tvm.testing.main()