[mlir] add VectorizeOp to structured transform ops

[ftynse]

Summary

Vectorization is a key transformation to achieve high performance on most
architectures. In the transform dialect, vectorization is implemented as a
parameterizable transform op. It currently applies to a scope of payload IR
delimited by some isolated-from-above op, mainly because several enabling
transformations (such as affine simplification) are needed to perform
vectorization and these transformation would apply to ops other than the "main"
computational payload op. A separate "navigation" transform op that obtains the
isolated-from-above ancestor of an op is introduced in the core transform
dialect. Even though it is currently only useful for vectorization,
isolated-from-above ops are a common anchor for transformations (usually
implemented as passes) that is likely to be reused in the future.

Depends On D126374

[nicolasvasilache]

typo in commit message: transfomrations.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

137

vectoriztaion

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

19

Rename to get_closest_isolated_parent?
I imagine in the future we may see a get_common_isolated_parent

30

typo: assocaited

31

same

[ftynse]

Review

[ftynse]

clang-format

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let arguments = (ins PDL_Operation:$target,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$sizes,
                    DefaultValuedAttr<I64ArrayAttr, "{}">:$interchange);
   let results = (outs PDL_Operation:$tiled_linalg_op,
                       Variadic<PDL_Operation>:$loops);
 
   let hasCustomAssemblyFormat = 1;
 }
 
+def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
+    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
+     TransformEachOpTrait, TransformOpInterface]> {
+  let description = [{
+    Indicates that the given `target` op all the ops it contains should be
+    vectorized with the configuration specified by the attributes of this op.
+    This vectorization only handles structured ops that operate on shaped types
+    and does not vectorize loops or straight-line. Internally, it applies a
+    set of rewrite patterns, some of which enable vectorization and some of
+    which clean up the results. Therefore, it can only be applied to an op with
+    the "isolated from above property". If finer granularity is required, it can
+    be achieved by outlining the target part of the payload IR into, e.g., a
+    function, performing the transformation, and inlining it back. This
+    transformation only fails if the entire pattern rewriting failed, i.e., it
+    does **not** fail when no ops were vectorized.
+
+    Note that this transformation is invalidating the handles to any payload IR
+    operation that is contained inside the vectoriztaion target.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                   DefaultValuedAttr<BoolAttr, "false">:$vectorize_padding);
+  let results = (outs PDL_Operation:$transformed);
+
+  let assemblyFormat = "$target attr-dict";
+
+  let extraClassDeclaration = [{
+    ::mlir::FailureOr<Operation *> applyToOne(::mlir::Operation *target);
+  }];
+}
+
 #endif // LINALG_TRANSFORM_OPS

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

 
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
 
+def GetClosestIsolatedParentOp : TransformDialectOp<"get_closest_isolated_parent",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+  let summary = "Gets handles to the closest isolated-from-above parents";
+  let description = [{
+    The handles defined by this Transform op correspond to the closest isolated
+    from above ancestor of the Payload IR operations associated with its
+    operand. If any of the given Payload IR ops has no such parent (unlikely as
+    there usually is a top-level ModuleOp), the transformation is considered to
+    have failed.
+
+    Ancestor ops follow the same order as the ops assocaited with the
+    operand, except for potential duplicates (multiple Payload IR ops assocaited
+    with the operand have the same parent) for which the ancestor will only be
+    listed once for the first time it occurs. For example, given the list
+    "(childof(A), childof(B), childof(B), childof(A), childof(B))", the
+    resulting list will be just "(A, B)". Note that no other semantic ordering
+    is applied, e.g., "B" may itself be a parent of "A". This may have an impact
+    on the further transformation applied to the handle produced here.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs PDL_Operation:$parent);
+  let assemblyFormat = "$target attr-dict";
+}
+
 def PDLMatchOp : TransformDialectOp<"pdl_match",
     [DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Finds ops that match the named PDL pattern";
   let description = [{
     Find Payload IR ops nested within the Payload IR op associated with the
     operand that match the PDL pattern identified by its name. The pattern is
     expected to be defined in the closest surrounding `WithPDLPatternsOp`.
 

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

add_mlir_dialect_library(MLIRLinalgTransformOps
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRLinalg
   MLIRLinalgTransforms
   MLIRParser
   MLIRPDL
   MLIRSideEffectInterfaces
   MLIRTransformDialect
+  MLIRVector
   )

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Parser/Parser.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
 /// Extracts a vector of int64_t from an array attribute. Asserts if the
 /// attribute contains values other than integers.
effects.emplace_back(MemoryEffects::Allocate::get(), r,
                          TransformMappingResource::get());
   }
 
   effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), PayloadIRResource::get());
 }
 
 //===----------------------------------------------------------------------===//
+// VectorizeOp
+//===----------------------------------------------------------------------===//
+
+FailureOr<Operation *> VectorizeOp::applyToOne(Operation *target) {
+  if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+    InFlightDiagnostic diag = emitOpError()
+                              << "applies only to isolated-from-above targets";
+    diag.attachNote(target->getLoc()) << "non-isolated target";
+    return diag;
+  }
+
+  MLIRContext *ctx = getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<LinalgVectorizationPattern>(ctx);
+
+  vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
+  vector::populateVectorReductionToContractPatterns(patterns);
+  patterns.add<linalg::LinalgCopyVTRForwardingPattern,
+               linalg::LinalgCopyVTWForwardingPattern>(ctx,
+                                                       /*benefit=*/2);
+  vector::TransferReadOp::getCanonicalizationPatterns(patterns, ctx);
+  vector::TransferWriteOp::getCanonicalizationPatterns(patterns, ctx);
+  if (getVectorizePadding())
+    linalg::populatePadOpVectorizationPatterns(patterns);
+
+  if (failed(applyPatternsAndFoldGreedily(target, std::move(patterns)))) {
+    InFlightDiagnostic diag = emitError() << "failed to apply";
+    diag.attachNote(target->getLoc()) << "target op";
+    return diag;
+  }
+  return target;
+}
+
+//===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the additional
 /// ops are using PDL types for operands and results.
 class LinalgTransformDialectExtension
     : public transform::TransformDialectExtension<
           LinalgTransformDialectExtension> {
 public:
   LinalgTransformDialectExtension() {
     declareDependentDialect<pdl::PDLDialect>();
     declareDependentDialect<scf::SCFDialect>();
+    declareDependentDialect<vector::VectorDialect>();
     registerTransformOps<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
         >();
   }
 };
 } // namespace
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"
 
 void mlir::linalg::registerTransformDialectExtension(
     DialectRegistry &registry) {
   registry.addExtensions<LinalgTransformDialectExtension>();
 }

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

if (succeeded(applicator.matchAndRewrite(op, rewriter)))
       results.push_back(op);
   });
 
   return success();
 }
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// GetClosestIsolatedParentOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult transform::GetClosestIsolatedParentOp::apply(
+    transform::TransformResults &results, transform::TransformState &state) {
+  SetVector<Operation *> parents;
+  for (Operation *target : state.getPayloadOps(getTarget())) {
+    Operation *parent =
+        target->getParentWithTrait<OpTrait::IsIsolatedFromAbove>();
+    if (!parent) {
+      InFlightDiagnostic diag =
+          emitError() << "could not find an isolated-from-above parent op";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    parents.insert(parent);
+  }
+  results.set(getResult().cast<OpResult>(), parents.getArrayRef());
+  return success();
+}
+
+void transform::GetClosestIsolatedParentOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
+                       TransformMappingResource::get());
+  effects.emplace_back(MemoryEffects::Allocate::get(), getParent(),
+                       TransformMappingResource::get());
+  effects.emplace_back(MemoryEffects::Write::get(), getParent(),
+                       TransformMappingResource::get());
+  effects.emplace_back(MemoryEffects::Read::get(), PayloadIRResource::get());
+}
+
+//===----------------------------------------------------------------------===//
 // PDLMatchOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult transform::PDLMatchOp::apply(transform::TransformResults &results,
                                            transform::TransformState &state) {
   auto *extension = state.getExtension<PatternApplicatorExtension>();
   assert(extension &&
          "expected PatternApplicatorExtension to be attached by the parent op");

mlir/test/Dialect/Linalg/transform-op-vectorize.mlir

+// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: @vectorize_matmul
+// CHECK-SAME: %[[A:.*]]: tensor<24x12xf32>
+// CHECK-SAME: %[[B:.*]]: tensor<12x25xf32>
+// CHECK-SAME: %[[C:.*]]: tensor<24x25xf32>
+func.func @vectorize_matmul(%arg0: tensor<24x12xf32>,
+                            %arg1: tensor<12x25xf32>,
+                            %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
+  // CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
+  // CHECK: %[[vB:.+]] = vector.transfer_read %[[B]]
+  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
+  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]]
+  // CHECK: %[[vS:.+]] = arith.addf %[[vR]], %[[vC]]
+  // CHECK: vector.transfer_write %[[vS]], %[[C]]
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1 = get_isolated_parent %0
+    %2 = transform.structured.vectorize %1
+  }
+}
+
+// -----
+
+#map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
+#map1 = affine_map<()[s0] -> (-s0 + 7)>
+
+// CHECK-LABEL: @vectorize_keep_pad
+// CHECK-SAME: %[[C:[a-zA-Z0-9_]+]]: tensor<24x25xf32>
+func.func @vectorize_keep_pad(
+    %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>,
+    %arg2: tensor<24x25xf32>, %arg3: index, %arg4: index,
+    %arg5: index) -> tensor<24x25xf32> {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = affine.min #map0()[%arg5]
+  %1 = tensor.extract_slice %arg0[%arg3, %arg5] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
+  %2 = tensor.extract_slice %arg1[%arg5, %arg4] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
+  %3 = tensor.extract_slice %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
+  %4 = affine.apply #map1()[%0]
+  // CHECK: %[[pA:.*]] = tensor.pad
+  %5 = tensor.pad %1 nofold low[%c0, %c0] high[%c0, %4] {
+  ^bb0(%arg6: index, %arg7: index):
+    tensor.yield %cst : f32
+  } : tensor<4x?xf32> to tensor<4x7xf32>
+  %6 = affine.apply #map1()[%0]
+  // CHECK: %[[pB:.*]] = tensor.pad
+  %7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] {
+  ^bb0(%arg6: index, %arg7: index):
+    tensor.yield %cst : f32
+  } : tensor<?x5xf32> to tensor<7x5xf32>
+  // CHECK: %[[vA:.+]] = vector.transfer_read %[[pA]]
+  // CHECK: %[[vB:.+]] = vector.transfer_read %[[pB]]
+  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
+  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]]
+  // CHECK: %[[vS:.+]] = arith.addf %[[vR]], %[[vC]]
+  // CHECK: vector.transfer_write %[[vS]], %[[C]]
+  %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
+  %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
+  return %9 : tensor<24x25xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1 = get_isolated_parent %0
+    %2 = transform.structured.vectorize %1
+  }
+}
+
+// -----
+
+#map0 = affine_map<()[s0] -> (-s0 + 12, 7)>
+#map1 = affine_map<()[s0] -> (-s0 + 7)>
+
+// CHECK-LABEL: @vectorize_pad
+// CHECK-SAME: %[[A:.+]]: tensor<24x12xf32>
+// CHECK-SAME: %[[B:.+]]: tensor<12x25xf32>
+// CHECK-SAME: %[[C:.+]]: tensor<24x25xf32>
+func.func @vectorize_pad(
+    %arg0: tensor<24x12xf32>, %arg1: tensor<12x25xf32>,
+    %arg2: tensor<24x25xf32>, %arg3: index, %arg4: index,
+    %arg5: index) -> tensor<24x25xf32> {    
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = affine.min #map0()[%arg5]
+  // CHECK: %[[sA:.+]] = tensor.extract_slice %[[A]]
+  // CHECK: %[[sB:.+]] = tensor.extract_slice %[[B]]
+  %1 = tensor.extract_slice %arg0[%arg3, %arg5] [4, %0] [1, 1] : tensor<24x12xf32> to tensor<4x?xf32>
+  %2 = tensor.extract_slice %arg1[%arg5, %arg4] [%0, 5] [1, 1] : tensor<12x25xf32> to tensor<?x5xf32>
+  %3 = tensor.extract_slice %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<24x25xf32> to tensor<4x5xf32>
+  // CHECK: %[[vA:.+]] = vector.transfer_read %[[sA]]
+  %4 = affine.apply #map1()[%0]
+  %5 = tensor.pad %1 nofold low[%c0, %c0] high[%c0, %4] {
+  ^bb0(%arg6: index, %arg7: index):
+    tensor.yield %cst : f32
+  } : tensor<4x?xf32> to tensor<4x7xf32>
+  %6 = affine.apply #map1()[%0]
+  // CHECK: %[[vB:.+]] = vector.transfer_read %[[sB]]
+  %7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] {
+  ^bb0(%arg6: index, %arg7: index):
+    tensor.yield %cst : f32
+  } : tensor<?x5xf32> to tensor<7x5xf32>
+  // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
+  // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]]
+  // CHECK: %[[vS:.+]] = arith.addf %[[vR]], %[[vC]]
+  // CHECK: vector.transfer_write %[[vS]], %[[C]]
+  %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
+  %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
+  return %9 : tensor<24x25xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    %1 = get_isolated_parent %0
+    %2 = transform.structured.vectorize %1 {vectorize_padding = true}
+  }
+}
+
+// -----
+
+func.func @vectorize(%arg0: tensor<24x12xf32>,
+                     %arg1: tensor<12x25xf32>,
+                     %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
+  // expected-note @below {{non-isolated target}}
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
+  func.return %0 : tensor<24x25xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @pdl_target : benefit(1) {
+    %args = operands
+    %results = types
+    %0 = pdl.operation "linalg.matmul"(%args : !pdl.range<value>) -> (%results : !pdl.range<type>)
+    // TODO: we don't want this, but it is the required terminator for pdl.pattern
+    rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = pdl_match @pdl_target in %arg1
+    // expected-error @below {{applies only to isolated-from-above targets}}
+    %2 = transform.structured.vectorize %0
+  }
+}

