[MLIR] Clarify (test-scf-)parallel-loop-collapsing

[tpopp]

Summary

1. parallel-loop-collapsing is renamed to test-scf-parallel-loop-collapsing.
2. The pass adds various checks to provide error messages instead of hitting assert failures.
3. Testing is added to verify these error messages

This is roughly an NFC. The name changes, but all checked behavior
previously would have resulted in an assertion failure. Almost no new
support is added, so this pass is still limited in scope to testing the
transform behaves correctly with input arguments that perfectly match
the ParallelLoop's iterator arg set. The one new piece of functionality
is that invalid operations will now be skipped with an error messages
instead of producing an assertion failure, so the pass can be used with
expected failures for pieces of the IR not cared about with a specific
RUN command.

[mehdi_amini]

mlir/include/mlir/Dialect/SCF/Transforms/Passes.td

52–53

Can you add a description indicating that this is a pass intended for unit testing the collapseParallelLoops utility? The summary does not make it clear.

mlir/include/mlir/Dialect/SCF/Transforms/Passes.h

 std::unique_ptr<Pass> createForLoopPeelingPass();
 
 /// Creates a pass that canonicalizes affine.min and affine.max operations
 /// inside of scf.for loops with known lower and upper bounds.
 std::unique_ptr<Pass> createSCFForLoopCanonicalizationPass();
 
 /// Creates a pass that transforms a single ParallelLoop over N induction
 /// variables into another ParallelLoop over less than N induction variables.
-std::unique_ptr<Pass> createParallelLoopCollapsingPass();
+std::unique_ptr<Pass> createTestSCFParallelLoopCollapsingPass();
 
 /// Creates a loop fusion pass which fuses parallel loops.
 std::unique_ptr<Pass> createParallelLoopFusionPass();
 
 /// Creates a pass that specializes parallel loop for unrolling and
 /// vectorization.
 std::unique_ptr<Pass> createParallelLoopSpecializationPass();
 

mlir/include/mlir/Dialect/SCF/Transforms/Passes.td

def SCFForLoopSpecialization : Pass<"scf-for-loop-specialization"> {
   let summary = "Specialize `for` loops for vectorization";
   let constructor = "mlir::createForLoopSpecializationPass()";
 }
 
 def SCFParallelLoopFusion : Pass<"scf-parallel-loop-fusion"> {
   let summary = "Fuse adjacent parallel loops";
   let constructor = "mlir::createParallelLoopFusionPass()";
 }
 
