Revert "[codegen] more consumer fusion (iree-org#21521)" (iree-org#21819)

This reverts commit 4d91ffb.

The above commit causes failure in compilation of llama 405B fp4 model.
Ticket to track the same iree-org#21814

Signed-off-by: Praveen G <[email protected]>
Signed-off-by: Ivan Ho <[email protected]>

Author: pravg-amd

compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.cpp

@@ -9,7 +9,6 @@
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
 #include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenInterfaces.h"
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "mlir/Analysis/TopologicalSortUtils.h"
@@ -88,51 +87,16 @@ void collectTiledAndFusedOps(Operation *rootOp,
   }
 }
 
-namespace {
-// Entry for the pseudo-priority queue of consumer fusion candidates. Contains
-// the consumer (fusableUser) that can be fused and the set of slice operations
-// in the loop to fuse into that feed the consumer.
-struct ConsumerFusionQueueEntry {
-  ConsumerFusionQueueEntry(SmallVector<Operation *> &&slices,
-                           Operation *fusableUser)
-      : slices(std::move(slices)), fusableUser(fusableUser) {}
-
-  SmallVector<Operation *> slices;
-  Operation *fusableUser;
-};
-} // namespace
-
 FailureOr<std::queue<Operation *>>
-fuseConsumersIntoForall(RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
+fuseConsumersIntoForall(RewriterBase &rewriter, Operation *tiledOp,
                         MutableArrayRef<LoopLikeOpInterface> loops,
                         std::function<bool(Operation *)> filterFn) {
   // Collect the candidate slices which can be potential consumers that can be
-  // fused. Keep them in a vector reverse-sorted by dominance: the candidate
-  // dominating others comes last (so it can be cheaply popped from the vector).
-  // The most-dominating candidate is to be fused first since not fusing it may
-  // prevent dominated candidates to be fused:
-  //
-  //     A
-  //     |
-  //     B
-  //   / |
-  //  |  D
-  //  | /
-  //  C
-  //
-  // here, B must be fused before both C and D, and D must be fused before C.
-  // Candidates are kept in a vector rather than a priority queue since we may
-  // update them as fusion happens, in particular, more slices may need to be
-  // handled. For example, fusing B with A will create a slice of B that will
-  // need to be handled correctly.
-  SmallVector<ConsumerFusionQueueEntry> candidates;
+  // fused.
+  std::queue<SmallVector<Operation *>> candidates;
   llvm::SmallDenseSet<tensor::ParallelInsertSliceOp> allCandidates;
   auto addCandidateSlices = [&candidates, &allCandidates,
                              &filterFn](Operation *fusedOp) {
-    // Dominance info recreated since op creation/movement in the fusion logic
-    // invalidates it anyway.
-    DominanceInfo dominanceInfo;
-
     for (auto *userOp : fusedOp->getResults().getUsers()) {
       auto sliceOp = dyn_cast<tensor::ParallelInsertSliceOp>(userOp);
       if (!sliceOp || allCandidates.contains(sliceOp)) {
@@ -171,38 +135,22 @@ fuseConsumersIntoForall(RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
         allCandidates.insert_range(slices);
       }
       if (!fusedSlices.empty()) {
-        ConsumerFusionQueueEntry entry(std::move(fusedSlices), fusableUser);
-
-        // Comparator that puts the dominating user last.
-        auto comp = [&](const ConsumerFusionQueueEntry &lhs,
-                        const ConsumerFusionQueueEntry &rhs) {
-          return dominanceInfo.properlyDominates(rhs.fusableUser,
-                                                 lhs.fusableUser);
-        };
-
-        // If the fusable user is already a candidate, update it with the new
-        // list of slices to handle. Otherwise, insert it into the right
-        // position based on dominance.
-        auto *it = llvm::lower_bound(candidates, entry, comp);
-        if (it != candidates.end() && it->fusableUser == fusableUser)
-          *it = std::move(entry);
-        else
-          candidates.insert(it, std::move(entry));
+        candidates.emplace(std::move(fusedSlices));
       }
     }
   };
 
-  // Add slices from all tiled ops, not only the "main" one.
-  for (Operation *tiledOp : tiledOps)
-    addCandidateSlices(tiledOp);
+  addCandidateSlices(tiledOp);
 
   std::queue<Operation *> newFusionOpportunities;
   while (!candidates.empty()) {
-    // Get the next candidate.
-    ConsumerFusionQueueEntry entry = candidates.pop_back_val();
+    // Traverse the slices in BFS fashion.
+    SmallVector<Operation *> candidateSlices = candidates.front();
+    candidates.pop();
 
     FailureOr<scf::SCFFuseConsumerOfSliceResult> fusedResult =
-        mlir::scf::tileAndFuseConsumerOfSlices(rewriter, entry.slices, loops);
+        mlir::scf::tileAndFuseConsumerOfSlices(rewriter, candidateSlices,
+                                               loops);
     if (failed(fusedResult)) {
       return failure();
     }

compiler/src/iree/compiler/Codegen/Common/TileAndFuseUtils.h

@@ -31,11 +31,12 @@ void fuseProducersOfSlices(RewriterBase &rewriter,
 void collectTiledAndFusedOps(Operation *rootOp,
                              llvm::SmallDenseSet<Operation *> &result);
 
-/// Fuse all consumers of the given `tiledOps` into the surrounding `scf.forall`
-/// unless specified otherwise by `filterFn`. Returns a list of new
-/// `tensor.extract_slice` ops with new fusion opportunities.
+/// Fuse all consumers of the given `tiledOp` into the surrounding `scf.forall`.
+/// Returns a list of new `tensor.extract_slice` ops with new fusion
+/// opportunities, as well as the new surrounding `scf.forall` (because consumer
+/// fusion replaces the loop).
 FailureOr<std::queue<Operation *>> fuseConsumersIntoForall(
-    RewriterBase &rewriter, ArrayRef<Operation *> tiledOps,
+    RewriterBase &rewriter, Operation *tiledOp,
     MutableArrayRef<LoopLikeOpInterface> loops,
     std::function<bool(Operation *)> filterFn = [](Operation *) {
       return true;

compiler/src/iree/compiler/Codegen/Common/TileDispatchUsingForall.cpp

@@ -210,18 +210,6 @@ static bool verifyComputeOpsAfterDistribution(FunctionOpInterface funcOp) {
 // Pass implementation.
 //===---------------------------------------------------------------------===//
 
-/// Returns true if any value produced by `producer` is used as an init value
-/// for the DPS `user`. Returns false if the user is not in DPS.
-static bool isUsedAsInit(Operation *producer, Operation *user) {
-  auto dpsIface = dyn_cast<DestinationStyleOpInterface>(user);
-  if (!dpsIface)
-    return false;
-  ValueRange results = producer->getResults();
-  return llvm::any_of(dpsIface.getDpsInits(), [&](Value operand) {
-    return llvm::is_contained(results, operand);
-  });
-}
-
 void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
   auto funcOp = getOperation();
   auto *context = &getContext();
@@ -244,14 +232,10 @@ void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
 
   llvm::DenseSet<Operation *> yieldReplacementsFor;
   for (auto op : tiledAndFusedOps) {
-    // Require replacement for values that are used after the main tilable op or
-    // by ops that will definitely not be fused. Note that if a value is used as
-    // an init of a DPS op, the user currently cannot be fused. Having a
-    // replacement for it would attempt fusion and fail, so avoid such cases.
+    // If tiledAndFused ops doesn't contain the user; add an replacement
+    // for that.
     if (llvm::any_of(op->getUsers(), [&](Operation *user) {
-          if (isUsedAsInit(op, user))
-            return false;
-          return dominanceInfo.properlyDominates(tilableOp, user) ||
+          return dominanceInfo.properlyDominates(tilableOp, user) &&
                  !tiledAndFusedOps.contains(user);
         })) {
       yieldReplacementsFor.insert(op);
@@ -333,18 +317,16 @@ void TileAndDistributeToWorkgroupsUsingForallOpPass::runOnOperation() {
       });
     }
     std::swap(tileAndFuseResult->loops, tilingLoops);
-
+    Operation *rootTiledOp = tileAndFuseResult->tiledAndFusedOps.front();
     FailureOr<std::queue<Operation *>> newFusionOpportunities =
-        fuseConsumersIntoForall(
-            rewriter, tileAndFuseResult->tiledAndFusedOps.getArrayRef(),
-            tilingLoops, [&tiledAndFusedOps](Operation *op) {
-              return tiledAndFusedOps.contains(op);
-            });
+        fuseConsumersIntoForall(rewriter, rootTiledOp, tilingLoops,
+                                [&tiledAndFusedOps](Operation *op) {
+                                  return tiledAndFusedOps.contains(op);
+                                });
     if (failed(newFusionOpportunities)) {
       // Continue the work if the failure is allowed.
       if (!verifyComputeOpsAfterDistribution(funcOp)) {
-        tileAndFuseResult->tiledAndFusedOps.front()->emitOpError(
-            "failed to fuse consumers");
+        rootTiledOp->emitOpError("failed to fuse consumers");
         return signalPassFailure();
       }
     } else {

compiler/src/iree/compiler/Codegen/Common/test/tile_and_distribute_workgroups_using_forall.mlir

@@ -1083,107 +1083,3 @@ func.func @infusible_pack(%arg0 : tensor<30xf32>) -> tensor<5x6xf32> {
 //       CHECK:     linalg.generic
 //       CHECK:   scf.forall.in_parallel {
 //       CHECK:   linalg.pack
-
-// -----
-
-// Adapted from layer normalization. The graph structure is as follows
-//
-//              %14
-//           /   |   \
-//         /    %15   %17
-//        |      |   / |
-//        |    [%19]   |
-//       %21     |    %22
-//        |      |     |
-//        v      v     v
-//
-// In particular, %21 and %22 are not users of the "main" tilable
-// operation but we still want them to be fused. %19, %21 and %22
-// all produce results returned from the function.
-//
-// Check that everything is fused and that there are three results
-// from the loop being produced and returned.
-//
-// CHECK-LABEL: @multi_result_consumer_fusion
-//   CHECK-NOT: linalg.generic
-//       CHECK: %[[LOOP:.+]]:3 = scf.forall (%[[I:.+]], %[[J:.+]]) in (16, 256) shared_outs(%[[OUT0:.+]] = %{{.+}}, %[[OUT1:.+]] = %{{.+}}, %[[OUT2:.+]] = %{{.+}})
-//       CHECK: %[[v14:.+]] = linalg.generic
-//       CHECK:     arith.divf
-//       CHECK:   %[[v15:.+]] = linalg.generic
-//       CHECK:     arith.subf
-//       CHECK:   %[[v17:.+]] = linalg.generic
-//       CHECK:     arith.divf
-//       CHECK:     math.rsqrt
-//       CHECK:   %[[RES0:.+]] = linalg.generic
-//       CHECK:     arith.mulf
-//       CHECK:     arith.extf
-//       CHECK:     arith.mulf
-//       CHECK:     arith.extf
-//       CHECK:     arith.addf
-//       CHECK:     arith.truncf
-//       CHECK:   %[[RES1:.+]] = linalg.generic {{.*}} ins(%[[v14]] :
-//       CHECK:     arith.truncf
-//       CHECK:   %[[RES2:.+]] = linalg.generic {{.*}} ins(%[[v17]] :
-//       CHECK:     arith.truncf
-//       CHECK:   scf.forall.in_parallel
-//       CHECK:     tensor.parallel_insert_slice %[[RES0]] into %[[OUT0]]
-//       CHECK:     tensor.parallel_insert_slice %[[RES1]] into %[[OUT1]]
-//       CHECK:     tensor.parallel_insert_slice %[[RES2]] into %[[OUT2]]
-//   CHECK-NOT: linalg.generic
-//       CHECK: return %[[LOOP]]#0, %[[LOOP]]#1, %[[LOOP]]#2
-func.func @multi_result_consumer_fusion(
-  %6: tensor<16x256x2048xbf16>,
-  %7: tensor<2048xbf16>,
-  %8: tensor<2048xbf16>,
-  %10: tensor<16x256x2048xf32>,
-  %13: tensor<16x256xf32>
-) -> (
-  tensor<16x256x2048xbf16>,
-  tensor<16x256xbf16>,
-  tensor<16x256xbf16>
-) {
-  %cst = arith.constant 0.000000e+00 : f32
-  %cst_0 = arith.constant 2.048000e+03 : f32
-  %c0 = arith.constant 0 : index
-  %9 = tensor.empty() : tensor<16x256x2048xf32>
-  %11 = tensor.empty() : tensor<16x256xf32>
-  %14 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%13 : tensor<16x256xf32>) outs(%11 : tensor<16x256xf32>) {
-  ^bb0(%in: f32, %out: f32):
-    %23 = arith.divf %in, %cst_0 : f32
-    linalg.yield %23 : f32
-  } -> tensor<16x256xf32>
-  %15 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%10, %14 : tensor<16x256x2048xf32>, tensor<16x256xf32>) outs(%9 : tensor<16x256x2048xf32>) {
-  ^bb0(%in: f32, %in_1: f32, %out: f32):
-    %23 = arith.subf %in, %in_1 : f32
-    linalg.yield %23 : f32
-  } -> tensor<16x256x2048xf32>
-  %17 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%14 : tensor<16x256xf32>) outs(%11 : tensor<16x256xf32>) {
-  ^bb0(%in: f32, %out: f32):
-    %23 = arith.divf %in, %cst_0 : f32
-    %24 = math.rsqrt %23 : f32
-    linalg.yield %24 : f32
-  } -> tensor<16x256xf32>
-  %18 = tensor.empty() : tensor<16x256x2048xbf16>
-  %19 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} ins(%15, %17, %7, %8 : tensor<16x256x2048xf32>, tensor<16x256xf32>, tensor<2048xbf16>, tensor<2048xbf16>) outs(%18 : tensor<16x256x2048xbf16>) attrs =  {lowering_config = #iree_gpu.lowering_config<{lane_basis = [[1, 1, 64], [0, 1, 2]], reduction = [0, 0, 256], subgroup_basis = [[1, 1, 1], [0, 1, 2]], thread = [0, 0, 4], workgroup = [1, 1, 0]}>} {
-  ^bb0(%in: f32, %in_1: f32, %in_2: bf16, %in_3: bf16, %out: bf16):
-    %23 = arith.mulf %in, %in_1 : f32
-    %24 = arith.extf %in_2 : bf16 to f32
-    %25 = arith.mulf %23, %24 : f32
-    %26 = arith.extf %in_3 : bf16 to f32
-    %27 = arith.addf %25, %26 : f32
-    %28 = arith.truncf %27 : f32 to bf16
-    linalg.yield %28 : bf16
-  } -> tensor<16x256x2048xbf16>
-  %20 = tensor.empty() : tensor<16x256xbf16>
-  %21 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%14 : tensor<16x256xf32>) outs(%20 : tensor<16x256xbf16>) {
-  ^bb0(%in: f32, %out: bf16):
-    %23 = arith.truncf %in : f32 to bf16
-    linalg.yield %23 : bf16
-  } -> tensor<16x256xbf16>
-  %22 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%17 : tensor<16x256xf32>) outs(%20 : tensor<16x256xbf16>) {
-  ^bb0(%in: f32, %out: bf16):
-    %23 = arith.truncf %in : f32 to bf16
-    linalg.yield %23 : bf16
-  } -> tensor<16x256xbf16>
-  return %19, %21, %22 : tensor<16x256x2048xbf16>, tensor<16x256xbf16>, tensor<16x256xbf16>
-}