[MLIR] Allow dialects to disable CSE for certain operations

[skatrak]

Summary

This patch adds the DialectCSEInterface, which dialects can implement and register to prevent the common sub-expression elimination (CSE) pass from modifying regions of certain operations.

The result is that these operations would be treated by CSE as if they were IsolatedFromAbove, but without the restrictions that come with that trait.

[jsjodin]

I'm not sure if we want to add an interface for every transform. Is there a way to have a generic solution to this, like a new trait or more general interface? It would be nice to avoid the case where a new transform will have to touch all the dialects, or a new op will require implementing/modifying a bunch of interfaces expressing basically the same thing.

[kiranchandramohan]

I think this requires an RFC in MLIR discourse to get some feedback from the experts.

[skatrak]

Thank you for your comments @jsjodin and @kiranchandramohan , I created an RFC here to try to get consensus on what the best solution should be.

[skatrak]

I'll close this patch now, since the agreed approach has been to try using the IsolatedFromAbove trait and not to add an interface modifying the CSE pass behavior.

clang/docs/tools/clang-formatted-files.txt

 mlir/include/mlir/ExecutionEngine/MemRefUtils.h
 mlir/include/mlir/ExecutionEngine/OptUtils.h
 mlir/include/mlir/ExecutionEngine/RunnerUtils.h
 mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h
 mlir/include/mlir/Interfaces/CallInterfaces.h
 mlir/include/mlir/Interfaces/CastInterfaces.h
 mlir/include/mlir/Interfaces/ControlFlowInterfaces.h
 mlir/include/mlir/Interfaces/CopyOpInterface.h
+mlir/include/mlir/Interfaces/CSEInterfaces.h
 mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
 mlir/include/mlir/Interfaces/DecodeAttributesInterfaces.h
 mlir/include/mlir/Interfaces/DerivedAttributeOpInterface.h
 mlir/include/mlir/Interfaces/FoldInterfaces.h
 mlir/include/mlir/Interfaces/InferTypeOpInterface.h
 mlir/include/mlir/Interfaces/LoopLikeInterface.h
 mlir/include/mlir/Interfaces/TilingInterface.h
 mlir/include/mlir/Interfaces/VectorInterfaces.h

mlir/include/mlir/Interfaces/CSEInterfaces.h

+//===- CSEInterfaces.h ------------------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_INTERFACES_CSEINTERFACES_H_
+#define MLIR_INTERFACES_CSEINTERFACES_H_
+
+#include "mlir/IR/DialectInterface.h"
+
+namespace mlir {
+class Operation;
+
+/// Define an interface to allow for dialects to control specific aspects of
+/// common subexpression elimination behavior for operations they define.
+class DialectCSEInterface : public DialectInterface::Base<DialectCSEInterface> {
+public:
+  DialectCSEInterface(Dialect *dialect) : Base(dialect) {}
+
+  /// Registered hook to check if an operation that is *not* isolated from
+  /// above, should allow common subexpressions to be extracted out of its
+  /// regions.
+  virtual bool subexpressionExtractionAllowed(Operation *op) const {
+    return true;
+  }
+};
+
+} // namespace mlir
+
+#endif // MLIR_INTERFACES_CSEINTERFACES_H_

