[mlir][irdl] Add IRDL registration

[math-fehr]

Summary

This patch add support for loading IRDL dialects at runtime
with mlir-opt.

Given the following dialect.irdl file:

mlir
module {
  irdl.dialect @cmath {
    irdl.type @complex {
      %0 = irdl.is f32
      %1 = irdl.is f64
      %2 = irdl.any_of(%0, %1)
      irdl.parameters(%2)
    }

    irdl.operation @norm {
      %0 = irdl.any
      %1 = irdl.parametric @complex<%0>
      irdl.operands(%1)
      irdl.results(%0)
    }
}

the IRDL file can be loaded with the mlir-opt --irdl-file=dialect.irdl
command, and the following file can then be parsed:

mlir
func.func @conorm(%p: !cmath.complex<f32>, %q: !cmath.complex<f32>) -> f32 {
  %norm_p = "cmath.norm"(%p) : (!cmath.complex<f32>) -> f32
  %norm_q = "cmath.norm"(%q) : (!cmath.complex<f32>) -> f32
  %pq = arith.mulf %norm_p, %norm_q : f32
  return %pq : f32
}

To minimize the size of this patch, the operation, attribute, and type verifier are all always returning success().

Depends on D144692

[Mogball]

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

327

please spell out this auto

329

please emit this as an error through the MLIR context, using an UnknownLoc

346

364

please use empty

[Mogball]

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

351

why is this here?

[rriddle]

This patch seems to be doing a lot of things at once. I would split out registration in mlir-opt from the various IRDL things being added, and make things a bit more incremental, given the IRDL constructs could be tested via a test-pass (for example).

[Mogball]

mlir/include/mlir/Dialect/IRDL/IRDLVerifiers.h

37 ↗

(On Diff #500034)

please put public first and then private class members

76 ↗

(On Diff #500034)

same here

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

29 ↗

(On Diff #500034)

please use static instead of placing file-scope functions in an anonymous namespace

35 ↗

(On Diff #500034)

why append instead of <<?

43 ↗

(On Diff #500034)

55 ↗

(On Diff #500034)

we typically prefer unsigned to size_t

58 ↗

(On Diff #500034)

please spell out this auto, auto should only be used if the type is present verbatim on the right hand side of the assignment. Please audit all your autos and apply this change

mlir/lib/Dialect/IRDL/IRDLVerifiers.cpp

27 ↗

(On Diff #500034)

function_ref can be null

[rriddle]

mlir/lib/Dialect/IRDL/IR/IRDLOps.cpp

27 ↗

(On Diff #500034)

Prefer std::optional over llvm::Optional.

40 ↗

(On Diff #500034)

Can you use enumerate here instead?

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

25 ↗

(On Diff #500034)

Use static instead of anonymous namespaces for functions.

423 ↗

(On Diff #500034)

Use namespace qualifiers instead of wrapping namespaces in .cpp files: https://llvm.org/docs/CodingStandards.html#use-namespace-qualifiers-to-implement-previously-declared-functions

Same throughout the patch.

[math-fehr]

In D144693#4172489, @rriddle wrote:

This patch seems to be doing a lot of things at once. I would split out registration in mlir-opt from the various IRDL things being added, and make things a bit more incremental, given the IRDL constructs could be tested via a test-pass (for example).

Do you mean removing the MLIROptMain change for another patch, and then test the rest of the patch in a test-pass (I guess I can add features one by one, like first registration of types/attributes, then of operations).

If I'm correct, you are proposing to have a pass similar to the PDL one that register rewrites that are part of the file (with the -test-pdl-bytecode-pass pass). However, I couldn't do something similar with IRDL, since we cannot parse the entire file until the attributes are registered from what I understand.

Is that what you meant? I probably misunderstood something

[Mogball]

In D144693#4172706, @math-fehr wrote:

In D144693#4172489, @rriddle wrote:

This patch seems to be doing a lot of things at once. I would split out registration in mlir-opt from the various IRDL things being added, and make things a bit more incremental, given the IRDL constructs could be tested via a test-pass (for example).

Do you mean removing the MLIROptMain change for another patch, and then test the rest of the patch in a test-pass (I guess I can add features one by one, like first registration of types/attributes, then of operations).

If I'm correct, you are proposing to have a pass similar to the PDL one that register rewrites that are part of the file (with the -test-pdl-bytecode-pass pass). However, I couldn't do something similar with IRDL, since we cannot parse the entire file until the attributes are registered from what I understand.

Is that what you meant? I probably misunderstood something

I think you should find a logical way to split this patch up in general. One way would be to do the MlirOptMain registration in a separate patch and slowly land the interfaces/verification stuff. At least, I think the op verification can be tested separately (unless the dynamic dialect has to be registered before the ops are loaded?)

[math-fehr]

I think you should find a logical way to split this patch up in general. One way would be to do the MlirOptMain registration in a separate patch and slowly land the interfaces/verification stuff. At least, I think the op verification can be tested separately (unless the dynamic dialect has to be registered before the ops are loaded?)

Is it possible to register new dialects during a pass?
Alternatively, what I can do is first add the MlirOptMain registration, which would register only the dialect, and then add type/attributes, then operations. That way, we still keep the same workflow, we just add features one by one.
Would that work for you?

Also, do you want me to split further D144692 as well?

[Mogball]

IMO that would make sense to me. Anything to land these changes incrementally.

[math-fehr]

Thanks! I'll fix these issues, and split that patch to smaller pieces, to make this more incremental. I'll do this first thing tomorrow!

[math-fehr]

Address comments, and remove AttributeWrapper

[math-fehr]

I addressed all the code style comments, and move some code away.
In order to make this increasingly incremental, I'll move out completely the any_of operation from irdl, which should remove quite a lot of code.

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

58 ↗

(On Diff #500034)

I should have changed all the ones in my code, besides these two exceptions:

• Lambdas, since their operands are spelled out anyway
• for (auto [i, value] : enumerate(myIterator))

Is that correct, or should these as well be spelled outs?

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

351

From my understanding, multithreading needs to be disabled when we register dialects.
I added a better comment at the location where we disable multithreading.
Otherwise, I can only inject the thread pool after the registerIRDL call, if you think this is better?

[math-fehr]

Remove any_of from irdl

[math-fehr]

I removed any_of from irdl, to simplify the registration.
I still plan on simplifying this patch further though, and split it up in more pieces.

[math-fehr]

Remove verification of constraints

This should make the patch way smaller.
Now, while operations, types, and attributes are registered, they
are unconstrained, so they will always verify.

I'll add the constraints on a separate patch.

[math-fehr]

The patch should be way smaller now, I removed the verification logic, and will add it in a subsequent patch.

[mehdi_amini]

In D144693#4174249, @math-fehr wrote:

Is it possible to register new dialects during a pass?

No it isn't. We should runtime_error on this.

[math-fehr]

In D144693#4183080, @mehdi_amini wrote:

In D144693#4174249, @math-fehr wrote:

Is it possible to register new dialects during a pass?

No it isn't. We should runtime_error on this.

Is there a way to add a check for this? Currently, I'm removing the multi-threading, but I guess that you would like a concrete check?

[mehdi_amini]

In D144693#4183083, @math-fehr wrote:

Is there a way to add a check for this?

The check exists for non IRDL cases here: https://github.com/llvm/llvm-project/blob/9e0c2626f3d4282a135da112320139d9fba32fb5/mlir/lib/IR/MLIRContext.cpp#L444-L451
(despite the name, it'll kick in even if one disables multi-threading)

but I guess that you would like a concrete check?

Yes :)

mlir/include/mlir/Dialect/IRDL/IRDLRegistration.h

23 ↗

(On Diff #503521)

I feel like you're using "register" when it is actually more "loading" here.
(I'm not sure IRDL has a concept equivalent to "registering" dialects?)

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

251

This method is deprecated and will be removed, can you rebase and use the MlirOptMainConfig instead?

mlir/lib/Dialect/IRDL/IR/IRDLOps.cpp

13 ↗

(On Diff #503521)

All these includes don't seem to belong to this patch anymore.

[mehdi_amini]

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

55 ↗

(On Diff #500034)

I strictly use signed int :)

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

351

You don't need to touch multi-threading to register/load dialects. These just aren't thread-safe APIs and shouldn't be used during PassManager execution

[math-fehr]

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

351

This is done during MlirOptMain though, before the PassManager is setup (from my understanding).
I recall having some non-deterministic bugs at some points, and disabling multithreading at the registration time was the solution.

So should this pass even with multithreading? Because what this function does is both a registration and a loading.

[mehdi_amini]

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

351

Multi-threading isn't a factor here: even if the thread pool has been created, there are no active threads in the context.
So unless your process is involving thread that's just fine.

[math-fehr]

Address additional comments

[mehdi_amini]

I feel like you're using "register" when it is actually more "loading" here. (I'm not sure IRDL has a concept equivalent to "registering" dialects?)

Coming back to this: how do we differentiate registering from loading in IRDL? Do you see some difference somehow?

[math-fehr]

In D144693#4186106, @mehdi_amini wrote:

I feel like you're using "register" when it is actually more "loading" here. (I'm not sure IRDL has a concept equivalent to "registering" dialects?)

Coming back to this: how do we differentiate registering from loading in IRDL? Do you see some difference somehow?

In IRDL, there are not really any registration, since the dialects are loaded right away, without going through the `DialectRegistry.
While there may be a way to allow registering without loading, I'm not completely sure what would be the actual benefits, and I'm not sure how easy this would be.

The way we *could* make this work with an actual benefit would be to keep the IRDL program in memory for the entire compilation process,
and to register our dialects in the DialectRegistry with just a pointer to the corresponding DialectOp. Then, we would actually do the
loading and the main registration logic only if the dialect is required.
This could work, though I didn't work on it since I'm not sure what would be the advantages yet.

mlir/include/mlir/Dialect/IRDL/IRDLRegistration.h

23 ↗

(On Diff #503521)

True, I changed it to "load'.
However, this is something I didn't really think about when writing the dynamic dialects,
so the dynamic dialects are still using the "register" terminology.

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

55 ↗

(On Diff #500034)

I moved them to unsigned in the meantime, what would be the best fit then?

[mehdi_amini]

In D144693#4186322, @math-fehr wrote:

In D144693#4186106, @mehdi_amini wrote:

I feel like you're using "register" when it is actually more "loading" here. (I'm not sure IRDL has a concept equivalent to "registering" dialects?)

Coming back to this: how do we differentiate registering from loading in IRDL? Do you see some difference somehow?

In IRDL, there are not really any registration, since the dialects are loaded right away, without going through the `DialectRegistry.
While there may be a way to allow registering without loading, I'm not completely sure what would be the actual benefits, and I'm not sure how easy this would be.

The main issue I was trying to solve when designing this in MLIR was to preserve consistency of the system: we can't run a pipeline that accidentally crash because a dialect wasn't loaded in the first place for example: passes will declare dialects they need as "dependent". We also didn't want to always load all dialects preemptively if not needed, and in particular to be allow to test everything (including the "getDependentDialect" for passes) we needed mlir-opt to not load any dialect that isn't in the input IR.

Now, I don't know how IRDL plays with passes actually? What is the kind of use-cases we'll have for IRDL in mlir-opt?

mlir/lib/Dialect/IRDL/IRDLRegistration.cpp

55 ↗

(On Diff #500034)

This is fine, you do as you prefer.

There is no rules in LLVM, so people like me will always used signed and others will prefer unsigned. So unfortunately we're inconsistent!

FYI, see my previous attempt to define a standard: https://discourse.llvm.org/t/rfc-coding-standards-prefer-int-for-regular-arithmetic-use-unsigned-only-for-bitmask-and-when-you-intend-to-rely-on-wrapping-behavior/52191

[math-fehr]

In D144693#4186106, @mehdi_amini wrote:

I feel like you're using "register" when it is actually more "loading" here. (I'm not sure IRDL has a concept equivalent to "registering" dialects?)

Coming back to this: how do we differentiate registering from loading in IRDL? Do you see some difference somehow?

In IRDL, there are not really any registration, since the dialects are loaded right away, without going through the `DialectRegistry.
While there may be a way to allow registering without loading, I'm not completely sure what would be the actual benefits, and I'm not sure how easy this would be.

The way we *could* make this work with an actual benefit would be to keep the IRDL program in memory for the entire compilation process,
and to register our dialects in the DialectRegistry with just a pointer to the corresponding DialectOp. Then, we would actually do the
loading and the main registration logic only if the dialect is required.
This could work, though I didn't work on it since I'm not sure what would be the advantages yet.

In D144693#4186792, @mehdi_amini wrote:

In D144693#4186322, @math-fehr wrote:

In D144693#4186106, @mehdi_amini wrote:

I feel like you're using "register" when it is actually more "loading" here. (I'm not sure IRDL has a concept equivalent to "registering" dialects?)

Coming back to this: how do we differentiate registering from loading in IRDL? Do you see some difference somehow?

In IRDL, there are not really any registration, since the dialects are loaded right away, without going through the `DialectRegistry.
While there may be a way to allow registering without loading, I'm not completely sure what would be the actual benefits, and I'm not sure how easy this would be.

The main issue I was trying to solve when designing this in MLIR was to preserve consistency of the system: we can't run a pipeline that accidentally crash because a dialect wasn't loaded in the first place for example: passes will declare dialects they need as "dependent". We also didn't want to always load all dialects preemptively if not needed, and in particular to be allow to test everything (including the "getDependentDialect" for passes) we needed mlir-opt to not load any dialect that isn't in the input IR.

Now, I don't know how IRDL plays with passes actually? What is the kind of use-cases we'll have for IRDL in mlir-opt?

My feeling is that we explicitly do not want a pass to depend on a specific IRDL dialect. In fact, this is already something that we do not want to have for DynamicDialect.
Having a pass relying on a dynamic dialect would be too dangerous and error-prone, and if you are already paying the cost of recompiling to write your pass, you are willing to pay the cost of compiling your dialect to C++ (in my opinion).
Also, manipulating IRDL dialects in C++ is a bit annoying, since you need to use the generic way of accessing operands/results/attributes, meaning that you have almost no guarantee that those operands/results/attributes exist in the first place.

However, having passes that use Traits or Interfaces to manipulate these dialects is something that is useful, though we would need to first add Traits and Interfaces to IRDL, which may be complex in some cases (since some heavily rely on C++ templates).

That being said, our objective with having IRDL in mlir-opt is to then be able to load PDL rewrites dynamically as well, to have a completely dynamic pipeline. Probably the transform dialect could be useful there as well.
This will of course require more work to make the interface nice enough for a user of mlir-opt.

[mehdi_amini]

My feeling is that we explicitly do not want a pass to depend on a specific IRDL dialect. In fact, this is already something that we do not want to have for DynamicDialect.
Having a pass relying on a dynamic dialect would be too dangerous and error-prone, and if you are already paying the cost of recompiling to write your pass, you are willing to pay the cost of compiling your dialect to C++ (in my opinion).
Also, manipulating IRDL dialects in C++ is a bit annoying, since you need to use the generic way of accessing operands/results/attributes, meaning that you have almost no guarantee that those operands/results/attributes exist in the first place.

Sure, but that means that no pass will every create entities (Operation/Type/Attribute) defined by an IRDL dialect? If you think otherwise: how would that happen?

That being said, our objective with having IRDL in mlir-opt is to then be able to load PDL rewrites dynamically as well, to have a completely dynamic pipeline. Probably the transform dialect could be useful there as well.
This will of course require more work to make the interface nice enough for a user of mlir-opt.

I'd have to think more about the whole consistency of the system, but we can punt on this for now :)

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

81

[mehdi_amini]

This LGTM, but @rriddle and @Mogball had concerned earlier, it would be great if they could take another look here!

[math-fehr]

In D144693#4187242, @mehdi_amini wrote:

Sure, but that means that no pass will every create entities (Operation/Type/Attribute) defined by an IRDL dialect? If you think otherwise: how would that happen?

For now, I'm fine with this, since there are currently no big benefit on defining a dialect in IRDL if you want to write a C++ pass.
However, I agree that in the future, we might want to lift this in order to allow passes to create entities defined in IRDL, especially once we manage to get traits and interfaces in IRDL.

From my understanding though, mlir::Pass require to specify the dependent dialects as class names. This would probably need to be changed (or we would need a separate function),
if we would like to support this for IRDL. This is because we do not have a class for IRDL dialects, so we cannot just insert the dialects in the registry without knowing in which file
the dialects are defined.

That being said, our objective with having IRDL in mlir-opt is to then be able to load PDL rewrites dynamically as well, to have a completely dynamic pipeline. Probably the transform dialect could be useful there as well.
This will of course require more work to make the interface nice enough for a user of mlir-opt.

I'd have to think more about the whole consistency of the system, but we can punt on this for now :)

[math-fehr]

Update comment (register -> load)

[rriddle]

Looks like a good start to me

mlir/include/mlir/Dialect/IRDL/IRDLLoading.h

17–18

Can you use forward declarations here instead?

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

86

Is the llvm:: here and above necessary?

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

50–51

Is this comment wrapped properly at 80? It looks prematurely split.

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

344–346

365–368

[math-fehr]

Address comments

[Mogball]

sorry. Forgot I was blocking on this. LGTM

[math-fehr]

Rebase to latests LLVM, and attempt to fix compilation

mlir/include/mlir/Dialect/IRDL/IRDLLoading.h

+//===- IRDLRegistration.h - IRDL registration -------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Manages the registration of MLIR objects from IRDL operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_IRDL_IRDLREGISTRATION_H
+#define MLIR_DIALECT_IRDL_IRDLREGISTRATION_H
+
+namespace mlir {
+struct LogicalResult;
+class ModuleOp;
+} // namespace mlir
+
+namespace mlir {
+namespace irdl {
+
+/// Load all the dialects defined in the module.
+LogicalResult loadDialects(ModuleOp op);
+
+} // namespace irdl
+} // namespace mlir
+
+#endif // MLIR_DIALECT_IRDL_IRDLREGISTRATION_H

mlir/include/mlir/IR/OpDefinition.h

 
 #ifndef MLIR_IR_OPDEFINITION_H
 #define MLIR_IR_OPDEFINITION_H
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
-#include <type_traits>
 #include <optional>
+#include <type_traits>
 
 namespace mlir {
 class Builder;
 class OpBuilder;
 
 /// This class implements `Optional` functionality for ParseResult. We don't
 /// directly use Optional here, because it provides an implicit conversion
 /// to 'bool' which we want to avoid. This class is used to implement tri-state

 class OneTypedResult {
 public:
   /// This class provides return value APIs for ops that are known to have a
   /// single result.  ResultType is the concrete type returned by getType().
   template <typename ConcreteType>
   class Impl
       : public TraitBase<ConcreteType, OneTypedResult<ResultType>::Impl> {
   public:
     mlir::TypedValue<ResultType> getResult() {
       return cast<mlir::TypedValue<ResultType>>(
           this->getOperation()->getResult(0));
     }
 
     /// If the operation returns a single value, then the Op can be implicitly
     /// converted to a Value. This yields the value of the only result.
     operator mlir::TypedValue<ResultType>() { return getResult(); }
 
static LogicalResult verifyTrait(Operation *op) {
       if (llvm::isa_and_nonnull<ParentOpTypes...>(op->getParentOp()))
         return success();
 
       return op->emitOpError()
              << "expects parent op "
              << (sizeof...(ParentOpTypes) != 1 ? "to be one of '" : "'")
              << llvm::ArrayRef({ParentOpTypes::getOperationName()...}) << "'";
     }
+
+    template <typename ParentOpType =
+                  std::tuple_element_t<0, std::tuple<ParentOpTypes...>>>
+    std::enable_if_t<sizeof...(ParentOpTypes) == 1, ParentOpType>
+    getParentOp() {
+      Operation *parent = this->getOperation()->getParentOp();
+      return llvm::cast<ParentOpType>(parent);
+    }
   };
 };
 
 /// A trait for operations that have an attribute specifying operand segments.
 ///
 /// Certain operations can have multiple variadic operands and their size
 /// relationship is not always known statically. For such cases, we need
 /// a per-op-instance specification to divide the operands into logical groups

mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h

public:
 
   /// Set the output format to bytecode instead of textual IR.
   MlirOptMainConfig &emitBytecode(bool emit) {
     emitBytecodeFlag = emit;
     return *this;
   }
   bool shouldEmitBytecode() const { return emitBytecodeFlag; }
 
+  /// Set the IRDL file to load before processing the input.
+  MlirOptMainConfig &setIrdlFile(StringRef file) {
+    irdlFileFlag = file;
+    return *this;
+  }
+  StringRef getIrdlFile() const { return irdlFileFlag; }
+
   /// Set the filename to use for logging actions, use "-" for stdout.
   MlirOptMainConfig &logActionsTo(StringRef filename) {
     logActionsToFlag = filename;
     return *this;
   }
   /// Get the filename to use for logging actions.
   StringRef getLogActionsTo() const { return logActionsToFlag; }
 
protected:
   bool allowUnregisteredDialectsFlag = false;
 
   /// Print the pipeline that will be run.
   bool dumpPassPipelineFlag = false;
 
   /// Emit bytecode instead of textual assembly when generating output.
   bool emitBytecodeFlag = false;
 
+  /// IRDL file to register before processing the input.
+  std::string irdlFileFlag = "";
+
   /// Log action execution to the given file (or "-" for stdout)
   std::string logActionsToFlag;
 
   /// Location Breakpoints to filter the action logging.
   std::vector<tracing::BreakpointManager *> logActionLocationFilter;
 
   /// The callback to populate the pass manager.
   std::function<LogicalResult(PassManager &)> passPipelineCallback;

 /// Perform the core processing behind `mlir-opt`.
 /// This API is deprecated, use the MlirOptMainConfig version above instead.
 LogicalResult MlirOptMain(
     llvm::raw_ostream &outputStream, std::unique_ptr<llvm::MemoryBuffer> buffer,
     PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
     bool splitInputFile, bool verifyDiagnostics, bool verifyPasses,
     bool allowUnregisteredDialects, bool preloadDialectsInContext = false,
     bool emitBytecode = false, bool explicitModule = true);
 
 /// Implementation for tools like `mlir-opt`.
 /// - toolName is used for the header displayed by `--help`.
 /// - registry should contain all the dialects that can be parsed in the source.
 /// - preloadDialectsInContext will trigger the upfront loading of all
 ///   dialects from the global registry in the MLIRContext. This option is
 ///   deprecated and will be removed soon.
 LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
                           DialectRegistry &registry,

mlir/lib/Dialect/IRDL/CMakeLists.txt

 add_mlir_dialect_library(MLIRIRDL
   IR/IRDL.cpp
+  IRDLLoading.cpp
 
   DEPENDS
   MLIRIRDLIncGen
   MLIRIRDLOpsIncGen
   MLIRIRDLTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRDialect
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/IRDL/IRDLLoading.cpp

+//===- IRDLLoading.cpp - IRDL dialect loading --------------------- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Manages the loading of MLIR objects from IRDL operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/IRDL/IRDLLoading.h"
+#include "mlir/Dialect/IRDL/IR/IRDL.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/SMLoc.h"
+
+using namespace mlir;
+using namespace mlir::irdl;
+
+/// Define and load an operation represented by a `irdl.operation`
+/// operation.
+static WalkResult loadOperation(OperationOp op, ExtensibleDialect *dialect) {
+  // IRDL does not support defining custom parsers or printers.
+  auto parser = [](OpAsmParser &parser, OperationState &result) {
+    return failure();
+  };
+  auto printer = [](Operation *op, OpAsmPrinter &printer, StringRef) {
+    printer.printGenericOp(op);
+  };
+
+  auto verifier = [](Operation *op) { return success(); };
+
+  // IRDL does not support defining regions.
+  auto regionVerifier = [](Operation *op) { return success(); };
+
+  auto opDef = DynamicOpDefinition::get(
+      op.getName(), dialect, std::move(verifier), std::move(regionVerifier),
+      std::move(parser), std::move(printer));
+  dialect->registerDynamicOp(std::move(opDef));
+
+  return WalkResult::advance();
+}
+
+/// Load all dialects in the given module, without loading any operation, type
+/// or attribute definitions.
+static DenseMap<DialectOp, ExtensibleDialect *> loadEmptyDialects(ModuleOp op) {
+  DenseMap<DialectOp, ExtensibleDialect *> dialects;
+  op.walk([&](DialectOp dialectOp) {
+    MLIRContext *ctx = dialectOp.getContext();
+    StringRef dialectName = dialectOp.getName();
+
+    DynamicDialect *dialect = ctx->getOrLoadDynamicDialect(
+        dialectName, [](DynamicDialect *dialect) {});
+
+    dialects.insert({dialectOp, dialect});
+  });
+  return dialects;
+}
+
+/// Preallocate type definitions objects with empty verifiers.
+/// This in particular allocates a TypeID for each type definition.
+static DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>>
+preallocateTypeDefs(ModuleOp op,
+                    DenseMap<DialectOp, ExtensibleDialect *> dialects) {
+  DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> typeDefs;
+  op.walk([&](TypeOp typeOp) {
+    ExtensibleDialect *dialect = dialects[typeOp.getParentOp()];
+    auto typeDef = DynamicTypeDefinition::get(
+        typeOp.getName(), dialect,
+        [](function_ref<InFlightDiagnostic()>, ArrayRef<Attribute>) {
+          return success();
+        });
+    typeDefs.try_emplace(typeOp, std::move(typeDef));
+  });
+  return typeDefs;
+}
+
+/// Preallocate attribute definitions objects with empty verifiers.
+/// This in particular allocates a TypeID for each attribute definition.
+static DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>>
+preallocateAttrDefs(ModuleOp op,
+                    DenseMap<DialectOp, ExtensibleDialect *> dialects) {
+  DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> attrDefs;
+  op.walk([&](AttributeOp attrOp) {
+    ExtensibleDialect *dialect = dialects[attrOp.getParentOp()];
+    auto attrDef = DynamicAttrDefinition::get(
+        attrOp.getName(), dialect,
+        [](function_ref<InFlightDiagnostic()>, ArrayRef<Attribute>) {
+          return success();
+        });
+    attrDefs.try_emplace(attrOp, std::move(attrDef));
+  });
+  return attrDefs;
+}
+
+LogicalResult mlir::irdl::loadDialects(ModuleOp op) {
+  // Preallocate all dialects, and type and attribute definitions.
+  // In particular, this allocates TypeIDs so type and attributes can have
+  // verifiers that refer to each other.
+  DenseMap<DialectOp, ExtensibleDialect *> dialects = loadEmptyDialects(op);
+  DenseMap<TypeOp, std::unique_ptr<DynamicTypeDefinition>> types =
+      preallocateTypeDefs(op, dialects);
+  DenseMap<AttributeOp, std::unique_ptr<DynamicAttrDefinition>> attrs =
+      preallocateAttrDefs(op, dialects);
+
+  // Define and load all operations.
+  WalkResult res = op.walk([&](OperationOp opOp) {
+    return loadOperation(opOp, dialects[opOp.getParentOp()]);
+  });
+  if (res.wasInterrupted())
+    return failure();
+
+  // Load all types in their dialects.
+  for (auto &pair : types) {
+    ExtensibleDialect *dialect = dialects[pair.first.getParentOp()];
+    dialect->registerDynamicType(std::move(pair.second));
+  }
+
+  // Load all attributes in their dialects.
+  for (auto &pair : attrs) {
+    ExtensibleDialect *dialect = dialects[pair.first.getParentOp()];
+    dialect->registerDynamicAttr(std::move(pair.second));
+  }
+
+  return success();
+}

mlir/lib/Tools/mlir-opt/CMakeLists.txt

 add_mlir_library(MLIROptLib
   MlirOptMain.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-opt
 
   LINK_LIBS PUBLIC
   MLIRBytecodeWriter
   MLIRDebug
   MLIRObservers
   MLIRPass
   MLIRParser
   MLIRPluginsLib
   MLIRSupport
+  MLIRIRDL
   )

mlir/lib/Tools/mlir-opt/MlirOptMain.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/Debug/Counter.h"
 #include "mlir/Debug/ExecutionContext.h"
 #include "mlir/Debug/Observers/ActionLogging.h"
+#include "mlir/Dialect/IRDL/IR/IRDL.h"
+#include "mlir/Dialect/IRDL/IRDLLoading.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Parser/Parser.h"
MlirOptMainConfigCLOptions() {
     static cl::opt<bool, /*ExternalStorage=*/true> dumpPassPipeline(
         "dump-pass-pipeline", cl::desc("Print the pipeline that will be run"),
         cl::location(dumpPassPipelineFlag), cl::init(false));
 
     static cl::opt<bool, /*ExternalStorage=*/true> emitBytecode(
         "emit-bytecode", cl::desc("Emit bytecode when generating output"),
         cl::location(emitBytecodeFlag), cl::init(false));
 
+    static cl::opt<std::string, /*ExternalStorage=*/true> irdlFile(
+        "irdl-file",
+        cl::desc("IRDL file to register before processing the input"),
+        cl::location(irdlFileFlag), cl::init(""), cl::value_desc("filename"));
+
     static cl::opt<bool, /*ExternalStorage=*/true> explicitModule(
         "no-implicit-module",
         cl::desc("Disable implicit addition of a top-level module op during "
                  "parsing"),
         cl::location(useExplicitModuleFlag), cl::init(false));
 
     static cl::opt<std::string, /*ExternalStorage=*/true> logActionsTo{
         "log-actions-to",
if (config.shouldEmitBytecode()) {
     AsmState asmState(op.get(), OpPrintingFlags(), /*locationMap=*/nullptr,
                       &fallbackResourceMap);
     op.get()->print(os, asmState);
     os << '\n';
   }
   return success();
 }
 
+LogicalResult loadIRDLDialects(StringRef irdlFile, MLIRContext &ctx) {
+  DialectRegistry registry;
+  registry.insert<irdl::IRDLDialect>();
+  ctx.appendDialectRegistry(registry);
+
+  // Set up the input file.
+  std::string errorMessage;
+  std::unique_ptr<MemoryBuffer> file = openInputFile(irdlFile, &errorMessage);
+  if (!file) {
+    emitError(UnknownLoc::get(&ctx)) << errorMessage;
+    return failure();
+  }
+
+  // Give the buffer to the source manager.
+  // This will be picked up by the parser.
+  SourceMgr sourceMgr;
+  sourceMgr.AddNewSourceBuffer(std::move(file), SMLoc());
+
+  SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &ctx);
+
+  // Parse the input file.
+  OwningOpRef<ModuleOp> module(parseSourceFile<ModuleOp>(sourceMgr, &ctx));
+
+  // Load IRDL dialects.
+  return irdl::loadDialects(module.get());
+}
+
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult processBuffer(raw_ostream &os,
                                    std::unique_ptr<MemoryBuffer> ownedBuffer,
                                    const MlirOptMainConfig &config,
                                    DialectRegistry &registry,
                                    llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   auto sourceMgr = std::make_shared<SourceMgr>();
   sourceMgr->AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
   if (threadPool)
     context.setThreadPool(*threadPool);
 
+  StringRef irdlFile = config.getIrdlFile();
+  if (!irdlFile.empty() && failed(loadIRDLDialects(irdlFile, context)))
+    return failure();
+
   // Parse the input file.
   if (config.shouldPreloadDialectsInContext())
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(config.shouldAllowUnregisteredDialects());
   if (config.shouldVerifyDiagnostics())
     context.printOpOnDiagnostic(false);
 
   InstallDebugHandler installDebugHandler(context, config);
 

mlir/test/Dialect/IRDL/test-cmath.mlir

+// RUN: mlir-opt %s --irdl-file=%S/cmath.irdl.mlir | mlir-opt --irdl-file=%S/cmath.irdl.mlir | FileCheck %s
+
+module {
+  // CHECK: func.func @conorm(%[[p:[^:]*]]: !cmath.complex<f32>, %[[q:[^:]*]]: !cmath.complex<f32>) -> f32 {
+  // CHECK:   %[[norm_p:[^ ]*]] = "cmath.norm"(%[[p]]) : (!cmath.complex<f32>) -> f32
+  // CHECK:   %[[norm_q:[^ ]*]] = "cmath.norm"(%[[q]]) : (!cmath.complex<f32>) -> f32
+  // CHECK:   %[[pq:[^ ]*]] = arith.mulf %[[norm_p]], %[[norm_q]] : f32
+  // CHECK:   return %[[pq]] : f32
+  // CHECK: }
+  func.func @conorm(%p: !cmath.complex<f32>, %q: !cmath.complex<f32>) -> f32 {
+    %norm_p = "cmath.norm"(%p) : (!cmath.complex<f32>) -> f32
+    %norm_q = "cmath.norm"(%q) : (!cmath.complex<f32>) -> f32
+    %pq = arith.mulf %norm_p, %norm_q : f32
+    return %pq : f32
+  }
+
+  // CHECK: func.func @conorm2(%[[p:[^:]*]]: !cmath.complex<f32>, %[[q:[^:]*]]: !cmath.complex<f32>) -> f32 {
+  // CHECK:   %[[pq:[^ ]*]] = "cmath.mul"(%[[p]], %[[q]]) : (!cmath.complex<f32>, !cmath.complex<f32>) -> !cmath.complex<f32>
+  // CHECK:   %[[conorm:[^ ]*]] = "cmath.norm"(%[[pq]]) : (!cmath.complex<f32>) -> f32
+  // CHECK:   return %[[conorm]] : f32
+  // CHECK: }
+  func.func @conorm2(%p: !cmath.complex<f32>, %q: !cmath.complex<f32>) -> f32 {
+    %pq = "cmath.mul"(%p, %q) : (!cmath.complex<f32>, !cmath.complex<f32>) -> !cmath.complex<f32>
+    %conorm = "cmath.norm"(%pq) : (!cmath.complex<f32>) -> f32
+    return %conorm : f32
+  }
+}

mlir/test/Dialect/IRDL/testd.mlir

+// RUN: mlir-opt %s --irdl-file=%S/testd.irdl.mlir -split-input-file -verify-diagnostics | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// Type or attribute constraint
+//===----------------------------------------------------------------------===//
+
+func.func @typeFitsType() {
+  // CHECK: "testd.any"() : () -> !testd.parametric<i32>
+  "testd.any"() : () -> !testd.parametric<i32>
+  return
+}
+
+// -----
+
+func.func @attrDoesntFitType() {
+  "testd.any"() : () -> !testd.parametric<"foo">
+  return
+}
+
+// -----
+
+func.func @attrFitsAttr() {
+  // CHECK: "testd.any"() : () -> !testd.attr_in_type_out<"foo">
+  "testd.any"() : () -> !testd.attr_in_type_out<"foo">
+  return
+}
+
+// -----
+
+func.func @typeFitsAttr() {
+  // CHECK: "testd.any"() : () -> !testd.attr_in_type_out<i32>
+  "testd.any"() : () -> !testd.attr_in_type_out<i32>
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Equality constraint
+//===----------------------------------------------------------------------===//
+
+func.func @succeededEqConstraint() {
+  // CHECK: "testd.eq"() : () -> i32
+  "testd.eq"() : () -> i32
+  return
+}
+
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Any constraint
+//===----------------------------------------------------------------------===//
+
+func.func @succeededAnyConstraint() {
+  // CHECK: "testd.any"() : () -> i32
+  "testd.any"() : () -> i32
+  // CHECK: "testd.any"() : () -> i64
+  "testd.any"() : () -> i64
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Dynamic base constraint
+//===----------------------------------------------------------------------===//
+
+func.func @succeededDynBaseConstraint() {
+  // CHECK: "testd.dynbase"() : () -> !testd.parametric<i32>
+  "testd.dynbase"() : () -> !testd.parametric<i32>
+  // CHECK: "testd.dynbase"() : () -> !testd.parametric<!testd.parametric<i32>>
+  "testd.dynbase"() : () -> !testd.parametric<!testd.parametric<i32>>
+  return
+}
+
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Dynamic parameters constraint
+//===----------------------------------------------------------------------===//
+
+func.func @succeededDynParamsConstraint() {
+  // CHECK: "testd.dynparams"() : () -> !testd.parametric<i32>
+  "testd.dynparams"() : () -> !testd.parametric<i32>
+  return
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Constraint variables
+//===----------------------------------------------------------------------===//
+
+func.func @succeededConstraintVars() {
+  // CHECK: "testd.constraint_vars"() : () -> (i32, i32)
+  "testd.constraint_vars"() : () -> (i32, i32)
+  return
+}
+
+// -----
+
+func.func @succeededConstraintVars2() {
+  // CHECK: "testd.constraint_vars"() : () -> (i64, i64)
+  "testd.constraint_vars"() : () -> (i64, i64)
+  return
+}