manupa-arm commented on code in PR #12029:
URL: https://github.com/apache/tvm/pull/12029#discussion_r916629542
##########
src/tir/contrib/ethosu/passes.cc:
##########
@@ -223,6 +224,572 @@ tvm::transform::Pass
CopyComputeReordering(Optional<Integer> max_copy_movements)
TVM_REGISTER_GLOBAL("tir.contrib.ethos-u.CopyComputeReordering")
.set_body_typed(CopyComputeReordering);
+/*!
+ * \brief This pass looks for the constants used by each compute operator
+ * and merges them into a single buffer.
+ * Constants written to a buffer with local scope are not merged.
+ */
+class MergeConstantsMutator : public StmtExprMutator {
+ public:
+ MergeConstantsMutator() {}
+
+ PrimFunc operator()(PrimFunc main_func, const Map<IntImm, runtime::NDArray>&
const_dict) {
+ // Analyze
+ Stmt new_body{this->VisitStmt(main_func->body)};
+
+ // Rewrite
+ analyze = false;
+ new_body = rewrite_prim_func_body(new_body);
+ std::set<ObjectRef> params_to_delete{};
+ auto new_buffer_map{make_new_buffer_map(main_func->buffer_map,
¶ms_to_delete)};
+ auto new_params{make_new_params(main_func->params, params_to_delete)};
+
+ // Make the new const dict
+ auto args_to_merge{get_args_to_merge(main_func->buffer_map,
main_func->params)};
+ auto buffers_to_merge{
+ get_args_to_merge_without_args_not_in_const_dict(args_to_merge,
const_dict)};
+ auto new_const_dict{make_new_const_dict(buffers_to_merge, const_dict)};
+
+ // Make the new prim func
+ auto prim_func_node{main_func.CopyOnWrite()};
+ prim_func_node->body = std::move(new_body);
+ prim_func_node->buffer_map = std::move(new_buffer_map);
+ prim_func_node->params = std::move(new_params);
+ prim_func_node->preflattened_buffer_map = {};
+ PrimFunc f{GetRef<PrimFunc>(prim_func_node)};
+
+ // Add the new const dict as an attribute
+ f = WithAttr(std::move(f), "ethos-u.const-dict", new_const_dict);
+
+ return f;
+ }
+
+ private:
+ /*! Indicates whether the pass is analyzing or rewriting */
+ bool analyze = true;
+
+ /*! A stack to store allocates as they are visited. */
+ std::vector<Allocate> allocates{};
+
+ /*! A list that contains in the i-th position the write buffer of the i-th
statement
+ * if that statement is a copy to a buffer with global scope */
+ std::vector<Optional<Buffer>> copy_write_buffers{};
+
+ /*! Maps a copy's write buffer to an index representing the
+ * new buffer and an offset in that buffer */
+ std::map<Buffer, std::pair<int /* new buffer index */, int /* offset */>>
+ old_to_new_write_buffer{};
+
+ /*! Maps an index representing a new buffer to the length of that buffer */
+ std::map<int /* new buffer index */, int /* length */> new_buffers_length{};
Review Comment:
Maps are ordered and therefore access time is not O(constant). Is there a
reason we need a ordering in the keys ?
If not I'd suggest we switch to unordered_maps.
--
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]
