diff --git a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp @@ -880,8 +880,10 @@ // Further check that the resulting index maps can be fused and // inverted. Without this the resultant op is not legal. - if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) - return op.emitRemark("fused op loop bound computation failed"); + if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) { + return rewriter.notifyMatchFailure( + op, "fused op loop bound computation failed"); + } rewriter.startRootUpdate(op); op->setOperands(fusedOperands); @@ -973,15 +975,19 @@ linearizeCollapsedDims(invMap, reshapeOp.getSrcType().getShape(), reshapeOp.getReassociationMaps()); for (AffineExpr expr : modifiedMap.getResults()) { - if (!expr.isPureAffine()) - return producer.emitRemark("fused op indexing map is not affine"); + if (!expr.isPureAffine()) { + return rewriter.notifyMatchFailure( + producer, "fused op indexing map is not affine"); + } } fusedIndexMaps.back() = modifiedMap; // Further check that the resulting index maps can be fused and // inverted. Without this the resultant op is not legal. - if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) - return reshapeOp.emitRemark("fused op loop bound computation failed"); + if (!inversePermutation(concatAffineMaps(fusedIndexMaps))) { + return rewriter.notifyMatchFailure( + producer, "fused op loop bound computation failed"); + } Location loc = producer.getLoc(); Value output = rewriter.create( diff --git a/mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir b/mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir --- a/mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir +++ b/mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt -split-input-file -linalg-fold-reshape-ops-by-linearization -verify-diagnostics %s | FileCheck %s +// RUN: mlir-opt -split-input-file -linalg-fold-reshape-ops-by-linearization %s | FileCheck %s #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> func @generic_op_reshape_producer_fusion(%arg0 : tensor, @@ -257,7 +257,6 @@ %arg1 : tensor) -> tensor { - // expected-remark @+1 {{fused op indexing map is not affine}} %0 = linalg.generic { indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} @@ -272,3 +271,10 @@ tensor into tensor return %1 : tensor } +// CHECK-LABEL: func @generic_op_reshape_consumer_nofusion +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor +// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor +// CHECK: %[[NOFUSE:.+]] = linalg.generic +// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] +// CHECK: %[[RESULT:.+]] = linalg.tensor_reshape %[[NOFUSE]] +// CHECK: return %[[RESULT]]