[mlir] Add apply_patterns op to transform dialect

[guraypp]

Summary

This revision introduces a transform.structured.apply_patterns operation that operates on an isolated from above op and applies a set of patterns while updating transform dialect handles.

It provides currently three patterns

• canonicalization: adds all the canonicalization patterns of all registered dialects and ops.
• rank_reducing: adds patterns that results in rank-reducing behavior on subset-based operations.
• vector_to_gpu: adds patterns that converts vector dialect to gpu dialect.

Depends on D134190

[nicolasvasilache]

unfortunately we cannot have this in upstream MLIR (@ftynse @mehdi_amini @Mogball in case there is a possible way forward for CSE / canonicalization with custom Rewriter/Driver that use Listeners)

[guraypp]

It looks like supporting this Op is not possible because of the reasons that @nicolasvasilache mentioned. So I abandon it.

[ftynse]

In D134525#3811134, @nicolasvasilache wrote:

unfortunately we cannot have this in upstream MLIR (@ftynse @mehdi_amini @Mogball in case there is a possible way forward for CSE / canonicalization with custom Rewriter/Driver that use Listeners)

Not sure about CSE, but for canonicalization what you need is the information of what the operation is being replaced with. We already have a hook that indicates that an operation is being replaced, adding an extra operand to that hook to indicate the replacement could be enough for your use case...

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

def MapNestedForeachThreadToGpuBlocks : Op<Transform_Dialect,
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::Operation *target, 
         ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def ApplyPatternsOp : Op<Transform_Dialect, "structured.apply_patterns",
+    [FunctionalStyleTransformOpTrait, 
+     MemoryEffectsOpInterface,
+     TransformEachOpTrait, 
+     TransformOpInterface]> {
+  let description = [{
+    Greedily applies patterns as specified by its attributes.
+    
+    Must be applied to an op with trait IsolatedFromAbove since the 
+    GreedyPatternRewriter asserts those.
+
+    Returns the IsolatedFromAbove op whose content it has modified for better
+    chaining APIs.
+
+    The following additive attributes can be set, they add patterns in an 
+    unspecified order:
+      - canonicalization: adds all the canonicalization patterns of all 
+      registered dialects and ops.
+      - rank_reducing: adds patterns that results in rank-reducing behavior on
+      subset-based operations.
+      - vector_to_gpu: adds patterns that converts vector dialect to gpu dialect.
+
+    Return modes:
+    =============
+    This operation applies a number of patterns to rewrite vector IR into 
+    distributed warp form. To apply these patterns, this operation must target 
+    an operation that is isolated from above, otherwise the transform definitely
+    fails.
+  
+    If the pattern application fails, or if the underlying listener fails to 
+    capture op handles, the transformation definitely fails.
+
+    Otherwise the transformation is successful and no result is returned.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                       UnitAttr:$canonicalization,
+                       UnitAttr:$rank_reducing,
+                       UnitAttr:$vector_to_gpu);
+  let results = (outs PDL_Operation:$result);
+
+  let assemblyFormat = "$target attr-dict";
+  
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::Operation *target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
+  }];  
+}
+
 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

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

 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/raw_ostream.h"
