[Codegen] Add canonicalization pattern for reshape into interface stores (#12471)
This pattern is similar to the existing one for loads except for stores. This is needed for implicit gemm where a reshape ends up at the bottom of the kernel.This pattern is similar to the existing one for loads except for stores. This is needed for implicit gemm where a reshape ends up at the bottom of the kernel.
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir b/compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir
index aeec7aa..deae76f 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/canonicalize_interface_load_store.mlir
@@ -1,9 +1,10 @@
// RUN: iree-opt --split-input-file --iree-codegen-cleanup-buffer-alloc-view %s | FileCheck %s
-// CHECK-LABEL: func.func @fold_reshape()
-func.func @fold_reshape() {
+// CHECK-LABEL: func.func @fold_reshape_load()
+func.func @fold_reshape_load() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
+ %cst = arith.constant 0.0 : f32
// CHECK: %[[ARG:.+]] = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<readonly:tensor<3x3x96xf32>>
%1 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<readonly:tensor<3x3x1x96xf32>>
%2 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<writeonly:tensor<3x3x96xf32>>
@@ -11,8 +12,31 @@
%3 = flow.dispatch.tensor.load %1, offsets=[0, 0, 0, 0], sizes =[3, 3, 1, 96], strides=[1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<3x3x1x96xf32>> -> tensor<3x3x1x96xf32>
%4 = tensor.collapse_shape %3 [[0, 1, 2, 3]] : tensor<3x3x1x96xf32> into tensor<864xf32>
%5 = tensor.expand_shape %4 [[0, 1, 2]] : tensor<864xf32> into tensor<3x3x96xf32>
- // CHECK: flow.dispatch.tensor.store %[[LOAD]], {{.*}}
- flow.dispatch.tensor.store %5, %2, offsets = [%c0, %c0, %c0], sizes = [%c1, %c1, %c1], strides = [%c1, %c1, %c1] : tensor<3x3x96xf32> -> !flow.dispatch.tensor<writeonly:tensor<3x3x96xf32>>
+ // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[LOAD]] : tensor<3x3x96xf32>)
+ %6 = linalg.fill ins(%cst : f32) outs(%5 : tensor<3x3x96xf32>) -> tensor<3x3x96xf32>
+ // CHECK: flow.dispatch.tensor.store %[[FILL]], {{.*}}
+ flow.dispatch.tensor.store %6, %2, offsets = [%c0, %c0, %c0], sizes = [%c1, %c1, %c1], strides = [%c1, %c1, %c1] : tensor<3x3x96xf32> -> !flow.dispatch.tensor<writeonly:tensor<3x3x96xf32>>
+ return
+}
+
+// -----
+
+// CHECK-LABEL: func.func @fold_reshape_store()
+func.func @fold_reshape_store() {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.0 : f32
+ %1 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<readonly:tensor<3x3x1x96xf32>>
+ %2 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<writeonly:tensor<3x3x96xf32>>
+ // CHECK: %[[OUT:.+]] = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) : !flow.dispatch.tensor<writeonly:tensor<3x3x1x96xf32>>
+ // CHECK: %[[LOAD:.+]] = flow.dispatch.tensor.load %{{.*}}, {{.*}}
+ %3 = flow.dispatch.tensor.load %1, offsets=[0, 0, 0, 0], sizes =[3, 3, 1, 96], strides=[1, 1, 1, 1] : !flow.dispatch.tensor<readonly:tensor<3x3x1x96xf32>> -> tensor<3x3x1x96xf32>
+ // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[LOAD]] : tensor<3x3x1x96xf32>)
+ %4 = linalg.fill ins(%cst : f32) outs(%3 : tensor<3x3x1x96xf32>) -> tensor<3x3x1x96xf32>
+ %5 = tensor.collapse_shape %4 [[0, 1, 2, 3]] : tensor<3x3x1x96xf32> into tensor<864xf32>
+ %6 = tensor.expand_shape %5 [[0, 1, 2]] : tensor<864xf32> into tensor<3x3x96xf32>
+ // CHECK: flow.dispatch.tensor.store %[[FILL]], %[[OUT]], {{.+}} : tensor<3x3x1x96xf32> -> !flow.dispatch.tensor<writeonly:tensor<3x3x1x96xf32>>
+ flow.dispatch.tensor.store %6, %2, offsets = [0, 0, 0], sizes = [3, 3, 96], strides = [1, 1, 1] : tensor<3x3x96xf32> -> !flow.dispatch.tensor<writeonly:tensor<3x3x96xf32>>
return
}
diff --git a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
index f8a1221..18a8089 100644
--- a/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
+++ b/compiler/src/iree/compiler/Codegen/Transforms/Transforms.cpp
@@ -254,6 +254,17 @@
namespace {
+// TODO(antigainst): enable dynamic shape support once they are needed.
+template <typename TensorReshapeOp>
+static Optional<Value> getStaticReshapeOpSrc(TensorReshapeOp reshapeOp) {
+ auto reshapeSrcType =
+ reshapeOp.getSrc().getType().template cast<ShapedType>();
+ auto reshapeDstType = reshapeOp.getType().template cast<ShapedType>();
+ if (!reshapeSrcType.hasStaticShape() || !reshapeDstType.hasStaticShape())
+ return std::nullopt;
+ return reshapeOp.getSrc();
+}
+
/// Folds tensor.expand/collapse_shape into the source
/// hal.interface.binding.subspan.
///
@@ -280,17 +291,12 @@
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
PatternRewriter &rewriter) const override {
- // TODO(antigainst): enable dynamic shape support once they are needed.
- auto reshapeSrcType =
- reshapeOp.getSrc().getType().template cast<ShapedType>();
- auto reshapeDstType = reshapeOp.getType().template cast<ShapedType>();
- if (!reshapeSrcType.hasStaticShape() || !reshapeDstType.hasStaticShape()) {
- return failure();
- }
+ Optional<Value> reshapeSrc =
+ getStaticReshapeOpSrc<TensorReshapeOp>(reshapeOp);
+ if (!reshapeSrc) return failure();
auto loadOp =
- reshapeOp.getSrc()
- .template getDefiningOp<IREE::Flow::DispatchTensorLoadOp>();
+ reshapeSrc->template getDefiningOp<IREE::Flow::DispatchTensorLoadOp>();
if (!loadOp) return failure();
// Make sure we are loading the full incoming subspan. Otherwise we cannot
@@ -324,12 +330,88 @@
return success();
}
};
+
+/// Folds tensor.expand/collapse_shape into the source
+/// hal.interface.binding.subspan.
+///
+/// For example, this matches the following pattern:
+///
+/// %subspan = hal.interface.binding.subspan ... :
+/// !flow.dispatch.tensor<writeonly:tensor<3x3x1x96xf32>>
+/// %0 = linalg.tensor_reshape %tensor [
+/// affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+/// ] : tensor<864xf32> into tensor<3x3x1x96xf32>
+/// %tensor = flow.dispatch.tensor.store %0, %subspan :
+/// !flow.dispatch.tensor<writeonly:tensor<3x3x1x96xf32>> ->
+/// tensor<3x3x1x96xf32>
+///
+/// And turns it into:
+///
+/// %subspan = hal.interface.binding.subspan ... :
+/// !flow.dispatch.tensor<writeonly:tensor<864xf32>>
+/// %0 = flow.dispatch.tensor.store %tensor, %subspan :
+/// !flow.dispatch.tensor<writeonly:tensor<864xf32>> -> tensor<864xf32>
+struct FoldReshapeIntoInterfaceTensorStore
+ : OpRewritePattern<IREE::Flow::DispatchTensorStoreOp> {
+ using OpRewritePattern<IREE::Flow::DispatchTensorStoreOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(IREE::Flow::DispatchTensorStoreOp storeOp,
+ PatternRewriter &rewriter) const override {
+ // Make sure we are storing the full incoming subspan. Otherwise we cannot
+ // simply adjust the subspan's resultant type later.
+ if (!storeOp.offsets().empty() || !storeOp.sizes().empty() ||
+ !storeOp.strides().empty())
+ return failure();
+
+ auto reshapeOp = storeOp.getValue().getDefiningOp();
+ if (!isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp>(reshapeOp))
+ return failure();
+
+ // Dynamic shapes are currently unsupported.
+ Optional<Value> reshapeSrc =
+ isa<tensor::CollapseShapeOp>(reshapeOp)
+ ? getStaticReshapeOpSrc<tensor::CollapseShapeOp>(
+ cast<tensor::CollapseShapeOp>(reshapeOp))
+ : getStaticReshapeOpSrc<tensor::ExpandShapeOp>(
+ cast<tensor::ExpandShapeOp>(reshapeOp));
+ if (!reshapeSrc) return failure();
+
+ auto subspanOp =
+ storeOp.getTarget()
+ .template getDefiningOp<IREE::HAL::InterfaceBindingSubspanOp>();
+ if (!subspanOp) return failure();
+ assert(subspanOp.getDynamicDims().empty());
+
+ auto tensorAccess = subspanOp.getType()
+ .template cast<IREE::Flow::DispatchTensorType>()
+ .getAccess();
+ auto newSubspanType = IREE::Flow::DispatchTensorType::get(
+ tensorAccess, reshapeSrc->getType());
+
+ Value newSubspanOp;
+ {
+ OpBuilder::InsertionGuard guard(rewriter);
+ rewriter.setInsertionPointAfter(subspanOp);
+ newSubspanOp = rewriter.create<IREE::HAL::InterfaceBindingSubspanOp>(
+ subspanOp.getLoc(), newSubspanType, subspanOp.getSet(),
+ subspanOp.getBinding(), subspanOp.getDescriptorType(),
+ subspanOp.getByteOffset(), subspanOp.getDynamicDims(),
+ subspanOp.getAlignmentAttr(), subspanOp.getDescriptorFlagsAttr());
+ }
+
+ rewriter.replaceOpWithNewOp<IREE::Flow::DispatchTensorStoreOp>(
+ storeOp, *reshapeSrc, newSubspanOp, storeOp.getTargetDims());
+
+ return success();
+ }
+};
} // namespace
void populateReshapeToInterfaceTensorPatterns(RewritePatternSet &patterns) {
patterns.insert<FoldReshapeIntoInterfaceTensorLoad<tensor::CollapseShapeOp>,
FoldReshapeIntoInterfaceTensorLoad<tensor::ExpandShapeOp>>(
patterns.getContext());
+ patterns.insert<FoldReshapeIntoInterfaceTensorStore>(patterns.getContext());
}
//===--------------------------------------------------------------------====//
