This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f8e8af1759 [Runtime] Register backend memory scopes with EnumDef
(#20725)
f8e8af1759 is described below
commit f8e8af175984ccf4522dfea8282a3233aa50bec1
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 14:54:48 2026 -0400
[Runtime] Register backend memory scopes with EnumDef (#20725)
Register backend memory scopes directly with EnumDef and explicit
allocation attributes, removing the backend registration helper and
redundant Vulkan WMMA registrations.
Add a C++ MemoryScope::IsDefined query that reads existing enum state
without initializing scopes. OpenCL and Hexagon guard each name while
retaining compiler and runtime registration entry points, preserving
canonical identities and existing attributes.
Inline the membership and registration shims while keeping lookup and
builtin initialization in the runtime. Expose allocation policy through
its typed attribute constant, with storage rewrite querying and
validating that metadata directly instead of using a MemoryScope member
accessor.
---
include/tvm/runtime/memory_scope.h | 35 +++++++++--
src/backend/cuda/codegen/target_kind.cc | 9 +--
src/backend/hexagon/memory_scope.h | 20 ++++---
src/backend/memory_scope.h | 70 ----------------------
src/backend/metal/codegen/target_kind.cc | 9 +--
src/backend/opencl/memory_scope.h | 28 +++++----
src/backend/trn/codegen/target_kind.cc | 9 +--
src/backend/vulkan/codegen/target_kind.cc | 6 --
src/codegen/llvm/codegen_cpu.cc | 8 ++-
src/runtime/memory_scope.cc | 36 +----------
src/tirx/transform/storage_rewrite.cc | 36 ++++++++---
tests/cpp/runtime/memory_scope_test.cc | 50 +++++++++++++---
.../test_tirx_transform_storage_rewrite.py | 16 +++--
13 files changed, 164 insertions(+), 168 deletions(-)
diff --git a/include/tvm/runtime/memory_scope.h
b/include/tvm/runtime/memory_scope.h
index ee933d85aa..4755b59099 100644
--- a/include/tvm/runtime/memory_scope.h
+++ b/include/tvm/runtime/memory_scope.h
@@ -21,12 +21,17 @@
#define TVM_RUNTIME_MEMORY_SCOPE_H_
#include <tvm/ffi/enum.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/runtime/base.h>
#include <ostream>
+#include <string_view>
namespace tvm {
namespace runtime {
+namespace detail {
+TVM_RUNTIME_DLL void RegisterBuiltinMemoryScopes();
+} // namespace detail
/*! \brief Allocation behavior associated with a memory scope. */
enum class MemoryAllocProperty : int {
@@ -37,7 +42,8 @@ enum class MemoryAllocProperty : int {
namespace attr {
/*! \brief Integer-valued MemoryAllocProperty; absence denotes kNormal. */
-inline constexpr const char* kMemoryAllocProperty =
"runtime.memory_alloc_property";
+inline constexpr ffi::reflection::AttrKey<MemoryAllocProperty>
kMemoryAllocProperty{
+ "runtime.memory_alloc_property"};
} // namespace attr
/*! \brief Canonical identity of an extensible memory space. */
@@ -61,12 +67,33 @@ class MemoryScope : public ffi::Enum {
explicit MemoryScope(const ffi::String& name) : ffi::Enum(Get(name)) {}
/*! \brief Return the canonical spelling. */
const ffi::String& name() const { return get()->_str_index; }
- /*! \brief Return the allocation property assigned by the scope owner. */
- TVM_RUNTIME_DLL MemoryAllocProperty alloc_property() const;
+ /*! \brief Test whether a name is registered, without initializing memory
scopes. */
+ static bool IsDefined(const ffi::String& name) {
+ namespace refl = ffi::reflection;
+ constexpr auto key = refl::AsByteArray(ffi::type_attr::kEnumState.name);
+ if (TVMFFIGetTypeAttrColumn(&key) == nullptr) return false;
+ auto value = refl::TypeAttrColumn(
+
ffi::type_attr::kEnumState)[MemoryScopeObj::_GetOrAllocRuntimeTypeIndex()];
+ if (value == nullptr) return false;
+ TVM_FFI_CHECK(value.type_index() ==
ffi::EnumStateObj::_GetOrAllocRuntimeTypeIndex(), TypeError)
+ << "Expected enum state for memory scopes";
+ // The derived EnumState type index may still be deferred during static
initialization.
+ auto* state = static_cast<const
ffi::EnumStateObj*>(value.cast<ffi::ObjectRef>().get());
+ return state->indexes.count(name) != 0;
+ }
/*! \brief Look up a registered name. Unknown names raise ValueError. */
TVM_RUNTIME_DLL static MemoryScope Get(const ffi::String& name);
/*! \brief Register a new name. Duplicate names raise ValueError. */
- TVM_RUNTIME_DLL static MemoryScope Register(const ffi::String& name);
+ static MemoryScope Register(const ffi::String& name) {
+ namespace refl = ffi::reflection;
+ detail::RegisterBuiltinMemoryScopes();
+ TVM_FFI_CHECK(!name.empty(), ValueError) << "Memory scope name must not be
empty";
+ TVM_FFI_CHECK(std::string_view(name.data(), name.size()).find('\0') ==
std::string_view::npos,
+ ValueError)
+ << "Memory scope name must not contain a null character";
+ refl::EnumDef<MemoryScopeObj>(name.c_str());
+ return Get(name);
+ }
/*! \brief Return the canonical global singleton. */
TVM_RUNTIME_DLL static const MemoryScope& Global();
bool operator==(const MemoryScope& other) const { return same_as(other); }
diff --git a/src/backend/cuda/codegen/target_kind.cc
b/src/backend/cuda/codegen/target_kind.cc
index 591b46517a..ac0732419f 100644
--- a/src/backend/cuda/codegen/target_kind.cc
+++ b/src/backend/cuda/codegen/target_kind.cc
@@ -23,17 +23,17 @@
*/
#include <dlpack/dlpack.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/ir/expr.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/logging.h>
+#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
#include <string>
-#include "../../memory_scope.h"
-
namespace tvm {
namespace backend {
@@ -119,8 +119,9 @@ ffi::Map<ffi::String, ffi::Any>
UpdateNVPTXAttrs(ffi::Map<ffi::String, ffi::Any>
TVM_FFI_STATIC_INIT_BLOCK() {
for (const char* name : {"tmem", "wmma.matrix_a", "wmma.matrix_b",
"wmma.accumulator",
"mma.matrixA", "mma.matrixB", "mma.matrixC"}) {
- tvm::backend::RegisterBackendMemoryScope(
- name, tvm::runtime::MemoryAllocProperty::kSeparateAllocations);
+ tvm::ffi::reflection::EnumDef<tvm::runtime::MemoryScopeObj>(name).set_attr(
+ tvm::runtime::attr::kMemoryAllocProperty,
+
static_cast<int>(tvm::runtime::MemoryAllocProperty::kSeparateAllocations));
}
using namespace tvm;
using namespace tvm::backend::cuda;
diff --git a/src/backend/hexagon/memory_scope.h
b/src/backend/hexagon/memory_scope.h
index 99eedfc56f..fe31721da5 100644
--- a/src/backend/hexagon/memory_scope.h
+++ b/src/backend/hexagon/memory_scope.h
@@ -20,19 +20,25 @@
#ifndef TVM_BACKEND_HEXAGON_MEMORY_SCOPE_H_
#define TVM_BACKEND_HEXAGON_MEMORY_SCOPE_H_
-#include "../memory_scope.h"
+#include <tvm/ffi/reflection/enum_def.h>
+#include <tvm/runtime/memory_scope.h>
namespace tvm {
namespace backend {
namespace hexagon {
inline void RegisterMemoryScopes() {
- static const bool initialized = [] {
- RegisterBackendMemoryScope("global.ddr",
runtime::MemoryAllocProperty::kSingleAllocation);
- RegisterBackendMemoryScope("global.vtcm",
runtime::MemoryAllocProperty::kNormal);
- return true;
- }();
- (void)initialized;
+ using runtime::MemoryAllocProperty;
+ if (!runtime::MemoryScope::IsDefined("global.ddr")) {
+ ffi::reflection::EnumDef<runtime::MemoryScopeObj>("global.ddr")
+ .set_attr(runtime::attr::kMemoryAllocProperty,
+ static_cast<int>(MemoryAllocProperty::kSingleAllocation));
+ }
+ if (!runtime::MemoryScope::IsDefined("global.vtcm")) {
+ ffi::reflection::EnumDef<runtime::MemoryScopeObj>("global.vtcm")
+ .set_attr(runtime::attr::kMemoryAllocProperty,
+ static_cast<int>(MemoryAllocProperty::kNormal));
+ }
}
} // namespace hexagon
diff --git a/src/backend/memory_scope.h b/src/backend/memory_scope.h
deleted file mode 100644
index ad9f5eee1f..0000000000
--- a/src/backend/memory_scope.h
+++ /dev/null
@@ -1,70 +0,0 @@
-/*
- * 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.
- */
-
-#ifndef TVM_BACKEND_MEMORY_SCOPE_H_
-#define TVM_BACKEND_MEMORY_SCOPE_H_
-
-#include <tvm/ffi/reflection/enum_def.h>
-#include <tvm/runtime/memory_scope.h>
-
-#include <initializer_list>
-#include <utility>
-
-namespace tvm {
-namespace backend {
-
-// A compiler backend and its runtime sidecar may define the same owned scopes.
-// Accept only matching explicit metadata when the other component loaded
first.
-inline void RegisterBackendMemoryScope(
- const char* name, runtime::MemoryAllocProperty property,
- std::initializer_list<std::pair<const char*, ffi::Any>> attributes = {}) {
- runtime::MemoryScope::Global();
- static ffi::reflection::TypeAttrColumn
state_column(ffi::type_attr::kEnumState);
- ffi::AnyView state_value =
state_column[runtime::MemoryScopeObj::_GetOrAllocRuntimeTypeIndex()];
- TVM_FFI_CHECK(state_value.type_index() ==
ffi::EnumStateObj::_GetOrAllocRuntimeTypeIndex(),
- TypeError)
- << "Expected enum state for memory scopes, but got " <<
state_value.GetTypeKey();
- // Backend registration runs before GCC initializes this DSO's inline type
- // indices. As in EnumDef, avoid a cast that reads the deferred EnumState
index.
- auto state =
ffi::details::ObjectUnsafe::ObjectRefFromObjectPtr<ffi::EnumState>(
- ffi::details::ObjectUnsafe::ObjectPtrFromUnowned<ffi::Object>(
- state_value.CopyToTVMFFIAny().v_obj));
- if (auto existing = state->indexes.Get(ffi::String(name))) {
- auto check_attribute = [&](const char* key, ffi::Any expected) {
- auto column = state->attrs.Get(key);
- TVM_FFI_CHECK(column.has_value(), ValueError)
- << "Memory scope " << name << " is missing backend attribute " <<
key;
- auto actual = column->Get(*existing);
- TVM_FFI_CHECK(actual.has_value() && actual->type_index() ==
expected.type_index() &&
- ffi::AnyEqual()(*actual, expected),
- ValueError)
- << "Memory scope " << name << " has incompatible backend attribute "
<< key;
- };
- check_attribute(runtime::attr::kMemoryAllocProperty,
static_cast<int>(property));
- for (const auto& [key, value] : attributes) check_attribute(key, value);
- return;
- }
- ffi::reflection::EnumDef<runtime::MemoryScopeObj> definition(name);
- definition.set_attr(runtime::attr::kMemoryAllocProperty,
static_cast<int>(property));
- for (const auto& [key, value] : attributes) definition.set_attr(key, value);
-}
-
-} // namespace backend
-} // namespace tvm
-#endif // TVM_BACKEND_MEMORY_SCOPE_H_
diff --git a/src/backend/metal/codegen/target_kind.cc
b/src/backend/metal/codegen/target_kind.cc
index 078128c57c..99dbeffa1f 100644
--- a/src/backend/metal/codegen/target_kind.cc
+++ b/src/backend/metal/codegen/target_kind.cc
@@ -23,12 +23,12 @@
*/
#include <dlpack/dlpack.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/runtime/base.h>
+#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
-#include "../../memory_scope.h"
-
namespace tvm {
namespace codegen {
@@ -41,8 +41,9 @@ void RegisterMetalIntrinRules();
TVM_FFI_STATIC_INIT_BLOCK() {
for (const char* name : {"metal.simdgroup", "metal.cooperative_tensor"}) {
- tvm::backend::RegisterBackendMemoryScope(
- name, tvm::runtime::MemoryAllocProperty::kSeparateAllocations);
+ tvm::ffi::reflection::EnumDef<tvm::runtime::MemoryScopeObj>(name).set_attr(
+ tvm::runtime::attr::kMemoryAllocProperty,
+
static_cast<int>(tvm::runtime::MemoryAllocProperty::kSeparateAllocations));
}
using namespace tvm;
namespace refl = tvm::ffi::reflection;
diff --git a/src/backend/opencl/memory_scope.h
b/src/backend/opencl/memory_scope.h
index 650625d3dd..70e087786b 100644
--- a/src/backend/opencl/memory_scope.h
+++ b/src/backend/opencl/memory_scope.h
@@ -20,7 +20,8 @@
#ifndef TVM_BACKEND_OPENCL_MEMORY_SCOPE_H_
#define TVM_BACKEND_OPENCL_MEMORY_SCOPE_H_
-#include "../memory_scope.h"
+#include <tvm/ffi/reflection/enum_def.h>
+#include <tvm/runtime/memory_scope.h>
namespace tvm {
namespace backend {
@@ -30,17 +31,20 @@ inline constexpr const char* kTextureStorage =
"opencl.texture_storage";
inline void RegisterMemoryScopes() {
using runtime::MemoryAllocProperty;
- static const bool initialized = [] {
- RegisterBackendMemoryScope("texture",
MemoryAllocProperty::kSeparateAllocations,
- {{kTextureStorage, true}});
- for (const char* name : {"global.texture", "global.texture-1d",
"global.texture-nchw",
- "global.texture-nhwc", "global.texture-weight"}) {
- RegisterBackendMemoryScope(name, MemoryAllocProperty::kSingleAllocation,
- {{kTextureStorage, true}});
- }
- return true;
- }();
- (void)initialized;
+ if (!runtime::MemoryScope::IsDefined("texture")) {
+ ffi::reflection::EnumDef<runtime::MemoryScopeObj>("texture")
+ .set_attr(runtime::attr::kMemoryAllocProperty,
+ static_cast<int>(MemoryAllocProperty::kSeparateAllocations))
+ .set_attr(kTextureStorage, true);
+ }
+ for (const char* name : {"global.texture", "global.texture-1d",
"global.texture-nchw",
+ "global.texture-nhwc", "global.texture-weight"}) {
+ if (runtime::MemoryScope::IsDefined(name)) continue;
+ ffi::reflection::EnumDef<runtime::MemoryScopeObj>(name)
+ .set_attr(runtime::attr::kMemoryAllocProperty,
+ static_cast<int>(MemoryAllocProperty::kSingleAllocation))
+ .set_attr(kTextureStorage, true);
+ }
}
} // namespace opencl
diff --git a/src/backend/trn/codegen/target_kind.cc
b/src/backend/trn/codegen/target_kind.cc
index f166a60c5e..573e71a2ce 100644
--- a/src/backend/trn/codegen/target_kind.cc
+++ b/src/backend/trn/codegen/target_kind.cc
@@ -23,12 +23,12 @@
*/
#include <dlpack/dlpack.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/runtime/base.h>
+#include <tvm/runtime/memory_scope.h>
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
-#include "../../memory_scope.h"
-
namespace tvm {
namespace codegen {
@@ -44,8 +44,9 @@ void RegisterTRNTransforms();
TVM_FFI_STATIC_INIT_BLOCK() {
for (const char* name : {"trn.sbuf", "trn.psum"}) {
- tvm::backend::RegisterBackendMemoryScope(
- name, tvm::runtime::MemoryAllocProperty::kSeparateAllocations);
+ tvm::ffi::reflection::EnumDef<tvm::runtime::MemoryScopeObj>(name).set_attr(
+ tvm::runtime::attr::kMemoryAllocProperty,
+
static_cast<int>(tvm::runtime::MemoryAllocProperty::kSeparateAllocations));
}
using namespace tvm;
TargetKindDef("trn")
diff --git a/src/backend/vulkan/codegen/target_kind.cc
b/src/backend/vulkan/codegen/target_kind.cc
index 28e96c3ae6..52ad669bc3 100644
--- a/src/backend/vulkan/codegen/target_kind.cc
+++ b/src/backend/vulkan/codegen/target_kind.cc
@@ -27,8 +27,6 @@
#include <tvm/target/target.h>
#include <tvm/target/target_kind.h>
-#include "../../memory_scope.h"
-
namespace tvm {
#ifdef TVM_ENABLE_SPIRV
@@ -42,10 +40,6 @@ void RegisterVulkanIntrinRules();
} // namespace tvm
TVM_FFI_STATIC_INIT_BLOCK() {
- for (const char* name : {"wmma.matrix_a", "wmma.matrix_b",
"wmma.accumulator"}) {
- tvm::backend::RegisterBackendMemoryScope(
- name, tvm::runtime::MemoryAllocProperty::kSeparateAllocations);
- }
using namespace tvm;
namespace refl = tvm::ffi::reflection;
diff --git a/src/codegen/llvm/codegen_cpu.cc b/src/codegen/llvm/codegen_cpu.cc
index 3de6a46130..64722f2718 100644
--- a/src/codegen/llvm/codegen_cpu.cc
+++ b/src/codegen/llvm/codegen_cpu.cc
@@ -48,11 +48,13 @@
#include <llvm/Target/TargetMachine.h>
#include <llvm/Transforms/Utils/ModuleUtils.h>
#include <tvm/ffi/extra/module.h>
+#include <tvm/ffi/reflection/enum_def.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/analysis.h>
#include <tvm/ir/function.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/logging.h>
+#include <tvm/runtime/memory_scope.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/abi.h>
#include <tvm/tirx/op/cpu.h>
@@ -65,7 +67,6 @@
#include <unordered_map>
#include <unordered_set>
-#include "../../backend/memory_scope.h"
#include "../../tirx/transform/ir_utils.h"
#include "llvm_instance.h"
@@ -1222,8 +1223,9 @@ void CodeGenCPU::Dispatch_(const ForNode* op) {
}
TVM_FFI_STATIC_INIT_BLOCK() {
- tvm::backend::RegisterBackendMemoryScope("amx.tmm",
-
tvm::runtime::MemoryAllocProperty::kSeparateAllocations);
+ ffi::reflection::EnumDef<runtime::MemoryScopeObj>("amx.tmm").set_attr(
+ runtime::attr::kMemoryAllocProperty,
+ static_cast<int>(runtime::MemoryAllocProperty::kSeparateAllocations));
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed("tvm.codegen.llvm.target_cpu",
[](const ffi::PackedArgs& targs, ffi::Any* rv) {
diff --git a/src/runtime/memory_scope.cc b/src/runtime/memory_scope.cc
index 18eb682e07..a2b58c9bb8 100644
--- a/src/runtime/memory_scope.cc
+++ b/src/runtime/memory_scope.cc
@@ -22,14 +22,11 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/memory_scope.h>
-#include <string_view>
-
namespace tvm {
namespace runtime {
-namespace {
namespace refl = ffi::reflection;
-void RegisterBuiltinScopes() {
+void detail::RegisterBuiltinMemoryScopes() {
// Lookups and defaults can run in other translation units' static
initializers.
// Initialize the complete set once, before any lookup, without relying on
order.
static const bool initialized = [] {
@@ -48,46 +45,19 @@ void RegisterBuiltinScopes() {
}();
(void)initialized;
}
-} // namespace
-
-MemoryAllocProperty MemoryScope::alloc_property() const {
- static ffi::reflection::TypeAttrColumn
state_column(ffi::type_attr::kEnumState);
- auto state =
state_column[MemoryScopeObj::RuntimeTypeIndex()].cast<ffi::EnumState>();
- if (auto column = state->attrs.Get(attr::kMemoryAllocProperty)) {
- if (auto value = column->Get(*this)) {
- int property = value->cast<int>();
- TVM_FFI_CHECK(property >= static_cast<int>(MemoryAllocProperty::kNormal)
&&
- property <=
static_cast<int>(MemoryAllocProperty::kSeparateAllocations),
- ValueError)
- << "Invalid memory allocation property for " << name();
- return static_cast<MemoryAllocProperty>(property);
- }
- }
- return MemoryAllocProperty::kNormal;
-}
MemoryScope MemoryScope::Get(const ffi::String& name) {
- RegisterBuiltinScopes();
+ detail::RegisterBuiltinMemoryScopes();
return
ffi::EnumObj::_GetByStrIndex<MemoryScopeObj>(name).as_or_throw<MemoryScope>();
}
-MemoryScope MemoryScope::Register(const ffi::String& name) {
- RegisterBuiltinScopes();
- TVM_FFI_CHECK(!name.empty(), ValueError) << "Memory scope name must not be
empty";
- TVM_FFI_CHECK(std::string_view(name.data(), name.size()).find('\0') ==
std::string_view::npos,
- ValueError)
- << "Memory scope name must not contain a null character";
- refl::EnumDef<MemoryScopeObj>(name.c_str());
- return Get(name);
-}
-
const MemoryScope& MemoryScope::Global() {
static const MemoryScope scope = Get("global");
return scope;
}
TVM_FFI_STATIC_INIT_BLOCK() {
- RegisterBuiltinScopes();
+ detail::RegisterBuiltinMemoryScopes();
refl::TypeAttrDef<MemoryScopeObj>()
.def(ffi::type_attr::kDataToJson.name,
[](const MemoryScopeObj* node) -> ffi::String { return
node->_str_index; })
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 144a2ede28..36c996143e 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -59,6 +59,22 @@ using runtime::MemoryScope;
namespace {
+MemoryAllocProperty GetMemoryAllocProperty(const MemoryScope& scope) {
+ static ffi::reflection::TypeAttrColumn
state_column(ffi::type_attr::kEnumState);
+ auto state =
state_column[runtime::MemoryScopeObj::RuntimeTypeIndex()].cast<ffi::EnumState>();
+ if (auto column =
state->attrs.Get(runtime::attr::kMemoryAllocProperty.name)) {
+ if (auto value = column->Get(scope)) {
+ int property = value->cast<int>();
+ TVM_FFI_CHECK(property >= static_cast<int>(MemoryAllocProperty::kNormal)
&&
+ property <=
static_cast<int>(MemoryAllocProperty::kSeparateAllocations),
+ ValueError)
+ << "Invalid memory allocation property for " << scope.name();
+ return static_cast<MemoryAllocProperty>(property);
+ }
+ }
+ return MemoryAllocProperty::kNormal;
+}
+
ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
if (auto var = data.as<Var>()) {
return var;
@@ -801,7 +817,7 @@ class StoragePlanRewriter : public StmtExprMutator {
static bool RequiresExactAllocation(const CallNode* call) {
auto type = call->ty_args[0].as_or_throw<TensorType>();
- if (type->scope.alloc_property() ==
MemoryAllocProperty::kSeparateAllocations) {
+ if (GetMemoryAllocProperty(type->scope) ==
MemoryAllocProperty::kSeparateAllocations) {
return true;
}
// Flat packing and lifetime reuse cannot reconstruct arbitrary physical
views.
@@ -923,7 +939,7 @@ class StoragePlanRewriter : public StmtExprMutator {
std::vector<StorageEntry*>& vec = kv.second;
for (size_t i = 0; i < vec.size(); ++i) {
StorageEntry* e = vec[i];
- if (e->scope.alloc_property() !=
MemoryAllocProperty::kSingleAllocation) continue;
+ if (GetMemoryAllocProperty(e->scope) !=
MemoryAllocProperty::kSingleAllocation) continue;
auto alignment = PackingAlignment(vec, e->scope);
if (!alignment) continue;
e->packing_alignment_bits = *alignment;
@@ -1071,7 +1087,7 @@ class StoragePlanRewriter : public StmtExprMutator {
}
// New allocation for merged data
void NewAllocMerged(StorageEntry* e) {
- TVM_FFI_ICHECK(e->scope.alloc_property() ==
MemoryAllocProperty::kSingleAllocation);
+ TVM_FFI_ICHECK(GetMemoryAllocProperty(e->scope) ==
MemoryAllocProperty::kSingleAllocation);
// allocate with element type.
TVM_FFI_ICHECK_NE(e->const_nbits, 0U);
uint64_t total_bits = e->const_nbits;
@@ -1199,7 +1215,7 @@ class StoragePlanRewriter : public StmtExprMutator {
StorageEntry* dst_entry = nullptr;
// inplace detection
if (detect_inplace && !RequiresExactAllocation(call) &&
- (scope.alloc_property() !=
MemoryAllocProperty::kSingleAllocation ||
+ (GetMemoryAllocProperty(scope) !=
MemoryAllocProperty::kSingleAllocation ||
alloc->var->ty.as_or_throw<TensorType>()->ConstantAllocationSize().value_or(0) >
0)) {
DLDataType dtype =
call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype;
@@ -1315,12 +1331,13 @@ class StoragePlanRewriter : public StmtExprMutator {
// disable reuse of small arrays, they will be lowered to registers in LLVM
// This rules only apply if we are using non special memory
- bool is_small_array = scope.alloc_property() ==
MemoryAllocProperty::kNormal && is_known_size &&
- const_nbits <= 32;
+ bool is_small_array = GetMemoryAllocProperty(scope) ==
MemoryAllocProperty::kNormal &&
+ is_known_size && const_nbits <= 32;
if (RequiresExactAllocation(call) || is_scalable_vector || !enable_reuse
|| is_small_array ||
!is_flat_memory_space ||
- (scope.alloc_property() == MemoryAllocProperty::kSingleAllocation &&
!is_known_size)) {
+ (GetMemoryAllocProperty(scope) ==
MemoryAllocProperty::kSingleAllocation &&
+ !is_known_size)) {
return NewAlloc(op, attach_scope, scope, const_nbits);
}
@@ -1381,10 +1398,11 @@ class StoragePlanRewriter : public StmtExprMutator {
const auto* call = e->allocs[0]->value.as<CallNode>();
PrimType
element_type(call->ty_args[0].as_or_throw<TensorType>()->dtype->dtype);
if (element_type.IsScalableVector()) return;
- if (e->scope.alloc_property() == MemoryAllocProperty::kNormal &&
e->const_nbits > 0 &&
+ if (GetMemoryAllocProperty(e->scope) == MemoryAllocProperty::kNormal &&
e->const_nbits > 0 &&
e->const_nbits <= 32)
return;
- if (e->scope.alloc_property() == MemoryAllocProperty::kSingleAllocation &&
e->const_nbits == 0)
+ if (GetMemoryAllocProperty(e->scope) ==
MemoryAllocProperty::kSingleAllocation &&
+ e->const_nbits == 0)
return;
// normal free.
if (e->const_nbits != 0) {
diff --git a/tests/cpp/runtime/memory_scope_test.cc
b/tests/cpp/runtime/memory_scope_test.cc
index 667f62aa4e..e62402560c 100644
--- a/tests/cpp/runtime/memory_scope_test.cc
+++ b/tests/cpp/runtime/memory_scope_test.cc
@@ -25,6 +25,16 @@
namespace tvm {
namespace runtime {
+namespace {
+std::optional<MemoryAllocProperty> AllocProperty(const MemoryScope& scope) {
+ auto state = ffi::reflection::TypeAttrColumn(
+
ffi::type_attr::kEnumState)[MemoryScopeObj::RuntimeTypeIndex()]
+ .cast<ffi::EnumState>();
+ auto column = state->attrs.Get(attr::kMemoryAllocProperty.name);
+ return attr::kMemoryAllocProperty.ValueFrom(column ?
column->Get(scope).value_or(ffi::Any())
+ : ffi::Any());
+}
+} // namespace
TEST(MemoryScope, CanonicalDefaultAndTypedBoundaries) {
MemoryScope global;
@@ -45,7 +55,7 @@ TEST(MemoryScope, ExtensionRegistrationAndAttributes) {
.set_attr("memory_scope_test.capacity", 1024);
auto scope = MemoryScope::Get("memory_scope_test.scratch");
EXPECT_EQ(scope.name(), "memory_scope_test.scratch");
- EXPECT_EQ(scope.alloc_property(), MemoryAllocProperty::kNormal);
+ EXPECT_FALSE(AllocProperty(scope).has_value());
EXPECT_EQ(scope, ffi::Any(scope).cast<MemoryScope>());
EXPECT_EQ(scope, PtrType(PrimType::Float(32), scope)->scope);
EXPECT_THROW(MemoryScope::Register(scope.name()), ffi::Error);
@@ -53,20 +63,44 @@ TEST(MemoryScope, ExtensionRegistrationAndAttributes) {
EXPECT_EQ(registered, MemoryScope::Get(registered.name()));
}
+TEST(MemoryScope, MembershipQueryPreservesRegistry) {
+ auto scope = MemoryScope::Register("memory_scope_test.membership");
+ auto state = ffi::reflection::TypeAttrColumn(
+
ffi::type_attr::kEnumState)[MemoryScopeObj::RuntimeTypeIndex()]
+ .cast<ffi::EnumState>();
+ auto num_entries = state->entries.size();
+ auto num_indexes = state->indexes.size();
+ auto num_attrs = state->attrs.size();
+ auto global = MemoryScope::Global();
+ auto property = AllocProperty(global);
+ EXPECT_TRUE(MemoryScope::IsDefined("global"));
+ EXPECT_TRUE(MemoryScope::IsDefined(scope.name()));
+ EXPECT_FALSE(MemoryScope::IsDefined("memory_scope_test.unknown"));
+ EXPECT_FALSE(MemoryScope::IsDefined(""));
+ EXPECT_FALSE(MemoryScope::IsDefined(ffi::String("global\0suffix", 13)));
+ EXPECT_EQ(state->entries.size(), num_entries);
+ EXPECT_EQ(state->indexes.size(), num_indexes);
+ EXPECT_EQ(state->attrs.size(), num_attrs);
+ EXPECT_TRUE(MemoryScope::Get(scope.name()).same_as(scope));
+ EXPECT_TRUE(MemoryScope::Global().same_as(global));
+ EXPECT_EQ(AllocProperty(global), property);
+ EXPECT_THROW(MemoryScope::Get("memory_scope_test.unknown"), ffi::Error);
+}
+
TEST(MemoryScope, AllocationProperties) {
- EXPECT_EQ(MemoryScope::Global().alloc_property(),
MemoryAllocProperty::kNormal);
- EXPECT_EQ(MemoryScope::Get("shared").alloc_property(),
MemoryAllocProperty::kNormal);
- EXPECT_EQ(MemoryScope::Get("shared.dyn").alloc_property(),
- MemoryAllocProperty::kSingleAllocation);
- EXPECT_EQ(MemoryScope::Get("local").alloc_property(),
MemoryAllocProperty::kSeparateAllocations);
+ EXPECT_EQ(AllocProperty(MemoryScope::Global()),
MemoryAllocProperty::kNormal);
+ EXPECT_EQ(AllocProperty(MemoryScope::Get("shared")),
MemoryAllocProperty::kNormal);
+ EXPECT_EQ(AllocProperty(MemoryScope::Get("shared.dyn")),
MemoryAllocProperty::kSingleAllocation);
+ EXPECT_EQ(AllocProperty(MemoryScope::Get("local")),
MemoryAllocProperty::kSeparateAllocations);
ffi::reflection::EnumDef<MemoryScopeObj>("memory_scope_test.single")
.set_attr(attr::kMemoryAllocProperty,
static_cast<int>(MemoryAllocProperty::kSingleAllocation));
- EXPECT_EQ(MemoryScope::Get("memory_scope_test.single").alloc_property(),
+ EXPECT_EQ(AllocProperty(MemoryScope::Get("memory_scope_test.single")),
MemoryAllocProperty::kSingleAllocation);
ffi::reflection::EnumDef<MemoryScopeObj>("memory_scope_test.invalid_property")
.set_attr(attr::kMemoryAllocProperty, 99);
-
EXPECT_THROW(MemoryScope::Get("memory_scope_test.invalid_property").alloc_property(),
ffi::Error);
+
EXPECT_EQ(AllocProperty(MemoryScope::Get("memory_scope_test.invalid_property")),
+ static_cast<MemoryAllocProperty>(99));
}
TEST(MemoryScope, GlobalTensorAllocation) {
diff --git a/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
b/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
index 035a2bb641..a1ee1b59c4 100644
--- a/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
+++ b/tests/python/tirx/transform/test_tirx_transform_storage_rewrite.py
@@ -30,7 +30,8 @@ for _scope_name in ("local.L0A", "local.L0A2"):
def _scope_with_property(value):
scope = tvm.runtime.MemoryScope.register("test.storage_rewrite." +
uuid.uuid4().hex)
- tvm.runtime.MemoryScope.def_attr("runtime.memory_alloc_property")[scope] =
value
+ if value != "absent":
+
tvm.runtime.MemoryScope.def_attr("runtime.memory_alloc_property")[scope] = value
return scope
@@ -42,7 +43,7 @@ def _rewrite(func):
return tvm.tirx.transform.StorageRewrite()(Before)["main"]
[email protected]("property_value", [0, 1, 2])
[email protected]("property_value", ["absent", None, False, True, 0, 1,
2, 99])
def test_allocation_property_controls_reuse(property_value):
scope = _scope_with_property(property_value)
@@ -70,8 +71,15 @@ def test_allocation_property_controls_reuse(property_value):
A_1[i_1] = T.Cast(i_1 + 1, "float32")
output[1] = A_1[127]
- expected = Expected["func"] if property_value != 2 else func
- tvm.ir.assert_structural_equal(_rewrite(func), expected)
+ if property_value is None:
+ with pytest.raises(TypeError):
+ _rewrite(func)
+ elif property_value == 99:
+ with pytest.raises(ValueError, match="Invalid memory allocation
property"):
+ _rewrite(func)
+ else:
+ expected = Expected["func"] if property_value != 2 else func
+ tvm.ir.assert_structural_equal(_rewrite(func), expected)
@pytest.mark.parametrize("property_value", [0, 1, 2])
