Revert "[VectorDistribute] Correctly find new dimensions during reduction config" (iree-org#21810)

Reverts iree-org#21797

This patch does some weird things with elementwise consumers, tiling
them serially, when they should be tiled by consumer fusion. This needs
more work before it can be landed.

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

Author: Groverkss

compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp

@@ -492,7 +492,6 @@ getVectorDistributeReductionConfig(
   return loweringConfig;
 }
 
-// TODO: Use IndexingMapInterface here instead of linalg::LinalgOp.
 static LogicalResult
 populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
                    IREE::GPU::TargetAttr target, int64_t workgroupSize,
@@ -524,28 +523,42 @@ populateConfigInfo(const llvm::SetVector<linalg::LinalgOp> &computeOps,
   // LinalgOp with only parallel dims. This is needed if the op cannot be fused
   // with a reduction or introduces new loop dimensions.
   auto shouldAttachLoweringConfig = [&](linalg::LinalgOp linalgOp) -> bool {
-    // We want to attach a lowering config to this operation if it introduces
-    // a new dimension, when going by topological order in the backward slice.
-    // The only two ways to introduce a new dimension are:
-    //
-    // 1. We have a reduction dimension.
-    if (hasReductionIterator(linalgOp)) {
-      return true;
-    }
-    // 2. There is no consumer which is a compute op (i.e., it already
-    // has some way of getting fused).
-    if (llvm::none_of(linalgOp->getUsers(), [&](Operation *user) {
+    // If the operation has a gather, we want to fuse it with the
+    // reduction.
+    if (hasExternalCapture(cast<linalg::GenericOp>(linalgOp))) {
+      return false;
+    }
+    // If some of the users are in computeOps and some are outside of
+    // computeOps; attach lowering config, since the op can't be fused.
+    if (llvm::any_of(linalgOp->getUsers(),
+                     [&](Operation *user) {
+                       auto linalgUser = dyn_cast<linalg::LinalgOp>(user);
+                       return linalgUser && computeOps.contains(linalgUser);
+                     }) &&
+        llvm::any_of(linalgOp->getUsers(), [&](Operation *user) {
           auto linalgUser = dyn_cast<linalg::LinalgOp>(user);
-          return linalgUser && computeOps.contains(linalgUser);
+          return !linalgUser;
         })) {
       return true;
     }
 
+    // If the indexing map introduces new dimensions (more inputs than results),
+    // attach a lowering config.
+    for (OpOperand *operand : linalgOp.getDpsInputOperands()) {
+      int64_t operandIdx = linalgOp.getIndexingMapIndex(operand);
+      AffineMap indexingMap = linalgOp.getIndexingMapsArray()[operandIdx];
+      if (indexingMap.getNumResults() > 0 &&
+          indexingMap.getNumInputs() > indexingMap.getNumResults()) {
+        return true;
+      }
+    }
+
     return false;
   };
 
   for (linalg::LinalgOp linalgOp : computeOps) {
-    if (shouldAttachLoweringConfig(linalgOp)) {
+    if (hasReductionIterator(linalgOp) ||
+        shouldAttachLoweringConfig(linalgOp)) {
       auto loweringConfig = getVectorDistributeReductionConfig(
           linalgOp, target, sharedWgpTiles, workgroupSize, subgroupSize,
           threadLoads);

compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp

@@ -755,6 +755,9 @@ void addGPUVectorDistributePassPipeline(OpPassManager &funcPassManager,
                                /*convertToDpsOptions=*/std::nullopt,
                                /*reorderStrategy=*/reorderStrategy);
 
+  // Some of the elementwise fusion can benefit from this pass.
+  funcPassManager.addPass(createRematerializeParallelOpsPass());
+
   funcPassManager.addPass(
       IREE::LinalgExt::createConvertAttentionToOnlineAttentionPass());
 

compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_vector_distribute_reduction_gfx942.mlir

@@ -245,7 +245,12 @@ func.func @test_multiple_stores(%arg0: !iree_tensor_ext.dispatch.tensor<readonly
 //       CHECK: func.func @test_multiple_stores
 //  CHECK-SAME:     translation_info = #[[$TRANSLATION]]
 //       CHECK:   linalg.generic
-//  CHECK-NOT:      lowering_config
+//  CHECK-SAME:      attrs =  {lowering_config = #iree_gpu.lowering_config<{
+//  CHECK-SAME:               lane_basis = {{\[}}[1, 64], [0, 1]],
+//  CHECK-SAME:               reduction = [0, 4096],
+//  CHECK-SAME:               subgroup_basis = {{\[}}[1, 16], [0, 1]],
+//  CHECK-SAME:               thread = [0, 4],
+//  CHECK-SAME:               workgroup = [1, 0]
 //       CHECK:   linalg.generic
 //  CHECK-SAME:      attrs =  {lowering_config = #iree_gpu.lowering_config<{
 //  CHECK-SAME:               lane_basis = {{\[}}[1, 64], [0, 1]],