[Triton] Fix LoopAwareCSE by removing the equivalence cache (#6894)

I made a classic blunder trying to cache context-sensitive DFS. The DFS
cache persisting from one context to another was causing miscompiles
because operations that were equal in one context are not equal in
another.

I just removed the cache for now because properly caching the DFS here
is very tricky. If compile time of this pass becomes a problem, we can
revisit it.

Author: Mogball

lib/Dialect/Triton/Transforms/LoopAwareCSE.cpp

@@ -43,7 +43,7 @@ struct LoopCSEDriver {
   bool areEqualInLoop(Value a, Value b);
 
   scf::ForOp loop;
-  ValueEquivalence equalValues;
+  SmallVector<std::pair<int, int>> argStack;
 };
 } // namespace
 
@@ -52,14 +52,15 @@ bool LoopCSEDriver::areIterArgsEqual(int i, int j) {
     return true;
   if (loop.getInitArgs()[i] != loop.getInitArgs()[j])
     return false;
+  if (llvm::is_contained(argStack, std::make_pair(i, j)))
+    return true;
   BlockArgument aArg = loop.getRegionIterArg(i);
   BlockArgument bArg = loop.getRegionIterArg(j);
   // First, assume the arguments are equal. This is how recursion is broken.
-  equalValues.setKnownEquivalence(aArg, bArg, true);
+  argStack.push_back({i, j});
   bool result =
       areEqualInLoop(loop.getYieldedValues()[i], loop.getYieldedValues()[j]);
-  // Now update the equivalence based on the actual result.
-  equalValues.setKnownEquivalence(aArg, bArg, result);
+  argStack.pop_back();
   return result;
 }
 
@@ -83,14 +84,10 @@ bool LoopCSEDriver::areEqualInLoop(Value a, Value b) {
   if (a == loop.getInductionVar() || b == loop.getInductionVar())
     return false;
 
-  if (std::optional<bool> eq = equalValues.getKnownEquivalence(a, b))
-    return *eq;
-
   if (auto aArg = dyn_cast<BlockArgument>(a)) {
     auto bArg = cast<BlockArgument>(b);
     bool result =
         areIterArgsEqual(aArg.getArgNumber() - 1, bArg.getArgNumber() - 1);
-    equalValues.setKnownEquivalence(a, b, result);
     return result;
   }
 
@@ -107,9 +104,7 @@ bool LoopCSEDriver::areEqualInLoop(Value a, Value b) {
   bool result = OperationEquivalence::isEquivalentTo(
       aDef, bDef,
       [&](Value a, Value b) { return success(areEqualInLoop(a, b)); },
-      [&](Value a, Value b) { equalValues.setKnownEquivalence(a, b, true); },
-      OperationEquivalence::IgnoreLocations);
-  equalValues.setKnownEquivalence(a, b, result);
+      /*markEquivalent=*/nullptr, OperationEquivalence::IgnoreLocations);
   return result;
 }
 

test/Triton/loop_cse.mlir

@@ -45,3 +45,28 @@ tt.func @loop_buffer_phase_args(%arg0: i32) {
   }
   tt.return
 }
+
+// CHECK-LABEL: @invalid_cache_test
+tt.func public @invalid_cache_test(%arg0: i32, %arg1: i32) -> (i32, i32) {
+  %c1_i32 = arith.constant 1 : i32
+  %c3_i32 = arith.constant 3 : i32
+  %c0_i32 = arith.constant 0 : i32
+  // CHECK: %0:4 = scf.for
+  %0:4 = scf.for %arg2 = %c0_i32 to %arg0 step %arg1 iter_args(%arg3 = %c0_i32, %arg4 = %c0_i32, %arg5 = %c0_i32, %arg6 = %c0_i32) -> (i32, i32, i32, i32)  : i32 {
+
+    %1 = arith.addi %arg5, %c1_i32 : i32
+    %2 = arith.xori %arg6, %c1_i32 : i32
+    %3 = arith.cmpi eq, %1, %c3_i32 : i32
+    %4 = arith.select %3, %2, %arg6 : i32
+    %5 = arith.select %3, %c1_i32, %1 : i32
+
+    %6 = arith.addi %arg3, %c1_i32 : i32
+    %7 = arith.xori %arg4, %c1_i32 : i32
+    %8 = arith.cmpi eq, %6, %c3_i32 : i32
+    %9 = arith.select %8, %c0_i32, %6 : i32
+    %10 = arith.select %8, %7, %arg4 : i32
+
+    scf.yield %9, %10, %5, %4 : i32, i32, i32, i32
+  }
+  tt.return %0#1, %0#3 : i32, i32
+}