#sdy add a canonicalization pattern to sdy.manual_computation that removes used block args and their corresponding operand.

shardy/dialect/sdy/ir/BUILD

@@ -113,6 +113,7 @@ gentbl_filegroup(
 cc_library(
     name = "dialect",
     srcs = [
+        "canonicalization.cc",
         "data_flow_utils.cc",
         "dialect.cc",
         "parsers.cc",

shardy/dialect/sdy/ir/canonicalization.cc

@@ -0,0 +1,73 @@
+/* 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 <cstdint>
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LLVM.h"
+#include "shardy/dialect/sdy/ir/dialect.h"
+
+namespace mlir {
+namespace sdy {
+
+namespace {
+
+// Pattern to remove unused block arguments and their corresponding operands
+// from  a `ManualComputationOp`.
+class ManualComputationUnusedInputsPattern
+    : public OpRewritePattern<ManualComputationOp> {
+ public:
+  using OpRewritePattern<ManualComputationOp>::OpRewritePattern;
+
+ private:
+  LogicalResult matchAndRewrite(ManualComputationOp manualComputationOp,
+                                PatternRewriter& rewriter) const override {
+    BitVector unusedArgs(manualComputationOp.getNumOperands());
+    for (BlockArgument arg : manualComputationOp.getRegion().getArguments()) {
+      if (arg.use_empty()) {
+        unusedArgs.set(arg.getArgNumber());
+      }
+    }
+    if (unusedArgs.none()) {
+      return failure();
+    }
+
+    manualComputationOp->eraseOperands(unusedArgs);
+    manualComputationOp.getRegion().front().eraseArguments(unusedArgs);
+
+    SmallVector<TensorShardingAttr> inShardings;
+    inShardings.reserve(manualComputationOp.getNumOperands());
+    for (int64_t index : unusedArgs.flip().set_bits()) {
+      inShardings.push_back(manualComputationOp.getInSharding(index));
+    }
+    manualComputationOp.setInShardingsAttr(TensorShardingPerValueAttr::get(
+        manualComputationOp.getContext(), inShardings));
+
+    return success();
+  }
+};
+
+}  // namespace
+
+void ManualComputationOp::getCanonicalizationPatterns(
+    RewritePatternSet& results, MLIRContext* context) {
+  results.add<ManualComputationUnusedInputsPattern>(context);
+}
+
+}  // namespace sdy
+}  // namespace mlir

shardy/dialect/sdy/ir/ops.td

@@ -154,6 +154,7 @@ def Sdy_ManualComputationOp : Sdy_Op<"manual_computation",
   }];
 
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 
   let extraClassDeclaration = [{
     TensorShardingAttr getInSharding(int64_t operandIndex) {

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

@@ -0,0 +1,18 @@
+// RUN: sdy_opt %s -canonicalize | FileCheck %s
+
+sdy.mesh @mesh = <"a"=2>
+
+// CHECK-LABEL: func @unused_args
+func.func @unused_args(%arg0: tensor<8xf32>, %arg1: tensor<32x32xf32>, %arg2: tensor<16xf32>) -> tensor<32x32xf32> {
+  // CHECK-NEXT: sdy.manual_computation(%arg1)
+  // CHECK-SAME:     in_shardings=[<@mesh, [{"a"}, {}]>]
+  // CHECK-SAME:     out_shardings=[<@mesh, [{"a"}, {}]>]
+  // CHECK-SAME:     manual_axes={"a"} (%arg3: tensor<16x32xf32>) {
+  // CHECK-NEXT:   sdy.return %arg3
+  // CHECK-NEXT: }
+  %0 = sdy.manual_computation(%arg0, %arg1, %arg2) in_shardings=[<@mesh, [{"a"}]>, <@mesh, [{"a"}, {}]>, <@mesh, [{"a"}]>] out_shardings=[<@mesh, [{"a"}, {}]>]
+      manual_axes={"a"} (%arg3: tensor<4xf32>, %arg4: tensor<16x32xf32>, %arg5: tensor<8xf32>) {
+    sdy.return %arg4 : tensor<16x32xf32>
+  } : (tensor<8xf32>, tensor<32x32xf32>, tensor<16xf32>) -> tensor<32x32xf32>
+  func.return %0: tensor<32x32xf32>
+}