mlir/test/Dialect/Transform/test-interpreter.mlir

 }
 
 // expected-remark @below {{matched}}
 "test.some_op"() : () -> ()
 "test.other_op"() : () -> ()
 // expected-remark @below {{matched}}
 "test.some_op"() : () -> ()
 
+// -----
+
+// expected-remark @below {{parent function}}
+func.func @foo() {
+  %0 = arith.constant 0 : i32
+  return
+}
+
+// expected-remark @below {{parent function}}
+func.func @bar() {
+  %0 = arith.constant 0 : i32
+  %1 = arith.constant 1 : i32
+  return
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @const : benefit(1) {
+    %r = pdl.types
+    %0 = pdl.operation "arith.constant" -> (%r : !pdl.range<type>)
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %f = pdl_match @const in %arg1
+    // CHECK: %{{.+}} = get_isolated_parent %{{.+}}
+    %m = get_isolated_parent %f
+    test_print_remark_at_operand %m, "parent function"
+  }
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":IR",
         ":LinalgOps",
         ":LinalgTransformOpsIncGen",
         ":LinalgTransforms",
         ":PDLDialect",
         ":Parser",
         ":SideEffectInterfaces",
         ":TransformDialect",
+        ":TransformUtils",
         "//llvm:Support",
     ],
 )
 
 gentbl_cc_library(
     name = "LinalgPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [