[DispatchCreation] Drop fusion restriction for stride != 1 conv (#19634) There was a restriction on aggressive fusion that prevented fusion with non-unit stride convs due to lack of support in the backend. The backend now supports such cases so we can re-enable those fusions.
diff --git a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp index 7cac15a..1b0d02e 100644 --- a/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp +++ b/compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp
@@ -425,21 +425,6 @@ return true; } -/// All operations in a dispatch should be vectorized, which isnt the case today -/// This is an explicit list of operations that arent vectorized for now -/// requiring special handling for now in dispatch region formation to avoid -/// large stack allocations. -static bool isVectorizedAlways(Operation *producer) { - // TODO(#17155) : This is a black list of operations that are not vectorized - // today (under the aggressive fusion flag). Remove this blacklist to return - // true always. - if (auto convOp = dyn_cast<linalg::Conv2DNhwcHwcfOp>(producer)) { - auto strides = convOp.getStrides(); - return strides.isSplat() && strides.getSplatValue<int64_t>() == 1; - } - return true; -} - /// Returns true if this is a fusable use, while fusing a root with its /// consumer. static bool @@ -554,9 +539,7 @@ // Under aggressive fusion assume that the dispatches are vectorized. In which // case we dont need to account for the subsequent stack allocation condition. if (options.aggressiveFusion) { - if (isVectorizedAlways(producer)) { - return true; - } + return true; } // While fusing with consumer, the result of the root might not be the final
diff --git a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir index acbeed9..16beb5b 100644 --- a/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir +++ b/compiler/src/iree/compiler/DispatchCreation/test/dispatch_linalg_on_tensors.mlir
@@ -2268,7 +2268,7 @@ // ----- -util.func @mixed_conv_unsupported(%arg0 : tensor<2x130x130x320xf16>, %arg1 : tensor<3x3x320x320xf16>) -> tensor<2x64x64x320xf16> { +util.func @mixed_conv_stride_2(%arg0 : tensor<2x130x130x320xf16>, %arg1 : tensor<3x3x320x320xf16>) -> tensor<2x64x64x320xf16> { %empty = tensor.empty() : tensor<2x64x64x320xf32> %cst = arith.constant 0.0 : f32 %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<2x64x64x320xf32>) -> tensor<2x64x64x320xf32> @@ -2288,16 +2288,15 @@ } -> tensor<2x64x64x320xf16> util.return %truncf : tensor<2x64x64x320xf16> } -// CHECK-LABEL: func public @mixed_conv_unsupported( -// CHECK: %[[DISPATCH0:.+]] = flow.dispatch.workgroups +// CHECK-LABEL: func public @mixed_conv_stride_2( +// CHECK: %[[DISPATCH:.+]] = flow.dispatch.workgroups // CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf // CHECK-SAME: outs(%[[FILL]] : -// CHECK: flow.dispatch.tensor.store %[[CONV]] -// CHECK: %[[DISPATCH1:.+]] = flow.dispatch.workgroups // CHECK: %[[GENERIC:.+]] = linalg.generic +// CHECK-SAME: ins(%[[CONV]] // CHECK: flow.dispatch.tensor.store %[[GENERIC]] -// CHECK: util.return %[[DISPATCH1]] +// CHECK: util.return %[[DISPATCH]] // -----