[mlir][Transform] Add `updateConversionTarget` to `ConversionPatternDescriptorOpInterface`

[christopherbate]

Summary

This change adds a method to modify the ConversionTarget used during
transform.apply_conversion_patterns to the
ConversionPatternDescriptorOpInterface. This is needed when the TypeConverter
is used to dictate the dynamic legality of operations, as in "structural"
conversion patterns present in, for example, the SCF and func dialects.

As a first use case/test, this change also adds a
transform.apply_patterns.scf.structural_conversions operation to the SCF
dialect.

[christopherbate]

Add missing check statements

[springerm]

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

338–339

I'd mention structural type conversions as an example use case here. Mention that type converters are set up such that legal types are converted to themselves and an op is considered illegal (a property of the conversion target) if one of its result/operand types maps to a type different from itself. (To clarify the connection between the conversion target and the type converter; at first glance they seem unrelated.)

342

nit: I'd call it populateConversionTargetRules. (update may sound like the existing rules are overwritten)

343

Let's use a const reference here.

mlir/test/Dialect/SCF/transform-ops.mlir

310

new line missing

[nicolasvasilache]

mlir/test/Dialect/SCF/transform-ops.mlir

288

hmm I am a bit confused, what triggers a tensor -> memref conversion here ?

[springerm]

mlir/test/Dialect/SCF/transform-ops.mlir

288

The TestTypeConverter in the test case does that.

addConversion([](RankedTensorType type) -> Type {
  return MemRefType::get(type.getShape(), type.getElementType());
});

[christopherbate]

Address review comments

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

let description = [{
     Collects patterns for canonicalizing operations inside SCF loop bodies.
     At the moment, only affine.min/max computations with iteration variables,
     loop bounds and loop steps are canonicalized.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+def ApplySCFStructuralConversionPatternsOp : Op<Transform_Dialect,
+    "apply_conversion_patterns.scf.structural_conversions",
+    [DeclareOpInterfaceMethods<ConversionPatternDescriptorOpInterface,
+      ["populateConversionTargetRules"]>]> {
+  let description = [{
+    Collects patterns for performing structural conversions of SCF operations.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def Transform_ScfForOp : Transform_ConcreteOpType<"scf.for">;
 
 def GetParentForOp : Op<Transform_Dialect, "loop.get_parent_for",
     [NavigationTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Gets a handle to the parent 'for' loop of the given operation";
   let description = [{
     Produces a handle to the n-th (default 1) parent `scf.for` or `affine.for`
def LoopCoalesceOp : Op<Transform_Dialect, "loop.coalesce", [
   }];
 }
 
 def TakeAssumedBranchOp : Op<Transform_Dialect, "scf.take_assumed_branch", [
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
   TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Given an scf.if conditional, inject user-defined information that it is
-    always safe to execute only the if or else branch. 
-    
+    always safe to execute only the if or else branch.
     This is achieved by just replacing the scf.if by the content of one of its
     branches.
 
     This is particularly useful for user-controlled rewriting of conditionals
     that exist solely to guard against out-of-bounds behavior.
 
     At the moment, no assume or assert operation is emitted as it is not always
     desirable. In the future, this may be controlled by a dedicated attribute.

mlir/include/mlir/Dialect/SCF/Transforms/Patterns.h

 /// completely agnostic to the actual types involved and simply need to update
 /// their types. An example of this is scf.if -- the scf.if op and the
 /// corresponding scf.yield ops need to update their types accordingly to the
 /// TypeConverter, but otherwise don't care what type conversions are happening.
 void populateSCFStructuralTypeConversionsAndLegality(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     ConversionTarget &target);
 
+/// Similar to `populateSCFStructuralTypeConversionsAndLegality` but does not
+/// populate the conversion target.
+void populateSCFStructuralTypeConversions(TypeConverter &typeConverter,
+                                          RewritePatternSet &patterns);
+
+/// Updates the ConversionTarget with dynamic legality of SCF operations based
+/// on the provided type converter.
+void populateSCFStructuralTypeConversionTarget(
+    const TypeConverter &typeConverter, ConversionTarget &target);
+
 /// Populates the provided pattern set with patterns that do 1:N type
 /// conversions on (some) SCF ops. This is intended to be used with
 /// applyPartialOneToNConversion.
 void populateSCFStructuralOneToNTypeConversions(TypeConverter &typeConverter,
                                                 RewritePatternSet &patterns);
 
 /// Populate patterns for SCF software pipelining transformation. See the
 /// ForLoopPipeliningPattern for the transformation details.

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

InterfaceMethod<
       }],
       /*returnType=*/"void",
       /*name=*/"populatePatterns",
       /*arguments=*/(ins "::mlir::TypeConverter &":$typeConverter,
                          "::mlir::RewritePatternSet &":$patterns)
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Populate the ConversionTarget using the final TypeConverter. The default
+        implementation is to do nothing. Overriding this method can be useful
+        in order to setup the ConversionTarget for structural type conversions.
+        In such a situation, an op's legality depends on using the TypeConverter
+        to determine whether the op's operand and result types are legal
+        (defined as converting to themselves).
+
+      }],
+      /*returnType=*/"void",
+      /*name=*/"populateConversionTargetRules",
+      /*arguments=*/(ins "const ::mlir::TypeConverter &":$typeConverter,
+                         "::mlir::ConversionTarget &":$conversionTarget),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"return;"
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the type converter to be used with this pattern set. If no
         type converter is specified, the default type converter of the enclosing
         "apply_conversion_patterns" op is used.
       }],
       /*returnType=*/"std::unique_ptr<::mlir::TypeConverter>",
       /*name=*/"getTypeConverter",
       /*arguments=*/(ins),
       /*methodBody=*/"",

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

 // Apply...PatternsOp
 //===----------------------------------------------------------------------===//
 
 void transform::ApplyForLoopCanonicalizationPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   scf::populateSCFForLoopCanonicalizationPatterns(patterns);
 }
 
