#sdy Add canonicalization rewrite pattern for ReshardOp.

Add a rewrite pattern to replace a chain of ReshardOp with one ReshardOp.

PiperOrigin-RevId: 653379478

shardy/dialect/sdy/ir/BUILD

@@ -8,6 +8,7 @@ td_library(
     name = "sdy_td_files",
     srcs = [
         "attrs.td",
+        "canonicalization.td",
         "dialect.td",
         "enums.td",
         "ops.td",
@@ -94,6 +95,19 @@ gentbl_cc_library(
     deps = [":sdy_td_files"],
 )
 
+gentbl_cc_library(
+    name = "canonicalization_inc",
+    tbl_outs = [
+        (
+            ["-gen-rewriters"],
+            "canonicalization.cc.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "canonicalization.td",
+    deps = [":sdy_td_files"],
+)
+
 gentbl_filegroup(
     name = "dialect_doc_gen",
     tbl_outs = [
@@ -113,6 +127,7 @@ gentbl_filegroup(
 cc_library(
     name = "dialect",
     srcs = [
+        "canonicalization.cc",
         "data_flow_utils.cc",
         "dialect.cc",
         "parsers.cc",
@@ -130,6 +145,7 @@ cc_library(
     ],
     deps = [
         ":attrs_inc",
+        ":canonicalization_inc",
         ":dialect_inc",
         ":enums_inc",
         ":ops_inc",

shardy/dialect/sdy/ir/canonicalization.cc

@@ -0,0 +1,36 @@
+/* Copyright 2024 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+
+namespace mlir {
+namespace sdy {
+
+namespace {
+
+#include "shardy/dialect/sdy/ir/canonicalization.cc.inc"
+
+}  // namespace
+
+void ReshardOp::getCanonicalizationPatterns(RewritePatternSet& results,
+                                            MLIRContext* context) {
+  results.add<ReshardOfReshardPattern>(context);
+}
+
+}  // namespace sdy
+}  // namespace mlir

shardy/dialect/sdy/ir/canonicalization.td

@@ -0,0 +1,28 @@
+/* Copyright 2024 The Shardy Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef SDY_CANONICALIZATION
+#define SDY_CANONICALIZATION
+
+include "shardy/dialect/sdy/ir/ops.td"
+include "mlir/IR/PatternBase.td"
+
+def HasOneUse: Constraint<CPred<"$_self.hasOneUse()">, "has one use">;
+
+def ReshardOfReshardPattern :
+    Pat<(Sdy_ReshardOp (Sdy_ReshardOp:$inner_reshard $tensor, $inner_attr), $outer_attr),
+        (Sdy_ReshardOp $tensor, $outer_attr), [(HasOneUse:$inner_reshard)]>;
+
+#endif  // SDY_CANONICALIZATION

shardy/dialect/sdy/ir/ops.td

@@ -104,6 +104,7 @@ def Sdy_ReshardOp : Sdy_Op<"reshard",
 
   let assemblyFormat = "$input $sharding attr-dict `:` type($result)";
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Sdy_ReturnOp : Sdy_Op<"return", [Pure, Terminator]> {

shardy/dialect/sdy/ir/test/reshard_canonicalization.mlir

@@ -0,0 +1,33 @@
+// RUN: sdy_opt %s -canonicalize | FileCheck %s
+
+sdy.mesh @mesh = <"a"=2, "b"=2>
+
+// CHECK-LABEL: func @reshard_of_reshard_no_other_uses
+// CHECK-NEXT: %0 = sdy.reshard %arg0 <@mesh, [{"a", ?}, {?}]>
+// CHECK-NEXT: return %0
+func.func @reshard_of_reshard_no_other_uses(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %0 = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+  %1 = sdy.reshard %0 <@mesh, [{"a", ?}, {?}]> : tensor<8x8xf32>
+  return %1 : tensor<8x8xf32>
+}
+
+// CHECK-LABEL: func @reshard_of_reshard_with_other_uses
+// CHECK-NEXT: %0 = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]>
+// CHECK-NEXT: %1 = sdy.reshard %0 <@mesh, [{"a", ?}, {?}]>
+// CHECK-NEXT: return %0, %1
+func.func @reshard_of_reshard_with_other_uses(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>, tensor<8x8xf32>)  {
+  %0 = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+  %1 = sdy.reshard %0 <@mesh, [{"a", ?}, {?}]> : tensor<8x8xf32>
+  return %0, %1 : tensor<8x8xf32>, tensor<8x8xf32>
+}
+
+// CHECK-LABEL: func @reshard_chain_of_three_no_other_uses
+// CHECK-NEXT: %0 = sdy.reshard %arg0 <@mesh, [{?}, {"a", ?}]>
+// CHECK-NEXT: return %0
+func.func @reshard_chain_of_three_no_other_uses(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  %0 = sdy.reshard %arg0 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+  %1 = sdy.reshard %0 <@mesh, [{"a", ?}, {?}]> : tensor<8x8xf32>
+  %2 = sdy.reshard %1 <@mesh, [{?}, {"a", ?}]> : tensor<8x8xf32>
+  return %2 : tensor<8x8xf32>
+}
+