[mlir][Transforms] teach CSE about recursive memory effects

[tblah]

Summary

Add support for reasoning about operations with recursive memory effects
to CSE. The recursive effects are gathered by a helper function. I
decided to allow returning duplicates from the helper function because
there's no benefit to spending the computation time to remove them in
the existing use case.

[aartbik]

lg for the sparse check changes

[kiranchandramohan]

I guess one reason for not doing this previously would be the time taken to recursively traverse operations. But as long as an operation is not visited more than once, i guess it is OK.

mlir/lib/Transforms/CSE.cpp

207

Would an alternative implementation be to recursively perform the check during op.hasEffect if the operation has recursiveEffects?

[kiranchandramohan]

mlir/lib/Interfaces/SideEffectInterfaces.cpp

198–200

Nit: spell the types here and above.

206–211

Nit: Braces here and above.

[vzakhari]

In D156805#4554310, @kiranchandramohan wrote:

I guess one reason for not doing this previously would be the time taken to recursively traverse operations. But as long as an operation is not visited more than once, i guess it is OK.

Actually, the operations are going to be visited multiple times because the CSE pass will apply to every operation. So we will be computing the effects on every runOnOperation invocation, and some of the computations will be redundant. I think we need to make the effects cache "global", but does it mean we have to make the pass a module pass?

mlir/lib/Interfaces/SideEffectInterfaces.cpp

210

typo: meemory

mlir/lib/Transforms/CSE.cpp

207

Not yours, but you may consider fixing the typo: assumes -> assume

[tblah]

In D156805#4554528, @vzakhari wrote:

In D156805#4554310, @kiranchandramohan wrote:

I guess one reason for not doing this previously would be the time taken to recursively traverse operations. But as long as an operation is not visited more than once, i guess it is OK.

Actually, the operations are going to be visited multiple times because the CSE pass will apply to every operation. So we will be computing the effects on every runOnOperation invocation, and some of the computations will be redundant. I think we need to make the effects cache "global", but does it mean we have to make the pass a module pass?

Thanks for taking a look. Yes you are right that there is a lot of opportunity for caching here. The existing implementation of wouldOpBeTriviallyDeadImpl works the same way and is already called on every operation (via isOpTriviallyDead). Therefore

1. There is even more opportunity for caching between these calls and getEffectsRecursively
2. The big-O complexity of the pass remains of the same order as without this patch. We just do the same expensive uncached lookup twice now

Would it be okay to keep the patch scope as it is and then address performance (of all functions in SideEffectInterfaces.cpp) separately, if it turns out to be a problem for real workloads?

mlir/lib/Transforms/CSE.cpp

207

Operations with recursive effects tend not to implement the memory effect interface so nextOpMemEffects would be null (recursive memory effects is a different interface).

Anyway, I think you mean, why not fix the normal ways of getting the side effects of an operation so that they always check recursively? I decided not to do that in this patch because it would have a far broader impact on other MLIR transformations.

[tblah]

Updating to fix nits. Thanks for review.

[vzakhari]

In D156805#4554769, @tblah wrote:

Would it be okay to keep the patch scope as it is and then address performance (of all functions in SideEffectInterfaces.cpp) separately, if it turns out to be a problem for real workloads?

It is fine by me. Please give it a couple of days for @rriddle and others to comment.

[unterumarmung]

mlir/lib/Transforms/CSE.cpp

216

Optional nit

[tblah]

Fix nit

[kiranchandramohan]

Thanks. LGTM. Please wait a day before you submit in case other reviewers have comments.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

338

mlir/lib/Transforms/CSE.cpp

206–207

Nit: I think this comment has to be expanded to say about the RecursiveMemoryEffectCase.

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

 /// example, it may implement `NoMemoryEffect` in ODS. Alternatively, if the
 /// operation has the `HasRecursiveMemoryEffects` trait, then it is free of
 /// memory effects if all of its nested operations are free of memory effects.
 ///
 /// If the operation has both, then it is free of memory effects if both
 /// conditions are satisfied.
 bool isMemoryEffectFree(Operation *op);
 