+void transform::ApplySCFStructuralConversionPatternsOp::populatePatterns(
+    TypeConverter &typeConverter, RewritePatternSet &patterns) {
+  scf::populateSCFStructuralTypeConversions(typeConverter, patterns);
+}
+
+void transform::ApplySCFStructuralConversionPatternsOp::
+    populateConversionTargetRules(const TypeConverter &typeConverter,
+                                  ConversionTarget &conversionTarget) {
+  scf::populateSCFStructuralTypeConversionTarget(typeConverter,
+                                                 conversionTarget);
+}
+
 //===----------------------------------------------------------------------===//
 // GetParentForOp
 //===----------------------------------------------------------------------===//
 DiagnosedSilenceableFailure
 transform::GetParentForOp::apply(transform::TransformRewriter &rewriter,
                                  transform::TransformResults &results,
                                  transform::TransformState &state) {
   SetVector<Operation *> parents;

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

for (Value operand : adaptor.getOperands())
       unpackUnrealizedConversionCast(operand, unpackedYield);
 
     rewriter.updateRootInPlace(op, [&]() { op->setOperands(unpackedYield); });
     return success();
   }
 };
 } // namespace
 
-void mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
-    TypeConverter &typeConverter, RewritePatternSet &patterns,
-    ConversionTarget &target) {
+void mlir::scf::populateSCFStructuralTypeConversions(
+    TypeConverter &typeConverter, RewritePatternSet &patterns) {
   patterns.add<ConvertForOpTypes, ConvertIfOpTypes, ConvertYieldOpTypes,
                ConvertWhileOpTypes, ConvertConditionOpTypes>(
       typeConverter, patterns.getContext());
+}
+
+void mlir::scf::populateSCFStructuralTypeConversionTarget(
+    const TypeConverter &typeConverter, ConversionTarget &target) {
   target.addDynamicallyLegalOp<ForOp, IfOp>([&](Operation *op) {
     return typeConverter.isLegal(op->getResultTypes());
   });
   target.addDynamicallyLegalOp<scf::YieldOp>([&](scf::YieldOp op) {
     // We only have conversions for a subset of ops that use scf.yield
     // terminators.
     if (!isa<ForOp, IfOp, WhileOp>(op->getParentOp()))
       return true;
     return typeConverter.isLegal(op.getOperandTypes());
   });
   target.addDynamicallyLegalOp<WhileOp, ConditionOp>(
       [&](Operation *op) { return typeConverter.isLegal(op); });
 }
+
+void mlir::scf::populateSCFStructuralTypeConversionsAndLegality(
+    TypeConverter &typeConverter, RewritePatternSet &patterns,
+    ConversionTarget &target) {
+  populateSCFStructuralTypeConversions(typeConverter, patterns);
+  populateSCFStructuralTypeConversionTarget(typeConverter, target);
+}

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

for (Operation &op : getPatterns().front()) {
                       << "pattern descriptor does not specify type "
                          "converter and apply_conversion_patterns op has "
                          "no default type converter";
           diag.attachNote(op.getLoc()) << "pattern descriptor op";
           return diag;
         }
         converter = defaultTypeConverter.get();
       }
