[mlir][conversion] NFC - Relax convert-cf-to-llvm and finalize-memref-to-llvm to also work on non-builtin module ops

[nicolasvasilache]

Summary

This is a prerequisite for unentangling LowerGpuOpsToNVVMOps which explicitly populates its conversion with
populateControlFlowToLLVMConversionPatterns and populateFinalizeMemRefToLLVMConversionPatterns.

However those patterns cannot be generally added to a pass anchored on a non-ModuleOp.

Instead, LowerGpuOpsToNVVMOps should stop including the world, in a future commit.

[ftynse]

Hmm, if some conversions need a symbol table, would it make sense to use that as a soft restriction rather than skipping the conversions entirely? Symbol table is not an interface (it's a trait and we don't have passes that run on traits), but it looks possible to check for the trait being early in runOnOperation and bail out?

[springerm]

mlir/include/mlir/Conversion/Passes.td

266

Can we make sure that the pass signals a failure in that case? (It probably already does, assuming that this is using a dialect conversion.)

[nicolasvasilache]

@ftynse @springerm

The use case here is to allow partial conversions depending on the anchor op.
During GPU lowering, such passes would apply on GPUModuleOp to lower what they can; later after GPUModuleOp is lowered away into a GPUModuleOp the rest of the operations can be lowered.

In D156860#4553040, @ftynse wrote:

Hmm, if some conversions need a symbol table, would it make sense to use that as a soft restriction rather than skipping the conversions entirely? Symbol table is not an interface (it's a trait and we don't have passes that run on traits), but it looks possible to check for the trait being early in runOnOperation and bail out?

I did start with this route but it is unclear to me that:

1. any op with a symbol table is good enough to support external functions?
2. we want the pass to fail/bail out

mlir/include/mlir/Conversion/Passes.td

266

I am unclear that we want to fail explicitly here: the current effect is that cf.assert is just ignored.
This is what happens already if one adds a nested pass that operates on a non-builtin module op: if the anchor op does not contain a module, this pass does not apply and no error is produced.

[ftynse]

Did this land accidentally?

In D156860#4553165, @nicolasvasilache wrote:

@ftynse @springerm

The use case here is to allow partial conversions depending on the anchor op.
During GPU lowering, such passes would apply on GPUModuleOp to lower what they can; later after GPUModuleOp is lowered away into a GPUModuleOp the rest of the operations can be lowered.

In D156860#4553040, @ftynse wrote:

Hmm, if some conversions need a symbol table, would it make sense to use that as a soft restriction rather than skipping the conversions entirely? Symbol table is not an interface (it's a trait and we don't have passes that run on traits), but it looks possible to check for the trait being early in runOnOperation and bail out?

I did start with this route but it is unclear to me that:

1. any op with a symbol table is good enough to support external functions?

Symbol tables have no knowledge of external functions, or even functions altogether. They only reason about symbols. It's the function op that has can present itself as external (empty region) or not (non-empty region). Nothing prevents one from creating such ops, symbol table is only there for symbol name resolution and deduplication.

2. we want the pass to fail/bail out

Passes can easily signal failures at any moment, this doesn't have to be a hardcoded precondition on the root operation...

void runOnOperation(Operation *op) {
  if (!op->hasTrait<OpTrait::SymbolTable>()) {
    op->emitError() << "pass can only run on ops with the SymbolTable trait";
    return signalPassFailure();
  }
  // do actual work here.
}

[nicolasvasilache]

In D156860#4560389, @ftynse wrote:

Did this land accidentally?

Ugh .. sorry yes .. it was stacked with another patch I landed.

Thanks for the revert(s)!

mlir/include/mlir/Conversion/Passes.td

