areusch commented on code in PR #10189:
URL: https://github.com/apache/tvm/pull/10189#discussion_r896959537
##########
include/tvm/ir/memory_pools.h:
##########
@@ -130,15 +129,159 @@ static const int kUnknownReadBandwidth = -1;
static const int kUnknownWriteBandwidth = -1;
class PoolInfo : public ObjectRef {
- public:
- TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
- Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
+ protected:
+ TVM_DLL PoolInfo(String pool_name, Integer size_hint_bytes =
kUnrestrictedPoolSizeHint,
Integer clock_frequency_hz = kUnknownClockFrequency,
Integer read_bandwidth_bytes_per_cycle =
kUnknownReadBandwidth,
Integer write_bandwidth_bytes_per_cycle =
kUnknownWriteBandwidth,
Integer read_latency_cycles = 0, Integer
write_latency_cycles = 0,
Map<Target, Integer> target_burst_bytes = {}, Bool
is_internal = Bool(false));
- TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+
+ public:
+ // TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+ TVM_DEFINE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+};
+
+/*!
+ * \brief Describes a pool of memory properties
+ */
+struct PoolInfoPropertiesNode : public Object {
+ /*! \brief The expected size hint to be used by the allocator.
+ * The size_hint_bytes is set to kUnrestrictedPoolSizeHint
+ * to indicate the pool is not size restricted.
+ */
+ Integer size_hint_bytes = kUnrestrictedPoolSizeHint;
+ /*! \brief The clock frequency of the memory in Hz */
+ Integer clock_frequency_hz = kUnknownClockFrequency;
+ /*! \brief The read bandwidth in bytes/cycle */
+ Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth;
+ /*! \brief The write bandwidth in bytes/cycle */
+ Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth;
+ /*! \brief The read latency in cycles */
+ Integer read_latency_cycles = 0;
+ /*! \brief The write latency in cycles */
+ Integer write_latency_cycles = 0;
+ /*! \brief The burst length in bytes for each Target */
+ Map<Target, Integer> target_burst_bytes{};
+ /*! \brief Whether pool is internally generated.
+ * The internal pools will be generated as part of
+ * the entry point code generation of the executor
+ */
+ bool is_internal = false;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("size_hint_bytes", &size_hint_bytes);
+ v->Visit("clock_frequency_hz", &clock_frequency_hz);
+ v->Visit("read_bandwidth_bytes_per_cycle",
&read_bandwidth_bytes_per_cycle);
+ v->Visit("write_bandwidth_bytes_per_cycle",
&write_bandwidth_bytes_per_cycle);
+ v->Visit("read_latency_cycles", &read_latency_cycles);
+ v->Visit("write_latency_cycles", &write_latency_cycles);
+ v->Visit("target_burst_bytes", &target_burst_bytes);
+ v->Visit("is_internal", &is_internal);
+ }
+
+ bool SEqualReduce(const PoolInfoPropertiesNode* other, SEqualReducer equal)
const {
+ return equal(size_hint_bytes, other->size_hint_bytes) &&
+ equal(clock_frequency_hz, other->clock_frequency_hz) &&
+ equal(read_bandwidth_bytes_per_cycle,
other->read_bandwidth_bytes_per_cycle) &&
+ equal(write_bandwidth_bytes_per_cycle,
other->write_bandwidth_bytes_per_cycle) &&
+ equal(read_latency_cycles, other->read_latency_cycles) &&
+ equal(write_latency_cycles, other->write_latency_cycles) &&
+ equal(target_burst_bytes, other->target_burst_bytes) &&
+ equal(is_internal, other->is_internal);
+ }
+
+ void SHashReduce(SHashReducer hash_reduce) const {
+ hash_reduce(size_hint_bytes);
+ hash_reduce(clock_frequency_hz);
+ hash_reduce(read_bandwidth_bytes_per_cycle);
+ hash_reduce(write_bandwidth_bytes_per_cycle);
+ hash_reduce(read_latency_cycles);
+ hash_reduce(write_latency_cycles);
+ hash_reduce(target_burst_bytes);
+ hash_reduce(is_internal);
+ }
+
+ static constexpr const char* _type_key = "ir.PoolInfoProperties";
+ TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoPropertiesNode, Object);
+};
+
+class PoolInfoProperties : public ObjectRef {
+ public:
+ TVM_DLL PoolInfoProperties(Integer size_hint_bytes,
+ Integer clock_frequency_hz =
kUnknownClockFrequency,
+ Integer read_bandwidth_bytes_per_cycle =
kUnknownReadBandwidth,
+ Integer write_bandwidth_bytes_per_cycle =
kUnknownWriteBandwidth,
+ Integer read_latency_cycles = 0, Integer
write_latency_cycles = 0,
+ Map<Target, Integer> target_burst_bytes = {},
+ Bool is_internal = Bool(false));
+ TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfoProperties, ObjectRef,
PoolInfoPropertiesNode);
+};
+
+struct WorkspacePoolInfoNode : public PoolInfoNode {
+ static constexpr const char* _type_key = "ir.WorkspacePoolInfo";
+ TVM_DECLARE_FINAL_OBJECT_INFO(WorkspacePoolInfoNode, PoolInfoNode);
+};
+
+class WorkspacePoolInfo : public PoolInfo {
Review Comment:
does this class exist only to provide a default for `properties` ctor arg?
if so, consider a static function instead?
##########
include/tvm/ir/memory_pools.h:
##########
@@ -130,15 +129,159 @@ static const int kUnknownReadBandwidth = -1;
static const int kUnknownWriteBandwidth = -1;
class PoolInfo : public ObjectRef {
- public:
- TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
- Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
+ protected:
+ TVM_DLL PoolInfo(String pool_name, Integer size_hint_bytes =
kUnrestrictedPoolSizeHint,
Integer clock_frequency_hz = kUnknownClockFrequency,
Integer read_bandwidth_bytes_per_cycle =
kUnknownReadBandwidth,
Integer write_bandwidth_bytes_per_cycle =
kUnknownWriteBandwidth,
Integer read_latency_cycles = 0, Integer
write_latency_cycles = 0,
Map<Target, Integer> target_burst_bytes = {}, Bool
is_internal = Bool(false));
- TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+
+ public:
+ // TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+ TVM_DEFINE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+};
+
+/*!
+ * \brief Describes a pool of memory properties
+ */
+struct PoolInfoPropertiesNode : public Object {
+ /*! \brief The expected size hint to be used by the allocator.
+ * The size_hint_bytes is set to kUnrestrictedPoolSizeHint
+ * to indicate the pool is not size restricted.
+ */
+ Integer size_hint_bytes = kUnrestrictedPoolSizeHint;
+ /*! \brief The clock frequency of the memory in Hz */
+ Integer clock_frequency_hz = kUnknownClockFrequency;
+ /*! \brief The read bandwidth in bytes/cycle */
+ Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth;
+ /*! \brief The write bandwidth in bytes/cycle */
+ Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth;
+ /*! \brief The read latency in cycles */
+ Integer read_latency_cycles = 0;
+ /*! \brief The write latency in cycles */
+ Integer write_latency_cycles = 0;
+ /*! \brief The burst length in bytes for each Target */
+ Map<Target, Integer> target_burst_bytes{};
+ /*! \brief Whether pool is internally generated.
+ * The internal pools will be generated as part of
+ * the entry point code generation of the executor
+ */
+ bool is_internal = false;
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("size_hint_bytes", &size_hint_bytes);
+ v->Visit("clock_frequency_hz", &clock_frequency_hz);
+ v->Visit("read_bandwidth_bytes_per_cycle",
&read_bandwidth_bytes_per_cycle);
+ v->Visit("write_bandwidth_bytes_per_cycle",
&write_bandwidth_bytes_per_cycle);
+ v->Visit("read_latency_cycles", &read_latency_cycles);
+ v->Visit("write_latency_cycles", &write_latency_cycles);
+ v->Visit("target_burst_bytes", &target_burst_bytes);
+ v->Visit("is_internal", &is_internal);
+ }
+
+ bool SEqualReduce(const PoolInfoPropertiesNode* other, SEqualReducer equal)
const {
+ return equal(size_hint_bytes, other->size_hint_bytes) &&
+ equal(clock_frequency_hz, other->clock_frequency_hz) &&
+ equal(read_bandwidth_bytes_per_cycle,
other->read_bandwidth_bytes_per_cycle) &&
+ equal(write_bandwidth_bytes_per_cycle,
other->write_bandwidth_bytes_per_cycle) &&
+ equal(read_latency_cycles, other->read_latency_cycles) &&
+ equal(write_latency_cycles, other->write_latency_cycles) &&
+ equal(target_burst_bytes, other->target_burst_bytes) &&
+ equal(is_internal, other->is_internal);
+ }
+
+ void SHashReduce(SHashReducer hash_reduce) const {
+ hash_reduce(size_hint_bytes);
+ hash_reduce(clock_frequency_hz);
+ hash_reduce(read_bandwidth_bytes_per_cycle);
+ hash_reduce(write_bandwidth_bytes_per_cycle);
+ hash_reduce(read_latency_cycles);
+ hash_reduce(write_latency_cycles);
+ hash_reduce(target_burst_bytes);
+ hash_reduce(is_internal);
+ }
+
+ static constexpr const char* _type_key = "ir.PoolInfoProperties";
+ TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoPropertiesNode, Object);
+};
+
+class PoolInfoProperties : public ObjectRef {
+ public:
+ TVM_DLL PoolInfoProperties(Integer size_hint_bytes,
+ Integer clock_frequency_hz =
kUnknownClockFrequency,
+ Integer read_bandwidth_bytes_per_cycle =
kUnknownReadBandwidth,
+ Integer write_bandwidth_bytes_per_cycle =
kUnknownWriteBandwidth,
+ Integer read_latency_cycles = 0, Integer
write_latency_cycles = 0,
+ Map<Target, Integer> target_burst_bytes = {},
+ Bool is_internal = Bool(false));
+ TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfoProperties, ObjectRef,
PoolInfoPropertiesNode);
+};
+
+struct WorkspacePoolInfoNode : public PoolInfoNode {
+ static constexpr const char* _type_key = "ir.WorkspacePoolInfo";
Review Comment:
be consistent in the prefix--either tir.usmp or ir. since it is in tvm/ir,
probably ir.
##########
python/tvm/relay/build_module.py:
##########
@@ -631,7 +652,7 @@ def _make_executor(self, expr=None):
# generated code.
temp_so_dir = contrib_utils.TempDirectory()
temp_so = temp_so_dir / "temp.so"
- mod.export_library(temp_so, cc="gcc", options=["-std=c11"])
+ mod.export_library(temp_so, cc="c++", options=["-std=gnu++14"])
Review Comment:
how come it's necessary to invoke the c++ compiler?
##########
src/runtime/aot_executor/aot_executor.cc:
##########
@@ -62,9 +64,38 @@ AotExecutor::AotExecutor(tvm::runtime::Module module, const
std::vector<Device>&
output->dtype(), devices_[0]));
}
- for (auto pool : metadata_->pools()) {
- args_.emplace_back(NDArray::Empty(ShapeTuple(pool->shape().begin(),
pool->shape().end()),
- pool->dtype(), devices_[0]));
+ auto get_array_byte_size = [](DataType data_type, auto shape) {
+ int64_t nbytes = (data_type.bits() * data_type.lanes() + 7) / 8;
Review Comment:
could you reuse one of the existing functions we have for computing array
extent size?
##########
python/tvm/ir/memory_pools.py:
##########
@@ -112,15 +112,155 @@ def __init__(
)
+@register_object("ir.PoolInfoProperties")
+class PoolInfoProperties(Object):
+ """PoolInfo object holds information related to memory pools
+ where the statically sized allocate nodes will pooled into.
+
+ Parameters
+ ----------
+ size_hint_bytes : Optional[int]
+ The expected size hint to be used by the allocator.
+ The default value would be -1 which means the pool
+ is not size restricted.
+
+ clock_frequency_hz : Optional[int]
+ The clock frequency that the memory pool runs at in Hz.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ read_bandwidth_bytes_per_cycle : Optional[int]
+ The read bandwidth of the memory pool in bytes/cycle.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ write_bandwidth_bytes_per_cycle : Optional[int]
+ The write bandwidth of the memory pool in bytes/cycle.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ read_latency_cycles : Optional[int]
+ The read latency of the memory pool in cycles.
+ If not specified/known, this will default to 0.
+
+ write_latency_cycles : Optional[int]
+ The write latency of the memory pool in cycles.
+ If not specified/known, this will default to 0.
+
+ target_burst_bytes : Optional[Union[Dict[Target, int], None]]
+ The burst length of the memory pool in bytes per target.
+ If not specified/known for a given target, a burst length
+ of 1 byte will be assumed.
+
+ """
+
+ def __init__(
+ self,
+ size_hint_bytes: Optional[int] = -1,
+ clock_frequency_hz: Optional[int] = -1,
+ read_bandwidth_bytes_per_cycle: Optional[int] = -1,
+ write_bandwidth_bytes_per_cycle: Optional[int] = -1,
+ read_latency_cycles: Optional[int] = 0,
+ write_latency_cycles: Optional[int] = 0,
+ target_burst_bytes=None, # Optional[Union[Dict[target.Target, int],
None]]
+ ):
+ if not target_burst_bytes:
+ target_burst_bytes = dict()
+
+ self.__init_handle_by_constructor__(
+ _ffi_api.PoolInfoProperties, # type: ignore # pylint:
disable=no-member
+ size_hint_bytes,
+ clock_frequency_hz,
+ read_bandwidth_bytes_per_cycle,
+ write_bandwidth_bytes_per_cycle,
+ read_latency_cycles,
+ write_latency_cycles,
+ target_burst_bytes,
+ )
+
+
+@register_object("ir.WorkspacePoolInfo")
+class WorkspacePoolInfo(PoolInfo):
+ """WorkspacePoolInfo object holds information related to RW memory pools
+ where the statically sized allocate nodes will pooled into.
+
+ Parameters
+ ----------
+ pool_name : str
+ The name of the memory pool
+
+ targets : list[Target]
+ A list of targets which could access the pool
+
+ pool_info_properties : PoolInfoProperties
+ The properties of the pool.
+ """
+
+ # pylint: disable=W0231
+ def __init__(
+ self,
+ pool_name: str,
+ targets, # list[Target]
Review Comment:
remove the commented-out code
##########
include/tvm/ir/memory_pools.h:
##########
@@ -130,15 +129,159 @@ static const int kUnknownReadBandwidth = -1;
static const int kUnknownWriteBandwidth = -1;
class PoolInfo : public ObjectRef {
Review Comment:
could you add/update docstrings for each of these classes?
##########
include/tvm/ir/memory_pools.h:
##########
@@ -130,15 +129,159 @@ static const int kUnknownReadBandwidth = -1;
static const int kUnknownWriteBandwidth = -1;
class PoolInfo : public ObjectRef {
- public:
- TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
- Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
+ protected:
+ TVM_DLL PoolInfo(String pool_name, Integer size_hint_bytes =
kUnrestrictedPoolSizeHint,
Integer clock_frequency_hz = kUnknownClockFrequency,
Integer read_bandwidth_bytes_per_cycle =
kUnknownReadBandwidth,
Integer write_bandwidth_bytes_per_cycle =
kUnknownWriteBandwidth,
Integer read_latency_cycles = 0, Integer
write_latency_cycles = 0,
Map<Target, Integer> target_burst_bytes = {}, Bool
is_internal = Bool(false));
- TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
+
+ public:
+ // TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
Review Comment:
remove
##########
tests/python/relay/aot/test_cpp_aot.py:
##########
@@ -92,31 +92,35 @@ def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight :
Tensor[(3, 3, 5, 5),
shape_dict = {p.name_hint: p.checked_type.concrete_shape for p in
main_func.params}
type_dict = {p.name_hint: p.checked_type.dtype for p in main_func.params}
- weight_data = np.ones(shape_dict["weight"]).astype(type_dict["weight"])
+ weight_data = np.random.randint(1, 255,
shape_dict["weight"]).astype(type_dict["weight"])
input_data = np.ones(shape_dict["data"]).astype(type_dict["data"])
-
params = {"weight": weight_data}
inputs = {"data": input_data}
ref_outputs = generate_ref_data(ir_mod, inputs, params)
with tvm.transform.PassContext(
- opt_level=3, config={"tir.disable_vectorize": True, "tir.usmp.enable":
enable_usmp}
+ opt_level=3,
+ config={
+ "tir.disable_vectorize": True,
+ "tir.usmp.enable": enable_usmp,
+ }, # , "tir.usmp.algorithm": "hill_climb"}
):
mod = tvm.relay.build(
ir_mod,
params=params,
- target=target_kind,
- executor=backend.Executor("aot", {"interface-api": "packed"}),
+ target=target_kind, # + " -constants-byte-alignment=8
-workspace-byte-alignment=8",
+ executor=backend.Executor(
+ "aot", {"interface-api": "packed", "unpacked-api": False}
+ ), # , "constant-byte-alignment":8,
"workspace-byte-alignment":8}),
)
-
temp_dir = tvm.contrib.utils.TempDirectory()
test_so_path = temp_dir / "test.so"
- mod.export_library(test_so_path, cc="gcc", options=["-std=c11"])
+ mod.export_library(test_so_path, cc="c++", options=["-std=gnu++14", "-g3",
"-O0"])
Review Comment:
same question here
##########
include/tvm/tir/usmp/utils.h:
##########
@@ -163,7 +162,7 @@ class BufferInfoAnalysis : public ObjectRef {
* \brief The pool allocation produced after the USMP algorithm
*/
struct PoolAllocationNode : public Object {
- /*! \brief The assigned PoolInfo object */
+ /*! \brief The assigned WorkspacePoolInfo object */
Review Comment:
comment doesn't match the type of the attr
##########
src/target/source/source_module.cc:
##########
@@ -283,6 +291,55 @@ class CSourceCrtMetadataModuleNode : public
runtime::ModuleNode {
code_ << "}\n\n";
}
+ void GenerateConstantBuffer(const ConstantPoolInfoNode* pool_info, size_t
allocated_size) {
+ size_t offset = 0;
+ if (pool_info->constant_info_array.size() > 0) {
+ // Pool is RO, form an initialized struct
+ code_ << "__attribute__((section(\".rodata.tvm\"), ";
+ code_ << "))\n";
+ code_ << "static struct " << pool_info->pool_name << " {\n";
+ // emit struct field names
+ std::vector<ConstantInfo>
const_info_vec(pool_info->constant_info_array.begin(),
+
pool_info->constant_info_array.end());
+ std::sort(const_info_vec.begin(), const_info_vec.end(),
+ [](const ConstantInfo& a, const ConstantInfo& b) {
+ return a->byte_offset->value < b->byte_offset->value;
+ });
+ for (const auto& const_info : const_info_vec) {
+ const auto& data = const_info->data;
+ const auto& offs = const_info->byte_offset;
+ int64_t num_elements = std::accumulate(data.Shape().begin(),
data.Shape().end(), 1,
+ std::multiplies<int64_t>());
+ code_ << " ";
+ codegen_c_base_.PrintType(data.DataType(), code_);
+ code_ << " " << const_info->name_hint << "[" << num_elements
+ << "] __attribute__((packed, aligned(" <<
metadata_->constant_alignment << ")));";
+ code_ << " // " << num_elements * data.DataType().bytes()
+ << " bytes, aligned offset: " << offs << "\n";
+ }
+ code_ << "} " << pool_info->pool_name << " = {\n";
+
+ // emit struct field initialization data
+ for (const auto& const_info : const_info_vec) {
+ code_ << " ." << const_info->name_hint << " = {\n";
+ codegen::NDArrayDataToC(const_info->data, 4, code_);
+ code_ << " },\n";
+ }
+ code_ << "};";
+ code_ << "// of total size " << allocated_size << " bytes, aligned: " <<
offset << " bytes\n";
+ } else {
+ LOG(FATAL) << "No constant data in constant pool found "
+ << PrettyPrint(GetRef<ObjectRef>(pool_info));
+ }
+ }
+
+ void GenerateWorkspaceBuffer(const WorkspacePoolInfoNode* pool_info, size_t
allocated_size) {
+ code_ << "__attribute__((section(\".bss.noinit.tvm\"), ";
Review Comment:
just musing here--i wonder if the section should become a property of the
PoolInfoNode, for this and constants
##########
python/tvm/ir/memory_pools.py:
##########
@@ -112,15 +112,155 @@ def __init__(
)
+@register_object("ir.PoolInfoProperties")
+class PoolInfoProperties(Object):
+ """PoolInfo object holds information related to memory pools
+ where the statically sized allocate nodes will pooled into.
+
+ Parameters
+ ----------
+ size_hint_bytes : Optional[int]
+ The expected size hint to be used by the allocator.
+ The default value would be -1 which means the pool
+ is not size restricted.
+
+ clock_frequency_hz : Optional[int]
+ The clock frequency that the memory pool runs at in Hz.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ read_bandwidth_bytes_per_cycle : Optional[int]
+ The read bandwidth of the memory pool in bytes/cycle.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ write_bandwidth_bytes_per_cycle : Optional[int]
+ The write bandwidth of the memory pool in bytes/cycle.
+ If not specified/known, this will default to -1 indicating
+ it hasn't been defined.
+
+ read_latency_cycles : Optional[int]
+ The read latency of the memory pool in cycles.
+ If not specified/known, this will default to 0.
+
+ write_latency_cycles : Optional[int]
+ The write latency of the memory pool in cycles.
+ If not specified/known, this will default to 0.
+
+ target_burst_bytes : Optional[Union[Dict[Target, int], None]]
+ The burst length of the memory pool in bytes per target.
+ If not specified/known for a given target, a burst length
+ of 1 byte will be assumed.
+
+ """
+
+ def __init__(
+ self,
+ size_hint_bytes: Optional[int] = -1,
+ clock_frequency_hz: Optional[int] = -1,
+ read_bandwidth_bytes_per_cycle: Optional[int] = -1,
+ write_bandwidth_bytes_per_cycle: Optional[int] = -1,
+ read_latency_cycles: Optional[int] = 0,
+ write_latency_cycles: Optional[int] = 0,
+ target_burst_bytes=None, # Optional[Union[Dict[target.Target, int],
None]]
Review Comment:
remove the commented out code
##########
src/tir/usmp/analysis/extract_buffer_info.cc:
##########
@@ -505,6 +572,41 @@ BufferInfoAnalysis BufferInfoExtractor::operator()(const
PrimFunc& main_func) {
open_set.erase(le_event.buffer_info);
}
}
+
+ // All ConstantPoolInfo items should have conflicts with each other
+ // as they will be placed in RO segment and pre-initialized
+
+ // split buffers to vars (WorkspacePoolInfo items) and constants
(ConstantPoolInfo items)
+ Array<BufferInfo> buffer_info_vars;
+ Array<BufferInfo> buffer_info_constants;
+ for (const auto& kv : this->buffer_info_map_) {
+ const auto& stmt = kv.second;
+ if (stmt->IsInstance<AllocateConstNode>()) {
+ buffer_info_constants.push_back(kv.first);
+ } else {
+ buffer_info_vars.push_back(kv.first);
+ }
+ }
+ ICHECK(buffer_info_map_.size() == buffer_info_vars.size() +
buffer_info_constants.size())
+ << "missing value";
+
+ Map<ObjectRef, ObjectRef> srch;
+ // intersect with each other, as all constants should exist at the same time
Review Comment:
did this comment get split from previous one?
##########
src/ir/memory_pools.cc:
##########
@@ -89,4 +223,21 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
<< "pools=" << node->pools << ")";
});
+ConstantMemoryPools::ConstantMemoryPools(Array<ConstantPoolInfo> pools) {
+ auto constant_memory_pools_node = make_object<ConstantMemoryPoolsNode>();
+ constant_memory_pools_node->pools = pools;
+ data_ = std::move(constant_memory_pools_node);
+}
+
+TVM_REGISTER_NODE_TYPE(ConstantMemoryPoolsNode);
+TVM_REGISTER_GLOBAL("ir.ConstantMemoryPools").set_body_typed([](Array<ConstantPoolInfo>
pools) {
+ return ConstantMemoryPools(pools);
+});
+
+TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
+ .set_dispatch<ConstantMemoryPoolsNode>([](const ObjectRef& ref,
ReprPrinter* p) {
+ auto* node = static_cast<const ConstantMemoryPoolsNode*>(ref.get());
+ p->stream << "ConstnatnMemoryPoolsNode(\n"
Review Comment:
nit: spelling
--
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: commits-unsubscr...@tvm.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