+/// Returns the side effects of an operation. If the operation has
+/// RecursiveMemoryEffects, include all side effects of child operations.
+///
+/// std::nullopt indicates that an option did not have a memory effect interface
+/// and so no result could be obtained. An empty vector indicates that there
+/// were no memory effects found (but every operation implemented the memory
+/// effect interface or has RecursiveMemoryEffects). If the vector contains
+/// multiple effects, these effects may be duplicates.
+std::optional<llvm::SmallVector<MemoryEffects::EffectInstance>>
+getEffectsRecursively(Operation *rootOp);
+
 /// Returns true if the given operation is speculatable, i.e. has no undefined
 /// behavior or other side effects.
 ///
 /// An operation can indicate that it is speculatable by implementing the
 /// getSpeculatability hook in the ConditionallySpeculatable op interface.
 bool isSpeculatable(Operation *op);
 
 /// Returns true if the given operation is pure, i.e., is speculatable that does

mlir/lib/Interfaces/SideEffectInterfaces.cpp

bool mlir::isMemoryEffectFree(Operation *op) {
   // free.
   for (Region &region : op->getRegions())
     for (Operation &op : region.getOps())
       if (!isMemoryEffectFree(&op))
         return false;
   return true;
 }
 
+// the returned vector may contain duplicate effects
+std::optional<llvm::SmallVector<MemoryEffects::EffectInstance>>
+mlir::getEffectsRecursively(Operation *rootOp) {
+  SmallVector<MemoryEffects::EffectInstance> effects;
+  SmallVector<Operation *> effectingOps(1, rootOp);
+  while (!effectingOps.empty()) {
+    Operation *op = effectingOps.pop_back_val();
+
+    // If the operation has recursive effects, push all of the nested
+    // operations on to the stack to consider.
+    bool hasRecursiveEffects =
+        op->hasTrait<OpTrait::HasRecursiveMemoryEffects>();
+    if (hasRecursiveEffects) {
+      for (Region &region : op->getRegions()) {
+        for (Block &block : region) {
+          for (Operation &nestedOp : block) {
+            effectingOps.push_back(&nestedOp);
+          }
+        }
+      }
+    }
+
+    if (auto effectInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+      effectInterface.getEffects(effects);
+    } else if (!hasRecursiveEffects) {
+      // the operation does not have recursive memory effects or implement
+      // the memory effect op interface. Its effects are unknown.
+      return std::nullopt;
+    }
+  }
+  return effects;
+}
+
 bool mlir::isSpeculatable(Operation *op) {
   auto conditionallySpeculatable = dyn_cast<ConditionallySpeculatable>(op);
   if (!conditionallySpeculatable)
     return false;
 
   switch (conditionallySpeculatable.getSpeculatability()) {
   case Speculation::RecursivelySpeculatable:
     for (Region &region : op->getRegions()) {

mlir/lib/Transforms/CSE.cpp

if (memEffectsCachePair.second == nullptr) {
       nextOp = memEffectsCachePair.first;
     } else {
       // MemoryEffects::Write has been detected before so there is no need to
       // check further.
       return true;
     }
   }
   while (nextOp && nextOp != toOp) {
-    auto nextOpMemEffects = dyn_cast<MemoryEffectOpInterface>(nextOp);
+    std::optional<SmallVector<MemoryEffects::EffectInstance>> effects =
+        getEffectsRecursively(nextOp);
+    if (!effects) {
       // TODO: Do we need to handle other effects generically?
       // If the operation does not implement the MemoryEffectOpInterface we
-    // conservatively assumes it writes.
-    if ((nextOpMemEffects &&
-         nextOpMemEffects.hasEffect<MemoryEffects::Write>()) ||
-        !nextOpMemEffects) {
+      // conservatively assume it writes.
       result.first->second =
           std::make_pair(nextOp, MemoryEffects::Write::get());
       return true;
     }
+
+    for (const MemoryEffects::EffectInstance &effect : *effects) {
+      if (isa<MemoryEffects::Write>(effect.getEffect())) {
+        result.first->second = {nextOp, MemoryEffects::Write::get()};
+        return true;
+      }
+    }
     nextOp = nextOp->getNextNode();
   }
   result.first->second = std::make_pair(toOp, nullptr);
   return false;
 }
 
 /// Attempt to eliminate a redundant operation.
 LogicalResult CSEDriver::simplifyOperation(ScopedMapTy &knownValues,

mlir/test/Dialect/SparseTensor/buffer_rewriting.mlir

 // CHECK:               %[[VAL_143:.*]] = arith.cmpi ult, %[[VAL_141]], %[[VAL_142]]
 // CHECK:               scf.condition(%[[VAL_143]]) %[[VAL_140]]
 // CHECK:             } do {
 // CHECK:             ^bb0(%[[VAL_144:.*]]: index):
 // CHECK:               %[[VAL_145:.*]] = arith.addi %[[VAL_144]], %[[VAL_7]]
 // CHECK:               scf.yield %[[VAL_145]]
 // CHECK:             }
 // CHECK:             %[[VAL_146:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_147:.*]]]
-// CHECK:             %[[VAL_148:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_127]]]
-// CHECK:             %[[VAL_149:.*]] = arith.cmpi eq, %[[VAL_146]], %[[VAL_148]]
+// CHECK:             %[[VAL_149:.*]] = arith.cmpi eq, %[[VAL_146]], %[[VAL_137]]
 // CHECK:             %[[VAL_150:.*]] = arith.cmpi ult, %[[VAL_136]], %[[VAL_147]]
 // CHECK:             %[[VAL_151:.*]]:3 = scf.if %[[VAL_150]]
 // CHECK:               %[[VAL_152:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_136]]]
 // CHECK:               %[[VAL_153:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_147]]]
 // CHECK:               memref.store %[[VAL_153]], %[[VAL_2]]{{\[}}%[[VAL_136]]]
 // CHECK:               memref.store %[[VAL_152]], %[[VAL_2]]{{\[}}%[[VAL_147]]]
 // CHECK:               %[[VAL_154:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_136]]]
 // CHECK:               %[[VAL_155:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_147]]]

 // verify correctness of the generated code.
 //
 // CHECK-DAG:     func.func private @_sparse_shift_down_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>, %arg5: index) {
 // CHECK-DAG:     func.func private @_sparse_heap_sort_2_index_coo_1_f32_i32(%arg0: index, %arg1: index, %arg2: memref<?xindex>, %arg3: memref<?xf32>, %arg4: memref<?xi32>) {
 // CHECK-LABEL:   func.func @sparse_sort_coo_heap
 func.func @sparse_sort_coo_heap(%arg0: index, %arg1: memref<100xindex>, %arg2: memref<?xf32>, %arg3: memref<10xi32>) -> (memref<100xindex>, memref<?xf32>, memref<10xi32>) {
   sparse_tensor.sort_coo heap_sort %arg0, %arg1 jointly %arg2, %arg3 {nx = 2 : index, ny = 1: index} : memref<100xindex> jointly memref<?xf32>, memref<10xi32>
   return %arg1, %arg2, %arg3 : memref<100xindex>, memref<?xf32>, memref<10xi32>
-}
+}
 

mlir/test/Dialect/SparseTensor/sparse_conv_2d_slice_based.mlir

 // CHECK:                   %[[VAL_128:.*]] = memref.load %[[VAL_17]]{{\[}}%[[VAL_124]]] : memref<?xi32>
 // CHECK:                   %[[VAL_129:.*]] = tensor.extract %[[VAL_1]]{{\[}}%[[VAL_109]], %[[VAL_127]]] : tensor<3x3xi32>
 // CHECK:                   %[[VAL_130:.*]] = arith.muli %[[VAL_128]], %[[VAL_129]] : i32
 // CHECK:                   %[[VAL_131:.*]] = arith.addi %[[VAL_125]], %[[VAL_130]] : i32
 // CHECK:                   %[[VAL_132:.*]] = arith.addi %[[VAL_124]], %[[VAL_7]] : index
 // CHECK:                   scf.yield %[[VAL_132]], %[[VAL_131]] : index, i32
 // CHECK:                 }
 // CHECK:                 %[[VAL_133:.*]] = arith.addi %[[VAL_105]], %[[VAL_7]] : index
-// CHECK:                 %[[VAL_134:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_7]]] : memref<11xindex>
-// CHECK:                 %[[VAL_135:.*]] = arith.addi %[[VAL_134]], %[[VAL_5]] : index
-// CHECK:                 memref.store %[[VAL_135]], %[[VAL_18]]{{\[}}%[[VAL_7]]] : memref<11xindex>
+// CHECK:                 memref.store %[[VAL_112]], %[[VAL_18]]{{\[}}%[[VAL_7]]] : memref<11xindex>
 // CHECK:                 scf.yield %[[VAL_133]], %[[VAL_136:.*]]#1, %[[VAL_2]] : index, i32, i1
 // CHECK:               }
 // CHECK:               %[[VAL_137:.*]] = scf.if %[[VAL_138:.*]]#2 -> (tensor<6x6xi32, #sparse_tensor.encoding<{{.*}}>>) {
 // CHECK:                 %[[VAL_139:.*]] = sparse_tensor.insert %[[VAL_138]]#1 into %[[VAL_89]]{{\[}}%[[VAL_33]], %[[VAL_88]]] : tensor<6x6xi32, #sparse_tensor.encoding<{{.*}}>>
 // CHECK:                 scf.yield %[[VAL_139]] : tensor<6x6xi32, #sparse_tensor.encoding<{{.*}}>>
 // CHECK:               } else {
 // CHECK:                 scf.yield %[[VAL_89]] : tensor<6x6xi32, #sparse_tensor.encoding<{{.*}}>>
 // CHECK:               }

mlir/test/Transforms/cse.mlir

 //       CHECK:   %[[if:.*]] = scf.if {{.*}} {
 //       CHECK:     tensor.empty
 //       CHECK:     scf.yield
 //       CHECK:   } else {
 //       CHECK:     scf.yield
 //       CHECK:   }
 //   CHECK-NOT:   scf.if
 //       CHECK:   return %[[if]], %[[if]]
+
+// CHECK-LABEL: @cse_recursive_effects_success
+func.func @cse_recursive_effects_success() -> (i32, i32, i32) {
+  // CHECK-NEXT: %[[READ_VALUE:.*]] = "test.op_with_memread"() : () -> i32
+  %0 = "test.op_with_memread"() : () -> (i32)
+
+  // do something with recursive effects, containing no side effects
+  %true = arith.constant true
+  // CHECK-NEXT: %[[TRUE:.+]] = arith.constant true
+  // CHECK-NEXT: %[[IF:.+]] = scf.if %[[TRUE]] -> (i32) {
+  %1 = scf.if %true -> (i32) {
+    %c42 = arith.constant 42 : i32
+    scf.yield %c42 : i32
+    // CHECK-NEXT: %[[C42:.+]] = arith.constant 42 : i32
+    // CHECK-NEXT: scf.yield %[[C42]]
+    // CHECK-NEXT: } else {
+  } else {
+    %c24 = arith.constant 24 : i32
+    scf.yield %c24 : i32
+    // CHECK-NEXT: %[[C24:.+]] = arith.constant 24 : i32
+    // CHECK-NEXT: scf.yield %[[C24]]
+    // CHECK-NEXT: }
+  }
+
+  // %2 can be removed
+  // CHECK-NEXT: return %[[READ_VALUE]], %[[READ_VALUE]], %[[IF]] : i32, i32, i32
+  %2 = "test.op_with_memread"() : () -> (i32)
+  return %0, %2, %1 : i32, i32, i32
+}
+
+// CHECK-LABEL: @cse_recursive_effects_failure
+func.func @cse_recursive_effects_failure() -> (i32, i32, i32) {
+  // CHECK-NEXT: %[[READ_VALUE:.*]] = "test.op_with_memread"() : () -> i32
+  %0 = "test.op_with_memread"() : () -> (i32)
+
+  // do something with recursive effects, containing a write effect
+  %true = arith.constant true
+  // CHECK-NEXT: %[[TRUE:.+]] = arith.constant true
+  // CHECK-NEXT: %[[IF:.+]] = scf.if %[[TRUE]] -> (i32) {
+  %1 = scf.if %true -> (i32) {
+    "test.op_with_memwrite"() : () -> ()
+    // CHECK-NEXT: "test.op_with_memwrite"() : () -> ()
+    %c42 = arith.constant 42 : i32
+    scf.yield %c42 : i32
+    // CHECK-NEXT: %[[C42:.+]] = arith.constant 42 : i32
+    // CHECK-NEXT: scf.yield %[[C42]]
+    // CHECK-NEXT: } else {
+  } else {
+    %c24 = arith.constant 24 : i32
+    scf.yield %c24 : i32
+    // CHECK-NEXT: %[[C24:.+]] = arith.constant 24 : i32
+    // CHECK-NEXT: scf.yield %[[C24]]
+    // CHECK-NEXT: }
+  }
+
+  // %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
+}