[TIR][USMP] Integrating USMP to AoT Executor

include/tvm/tir/transform.h

@@ -484,6 +484,14 @@ TVM_DLL Pass MergeDynamicSharedMemoryAllocations();
  */
 TVM_DLL Pass ConvertForLoopsToSerial();
 
+/*!
+ * \brief This is the unified static memory planner pass that will
+ * plan for memory intra- and inter- PrimFuncs together. The pass
+ * requires all the function to be PrimFuncs including the main.
+ * \return The pass.
+ */
+TVM_DLL Pass UnifiedStaticMemoryPlanner();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

include/tvm/tir/usmp/algorithms.h

@@ -0,0 +1,62 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tir/usmp/algorithms.h
+ * \brief The memory planning algorithm for USMP
+ */
+
+#ifndef TVM_TIR_USMP_ALGORITHMS_H_
+#define TVM_TIR_USMP_ALGORITHMS_H_
+
+#include <tvm/tir/usmp/utils.h>
+
+namespace tvm {
+namespace tir {
+namespace usmp {
+namespace algo {
+
+/*!
+ * \brief The Greedy-by-Size algorithm to plan memory
+ *
+ * This will perform a greedy algorithm in deciding the offsets
+ * within provided Pools, using the size of the buffer.
+ *
+ * \return A Map of BufferInfo objects and their associated PoolAllocation
+ */
+Map<BufferInfo, PoolAllocation> GreedyBySize(const Array<BufferInfo>& buffer_info_arr,
+                                             const Integer& memory_pressure);
+
+/*!
+ * \brief The Greedy-by-Conflicts algorithm to plan memory
+ *
+ * This will perform a greedy algorithm in deciding the offsets
+ * within provided Pools, using the number of liveness conflicts of the buffer.
+ *
+ * \return A Map of BufferInfo objects and their associated PoolAllocation
+ */
+Map<BufferInfo, PoolAllocation> GreedyByConflicts(const Array<BufferInfo>& buffer_info_arr,
+                                                  const Integer& memory_pressure);
+
+}  // namespace algo
+}  // namespace usmp
+}  // namespace tir
+}  // namespace tvm
+
+#endif  // TVM_TIR_USMP_ALGORITHMS_H_

include/tvm/tir/usmp/analysis.h

@@ -0,0 +1,49 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tir/usmp/analysis.h
+ * \brief The analysis passes for TIR-based Unified Static Memory Planner
+ */
+
+#ifndef TVM_TIR_USMP_ANALYSIS_H_
+#define TVM_TIR_USMP_ANALYSIS_H_
+
+#include <tvm/tir/function.h>
+#include <tvm/tir/usmp/utils.h>
+
+namespace tvm {
+namespace tir {
+namespace usmp {
+
+/*!
+ * \brief Extract BufferInfo objects from a TIR IRModule
+ *
+ * This pass would extract the buffer information of allocate nodes
+ * including liveness conflict with other buffer info objects.
+ *
+ * \return A Map of BufferInfo objects and their associated Stmts
+ */
+BufferInfoAnalysis ExtractBufferInfo(const PrimFunc& main_func, const IRModule& mod);
+
+}  // namespace usmp
+}  // namespace tir
+}  // namespace tvm
+
+#endif  // TVM_TIR_USMP_ANALYSIS_H_

include/tvm/tir/usmp/transform.h

@@ -0,0 +1,64 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tir/usmp/transform.h
+ * \brief The transform passes for TIR-based Unified Static Memory Planner
+ */
+
+#ifndef TVM_TIR_USMP_TRANSFORM_H_
+#define TVM_TIR_USMP_TRANSFORM_H_
+
+#include <tvm/tir/usmp/utils.h>
+
+namespace tvm {
+namespace tir {
+namespace usmp {
+namespace transform {
+
+using Pass = tvm::transform::Pass;
+
+/*!
+ * \brief Convert the analyzed PoolAllocation to offsets from pool variables
+ *
+ * This pass would convert the main function to accept pool variables as an input
+ * that get passed onto the operator PrimFuncs. Furthermore, the static allocations
+ * will be converted to offsets within the pool variable.
+ *
+ * \return the pass
+ */
+TVM_DLL Pass ConvertPoolAllocationsToOffsets(const Map<tir::Stmt, PoolAllocation>& pool_allocations,
+                                             Bool emit_tvmscript_printable = Bool(false));
+
+/*!
+ * \brief Assign PoolInfo objects to tir.allocate nodes depending on the PrimFunc's target
+ *
+ * This pass would assign default PoolInfo objects to allocate nodes that are not otherwise
+ * annotated, depending on pool info supplied for each target.
+ *
+ * \return the pass
+ */
+TVM_DLL Pass AssignPoolInfo();
+
+}  // namespace transform
+}  // namespace usmp
+}  // namespace tir
+}  // namespace tvm
+
+#endif  // TVM_TIR_USMP_TRANSFORM_H_

include/tvm/tir/usmp/utils.h

@@ -26,10 +26,21 @@
 #define TVM_TIR_USMP_UTILS_H_
 
 #include <tvm/ir/expr.h>
+#include <tvm/runtime/device_api.h>
 #include <tvm/target/target.h>
 #include <tvm/tir/stmt.h>
 
 namespace tvm {
+
+/*!
+ * \brief PassContext option to enable the USMP
+ */
+constexpr const char* kUSMPEnableOption = "tir.usmp.enable";
+/*!
+ * \brief PassContext option to select the memory planning algorithm in USMP
+ */
+constexpr const char* kUSMPAlgorithmOption = "tir.usmp.algorithm";
+
 namespace tir {
 namespace usmp {
 
@@ -59,22 +70,28 @@ struct PoolInfoNode : public Object {
   Integer size_hint_bytes;
   /*! \brief The accessibility from each Target*/
   Map<Target, String> target_access;  // 'rw' or 'ro'
+  /*! \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("pool_name", &pool_name);
     v->Visit("size_hint_bytes", &size_hint_bytes);
     v->Visit("target_access", &target_access);
+    v->Visit("is_internal", &is_internal);
   }
 
   bool SEqualReduce(const PoolInfoNode* other, SEqualReducer equal) const {
     return equal(pool_name, other->pool_name) && equal(size_hint_bytes, other->size_hint_bytes) &&
-           equal(target_access, other->target_access);
+           equal(target_access, other->target_access) && equal(is_internal, other->is_internal);
   }
 
   void SHashReduce(SHashReducer hash_reduce) const {
     hash_reduce(pool_name);
     hash_reduce(size_hint_bytes);
     hash_reduce(target_access);
+    hash_reduce(is_internal);
   }
 
   static constexpr const char* _type_key = "tir.usmp.PoolInfo";
@@ -89,7 +106,8 @@ static const int kUnrestrictedPoolSizeHint = -1;
 class PoolInfo : public ObjectRef {
  public:
   TVM_DLL PoolInfo(String pool_name, Map<Target, String> target_access,
-                   Integer size_hint_bytes = kUnrestrictedPoolSizeHint);
+                   Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
+                   Bool is_internal = Bool(false));
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
 };
 
@@ -268,7 +286,14 @@ class AllocatedPoolInfo : public ObjectRef {
  *
  * \param buffer_info_map IR-bound BufferInfo map
  */
-Array<BufferInfo> CreateArrayBufferInfo(const Map<Stmt, BufferInfo>& buffer_info_map);
+Array<BufferInfo> CreateArrayBufferInfo(const Map<BufferInfo, Stmt>& buffer_info_map);
+
+/*!
+ * \brief Calculate workspace required to execute a IRModule with main expressed in TIR
+ *
+ * \param mod the IRModule with TIR-based main function
+ */
+Integer CalculateModuleWorkspaceSize(const IRModule& mod);
 
 /*!
  * \brief The allocate node attribute to indicate candidate memory pools.
@@ -284,6 +309,16 @@ static constexpr const char* kPoolCandidatesAllocateAttr = "candidate_memory_poo
  */
 Integer CalculateExtentsSize(const AllocateNode* op);
 
+/*!
+ * \brief Joins the Stmt nodes with PoolAllocation objects
+ *
+ * \param buffer_info_to_stmt the map of BufferInfo objects to Stmt nodes
+ * \param buffer_info_to_pool_allocation the map of BufferInfo objects to PoolAllocation objects
+ */
+Map<Stmt, PoolAllocation> AssignStmtPoolAllocations(
+    const Map<BufferInfo, Stmt>& buffer_info_to_stmt,
+    const Map<BufferInfo, PoolAllocation>& buffer_info_to_pool_allocation);
+
 }  // namespace usmp
 }  // namespace tir
 
@@ -294,6 +329,12 @@ namespace attr {
  */
 static constexpr const char* kPoolArgs = "pool_args";
 
+/*!
+ * \brief This is a IRModule attribute that contains all the PoolInfo objects
+ * as an Array.
+ */
+static constexpr const char* kPoolInfoIRModuleAttr = "pool_infos";
+
 }  // namespace attr
 
 }  // namespace tvm

src/driver/driver_api.cc

@@ -44,6 +44,7 @@ TVM_REGISTER_PASS_CONFIG_OPTION("tir.detect_global_barrier", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.instrument_bound_checkers", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_assert", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_vectorize", Bool);
+TVM_REGISTER_PASS_CONFIG_OPTION("tir.disable_storage_rewrite", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.is_entry_func", Bool);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.add_lower_pass", Array<Array<ObjectRef>>);
 TVM_REGISTER_PASS_CONFIG_OPTION("tir.debug_keep_trivial_loop", Bool);
@@ -191,6 +192,8 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   transform::PassContext pass_ctx = transform::PassContext::Current();
 
   bool disable_vectorize = pass_ctx->GetConfig<Bool>("tir.disable_vectorize", Bool(false)).value();
+  bool disable_storage_rewrite =
+      pass_ctx->GetConfig<Bool>("tir.disable_storage_rewrite", Bool(false)).value();
   bool instrument_bound_checkers =
       pass_ctx->GetConfig<Bool>("tir.instrument_bound_checkers", Bool(false)).value();
 
@@ -260,7 +263,9 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   pass_list.push_back(tir::transform::VectorizeLoop(!disable_vectorize));
   pass_list.push_back(tir::transform::InjectVirtualThread());
   pass_list.push_back(tir::transform::InjectDoubleBuffer());
-  pass_list.push_back(tir::transform::StorageRewrite());
+  if (!disable_storage_rewrite) {
+    pass_list.push_back(tir::transform::StorageRewrite());
+  }
   pass_list.push_back(tir::transform::UnrollLoop());
 
   // Add user-defined phase-2 passes