-def SCFParallelLoopCollapsing : Pass<"scf-parallel-loop-collapsing"> {
-  let summary = "Collapse parallel loops to use less induction variables";
-  let constructor = "mlir::createParallelLoopCollapsingPass()";
+def TestSCFParallelLoopCollapsing : Pass<"test-scf-parallel-loop-collapsing"> {
+  let summary = "Test parallel loops collapsing transformation";
+  let constructor = "mlir::createTestSCFParallelLoopCollapsingPass()";
+  let description = [{
+      This pass is purely for testing the scf::collapseParallelLoops
+      transformation. The transformation does not have opinions on how a
+      parallel loop should be collapsed, so this pass is structured for the
+      common case on GPUs of collapsing to a 3d parallel loop. 3 lists can be
+      provided to collapsed-indices-{0,1,2} to represent how the loop should be
+      collapsed and must reference evrey iterator in the original parallel loop.
+
+    ```mlir
+    # Before:
+    scf.parallel (%arg0, %arg1)
+                 = (%c0, %c0) to (%c2, %c2) step (%c1, %c1) {
+      "test.sink"(%5, %3) : (index, index) -> ()
+      scf.yield
+    }
+
+    # After:
+    scf.parallel (%arg0) = (%c0) to (%c4) step (%c1) {
+      %0 = arith.remsi %arg0, %c2 : index
+      %1 = arith.divsi %arg0, %c2 : index
+      %2 = arith.muli %0, %c7 : index
+      %3 = arith.addi %2, %c3 : index
+      %4 = arith.muli %1, %c7 : index
+      %5 = arith.addi %4, %c3 : index
+      "test.sink"(%5, %3) : (index, index) -> ()
+    }
+    ```
+  }];
+
   let options = [
     ListOption<"clCollapsedIndices0", "collapsed-indices-0", "unsigned",
                "Which loop indices to combine 0th loop index">,
     ListOption<"clCollapsedIndices1", "collapsed-indices-1", "unsigned",
                "Which loop indices to combine into the position 1 loop index">,
     ListOption<"clCollapsedIndices2", "collapsed-indices-2", "unsigned",
                "Which loop indices to combine into the position 2 loop index">,
   ];

mlir/lib/Dialect/SCF/Transforms/ParallelLoopCollapsing.cpp

 //===- ParallelLoopCollapsing.cpp - Pass collapsing parallel loop indices -===//
 //
 // 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 "mlir/Dialect/SCF/Transforms/Passes.h"
 
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
-#define GEN_PASS_DEF_SCFPARALLELLOOPCOLLAPSING
+#define GEN_PASS_DEF_TESTSCFPARALLELLOOPCOLLAPSING
 #include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
 } // namespace mlir
 
 #define DEBUG_TYPE "parallel-loop-collapsing"
 
 using namespace mlir;
 
 namespace {
-struct ParallelLoopCollapsing
-    : public impl::SCFParallelLoopCollapsingBase<ParallelLoopCollapsing> {
+struct TestSCFParallelLoopCollapsing
+    : public impl::TestSCFParallelLoopCollapsingBase<
+          TestSCFParallelLoopCollapsing> {
   void runOnOperation() override {
     Operation *module = getOperation();
 
-    module->walk([&](scf::ParallelOp op) {
-      // The common case for GPU dialect will be simplifying the ParallelOp to 3
-      // arguments, so we do that here to simplify things.
-      llvm::SmallVector<std::vector<unsigned>, 3> combinedLoops;
-      if (!clCollapsedIndices0.empty())
-        combinedLoops.push_back(clCollapsedIndices0);
-      if (!clCollapsedIndices1.empty())
-        combinedLoops.push_back(clCollapsedIndices1);
-      if (!clCollapsedIndices2.empty())
-        combinedLoops.push_back(clCollapsedIndices2);
+
+    // Gather the input args into the format required by
+    // `collapseParallelLoops`.
+    if (!clCollapsedIndices1.empty()) {
+      if (clCollapsedIndices0.empty()) {
+        llvm::errs()
+            << "collapsed-indices-1 specified but not collapsed-indices-0";
+        signalPassFailure();
+        return;
+      }
+    }
+    if (!clCollapsedIndices2.empty()) {
+      if (clCollapsedIndices1.empty()) {
+        llvm::errs()
+            << "collapsed-indices-2 specified but not collapsed-indices-1";
+        signalPassFailure();
+        return;
+      }
+    }
+
+    if (combinedLoops.empty()) {
+      llvm::errs() << "No collapsed-indices were specified. This pass is only "
+                      "for testing and does not automatically collapse all "
+                      "parallel loops or similar.";
+      signalPassFailure();
+      return;
+    }
+
+    // Confirm that the specified loops are [0,N) by testing that N values exist
+    // with the maximum value being N-1.
+    llvm::SmallSet<unsigned, 8> flattenedCombinedLoops;
+    unsigned maxCollapsedIndex = 0;
+    for (auto &loops : combinedLoops) {
+      for (auto &loop : loops) {
+        flattenedCombinedLoops.insert(loop);
+        maxCollapsedIndex = std::max(maxCollapsedIndex, loop);
+      }
+    }
+
+    if (maxCollapsedIndex != flattenedCombinedLoops.size() - 1 ||
+        !flattenedCombinedLoops.contains(maxCollapsedIndex)) {
+      llvm::errs()
+          << "collapsed-indices arguments must include all values [0,N).";
+      signalPassFailure();
+      return;
+    }
+
+    // Only apply the transformation on parallel loops where the specified
+    // transformation is valid, but do NOT early abort in the case of invalid
+    // loops.
+    module->walk([&](scf::ParallelOp op) {
+      if (flattenedCombinedLoops.size() != op.getNumLoops()) {
+        op.emitOpError("has ")
+            << op.getNumLoops()
+            << " iter args while this limited functionality testing pass was "
+               "configured only for loops with exactly "
+            << flattenedCombinedLoops.size() << " iter args.";
+        return;
+      }
       collapseParallelLoops(op, combinedLoops);
     });
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::createParallelLoopCollapsingPass() {
-  return std::make_unique<ParallelLoopCollapsing>();
+std::unique_ptr<Pass> mlir::createTestSCFParallelLoopCollapsingPass() {
+  return std::make_unique<TestSCFParallelLoopCollapsing>();
 }

mlir/test/Transforms/invalid-parallel-loop-collapsing.mlir

+// First test various sets of invalid arguments
+// RUN: not mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing))' 2>&1 | FileCheck %s --check-prefix=CL0
+// CL0: No collapsed-indices were specified. This pass is only for testing and does not automatically collapse all parallel loops or similar
+
+// RUN: not mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-1=1}))' 2>&1 | FileCheck %s --check-prefix=CL1
+// CL1: collapsed-indices-1 specified but not collapsed-indices-0
+
+// RUN: not mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=1  collapsed-indices-2=2}))' 2>&1 | FileCheck %s --check-prefix=CL2
+// CL2: collapsed-indices-2 specified but not collapsed-indices-1
+
+// RUN: not mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=1  collapsed-indices-1=2}))' 2>&1 | FileCheck %s --check-prefix=NON-ZERO
+// NON-ZERO: collapsed-indices arguments must include all values [0,N).
+
+// RUN: not mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=0  collapsed-indices-1=2}))' 2>&1 | FileCheck %s --check-prefix=NON-CONTIGUOUS
+// NON-CONTIGUOUS: collapsed-indices arguments must include all values [0,N).
+
+
+// Then test for invalid combinations of argument+input-ir
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=0,1}))' -verify-diagnostics
+func.func @too_few_iters(%arg0: index, %arg1: index, %arg2: index) {
+  // expected-error @+1 {{op has 1 iter args while this limited functionality testing pass was configured only for loops with exactly 2 iter args.}}
+  scf.parallel (%arg3) = (%arg0) to (%arg1) step (%arg2) {
+    scf.yield
+  }
+  return
+}
+
+func.func @too_many_iters(%arg0: index, %arg1: index, %arg2: index) {
+  // expected-error @+1 {{op has 3 iter args while this limited functionality testing pass was configured only for loops with exactly 2 iter args.}}
+  scf.parallel (%arg3, %arg4, %arg5) = (%arg0, %arg0, %arg0) to (%arg1, %arg1, %arg1) step (%arg2, %arg2, %arg2) {
+    scf.yield
+  }
+  return
+}

mlir/test/Transforms/parallel-loop-collapsing.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(scf-parallel-loop-collapsing{collapsed-indices-0=0,3 collapsed-indices-1=1,4 collapsed-indices-2=2}, canonicalize))' | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=0,3 collapsed-indices-1=1,4 collapsed-indices-2=2}, canonicalize))' | FileCheck %s
 
 // CHECK-LABEL: func @parallel_many_dims() {
 func.func @parallel_many_dims() {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c2 = arith.constant 2 : index
   %c3 = arith.constant 3 : index
   %c4 = arith.constant 4 : index

mlir/test/Transforms/single-parallel-loop-collapsing.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(scf-parallel-loop-collapsing{collapsed-indices-0=0,1}, canonicalize))' | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='builtin.module(func.func(test-scf-parallel-loop-collapsing{collapsed-indices-0=0,1}, canonicalize))' | FileCheck %s
 
 func.func @collapse_to_single() {
   %c0 = arith.constant 3 : index
   %c1 = arith.constant 7 : index
   %c2 = arith.constant 11 : index
   %c3 = arith.constant 29 : index
   %c4 = arith.constant 3 : index
   %c5 = arith.constant 4 : index