[mlir] Add support for "promised" interfaces

Promised interfaces allow for a dialect to "promise" the implementation of an interface, i.e. declare that it supports an interface, but have the interface defined in an extension in a library separate from the dialect itself. A promised interface is powerful in that it alerts the user when the interface is attempted to be used (e.g. via cast/dyn_cast/etc.) and the implementation has not yet been provided. This makes the system much more robust against misconfiguration, and ensures that we do not lose the benefit we currently have of defining the interface in the dialect library. Differential Revision: https://reviews.llvm.org/D120368
llvm · Jun 9, 2023 · a5ef51d · a5ef51d
1 parent fb19fa2
commit a5ef51d
Show file tree

Hide file tree

Showing 40 changed files with 411 additions and 75 deletions.
diff --git a/flang/include/flang/Optimizer/Support/InitFIR.h b/flang/include/flang/Optimizer/Support/InitFIR.h
@@ -18,6 +18,7 @@
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassRegistry.h"
@@ -34,27 +35,39 @@ namespace fir::support {
       mlir::vector::VectorDialect, mlir::math::MathDialect,                    \
       mlir::complex::ComplexDialect, mlir::DLTIDialect
 
+#define FLANG_CODEGEN_DIALECT_LIST FIRCodeGenDialect, mlir::LLVM::LLVMDialect
+
 // The definitive list of dialects used by flang.
 #define FLANG_DIALECT_LIST                                                     \
-  FLANG_NONCODEGEN_DIALECT_LIST, FIRCodeGenDialect, mlir::LLVM::LLVMDialect
+  FLANG_NONCODEGEN_DIALECT_LIST, FLANG_CODEGEN_DIALECT_LIST
 
 inline void registerNonCodegenDialects(mlir::DialectRegistry &registry) {
   registry.insert<FLANG_NONCODEGEN_DIALECT_LIST>();
+  mlir::func::registerInlinerExtension(registry);
 }
 
 /// Register all the dialects used by flang.
 inline void registerDialects(mlir::DialectRegistry &registry) {
-  registry.insert<FLANG_DIALECT_LIST>();
+  registerNonCodegenDialects(registry);
+  registry.insert<FLANG_CODEGEN_DIALECT_LIST>();
 }
 
 inline void loadNonCodegenDialects(mlir::MLIRContext &context) {
+  mlir::DialectRegistry registry;
+  registerNonCodegenDialects(registry);
+  context.appendDialectRegistry(registry);
+
   context.loadDialect<FLANG_NONCODEGEN_DIALECT_LIST>();
 }
 
 /// Forced load of all the dialects used by flang.  Lowering is not an MLIR
 /// pass, but a producer of FIR and MLIR. It is therefore a requirement that the
 /// dialects be preloaded to be able to build the IR.
 inline void loadDialects(mlir::MLIRContext &context) {
+  mlir::DialectRegistry registry;
+  registerDialects(registry);
+  context.appendDialectRegistry(registry);
+
   context.loadDialect<FLANG_DIALECT_LIST>();
 }
 

diff --git a/flang/lib/Frontend/CMakeLists.txt b/flang/lib/Frontend/CMakeLists.txt
@@ -1,4 +1,5 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 add_flang_library(flangFrontend
   CompilerInstance.cpp
@@ -18,6 +19,7 @@ add_flang_library(flangFrontend
   HLFIRDialect
   MLIRIR
   ${dialect_libs}
+  ${extension_libs}
 
   LINK_LIBS
   FortranParser
@@ -39,6 +41,7 @@ add_flang_library(flangFrontend
   MLIRSCFToControlFlow
   MLIRTargetLLVMIRImport
   ${dialect_libs}
+  ${extension_libs}
 
   LINK_COMPONENTS
   Passes

diff --git a/flang/lib/Lower/CMakeLists.txt b/flang/lib/Lower/CMakeLists.txt
@@ -1,4 +1,5 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 add_flang_library(FortranLower
   Allocatable.cpp
@@ -33,6 +34,7 @@ add_flang_library(FortranLower
   FIRTransforms
   HLFIRDialect
   ${dialect_libs}
+  ${extension_libs}
 
   LINK_LIBS
   FIRDialect
@@ -42,6 +44,7 @@ add_flang_library(FortranLower
   FIRTransforms
   HLFIRDialect
   ${dialect_libs}
+  ${extension_libs}
   FortranCommon
   FortranParser
   FortranEvaluate

diff --git a/flang/lib/Optimizer/Builder/CMakeLists.txt b/flang/lib/Optimizer/Builder/CMakeLists.txt
@@ -1,4 +1,5 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 add_flang_library(FIRBuilder
   BoxValue.cpp
@@ -31,11 +32,13 @@ add_flang_library(FIRBuilder
   FIRDialect
   HLFIRDialect
   ${dialect_libs}
+  ${extension_libs}
 
   LINK_LIBS
   FIRDialect
   FIRDialectSupport
   FIRSupport
   HLFIRDialect
   ${dialect_libs}
+  ${extension_libs}
 )
diff --git a/flang/lib/Optimizer/Support/CMakeLists.txt b/flang/lib/Optimizer/Support/CMakeLists.txt
@@ -1,4 +1,5 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 add_flang_library(FIRSupport
   InitFIR.cpp
@@ -9,9 +10,11 @@ add_flang_library(FIRSupport
   HLFIROpsIncGen
   MLIRIR
   ${dialect_libs}
+  ${extension_libs}
 
   LINK_LIBS
   ${dialect_libs}
+  ${extension_libs}
   MLIRBuiltinToLLVMIRTranslation
   MLIROpenACCToLLVMIRTranslation
   MLIROpenMPToLLVMIRTranslation

diff --git a/flang/tools/bbc/CMakeLists.txt b/flang/tools/bbc/CMakeLists.txt
@@ -10,6 +10,7 @@ FIROptCodeGenPassIncGen
 
 llvm_update_compile_flags(bbc)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 target_link_libraries(bbc PRIVATE
 FIRDialect
 FIRDialectSupport
@@ -19,6 +20,7 @@ FIRBuilder
 HLFIRDialect
 HLFIRTransforms
 ${dialect_libs}
+${extension_libs}
 MLIRAffineToStandard
 MLIRSCFToControlFlow
 FortranCommon

diff --git a/flang/tools/fir-opt/CMakeLists.txt b/flang/tools/fir-opt/CMakeLists.txt
@@ -1,6 +1,7 @@
 add_flang_tool(fir-opt fir-opt.cpp)
 llvm_update_compile_flags(fir-opt)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 if(FLANG_INCLUDE_TESTS)
   set(test_libs
@@ -18,6 +19,7 @@ target_link_libraries(fir-opt PRIVATE
   FIRAnalysis
   ${test_libs}
   ${dialect_libs}
+  ${extension_libs}
 
   # TODO: these should be transitive dependencies from a target providing
   # "registerFIRPasses()"

diff --git a/flang/tools/tco/CMakeLists.txt b/flang/tools/tco/CMakeLists.txt
@@ -5,6 +5,7 @@ set(LLVM_LINK_COMPONENTS
 add_flang_tool(tco tco.cpp)
 llvm_update_compile_flags(tco)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 target_link_libraries(tco PRIVATE
   FIRCodeGen
   FIRDialect
@@ -15,13 +16,13 @@ target_link_libraries(tco PRIVATE
   HLFIRDialect
   HLFIRTransforms
   ${dialect_libs}
+  ${extension_libs}
   MLIRIR
   MLIRLLVMDialect
   MLIRBuiltinToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRTargetLLVMIRExport
   MLIRPass
-  MLIRFuncToLLVM
   MLIRTransforms
   MLIRAffineToStandard
   MLIRAnalysis

diff --git a/flang/unittests/Optimizer/CMakeLists.txt b/flang/unittests/Optimizer/CMakeLists.txt
@@ -1,4 +1,5 @@
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 set(LIBS
   FIRBuilder
@@ -8,6 +9,7 @@ set(LIBS
   FIRSupport
   HLFIRDialect
   ${dialect_libs}
+  ${extension_libs}
   LLVMTargetParser
 )
 

diff --git a/mlir/cmake/modules/AddMLIR.cmake b/mlir/cmake/modules/AddMLIR.cmake
@@ -588,6 +588,12 @@ function(add_mlir_conversion_library name)
   add_mlir_library(${ARGV} DEPENDS mlir-headers)
 endfunction(add_mlir_conversion_library)
 
+# Declare the library associated with an extension.
+function(add_mlir_extension_library name)
+  set_property(GLOBAL APPEND PROPERTY MLIR_EXTENSION_LIBS ${name})
+  add_mlir_library(${ARGV} DEPENDS mlir-headers)
+endfunction(add_mlir_extension_library)
+
 # Declare the library associated with a translation.
 function(add_mlir_translation_library name)
   set_property(GLOBAL APPEND PROPERTY MLIR_TRANSLATION_LIBS ${name})

diff --git a/mlir/cmake/modules/CMakeLists.txt b/mlir/cmake/modules/CMakeLists.txt
@@ -24,6 +24,7 @@ export(TARGETS ${MLIR_EXPORTS} FILE ${mlir_cmake_builddir}/MLIRTargets.cmake)
 get_property(MLIR_ALL_LIBS GLOBAL PROPERTY MLIR_ALL_LIBS)
 get_property(MLIR_DIALECT_LIBS GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(MLIR_CONVERSION_LIBS GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+get_property(MLIR_EXTENSION_LIBS GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 get_property(MLIR_TRANSLATION_LIBS GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
 
 # Generate MlirConfig.cmake for the build tree.

diff --git a/mlir/cmake/modules/MLIRConfig.cmake.in b/mlir/cmake/modules/MLIRConfig.cmake.in
@@ -21,6 +21,7 @@ set(MLIR_MAIN_SRC_DIR "@MLIR_MAIN_SRC_DIR@")
 set_property(GLOBAL PROPERTY MLIR_ALL_LIBS "@MLIR_ALL_LIBS@")
 set_property(GLOBAL PROPERTY MLIR_DIALECT_LIBS "@MLIR_DIALECT_LIBS@")
 set_property(GLOBAL PROPERTY MLIR_CONVERSION_LIBS "@MLIR_CONVERSION_LIBS@")
+set_property(GLOBAL PROPERTY MLIR_EXTENSION_LIBS "@MLIR_EXTENSION_LIBS@")
 set_property(GLOBAL PROPERTY MLIR_TRANSLATION_LIBS "@MLIR_TRANSLATION_LIBS@")
 
 # Provide all our library targets to users.

diff --git a/mlir/examples/toy/Ch5/CMakeLists.txt b/mlir/examples/toy/Ch5/CMakeLists.txt
@@ -28,9 +28,11 @@ add_toy_chapter(toyc-ch5
 include_directories(${CMAKE_CURRENT_BINARY_DIR})
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include/)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 target_link_libraries(toyc-ch5
   PRIVATE
     ${dialect_libs}
+    ${extension_libs}
     MLIRAnalysis
     MLIRCallInterfaces
     MLIRCastInterfaces

diff --git a/mlir/examples/toy/Ch5/toyc.cpp b/mlir/examples/toy/Ch5/toyc.cpp
@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
 #include "toy/Dialect.h"
 #include "toy/MLIRGen.h"
 #include "toy/Parser.h"
@@ -107,7 +108,10 @@ int loadMLIR(llvm::SourceMgr &sourceMgr, mlir::MLIRContext &context,
 }
 
 int dumpMLIR() {
-  mlir::MLIRContext context;
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
+
+  mlir::MLIRContext context(registry);
   // Load our Dialect in this MLIR Context.
   context.getOrLoadDialect<mlir::toy::ToyDialect>();
 

diff --git a/mlir/examples/toy/Ch6/CMakeLists.txt b/mlir/examples/toy/Ch6/CMakeLists.txt
@@ -39,10 +39,12 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR})
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include/)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 target_link_libraries(toyc-ch6
   PRIVATE
     ${dialect_libs}
     ${conversion_libs}
+    ${extension_libs}
     MLIRAnalysis
     MLIRBuiltinToLLVMIRTranslation
     MLIRCallInterfaces

diff --git a/mlir/examples/toy/Ch6/toyc.cpp b/mlir/examples/toy/Ch6/toyc.cpp
@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
 #include "toy/Dialect.h"
 #include "toy/MLIRGen.h"
 #include "toy/Parser.h"
@@ -289,8 +290,10 @@ int main(int argc, char **argv) {
     return dumpAST();
 
   // If we aren't dumping the AST, then we are compiling with/to MLIR.
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
 
-  mlir::MLIRContext context;
+  mlir::MLIRContext context(registry);
   // Load our Dialect in this MLIR Context.
   context.getOrLoadDialect<mlir::toy::ToyDialect>();
 

diff --git a/mlir/examples/toy/Ch7/CMakeLists.txt b/mlir/examples/toy/Ch7/CMakeLists.txt
@@ -38,10 +38,12 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR})
 include_directories(${CMAKE_CURRENT_BINARY_DIR}/include/)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 target_link_libraries(toyc-ch7
   PRIVATE
     ${dialect_libs}
     ${conversion_libs}
+    ${extension_libs}
     MLIRAnalysis
     MLIRBuiltinToLLVMIRTranslation
     MLIRCallInterfaces

diff --git a/mlir/examples/toy/Ch7/toyc.cpp b/mlir/examples/toy/Ch7/toyc.cpp
@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Func/Extensions/AllExtensions.h"
 #include "toy/Dialect.h"
 #include "toy/MLIRGen.h"
 #include "toy/Parser.h"
@@ -290,8 +291,10 @@ int main(int argc, char **argv) {
     return dumpAST();
 
   // If we aren't dumping the AST, then we are compiling with/to MLIR.
+  mlir::DialectRegistry registry;
+  mlir::func::registerAllExtensions(registry);
 
-  mlir::MLIRContext context;
+  mlir::MLIRContext context(registry);
   // Load our Dialect in this MLIR Context.
   context.getOrLoadDialect<mlir::toy::ToyDialect>();
 

diff --git a/mlir/include/mlir/Dialect/Func/Extensions/AllExtensions.h b/mlir/include/mlir/Dialect/Func/Extensions/AllExtensions.h
@@ -0,0 +1,30 @@
+//===- AllExtensions.h - All Func Extensions --------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines a common entry point for registering all extensions to the
+// func dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_FUNC_EXTENSIONS_ALLEXTENSIONS_H
+#define MLIR_DIALECT_FUNC_EXTENSIONS_ALLEXTENSIONS_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace func {
+/// Register all extensions of the func dialect. This should generally only be
+/// used by tools, or other use cases that really do want *all* extensions of
+/// the dialect. All other cases should prefer to instead register the specific
+/// extensions they intend to take advantage of.
+void registerAllExtensions(DialectRegistry &registry);
+} // namespace func
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_FUNC_EXTENSIONS_ALLEXTENSIONS_H
diff --git a/mlir/include/mlir/Dialect/Func/Extensions/InlinerExtension.h b/mlir/include/mlir/Dialect/Func/Extensions/InlinerExtension.h
@@ -0,0 +1,27 @@
+//===- InlinerExtension.h - Func Inliner Extension 0000----------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines an extension for the func dialect that implements the
+// interfaces necessary to support inlining.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_FUNC_EXTENSIONS_INLINEREXTENSION_H
+#define MLIR_DIALECT_FUNC_EXTENSIONS_INLINEREXTENSION_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace func {
+/// Register the extension used to support inlining the func dialect.
+void registerInlinerExtension(DialectRegistry &registry);
+} // namespace func
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_FUNC_EXTENSIONS_INLINEREXTENSION_H
diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -21,7 +21,6 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 def Func_Dialect : Dialect {
   let name = "func";
   let cppNamespace = "::mlir::func";
-  let dependentDialects = ["cf::ControlFlowDialect"];
   let hasConstantMaterializer = 1;
   let usePropertiesForAttributes = 1;
 }