adding tests for forward and forward-backward sharding propagation

Author: fschlimb

mlir/test/Dialect/Mesh/backward-sharding-propagation.mlir

@@ -0,0 +1,26 @@
+// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=backward}))" %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+  mesh.mesh @mesh(shape = 1) {sym_visibility = "private"}
+  func.func @test_forward() -> tensor<6x6xi32> {
+    %c1_i32 = arith.constant 1 : i32
+    // CHECK: tensor.empty()
+    %0 = tensor.empty() : tensor<6x6xi32>
+    %sharding = mesh.sharding @mesh split_axes = [[0]] : !mesh.sharding
+    // CHECK-COUNT-2: mesh.shard
+    %sharding_annotated = mesh.shard %0 to %sharding : tensor<6x6xi32>
+    %1 = linalg.fill ins(%c1_i32 : i32) outs(%sharding_annotated : tensor<6x6xi32>) -> tensor<6x6xi32>
+    // CHECK: tensor.empty()
+    // CHECK-NOT: mesh.shard @
+    %2 = tensor.empty() : tensor<6x6xi32>
+    %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%1, %1
+        : tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {
+    ^bb0(%in: i32, %in_2: i32, %out: i32):
+      %9 = arith.addi %in, %in_2 : i32
+      linalg.yield %9 : i32
+    } -> tensor<6x6xi32>
+    // CHECK: return
+    return %3 : tensor<6x6xi32>
+  }
+}

mlir/test/Dialect/Mesh/forward-backward-sharding-propagation.mlir

@@ -0,0 +1,27 @@
+// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward-backward}))" %s | FileCheck %s
+
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+  mesh.mesh @mesh(shape = 1) {sym_visibility = "private"}
+  func.func @test_forward() -> tensor<6x6xi32> {
+    %c1_i32 = arith.constant 1 : i32
+    // CHECK: tensor.empty()
+    %0 = tensor.empty() : tensor<6x6xi32>
+    // CHECK-COUNT-3: mesh.sharding @mesh split_axes = {{\[\[0}}]]
+    %sharding_row = mesh.sharding @mesh split_axes = [[0]] : !mesh.sharding
+    %annotated_row = mesh.shard %0 to %sharding_row : tensor<6x6xi32>
+    %1 = linalg.fill ins(%c1_i32 : i32) outs(%annotated_row : tensor<6x6xi32>) -> tensor<6x6xi32>
+    %2 = tensor.empty() : tensor<6x6xi32>
+    // CHECK-COUNT-4: mesh.sharding @mesh split_axes = {{\[\[1}}]]
+    %3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%2, %1
+        : tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {
+    ^bb0(%in: i32, %in_2: i32, %out: i32):
+      %9 = arith.addi %in, %in_2 : i32
+      linalg.yield %9 : i32
+    } -> tensor<6x6xi32>
+    %sharding_col = mesh.sharding @mesh split_axes = [[1]] : !mesh.sharding
+    %annotated_col = mesh.shard %3 to %sharding_col : tensor<6x6xi32>
+    // CHECK: return
+    return %annotated_col : tensor<6x6xi32>
+  }
+}