[MLIR][Transform] Transform dialect support for affine

[kaitingwang]

Summary

Create the directory and file structures to support transform ops for
the affine dialect in the future. Enable affine.unroll and affine.get_parent_for
transform operations in this first batch of check in.

[chelini]

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

55

nit: please add a space between # and 'Return modes'.

59

nit: space before 'Otherwise'.

mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt

7

I would rename AffineNewTransformOps to something else, maybe: AffineLoopTransformOps.

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

16

nit: new line before using namespace mlir?

65

nullptr should not be needed. The callback is defaulted to nullptr.

[ftynse]

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

2

Nit: this looks 1 character too long (or the bottom one too short, I didn't count).

25

Please specify that this means the _affine_ for loop.

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

86

Where is this generating the func dialect?

[kaitingwang]

All comments are addressed. Much appreciated! Please let me know if there are further comments.

[kaitingwang]

mlir/include/mlir/Dialect/Affine/TransformOps/CMakeLists.txt

7

Name changed as suggested.

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

65

Removed.

86

Removed.

[kaitingwang]

Clang-format causing Debian build to fail. Addressed in this patch.

[kaitingwang]

@ftynse @chelini Thanks for the reviews. Would you be able to accept this revision? Much appreciated!

[nicolasvasilache]

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

21

I'd strongly prefer to see this implemented on LoopLike interface (with maybe a name attribute to catch a particular op type) and have a single more general transform than duplicating this op in the affine dialect.

46

I'd strongly prefer to see this implemented on LoopLike interface and have a single more general transform than duplicating this op in the affine dialect.

[kaitingwang]

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

21

I can merge these two TransformOps into SCF/SCFTransformOps.cpp and in the respective apply methods in SCFTransformOps.cpp, make a determination (base on a newly introduced parameter), whether to handle in a SCF way or an Affine way. For instance, to invoke SCF's unroll method, or Affine's unroll method. And in the event a matched ops is surrounded by both affine and scf.for's, which parents are to be returned.

Please advise if this is what you have in mind. I have two more affine TransformOps in an upcoming patch: affine.tile and affine.stripmine. Should we consolidate affine's TransformOps into SCF's going forward? Thank you!

[nicolasvasilache]

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

21

I can merge these two TransformOps into SCF/SCFTransformOps.cpp and in the respective apply methods in SCFTransformOps.cpp, make a determination (base on a newly introduced parameter), whether to handle in a SCF way or an Affine way. For instance, to invoke SCF's unroll method, or Affine's unroll method. And in the event a matched ops is surrounded by both affine and scf.for's, which parents are to be returned.
Please advise if this is what you have in mind.

Yes, this would make sense to me and in the future we can continue extending to the Parallel flavor of ops and While loops etc as the need grows. I'll let @ftynse also chime in here.

I have two more affine TransformOps in an upcoming patch: affine.tile and affine.stripmine. Should we consolidate affine's TransformOps into SCF's going forward? Thank you!

These ones seem specific enough to affine that they make sense to be in the affine dialect.

Thanks!

[kaitingwang]

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

21

Thank you for the advice! Hopefully the patch https://reviews.llvm.org/D138980 addresses the concern.

mlir/include/mlir/Dialect/Affine/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(TransformOps)
 
 set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name Affine)
 add_public_tablegen_target(MLIRAffinePassIncGen)
 
 add_mlir_doc(Passes AffinePasses ./ -gen-pass-doc)

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

+//===- AffineTransformOps.h - Affine transformation ops ---------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
+#define MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H
+
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
+#include "mlir/IR/OpImplementation.h"
+
+namespace mlir {
+class AffineForOp;
+namespace func {
+class FuncOp;
+} // namespace func
+namespace affine {
+class ForOp;
+} // namespace affine
+} // namespace mlir
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h.inc"
+
+namespace mlir {
+class DialectRegistry;
+
+namespace affine {
+void registerTransformDialectExtension(DialectRegistry &registry);
+} // namespace affine
+} // namespace mlir
+
+#endif // MLIR_DIALECT_AFFINE_TRANSFORMOPS_AFFINETRANSFORMOPS_H

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

+//===- AffineTransformOps.td - Affine transformation ops ---*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef AFFINE_TRANSFORM_OPS
+#define AFFINE_TRANSFORM_OPS
+
+include "mlir/Dialect/Transform/IR/TransformDialect.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/Affine/TransformOps/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS AffineTransformOps.td)
+mlir_tablegen(AffineTransformOps.h.inc -gen-op-decls)
+mlir_tablegen(AffineTransformOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRAffineTransformOpsIncGen)
+
+add_mlir_doc(AffineTransformOps AffineLoopTransformOps Dialects/ -gen-op-doc)

mlir/include/mlir/InitAllDialects.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INITALLDIALECTS_H_
 #define MLIR_INITALLDIALECTS_H_
 
 #include "mlir/Dialect/AMDGPU/AMDGPUDialect.h"
 #include "mlir/Dialect/AMX/AMXDialect.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/ArmNeon/ArmNeonDialect.h"
 #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h"
 #include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h"
 #include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
inline void registerAllDialects(DialectRegistry &registry) {
   // clang-format on
 
   // Register all dialect extensions.
   bufferization::registerTransformDialectExtension(registry);
   linalg::registerTransformDialectExtension(registry);
   memref::registerTransformDialectExtension(registry);
   scf::registerTransformDialectExtension(registry);
   gpu::registerTransformDialectExtension(registry);
+  affine::registerTransformDialectExtension(registry);
 
   // Register all external models.
   arith::registerBufferizableOpInterfaceExternalModels(registry);
   bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
       registry);
   linalg::registerBufferizableOpInterfaceExternalModels(registry);
   linalg::registerTilingInterfaceExternalModels(registry);
   scf::registerBufferizableOpInterfaceExternalModels(registry);

mlir/lib/Dialect/Affine/CMakeLists.txt

 add_subdirectory(Analysis)
 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(TransformOps)
 add_subdirectory(Utils)

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

+//=== AffineTransformOps.cpp - Implementation of Affine 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/Affine/TransformOps/AffineTransformOps.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/LoopUtils.h"
+#include "mlir/Dialect/Transform/IR/TransformDialect.h"
+#include "mlir/Dialect/Transform/IR/TransformInterfaces.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
+
+//===----------------------------------------------------------------------===//
+// 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:
+  using Base::Base;
+
+  void init() {
+    declareGeneratedDialect<AffineDialect>();
+
+    registerTransformOps<
+#define GET_OP_LIST
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
+        >();
+  }
+};
+} // namespace
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/Affine/TransformOps/AffineTransformOps.cpp.inc"
+
+void mlir::affine::registerTransformDialectExtension(
+    DialectRegistry &registry) {
+  registry.addExtensions<AffineTransformDialectExtension>();
+}

mlir/lib/Dialect/Affine/TransformOps/CMakeLists.txt

+add_mlir_dialect_library(MLIRAffineTransformOps
+  AffineTransformOps.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Affine/TransformOps
+
+  DEPENDS
+  MLIRAffineTransformOpsIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRAffineDialect
+  MLIRFuncDialect
+  MLIRIR
+  MLIRPDLDialect
+  MLIRAffineDialect
+  MLIRAffineTransforms
+  MLIRAffineUtils
+  MLIRTransformDialect
+  MLIRVectorDialect
+)

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">
+}
+