Func.call support for nested references

[chenghuaWang]

func.call is not able to call symbol across modules. Like ‘func’ Dialect - MLIR (llvm.org) said

The func.call operation represents a direct call to a function that is within the same symbol scope as the call.

However, nested symbol call is a quite useful op. As described in the example below:

in foo.mlir

module @module_foo {
  func.func @nested_symbol() -> f32 {
    %1 = arith.constant 42.0 : f32
    return %1 : f32
  }
}

in bar.mlir

module @module_bar {
  func.func @nested_symbol() -> f32 {
    %1 = arith.constant 42.0 : f32
    return %1 : f32
  }
}

in main.mlir , there is a function that depends on modules module_foo and module_bar

module @module_main {
  func.func @_exec_main() -> () {
    %0 = func.call @module_foo::@nested_symbol() : () -> f32
    %1 = func.call @module_bar::@nested_symbol() : () -> f32
    return
  }
}

When process main.mlir file, I can just copy module module_foo and module_bar into module module_main .

module @module_main {

  module @moudle_foo {
    ...
  }

  module @module_bar {
    ...
  }

  func.func @_exec_main() -> () {
    %0 = func.call @module_foo::@nested_symbol() : () -> f32
    %1 = func.call @module_bar::@nested_symbol() : () -> f32
    return
  }
}

Has anyone implemented a similar method? Or there is a similar Op that can replace func.call ?

[chenghuaWang]

I plan to expand func.call to support nested reference, I am not very familiar with SymbolTable in mlir, I think I may only need to modify these two methods?

llvm/llvm-project/blob/8be07adfb42d8d5d060b807d5445b6fefc949bee/mlir/lib/Dialect/Func/IR/FuncOps.cpp#L60

llvm/llvm-project/blob/365d6eb1f7d86cf28dc7d4995c3949e9d8bead58/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp#L596C6-L596C6

gpu.launch_func has implemented nested reference, I think I can learn something from it. I still hope someone can give me a little hint.

llvm/llvm-project/blob/efcd4d4ddd603e9689218c85815bfc2efa152b56/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp#L245

[mehdi_amini]

Have you looked at how the callee arguments is defined in ODS between func.call and gpu.launch_func? That should provide some good info I think (I suspect func.call is a FlatSymbolRef)

[chenghuaWang]

Yes, I’ve looked at the implementation of the gpu.launch_func. func.call using FlatSymbolRef.

I think I need to modify/implement these:

1. Using SymbolRef instead FlatSymbolRef
2. According to the number of layers of nested reference, continuously look up the symbol of parent op. The logic is like this:
   • Find the topmost ModuleOp where the current Op is located.
   • Using ArrayRef<FlatSymbolRef> to check the symbols from outside to inside to see if they are correct.

The logic is same in struct CallOpLowering and LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) .

I think I only need to modify the logic for symbol table checking. It is not necessary to modify the code for lowering to llvm ir ?

btw, Is there a quick way to compile it? I use cmake --build . ---target tools/mlir/test/check-mlir. This command is very time consuming.

[mehdi_amini]

@chenghuaWang

• Find the topmost ModuleOp where the current Op is located.

I don’t know what component you’re saying you need to modify, but in general “module” has nothing specific: you should get the closest enclosing SymbolTable operation, it may or may not be a ModuleOp.

@chenghuaWang

I think I only need to modify the logic for symbol table checking. It is not necessary to modify the code for lowering to llvm ir ?

Lowering to LLVM IR should not touch a call to a nested symbol and leave it as-is: it is the kind of thing to handle at a higher-level.

@chenghuaWang

This command is very time consuming.

It all depends on the cmake configuration… For the basics: are you using clang, lld, ninja, and ccache? Are you building in Release+Assert mode?

[stamnug]

有没有一种可能, 跨模块调用是链接器的职权.

[chenghuaWang]

这不仅仅只是符号链接吧. Module 可以用 不同的符号名 来区分出各个 Op 所在的命名空间啊。当然如果 func.call 需要支持这样的调用，应该在 lowering 时候需要对函数做重命名。

如果能在 high-level 上做出命名空间这样类似的概念，比直接手动做函数重命名方便点。

也有人提出想用这个来做 Module System

Calling functions on different modules

Just noticed this had an update 9 days ago, and thought I should update with the conclusion back then. This was for the Verona language, where we were lowering AST to MLIR. Verona had a module system, so it “felt natural” to use MLIR’s module system to represent it. What I was looking for is if there was any intent for this to be supported (and I was doing something wrong) or if there wasn’t (and support would need to be added/worked around). Turns out it was the latter, which is perfectly fin…

