[mlir] DialectConversion: fix OperationLegalizer::isIllegal result when legality callback returns None

[Hardcode84]

Summary

OperationLegalizer::isIllegal returns false if operation legality wasn't
registered by user and we expect same behaviour when dynamic legality
callback return None, but instead true was returned.

[Hardcode84]

fix comment

[rriddle]

mlir/lib/Transforms/Utils/DialectConversion.cpp

3139–3143

Does this work?

[Hardcode84]

mlir/lib/Transforms/Utils/DialectConversion.cpp

3139–3143

This seems to work, but IMO this is confusing. Without knowing the context someone can try to do 'obvious' refactoring to return !info->legalityFn(op) which will break this code.

mlir/include/mlir/Transforms/DialectConversion.h

public:
   // Legality Querying
   //===--------------------------------------------------------------------===//
 
   /// Get the legality action for the given operation.
   Optional<LegalizationAction> getOpAction(OperationName op) const;
 
   /// If the given operation instance is legal on this target, a structure
   /// containing legality information is returned. If the operation is not
-  /// legal, None is returned.
+  /// legal, None is returned. Also returns None is operation legality wasn't
+  /// registered by user or dynamic legality callbacks returned None.
+  ///
+  /// Note: Legality is actually a 4-state: Legal(recursive=true),
+  /// Legal(recursive=false), Illegal or Unknown, where Unknown is treated
+  /// either as Legal or Illegal depending on context.
   Optional<LegalOpDetails> isLegal(Operation *op) const;
 
+  /// Returns true is operation instance is illegal on this target. Returns
+  /// false if operation is legal, operation legality wasn't registered by user
+  /// or dynamic legality callbacks returned None.
+  bool isIllegal(Operation *op) const;
+
 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);

mlir/lib/Transforms/Utils/DialectConversion.cpp

OperationLegalizer::OperationLegalizer(ConversionTarget &targetInfo,
   DenseMap<OperationName, LegalizationPatterns> legalizerPatterns;
   LegalizationPatterns anyOpLegalizerPatterns;
 
   buildLegalizationGraph(anyOpLegalizerPatterns, legalizerPatterns);
   computeLegalizationGraphBenefit(anyOpLegalizerPatterns, legalizerPatterns);
 }
 
 bool OperationLegalizer::isIllegal(Operation *op) const {
-  // Check if the target explicitly marked this operation as illegal.
-  if (auto info = target.getOpAction(op->getName())) {
-    if (*info == LegalizationAction::Dynamic)
-      return !target.isLegal(op);
-    return *info == LegalizationAction::Illegal;
-  }
-
-  return false;
+  return target.isIllegal(op);
 }
 
 LogicalResult
 OperationLegalizer::legalize(Operation *op,
                              ConversionPatternRewriter &rewriter) {
 #ifndef NDEBUG
   const char *logLineComment =
       "//===-------------------------------------------===//\n";
if (legalityFnIt != opRecursiveLegalityFns.end()) {
           legalityFnIt->second(op).getValueOr(true);
     } else {
       legalityDetails.isRecursivelyLegal = true;
     }
   }
   return legalityDetails;
 }
 
+bool ConversionTarget::isIllegal(Operation *op) const {
+  Optional<LegalizationInfo> info = getOpInfo(op->getName());
+  if (!info)
+    return false;
+
+  if (info->action == LegalizationAction::Dynamic) {
+    Optional<bool> result = info->legalityFn(op);
+    if (!result)
+      return false;
+
+    return !(*result);
+  }
+
+  return info->action == LegalizationAction::Illegal;
+}
+
 static ConversionTarget::DynamicLegalityCallbackFn composeLegalityCallbacks(
     ConversionTarget::DynamicLegalityCallbackFn oldCallback,
     ConversionTarget::DynamicLegalityCallbackFn newCallback) {
   if (!oldCallback)
     return newCallback;
 
   auto chain = [oldCl = std::move(oldCallback), newCl = std::move(newCallback)](
                    Operation *op) -> Optional<bool> {

mlir/unittests/Transforms/DialectConversion.cpp

target.addDynamicallyLegalOp<DummyOp>([&](Operation *) -> Optional<bool> {
     callbackCalled2 = ++index;
     return llvm::None;
   });
 
   auto *op = createOp(&context);
   EXPECT_TRUE(target.isLegal(op));
   EXPECT_EQ(2, callbackCalled1);
   EXPECT_EQ(1, callbackCalled2);
+  EXPECT_FALSE(target.isIllegal(op));
+  EXPECT_EQ(4, callbackCalled1);
+  EXPECT_EQ(3, callbackCalled2);
   op->destroy();
 }
 
 TEST(DialectConversionTest, DynamicallyLegalOpCallbackSkip) {
   MLIRContext context;
   ConversionTarget target(context);
 
   int index = 0;
   int callbackCalled = 0;
   target.addDynamicallyLegalOp<DummyOp>([&](Operation *) -> Optional<bool> {
     callbackCalled = ++index;
     return llvm::None;
   });
 
   auto *op = createOp(&context);
   EXPECT_FALSE(target.isLegal(op));
   EXPECT_EQ(1, callbackCalled);
+  EXPECT_FALSE(target.isIllegal(op));
+  EXPECT_EQ(2, callbackCalled);
   op->destroy();
 }
 
 TEST(DialectConversionTest, DynamicallyLegalUnknownOpCallbackOrder) {
   MLIRContext context;
   ConversionTarget target(context);
 
   int index = 0;
   int callbackCalled1 = 0;
   target.markUnknownOpDynamicallyLegal([&](Operation *) {
     callbackCalled1 = ++index;
     return true;
   });
 
   int callbackCalled2 = 0;
   target.markUnknownOpDynamicallyLegal([&](Operation *) -> Optional<bool> {
     callbackCalled2 = ++index;
     return llvm::None;
   });
 
   auto *op = createOp(&context);
   EXPECT_TRUE(target.isLegal(op));
   EXPECT_EQ(2, callbackCalled1);
   EXPECT_EQ(1, callbackCalled2);
+  EXPECT_FALSE(target.isIllegal(op));
+  EXPECT_EQ(4, callbackCalled1);
+  EXPECT_EQ(3, callbackCalled2);
+  op->destroy();
+}
+
+TEST(DialectConversionTest, DynamicallyLegalReturnNone) {
+  MLIRContext context;
+  ConversionTarget target(context);
+
+  target.addDynamicallyLegalOp<DummyOp>(
+      [&](Operation *) -> Optional<bool> { return llvm::None; });
+
+  auto *op = createOp(&context);
+  EXPECT_FALSE(target.isLegal(op));
+  EXPECT_FALSE(target.isIllegal(op));
+
+  EXPECT_TRUE(succeeded(applyPartialConversion(op, target, {})));
+  EXPECT_TRUE(failed(applyFullConversion(op, target, {})));
+
+  op->destroy();
+}
+
+TEST(DialectConversionTest, DynamicallyLegalUnknownReturnNone) {
+  MLIRContext context;
+  ConversionTarget target(context);
+
+  target.markUnknownOpDynamicallyLegal(
+      [&](Operation *) -> Optional<bool> { return llvm::None; });
+
+  auto *op = createOp(&context);
+  EXPECT_FALSE(target.isLegal(op));
+  EXPECT_FALSE(target.isIllegal(op));
+
+  EXPECT_TRUE(succeeded(applyPartialConversion(op, target, {})));
+  EXPECT_TRUE(failed(applyFullConversion(op, target, {})));
+
   op->destroy();
 }
 } // namespace