+#include "mlir/Conversion/VectorToGPU/VectorToGPU.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
 /// Extracts a vector of unsigned from an array attribute. Asserts if the
 /// attribute contains values other than intergers. May truncate.
 static SmallVector<unsigned> extractUIntArray(ArrayAttr attr) {
if (failed(alterGpuLaunch(rewriter, gpuLaunch, gridDim[0], gridDim[1],
                             gridDim[2])))
     return DiagnosedSilenceableFailure::definiteFailure();
 
   results.assign({gpuLaunch});
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
+// ApplyPatternsOp
+//===----------------------------------------------------------------------===//
+static void addRankReducingPatterns(RewritePatternSet &patterns) {
+  vector::populateCastAwayVectorLeadingOneDimPatterns(patterns);
+  linalg::populateFoldUnitExtentDimsPatterns(patterns);
+}
+
+static void addAllRegisteredCanonicalizationPatterns(
+    RewritePatternSet &patterns) {
+  MLIRContext *ctx = patterns.getContext();
+  for (Dialect *dialect : ctx->getLoadedDialects())
+    dialect->getCanonicalizationPatterns(patterns);
+  for (RegisteredOperationName op : ctx->getRegisteredOperations())
+    op.getCanonicalizationPatterns(patterns, ctx);
+  FrozenRewritePatternSet();
+}
+
+static void addVectorToGPU(RewritePatternSet &patterns) {
+  populatePrepareVectorToMMAPatterns(patterns, false);
+}
+
+static void postVectorToGPU(Operation* target) {
+  (void)convertVectorToMMAOps(target);
+}
+
+DiagnosedSilenceableFailure
+transform::ApplyPatternsOp::applyToOne(
+    Operation *target, SmallVectorImpl<Operation *> &results,
+    transform::TransformState &state) {
+  
+  if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {
+    target->emitOpError(
+        "applies only to isolated-from-above targets because it needs to apply "
+        "patterns greedily");
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+  }
+
+  MLIRContext *ctx = target->getContext();
+  RewritePatternSet patterns(ctx);  
+  if (getVectorToGpu()) addVectorToGPU(patterns);
+  if (getCanonicalization()) addAllRegisteredCanonicalizationPatterns(patterns);
+  if (getRankReducing()) addRankReducingPatterns(patterns);
+
+  GreedyRewriteConfig config;
+  LogicalResult result = applyPatternsAndFoldGreedily(
+      target, std::move(patterns), config);
+  if (failed(result))
+    return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
+
+  if (getVectorToGpu()) postVectorToGPU(target);
+
+  results.assign({target});
+  return DiagnosedSilenceableFailure(success());
+}
+//===----------------------------------------------------------------------===//
 // TileToForeachThreadOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::tileToForeachThreadOpImpl(
     RewriterBase &rewriter, transform::TransformState &state,
     TransformOpInterface transformOp, ArrayRef<Operation *> targets,
     ArrayRef<OpFoldResult> mixedNumThreads,
     ArrayRef<OpFoldResult> mixedTileSizes, Optional<ArrayAttr> threadDimMapping,

mlir/test/Dialect/Linalg/transform-gpu.mlir

-// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file -canonicalize -cse %s | FileCheck %s
 
 !type = memref<2 x 32 x f32>
 !type1d = memref<32 x f32>
 
 // CHECK-LABEL: func.func @saxpy2dblock(
 // CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGT:[0-9a-z]+]]: memref<32xf32>

 // CHECK-SAME:    %[[ARGX:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGY:[0-9a-z]+]]: memref<2x32xf32>
 // CHECK-SAME:    %[[ARGT:[0-9a-z]+]]: memref<32xf32>
 func.func @saxpy2d(%x: !type, %y: !type, %t: !type1d, %alpha : f32, %stream : !gpu.async.token) -> !type {
   %one = arith.constant 1 : index
   %c12 = arith.constant 12 : index
   %c9 = arith.constant 9 : index
   %c7 = arith.constant 7 : index
+//      CHECK:   %[[C1:.*]] = arith.constant 1 : index
+//      CHECK:   %[[C9:.*]] = arith.constant 9 : index
+//      CHECK:   %[[C7:.*]] = arith.constant 7 : index
 //      CHECK:   gpu.launch
 //      CHECK:   %[[TIDX:.*]] = gpu.thread_id  x
 //      CHECK:   %[[TIDY:.*]] = gpu.thread_id  y
-//      CHECK:   %[[C9:.*]] = arith.constant 9 : index
 //      CHECK:   arith.cmpi ult, %[[TIDX]], %[[C9]] : index
-//      CHECK:   %[[C7:.*]] = arith.constant 7 : index
 //      CHECK:   arith.cmpi ult, %[[TIDY]], %[[C7]] : index
 //      CHECK:   memref.load %[[ARGX]][%[[TIDY]], %[[TIDX]]]
 //      CHECK:   memref.load %[[ARGY]][%[[TIDY]], %[[TIDX]]]
 //      CHECK:   gpu.barrier
 //      CHECK:   %[[TIDX2:.*]] = gpu.thread_id  x
 //      CHECK:   %[[TIDY2:.*]] = gpu.thread_id  y
-//      CHECK:   %[[C1:.*]] = arith.constant 1 : index
 //      CHECK:   arith.cmpi ult, %[[TIDY2]], %[[C1]] : index
 //      CHECK:   memref.load %[[ARGT]][%[[TIDX2]]]
 //      CHECK:   gpu.barrier
   %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)
             threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) 
   {
     scf.foreach_thread (%i, %j) in (%c7, %c9) {
         %4 = memref.load %x[%i, %j] : !type        

 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %funcop = transform.structured.match ops{["func.func"]} in %arg0
     %gpuLaunch = transform.structured.map_nested_foreach_thread_to_gpu_blocks %funcop { generate_gpu_launch }
     transform.structured.map_nested_foreach_thread_to_gpu_threads %gpuLaunch { blockDim = [32, 4, 1] }
   }
 }
