[mlir] Commutative Operand Sorting Pass

[amandatang]

Summary

A characteristic of commutative operations is that such operations with swapped operands are identical.

This patch creates a pass that applies the SortCommutativeOperands rewrite pattern. The intention is for the sorting pass to be applied first to reorder operands in a deterministic way. This is followed by CSE which simply compares operands in order.

[jpienaar]

mlir/include/mlir/Transforms/Passes.h

49

Lets move this below canonicalize for some alphabetical order

mlir/include/mlir/Transforms/Passes.td

22

s/implements/uses/

This is calling the helping method that implements the sorting algo.

23

Could you add more about the sort here? Even if it is just to point to the helper method for more info, I feel that gets lost here.

mlir/lib/Transforms/CommutativeOperandSort.cpp

36

Could we freeze the pattern set in the pass initializer?

mlir/test/Transforms/commutative-operand-sort.mlir

2

Lets check with just this pass (fine to add a test with cse too if you want, but lets unit test this one in isolation).

[amandatang]

Address comments

[amandatang]

mlir/test/Transforms/commutative-operand-sort.mlir

2

In https://reviews.llvm.org/D157528, there are many additional tests in mlir/test/Transforms/test-commutativity-utils.mlir that check the sorting rewrite pattern alone. Currently, a test pass that applies the rewrite pattern in mlir/test/lib/Transforms/TestCommutativityUtils.cpp is used.

I was thinking the CommutativeOperandSort pass can be used there instead of the test pass to check the pass and sorting logic in isolation.

[srishti-pm]

Thanks for adding this. An initial comment:-

mlir/include/mlir/Transforms/Passes.td

55

The meaning of "ancestors" is not explained here. Thus, the right word to use is "backward slice", which is a known term in MLIR.

[amandatang]

Address comment

[amandatang]

@srishti-pm @jpienaar Any other followups on this?

[srishti-pm]

Please update the commit summary. It only contains the title currently. Also, change your revision and commit titles to [MLIR][transforms] Add a pass to sort commutative operands.

[jpienaar]

In D158444#4611742, @amandatang wrote:

@srishti-pm @jpienaar Any other followups on this?

Looks good in general.

In D158444#4618868, @srishti-pm wrote:

Please update the commit summary. It only contains the title currently. Also, change your revision and commit titles to [MLIR][transforms] Add a pass to sort commutative operands.

I see a summary was added Aug 21. Perhaps it didn't refresh correctly your side?

mlir/include/mlir/Transforms/Passes.h

65

Nit: commutative and end sentence with period.

mlir/lib/Transforms/CommutativeOperandSort.cpp

36

This is marked done but I don't see the change.

mlir/test/Transforms/commutative-operand-sort.mlir

2

Thats fine but unrelated to running only one pass here. While it's good to CSE applying post, this wouldn't be a surprise if one saw the operands are similarly sorted. Just test that the sorting here resulted in the same order using filecheck should be sufficient to see pass is registered and behaving as expected. If more tests can reuse same pass then that's good.

