From e30bd804dbf164d07bfb5d67be3e94a1a3ef8db9 Mon Sep 17 00:00:00 2001 From: lorenzo chelini Date: Wed, 17 Jan 2024 05:12:08 -0600 Subject: [PATCH] Refresh transform dialect (#868) --- include/TPP/Passes.td | 12 --- lib/TPP/DefaultTppPasses.cpp | 3 - lib/TPP/Transforms/CMakeLists.txt | 1 - .../TransformDialectInterpreter.cpp | 60 -------------- .../Dialect/Transform/transform-collapse.mlir | 40 ++++++---- .../Transform/transform-convolutions.mlir | 50 +++++++----- .../Transform/transform-drop-schedule.mlir | 20 ----- test/Dialect/Transform/transform-matmuls.mlir | 23 +++--- test/Dialect/Transform/transform-mlp.mlir | 22 +++--- test/Dialect/Transform/transform-pack.mlir | 50 +++++++----- .../Transform/transform-propagation.mlir | 78 ++++++++++++------- .../transform-rewrite-to-brgemm.mlir | 38 +++++---- .../Transform/transform-rewrite-to-gemm.mlir | 15 ++-- .../transform/transform-convolution.mlir | 14 ++-- .../transform/transform-convolutions.mlir | 14 ++-- .../transform/transform-matmul.mlir | 12 +-- .../transform-rewrite-conv-to-gemm.mlir | 13 ++-- 17 files changed, 232 insertions(+), 233 deletions(-) delete mode 100644 lib/TPP/Transforms/TransformDialectInterpreter.cpp delete mode 100644 test/Dialect/Transform/transform-drop-schedule.mlir diff --git a/include/TPP/Passes.td b/include/TPP/Passes.td index a41dbeebe..895512b39 100644 --- a/include/TPP/Passes.td +++ b/include/TPP/Passes.td @@ -62,14 +62,6 @@ def ConvertCheckToLoops : Pass<"convert-check-to-loops", "func::FuncOp"> { let dependentDialects = ["scf::SCFDialect"]; } -def TransformDialectInterpreter : Pass<"transform-dialect-interpreter", "ModuleOp"> { - let summary = "Apply transform dialect operations one by one"; - let description = [{ - Copy and paste from 'TestTransformDialectInterpreter.cpp'. Apply the transform - schedule. - }]; -} - def ConvertPerfToLoops : Pass<"convert-perf-to-loops", "func::FuncOp"> { let summary = "Convert perf to loops"; let description = [{ @@ -89,10 +81,6 @@ def ConvertPerfToFunc : Pass<"convert-perf-to-func", "ModuleOp"> { "tensor::TensorDialect"]; } -def TransformDropSchedule : Pass<"transform-drop-schedule", "ModuleOp"> { - let summary = "Drop the transform schedule"; -} - def PackVNNI : Pass<"pack-vnni", "func::FuncOp"> { let summary = "Convert matmul/brgemm to vnni layout"; let description = [{ diff --git a/lib/TPP/DefaultTppPasses.cpp b/lib/TPP/DefaultTppPasses.cpp index b31e9f03a..b7905cf24 100644 --- a/lib/TPP/DefaultTppPasses.cpp +++ b/lib/TPP/DefaultTppPasses.cpp @@ -272,9 +272,6 @@ struct DefaultTppPasses void constructPipeline() override { pm.clear(); - // Default pipeline does not support transforms yet - pm.addPass(createTransformDropSchedule()); - if (linalgToLoops) { // Lower linalg directly to loops. // Skip all TPP transformations. diff --git a/lib/TPP/Transforms/CMakeLists.txt b/lib/TPP/Transforms/CMakeLists.txt index 65ab9791a..16916c87a 100644 --- a/lib/TPP/Transforms/CMakeLists.txt +++ b/lib/TPP/Transforms/CMakeLists.txt @@ -14,7 +14,6 @@ add_mlir_library(TPPTransforms RewriteToBatchReduceGemm.cpp TileConsumerAndFuseProducers.cpp ToBlockLayoutAndBack.cpp - TransformDialectInterpreter.cpp TransformUtils.cpp CombineXsmmPass.cpp diff --git a/lib/TPP/Transforms/TransformDialectInterpreter.cpp b/lib/TPP/Transforms/TransformDialectInterpreter.cpp deleted file mode 100644 index c39477b39..000000000 --- a/lib/TPP/Transforms/TransformDialectInterpreter.cpp +++ /dev/null @@ -1,60 +0,0 @@ -//===- TransformDialectInterpreter.cpp ------------------------------------===// -// -// 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 -// -//===----------------------------------------------------------------------===// -// -// This file defines a test pass that interprets Transform dialect operations in -// the module. -// -//===----------------------------------------------------------------------===// - -#include "TPP/Passes.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Transform/IR/TransformInterfaces.h" -#include "mlir/IR/BuiltinOps.h" - -using namespace mlir; -using namespace mlir::tpp; - -namespace mlir { -namespace tpp { -#define GEN_PASS_DEF_TRANSFORMDIALECTINTERPRETER -#include "TPP/Passes.h.inc" -#define GEN_PASS_DEF_TRANSFORMDROPSCHEDULE -#include "TPP/Passes.h.inc" -} // namespace tpp -} // namespace mlir - -namespace { - -struct TransformDialectInterpreter - : tpp::impl::TransformDialectInterpreterBase { - void runOnOperation() override { - ModuleOp module = getOperation(); - for (auto op : - module.getBody()->getOps()) { - if (failed(transform::applyTransforms( - module, op, {}, - transform::TransformOptions().enableExpensiveChecks(true)))) - return signalPassFailure(); - } - } -}; - -struct TransformDropSchedule - : tpp::impl::TransformDropScheduleBase { - void runOnOperation() override { - getOperation()->walk([&](Operation *nestedOp) { - if (isa<::mlir::transform::TransformOpInterface>(nestedOp)) { - nestedOp->erase(); - return WalkResult::skip(); - } - return WalkResult::advance(); - }); - } -}; - -} // namespace diff --git a/test/Dialect/Transform/transform-collapse.mlir b/test/Dialect/Transform/transform-collapse.mlir index 67e832026..b54fca34d 100644 --- a/test/Dialect/Transform/transform-collapse.mlir +++ b/test/Dialect/Transform/transform-collapse.mlir @@ -1,9 +1,11 @@ -// RUN: tpp-opt -transform-dialect-interpreter -split-input-file -verify-diagnostics %s | FileCheck %s +// RUN: tpp-opt -transform-interpreter -split-input-file -verify-diagnostics %s | FileCheck %s -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1], [2], [3, 4]] : !transform.any_op -> !transform.any_op + transform.yield + } } // CHECK: #[[MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> @@ -28,11 +30,13 @@ func.func @parallel(%arg0: tensor<5x5x4x3x3xf32>, %arg1: tensor<5x5x4x3x3xf32>) // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1], [2]] : !transform.any_op -> !transform.any_op -} + transform.yield + } +} // CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> @@ -59,10 +63,12 @@ func.func @parallel(%arg0: tensor<5x5x5xf32>, %arg1: tensor<5x5x5xf32>) -> tenso // ----- // This must fail as we attempt to collapse dimensions of different types. -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1, 2]] : !transform.any_op -> !transform.any_op + transform.yield + } } #map0 = affine_map<(i, j, k) -> (i, j)> @@ -83,10 +89,12 @@ func.func @matmul(%arg0: tensor<3x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor< // ----- // This must fail as the reassociation dimensions do not match the number of loops. -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1]] : !transform.any_op -> !transform.any_op + transform.yield + } } #map0 = affine_map<(i, j, k) -> (i, j)> @@ -106,10 +114,12 @@ func.func @matmul(%arg0: tensor<3x2xf32>, %arg1: tensor<2x3xf32>, %arg2: tensor< // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1], [2]] : !transform.any_op -> !transform.any_op + transform.yield + } } #map0 = affine_map<(i, j, k) -> (i, j, k)> @@ -130,10 +140,12 @@ func.func @parallel(%arg0: tensor<3x3x3xf32> , %arg1: tensor<3x3x3xf32>) -> tens // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.collapse %0 [[0, 1], [2]] : !transform.any_op -> !transform.any_op + transform.yield + } } #map0 = affine_map<(i, j, k) -> (i, j)> diff --git a/test/Dialect/Transform/transform-convolutions.mlir b/test/Dialect/Transform/transform-convolutions.mlir index bc176503d..31e6ea621 100644 --- a/test/Dialect/Transform/transform-convolutions.mlir +++ b/test/Dialect/Transform/transform-convolutions.mlir @@ -1,9 +1,9 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s +// RUN: tpp-opt %s -transform-interpreter -canonicalize -split-input-file | FileCheck %s // Map a linalg.conv_2d_nhwc_hwcf to a matmul operation. // Unit filter. -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.generalize %0 : (!transform.any_op) -> !transform.any_op @@ -33,6 +33,8 @@ transform.sequence failures(propagate) { %2 = transform.structured.interchange %1 iterator_interchange = [ 0, 1, 4, 5, 2, 3, 6 ] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv1(%arg0: memref<1x4x4x3xf32>, %arg1: memref<1x1x3x8xf32>, %arg2: memref<1x4x4x8xf32>) { @@ -57,13 +59,15 @@ func.func @conv1(%arg0: memref<1x4x4x3xf32>, %arg1: memref<1x1x3x8xf32>, %arg2: // Map a linalg.conv_2d_nhwc_hwcf to a matmul operation. // Non-unit filter. -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.generalize %0 : (!transform.any_op) -> !transform.any_op %2 = transform.structured.interchange %1 iterator_interchange = [ 0, 1, 4, 5, 2, 3, 6 ] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv2(%arg0: memref<1x4x4x3xf32>, %arg1: memref<2x2x3x8xf32>, %arg2: memref<1x3x3x8xf32>) { @@ -92,13 +96,15 @@ func.func @conv2(%arg0: memref<1x4x4x3xf32>, %arg1: memref<2x2x3x8xf32>, %arg2: // ----- // Unit filter but non-static dims. -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.generalize %0 : (!transform.any_op) -> !transform.any_op %2 = transform.structured.interchange %1 iterator_interchange = [ 0, 1, 4, 5, 2, 3, 6 ] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv3(%arg0: memref, %arg1: memref<1x1x?x?xf32>, %arg2: memref) { @@ -130,8 +136,8 @@ func.func @conv3(%arg0: memref, %arg1: memref<1x1x?x?xf32>, %arg2: // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Original layout: [N][K][P][Q] = [N][C][H][W] * [K][C][R][S] // New layout: [N][K'][P][Q][k] = [N][C'][H][W][c] * [K'][C'][R][S][c][k] @@ -163,6 +169,8 @@ transform.sequence failures(propagate) { %4 = transform.structured.interchange %3 iterator_interchange = [0, 1, 4, 2, 3, 5] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %4 : !transform.any_op + transform.yield + } } func.func @conv(%i: tensor<14x512x28x28xf32>, %f: tensor<1024x512x1x1xf32>, @@ -272,8 +280,8 @@ func.func @walk(%arg0: tensor<1x1x64x64xf32>, %arg1: tensor<3x3x64x64xf32>, %arg return %9 : tensor<1x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Blocks all the convs @@ -287,7 +295,7 @@ transform.sequence failures(propagate) { : (!transform.any_op) -> !transform.any_op %4 = transform.structured.get_blocked_convolutions %3 : (!transform.any_op) -> (!transform.op<"linalg.generic">) - %blocked_matmuls:2 = split_handle %4 + %blocked_matmuls:2 = transform.split_handle %4 : (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.generic">, !transform.op<"linalg.generic">) %first_relu = transform.get_consumers_of_result %blocked_matmuls#0[0] @@ -308,8 +316,8 @@ transform.sequence failures(propagate) { !transform.any_op, !transform.any_op, !transform.any_op) - %6 = get_producer_of_operand %5[0] : (!transform.any_op) -> !transform.any_op - %convs:2 = split_handle %6 : (!transform.any_op) -> (!transform.any_op, !transform.any_op) + %6 = transform.get_producer_of_operand %5[0] : (!transform.any_op) -> !transform.any_op + %convs:2 = transform.split_handle %6 : (!transform.any_op) -> (!transform.any_op, !transform.any_op) // Map the conv to linalg.matmul // With R = S = 3 we map to linalg.matmul @@ -325,17 +333,21 @@ transform.sequence failures(propagate) { %9 = transform.structured.interchange %8 iterator_interchange = [0, 1, 4, 2, 3, 5] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %9 : !transform.any_op + transform.yield + } } // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.generalize %0 : (!transform.any_op) -> !transform.any_op %2 = transform.structured.interchange %1 iterator_interchange = [0, 1, 4, 5, 2, 3, 6] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv2d_stride(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x64x256xf32>, %arg2: tensor<1x56x56x256xf32>) -> tensor<1x56x56x256xf32> { @@ -367,13 +379,15 @@ func.func @conv2d_stride(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x64x2 // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op %2 = transform.structured.interchange %1 iterator_interchange = [0, 1, 2, 5, 6, 7, 3, 4, 8] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv2d_stride(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x64x256xf32>, %arg2: tensor<1x56x56x256xf32>) -> tensor<1x56x56x256xf32> { diff --git a/test/Dialect/Transform/transform-drop-schedule.mlir b/test/Dialect/Transform/transform-drop-schedule.mlir deleted file mode 100644 index fc6907fb3..000000000 --- a/test/Dialect/Transform/transform-drop-schedule.mlir +++ /dev/null @@ -1,20 +0,0 @@ -// RUN: tpp-opt -transform-dialect-interpreter -transform-drop-schedule %s | FileCheck %s - -#map = affine_map<(d0, d1) -> (d0, d1)> - -func.func @relu(%arg0: tensor<128x128xf32>) -> tensor<128x128xf32> { - %c0 = arith.constant 0.0 : f32 - %0 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel"]} outs(%arg0: tensor<128x128xf32>) { - ^bb0(%out: f32): - %1 = arith.maximumf %out, %c0 : f32 - linalg.yield %1 : f32 - } -> tensor<128x128xf32> - return %0 : tensor<128x128xf32> -} - -// CHECK-NOT: transform.sequence -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): - %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 - : (!transform.any_op) -> !transform.any_op -} diff --git a/test/Dialect/Transform/transform-matmuls.mlir b/test/Dialect/Transform/transform-matmuls.mlir index 8cfcd9e44..a15c16dac 100644 --- a/test/Dialect/Transform/transform-matmuls.mlir +++ b/test/Dialect/Transform/transform-matmuls.mlir @@ -1,9 +1,11 @@ -// RUN: tpp-opt -transform-dialect-interpreter -canonicalize -split-input-file %s | FileCheck %s +// RUN: tpp-opt -transform-interpreter -canonicalize -split-input-file %s | FileCheck %s -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1, %loops:3 = transform.structured.tile_using_for %0 [4, 4, 4] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) + transform.yield + } } // CHECK-LABEL: func @tile_linalg_matmul( @@ -36,10 +38,12 @@ func.func @tile_linalg_matmul( // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] : !transform.any_op -> !transform.any_op + transform.yield + } } func.func @block_linalg_matmul( @@ -88,15 +92,15 @@ func.func @matmul_and_relu(%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32 } // Cooking recipe for matmul -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { // Get the matmul %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Pack the matmul with blocking factors of 32 along i, j and k %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] : !transform.any_op -> !transform.any_op // Get parent operation (aka func.func) - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op // Propagate the packing down through the relu transform.structured.packing_propagation %2 : !transform.any_op @@ -130,7 +134,8 @@ transform.sequence failures(propagate) { : (!transform.any_op) -> (!transform.op<"linalg.generic">) %blk_casted = transform.cast %blk_gemm : !transform.op<"linalg.generic"> to !transform.any_op transform.structured.rewrite_to_brgemm %blk_casted : !transform.any_op - + transform.yield + } } // CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)> diff --git a/test/Dialect/Transform/transform-mlp.mlir b/test/Dialect/Transform/transform-mlp.mlir index 58ba33c49..1f6097ae8 100644 --- a/test/Dialect/Transform/transform-mlp.mlir +++ b/test/Dialect/Transform/transform-mlp.mlir @@ -1,4 +1,4 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter -transform-drop-schedule -loop-invariant-code-motion -canonicalize -cse -canonicalize -split-input-file | FileCheck %s +// RUN: tpp-opt %s -transform-interpreter -loop-invariant-code-motion -canonicalize -cse -canonicalize -split-input-file | FileCheck %s #map0 = affine_map<(d0, d1) -> (d1)> #map1 = affine_map<(d0, d1) -> (d0, d1)> @@ -51,8 +51,8 @@ func.func @mlp(%arg0: tensor<128x256xf32>, %arg1: tensor<256x512xf32>, %arg2: te return %15 : tensor<128x1024xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Block matmul i, j and k @@ -70,7 +70,7 @@ transform.sequence failures(propagate) { : (!transform.any_op) -> (!transform.op<"linalg.generic">) // Get the last one, and fuse the outermost dimensions with all the producers - %blocked_matmuls:4 = split_handle %4 : (!transform.op<"linalg.generic">) + %blocked_matmuls:4 = transform.split_handle %4 : (!transform.op<"linalg.generic">) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) %relu = transform.get_consumers_of_result %blocked_matmuls#3[0] : (!transform.any_op) -> (!transform.any_op) @@ -89,7 +89,7 @@ transform.sequence failures(propagate) { : (!transform.any_op) -> !transform.any_op %8 = transform.structured.get_blocked_matmuls %7 : (!transform.any_op) -> (!transform.op<"linalg.generic">) - %rematch_blocked_matmuls:4 = split_handle %8 + %rematch_blocked_matmuls:4 = transform.split_handle %8 : (!transform.op<"linalg.generic">) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) @@ -107,13 +107,15 @@ transform.sequence failures(propagate) { // Fuse matmul + relu and map the matmul to BRGEMM %9, %loop1 = transform.structured.fuse %first_second { tile_sizes = [1, 0, 0] } : (!transform.any_op) -> (!transform.any_op, !transform.any_op) - %10 = get_producer_of_operand %9[0] : (!transform.any_op) -> !transform.any_op + %10 = transform.get_producer_of_operand %9[0] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %10 : !transform.any_op %11, %loop2 = transform.structured.fuse %third_fourth { tile_sizes = [1, 0, 0] } : (!transform.any_op) -> (!transform.any_op, !transform.any_op) - %12 = get_producer_of_operand %11[0] : (!transform.any_op) -> !transform.any_op + %12 = transform.get_producer_of_operand %11[0] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %12 : !transform.any_op + transform.yield + } } // We have 4 layers. 1 loop for each layer and 1 outermost loop for all the layers @@ -129,8 +131,8 @@ transform.sequence failures(propagate) { #map0 = affine_map<(d0, d1) -> (d1)> #map1 = affine_map<(d0, d1) -> (d0, d1)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { // Pack matmul %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] @@ -162,6 +164,8 @@ transform.sequence failures(propagate) { %7 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %7 : !transform.any_op + transform.yield + } } func.func @mlp_single_layer_with_fusion(%A : !A_tensor_t, %B : !B_tensor_t, %C : !C_tensor_t, %Bias: !Bias_tensor_t) -> !C_tensor_t { diff --git a/test/Dialect/Transform/transform-pack.mlir b/test/Dialect/Transform/transform-pack.mlir index 96a2cf94a..2484b0649 100644 --- a/test/Dialect/Transform/transform-pack.mlir +++ b/test/Dialect/Transform/transform-pack.mlir @@ -1,10 +1,12 @@ -// RUN: tpp-opt -transform-dialect-interpreter -verify-diagnostics -split-input-file %s | FileCheck %s +// RUN: tpp-opt -transform-interpreter -verify-diagnostics -split-input-file %s | FileCheck %s -transform.sequence failures(propagate) { -^bb1(%arg0: !transform.any_op): - %0 = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op - // CHECK: transform.structured.pack_ext - %1 = transform.structured.pack_ext %0 blocking_factors = [2, 2, 2] : !transform.any_op -> !transform.any_op +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op + // CHECK: transform.structured.pack_ext + %1 = transform.structured.pack_ext %0 blocking_factors = [2, 2, 2] : !transform.any_op -> !transform.any_op + transform.yield + } } // ----- @@ -21,20 +23,24 @@ func.func @parallel(%arg0: tensor<5x5x5xf32>, %arg1: tensor<5x5x5xf32>) -> tenso return %0 : tensor<5x5x5xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op // expected-error @below {{Could not pack op}} %1 = transform.structured.pack_ext %0 blocking_factors = [2, 2, 2] : !transform.any_op -> !transform.any_op + transform.yield + } } // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op // expected-error @below {{Could not pack op:}} %1 = transform.structured.pack_ext %0 blocking_factors = [200, 200, 200] : !transform.any_op -> !transform.any_op + transform.yield + } } func.func @block_linalg_matmul( @@ -49,10 +55,12 @@ func.func @block_linalg_matmul( // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op + transform.yield + } } func.func @conv_2d_nchw_fchw(%i: tensor<14x512x28x28xf32>, %f: tensor<1024x512x1x1xf32>, @@ -83,10 +91,12 @@ func.func @conv_2d_nchw_fchw(%i: tensor<14x512x28x28xf32>, %f: tensor<1024x512x1 // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op + transform.yield + } } func.func @conv_2d_nhwc_hwcf(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x64x256xf32>, %arg2: tensor<1x56x56x256xf32>) -> tensor<1x56x56x256xf32> { @@ -117,10 +127,12 @@ func.func @conv_2d_nhwc_hwcf(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op + transform.yield + } } // test we properly pass stride information. @@ -152,10 +164,12 @@ func.func @conv_2d_nhwc_hwcf(%arg0: tensor<1x113x113x64xf32>, %arg1: tensor<3x3x // ----- -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] : !transform.any_op -> !transform.any_op + transform.yield + } } func.func @matmul(%arg0: tensor<128x512xf32>, %arg1: tensor<512x256xf32>, %arg2: tensor<128x256xf32>) -> tensor<128x256xf32> { diff --git a/test/Dialect/Transform/transform-propagation.mlir b/test/Dialect/Transform/transform-propagation.mlir index da2c4c6c8..194e08f10 100644 --- a/test/Dialect/Transform/transform-propagation.mlir +++ b/test/Dialect/Transform/transform-propagation.mlir @@ -1,4 +1,4 @@ -// RUN: tpp-opt -transform-dialect-interpreter -verify-diagnostics -split-input-file %s | FileCheck %s +// RUN: tpp-opt -transform-interpreter -verify-diagnostics -split-input-file %s | FileCheck %s #map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> func.func @propagation(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf32> { @@ -13,10 +13,12 @@ func.func @propagation(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf return %2 : tensor<12x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %0 : !transform.any_op + transform.yield + } } // CHECK: func.func @propagation( @@ -42,11 +44,13 @@ func.func @propagation1(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64x return %2 : tensor<12x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op - %1 = get_parent_op %0 : (!transform.any_op) -> !transform.any_op + %1 = transform.get_parent_op %0 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %1 : !transform.any_op + transform.yield + } } // CHECK: func.func @propagation1( @@ -73,11 +77,13 @@ func.func @main(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf32> { return %2 : tensor<12x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op // expected-error @below {{op requires isolated-from-above targets}} transform.structured.packing_propagation %0 : !transform.any_op + transform.yield + } } // ----- @@ -95,12 +101,14 @@ func.func @matmul(%arg0: tensor<128x512xf32>, %arg1: tensor<512x256xf32>, %arg2: return %1 : tensor<128x256xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)> @@ -137,12 +145,14 @@ func.func @matmul(%arg0: tensor<128x512xf32>, %arg1: tensor<512x256xf32>, %arg2: return %1 : tensor<128x256xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)> @@ -182,12 +192,14 @@ func.func @conv(%arg0: tensor<1x56x56x64xf32>, %arg1: tensor<1x1x64x64xf32>, %ar return %1 : tensor<1x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d5, d2 + d6, d3 + d7, d8)> @@ -223,12 +235,14 @@ func.func @conv(%arg0: tensor<1x56x56x64xf32>, %arg1: tensor<1x1x64x64xf32>, %ar return %1 : tensor<1x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d5, d2 + d6, d3 + d7, d8)> @@ -268,12 +282,14 @@ func.func @conv(%arg0: tensor<1x56x56x64xf32>, %arg1: tensor<1x1x64x64xf32>, %ar return %1 : tensor<1x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d5, d2 + d6, d3 + d7, d8)> @@ -315,12 +331,14 @@ func.func @conv(%arg0: tensor<1x56x56x64xf32>, %arg1: tensor<1x1x64x64xf32>, %ar return %1 : tensor<1x56x56x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d5, d2 + d6, d3 + d7, d8)> @@ -367,12 +385,14 @@ func.func @conv(%arg0: tensor<1x56x56x64xf32>, %arg1: tensor<1x1x64x64xf32>, %ar return %2 : tensor<1x58x58x64xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg0 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [32, 32] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op + transform.yield + } } // CONV-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7, d8) -> (d0, d5, d2 + d6, d3 + d7, d8)> diff --git a/test/Dialect/Transform/transform-rewrite-to-brgemm.mlir b/test/Dialect/Transform/transform-rewrite-to-brgemm.mlir index 270659d5f..4ef1ae7a2 100644 --- a/test/Dialect/Transform/transform-rewrite-to-brgemm.mlir +++ b/test/Dialect/Transform/transform-rewrite-to-brgemm.mlir @@ -1,13 +1,15 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s +// RUN: tpp-opt %s -transform-interpreter -canonicalize -split-input-file | FileCheck %s #map3 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d2, d3, d5)> #map4 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d5, d4)> #map5 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } // CHECK-LABEL: func.func @blocked_matmul( @@ -46,10 +48,12 @@ func.func @blocked_matmul(%arg0: tensor<4x16x32x32xf32>, %arg1: tensor<8x16x32x3 #map4 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d5, d4)> #map5 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d3, d4)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } // CHECK-LABEL: func.func @blocked_matmul( @@ -83,10 +87,12 @@ func.func @blocked_matmul(%arg0: memref<4x16x32x32xf32>, %arg1: memref<8x16x32x3 #map6 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)> #map7 = affine_map<(d0, d1, d2, d3) -> (d1, d2)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } // CHECK-LABEL: func.func @blocked_matmul( @@ -109,10 +115,12 @@ func.func @blocked_matmul(%arg0: tensor<8x32x32xf32>, %arg1: tensor<8x32x32xf32> #map2 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d5, d4)> #map3 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } // CHECK-LABEL: func.func @blocked_matmul( @@ -148,10 +156,12 @@ func.func @blocked_matmul(%arg0: tensor<4x8x32x32xf32>, %arg1: tensor<16x8x32x32 #map2 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d2, d5, d4)> #map3 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1, d3, d4)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } // CHECK-LABEL: func.func @blocked_matmul( @@ -190,10 +200,12 @@ func.func @blocked_matmul(%arg0: tensor<4x8x32x32xf32>, %arg1: tensor<16x8x32x32 #map6 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)> #map7 = affine_map<(d0, d1, d2, d3) -> (d1, d2)> -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %0 : !transform.any_op + transform.yield + } } func.func @blocked_matmul(%arg0: tensor, %arg1: tensor, %arg2: tensor<32x32xf32>) -> tensor<32x32xf32> { diff --git a/test/Dialect/Transform/transform-rewrite-to-gemm.mlir b/test/Dialect/Transform/transform-rewrite-to-gemm.mlir index b0a8abe09..41ecc05f7 100644 --- a/test/Dialect/Transform/transform-rewrite-to-gemm.mlir +++ b/test/Dialect/Transform/transform-rewrite-to-gemm.mlir @@ -1,4 +1,4 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter -canonicalize | FileCheck %s +// RUN: tpp-opt %s -transform-interpreter -canonicalize | FileCheck %s #map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d2, d4 + d3, d6)> #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d2, d3, d6, d5)> @@ -6,7 +6,10 @@ func.func @conv(%arg0: tensor<1x4x4x3xi64>, %arg1: tensor<2x2x3x8xi64>, %arg2: tensor<1x3x3x8xi64>) -> tensor<1x3x3x8xi64> { // CHECK: linalg.matmul - %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction", "reduction", "parallel", "parallel", "reduction"]} ins(%arg0, %arg1 : tensor<1x4x4x3xi64>, tensor<2x2x3x8xi64>) outs(%arg2 : tensor<1x3x3x8xi64>) { + %0 = linalg.generic { + indexing_maps = [#map, #map1, #map2], + iterator_types = ["parallel", "parallel", "reduction", "reduction", "parallel", "parallel", "reduction"]} + ins(%arg0, %arg1 : tensor<1x4x4x3xi64>, tensor<2x2x3x8xi64>) outs(%arg2 : tensor<1x3x3x8xi64>) { ^bb0(%in: i64, %in_0: i64, %out: i64): %1 = arith.muli %in, %in_0 : i64 %2 = arith.addi %out, %1 : i64 @@ -15,8 +18,10 @@ func.func @conv(%arg0: tensor<1x4x4x3xi64>, %arg1: tensor<2x2x3x8xi64>, %arg2: t return %0 : tensor<1x3x3x8xi64> } -transform.sequence failures(propagate) { - ^bb0(%arg0: !transform.any_op): - %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { + %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %0 : !transform.any_op + transform.yield + } } diff --git a/test/Integration/transform/transform-convolution.mlir b/test/Integration/transform/transform-convolution.mlir index c2dd73674..5464d27f3 100644 --- a/test/Integration/transform/transform-convolution.mlir +++ b/test/Integration/transform/transform-convolution.mlir @@ -1,7 +1,7 @@ // This should really be in the passes directory, not here -// RUN: tpp-opt %s -transform-dialect-interpreter | FileCheck %s -check-prefix=IR +// RUN: tpp-opt %s -transform-interpreter | FileCheck %s -check-prefix=IR -// RUN: tpp-opt %s -transform-drop-schedule | \ +// RUN: tpp-opt %s -transform-interpreter | \ // RUN: tpp-run -print \ // RUN: -e entry -entry-point-result=void | \ // RUN: FileCheck %s @@ -38,18 +38,18 @@ func.func @conv(%arg0: tensor<1x1x8x8xf32>, %arg1: tensor<8xf32>, %conv_out: ten return %3 : tensor<1x6x6x8xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Blocks all the convs %1 = transform.structured.pack_ext %0 blocking_factors = [2, 2] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 : (!transform.any_op) -> !transform.any_op // Propagate all the packs transform.structured.packing_propagation %2 : !transform.any_op %3 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Here we match the generic ops in the entire module. The conv is the last one. - %generic:3 = split_handle %3 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op) + %generic:3 = transform.split_handle %3 : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op) %4 = transform.structured.collapse %generic#2 [[0], [1], [2], [3], [4], [5, 6, 7], [8]] : !transform.any_op -> !transform.any_op %5 = transform.structured.collapse %4 [[0], [1], [2, 3], [4], [5], [6]] @@ -57,6 +57,8 @@ transform.sequence failures(propagate) { %6 = transform.structured.interchange %5 iterator_interchange = [0, 1, 4, 2, 3, 5] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %6 : !transform.any_op + transform.yield + } } func.func @entry() { diff --git a/test/Integration/transform/transform-convolutions.mlir b/test/Integration/transform/transform-convolutions.mlir index 07f44460e..c30265b24 100644 --- a/test/Integration/transform/transform-convolutions.mlir +++ b/test/Integration/transform/transform-convolutions.mlir @@ -1,4 +1,4 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter | FileCheck %s -check-prefix=IR +// RUN: tpp-opt %s -transform-interpreter | FileCheck %s -check-prefix=IR // RUN: tpp-run %s -print \ // RUN: -e entry -entry-point-result=void | \ @@ -58,8 +58,8 @@ func.func @walk(%arg0: tensor<1x1x8x8xf32>, %arg1: tensor<3x3x8x8xf32>, %arg2: t return %9 : tensor<1x6x6x8xf32> } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op // Blocks all the convs @@ -74,7 +74,7 @@ transform.sequence failures(propagate) { : (!transform.any_op) -> !transform.any_op %4 = transform.structured.get_blocked_convolutions %3 : (!transform.any_op) -> (!transform.op<"linalg.generic">) - %blocked_matmuls:2 = split_handle %4 + %blocked_matmuls:2 = transform.split_handle %4 : (!transform.op<"linalg.generic">) -> (!transform.op<"linalg.generic">, !transform.op<"linalg.generic">) %first_relu = transform.get_consumers_of_result %blocked_matmuls#0[0] @@ -91,8 +91,8 @@ transform.sequence failures(propagate) { // Fuse relu and conv on the three outermost loops %5, %loop:3 = transform.structured.fuse %relus { tile_sizes = [1, 1, 1, 0, 0] } : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op) - %6 = get_producer_of_operand %5[0] : (!transform.any_op) -> !transform.any_op - %convs:2 = split_handle %6 : (!transform.any_op) -> (!transform.any_op, !transform.any_op) + %6 = transform.get_producer_of_operand %5[0] : (!transform.any_op) -> !transform.any_op + %convs:2 = transform.split_handle %6 : (!transform.any_op) -> (!transform.any_op, !transform.any_op) // Map the conv to linalg.matmul // With R = S = 3 we map to linalg.matmul @@ -109,6 +109,8 @@ transform.sequence failures(propagate) { %9 = transform.structured.interchange %8 iterator_interchange = [0, 1, 4, 2, 3, 5] : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %9 : !transform.any_op + transform.yield + } } func.func @entry() { diff --git a/test/Integration/transform/transform-matmul.mlir b/test/Integration/transform/transform-matmul.mlir index d8d672222..29e9aa5c4 100644 --- a/test/Integration/transform/transform-matmul.mlir +++ b/test/Integration/transform/transform-matmul.mlir @@ -1,6 +1,6 @@ -// RUN: tpp-opt %s -transform-dialect-interpreter | FileCheck %s -check-prefix=IR +// RUN: tpp-opt %s -transform-interpreter | FileCheck %s -check-prefix=IR -// RUN: tpp-opt %s -transform-drop-schedule | \ +// RUN: tpp-opt %s -transform-interpreter | \ // RUN: tpp-run -print \ // RUN: -e entry -entry-point-result=void | \ // RUN: FileCheck %s @@ -52,15 +52,17 @@ func.func @matmul_static( return } -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.pack_ext %0 blocking_factors = [2, 2, 2] : !transform.any_op -> !transform.any_op - %2 = get_parent_op %1 {isolated_from_above} : (!transform.any_op) -> !transform.any_op + %2 = transform.get_parent_op %1 {isolated_from_above} : (!transform.any_op) -> !transform.any_op transform.structured.packing_propagation %2 : !transform.any_op %3 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_to_brgemm %3 : !transform.any_op + transform.yield + } } func.func @entry() { diff --git a/test/Integration/transform/transform-rewrite-conv-to-gemm.mlir b/test/Integration/transform/transform-rewrite-conv-to-gemm.mlir index b356cbc4f..7bc367708 100644 --- a/test/Integration/transform/transform-rewrite-conv-to-gemm.mlir +++ b/test/Integration/transform/transform-rewrite-conv-to-gemm.mlir @@ -1,7 +1,7 @@ // Make sure we map to linalg.matmul -// RUN: tpp-opt %s -transform-dialect-interpreter -transform-drop-schedule | FileCheck -check-prefix=LINALG %s +// RUN: tpp-opt %s -transform-interpreter | FileCheck -check-prefix=LINALG %s -// RUN: tpp-opt %s -transform-drop-schedule | \ +// RUN: tpp-opt %s -transform-interpreter | \ // RUN: tpp-run -print \ // RUN: -e entry -entry-point-result=void | \ // RUN: FileCheck %s @@ -10,13 +10,16 @@ // RUN: -e entry -entry-point-result=void | \ // RUN: FileCheck %s -transform.sequence failures(propagate) { - ^bb0(%arg1: !transform.any_op): +module attributes {transform.with_named_sequence} { + transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) { %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op %1 = transform.structured.generalize %0 : (!transform.any_op) -> !transform.any_op - %2 = transform.structured.interchange %1 iterator_interchange = [ 0, 1, 4, 5, 2, 3, 6 ] : (!transform.any_op) -> !transform.any_op + %2 = transform.structured.interchange %1 iterator_interchange = [ 0, 1, 4, 5, 2, 3, 6 ] + : (!transform.any_op) -> !transform.any_op transform.structured.rewrite_conv_to_matmul %2 : !transform.any_op + transform.yield + } } func.func @conv(%img: tensor<1x4x4x3xf32>, %filt: tensor<2x2x3x8xf32>,