def ConvertComplexToStandard : Pass<"convert-complex-to-standard"> {
   let constructor = "mlir::createConvertComplexToStandardPass()";
   let dependentDialects = ["math::MathDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // ControlFlowToLLVM
 //===----------------------------------------------------------------------===//
 
-def ConvertControlFlowToLLVMPass : Pass<"convert-cf-to-llvm", "ModuleOp"> {
+def ConvertControlFlowToLLVMPass : Pass<"convert-cf-to-llvm"> {
   let summary = "Convert ControlFlow operations to the LLVM dialect";
   let description = [{
     Convert ControlFlow operations into LLVM IR dialect operations.
 
     If other operations are present and their results are required by the LLVM
     IR dialect operations, the pass will fail.  Any LLVM IR operations or types
     already present in the IR will be kept as is.
+
+    This pass is not restricted to run to any op but certain conversions are
+    omitted when the pass is run on a non-builtin ModuleOp. The omitted 
+    conversions are the ones that assume a ModuleOp is accessible, for 
+    manipulating its symbol table (for the purpose of calling into external
+    libraries), namely:
+      1. cf.assert lowering
+
+    Such decoupling of partial lowerings is particularly useful in the context
+    of accelerators that may depend on custom module operations.
   }];
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
     Option<"useOpaquePointers", "use-opaque-pointers", "bool",
                    /*default=*/"true", "Generate LLVM IR using opaque pointers "
Option<"convertCtlz", "convert-ctlz", "bool", /*default=*/"false",
            "for targets that do not natively support ctlz.">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefToLLVM
 //===----------------------------------------------------------------------===//
 
-def FinalizeMemRefToLLVMConversionPass :
-    Pass<"finalize-memref-to-llvm", "ModuleOp"> {
+def FinalizeMemRefToLLVMConversionPass : Pass<"finalize-memref-to-llvm"> {
   let summary = "Finalize MemRef dialect to LLVM dialect conversion";
   let description = [{
-    Finalize the conversion of the operations from the MemRef
-    dialect to the LLVM dialect.
-    This conversion will not convert some complex MemRef
-    operations. Make sure to run `expand-strided-metadata`
-    beforehand for these.
+    Finalize the conversion of the operations from the MemRef dialect to the
+    LLVM dialect.
+    This conversion will not convert some complex MemRef operations. Make sure
+    to run `expand-strided-metadata` beforehand for these.
+
+    This pass is not restricted to run to any op but certain conversions are
+    omitted when the pass is run on a non-builtin ModuleOp. The omitted 
+    conversions are the ones that assume a ModuleOp is accessible, for 
+    manipulating its symbol table (for the purpose of calling into external
+    libraries), namely:
+      1. memref.copy lowering
+      2. memref.alloc lowering
+      3. memref.dealloc lowering
+
+    Such decoupling of partial lowerings is particularly useful in the context
+    of accelerators that may depend on custom module operations.
   }];
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"useAlignedAlloc", "use-aligned-alloc", "bool", /*default=*/"false",
            "Use aligned_alloc in place of malloc for heap allocations">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

matchAndRewrite(cf::SwitchOp op, typename cf::SwitchOp::Adaptor adaptor,
     rewriter.replaceOpWithNewOp<LLVM::SwitchOp>(
         op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs());
     return success();
   }
 };
 
 } // namespace
 
-void mlir::cf::populateControlFlowToLLVMConversionPatterns(
+static void populateModuleIndependentControlFlowToLLVMConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
-      AssertOpLowering,
       BranchOpLowering,
       CondBranchOpLowering,
       SwitchOpLowering>(converter);
   // clang-format on
 }
 
+void mlir::cf::populateControlFlowToLLVMConversionPatterns(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
+  // clang-format off
+  patterns.add<AssertOpLowering>(converter);
+  // clang-format on
+  populateModuleIndependentControlFlowToLLVMConversionPatterns(converter,
+                                                               patterns);
+}
+
 void mlir::cf::populateAssertToLLVMConversionPattern(
     LLVMTypeConverter &converter, RewritePatternSet &patterns,
     bool abortOnFailure) {
   patterns.add<AssertOpLowering>(converter, abortOnFailure);
 }
 
 //===----------------------------------------------------------------------===//
 // Pass Definition
void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
 
     LowerToLLVMOptions options(&getContext());
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
     options.useOpaquePointers = useOpaquePointers;
 
     LLVMTypeConverter converter(&getContext(), options);
-    mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns);
-
+    if (isa<ModuleOp>(getOperation())) {
+      mlir::cf::populateControlFlowToLLVMConversionPatterns(converter,
+                                                            patterns);
+    } else {
+      populateModuleIndependentControlFlowToLLVMConversionPatterns(converter,
+                                                                   patterns);
+    }
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
 
     rewriter.replaceOp(extractStridedMetadataOp, results);
     return success();
   }
 };
 
 } // namespace
 
-void mlir::populateFinalizeMemRefToLLVMConversionPatterns(
+static void populateModuleIndependentFinalizeMemRefToLLVMConversionPatterns(
     LLVMTypeConverter &converter, RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
       AllocaOpLowering,
       AllocaScopeOpLowering,
       AtomicRMWOpLowering,
       AssumeAlignmentOpLowering,
       ConvertExtractAlignedPointerAsIndex,
       DimOpLowering,
       ExtractStridedMetadataOpLowering,
       GenericAtomicRMWOpLowering,
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
-      MemRefCopyOpLowering,
       MemorySpaceCastOpLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
       PrefetchOpLowering,
       RankOpLowering,
       ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
       ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
       StoreOpLowering,
       SubViewOpLowering,
       TransposeOpLowering,
       ViewOpLowering>(converter);
   // clang-format on
+}
+
+void mlir::populateFinalizeMemRefToLLVMConversionPatterns(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns) {
+  // clang-format off
+  patterns.add<
+      MemRefCopyOpLowering>(converter);
+  // clang-format on
+
   auto allocLowering = converter.getOptions().allocLowering;
   if (allocLowering == LowerToLLVMOptions::AllocLowering::AlignedAlloc)
     patterns.add<AlignedAllocOpLowering, AlignedReallocOpLowering,
                  DeallocOpLowering>(converter);
   else if (allocLowering == LowerToLLVMOptions::AllocLowering::Malloc)
     patterns.add<AllocOpLowering, ReallocOpLowering, DeallocOpLowering>(
         converter);
+
+  populateModuleIndependentFinalizeMemRefToLLVMConversionPatterns(converter,
+                                                                  patterns);
 }
 
 namespace {
 struct FinalizeMemRefToLLVMConversionPass
     : public impl::FinalizeMemRefToLLVMConversionPassBase<
           FinalizeMemRefToLLVMConversionPass> {
   using FinalizeMemRefToLLVMConversionPassBase::
       FinalizeMemRefToLLVMConversionPassBase;
void runOnOperation() override {
 
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
 
     LLVMTypeConverter typeConverter(&getContext(), options,
                                     &dataLayoutAnalysis);
     RewritePatternSet patterns(&getContext());
     populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns);
+    if (isa<ModuleOp>(getOperation())) {
+      populateFinalizeMemRefToLLVMConversionPatterns(typeConverter, patterns);
+    } else {
+      populateModuleIndependentFinalizeMemRefToLLVMConversionPatterns(
+          typeConverter, patterns);
+    }
     LLVMConversionTarget target(getContext());
     target.addLegalOp<func::FuncOp>();
     if (failed(applyPartialConversion(op, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace