[Mlir-commits] [mlir] 2a3878e - [mlir] DialectConversion: fix OperationLegalizer::isIllegal result when legality callback returns None

[llvmlistbot at llvm.org]

Author: Butygin
Date: 2021-11-15T14:53:06+03:00
New Revision: 2a3878ea164462caf0b37c46767242b77b66736f

URL: https://github.com/llvm/llvm-project/commit/2a3878ea164462caf0b37c46767242b77b66736f
DIFF: https://github.com/llvm/llvm-project/commit/2a3878ea164462caf0b37c46767242b77b66736f.diff

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

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.

Differential Revision: https://reviews.llvm.org/D113267

Added: 
    

Modified: 
    mlir/include/mlir/Transforms/DialectConversion.h
    mlir/lib/Transforms/Utils/DialectConversion.cpp
    mlir/unittests/Transforms/DialectConversion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Transforms/DialectConversion.h b/mlir/include/mlir/Transforms/DialectConversion.h
index 849d5427ed7c..2953a5d43c9a 100644
--- a/mlir/include/mlir/Transforms/DialectConversion.h
+++ b/mlir/include/mlir/Transforms/DialectConversion.h
@@ -793,9 +793,19 @@ class ConversionTarget {
 
   /// 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,

diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index cea4b2aaf80b..cc83ab353b06 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -1816,14 +1816,7 @@ OperationLegalizer::OperationLegalizer(ConversionTarget &targetInfo,
 }
 
 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
@@ -3137,6 +3130,22 @@ auto ConversionTarget::isLegal(Operation *op) const
   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) {

diff  --git a/mlir/unittests/Transforms/DialectConversion.cpp b/mlir/unittests/Transforms/DialectConversion.cpp
index d3e7ff6909e4..13ce43160f4e 100644
--- a/mlir/unittests/Transforms/DialectConversion.cpp
+++ b/mlir/unittests/Transforms/DialectConversion.cpp
@@ -44,6 +44,9 @@ TEST(DialectConversionTest, DynamicallyLegalOpCallbackOrder) {
   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();
 }
 
@@ -61,6 +64,8 @@ TEST(DialectConversionTest, DynamicallyLegalOpCallbackSkip) {
   auto *op = createOp(&context);
   EXPECT_FALSE(target.isLegal(op));
   EXPECT_EQ(1, callbackCalled);
+  EXPECT_FALSE(target.isIllegal(op));
+  EXPECT_EQ(2, callbackCalled);
   op->destroy();
 }
 
@@ -85,6 +90,43 @@ TEST(DialectConversionTest, DynamicallyLegalUnknownOpCallbackOrder) {
   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