[TIR][USMP] adding the pass to convert to pool offsets

include/tvm/tir/usmp/utils.h

@@ -225,6 +225,44 @@ class PoolAllocation : public ObjectRef {
   TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolAllocation, ObjectRef, PoolAllocationNode);
 };
 
+/*!
+ * \brief This object contains information post-allocation for PoolInfo objects
+ */
+struct AllocatedPoolInfoNode : public Object {
+  /*! \brief The assigned PoolInfo object */
+  PoolInfo pool_info;
+  /*! \brief The allocated size into this pool */
+  Integer allocated_size;
+  /*! \brief An optional associated pool Var*/
+  Optional<Var> pool_var;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("pool_info", &pool_info);
+    v->Visit("allocated_size", &allocated_size);
+    v->Visit("pool_var", &pool_var);
+  }
+
+  bool SEqualReduce(const AllocatedPoolInfoNode* other, SEqualReducer equal) const {
+    return equal(pool_info, other->pool_info) && equal(allocated_size, other->allocated_size) &&
+           equal(pool_var, other->pool_var);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce(pool_info);
+    hash_reduce(allocated_size);
+    hash_reduce(pool_var);
+  }
+
+  static constexpr const char* _type_key = "tir.usmp.AllocatedPoolInfo";
+  TVM_DECLARE_FINAL_OBJECT_INFO(AllocatedPoolInfoNode, Object);
+};
+
+class AllocatedPoolInfo : public ObjectRef {
+ public:
+  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size, Var pool_var = Var());
+  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AllocatedPoolInfo, ObjectRef, AllocatedPoolInfoNode);
+};
+
 /*!
  * \brief Convert the IR-bound BufferInfo map to an array of BufferInfo
  *
@@ -248,6 +286,16 @@ Integer CalculateExtentsSize(const AllocateNode* op);
 
 }  // namespace usmp
 }  // namespace tir
+
+namespace attr {
+/*!
+ * \brief This is a BaseFunc attribute to indicate which input var represent
+ * a PoolInfo Object in the form of a Map<Var, PoolInfo>.
+ */
+static constexpr const char* kPoolArgs = "pool_args";
+
+}  // namespace attr
+
 }  // namespace tvm
 
 #endif  // TVM_TIR_USMP_UTILS_H_

python/tvm/script/tir/__init__.py

@@ -17,7 +17,7 @@
 """TVMScript for TIR"""
 
 # Type system
-from .ty import int8, int16, int32, int64, float16, float32, float64
+from .ty import uint8, int8, int16, int32, int64, float16, float32, float64
 from .ty import boolean, handle, Ptr, Tuple, Buffer
 
 from .prim_func import prim_func

python/tvm/script/tir/ty.py

@@ -137,6 +137,7 @@ def __getitem__(self, args):
         pass  # pylint: disable=unnecessary-pass
 
 
+uint8 = ConcreteType("uint8")
 int8 = ConcreteType("int8")
 int16 = ConcreteType("int16")
 int32 = ConcreteType("int32")

python/tvm/tir/usmp/__init__.py

@@ -18,4 +18,5 @@
 """Namespace for Unified Static Memory Planner"""
 
 from . import analysis
+from . import transform
 from .utils import BufferInfo

python/tvm/tir/usmp/transform/__init__.py

@@ -0,0 +1,20 @@
+# 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=unused-import, redefined-builtin
+"""Namespace for Unified Static Memory Planner"""
+
+from .transform import convert_pool_allocations_to_offsets

python/tvm/tir/usmp/transform/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.tir.usmp.analysis"""
+import tvm._ffi
+
+
+tvm._ffi._init_api("tir.usmp.transform", __name__)

python/tvm/tir/usmp/transform/transform.py

@@ -0,0 +1,46 @@
+# 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.
+"""USMP Transform Python API for passes"""
+# pylint: disable=invalid-name
+
+from typing import Dict
+
+import tvm
+from tvm.tir import Stmt
+from tvm.tir.usmp.utils import PoolAllocation
+from . import _ffi_api
+
+
+def convert_pool_allocations_to_offsets(
+    pool_allocations: Dict[Stmt, PoolAllocation], emit_tvmscript_printable: bool = False
+) -> tvm.transform.Pass:
+    """Convert pool allocations to Load nodes with offsets from pools.
+
+    Parameters
+    ----------
+    pool_allocations : Dict[Stmt, PoolAllocation]
+        Allocate or AllocateConst node to pool allocation mapping
+    emit_tvmscript_printable : bool
+        A toggle to emit TVMScript printable IRModule for unit tests
+        removing all attributes that should be attached for integration
+
+    Returns
+    -------
+    ret: tvm.transform.Pass
+        The registered pass that converts the allocations to offsets.
+    """
+    return _ffi_api.ConvertPoolAllocationsToOffsets(pool_allocations, emit_tvmscript_printable)

src/printer/text_printer.h

@@ -449,10 +449,11 @@ class TextPrinter {
 
   Doc PrintFinal(const ObjectRef& node) {
     Doc doc;
-    if (node->IsInstance<IRModuleNode>()) {
+    if (node.defined() && node->IsInstance<IRModuleNode>()) {
       doc << PrintMod(Downcast<IRModule>(node));
-    } else if (node->IsInstance<tir::PrimFuncNode>() || node->IsInstance<PrimExprNode>() ||
-               node->IsInstance<tir::StmtNode>()) {
+    } else if (node.defined() &&
+               (node->IsInstance<tir::PrimFuncNode>() || node->IsInstance<PrimExprNode>() ||
+                node->IsInstance<tir::StmtNode>())) {
       doc << tir_text_printer_.Print(node);
     } else {
       doc << relay_text_printer_.PrintFinal(node);

src/tir/ir/stmt.cc

@@ -35,7 +35,14 @@ namespace tir {
 LetStmt::LetStmt(Var var, PrimExpr value, Stmt body, Span span) {
   ICHECK(value.defined());
   ICHECK(body.defined());
-  ICHECK_EQ(value.dtype(), var.dtype());
+  auto vdtype = value.dtype();
+  // It is still valid to bind a pointer type
+  // var to a value that is of type handle.
+  if (var->type_annotation.as<PointerTypeNode>()) {
+    ICHECK(vdtype.is_handle());
+  } else {
+    ICHECK_EQ(value.dtype(), var.dtype());
+  }
 
   ObjectPtr<LetStmtNode> node = make_object<LetStmtNode>();
   node->var = std::move(var);