+
+// -----
+// CHECK-LABEL: func.func @tiled_matmul(
+// CHECK-SAME:    %[[INX:[0-9a-z]+]]: memref<8192x4096xf16>
+// CHECK-SAME:    %[[INY:[0-9a-z]+]]: memref<4096x8192xf16>
+// CHECK-SAME:    %[[OUT:[0-9a-z]+]]: memref<8192x8192xf16>
+func.func @tiled_matmul(%arg0: memref<8192x4096xf16>, %arg1: memref<4096x8192xf16>, %arg2: memref<8192x8192xf16>) -> memref<8192x8192xf16> {
+  %c0 = arith.constant 0 : index
+  %c8192 = arith.constant 8192 : index
+  %c256 = arith.constant 256 : index
+  %c1 = arith.constant 1 : index
+//      CHECK: gpu.launch
+//      CHECK: %[[BLKX:.*]] = gpu.block_id  x
+//      CHECK: %[[BLKY:.*]] = gpu.block_id  y
+//      CHECK: %[[BIDINX:[0-9a-z]+]] = memref.subview %[[INX]][%[[BLKX]], 0]
+//      CHECK: %[[BIDINY:[0-9a-z]+]] = memref.subview %[[INY]][0, %[[BLKY]]]
+//      CHECK: %[[BIDOUT:[0-9a-z]+]] = memref.subview %[[OUT]][%[[BLKX]], %[[BLKY]]]
+//      CHECK: %[[TIDX:.*]] = gpu.thread_id  x
+//      CHECK: %[[TIDY:.*]] = gpu.thread_id  y
+//      CHECK: %[[TIDINX:[0-9a-z]+]] = memref.subview %[[BIDINX]][%[[TIDX]], 0]
+//      CHECK: %[[TIDINY:[0-9a-z]+]] = memref.subview %[[BIDINY]][0, %[[TIDY]]]
+//      CHECK: %[[TIDOUT:[0-9a-z]+]] = memref.subview %[[BIDOUT]][%[[TIDX]], %[[TIDY]]]
+//      CHECK: scf.for %[[IV:[0-9a-z]+]]
+//      CHECK: %[[TILEDINX:[0-9a-z]+]] = memref.subview %[[TIDINX]][0, %[[IV]]]
+//      CHECK: %[[TILEDINY:[0-9a-z]+]] = memref.subview %[[TIDINY]][%[[IV]], 0]
+//      CHECK: %[[SUBINX:[0-9a-z]+]] = gpu.subgroup_mma_load_matrix %[[TILEDINX]]
+//      CHECK: %[[SUBINY:[0-9a-z]+]] = gpu.subgroup_mma_load_matrix %[[TILEDINY]]
+//      CHECK: %[[SUBOUT:[0-9a-z]+]] = gpu.subgroup_mma_load_matrix %[[TIDOUT]]
+//      CHECK: %[[RES:[0-9a-z]+]] = gpu.subgroup_mma_compute %[[SUBINX]], %[[SUBINY]], %[[SUBOUT]] : !gpu.mma_matrix<4x4xf16, "AOp">, !gpu.mma_matrix<4x4xf16, "BOp"> -> !gpu.mma_matrix<4x4xf16, "COp">
+//      CHECK: gpu.subgroup_mma_store_matrix %[[RES]], %[[TIDOUT]]
+  scf.foreach_thread (%i, %j) in (%c256, %c256) {  
+    %0 = memref.subview %arg0[%i, 0] [256, 4096] [1, 1] : memref<8192x4096xf16> to memref<256x4096xf16, strided<[4096, 1], offset: ?>>
+    %1 = memref.subview %arg1[0, %j] [4096, 256] [1, 1] : memref<4096x8192xf16> to memref<4096x256xf16, strided<[8192, 1], offset: ?>>
+    %2 = memref.subview %arg2[%i, %j] [256, 256] [1, 1] : memref<8192x8192xf16> to memref<256x256xf16, strided<[8192, 1], offset: ?>>
+    %c4 = arith.constant 4 : index
+    scf.foreach_thread (%ti, %tj) in (%c4, %c4) {  
+      %3 = memref.subview %0[%ti, 0] [4, 4096] [1, 1] : memref<256x4096xf16, strided<[4096, 1], offset: ?>> to memref<4x4096xf16, strided<[4096, 1], offset: ?>>
+      %4 = memref.subview %1[0, %tj] [4096, 4] [1, 1] : memref<4096x256xf16, strided<[8192, 1], offset: ?>> to memref<4096x4xf16, strided<[8192, 1], offset: ?>>
+      %5 = memref.subview %2[%ti, %tj] [4, 4] [1, 1] : memref<256x256xf16, strided<[8192, 1], offset: ?>> to memref<4x4xf16, strided<[8192, 1], offset: ?>>
+      %c4096 = arith.constant 4096 : index
+      scf.for %arg7 = %c0 to %c4096 step %c4 {
+        %6 = memref.subview %3[0, %arg7] [4, 4] [1, 1] : memref<4x4096xf16, strided<[4096, 1], offset: ?>> to memref<4x4xf16, strided<[4096, 1], offset: ?>>
+        %7 = memref.subview %4[%arg7, 0] [4, 4] [1, 1] : memref<4096x4xf16, strided<[8192, 1], offset: ?>> to memref<4x4xf16, strided<[8192, 1], offset: ?>>
+        linalg.matmul ins(%6, %7 : memref<4x4xf16, strided<[4096, 1], offset: ?>>, memref<4x4xf16, strided<[8192, 1], offset: ?>>) outs(%5 : memref<4x4xf16, strided<[8192, 1], offset: ?>>)
+      }
+    }
+  }
+  return %arg2 : memref<8192x8192xf16>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %funcop = transform.structured.match ops{["func.func"]} in %arg1
+    
+    // This op transforms as follows
+    // scf.foreach_thread(i, j)     --> gpu.launch (blockIdx.x, blockIdx.y) 
+    //  scf.foreach_thread(ti, tj)
+    //   body
+    %gpuLaunch = transform.structured.map_nested_foreach_thread_to_gpu_blocks %funcop { generate_gpu_launch }
+    
+    // This op transforms as follows
+    // gpu.launch (...)
+    //  i = blockIdx.x
+    //  j = blockIdx.y
+    //  scf.foreach_thread(ti, tj)  --> (threadIdx.x, threadIdx.y) 
+    //   body
+    %gpuLaunch2 = transform.structured.map_nested_foreach_thread_to_gpu_threads %gpuLaunch { blockDim = [4, 4, 1] }
+
+    // linalg.matmul -> vector dialect    
+    %isolatedFunc = get_closest_isolated_parent %gpuLaunch2
+    %vectorizedFunc = transform.structured.vectorize %isolatedFunc
+    
+    // vector dialect -> GPU warp synchronous matrix multiply accumulate
+    transform.structured.apply_patterns %vectorizedFunc { vector_to_gpu }
+  }
+}