[RELAX][LAYOUT] Support multiple axis paching.

Like OIHW[4o4i] where we can pack multiple axis.
Helpful while handling complex target layouts.
This PR covers layout representation and transforms for these.

Author: srkreddy1238

include/tvm/tir/data_layout.h

@@ -35,6 +35,8 @@
 #include <utility>
 #include <vector>
 
+#include "tvm/tir/var.h"
+
 namespace tvm {
 namespace tir {
 
@@ -158,6 +160,22 @@ class Layout : public ObjectRef {
     return undef;
   }
 
+  /*!
+   * \brief Packs the Given Array of IterVars into a Single IterVar. Each IterVar in the Array
+   *        should represent either a single primal axis or one or more subordinate axis
+   * \param iters Array of iter vars to be packed
+   * \return A packed iter var
+   */
+  static IterVar PackIterVar(ffi::Array<IterVar> iters);
+
+  /*!
+   * \brief Unpacks a Packed IterVar into its constituents
+   * \param packed_iter A Packed IterVar containing a single primal axis or one or more subordinate
+   *                    axis
+   * \return Constituent IterVars
+   */
+  static ffi::Array<IterVar> UnpackIterVar(IterVar packed_iter);
+
   /*!
    * \brief Returns a sub-layout which is the portion of the object
    *        that starts at dimension \p pos and spans \p len dimensions
@@ -187,9 +205,12 @@ class Layout : public ObjectRef {
   inline size_t ndim_primal() const {
     if (!defined()) return 0;
     size_t ct = 0;
-    for (auto x : operator->()->axes) {
-      if (LayoutAxis::Get(x).IsPrimal()) {
-        ct++;
+    for (auto px : operator->()->axes) {
+      auto iter_vars = UnpackIterVar(px);
+      for (auto x : iter_vars) {
+        if (LayoutAxis::Get(x).IsPrimal()) {
+          ct++;
+        }
       }
     }
     return ct;
@@ -204,10 +225,13 @@ class Layout : public ObjectRef {
     Layout new_src_layout;
     // 1) Find the axis which are missing in the current layout. Make them the prefix.
     std::string new_src_layout_str = "";
-    for (auto dst_axis : dst_layout->axes) {
-      if (LayoutAxis::Get(dst_axis).IsPrimal()) {
-        if (!this->Contains(LayoutAxis::Get(dst_axis))) {
-          new_src_layout_str += dst_axis->var->name_hint;
+    for (auto packed_axis : dst_layout->axes) {
+      auto iter_vars = UnpackIterVar(packed_axis);
+      for (auto dst_axis : iter_vars) {
+        if (LayoutAxis::Get(dst_axis).IsPrimal()) {
+          if (!this->Contains(LayoutAxis::Get(dst_axis))) {
+            new_src_layout_str += dst_axis->var->name_hint;
+          }
         }
       }
     }
@@ -221,18 +245,36 @@ class Layout : public ObjectRef {
    * \brief return the index of the input axis.
    *        If it is not found in the layout or the layout is undefined,
    *        return -1.
-   * \param axis the input axis.
+   * \param axis The input axis either a layout axis, or a packed axis
    * \return the index or -1 if not found.
    */
-  inline int32_t IndexOf(const LayoutAxis& axis) const {
+  inline int32_t IndexOf(const std::string& axis) const {
     if (!this->defined()) return -1;
     const auto axes = operator->()->axes;
     for (size_t i = 0; i < axes.size(); ++i) {
-      if (axes[i]->var->name_hint == axis.name()) return static_cast<int32_t>(i);
+      if (axes[i]->var->name_hint == axis) return static_cast<int32_t>(i);
     }
     return -1;
   }
 
+  /*!
+   * \brief return the index of the input axis.
+   *        If it is not found in the layout or the layout is undefined,
+   *        return -1.
+   * \param axis the input layout axis.
+   * \return the index or -1 if not found.
+   */
+  inline int32_t IndexOf(const LayoutAxis& axis) const { return IndexOf(axis.name()); }
+
+  /*!
+   * \brief return the index of the input axis.
+   *        If it is not found in the layout or the layout is undefined,
+   *        return -1.
+   * \param iter the input iter var.
+   * \return the index or -1 if not found.
+   */
+  inline int32_t IndexOf(const tir::IterVar& iter) const { return IndexOf(iter->var->name_hint); }
+
   /*!
    * \brief Get the factor size of the subordinate axis.
    * \param axis the input primal-axis or subordinate-axis.
@@ -249,20 +291,23 @@ class Layout : public ObjectRef {
    */
   bool Contains(const LayoutAxis& axis) const {
     if (!defined()) return false;
-    for (const tir::IterVar var : operator->()->axes) {
-      if (var->var->name_hint == axis.name()) {
-        return true;
+    for (const tir::IterVar packed_var : operator->()->axes) {
+      auto iter_vars = UnpackIterVar(packed_var);
+      for (auto var : iter_vars) {
+        if (var->var->name_hint == axis.name()) {
+          return true;
+        }
       }
     }
     return false;
   }
 
-  const LayoutAxis& operator[](int32_t i) const {
+  IterVar operator[](int32_t i) const {
     ICHECK(defined()) << "Try to access axis from an undefined layout.";
     int32_t index = i < 0 ? static_cast<int32_t>(ndim() + i) : i;
     ICHECK(index >= 0 && static_cast<size_t>(index) < ndim()) << "Invalid index " << i;
     const tir::IterVar axis = operator->()->axes[index];
-    return LayoutAxis::Get(axis);
+    return axis;
   }
 
   /*! \return the string description of the layout */

python/tvm/tir/data_layout.py

@@ -41,7 +41,8 @@ def __len__(self):
         return _ffi_api.LayoutNdim(self)  # type: ignore
 
     def __contains__(self, axis):
-        return len(axis) == 1 and axis[0].isalpha() and axis[0] in self.name
+        # Note: We do a weaker check for packed axis assuming layout is valid
+        return not any(bkt in axis for bkt in "[]") and axis in self.name
 
     def __getitem__(self, index):
         if index >= len(self):
@@ -54,7 +55,7 @@ def index_of(self, axis):
         Parameters
         ----------
         axis : str
-            The axis name, need to be [a-z,A-Z]
+            The axis name, needs to be [a-z,A-Z] or a packed axis
 
         Returns
         -------

src/contrib/msc/core/ir/graph_builder.cc

@@ -454,7 +454,8 @@ const MSCJoint GraphBuilder::AddNode(const Expr& expr, const ffi::Optional<Expr>
           } else if (input_types[i] == "weight" &&
                      (optype == "msc.linear" || optype == "msc.linear_bias")) {
             if (ref->layout.name() == "IO") {
-              ffi::String valid_layout = ref->layout[1].name() + ref->layout[0].name();
+              ffi::String valid_layout =
+                  ref->layout[1]->var->name_hint + ref->layout[0]->var->name_hint;
               const auto& valid_shape = ffi::Array<Integer>({ref->shape[1], ref->shape[0]});
               weight = MSCTensor(weight_name, ref->dtype, valid_layout, valid_shape);
             } else {

src/contrib/msc/core/transform/layout_utils.cc

@@ -158,7 +158,7 @@ const LayoutDecision LayoutUtils::ExpandLayout(const LayoutDecision& src_layout,
       << "Only support normal layout, get " << src_layout->layout;
   std::set<std::string> used_axes;
   for (size_t i = 0; i < src_layout->layout.ndim(); i++) {
-    used_axes.insert(src_layout->layout[i].name());
+    used_axes.insert(src_layout->layout[i]->var->name_hint);
   }
   std::vector<std::string> prefer_axes{"N", "C", "H", "W", "D"};
   for (const auto& a : axes) {
@@ -198,7 +198,7 @@ const LayoutDecision LayoutUtils::ReduceLayout(const LayoutDecision& src_layout,
     if (reduce_axes_set.count(i)) {
       continue;
     }
-    new_layout += src_layout->layout[i].name();
+    new_layout += src_layout->layout[i]->var->name_hint;
   }
   return LayoutDecision(new_layout);
 }
@@ -207,7 +207,7 @@ const LayoutDecision LayoutUtils::PermuteLayout(const LayoutDecision& src_layout
                                                 const ffi::Array<Integer>& axes) {
   ffi::String layout_str;
   for (const auto& a : axes) {
-    layout_str = layout_str + src_layout->layout[a->value].name();
+    layout_str = layout_str + src_layout->layout[a->value]->var->name_hint;
   }
   return LayoutDecision(layout_str);
 }
@@ -216,7 +216,7 @@ const LayoutDecision LayoutUtils::PermuteLayout(const LayoutDecision& src_layout
                                                 const std::vector<size_t>& axes) {
   ffi::String layout_str;
   for (const auto& a : axes) {
-    layout_str = layout_str + src_layout->layout[a].name();
+    layout_str = layout_str + src_layout->layout[a]->var->name_hint;
   }
   return LayoutDecision(layout_str);
 }
@@ -226,7 +226,7 @@ int LayoutUtils::InferBatchDim(const LayoutDecision& layout) {
     return -1;
   }
   for (size_t i = 0; i < layout->layout.ndim(); i++) {
-    if (layout->layout[i].name() == "N") {
+    if (layout->layout[i]->var->name_hint == "N") {
       return static_cast<int>(i);
     }
   }

src/contrib/msc/core/transform/set_expr_layout.cc

@@ -246,7 +246,7 @@ InferLayoutOutput ForwardInferLayoutBinary(
         input_layouts.push_back(LayoutDecision(""));
       } else if (t_info->ndim == 1) {
         const auto& ref_layout = output->output_layouts[0].LeafValue()->layout;
-        input_layouts.push_back(LayoutDecision(ref_layout[ref_layout.ndim() - 1].name()));
+        input_layouts.push_back(LayoutDecision(ref_layout[ref_layout.ndim() - 1]->var->name_hint));
       } else {
         input_layouts.push_back(output->input_layouts[i]);
       }
@@ -361,7 +361,7 @@ InferLayoutOutput ForwardInferLayoutMatmul(
   size_t start = a_layout->layout.ndim() - b_shape.size();
   ffi::String pre_layout;
   for (size_t i = start; i < a_layout->layout.ndim() - 2; i++) {
-    pre_layout = pre_layout + a_layout->layout[i].name();
+    pre_layout = pre_layout + a_layout->layout[i]->var->name_hint;
   }
   LayoutDecision b_layout = LayoutDecision(pre_layout + "IO");
   return InferLayoutOutput({a_layout, b_layout}, {a_layout}, Attrs());
@@ -671,7 +671,7 @@ InferLayoutOutput BackwardInferLayoutBinary(
         input_layouts.push_back(LayoutDecision(""));
       } else if (t_info->ndim == 1) {
         const auto& ref_layout = output->output_layouts[0].LeafValue()->layout;
-        input_layouts.push_back(LayoutDecision(ref_layout[ref_layout.ndim() - 1].name()));
+        input_layouts.push_back(LayoutDecision(ref_layout[ref_layout.ndim() - 1]->var->name_hint));
       } else {
         input_layouts.push_back(output->input_layouts[i]);
       }
@@ -766,7 +766,7 @@ InferLayoutOutput BackwardInferLayoutMatmul(
   size_t start = output_layout->layout.ndim() - b_shape.size();
   ffi::String pre_layout;
   for (size_t i = start; i < output_layout->layout.ndim() - 2; i++) {
-    pre_layout = pre_layout + output_layout->layout[i].name();
+    pre_layout = pre_layout + output_layout->layout[i]->var->name_hint;
   }
   LayoutDecision b_layout = LayoutDecision(pre_layout + "IO");
   return InferLayoutOutput({output_layout, b_layout}, {output_layout}, Attrs());

src/contrib/msc/framework/tensorrt/tensorrt_opcode.cc

@@ -329,7 +329,7 @@ class TensorRTConvCodeGen : public TensorRTOpCode {
     const auto& weight = node()->WeightAt("weight");
     std::vector<int64_t> kernel_size;
     for (size_t i = 0; i < weight->Ndim(); i++) {
-      if (weight->layout[i].name() == "I" || weight->layout[i].name() == "O") {
+      if (weight->layout[i]->var->name_hint == "I" || weight->layout[i]->var->name_hint == "O") {
         continue;
       }
       kernel_size.push_back(weight->DimAt(i)->value);
@@ -442,10 +442,10 @@ class TensorRTPadCodeGen : public TensorRTOpCode {
     std::vector<int> pre_padding{2, 0}, post_padding{2, 0};
     const auto& input = node()->InputAt(0);
     for (size_t i = 0; i < input->Ndim(); i++) {
-      if (input->layout[i].name() == "H") {
+      if (input->layout[i]->var->name_hint == "H") {
         pre_padding[0] = pad_width[i * 2];
         post_padding[0] = pad_width[i * 2 + 1];
-      } else if (input->layout[i].name() == "W") {
+      } else if (input->layout[i]->var->name_hint == "W") {
         pre_padding[1] = pad_width[i * 2];
         post_padding[1] = pad_width[i * 2 + 1];
       }

src/contrib/msc/framework/torch/torch_opcode.cc

@@ -276,7 +276,7 @@ class TorchConvCodeGen : public TorchOpCode {
     const auto& weight = node()->WeightAt("weight");
     std::vector<int64_t> kernel_size;
     for (size_t i = 0; i < weight->Ndim(); i++) {
-      if (weight->layout[i].name() == "I" || weight->layout[i].name() == "O") {
+      if (weight->layout[i]->var->name_hint == "I" || weight->layout[i]->var->name_hint == "O") {
         continue;
       }
       kernel_size.push_back(weight->DimAt(i)->value);

src/contrib/msc/framework/tvm/relax_opcode.cc

@@ -288,7 +288,7 @@ class RelaxConvCodeGen : public RelaxOpCode {
       const auto& out_layout = tir::Layout(out_layout_str);
       ffi::Array<Integer> expand_shape;
       for (size_t i = 0; i < node()->OutputAt(0)->Ndim(); i++) {
-        if (out_layout[i].name() == "C") {
+        if (out_layout[i]->var->name_hint == "C") {
           expand_shape.push_back(node()->OutputAt(0)->DimAt(i));
         } else {
           expand_shape.push_back(Integer(1));

src/relax/op/tensor/manipulate.cc

@@ -1968,7 +1968,7 @@ InferLayoutOutput InferLayoutTile(
     // Same dimension: reorder repeats according to layout transformation.
     // If len(repeats) < ndim, it's padded with 1s at the beginning.
     for (int i = 0; i < ndim; ++i) {
-      const tir::LayoutAxis& axis = existing_layout_obj[i];
+      const tir::LayoutAxis& axis = tir::LayoutAxis::Get(existing_layout_obj[i]);
       int pos_in_initial = initial_layout.IndexOf(axis);
       ICHECK_NE(pos_in_initial, -1) << "Axis not found in initial layout";
       // If len(repeats) < ndim, repeats are right-aligned.
@@ -1990,7 +1990,7 @@ InferLayoutOutput InferLayoutTile(
     }
     // Repeats for existing dimensions need to be permuted.
     for (int i = 0; i < ndim; ++i) {
-      const tir::LayoutAxis& axis = existing_layout_obj[i];
+      const tir::LayoutAxis& axis = tir::LayoutAxis::Get(existing_layout_obj[i]);
       int pos_in_initial = initial_layout.IndexOf(axis);
       ICHECK_NE(pos_in_initial, -1) << "Axis not found in initial layout";
       new_repeats.push_back(attrs->repeats[pos_in_initial + num_new_dims]);