[MLIR][Transform] Consolidate the transform ops of get_parent_for and loop unroll from affine and scf dialects.

[kaitingwang]

Summary

This patch consolidates the two transform ops from the affine dialect and the scf dialect to avoid code duplication.

This is to address the review comments from https://reviews.llvm.org/D137997.

The transform ops directory / file structure for the affine dialect is kept for the purpose of forth-coming transform ops
for affine, but get_parent_for and unroll are removed.

[nicolasvasilache]

Nice red diff :)

Thanks!

[kaitingwang]

In D138980#3959192, @nicolasvasilache wrote:

Nice red diff :)

Thanks!

@nicolasvasilache Thank you accepting the patch. Your comments were very helpful! Looking forward to working with you on the transform dialect going forward!

mlir/include/mlir/Dialect/Affine/TransformOps/AffineTransformOps.td

 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def Transform_AffineForOp : Transform_ConcreteOpType<"affine.for">;
 
-def AffineGetParentForOp : Op<Transform_Dialect, "affine.get_parent_for", [
-  NavigationTransformOpTrait, MemoryEffectsOpInterface,
-  DeclareOpInterfaceMethods<TransformOpInterface>]> {
-  let summary =
-      "Gets a handle to the parent 'affine.for' loop of the given operation";
-  let description = [{
-    Produces a handle to the n-th (default 1) parent `affine.for` loop for each
-    Payload IR operation associated with the operand. Fails if such a loop
-    cannot be found. The list of operations associated with the handle contains
-    parent operations in the same order as the list associated with the operand,
-    except for operations that are parents to more than one input which are only
-    present once.
-  }];
-
-  let arguments =
-    (ins TransformTypeInterface:$target,
-         DefaultValuedAttr<ConfinedAttr<I64Attr, [IntPositive]>,
-                           "1">:$num_loops);
-  let results = (outs TransformTypeInterface:$parent);
-
-  let assemblyFormat =
-      "$target attr-dict `:` functional-type(operands, results)";
-}
-
-
-def AffineLoopUnrollOp : Op<Transform_Dialect, "affine.unroll", [
-  FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
-  TransformOpInterface, TransformEachOpTrait]> {
-  let summary = "Unrolls the given loop with the given unroll factor";
-  let description = [{
-    Unrolls each loop associated with the given handle to have up to the given
-    number of loop body copies per iteration. If the unroll factor is larger
-    than the loop trip count, the latter is used as the unroll factor instead.
-
-    #### Return modes
-
-    This operation ignores non-affine::For ops and drops them in the return.
-    If all the operations referred to by the `target` PDLOperation unroll
-    properly, the transform succeeds. Otherwise the transform silently fails.
-
-    Does not return handles as the operation may result in the loop being
-    removed after a full unrolling.
-  }];
-
-  let arguments = (ins Transform_AffineForOp:$target,
-                       ConfinedAttr<I64Attr, [IntPositive]>:$factor);
-
-  let assemblyFormat = "$target attr-dict `:` type($target)";
-
-  let extraClassDeclaration = [{
-    ::mlir::DiagnosedSilenceableFailure applyToOne(
-        ::mlir::AffineForOp target,
-        ::llvm::SmallVector<::mlir::Operation *> & results,
-        ::mlir::transform::TransformState & state);
-  }];
-}
-
 #endif // Affine_TRANSFORM_OPS

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

 
 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` loop for each
-    Payload IR operation associated with the operand. Fails if such a loop
-    cannot be found. The list of operations associated with the handle contains
-    parent operations in the same order as the list associated with the operand,
-    except for operations that are parents to more than one input which are only
-    present once.
+    Produces a handle to the n-th (default 1) parent `scf.for` or `affine.for`
+    (when the affine flag is true) loop for each Payload IR operation
+    associated with the operand. Fails if such a loop cannot be found. The list
+    of operations associated with the handle contains parent operations in the
+    same order as the list associated with the operand, except for operations
+    that are parents to more than one input which are only present once.
   }];
 
   let arguments =
     (ins TransformTypeInterface:$target,
          DefaultValuedAttr<ConfinedAttr<I64Attr, [IntPositive]>,
-                           "1">:$num_loops);
-  let results = (outs TransformTypeInterface:$parent);
+                           "1">:$num_loops,
+         DefaultValuedAttr<BoolAttr, "false">:$affine);
+  let results = (outs TransformTypeInterface : $parent);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 }
 
 def LoopOutlineOp : Op<Transform_Dialect, "loop.outline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
   let summary = "Unrolls the given loop with the given unroll factor";
   let description = [{
     Unrolls each loop associated with the given handle to have up to the given
     number of loop body copies per iteration. If the unroll factor is larger
     than the loop trip count, the latter is used as the unroll factor instead.
 
     #### Return modes
 
-    This operation ignores non-scf::For ops and drops them in the return.
-    If all the operations referred to by the `target` PDLOperation unroll
-    properly, the transform succeeds. Otherwise the transform silently fails.
+    This operation ignores non-scf::For, non-affine::For ops and drops them in
+    the return.  If all the operations referred to by the `target` PDLOperation
+    unroll properly, the transform succeeds. Otherwise the transform silently
+    fails.
 
     Does not return handles as the operation may result in the loop being
     removed after a full unrolling.
   }];
 
-  let arguments = (ins Transform_ScfForOp:$target,
+  let arguments = (ins TransformTypeInterface:$target,
                        ConfinedAttr<I64Attr, [IntPositive]>:$factor);
 
   let assemblyFormat = "$target attr-dict `:` type($target)";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
-        ::mlir::scf::ForOp target,
+        ::mlir::Operation *target,
         ::llvm::SmallVector<::mlir::Operation *> &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
 #endif // SCF_TRANSFORM_OPS

mlir/lib/Dialect/Affine/TransformOps/AffineTransformOps.cpp

 /// A simple pattern rewriter that implements no special logic.
 class SimpleRewriter : public PatternRewriter {
 public:
   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
-// AffineGetParentForOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::AffineGetParentForOp::apply(transform::TransformResults &results,
-                                       transform::TransformState &state) {
-  SetVector<Operation *> parents;
-  for (Operation *target : state.getPayloadOps(getTarget())) {
-    AffineForOp loop;
-    Operation *current = target;
-    for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
-      loop = current->getParentOfType<AffineForOp>();
-      if (!loop) {
-        DiagnosedSilenceableFailure diag = emitSilenceableError()
-                                           << "could not find an '"
-                                           << AffineForOp::getOperationName()
-                                           << "' parent";
-        diag.attachNote(target->getLoc()) << "target op";
-        results.set(getResult().cast<OpResult>(), {});
-        return diag;
-      }
-      current = loop;
-    }
-    parents.insert(loop);
-  }
-  results.set(getResult().cast<OpResult>(), parents.getArrayRef());
-  return DiagnosedSilenceableFailure::success();
-}
-
-//===----------------------------------------------------------------------===//
-// LoopUnrollOp
-//===----------------------------------------------------------------------===//
-
-DiagnosedSilenceableFailure
-transform::AffineLoopUnrollOp::applyToOne(AffineForOp target,
-                                          SmallVector<Operation *> &results,
-                                          transform::TransformState &state) {
-  if (failed(loopUnrollByFactor(target, getFactor()))) {
-    Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);
-    diag << "op failed to unroll";
-    return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
-  }
-  return DiagnosedSilenceableFailure(success());
-}
-
-//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 class AffineTransformDialectExtension
     : public transform::TransformDialectExtension<
           AffineTransformDialectExtension> {
 public:

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

 //===- SCFTransformOps.cpp - Implementation of SCF transformation ops -----===//
 //
 // 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/Dialect/SCF/TransformOps/SCFTransformOps.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Transforms/Patterns.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 
 using namespace mlir;
 
 namespace {
 /// A simple pattern rewriter that implements no special logic.
 class SimpleRewriter : public PatternRewriter {
 public:
   SimpleRewriter(MLIRContext *context) : PatternRewriter(context) {}
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // GetParentForOp
 //===----------------------------------------------------------------------===//
-
 DiagnosedSilenceableFailure
 transform::GetParentForOp::apply(transform::TransformResults &results,
                                  transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {
-    scf::ForOp loop;
-    Operation *current = target;
+    Operation *loop, *current = target;
     for (unsigned i = 0, e = getNumLoops(); i < e; ++i) {
-      loop = current->getParentOfType<scf::ForOp>();
+      loop = getAffine() ? current->getParentOfType<AffineForOp>()
+                         : current->getParentOfType<scf::ForOp>();
+
       if (!loop) {
-        DiagnosedSilenceableFailure diag = emitSilenceableError()
-                                           << "could not find an '"
-                                           << scf::ForOp::getOperationName()
-                                           << "' parent";
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << (getAffine() ? AffineForOp::getOperationName()
+                            : scf::ForOp::getOperationName())
         diag.attachNote(target->getLoc()) << "target op";
         results.set(getResult().cast<OpResult>(), {});
         return diag;
       }
       current = loop;
     }
     parents.insert(loop);
   }
transform::LoopPipelineOp::applyToOne(scf::ForOp target,
   return emitDefaultSilenceableFailure(target);
 }
 
 //===----------------------------------------------------------------------===//
 // LoopUnrollOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
-transform::LoopUnrollOp::applyToOne(scf::ForOp target,
+transform::LoopUnrollOp::applyToOne(Operation *op,
                                     SmallVector<Operation *> &results,
                                     transform::TransformState &state) {
-  if (failed(loopUnrollByFactor(target, getFactor()))) {
-    Diagnostic diag(target->getLoc(), DiagnosticSeverity::Note);
-    diag << "op failed to unroll";
+  LogicalResult result(failure());
+  if (scf::ForOp scfFor = dyn_cast<scf::ForOp>(op))
+    result = loopUnrollByFactor(scfFor, getFactor());
+  else if (AffineForOp affineFor = dyn_cast<AffineForOp>(op))
+    result = loopUnrollByFactor(affineFor, getFactor());
+
+  if (failed(result)) {
+    Diagnostic diag(op->getLoc(), DiagnosticSeverity::Note);
+    diag << "Op failed to unroll";
     return DiagnosedSilenceableFailure::silenceableFailure(std::move(diag));
   }
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//

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

-// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
-
-// CHECK-LABEL: @get_parent_for_op
-func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
-  // expected-remark @below {{first loop}}
-  affine.for %i = %arg0 to %arg1 {
-    // expected-remark @below {{second loop}}
-    affine.for %j = %arg0 to %arg1 {
-      // expected-remark @below {{third loop}}
-      affine.for %k = %arg0 to %arg1 {
-        arith.addi %i, %j : index
-      }
-    }
-  }
-  return
-}
-
-transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-  // CHECK: = transform.affine.get_parent_for
-  %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
-  %2 = transform.affine.get_parent_for %0 { num_loops = 2 } : (!pdl.operation) -> !transform.op<"affine.for">
-  %3 = transform.affine.get_parent_for %0 { num_loops = 3 } : (!pdl.operation) -> !transform.op<"affine.for">
-  transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"affine.for">
-  transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"affine.for">
-  transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"affine.for">
-}
-
-// -----
-
-func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
-  // expected-note @below {{target op}}
-  arith.addi %arg0, %arg1 : index
-  return
-}
-
-transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-  // expected-error @below {{could not find an 'affine.for' parent}}
-  %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
-}
-
-// -----
-
-func.func @loop_unroll_op() {
-  %c0 = arith.constant 0 : index
-  %c42 = arith.constant 42 : index
-  %c5 = arith.constant 5 : index
-  // CHECK: affine.for %[[I:.+]] =
-  // expected-remark @below {{affine for loop}}
-  affine.for %i = %c0 to %c42 {
-    // CHECK-COUNT-4: arith.addi
-    arith.addi %i, %i : index
-  }
-  return
-}
-
-transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
-  %1 = transform.affine.get_parent_for %0 : (!pdl.operation) -> !transform.op<"affine.for">
-  transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
-  transform.affine.unroll %1 { factor = 4 } : !transform.op<"affine.for">
-}
-

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

 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["arith.addi"]} in %arg1
   %1 = transform.loop.get_parent_for %0 : (!pdl.operation) -> !transform.op<"scf.for">
   transform.loop.unroll %1 { factor = 4 } : !transform.op<"scf.for">
 }
 
+// -----
+
+// CHECK-LABEL: @get_parent_for_op
+func.func @get_parent_for_op(%arg0: index, %arg1: index, %arg2: index) {
+  // expected-remark @below {{first loop}}
+  affine.for %i = %arg0 to %arg1 {
+    // expected-remark @below {{second loop}}
+    affine.for %j = %arg0 to %arg1 {
+      // expected-remark @below {{third loop}}
+      affine.for %k = %arg0 to %arg1 {
+        arith.addi %i, %j : index
+      }
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+  // CHECK: = transform.loop.get_parent_for
+  %1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  %2 = transform.loop.get_parent_for %0 { num_loops = 2, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  %3 = transform.loop.get_parent_for %0 { num_loops = 3, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  transform.test_print_remark_at_operand %1, "third loop" : !transform.op<"affine.for">
+  transform.test_print_remark_at_operand %2, "second loop" : !transform.op<"affine.for">
+  transform.test_print_remark_at_operand %3, "first loop" : !transform.op<"affine.for">
+}
+
+// -----
+
+func.func @get_parent_for_op_no_loop(%arg0: index, %arg1: index) {
+  // expected-note @below {{target op}}
+  arith.addi %arg0, %arg1 : index
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+  // expected-error @below {{could not find an 'affine.for' parent}}
+  %1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+}
+
+// -----
+
+func.func @loop_unroll_op() {
+  %c0 = arith.constant 0 : index
+  %c42 = arith.constant 42 : index
+  %c5 = arith.constant 5 : index
+  // CHECK: affine.for %[[I:.+]] =
+  // expected-remark @below {{affine for loop}}
+  affine.for %i = %c0 to %c42 {
+    // CHECK-COUNT-4: arith.addi
+    arith.addi %i, %i : index
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+  %1 = transform.loop.get_parent_for %0 { affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
+  transform.loop.unroll %1 { factor = 4, affine = true } : !transform.op<"affine.for">
+}
+
+// -----
+
+func.func @test_mixed_loops() {
+  %c0 = arith.constant 0 : index
+  %c42 = arith.constant 42 : index
+  %c5 = arith.constant 5 : index
+  scf.for %j = %c0 to %c42 step %c5 {
+    // CHECK: affine.for %[[I:.+]] =
+    // expected-remark @below {{affine for loop}}
+    affine.for %i = %c0 to %c42 {
+      // CHECK-COUNT-4: arith.addi
+      arith.addi %i, %i : index
+    }
+  }
+  return
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1
+  %1 = transform.loop.get_parent_for %0 { num_loops = 1, affine = true } : (!pdl.operation) -> !transform.op<"affine.for">
+  transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
+  transform.loop.unroll %1 { factor = 4 } : !transform.op<"affine.for">
+}