Inline all functions that are called.

[LPanosTT]

• Register FuncDialect Inliner interface
• Create and register TTIRDialect Inliner interface
• Added pass that will make-private all functions which are called
  • The inliners will inline all functions - however it will leave those functions in the module if they are public even when they are dead code. Setting them to private ensures they are deleted from the module.

Since we are inlining all functions that are called, if the module has multiple distinct programs they will all remain afterward. This is because the topmost function in the call stack is the only non-dead-code function that doesn't get called.

Disabling the inliner pass is as simple as removing the added lines in lib/Dialect/TTNN/Pipelines/TTNNPipelines.cpp

If you wish to not register them altogether you'll have to delete their invocations in lib/RegisterAll.cpp::registerAllExtensions

[mtopalovicTT]

Out of curiosity, will forge ever generate something like below?

%1 = call @do_mult(%arg0, %arg1, %0) : (tensor<64x128xf32>, tensor<64x128xf32>, tensor<64x128xf32>) -> tensor<64x128xf32>

[rpavlovicTT]

Does this have to be defined as an extension? Or could you just do

addInterfaces<TTIRInlinerInterface>();

at the dialect initialization.

[LPanosTT]

I could do that. But for FuncDialect I’m forced to register it since its initialize method is inside the toolchain. So, I thought it would be cleaner to just explicitly register them both rather than register only the FuncDialect Inliner and add the TTIRDialect Inliner in Initialize()

[rpavlovicTT]

I see your point. I am not sure should we do it this way, as extensions are used to extend some dialect that is already built for you. Maybe I am missing knowledge on dialect extensions. @sdjordjevicTT what do you think?

Btw func inliner depends on CF dialect, and creates CF branch op.

  void handleTerminator(Operation *op, Block *newDest) const final {
...
    builder.create<cf::BranchOp>(op->getLoc(), newDest, returnOp.getOperands());

are we sure this op won't be inserted in our IR? We have no lowering for CF dialect for any backend.

[LPanosTT]

I believe that only hits for functions defined inside other functions or as lambdas to certain ops. Which I guess could be a problem. Another idea is to create our own Call, Func, and Return ops in TTIR so we don't need to rely on the Func dialect at all, and we can just use the Inliner I've provided. When lowering to TTIR from whatever frontend we'd just need to replace the Func::FuncOp generation with TTIR::FuncOp etc.

[LPanosTT]

I think I can also define my own Inliner for the func dialect (which wouldn't have the overloaded handleTerminator that creates cf ops) and register that instead.

[nsmithtt]

It does seem like arith dialect uses the interface approach, can we experiment with that @LPanosTT?

 namespace {                                         
  /// This class defines the interface for handling inlining for arithmetic         
  /// dialect operations.                         
  struct ArithInlinerInterface : public DialectInlinerInterface {    
    using DialectInlinerInterface::DialectInlinerInterface;    
                                                  
    /// All arithmetic dialect ops can be inlined.                              
    bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {    
      return true;                                                              
    }                                                                           
  };                                                                            
  } // namespace 

  void arith::ArithDialect::initialize() {                                      
    ...                                                                                                                                                                                                                               
    addInterfaces<ArithInlinerInterface>();

[nsmithtt]

Yeah torch mlir does it this way too: https://github.com/llvm/torch-mlir/blob/8787970afed3c4e1497fb24c4fdeec179fcb61f6/lib/Dialect/Torch/IR/TorchDialect.cpp#L32

I think it seems like it can be done with only changes to TTIRDialect.cpp. I'm not sure we need to register the inliner on behalf of func.

[nsmithtt]

See their InitAll for reference too https://github.com/llvm/torch-mlir/blob/8787970afed3c4e1497fb24c4fdeec179fcb61f6/lib/InitAll.cpp#L13.

It seems like all they do is:

• Implement the dialect interface
• addInterfaces
• Register func dialect inliner extension

[LPanosTT]

@nsmithtt it looks like they do register the inliner on behalf of func though

void mlir::torch::registerAllExtensions(mlir::DialectRegistry &registry) {
  mlir::func::registerInlinerExtension(registry);
  tensor::registerInferTypeOpInterfaceExternalModels(registry);
}

I can move the registration of the TTIR Inliner interface to TTIRDialect.cpp and use addInterfaces I've tried and it works. So I'll do that

[LPanosTT]

@nsmithtt @rpavlovicTT Should I create my own inliner for func? One that doesn't override the handleTerminator for Block? Just so that we leave the cf dialect entirely out of the picture.

[LPanosTT]

Out of curiosity, will forge ever generate something like below?

@mtopalovicTT I don’t think forge will since we only have one forge graph and no 'func' ops. However JAX certainly can (I.e relu can be put in a separate function defined as a maximum against 0). I think it depends on the actual structure of the model in JAX and if there are functions that are called in the python definition

[nsmithtt]

Yeah torch mlir does it this way too: https://github.com/llvm/torch-mlir/blob/8787970afed3c4e1497fb24c4fdeec179fcb61f6/lib/Dialect/Torch/IR/TorchDialect.cpp#L32

I think it seems like it can be done with only changes to TTIRDialect.cpp. I'm not sure we need to register the inliner on behalf of func.

[nsmithtt]

Thanks @LPanosTT! This looks great

[nsmithtt]

@LPanosTT, I think I have a build fix over here:

• Static link all targets and tools #949

I've only tested on macOS so far, still need to get ubuntu to sign off and ideally one of the FE CI's to run against it.

[nsmithtt]

fyi @wooseokTT