[mlir] Add support for marking 'unknown' operations as dynamically legal.

[rriddle]

Summary

This allows for providing a default "catchall" legality check that is not dependent on specific operations or dialects. For example, this can be useful to check legality based on the specific types of operation operands or results.

[merge_guards_bot]

Unit tests: fail. 62150 tests passed, 5 failed and 811 were skipped.

failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/docs/DialectConversion.md

MyTarget(MLIRContext &ctx) : ConversionTarget(ctx) {
 
     /// Mark `std.return` as dynamically legal.
     addDynamicallyLegalOp<ReturnOp>();
 
     /// Mark `std.return` as dynamically legal, but provide a specific legality
     /// callback.
     addDynamicallyLegalOp<ReturnOp>([](ReturnOp op) { ... });
 
+    /// Treat unknown operations, i.e. those without a legalization action
+    /// directly set, as dynamically legal.
+    markUnknownOpDynamicallyLegal();
+    markUnknownOpDynamicallyLegal([](Operation *op) { ... });
+
     //--------------------------------------------------------------------------
     // Marking an operation as illegal.
 
     /// All operations within the GPU dialect are illegal.
     addIllegalDialect<GPUDialect>();
 
     /// Mark `std.br` and `std.cond_br` as illegal.
     addIllegalOp<BranchOp, CondBranchOp>();

mlir/include/mlir/Transforms/DialectConversion.h

struct LegalOpDetails {
     /// all of the operations nested within are also considered legal.
     bool isRecursivelyLegal = false;
   };
 
   /// The signature of the callback used to determine if an operation is
   /// dynamically legal on the target.
   using DynamicLegalityCallbackFn = std::function<bool(Operation *)>;
 
-  ConversionTarget(MLIRContext &ctx) : ctx(ctx) {}
+  ConversionTarget(MLIRContext &ctx)
+      : unknownOpsDynamicallyLegal(false), ctx(ctx) {}
   virtual ~ConversionTarget() = default;
 
   //===--------------------------------------------------------------------===//
   // Legality Registration
   //===--------------------------------------------------------------------===//
 
   /// Register a legality action for the given operation.
   void setOpAction(OperationName op, LegalizationAction action);
public:
   void addDynamicallyLegalDialect(
       Optional<DynamicLegalityCallbackFn> callback = llvm::None) {
     SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
     setDialectAction(dialectNames, LegalizationAction::Dynamic);
     if (callback)
       setLegalityCallback(dialectNames, *callback);
   }
 
+  /// Register unknown operations as dynamically legal. For operations(and
+  /// dialects) that do not have a set legalization action, treat them as
+  /// dynamically legal and invoke the given callback if valid or
+  /// 'isDynamicallyLegal'.
+  void markUnknownOpDynamicallyLegal(const DynamicLegalityCallbackFn &fn) {
+    unknownOpsDynamicallyLegal = true;
+    unknownLegalityFn = fn;
+  }
+  void markUnknownOpDynamicallyLegal() { unknownOpsDynamicallyLegal = true; }
+
   /// Register the operations of the given dialects as illegal, i.e.
   /// operations of this dialect are not supported by the target.
   template <typename... Names>
   void addIllegalDialect(StringRef name, Names... names) {
     SmallVector<StringRef, 2> dialectNames({name, names...});
     setDialectAction(dialectNames, LegalizationAction::Illegal);
   }
   template <typename... Args> void addIllegalDialect() {
private:
 
   /// The set of information that configures the legalization of an operation.
   struct LegalizationInfo {
     /// The legality action this operation was given.
     LegalizationAction action;
 
     /// If some legal instances of this operation may also be recursively legal.
     bool isRecursivelyLegal;
+
+    /// The legality callback if this operation is dynamically legal.
+    Optional<DynamicLegalityCallbackFn> legalityFn;
   };
 
   /// Get the legalization information for the given operation.
   Optional<LegalizationInfo> getOpInfo(OperationName op) const;
 
   /// A deterministic mapping of operation name and its respective legality
   /// information.
   llvm::MapVector<OperationName, LegalizationInfo> legalOperations;
 
-  /// A set of dynamic legality callbacks for given operation names.
-  DenseMap<OperationName, DynamicLegalityCallbackFn> opLegalityFns;
-
   /// A set of legality callbacks for given operation names that are used to
   /// check if an operation instance is recursively legal.
   DenseMap<OperationName, DynamicLegalityCallbackFn> opRecursiveLegalityFns;
 
   /// A deterministic mapping of dialect name to the specific legality action to
   /// take.
   llvm::StringMap<LegalizationAction> legalDialects;
 
   /// A set of dynamic legality callbacks for given dialect names.
   llvm::StringMap<DynamicLegalityCallbackFn> dialectLegalityFns;
 
+  /// An optional legality callback for unknown operations.
+  Optional<DynamicLegalityCallbackFn> unknownLegalityFn;
+
+  /// Flag indicating if unknown operations should be treated as dynamically
+  /// legal.
+  bool unknownOpsDynamicallyLegal;
+
   /// The current context this target applies to.
   MLIRContext &ctx;
 };
 
 //===----------------------------------------------------------------------===//
 // Op Conversion Entry Points
 //===----------------------------------------------------------------------===//
 

mlir/lib/Transforms/DialectConversion.cpp

 auto ConversionTarget::isLegal(Operation *op) const
     -> Optional<LegalOpDetails> {
   Optional<LegalizationInfo> info = getOpInfo(op->getName());
   if (!info)
     return llvm::None;
 
   // Returns true if this operation instance is known to be legal.
   auto isOpLegal = [&] {
-    // Handle dynamic legality.
-    if (info->action == LegalizationAction::Dynamic) {
-      // Check for callbacks on the operation or dialect.
-      auto opFn = opLegalityFns.find(op->getName());
-      if (opFn != opLegalityFns.end())
-        return opFn->second(op);
-      auto dialectFn = dialectLegalityFns.find(op->getName().getDialect());
-      if (dialectFn != dialectLegalityFns.end())
-        return dialectFn->second(op);
-
-      // Otherwise, invoke the hook on the derived instance.
-      return isDynamicallyLegal(op);
-    }
+    // Handle dynamic legality either with the provided legality function, or
+    // the default hook on the derived instance.
+    if (info->action == LegalizationAction::Dynamic)
+      return info->legalityFn ? (*info->legalityFn)(op)
+                              : isDynamicallyLegal(op);
 
     // Otherwise, the operation is only legal if it was marked 'Legal'.
     return info->action == LegalizationAction::Legal;
   };
   if (!isOpLegal())
     return llvm::None;
 
   // This operation is legal, compute any additional legality information.
   LegalOpDetails legalityDetails;
-
   if (info->isRecursivelyLegal) {
     auto legalityFnIt = opRecursiveLegalityFns.find(op->getName());
     if (legalityFnIt != opRecursiveLegalityFns.end())
       legalityDetails.isRecursivelyLegal = legalityFnIt->second(op);
     else
       legalityDetails.isRecursivelyLegal = true;
   }
   return legalityDetails;
 }
 
 /// Set the dynamic legality callback for the given operation.
 void ConversionTarget::setLegalityCallback(
     OperationName name, const DynamicLegalityCallbackFn &callback) {
   assert(callback && "expected valid legality callback");
-  opLegalityFns[name] = callback;
+  auto infoIt = legalOperations.find(name);
+  assert(infoIt != legalOperations.end() &&
+         infoIt->second.action == LegalizationAction::Dynamic &&
+         "expected operation to already be marked as dynamically legal");
+  infoIt->second.legalityFn = callback;
 }
 
 /// Set the recursive legality callback for the given operation and mark the
 /// operation as recursively legal.
 void ConversionTarget::markOpRecursivelyLegal(
     OperationName name, const DynamicLegalityCallbackFn &callback) {
   auto infoIt = legalOperations.find(name);
   assert(infoIt != legalOperations.end() &&

 
 /// Get the legalization information for the given operation.
 auto ConversionTarget::getOpInfo(OperationName op) const
     -> Optional<LegalizationInfo> {
   // Check for info for this specific operation.
   auto it = legalOperations.find(op);
   if (it != legalOperations.end())
     return it->second;
-  // Otherwise, default to checking on the parent dialect.
+  // Check for info for the parent dialect.
   auto dialectIt = legalDialects.find(op.getDialect());
-  if (dialectIt != legalDialects.end())
-    return LegalizationInfo{dialectIt->second, /*isRecursivelyLegal=*/false};
+  if (dialectIt != legalDialects.end()) {
+    Optional<DynamicLegalityCallbackFn> callback;
+    auto dialectFn = dialectLegalityFns.find(op.getDialect());
+    if (dialectFn != dialectLegalityFns.end())
+      callback = dialectFn->second;
+    return LegalizationInfo{dialectIt->second, /*isRecursivelyLegal=*/false,
+                            callback};
+  }
+  // Otherwise, check if we mark unknown operations as dynamic.
+  if (unknownOpsDynamicallyLegal)
+    return LegalizationInfo{LegalizationAction::Dynamic,
+                            /*isRecursivelyLegal=*/false, unknownLegalityFn};
   return llvm::None;
 }
 
 //===----------------------------------------------------------------------===//
 // Op Conversion Entry Points
 //===----------------------------------------------------------------------===//
 
 /// Apply a partial conversion on the given operations, and all nested

mlir/test/Transforms/test-legalizer-full.mlir

"test.region"() ({
     ^bb1(%i0: i64):
       "test.invalid"(%i0) : (i64) -> ()
   }) {legalizer.should_clone} : () -> ()
 
   // expected-error@+1 {{failed to legalize operation 'test.illegal_op_f'}}
   %ignored = "test.illegal_op_f"() : () -> (i32)
   "test.return"() : () -> ()
 }
+
+// -----
+
+// Test that unknown operations can be dynamically legal.
+func @test_unknown_dynamically_legal() {
+  "foo.unknown_op"() {test.dynamically_legal} : () -> ()
+
+  // expected-error@+1 {{failed to legalize operation 'foo.unknown_op'}}
+  "foo.unknown_op"() {} : () -> ()
+  "test.return"() : () -> ()
+}

mlir/test/lib/TestDialect/TestPatterns.cpp

void runOnModule() override {
     // Handle a partial conversion.
     if (mode == ConversionMode::Partial) {
       (void)applyPartialConversion(getModule(), target, patterns, &converter);
       return;
     }
 
     // Handle a full conversion.
     if (mode == ConversionMode::Full) {
+      // Check support for marking unknown operations as dynamically legal.
+      target.markUnknownOpDynamicallyLegal([](Operation *op) {
+        return (bool)op->getAttrOfType<UnitAttr>("test.dynamically_legal");
+      });
+
       (void)applyFullConversion(getModule(), target, patterns, &converter);
       return;
     }
 
     // Otherwise, handle an analysis conversion.
     assert(mode == ConversionMode::Analysis);
 
     // Analyze the convertible operations.