[stamnug]

跨文件没有其他文件模块信息, 你要么include其他文件, 这本质还是一个模块, 只不过强行分成了多个文件. 要么忽略其他模块信息, 这本质又是链接器进行链接. 在lowering的时候重命名就好.

所以你想做的不是要添加个include到mlir吧

[chenghuaWang]

是的 ，除了 import/include 这样的功能，还有 class 类其实也可以用 ModuleOp 来表示。我是想都包在 ModuleOp 内，这样比直接打平到一个 module 里面更加清晰。

但是 ModuleOp 似乎不能做这样的事情，从一个子模块调用另一个子模块的 Symbol 貌似是不被允许的。

Calling functions on different modules

No one looked into changing the func.func call, I haven’t through about the implications there, but note that Why is “small feature” not available in MLIR? applies I don’t know about your use case, but why not use some custom operation? (like the gpu.launc_func op) I would not that your example has two different cases, the call to @Bar::@bar from the top-level module is a proper nested reference (unsupported by func.call though), while the call to @Foo::@foo from a sibling is …

看来我还是在 func.call 上包一层，写一个新的 Op，然后这个 Op 接受一些模块标识的 Attribute，然后 Lowering 的时候对函数做重命名就行。然后所有的 Op 都打平到一个 Module 里面。

谢谢你的回答

[chenghuaWang]

Conclusion of this Post:

Thanks to @mehdi_amini @stamnug , I figure out what I need to do next.

why I need a nested references ?

I am translating AST(or CST) to MLIR. And I need to translate classAst and manage a module system . Organizing them with Module Op and then using nested symbols to access the Op/Value in each Module looks pretty nice! Just like what I mentioned before.

@chenghuaWang

module @module_main {

  module @moudle_foo {
    ...
  }

  module @module_bar {
    ...
  }

  func.func @_exec_main() -> () {
    %0 = func.call @module_foo::@nested_symbol() : () -> f32
    %1 = func.call @module_bar::@nested_symbol() : () -> f32
    return
  }
}

Solution

Just like @mehdi_amini said in post:

Calling functions on different modules

No one looked into changing the func.func call, I haven’t through about the implications there, but note that Why is “small feature” not available in MLIR? applies I don’t know about your use case, but why not use some custom operation? (like the gpu.launc_func op) I would not that your example has two different cases, the call to @Bar::@bar from the top-level module is a proper nested reference (unsupported by func.call though), while the call to @Foo::@foo from a sibling is …

The call to nested reference (@Foo::@foo) is not correct in the code below.

@jwkim98

module @__  {
  module @Foo  {
    func.func @foo() -> i64 {
      %c21_i64 = arith.constant 21 : i64
      func.return %c21_i64 : i64
    }
  }
  module @Bar  {
    func.func @bar(%arg0: i64) -> i64 {
      %0 = func.call @Foo::@foo() : () -> i64
      %1 = arith.addi %arg0, %0 : i64
      func.return %1 : i64
    }
  }
  func.func @main() -> i64 {
    %c21_i64 = arith.constant 21 : i64
    %0 = func.call @Bar::@bar(%c21_i64) : (i64) -> i64
    func.return %0 : i64
  }
}

This is contrary to what I was trying to accomplish (e.g., calling a function of class b from within class a). If you want to do this, you might as well rename the function.

Option 1

So why not just design a new FuncCallOp on the top of func.call.
The new FuncCallOp will probably require these parameter inputs:

1. func name, arguments Type, Result Type.
2. Visibility
3. ArrayRef< StringAttr> (nested moduleName, ClassName)

When lowring, rename function name to _moduleA_moduleB_classC_funcName.

Or just define a new SymbolType which will do renaming by default, and let FuncCallOp use this new type as input.

Option 2

Manually manage a table of symbols and use the information recorded in this table directly to generate the renamed function.

More

For compiling multiple mlir files, I think it’s better to just copy everything from the other modules into the module where the entry function is located. Then use a few passes to eliminate unused functions and other components.

[chenghuaWang]

@mehdi_amini

It all depends on the cmake configuration… For the basics: are you using clang, lld, ninja, and ccache? Are you building in Release+Assert mode?

I’m using MSVC Build system. Release+Assert mode. I found that cmake will rebuild all components, it seems like it is not cached. Maybe using clang and set -DLLVM_CCACHE_BUILD=ON will speed things up.