[MLIR] Allow nested operations to be transformed

[Whitney]

Summary

We noticed that transformations that operate on func::FuncOp, e.g., AffineScalarReplacementPass, don't get invoked for func::FuncOp nested inside another region, e.g., gpu.module.

This commit allows transformations that operate on func::FuncOp to work as expected even if the func::FuncOp operation is nested in another region.

[mehdi_amini]

mlir/test/Pass/nested-operations.mlir

5

I'm not sure I understand what you're trying to do.

If you want -affine-scalrep to be scheduled on the function, you should do so explicitly.

mlir-opt --pass-pipeline="gpu.module(func.func(affine-scalrep))"

[Whitney]

mlir/test/Pass/nested-operations.mlir

5

We have a driver that add the pass like PM.addNestedPass<func::FuncOp>(mlir::createAffineScalarReplacementPass());.
We noticed that when we nest the func::FuncOp under gpu.module, the pass is no longer invoked.
Is there a better way to add the pass in the pipeline?

[Whitney]

mlir/test/Pass/nested-operations.mlir

5

We are trying to have a pipeline that AffineScalarReplacementPass is invoked no matter the func::FuncOp is nested in gpu.module or not.

[mehdi_amini]

mlir/test/Pass/nested-operations.mlir

5

We are trying to have a pipeline that AffineScalarReplacementPass is invoked no matter the func::FuncOp is nested in gpu.module or not.

In general you can't.
What you can do is something equivalent to:

--pass-pipeline="gpu.module(func.func(affine-scalrep)),func.func(affine-scalrep)"

or (I think):

--pass-pipeline="any(func.func(affine-scalrep)),func.func(affine-scalrep)"

[Whitney]

Given that there is a workaround, closing this review.

mlir/test/Pass/nested-operations.mlir

5

That's unfortunate. It would be ideal to have a way to run a func::FuncOp pass no matter if it is nested or not.

[mehdi_amini]

mlir/test/Pass/nested-operations.mlir

5

You can implement this kind of conditional logic in a pass that you write. The pass would run a pipeline dynamically: https://mlir.llvm.org/docs/PassManagement/#dynamic-pass-pipelines

In general a compiler pipeline has full control on the structure of the IR and even more: it expects a specific form.
It's pretty rare that you have to take some IR and find functions without knowing where they are supposed to be.

mlir/lib/Pass/Pass.cpp

struct OpPMInfo {
     AnalysisManager am;
   };
 
   // Run a prepass over the operation to collect the nested operations to
   // execute over. This ensures that an analysis manager exists for each
   // operation, as well as providing a queue of operations to execute over.
   std::vector<OpPMInfo> opInfos;
   DenseMap<OperationName, std::optional<unsigned>> knownOpPMIdx;
-  for (auto &region : getOperation()->getRegions()) {
+    std::function<void(Operation &)> collectOpInfos = [&](Operation &regionOp) {
+    for (auto &region : regionOp.getRegions()) {
       for (Operation &op : region.getOps()) {
         // Get the pass manager index for this operation type.
         auto pmIdxIt = knownOpPMIdx.try_emplace(op.getName(), std::nullopt);
         if (pmIdxIt.second) {
           if (auto *mgr = findPassManagerFor(mgrs, op.getName(), *context))
             pmIdxIt.first->second = std::distance(mgrs.begin(), mgr);
         }
 
         // If this operation can be scheduled, add it to the list.
         if (pmIdxIt.first->second)
           opInfos.emplace_back(*pmIdxIt.first->second, &op, am.nest(&op));
+        else
+          collectOpInfos(op);
       }
     }
+  };
+  collectOpInfos(*getOperation());
 
   // Get the current thread for this adaptor.
   PassInstrumentation::PipelineParentInfo parentInfo = {llvm::get_threadid(),
                                                         this};
   auto *instrumentor = am.getPassInstrumentor();
 
   // An atomic failure variable for the async executors.
   std::vector<std::atomic<bool>> activePMs(asyncExecutors.size());

mlir/test/Pass/nested-operations.mlir

+// RUN: mlir-opt -allow-unregistered-dialect %s -affine-scalrep | FileCheck %s
+
+// Affine scalar replacement pass is a func::FuncOp pass.
+// Ensure that simple_store_load function nested in gpu.module can be transformed by -affine-scalrep.
+
+gpu.module @functions {
+  // CHECK-LABEL: func @simple_store_load() {
+  func.func @simple_store_load() {
+    %cf7 = arith.constant 7.0 : f32
+    %m = memref.alloc() : memref<10xf32>
+    affine.for %i0 = 0 to 10 {
+      affine.store %cf7, %m[%i0] : memref<10xf32>
+      %v0 = affine.load %m[%i0] : memref<10xf32>
+      %v1 = arith.addf %v0, %v0 : f32
+    }
+    memref.dealloc %m : memref<10xf32>
+    return
+  // CHECK:       %[[C7:.*]] = arith.constant 7.000000e+00 : f32
+  // CHECK-NEXT:  affine.for %{{.*}} = 0 to 10 {
+  // CHECK-NEXT:    arith.addf %[[C7]], %[[C7]] : f32
+  // CHECK-NEXT:  }
+  // CHECK-NEXT:  return
+  }
+}