+
+      // Add descriptor-specific updates to the conversion target, which may
+      // depend on the final type converter. In structural converters, the
+      // legality of types dictates the dynamic legality of an operation.
+      descriptor.populateConversionTargetRules(*converter, conversionTarget);
+
       descriptor.populatePatterns(*converter, patterns);
     }
   }
 
   FrozenRewritePatternSet frozenPatterns(std::move(patterns));
   for (Operation *target : state.getPayloadOps(getTarget())) {
     // Make sure that this transform is not applied to itself. Modifying the
     // transform IR while it is being interpreted is generally dangerous.

mlir/test/Dialect/SCF/transform-ops.mlir

func.func @test_promote_if_one_iteration(%a: index) -> index {
   return %0 : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   transform.loop.promote_if_one_iteration %0 : !transform.any_op
 }
+
+
+// -----
+
+// CHECK-LABEL: func @test_structural_conversion_patterns(
+// CHECK: scf.for {{.*}} -> (memref<f32>) {
+
+func.func @test_structural_conversion_patterns(%a: tensor<f32>) -> tensor<f32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0 = scf.for %j = %c0 to %c10 step %c1 iter_args(%arg0 = %a) -> tensor<f32> {
+    %1 = "test.foo"(%arg0) : (tensor<f32>) -> (tensor<f32>)
+    scf.yield %1 : tensor<f32>
+  }
+  return %0 : tensor<f32>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.apply_conversion_patterns to %0 {
+    transform.apply_conversion_patterns.scf.structural_conversions
+  } with type_converter {
+    transform.apply_conversion_patterns.transform.test_type_converter
+  } {  partial_conversion  } : !transform.any_op
+}

mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp

patterns.insert<ReplaceWithNewOpConversion>(typeConverter,
                                               patterns.getContext());
 }
 
 namespace {
 /// Test type converter that converts tensor types to memref types.
 class TestTypeConverter : public TypeConverter {
 public:
   TestTypeConverter() {
+    addConversion([](Type t) { return t; });
     addConversion([](RankedTensorType type) -> Type {
       return MemRefType::get(type.getShape(), type.getElementType());
     });
-    addSourceMaterialization([&](OpBuilder &builder, Type resultType,
+    auto unrealizedCastConverter = [&](OpBuilder &builder, Type resultType,
                                        ValueRange inputs,
                                        Location loc) -> std::optional<Value> {
       if (inputs.size() != 1)
         return std::nullopt;
       return builder.create<UnrealizedConversionCastOp>(loc, resultType, inputs)
           .getResult(0);
-    });
+    };
+    addSourceMaterialization(unrealizedCastConverter);
+    addTargetMaterialization(unrealizedCastConverter);
   }
 };
 } // namespace
 
 std::unique_ptr<::mlir::TypeConverter>
 mlir::test::TestTypeConverterOp::getTypeConverter() {
   return std::make_unique<TestTypeConverter>();
 }