mlir/lib/Transforms/CSE.cpp

 // algorithm on operations within a region.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/CSE.h"
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/CSEInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseMapInfo.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/RecyclingAllocator.h"

 };
 } // namespace
 
 namespace {
 /// Simple common sub-expression elimination.
 class CSEDriver {
 public:
   CSEDriver(RewriterBase &rewriter, DominanceInfo *domInfo)
-      : rewriter(rewriter), domInfo(domInfo) {}
+      : rewriter(rewriter), domInfo(domInfo),
+        interfaces(rewriter.getContext()) {}
 
   /// Simplify all operations within the given op.
   void simplify(Operation *op, bool *changed = nullptr);
 
   int64_t getNumCSE() const { return numCSE; }
   int64_t getNumDCE() const { return numDCE; }
 
 private:
private:
   /// A rewriter for modifying the IR.
   RewriterBase &rewriter;
 
   /// Operations marked as dead and to be erased.
   std::vector<Operation *> opsToErase;
   DominanceInfo *domInfo = nullptr;
   MemEffectsCache memEffectsCache;
 
+  /// CSE interfaces in the present context that can modify CSE behavior.
+  DialectInterfaceCollection<DialectCSEInterface> interfaces;
+
   // Various statistics.
   int64_t numCSE = 0;
   int64_t numDCE = 0;
 };
 } // namespace
 
 void CSEDriver::replaceUsesAndDelete(ScopedMapTy &knownValues, Operation *op,
                                      Operation *existing,

 void CSEDriver::simplifyBlock(ScopedMapTy &knownValues, Block *bb,
                               bool hasSSADominance) {
   for (auto &op : *bb) {
     // Most operations don't have regions, so fast path that case.
     if (op.getNumRegions() != 0) {
       // If this operation is isolated above, we can't process nested regions
       // with the given 'knownValues' map. This would cause the insertion of
       // implicit captures in explicit capture only regions.
-      if (op.mightHaveTrait<OpTrait::IsIsolatedFromAbove>()) {
+      // Also, avoid capturing known values from parent regions if this behavior
+      // is explicitly disabled for the given operation.
+      const DialectCSEInterface *cseInterface = interfaces.getInterfaceFor(&op);
+      if (op.mightHaveTrait<OpTrait::IsIsolatedFromAbove>() ||
+          LLVM_UNLIKELY(cseInterface &&
+                        !cseInterface->subexpressionExtractionAllowed(&op))) {
         ScopedMapTy nestedKnownValues;
         for (auto &region : op.getRegions())
           simplifyRegion(nestedKnownValues, region);
       } else {
         // Otherwise, process nested regions normally.
         for (auto &region : op.getRegions())
           simplifyRegion(knownValues, region);
       }

mlir/test/Transforms/cse.mlir

func.func @cse_recursive_effects_failure() -> (i32, i32, i32) {
   }
 
   // %2 can not be be removed because of the write
   // CHECK-NEXT: %[[READ_VALUE2:.*]] = "test.op_with_memread"() : () -> i32
   // CHECK-NEXT: return %[[READ_VALUE]], %[[READ_VALUE2]], %[[IF]] : i32, i32, i32
   %2 = "test.op_with_memread"() : () -> (i32)
   return %0, %2, %1 : i32, i32, i32
 }
+
+// CHECK-LABEL: @no_cse_across_disabled_op
+func.func @no_cse_across_disabled_op() -> (i32) {
+  // CHECK-NEXT: %[[CONST1:.+]] = arith.constant 1 : i32
+  %0 = arith.constant 1 : i32
+
+  // CHECK-NEXT: test.no_cse_one_region_op
+  "test.no_cse_one_region_op"() ({
+    %1 = arith.constant 1 : i32
+    %2 = arith.addi %1, %1 : i32
+    "foo.yield"(%2) : (i32) -> ()
+
+    // CHECK-NEXT: %[[CONST2:.+]] = arith.constant 1 : i32
+    // CHECK-NEXT: %[[SUM:.+]] = arith.addi %[[CONST2]], %[[CONST2]] : i32
+    // CHECK-NEXT: "foo.yield"(%[[SUM]]) : (i32) -> ()
+  }) : () -> ()
+
+  // CHECK: return %[[CONST1]] : i32
+  return %0 : i32
+}

mlir/test/lib/Dialect/Test/TestDialectInterfaces.cpp

 //===- TestDialectInterfaces.cpp - Test dialect interface definitions -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
+#include "mlir/Interfaces/CSEInterfaces.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 
 using namespace mlir;
 using namespace test;
 
 //===----------------------------------------------------------------------===//
struct TestDialectFoldInterface : public DialectFoldInterface {
   /// operation that is *not* isolated from above, should be used when
   /// materializing constants.
   bool shouldMaterializeInto(Region *region) const final {
     // If this is a one region operation, then insert into it.
     return isa<OneRegionOp>(region->getParentOp());
   }
 };
 
+struct TestDialectCSEInterface : public DialectCSEInterface {
+  using DialectCSEInterface::DialectCSEInterface;
+
+  bool subexpressionExtractionAllowed(Operation *op) const final {
+    // Don't allow extracting common subexpressions from the region of these
+    // operations.
+    return !isa<NoCSEOneRegionOp>(op);
+  }
+};
+
 /// This class defines the interface for handling inlining with standard
 /// operations.
 struct TestInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//

 };
 
 } // namespace
 
 void TestDialect::registerInterfaces() {
   auto &blobInterface = addInterface<TestResourceBlobManagerInterface>();
   addInterface<TestOpAsmInterface>(blobInterface);
 
-  addInterfaces<TestDialectFoldInterface, TestInlinerInterface,
-                TestReductionPatternInterface, TestBytecodeDialectInterface>();
+  addInterfaces<TestDialectFoldInterface, TestDialectCSEInterface,
+                TestInlinerInterface, TestReductionPatternInterface,
+                TestBytecodeDialectInterface>();
 }

mlir/test/lib/Dialect/Test/TestOps.td

def TestCSEOfSingleBlockOp : TEST_Op<"cse_of_single_block_op",
   let results = (outs Variadic<AnyType>:$outputs);
   let regions = (region SizedRegion<1>:$region);
   let assemblyFormat = [{
     attr-dict `inputs` `(` $inputs `)`
     $region `:` type($inputs)  `->` type($outputs)
   }];
 }
 
+def NoCSEOneRegionOp : TEST_Op<"no_cse_one_region_op", []> {
+  let regions = (region AnyRegion);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Ops to upgrade base on the dialect versions
 //===----------------------------------------------------------------------===//
 
 def TestVersionedOpA : TEST_Op<"versionedA"> {
   // A previous version of the dialect (let's say 1.*) supported an attribute
   // named "dimensions":
   // let arguments = (ins