[mlir] Make ConversionTarget dynamic legality callbacks composable

[Hardcode84]

Summary

• Change callback signature bool(Operation *) -> Optional<bool>(Operation *)
• addDynamicallyLegalOp add callback to the chain
• If callback returned empty Optional next callback in chain will be called

[Hardcode84]

Old discourse thread https://llvm.discourse.group/t/markunknownopdynamicallylegal-is-not-composable/3800

Initially I wanted to do this only for markUnknownOpDynamicallyLegal, but now I have usecase for addDynamicallyLegalOp too:
I have getitem op in my dialect which can be applied to arrays or tuples with different set of conversions and legality setup, and I want to setup them from different places.

[Hardcode84]

Remove debug print

[Hardcode84]

fix clang-tidy warnings

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

664

I would consider Optional<bool> to be a bit of anti-pattern. Can we switch to an enum/Optional<enum> instead?

mlir/lib/Transforms/Utils/DialectConversion.cpp

2728–2732

2745

nit: Spell out auto here.

[Hardcode84]

mlir/include/mlir/Transforms/DialectConversion.h

664

I think Optional<smth> approach is in line with other APIs (e.g. TypeConverter) and Optional<bool> allow to use existing bool callbacks as is

[Hardcode84]

rebase, address some review comments, more tests

[Hardcode84]

@rriddle ping

[rriddle]

I am still iffy about the Optional<bool>, but am fine with it for now for consistency sake.

mlir/lib/Transforms/Utils/DialectConversion.cpp

2712–2714

nit: Spell out auto here, and also merge the variable into the if.

mlir/unittests/Transforms/DialectConversion.cpp

19–24

nit: Pull the function out of the anonymous namespace.

mlir/include/mlir/Transforms/DialectConversion.h

struct LegalOpDetails {
     /// A flag that indicates if this operation is 'recursively' legal. This
     /// means that if an operation is legal, either statically or dynamically,
     /// 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 *)>;
+  using DynamicLegalityCallbackFn = std::function<Optional<bool>(Operation *)>;
 
   ConversionTarget(MLIRContext &ctx) : ctx(ctx) {}
   virtual ~ConversionTarget() = default;
 
   //===--------------------------------------------------------------------===//
   // Legality Registration
   //===--------------------------------------------------------------------===//
 
private:
   /// 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;
+    LegalizationAction action = LegalizationAction::Illegal;
 
     /// If some legal instances of this operation may also be recursively legal.
-    bool isRecursivelyLegal;
+    bool isRecursivelyLegal = false;
 
     /// The legality callback if this operation is dynamically legal.
     DynamicLegalityCallbackFn legalityFn;
   };
 
   /// Get the legalization information for the given operation.
   Optional<LegalizationInfo> getOpInfo(OperationName op) const;
 

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, nullptr};
+  legalOperations[op].action = action;
 }
 
 /// 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
     -> 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 either with the provided legality function.
-    if (info->action == LegalizationAction::Dynamic)
-      return info->legalityFn(op);
+    if (info->action == LegalizationAction::Dynamic) {
+      auto result = info->legalityFn(op);
+      if (result)
+        return *result;
+    }
 
     // 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
+    if (legalityFnIt != opRecursiveLegalityFns.end()) {
+      legalityDetails.isRecursivelyLegal =
+          legalityFnIt->second(op).getValueOr(true);
+    } else {
       legalityDetails.isRecursivelyLegal = true;
     }
+  }
   return legalityDetails;
 }
 
+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> {
+    if (Optional<bool> result = newCl(op))
+      return *result;
+
+    return oldCl(op);
+  };
+  return chain;
+}
+
 /// Set the dynamic legality callback for the given operation.
 void ConversionTarget::setLegalityCallback(
     OperationName name, const DynamicLegalityCallbackFn &callback) {
   assert(callback && "expected valid legality 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;
+  infoIt->second.legalityFn =
+      composeLegalityCallbacks(std::move(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() &&
          infoIt->second.action != LegalizationAction::Illegal &&
          "expected operation to already be marked as legal");
   infoIt->second.isRecursivelyLegal = true;
   if (callback)
-    opRecursiveLegalityFns[name] = callback;
+    opRecursiveLegalityFns[name] = composeLegalityCallbacks(
+        std::move(opRecursiveLegalityFns[name]), callback);
   else
     opRecursiveLegalityFns.erase(name);
 }
 
 /// 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;
+    dialectLegalityFns[dialect] = composeLegalityCallbacks(
+        std::move(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;
+  unknownLegalityFn = composeLegalityCallbacks(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())

mlir/unittests/CMakeLists.txt

 set_target_properties(MLIRUnitTests PROPERTIES FOLDER "MLIR Tests")
 
 function(add_mlir_unittest test_dirname)
   add_unittest(MLIRUnitTests ${test_dirname} ${ARGN})
 endfunction()
 
 add_subdirectory(Analysis)
 add_subdirectory(Dialect)
 add_subdirectory(ExecutionEngine)
 add_subdirectory(Interfaces)
 add_subdirectory(IR)
 add_subdirectory(Pass)
 add_subdirectory(Rewrite)
 add_subdirectory(TableGen)
+add_subdirectory(Transforms)

mlir/unittests/Transforms/CMakeLists.txt

+add_mlir_unittest(MLIRTransformsTests
+  DialectConversion.cpp
+)
+target_link_libraries(MLIRTransformsTests
+  PRIVATE
+  MLIRTransforms)

mlir/unittests/Transforms/DialectConversion.cpp

+//===- DialectConversion.cpp - Dialect conversion unit tests --------------===//
+//
+// Part of the LLVM Project, 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/DialectConversion.h"
+#include "gtest/gtest.h"
+
+using namespace mlir;
+
+namespace {
+struct DummyOp {
+  static StringRef getOperationName() { return "foo.bar"; }
+};
+
+static Operation *createOp(MLIRContext *context) {
+  context->allowUnregisteredDialects();
+  return Operation::create(UnknownLoc::get(context),
+                           OperationName("foo.bar", context), llvm::None,
+                           llvm::None, llvm::None, llvm::None, 0);
+}
+
+TEST(DialectConversionTest, DynamicallyLegalOpCallbackOrder) {
+  MLIRContext context;
+  ConversionTarget target(context);
+
+  int index = 0;
+  int callbackCalled1 = 0;
+  target.addDynamicallyLegalOp<DummyOp>([&](Operation *) {
+    callbackCalled1 = ++index;
+    return true;
+  });
+
+  int callbackCalled2 = 0;
+  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);
+  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);
+  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);
+  op->destroy();
+}
+} // namespace