[Mlir-commits] [mlir] [mlir-c] Add ConversionTarget dynamic legality C API (PR #206161)

[Maksim Levental]

================
@@ -623,6 +623,185 @@ void testCloneWithMapping(MlirContext ctx) {
   fprintf(stderr, "testCloneWithMapping: PASSED\n");
 }
 
+static MlirConversionTargetLegality dynamicLegalityAlwaysLegal(MlirOperation op,
+                                                               void *userData) {
+  (void)op;
+  intptr_t *counter = (intptr_t *)userData;
+  (*counter)++;
+  return MLIR_CONVERSION_TARGET_LEGALITY_LEGAL;
+}
+
+static MlirConversionTargetLegality
+dynamicLegalityAlwaysIllegal(MlirOperation op, void *userData) {
+  (void)op;
+  intptr_t *counter = (intptr_t *)userData;
+  (*counter)++;
+  return MLIR_CONVERSION_TARGET_LEGALITY_ILLEGAL;
+}
+
+static MlirConversionTargetLegality dynamicLegalityNoOpinion(MlirOperation op,
+                                                             void *userData) {
+  (void)op;
+  intptr_t *counter = (intptr_t *)userData;
+  (*counter)++;
+  return MLIR_CONVERSION_TARGET_LEGALITY_NO_OPINION;
+}
+
+// Runs a partial conversion of `moduleString` against `target` with an empty
+// pattern set and returns whether it succeeded. This is what actually drives
+// the registered dynamic-legality callbacks.
+static bool runPartialConversion(MlirContext ctx, const char *moduleString,
+                                 MlirConversionTarget target) {
+  MlirModule module =
+      mlirModuleCreateParse(ctx, mlirStringRefCreateFromCString(moduleString));
+  assert(!mlirModuleIsNull(module) && "expected module to parse");
+  MlirOperation moduleOp = mlirModuleGetOperation(module);
+
+  MlirRewritePatternSet patterns = mlirRewritePatternSetCreate(ctx);
+  MlirFrozenRewritePatternSet frozen = mlirFreezeRewritePattern(patterns);
+  MlirConversionConfig config = mlirConversionConfigCreate();
+
+  MlirLogicalResult result =
+      mlirApplyPartialConversion(moduleOp, target, frozen, config);
+
+  mlirConversionConfigDestroy(config);
+  mlirFrozenRewritePatternSetDestroy(frozen);
+  mlirModuleDestroy(module);
+
+  return mlirLogicalResultIsSuccess(result);
+}
+
+void testConversionTargetDynamicLegality(MlirContext ctx) {
+  // CHECK-LABEL: @testConversionTargetDynamicLegality
+  fprintf(stderr, "@testConversionTargetDynamicLegality\n");
+
+  const char *opModule = "\"dialect.op1\"() : () -> ()\n";
+
+  // addDynamicallyLegalOp: callback returning true makes the op legal, so the
+  // (pattern-free) partial conversion succeeds and the callback is invoked.
+  {
+    MlirConversionTarget target = mlirConversionTargetCreate(ctx);
+    intptr_t counter = 0;
+    mlirConversionTargetAddDynamicallyLegalOp(
+        target, mlirStringRefCreateFromCString("dialect.op1"),
+        dynamicLegalityAlwaysLegal, &counter);
+    assert(runPartialConversion(ctx, opModule, target));
+    assert(counter > 0 && "legality callback must be invoked");
+    mlirConversionTargetDestroy(target);
+  }
+
+  // addDynamicallyLegalOp: callback returning false makes the op illegal. With
+  // no pattern to legalize it, the partial conversion fails -- proving the
+  // callback's return value actually drives the result.
+  {
+    MlirConversionTarget target = mlirConversionTargetCreate(ctx);
+    intptr_t counter = 0;
+    mlirConversionTargetAddDynamicallyLegalOp(
+        target, mlirStringRefCreateFromCString("dialect.op1"),
+        dynamicLegalityAlwaysIllegal, &counter);
+    assert(!runPartialConversion(ctx, opModule, target));
+    assert(counter > 0 && "legality callback must be invoked");
+    mlirConversionTargetDestroy(target);
+  }
+
+  // addDynamicallyLegalOp composition: callbacks registered for the same op are
+  // chained, most-recent first. A callback returning NoOpinion abstains and
+  // defers to the previously-registered callback. Here the first callback marks
+  // the op illegal and the second abstains, so the op stays illegal (conversion
+  // fails) and BOTH callbacks are invoked.
+  {
+    MlirConversionTarget target = mlirConversionTargetCreate(ctx);
+    intptr_t illegalCounter = 0;
+    intptr_t noOpinionCounter = 0;
+    mlirConversionTargetAddDynamicallyLegalOp(
+        target, mlirStringRefCreateFromCString("dialect.op1"),
+        dynamicLegalityAlwaysIllegal, &illegalCounter);
+    mlirConversionTargetAddDynamicallyLegalOp(
+        target, mlirStringRefCreateFromCString("dialect.op1"),
+        dynamicLegalityNoOpinion, &noOpinionCounter);
+    assert(!runPartialConversion(ctx, opModule, target));
+    assert(noOpinionCounter > 0 && "abstaining callback must be invoked");
+    assert(illegalCounter > 0 && "deferred-to callback must be invoked");
+    mlirConversionTargetDestroy(target);
+  }
----------------
makslevental wrote:

@PragmaTwice here's the test for the tri-state

https://github.com/llvm/llvm-project/pull/206161