[TIR] tir.const extraction

This commit tries to implement an amendment to tir.constant RFC
with centralized storage of constant data within the IRModule
Please note that data and irmod_storage_idx are not mutual exclisive
further more the irmod_storage_idx is valid only immediatly after
prim func addition to the mod or after update within the mod.
If prim func is out of the the module scope then the index become
meangless. irmod_storage_idx also is not used in calculation of hash
function of the tir.constant node.

Change-Id: I40742ed580468b0252ea3fec02184cba65e20871

Author: d-smirnov

include/tvm/ir/module.h

@@ -32,6 +32,7 @@
 #include <tvm/runtime/container/array.h>
 #include <tvm/runtime/container/map.h>
 #include <tvm/runtime/container/string.h>
+#include <tvm/tir/function.h>
 
 #include <string>
 #include <unordered_map>
@@ -343,6 +344,9 @@ class IRModuleNode : public Object {
    */
   std::unordered_set<String> import_set_;
   friend class IRModule;
+
+ public:
+  void ExtractPrimFuncConstants(tir::PrimFunc func);
 };
 
 /*!
@@ -351,6 +355,8 @@ class IRModuleNode : public Object {
  */
 class IRModule : public ObjectRef {
  public:
+  static constexpr const char* _constants_attrs_key = "Constants";
+
   /*!
    * \brief constructor
    * \param functions Functions in the module.

include/tvm/tir/stmt.h

@@ -628,8 +628,14 @@ class AllocateConstNode : public StmtNode {
  public:
   /*! \brief The buffer variable. */
   Var buffer_var;
-  /*! \brief The data associated to the constant. */
-  ::tvm::runtime::NDArray data;
+  /*! \brief The optional data associated to the constant.
+   */
+  Optional<runtime::NDArray> data;
+  /*! \brief If the PrimFunc containing the Stmt is added to IRModule,
+       this is an optional index to indicate the index within
+       "Constants" attribute, that is a Array<NDArray> of IRModule.
+   */
+  Optional<Integer> irmod_storage_idx;
   /*! \brief The type of the buffer. */
   DataType dtype;
   /*! \brief The extents of the buffer. */
@@ -677,14 +683,17 @@ class AllocateConstNode : public StmtNode {
 };
 
 /*!
- * \brief Managed reference to AllocateNode.
- * \sa AllocateNode
+ * \brief Managed reference to AllocateConstNode.
+ * \sa AllocateConstNode
  */
 class AllocateConst : public Stmt {
  public:
-  TVM_DLL AllocateConst(Var buffer_var, ::tvm::runtime::NDArray data, DataType dtype,
-                        Array<PrimExpr> extents, Stmt body, Span span = Span());
-
+  /* The constructor to create a IRNode with constant data
+   * depending on the type of ObjectRef, it will either
+   * create AllocateConstNode with irmod_storage_idx or data
+   */
+  TVM_DLL AllocateConst(Var buffer_var, DataType dtype, Array<PrimExpr> extents,
+                        ObjectRef data_or_idx, Stmt body, Span span = Span());
   TVM_DEFINE_OBJECT_REF_METHODS(AllocateConst, Stmt, AllocateConstNode);
 };
 

python/tvm/script/tir/scope_handler.py

@@ -18,12 +18,12 @@
 # pylint: disable=redefined-builtin, unused-argument, invalid-name, relative-beyond-top-level
 from typing import Tuple, Any, Callable, Optional, List, Union, Mapping
 
+import numpy as np
 import synr
 import tvm.tir
 from tvm.runtime import Object
 from tvm.ir import Span, Range
 from tvm.tir import Stmt, PrimExpr, IterVar, Var, Buffer, BufferRegion, ForKind
-import numpy as np
 
 from .node import BufferSlice
 from .utils import buffer_slice_to_region
@@ -159,16 +159,15 @@ def setup_buffer_var(
 
 @register
 class AllocateConst(WithScopeHandler):
-    """With scope handler tir.allocate(data, extents, dtype, condition)"""
+    """With scope handler tir.allocate_const(data, extents, dtype, condition)"""
 
     def __init__(self):
         def allocate_const(raw_data, dtype, shape, span=None):
             list_data = []
             for i in raw_data:
                 list_data.append(i.value)
             nd_data = tvm.nd.array(np.asarray(list_data, dtype=dtype))
-
-            n = tvm.tir.AllocateConst(self.buffer_var, nd_data, dtype, shape, self.body, span=span)
+            n = tvm.tir.AllocateConst(self.buffer_var, dtype, shape, nd_data, self.body, span=span)
             return n
 
         super().__init__(allocate_const, concise_scope=True, def_symbol=True)
@@ -182,15 +181,17 @@ def enter_scope(
         span: synr.ast.Span,
     ):
         # define buffer vars in symbol table
-        if isinstance(node, ast.With):
+        if isinstance(node, synr.ast.With):
             vars = WithScopeHandler.get_optional_vars(node, context)
             if len(vars) != 1:
-                context.report_error("Unexpected number of vars", node.span)
+                context.report_error(f"Unexpected number of vars: 1 vs. {len(vars)}", node.span)
             name = vars[0].id.name
             var_span = vars[0].id.span
-        elif isinstance(node, ast.Assign):
-            name = node.lhs.id.name
-            var_span = node.lhs.id.span
+        elif isinstance(node, synr.ast.Assign):
+            if len(node.lhs) != 1:
+                context.report_error(f"Unexpected number of vars: 1 vs. {len(node.lhs)}", node.span)
+            name = node.lhs[0].id.name
+            var_span = node.lhs[0].id.span
         else:
             raise Exception("Internal Bug")
 
@@ -214,11 +215,7 @@ def launch_thread(env_var, extent, span):
             attr_key = "virtual_thread" if thread_id == "vthread" else "thread_extent"
             return tvm.tir.AttrStmt(
                 IterVar(
-                    (0, extent),
-                    env_var,
-                    getattr(IterVar, "ThreadIndex"),
-                    thread_id,
-                    span=span,
+                    (0, extent), env_var, getattr(IterVar, "ThreadIndex"), thread_id, span=span,
                 ),
                 attr_key,
                 extent,
@@ -545,9 +542,7 @@ class Serial(ForScopeHandler):
 
     def __init__(self):
         def serial(
-            begin: PrimExpr,
-            end: PrimExpr,
-            annotations: Optional[Mapping[str, Object]] = None,
+            begin: PrimExpr, end: PrimExpr, annotations: Optional[Mapping[str, Object]] = None,
         ):
             self.create_loop_info(begin, end, ForKind.SERIAL, annotations=annotations)
 
@@ -560,9 +555,7 @@ class Parallel(ForScopeHandler):
 
     def __init__(self):
         def parallel(
-            begin: PrimExpr,
-            end: PrimExpr,
-            annotations: Optional[Mapping[str, Object]] = None,
+            begin: PrimExpr, end: PrimExpr, annotations: Optional[Mapping[str, Object]] = None,
         ):
             self.create_loop_info(begin, end, ForKind.PARALLEL, annotations=annotations)
 
@@ -575,9 +568,7 @@ class Vectorized(ForScopeHandler):
 
     def __init__(self):
         def vectorized(
-            begin: PrimExpr,
-            end: PrimExpr,
-            annotations: Optional[Mapping[str, Object]] = None,
+            begin: PrimExpr, end: PrimExpr, annotations: Optional[Mapping[str, Object]] = None,
         ):
             self.create_loop_info(begin, end, ForKind.VECTORIZED, annotations=annotations)
 
@@ -590,9 +581,7 @@ class Unroll(ForScopeHandler):
 
     def __init__(self):
         def unroll(
-            begin: PrimExpr,
-            end: PrimExpr,
-            annotations: Optional[Mapping[str, Object]] = None,
+            begin: PrimExpr, end: PrimExpr, annotations: Optional[Mapping[str, Object]] = None,
         ):
             self.create_loop_info(begin, end, ForKind.UNROLLED, annotations=annotations)
 

src/ir/module.cc

@@ -34,19 +34,20 @@
 #include <tvm/relay/analysis.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/transform.h>
+#include <tvm/tir/stmt.h>
+#include <tvm/tir/stmt_functor.h>
 
 #include <fstream>
 #include <sstream>
 #include <unordered_set>
 
 namespace tvm {
-
+constexpr const char* IRModule::_constants_attrs_key;
 IRModule::IRModule(tvm::Map<GlobalVar, BaseFunc> functions,
                    tvm::Map<GlobalTypeVar, TypeData> type_definitions,
                    std::unordered_set<String> import_set, parser::SourceMap source_map,
                    DictAttrs attrs) {
   auto n = make_object<IRModuleNode>();
-  n->functions = std::move(functions);
   n->type_definitions = std::move(type_definitions);
   n->global_type_var_map_ = {};
   n->global_var_map_ = {};
@@ -55,11 +56,10 @@ IRModule::IRModule(tvm::Map<GlobalVar, BaseFunc> functions,
   n->source_map = source_map;
   n->attrs = std::move(attrs);
 
-  for (const auto& kv : n->functions) {
-    // set global var map
-    ICHECK(n->global_var_map_.count(kv.first->name_hint) == 0)
-        << "Duplicate global function name " << kv.first->name_hint;
-    n->global_var_map_.Set(kv.first->name_hint, kv.first);
+  if (functions.defined()) {
+    for (const auto& kv : functions) {
+      n->Add(kv.first, kv.second);
+    }
   }
 
   for (const auto& kv : n->type_definitions) {
@@ -201,6 +201,10 @@ void IRModuleNode::Add(const GlobalVar& var, const BaseFunc& f, bool update) {
     WarnIfMalformed(GetRef<IRModule>(this), GetRef<relay::Function>(ptr));
   }
 
+  if (f->IsInstance<tir::PrimFuncNode>()) {
+    ExtractPrimFuncConstants(Downcast<tir::PrimFunc>(f));
+  }
+
   AddUnchecked(var, checked_func);
 }
 
@@ -218,6 +222,63 @@ void IRModuleNode::AddUnchecked(const GlobalVar& var, const BaseFunc& func) {
   global_var_map_.Set(var->name_hint, var);
 }
 
+// Replaces constant data to index into mod's "Constants" attrs array.
+void IRModuleNode::ExtractPrimFuncConstants(tir::PrimFunc func) {
+  using ConstArrayType = Array<runtime::NDArray>;
+  class Applicator : public tir::StmtExprVisitor {
+   protected:
+    // returns index of the a in constant_array_, if not found - appends
+    // TODO(@d-smirnov): make real content comparision with already existing NDArrays
+    // instead of reference comparision
+    size_t deDup(const runtime::NDArray& a) {
+      auto it = std::find(constant_array_.begin(), constant_array_.end(), a);
+      if (it != constant_array_.end()) {
+        return it - constant_array_.begin();
+      }
+      constant_array_.push_back(std::move(a));
+      return constant_array_.size() - 1;
+    }
+
+   public:
+    ConstArrayType Apply(tir::Stmt body, const ConstArrayType& constant_array) {
+      constant_array_ = constant_array;
+      this->VisitStmt(body);
+      return constant_array_;
+    }
+
+    void VisitStmt_(const tir::AllocateConstNode* acn) override {
+      tir::AllocateConstNode* node = const_cast<tir::AllocateConstNode*>(acn);
+      // Check whether the data already defined within the module's attrs
+      // and replace it with array index;
+      ICHECK(node->data) << "data field should be defined";
+      if (node->data) {
+        node->irmod_storage_idx = Optional<Integer>(Integer(deDup(node->data.value())));
+      }
+      tir::StmtExprVisitor::VisitStmt_(acn);
+    }
+
+   private:
+    ConstArrayType constant_array_;
+  };
+
+  std::pair<const char*, const ObjectRef> default_value = {IRModule::_constants_attrs_key,
+                                                           Array<runtime::NDArray>()};
+  if (attrs.defined()) {
+    if (!attrs->dict.count(IRModule::_constants_attrs_key))
+      attrs.CopyOnWrite()->dict.Set(default_value.first, default_value.second);
+  } else {
+    Map<String, ObjectRef> dict = {default_value};
+    attrs = DictAttrs(dict);
+  }
+
+  ConstArrayType constant_array_ =
+      Downcast<ConstArrayType>(attrs->dict[IRModule::_constants_attrs_key]);
+
+  const ConstArrayType constant_list =
+      Applicator().Apply(func.CopyOnWrite()->body, constant_array_);
+  attrs.CopyOnWrite()->dict.Set(IRModule::_constants_attrs_key, constant_list);
+}
+
 void IRModuleNode::RegisterConstructors(const GlobalTypeVar& var, const TypeData& type) {
   // We hash the global type var name to use as a globally unique prefix for tags.
   // The hash will be used as the most significant byte of the tag, with the index of

src/printer/tvmscript_printer.cc

@@ -899,24 +899,26 @@ Doc TVMScriptPrinter::VisitStmt_(const AllocateNode* op) {
 
 Doc TVMScriptPrinter::VisitStmt_(const AllocateConstNode* alloc) {
   std::stringstream ss;
+  ICHECK(alloc->data) << "Should be presented";
+  const auto& data = alloc->data.value();
 
   if (alloc->dtype.is_int()) {
     if (alloc->dtype.bits() == 8) {
-      NDArrayToTIR<int8_t>(alloc->data, ss);
+      NDArrayToTIR<int8_t>(data, ss);
     } else if (alloc->dtype.bits() == 16) {
-      NDArrayToTIR<int16_t>(alloc->data, ss);
+      NDArrayToTIR<int16_t>(data, ss);
     } else if (alloc->dtype.bits() == 32) {
-      NDArrayToTIR<int32_t>(alloc->data, ss);
+      NDArrayToTIR<int32_t>(data, ss);
     } else {
       LOG(FATAL) << "DataType not supported";
     }
   } else if (alloc->dtype.is_float()) {
     if (alloc->dtype.bits() == 16) {
-      NDArrayToTIR<int16_t>(alloc->data, ss);
+      NDArrayToTIR<int16_t>(data, ss);
     } else if (alloc->dtype.bits() == 32) {
-      NDArrayToTIR<float>(alloc->data, ss);
+      NDArrayToTIR<float>(data, ss);
     } else if (alloc->dtype.bits() == 64) {
-      NDArrayToTIR<double>(alloc->data, ss);
+      NDArrayToTIR<double>(data, ss);
     } else {
       LOG(FATAL) << "DataType not supported";
     }

src/relay/backend/aot_executor_codegen.cc

@@ -631,7 +631,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
         int shape = kv.second->data->shape[i];
         extents.push_back(tir::make_const(DataType::Int(32), shape));
       }
-      body = tir::AllocateConst(buffer_var, kv.second->data, dtype, extents, body);
+      body = tir::AllocateConst(buffer_var, dtype, extents, kv.second->data, body);
     }
 
     // Define the attributes

src/target/llvm/codegen_llvm.cc

@@ -1434,7 +1434,8 @@ void CodeGenLLVM::VisitStmt_(const IfThenElseNode* op) {
 }
 
 void CodeGenLLVM::VisitStmt_(const AllocateConstNode* op) {
-  auto array = NDArrayToLLVMArray(ctx_, op->data);
+  auto data = op->data.value();
+  auto array = NDArrayToLLVMArray(ctx_, data);
   std::string symbol_name = op->buffer_var->name_hint;
   llvm::GlobalVariable* param_symbol = new llvm::GlobalVariable(
       *module_, array->getType(), true, llvm::GlobalValue::InternalLinkage, array, symbol_name);

src/target/source/codegen_c.cc

@@ -705,7 +705,8 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op, DataType t, PrimExpr lhs,
 void CodeGenC::VisitStmt_(const AllocateConstNode* op) {
   std::string symbol_name = op->buffer_var->name_hint;
   int64_t num_elements = 1;
-  for (int64_t dim : op->data.Shape()) {
+  const auto& data = op->data.value();
+  for (int64_t dim : data.Shape()) {
     num_elements *= dim;
   }
 
@@ -715,11 +716,11 @@ void CodeGenC::VisitStmt_(const AllocateConstNode* op) {
               << "#endif\n"
               << "static const ";
 
-  PrintType(op->data.DataType(), decl_stream);
+  PrintType(data.DataType(), decl_stream);
 
   // Allocate the global static variable
   decl_stream << " " << symbol_name << "[" << num_elements << "] = {\n";
-  NDArrayDataToC(op->data, 4, decl_stream);
+  NDArrayDataToC(data, 4, decl_stream);
 
   decl_stream << "};\n"
               << "#ifdef __cplusplus\n"