[TIR] tir.transform.StorageFlatten refactor

include/tvm/tir/buffer.h

@@ -121,6 +121,14 @@ class BufferNode : public Object {
     return shape.size() != 0 ? shape[0].dtype() : DataType::Int(32);
   }
 
+  /*! \brief Determine the offset in the buffer of the given index.
+   *
+   * Returns the buffer offset, in number of elements of type dtype,
+   * without adjusting for number of lanes.  (e.g. The number of
+   * float16x4 elements in a buffer of type float16x4.)
+   */
+  PrimExpr ElemOffset(Array<PrimExpr> index) const;
+
   static constexpr const char* _type_key = "tir.Buffer";
   static constexpr const bool _type_has_method_sequal_reduce = true;
   static constexpr const bool _type_has_method_shash_reduce = true;

src/tir/ir/buffer.cc

@@ -246,41 +246,41 @@ inline PrimExpr MergeMulMod(arith::Analyzer* analyzer, const PrimExpr& base) {
 // The buffer offset in convention of number of elements of
 // original data ignoring number of lanes.
 // We also perform optimization to simplify the indexing expression.
-inline PrimExpr ElemOffset(const BufferNode* n, Array<PrimExpr> index) {
-  PrimExpr base = n->elem_offset;
+PrimExpr BufferNode::ElemOffset(Array<PrimExpr> index) const {
+  PrimExpr base = this->elem_offset;
   arith::Analyzer ana;
-  if (n->strides.size() == 0) {
+  if (this->strides.size() == 0) {
     // Scalar case
-    if (n->shape.size() == 0 && index.size() == 1) {
+    if (this->shape.size() == 0 && index.size() == 1) {
       auto is_int = index[0].as<IntImmNode>();
       ICHECK(is_int && is_int->value == 0);
       base = base + index[0];
     } else {
-      ICHECK_EQ(n->shape.size(), index.size());
+      ICHECK_EQ(this->shape.size(), index.size());
       if (index.size() > 0) {
         PrimExpr offset = index[0];
         for (size_t i = 1; i < index.size(); ++i) {
-          offset = MergeMulMod(&ana, offset * n->shape[i] + index[i]);
+          offset = MergeMulMod(&ana, offset * this->shape[i] + index[i]);
         }
         base = base + offset;
       }
     }
   } else {
-    ICHECK_EQ(n->strides.size(), index.size());
+    ICHECK_EQ(this->strides.size(), index.size());
     if (is_zero(base)) {
-      base = MergeMulMod(&ana, index[0] * n->strides[0]);
+      base = MergeMulMod(&ana, index[0] * this->strides[0]);
     } else {
-      base = MergeMulMod(&ana, base + index[0] * n->strides[0]);
+      base = MergeMulMod(&ana, base + index[0] * this->strides[0]);
     }
     for (size_t i = 1; i < index.size(); ++i) {
-      base = MergeMulMod(&ana, base + index[i] * n->strides[i]);
+      base = MergeMulMod(&ana, base + index[i] * this->strides[i]);
     }
   }
   return base;
 }
 
 inline PrimExpr BufferOffset(const BufferNode* n, Array<PrimExpr> index, DataType dtype) {
-  PrimExpr offset = ElemOffset(n, index);
+  PrimExpr offset = n->ElemOffset(index);
   if (n->dtype.lanes() != 1) {
     offset = offset * make_const(offset.dtype(), dtype.lanes());
   }
@@ -353,7 +353,7 @@ Buffer Buffer::MakeSlice(Array<PrimExpr> begins, Array<PrimExpr> extents) const
   ICHECK(n != nullptr);
   arith::Analyzer ana;
   begins = SimplifyArray(&ana, begins);
-  PrimExpr elem_offset = ana.Simplify(ElemOffset(n, begins));
+  PrimExpr elem_offset = ana.Simplify(n->ElemOffset(begins));
   Array<PrimExpr> strides = n->strides;
   if (strides.size() == 0) {
     bool can_relax = true;

src/tir/transforms/arg_binder.cc

@@ -204,7 +204,7 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
   def_handle_dtype_.Set(v_strides, tir::TypeAnnotation(tvm_shape_type));
   init_nest_.emplace_back(
       LetStmt(v_strides, TVMArrayGet(DataType::Handle(), handle, builtin::kArrStrides), nop));
-  PrimExpr is_null = Call(DataType::Bool(1), builtin::isnullptr(), {v_strides});
+  PrimExpr v_strides_is_null = Call(DataType::Bool(1), builtin::isnullptr(), {v_strides});
   if (buffer->strides.size() == 0) {
     // Assert the buffer is compact
     DataType stype = buffer->DefaultIndexType();
@@ -226,7 +226,7 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
           foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); },
                 const_true(1), conds),
           stride_msg, Evaluate(0));
-      check = IfThenElse(Not(is_null), check, Stmt());
+      check = IfThenElse(Not(v_strides_is_null), check, Stmt());
       asserts_.emplace_back(SeqStmt({check, Evaluate(0)}));
     }
   } else if (buffer->buffer_type == kAutoBroadcast) {
@@ -239,24 +239,29 @@ void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
       PrimExpr value =
           cast(buffer->shape[k].dtype(),
                Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1)));
-      value = tvm::if_then_else(is_null, stride, value);
+      value = tvm::if_then_else(v_strides_is_null, stride, value);
       value = tvm::if_then_else(buffer->shape[k] == 1, 0, value);
       Bind_(buffer->strides[k], value, field_name.str(), true);
       stride = analyzer_.Simplify(stride * buffer->shape[k]);
     }
   } else {
-    std::ostringstream stride_null_err_msg;
-    stride_null_err_msg << arg_name << ".strides: expected non-null strides.";
-    asserts_.emplace_back(
-        AssertStmt(Not(is_null), tvm::tir::StringImm(stride_null_err_msg.str()), nop));
+    PrimExpr stride_from_shape = 1;
 
-    for (size_t k = 0; k < buffer->strides.size(); ++k) {
+    for (int k = buffer->strides.size() - 1; k >= 0; k--) {
       std::ostringstream field_name;
       field_name << v_strides->name_hint << '[' << k << ']';
+
+      PrimExpr explicit_stride =
+          cast(buffer->shape[k].dtype(),
+               Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1)));
+
       Bind_(buffer->strides[k],
-            cast(buffer->shape[k].dtype(),
-                 Load(tvm_shape_type, v_strides, IntImm(DataType::Int(32), k), const_true(1))),
+            tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride),
             field_name.str(), true);
+
+      stride_from_shape *=
+          cast(buffer->shape[k].dtype(),
+               Load(tvm_shape_type, v_shape, IntImm(DataType::Int(32), k), const_true(1)));
     }
   }
   // Byte_offset field.