PassManager: What happens if there is no ModuleOp in input MLIR? Just the FuncOp only?

Hello,

I am a beginner and currently exploring the purpose of the PassManager and getOperation() methods in MLIR.

I was reading from this tutorial. There was a line: “the MLIR PassManager is actually limited to operation on a top-level ModuleOp.”

If the PassManager operates at the ModuleOp level, then what happens, for example, if my input MLIR is like the example below?

func @matmul() {
  %A = memref.alloc() : memref<64x64xf32>
  %B = memref.alloc() : memref<64x64xf32>
  %C = memref.alloc() : memref<64x64xf32>

  affine.for %i = 0 to 64 {
    affine.for %j = 0 to 64 {
      affine.for %k = 0 to 64 {
        %0 = affine.load %A[%i, %k] : memref<64x64xf32>
        %1 = affine.load %B[%k, %j] : memref<64x64xf32>
        %2 = arith.mulf %0, %1 : f32
        %3 = affine.load %C[%i, %j] : memref<64x64xf32>
        %4 = arith.addf %2, %3 : f32
        affine.store %4, %C[%i, %j] : memref<64x64xf32>
      }
    }
  }
  return
}

In such a case, I believe that to use the PassManager, we need to wrap the FuncOp with a ModuleOp. Am I right?

Thanks in advance for your time and help!

Hey :wave:

My understanding is that every Op in MLIR is either directly or indirctly a child of the builtin ModuleOp, so even in your example if you have the MLIR in a file and you run it through mlir-opt you will see that when it round trips the MLIR it sticks it inside a module.

Before:

func.func @matmul() {
  %A = memref.alloc() : memref<64x64xf32>
  %B = memref.alloc() : memref<64x64xf32>
  %C = memref.alloc() : memref<64x64xf32>

  affine.for %i = 0 to 64 {
    affine.for %j = 0 to 64 {
      affine.for %k = 0 to 64 {
        %0 = affine.load %A[%i, %k] : memref<64x64xf32>
        %1 = affine.load %B[%k, %j] : memref<64x64xf32>
        %2 = arith.mulf %0, %1 : f32
        %3 = affine.load %C[%i, %j] : memref<64x64xf32>
        %4 = arith.addf %2, %3 : f32
        affine.store %4, %C[%i, %j] : memref<64x64xf32>
      }
    }
  }
  return
}

After running mlir-opt test-file.mlir

module {
  func.func @matmul() {
    %alloc = memref.alloc() : memref<64x64xf32>
    %alloc_0 = memref.alloc() : memref<64x64xf32>
    %alloc_1 = memref.alloc() : memref<64x64xf32>
    affine.for %arg0 = 0 to 64 {
      affine.for %arg1 = 0 to 64 {
        affine.for %arg2 = 0 to 64 {
          %0 = affine.load %alloc[%arg0, %arg2] : memref<64x64xf32>
          %1 = affine.load %alloc_0[%arg2, %arg1] : memref<64x64xf32>
          %2 = arith.mulf %0, %1 : f32
          %3 = affine.load %alloc_1[%arg0, %arg1] : memref<64x64xf32>
          %4 = arith.addf %2, %3 : f32
          affine.store %4, %alloc_1[%arg0, %arg1] : memref<64x64xf32>
        }
      }
    }
    return
  }
}

(btw you need to prefix your func with the func. dialect namespace in your original example).

I think this is just a convience feature so when you’re writing MLIR out by hand or reading it you don’t need to explicitly define the module; but it is always there, just hidden :slight_smile:

You’re correct in that your pass manager will run on the top level ModuleOp though, if you want to run it on the FuncOp nested inside the module (which remember is there even if you didn’t define it yourself) you’ll need to add a nested pass which matches the recursive depth of the operation you wish to run on which in this case is just 1 level.

For example, assuming you have a top level pass manager pm, then to run a pass on a func::FuncOp nested inside a mlir::ModuleOp you’d do something along the lines of pm.addNestedPass<func::FuncOp>(YourNameSpace::createYourFunctionPass()).

If you want to try this out on the command line via mlir-opt to better understand the nested topology of the IR and how it interacts with the nested pass manager structure I think there is a good example here in the docs: Pass Infrastructure - MLIR

$ mlir-opt foo.mlir -pass-pipeline='builtin.module(func.func(cse,canonicalize),convert-func-to-llvm{use-bare-ptr-memref-call-conv=1})'

The passes aren’t important here, but you can see we have a pipeline with a “top level” builtin.module pass manager and within that, nested at 1 level down is a pass manager that will run cse and canonicalization on FuncOps which is matching the recursive nesting of the input IR.

You could always make your pass a module pass and just iterate over the child ops finding the function you want to operate on, however this means your pass can’t be run in parallel on different functions since it could in theory modify the module scope. This I think is one of the motivations for designing the pass manager in this way.

I hope that helps somewhat :slight_smile:

2 Likes

That sentence does not seem right to me. Maybe it was the case a while ago, I don’t remember.

Your post is globally correct, but this sentence isn’t entirely accurate actually: as you mentioned later it is just a convenience of mlir-opt, which by default will add an enclosing module implicitly when not present.
This behavior can be disabled, and so it’s not a limitation of MLIR (you can build IR programmatically without using mlir-opt).

So the example can also be processed with mlir-opt matmul.mlir --no-implicit-module --pass-pipeline="func.func(canonicalize)" just fine.

Great!! Thanks for letting me know :blush:

But I have one question. what did you mean by “you can build IR programmatically without using mlir-opt”?

Does it mean, to have the ability to construct and manipulate MLIR directly through code, using APIs provided by the MLIR framework, rather than relying on the mlir-opt command?

It would be great if this would be mentioned in the documentation. I am not sure if it is there anywhere else though :blush:

Thanks a lot.

Yes: this is what is done in the tutorial starting Chapter 2 :slight_smile:

1 Like

Thanks for the reply :blush:

I have to go through Chapter 2 very soon…