[mlir] ConversionTarget legality callbacks refactoring

[Hardcode84]

Summary

• Get rid of Optional<std::function> as std::function already have a null state
• Add private setLegalityCallback function to set legality callback for unknown ops
• Get rid of unknownOpsDynamicallyLegal flag, use unknownLegalityFn state insted. This causes behavior change when user first calls markUnknownOpDynamicallyLegal with callback and then without but I am not sure is the original behavior was really a 'feature', or just oversignt in the original implementation.

[Hardcode84]

@rriddle ping

[rriddle]

Sorry, I've been OOO for a while. Looks great, thanks!

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)
-      : unknownOpsDynamicallyLegal(false), ctx(ctx) {}
+  ConversionTarget(MLIRContext &ctx) : ctx(ctx) {}
   virtual ~ConversionTarget() = default;
 
   //===--------------------------------------------------------------------===//
   // Legality Registration
   //===--------------------------------------------------------------------===//
 
   /// Register a legality action for the given operation.
   void setOpAction(OperationName op, LegalizationAction action);
public:
   /// Register the operations of the given dialects as dynamically legal, i.e.
   /// requiring custom handling by the target via 'isDynamicallyLegal'.
   template <typename... Names>
   void addDynamicallyLegalDialect(StringRef name, Names... names) {
     SmallVector<StringRef, 2> dialectNames({name, names...});
     setDialectAction(dialectNames, LegalizationAction::Dynamic);
   }
   template <typename... Args>
-  void addDynamicallyLegalDialect(
-      Optional<DynamicLegalityCallbackFn> callback = llvm::None) {
+  void addDynamicallyLegalDialect(DynamicLegalityCallbackFn callback = {}) {
     SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
     setDialectAction(dialectNames, LegalizationAction::Dynamic);
     if (callback)
-      setLegalityCallback(dialectNames, *callback);
-  }
-  template <typename... Args>
-  void addDynamicallyLegalDialect(DynamicLegalityCallbackFn callback) {
-    SmallVector<StringRef, 2> dialectNames({Args::getDialectNamespace()...});
-    setDialectAction(dialectNames, LegalizationAction::Dynamic);
       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;
+    setLegalityCallback(fn);
+  }
+  void markUnknownOpDynamicallyLegal() {
+    setLegalityCallback([](Operation *) { return true; });
   }
-  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);
   }
private:
   /// Set the dynamic legality callback for the given operation.
   void setLegalityCallback(OperationName name,
                            const DynamicLegalityCallbackFn &callback);
 
   /// Set the dynamic legality callback for the given dialects.
   void setLegalityCallback(ArrayRef<StringRef> dialects,
                            const DynamicLegalityCallbackFn &callback);
 
+  /// Set the dynamic legality callback for the unknown ops.
+  void setLegalityCallback(const DynamicLegalityCallbackFn &callback);
+
   /// Set the recursive legality callback for the given operation and mark the
   /// operation as recursively legal.
   void markOpRecursivelyLegal(OperationName name,
                               const DynamicLegalityCallbackFn &callback);
 
   /// 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;
+    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 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;
+  DynamicLegalityCallbackFn unknownLegalityFn;
 
   /// The current context this target applies to.
   MLIRContext &ctx;
 };
 
 //===----------------------------------------------------------------------===//
 // Op Conversion Entry Points
 //===----------------------------------------------------------------------===//

mlir/lib/Transforms/Utils/DialectConversion.cpp

 
 //===----------------------------------------------------------------------===//
 // ConversionTarget
 //===----------------------------------------------------------------------===//
 
 /// Register a legality action for the given operation.
 void ConversionTarget::setOpAction(OperationName op,
                                    LegalizationAction action) {
-  legalOperations[op] = {action, /*isRecursivelyLegal=*/false, llvm::None};
+  legalOperations[op] = {action, /*isRecursivelyLegal=*/false, nullptr};
 }
 
 /// Register a legality action for the given dialects.
 void ConversionTarget::setDialectAction(ArrayRef<StringRef> dialectNames,
                                         LegalizationAction action) {
   for (StringRef dialect : dialectNames)
     legalDialects[dialect] = action;
 }
auto ConversionTarget::isLegal(Operation *op) const
   if (!info)
     return llvm::None;
 
   // Returns true if this operation instance is known to be legal.
   auto isOpLegal = [&] {
     // 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);
+      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.

 /// Set the dynamic legality callback for the given dialects.
 void ConversionTarget::setLegalityCallback(
     ArrayRef<StringRef> dialects, const DynamicLegalityCallbackFn &callback) {
   assert(callback && "expected valid legality callback");
   for (StringRef dialect : dialects)
     dialectLegalityFns[dialect] = callback;
 }
 
+/// Set the dynamic legality callback for the unknown ops.
+void ConversionTarget::setLegalityCallback(
+    const DynamicLegalityCallbackFn &callback) {
+  assert(callback && "expected valid legality callback");
+  unknownLegalityFn = callback;
+}
+
 /// 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;
   // Check for info for the parent dialect.
   auto dialectIt = legalDialects.find(op.getDialectNamespace());
   if (dialectIt != legalDialects.end()) {
-    Optional<DynamicLegalityCallbackFn> callback;
+    DynamicLegalityCallbackFn callback;
     auto dialectFn = dialectLegalityFns.find(op.getDialectNamespace());
     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)
+  if (unknownLegalityFn)
     return LegalizationInfo{LegalizationAction::Dynamic,
                             /*isRecursivelyLegal=*/false, unknownLegalityFn};
   return llvm::None;
 }
 
 //===----------------------------------------------------------------------===//
 // Op Conversion Entry Points
 //===----------------------------------------------------------------------===//