[Layout] Layout API: Shard (apache#35)

* dev

* dev

* dev

Author: spectrometerHBH

include/tvm/tir/exec_scope.h

@@ -270,6 +270,10 @@ bool Higher(const ExecScope& lhs, const ExecScope& rhs);
 
 bool Higher(const String& lhs, const String& rhs);
 
+bool Equal(const ExecScope& lhs, const ExecScope& rhs);
+
+bool Equal(const String& lhs, const String& rhs);
+
 bool ValideScope(const ExecScope& scope);
 
 bool ValideScope(const String& scope);

include/tvm/tir/layout.h

@@ -350,6 +350,18 @@ class TileLayout : public TLayout {
 /*! \brief Construct a new layout by tiling the ouer layout over the inner layout */
 TileLayout Tile(TileLayout outer, TileLayout inner);
 
+/*!
+ * \brief Construct a new layout to express the sharding strategy of a tensor.
+ * \param shape The shape of the tensor.
+ * \param mesh The device mesh
+ * \param strategy The sharding strategy of the tensor.
+ * \param inner The layout of the sharded partition of the tensor.
+ * \param from The source scope of the layout.
+ * \param to The target scope of the layout.
+ */
+TileLayout Shard(Array<PrimExpr> shape, IterTree mesh, String strategy, TileLayout inner,
+                 ExecScope from, ExecScope to);
+
 /*! \brief Layout normalization
     1. Deduplicate the split nodes in the tree, such that no two split nodes share the same child
    node.

python/tvm/tir/layout.py

@@ -176,9 +176,9 @@ def _construct_data_iter_tree(
 
     @staticmethod
     def from_nested_tuple(
-        data: Tuple,
-        strides: Tuple,
-        device: Optional[Tuple] = None,
+        data: Union[Tuple, int],
+        strides: Union[Tuple, int],
+        device: Optional[Union[Tuple, int]] = None,
         exclusive: Optional[Tuple] = None,
         from_to: Optional[Tuple[str]] = None,
     ) -> "TileLayout":
@@ -189,6 +189,11 @@ def inc_leaf_cnt():
             leaf_cnt += 1
             return leaf_cnt - 1
 
+        if not isinstance(data, tuple):
+            data = (data,)
+        if not isinstance(strides, tuple):
+            strides = (strides,)
+
         if device is None:
             assert exclusive is None, "exclusive must be None if device is None"
             assert from_to is None, "from_to must be None if device is None"
@@ -199,6 +204,9 @@ def inc_leaf_cnt():
             return TileLayout(data_tree=data_tree)
 
         else:
+            if not isinstance(device, tuple):
+                device = (device,)
+
             assert from_to is not None, "from_to must be provided if device is provided"
             assert isinstance(from_to, tuple) and len(from_to) == 2, "from_to must be a tuple of 2"
 
@@ -227,6 +235,28 @@ def inc_leaf_cnt():
     def tile(outer: "TileLayout", inner: "TileLayout") -> "TileLayout":
         return get_global_func("tir.TileLayoutTile")(outer, inner)
 
+    @staticmethod
+    def shard(
+        shape: Tuple[PrimExpr, int],
+        mesh: Tuple,
+        strategy: str,
+        inner: "TileLayout",
+        from_to: Optional[Tuple[str]] = None,
+    ) -> "TileLayout":
+        assert from_to is not None, "from_to must be provided if device is provided"
+        assert isinstance(from_to, tuple) and len(from_to) == 2, "from_to must be a tuple of 2"
+
+        f = get_global_func("tir.IterTreeFromTuple")
+        iter_tree, _ = f(convert_to_object(mesh))
+        return get_global_func("tir.TileLayoutShard")(
+            shape,
+            iter_tree,
+            strategy,
+            inner,
+            ExecScope.create(from_to[0]) if from_to else None,
+            ExecScope.create(from_to[1]) if from_to else None,
+        )
+
     @staticmethod
     def normalize(layout: "TileLayout") -> "TileLayout":
         return get_global_func("tir.NormalizeTileLayout")(layout)

src/tir/ir/exec_scope.cc

@@ -112,9 +112,7 @@ ExecScope ExecScope::Create(String name) {
 
 TVM_REGISTER_NODE_TYPE(ExecScopeNode);
 
-TVM_REGISTER_GLOBAL("tir.ExecScope").set_body_typed([](String name) {
-  return ExecScope(name);
-});
+TVM_REGISTER_GLOBAL("tir.ExecScope").set_body_typed([](String name) { return ExecScope(name); });
 
 TVM_REGISTER_GLOBAL("tir.ExecScopeCreate").set_body_typed([](String name) {
   return ExecScope::Create(name);
@@ -178,6 +176,10 @@ bool Higher(const String& lhs, const String& rhs) {
   return ScopeOrder.at(lhs) < ScopeOrder.at(rhs);
 }
 
+bool Equal(const ExecScope& lhs, const ExecScope& rhs) { return lhs->name == rhs->name; }
+
+bool Equal(const String& lhs, const String& rhs) { return lhs == rhs; }
+
 bool ValideScope(const ExecScope& scope) { return ValideScope(scope->name); }
 
 bool ValideScope(const String& scope) { return ScopeOrder.count(scope) > 0; }