diff --git a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp @@ -34,9 +34,6 @@ // Check if the operation is a LinalgOp but not a GenericOp. if (!namedOp || isa(op)) return failure(); - // Check if the operation has a region builder. - if (!namedOp.getRegionBuilder()) - return failure(); return success(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp b/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/LinalgStrategyPasses.cpp @@ -176,9 +176,8 @@ generalizationPattern.add( funcOp.getContext(), filter); } - if (failed(applyPatternsAndFoldGreedily(funcOp, - std::move(generalizationPattern)))) - signalPassFailure(); + (void)applyPatternsAndFoldGreedily(funcOp, + std::move(generalizationPattern)); } LinalgTransformationFilter filter; diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -643,6 +643,10 @@ Operation *op, PatternRewriter &rewriter) const { if (failed(filter.checkAndNotify(rewriter, op))) return failure(); + if (isa(op)) { + filter.replaceLinalgTransformationFilter(rewriter, op); + return success(); + } if (failed(generalizeNamedOpPrecondition(op))) return failure(); diff --git a/mlir/test/Dialect/Linalg/codegen-strategy.mlir b/mlir/test/Dialect/Linalg/codegen-strategy.mlir --- a/mlir/test/Dialect/Linalg/codegen-strategy.mlir +++ b/mlir/test/Dialect/Linalg/codegen-strategy.mlir @@ -1,8 +1,9 @@ // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=2,4,8 vectorize vectorize-contraction-to=matrixintrinsics unroll-vector-transfers=true" -split-input-file | FileCheck %s --check-prefix=CHECK-INTRINSIC // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 promote promote-full-tile-pad register-tile-sizes=2,4,8 vectorize vectorize-contraction-to=outerproduct split-transfers=true unroll-vector-transfers=false" -split-input-file | FileCheck %s --check-prefix=CHECK-OUTER -// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-INTERCHANGE +// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 tile-interchange=1,2,0 generalize iterator-interchange=0,2,1" -split-input-file | FileCheck %s --check-prefix=CHECK-TILE-AND-INTERCHANGE // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 pad pack-paddings=1,1,0 hoist-paddings=3,3,0" -split-input-file | FileCheck %s --check-prefix=CHECK-PAD // RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matmul anchor-op=linalg.matmul tile-sizes=16,32,64 fuse pad vectorize" -split-input-file | FileCheck %s --check-prefix=CHECK-FUSE +// RUN: mlir-opt %s -test-linalg-codegen-strategy="anchor-func=matvec generalize iterator-interchange=1,0" -split-input-file | FileCheck %s --check-prefix=CHECK-GENERALIZE // CHECK-INTRINSIC: func @matmul( // CHECK-OUTER: func @matmul( @@ -21,20 +22,20 @@ // ----- -// CHECK-INTERCHANGE: func @matmul( +// CHECK-TILE-AND-INTERCHANGE: func @matmul( func @matmul(%arg0: tensor<72x72xf32>, %arg1: tensor<72x72xf32>, %arg2: tensor<72x72xf32>) -> tensor<72x72xf32> { - // CHECK-INTERCHANGE-DAG: %[[C16:.*]] = arith.constant 16 - // CHECK-INTERCHANGE-DAG: %[[C32:.*]] = arith.constant 32 - // CHECK-INTERCHANGE-DAG: %[[C64:.*]] = arith.constant 64 + // CHECK-TILE-AND-INTERCHANGE-DAG: %[[C16:.*]] = arith.constant 16 + // CHECK-TILE-AND-INTERCHANGE-DAG: %[[C32:.*]] = arith.constant 32 + // CHECK-TILE-AND-INTERCHANGE-DAG: %[[C64:.*]] = arith.constant 64 // Check the tile loops are interchanged. - // CHECK-INTERCHANGE: scf.for {{.*}} step %[[C32]] - // CHECK-INTERCHANGE: scf.for {{.*}} step %[[C64]] - // CHECK-INTERCHANGE: scf.for {{.*}} step %[[C16]] + // CHECK-TILE-AND-INTERCHANGE: scf.for {{.*}} step %[[C32]] + // CHECK-TILE-AND-INTERCHANGE: scf.for {{.*}} step %[[C64]] + // CHECK-TILE-AND-INTERCHANGE: scf.for {{.*}} step %[[C16]] // Check the operation has been generalized and interchanged. - // CHECK-INTERCHANGE: linalg.generic - // CHECK-INTERCHANGE-SAME: iterator_types = ["parallel", "reduction", "parallel"] + // CHECK-TILE-AND-INTERCHANGE: linalg.generic + // CHECK-TILE-AND-INTERCHANGE-SAME: iterator_types = ["parallel", "reduction", "parallel"] %0 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%arg2: tensor<72x72xf32>) -> tensor<72x72xf32> return %0 : tensor<72x72xf32> } @@ -74,3 +75,29 @@ %1 = linalg.matmul ins(%arg0, %arg1: tensor<72x72xf32>, tensor<72x72xf32>) outs(%0: tensor<72x72xf32>) -> tensor<72x72xf32> return %1 : tensor<72x72xf32> } + +// ----- + +#map0 = affine_map<(d0, d1) -> (d0, d1)> +#map1 = affine_map<(d0, d1) -> (d0)> + +// CHECK-GENERALIZE: func @matvec( +func @matvec(%arg0: tensor<72x72xf32>, %arg1: tensor<72xf32>, %arg2: tensor<72xf32>) -> tensor<72xf32> { + // Check the generic op iterators are interchanged although generalization was not needed. + // CHECK-GENERALIZE: linalg.generic + // CHECK-GENERALIZE-SAME: iterator_types = ["reduction", "parallel"] + %0 = linalg.generic {indexing_maps = [#map0, #map1], + iterator_types = ["parallel", "reduction"]} + ins(%arg0 : tensor<72x72xf32>) + outs(%arg1 : tensor<72xf32>) { + ^bb0(%arg3: f32, %arg5: f32): // no predecessors + %2 = arith.addf %arg3, %arg5 : f32 + linalg.yield %2 : f32 + } -> tensor<72xf32> + + // Check matvec is generalized and its iterators are interchanged. + // CHECK-GENERALIZE: linalg.generic + // CHECK-GENERALIZE-SAME: iterator_types = ["reduction", "parallel"] + %1 = linalg.matvec ins(%arg0, %0: tensor<72x72xf32>, tensor<72xf32>) outs(%arg2: tensor<72xf32>) -> tensor<72xf32> + return %1 : tensor<72xf32> +}