mlir/include/mlir/Transforms/Passes.h

 namespace mlir {
 
 class GreedyRewriteConfig;
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
+#define GEN_PASS_DECL_COMMUTATIVEOPERANDSORT
 #define GEN_PASS_DECL_CANONICALIZER
 #define GEN_PASS_DECL_CONTROLFLOWSINK
 #define GEN_PASS_DECL_CSEPASS
 #define GEN_PASS_DECL_INLINER
 #define GEN_PASS_DECL_LOOPINVARIANTCODEMOTION
 #define GEN_PASS_DECL_MEM2REG
 #define GEN_PASS_DECL_PRINTIRPASS
 #define GEN_PASS_DECL_PRINTOPSTATS
 #define GEN_PASS_DECL_SROA
 #define GEN_PASS_DECL_STRIPDEBUGINFO
 #define GEN_PASS_DECL_SCCP
 #define GEN_PASS_DECL_SYMBOLDCE
 #define GEN_PASS_DECL_SYMBOLPRIVATIZE
 #define GEN_PASS_DECL_TOPOLOGICALSORT
 #include "mlir/Transforms/Passes.h.inc"
 
 /// Creates an instance of the Canonicalizer pass, configured with default
 /// settings (which can be overridden by pass options on the command line).
 std::unique_ptr<Pass> createCanonicalizerPass();
 
 /// Creates an instance of the Canonicalizer pass with the specified config.
 /// `disabledPatterns` is a set of labels used to filter out input patterns with
 /// a debug label or debug name in this set. `enabledPatterns` is a set of
 /// labels used to filter out input patterns that do not have one of the labels
 /// in this set. Debug labels must be set explicitly on patterns or when adding
 /// them with `RewritePatternSet::addWithLabel`. Debug names may be empty, but
 /// patterns created with `RewritePattern::create` have their default debug name
 /// set to their type name.
 std::unique_ptr<Pass>
 createCanonicalizerPass(const GreedyRewriteConfig &config,
                         ArrayRef<std::string> disabledPatterns = std::nullopt,
                         ArrayRef<std::string> enabledPatterns = std::nullopt);
 
+/// Creates a pass to deterministically sort Commutative operation operands
+std::unique_ptr<Pass> createCommutativeOperandSortPass();
+
 /// Creates a pass to perform control-flow sinking.
 std::unique_ptr<Pass> createControlFlowSinkPass();
 
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
 /// Creates a pass to print IR on the debug stream.
 std::unique_ptr<Pass> createPrintIRPass(const PrintIRPassOptions & = {});

mlir/include/mlir/Transforms/Passes.td

 #define MLIR_TRANSFORMS_PASSES
 
 include "mlir/Pass/PassBase.td"
 include "mlir/Rewrite/PassUtil.td"
 
 def Canonicalizer : Pass<"canonicalize"> {
   let summary = "Canonicalize operations";
   let description = [{
     This pass performs various types of canonicalizations over a set of
     operations by iteratively applying the canonicalization patterns of all
     loaded dialects until either a fixpoint is reached or the maximum number of
     iterations/rewrites is exhausted. Canonicalization is best-effort and does
     not guarantee that the entire IR is in a canonical form after running this
     pass. See [Operation Canonicalization](Canonicalization.md) for more
     details.
   }];
   let constructor = "mlir::createCanonicalizerPass()";
   let options = [
     Option<"topDownProcessingEnabled", "top-down", "bool",
            /*default=*/"true",
            "Seed the worklist in general top-down order">,
     Option<"enableRegionSimplification", "region-simplify", "bool",
            /*default=*/"true",
            "Perform control flow optimizations to the region tree">,
     Option<"maxIterations", "max-iterations", "int64_t",
            /*default=*/"10",
            "Max. iterations between applying patterns / simplifying regions">,
     Option<"maxNumRewrites", "max-num-rewrites", "int64_t", /*default=*/"-1",
            "Max. number of pattern rewrites within an iteration">,
     Option<"testConvergence", "test-convergence", "bool", /*default=*/"false",
            "Test only: Fail pass on non-convergence to detect cyclic pattern">
   ] # RewritePassUtils.options;
 }
 
+def CommutativeOperandSort : Pass<"commutative-operand-sort"> {
+  let summary = "Sort operands of commutative operands in a deterministic way";
+  let description = [{
+    This pass calls a deterministic sorting algorithm for commutative
+    operands. This way commutative operations with swapped operands can be
+    treated as equal and removed by the CSE pass.
+
+    The sorting algorithm looks at the ancestors of operands
+    to establish an ordering. See 
+    mlir/lib/Transforms/Utils/CommutativityUtils.cpp for more details. 
+  }];
+  let constructor = "mlir::createCommutativeOperandSortPass()";
+}
+
 def ControlFlowSink : Pass<"control-flow-sink"> {
   let summary = "Sink operations into conditional blocks";
   let description = [{
     This pass implements control-flow sink on operations that implement
     `RegionBranchOpInterface` by moving dominating operations whose only uses
     are in a conditionally-executed regions into those regions so that
     executions paths where their results are not needed do not perform
     unnecessary computations.

mlir/lib/Transforms/CMakeLists.txt

 add_subdirectory(Utils)
 
 add_mlir_library(MLIRTransforms
   Canonicalizer.cpp
+  CommutativeOperandSort.cpp
   ControlFlowSink.cpp
   CSE.cpp
   GenerateRuntimeVerification.cpp
   Inliner.cpp
   LocationSnapshot.cpp
   LoopInvariantCodeMotion.cpp
   Mem2Reg.cpp
   OpStats.cpp

mlir/lib/Transforms/CommutativeOperandSort.cpp

+//===- CommutativeOperandSort.cpp - Commutative Operand Sorting Pass ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This transformation pass sorts operands of commutative operations in a
+// deterministic manner so they can be compared and removed in CSE.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/Passes.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/CommutativityUtils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_COMMUTATIVEOPERANDSORT
+#include "mlir/Transforms/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+struct CommutativeOperandSortPass
+    : public impl::CommutativeOperandSortBase<CommutativeOperandSortPass> {
+  CommutativeOperandSortPass() = default;
+
+  void runOnOperation() override {
+    // Populate pattern
+    RewritePatternSet patterns(&getContext());
+    populateCommutativityUtilsPatterns(patterns);
+
+    if (failed(
+            applyPatternsAndFoldGreedily(getOperation(), std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // end anonymous namespace
+
+// Create a Commutative operand sorting pass
+std::unique_ptr<Pass> mlir::createCommutativeOperandSortPass() {
+  return std::make_unique<CommutativeOperandSortPass>();
+}

mlir/test/Transforms/commutative-operand-sort.mlir

+// RUN: mlir-opt %s -commutative-operand-sort -cse | FileCheck %s
+
+// CHECK-LABEL: func @check_commutative_cse
+func.func @check_commutative_cse(%a : i32, %b : i32) -> i32 {
+  %0 = arith.subi %a, %b : i32
+  %1 = arith.divsi %a, %b : i32
+  // CHECK: %[[ADD1:.*]] = arith.addi %{{.*}}, %{{.*}} : i32
+  %2 = arith.addi %0, %1 : i32
+  %3 = arith.addi %1, %0 : i32
+  // CHECK-NEXT:  arith.muli %[[ADD1]], %[[ADD1]] : i32
+  %4 = arith.muli %2, %3 : i32
+  return %4 : i32
+}

mlir/test/lib/Transforms/TestCommutativityUtils.cpp

struct CommutativityUtils
 
   void runOnOperation() override {
     auto func = getOperation();
     auto *context = &getContext();
 
     RewritePatternSet patterns(context);
     populateCommutativityUtilsPatterns(patterns);
 
-    (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
+    FrozenRewritePatternSet frozenPatterns =
+        FrozenRewritePatternSet(std::move(patterns));
+
+    (void)applyPatternsAndFoldGreedily(func, frozenPatterns);
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
 void registerCommutativityUtils() { PassRegistration<CommutativityUtils>(); }
 } // namespace test
 } // namespace mlir