mirror of
https://github.com/intel/llvm.git
synced 2026-01-30 14:07:28 +08:00
[mlir][vector] LISH: Implement SubsetOpInterface for transfer_read/write (#70629)
- Implement `SubsetOpInterface`, `SubsetExtractionOpInterface`, `SubsetInsertionOpInterface` for `vector.transfer_read` and `vector.transfer_write`. - Move all tensor subset hoisting test cases from `Linalg` to `loop-invariant-subset-hoisting.mlir`. (Removing 1 duplicate test case.)
This commit is contained in:
Matthias Springer
committed by
GitHub
GitHub
parent
773caac711
commit
1df6504ac2
@@ -0,0 +1,20 @@
|
||||
//===- SubsetOpInterfaceImpl.h - Tensor subsets -----------------*- C++ -*-===//
|
||||
//
|
||||
// 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#ifndef MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H
|
||||
#define MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H
|
||||
|
||||
namespace mlir {
|
||||
class DialectRegistry;
|
||||
|
||||
namespace vector {
|
||||
void registerSubsetOpInterfaceExternalModels(DialectRegistry ®istry);
|
||||
} // namespace vector
|
||||
} // namespace mlir
|
||||
|
||||
#endif // MLIR_DIALECT_VECTOR_SUBSETOPINTERFACEIMPL_H
|
||||
@@ -85,6 +85,7 @@
|
||||
#include "mlir/Dialect/UB/IR/UBOps.h"
|
||||
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
||||
#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
|
||||
#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
|
||||
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
|
||||
#include "mlir/IR/Dialect.h"
|
||||
#include "mlir/Interfaces/CastInterfaces.h"
|
||||
@@ -171,6 +172,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
|
||||
tensor::registerTilingInterfaceExternalModels(registry);
|
||||
tensor::registerValueBoundsOpInterfaceExternalModels(registry);
|
||||
vector::registerBufferizableOpInterfaceExternalModels(registry);
|
||||
vector::registerSubsetOpInterfaceExternalModels(registry);
|
||||
NVVM::registerNVVMTargetInterfaceExternalModels(registry);
|
||||
ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
|
||||
LowerVectorShapeCast.cpp
|
||||
LowerVectorTransfer.cpp
|
||||
LowerVectorTranspose.cpp
|
||||
SubsetOpInterfaceImpl.cpp
|
||||
VectorDistribute.cpp
|
||||
VectorDropLeadUnitDim.cpp
|
||||
VectorEmulateNarrowType.cpp
|
||||
@@ -40,6 +41,7 @@ add_mlir_dialect_library(MLIRVectorTransforms
|
||||
MLIRMemRefUtils
|
||||
MLIRSCFDialect
|
||||
MLIRSideEffectInterfaces
|
||||
MLIRSubsetOpInterface
|
||||
MLIRTensorDialect
|
||||
MLIRTransforms
|
||||
MLIRVectorDialect
|
||||
|
||||
82
mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
Normal file
82
mlir/lib/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.cpp
Normal file
@@ -0,0 +1,82 @@
|
||||
//===- SubsetOpInterfaceImpl.cpp - Tensor subsets -------------------------===//
|
||||
//
|
||||
// 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
|
||||
//
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
#include "mlir/Dialect/Vector/Transforms/SubsetOpInterfaceImpl.h"
|
||||
|
||||
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
||||
#include "mlir/Interfaces/SubsetOpInterface.h"
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::vector;
|
||||
|
||||
namespace {
|
||||
|
||||
template <typename OpTy>
|
||||
struct XferOpSubsetOpInterface
|
||||
: public SubsetOpInterface::ExternalModel<XferOpSubsetOpInterface<OpTy>,
|
||||
OpTy> {
|
||||
FailureOr<HyperrectangularSlice>
|
||||
getAccessedHyperrectangularSlice(Operation *op) const {
|
||||
auto xferOp = cast<OpTy>(op);
|
||||
Builder b(xferOp->getContext());
|
||||
SmallVector<OpFoldResult> offsets = llvm::map_to_vector(
|
||||
xferOp.getIndices(), [](Value v) -> OpFoldResult { return v; });
|
||||
SmallVector<OpFoldResult> sizes = llvm::map_to_vector(
|
||||
xferOp.getTransferChunkAccessed(),
|
||||
[&](int64_t sz) -> OpFoldResult { return b.getIndexAttr(sz); });
|
||||
return HyperrectangularSlice(offsets, sizes);
|
||||
}
|
||||
};
|
||||
|
||||
struct TransferReadOpSubsetExtractionOpInterface
|
||||
: public SubsetExtractionOpInterface::ExternalModel<
|
||||
TransferReadOpSubsetExtractionOpInterface, vector::TransferReadOp> {
|
||||
OpOperand &getSourceOperand(Operation *op) const {
|
||||
return cast<vector::TransferReadOp>(op).getSourceMutable();
|
||||
}
|
||||
};
|
||||
|
||||
struct TransferWriteOpSubsetInsertionOpInterface
|
||||
: public SubsetInsertionOpInterface::ExternalModel<
|
||||
TransferWriteOpSubsetInsertionOpInterface, vector::TransferWriteOp> {
|
||||
OpOperand &getSourceOperand(Operation *op) const {
|
||||
return cast<vector::TransferWriteOp>(op).getVectorMutable();
|
||||
}
|
||||
|
||||
OpOperand &getDestinationOperand(Operation *op) const {
|
||||
return cast<vector::TransferWriteOp>(op).getSourceMutable();
|
||||
}
|
||||
|
||||
Value buildSubsetExtraction(Operation *op, OpBuilder &builder,
|
||||
Location loc) const {
|
||||
// TODO: Implement when needed.
|
||||
return Value();
|
||||
}
|
||||
|
||||
SmallVector<Value>
|
||||
getValuesNeededToBuildSubsetExtraction(Operation *op) const {
|
||||
// TODO: Implement when needed.
|
||||
return {};
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace
|
||||
|
||||
void mlir::vector::registerSubsetOpInterfaceExternalModels(
|
||||
DialectRegistry ®istry) {
|
||||
registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
|
||||
TransferReadOp::attachInterface<XferOpSubsetOpInterface<TransferReadOp>>(
|
||||
*ctx);
|
||||
TransferReadOp::attachInterface<TransferReadOpSubsetExtractionOpInterface>(
|
||||
*ctx);
|
||||
TransferWriteOp::attachInterface<XferOpSubsetOpInterface<TransferWriteOp>>(
|
||||
*ctx);
|
||||
TransferWriteOp::attachInterface<TransferWriteOpSubsetInsertionOpInterface>(
|
||||
*ctx);
|
||||
});
|
||||
}
|
||||
@@ -224,477 +224,6 @@ module attributes {transform.with_named_sequence} {
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor
|
||||
func.func @hoist_vector_transfer_pairs_tensor(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>, %tensor2: tensor<?x?xf32>,
|
||||
%tensor3: tensor<?x?xf32>, %tensor4: tensor<?x?xf32>, %tensor5: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>) {
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<2xf32>, vector<1xf32>) {
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: "some_crippling_use"(%{{.*}}) : (tensor<?x?xf32>) -> ()
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<5xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (tensor<?x?xf32>) -> vector<3xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<5xf32>) -> vector<5xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<5xf32>, tensor<?x?xf32>
|
||||
// CHECK: "some_crippling_use"(%{{.*}}) : (tensor<?x?xf32>) -> ()
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<2xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<1xf32>, tensor<?x?xf32>
|
||||
%0:6 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2,
|
||||
%arg3 = %tensor3, %arg4 = %tensor4, %arg5 = %tensor5)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%1:6 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2,
|
||||
%arg9 = %arg3, %arg10 = %arg4, %arg11 = %arg5)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%r0 = vector.transfer_read %arg7[%c0, %c0], %cst: tensor<?x?xf32>, vector<1xf32>
|
||||
%r1 = vector.transfer_read %arg6[%i, %i], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r3 = vector.transfer_read %arg9[%c0, %c0], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
"some_crippling_use"(%arg10) : (tensor<?x?xf32>) -> ()
|
||||
%r4 = vector.transfer_read %arg10[%c0, %c0], %cst: tensor<?x?xf32>, vector<5xf32>
|
||||
%r5 = vector.transfer_read %arg11[%c0, %c0], %cst: tensor<?x?xf32>, vector<6xf32>
|
||||
"some_crippling_use"(%arg11) : (tensor<?x?xf32>) -> ()
|
||||
%u0 = "some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
|
||||
%u1 = "some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u2 = "some_use"(%arg8) : (tensor<?x?xf32>) -> vector<3xf32>
|
||||
%u3 = "some_use"(%r3) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u4 = "some_use"(%r4) : (vector<5xf32>) -> vector<5xf32>
|
||||
%u5 = "some_use"(%r5) : (vector<6xf32>) -> vector<6xf32>
|
||||
%w1 = vector.transfer_write %u0, %arg7[%c0, %c0] : vector<1xf32>, tensor<?x?xf32>
|
||||
%w0 = vector.transfer_write %u1, %arg6[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w2 = vector.transfer_write %u2, %arg8[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w3 = vector.transfer_write %u3, %arg9[%c0, %c0] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w4 = vector.transfer_write %u4, %arg10[%c0, %c0] : vector<5xf32>, tensor<?x?xf32>
|
||||
%w5 = vector.transfer_write %u5, %arg11[%c0, %c0] : vector<6xf32>, tensor<?x?xf32>
|
||||
"some_crippling_use"(%w3) : (tensor<?x?xf32>) -> ()
|
||||
scf.yield %w0, %w1, %w2, %w3, %w4, %w5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
scf.yield %1#0, %1#1, %1#2, %1#3, %1#4, %1#5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
module attributes {transform.with_named_sequence} {
|
||||
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
||||
%0 = transform.structured.match ops{["func.func"]} in %arg1
|
||||
: (!transform.any_op) -> !transform.any_op
|
||||
transform.structured.hoist_redundant_tensor_subsets %0
|
||||
: (!transform.any_op) -> ()
|
||||
transform.yield
|
||||
}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_disjoint_tensor(
|
||||
// CHECK-SAME: %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR3:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
func.func @hoist_vector_transfer_pairs_disjoint_tensor(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>,
|
||||
%tensor2: tensor<?x?xf32>, %tensor3: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index,
|
||||
%random_index : index) ->
|
||||
(tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%c1 = arith.constant 1 : index
|
||||
%c3 = arith.constant 3 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: vector.transfer_read %[[TENSOR2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR3]]{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR3]]{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: %[[R:.*]]:6 = scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>) {
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>) {
|
||||
// CHECK: vector.transfer_read %[[TENSOR1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<3xf32>) -> vector<3xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<3xf32>) -> vector<3xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}}, %{{.*}}{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}}, %{{.*}}{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>
|
||||
// CHECK: }
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>
|
||||
// CHECK: }
|
||||
// CHECK: %[[TENSOR4:.*]] = vector.transfer_write %[[R]]#5, %[[TENSOR3]]{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %[[R]]#4, %[[TENSOR4]]{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: %[[TENSOR5:.*]] = vector.transfer_write %[[R]]#3, %[[TENSOR2]]{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %[[R]]#2, %[[TENSOR5]]{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
%0:4 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2,
|
||||
%arg3 = %tensor3)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%1:4 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg4 = %arg0, %arg5 = %arg1, %arg6 = %arg2,
|
||||
%arg7 = %arg3)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%r00 = vector.transfer_read %arg5[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r01 = vector.transfer_read %arg5[%c0, %c1], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r20 = vector.transfer_read %arg6[%c0, %c0], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
%r21 = vector.transfer_read %arg6[%c0, %c3], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
%r30 = vector.transfer_read %arg7[%c0, %random_index], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
%r31 = vector.transfer_read %arg7[%c1, %random_index], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
%r10 = vector.transfer_read %arg4[%i, %i], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r11 = vector.transfer_read %arg4[%random_index, %random_index], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u00 = "some_use"(%r00) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u01 = "some_use"(%r01) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u20 = "some_use"(%r20) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u21 = "some_use"(%r21) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u30 = "some_use"(%r30) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u31 = "some_use"(%r31) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u10 = "some_use"(%r10) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u11 = "some_use"(%r11) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w10 = vector.transfer_write %u00, %arg5[%c0, %c0] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w11 = vector.transfer_write %u01, %w10[%c0, %c1] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w20 = vector.transfer_write %u20, %arg6[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w21 = vector.transfer_write %u21, %w20[%c0, %c3] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w30 = vector.transfer_write %u30, %arg7[%c0, %random_index] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w31 = vector.transfer_write %u31, %w30[%c1, %random_index] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w00 = vector.transfer_write %u10, %arg4[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w01 = vector.transfer_write %u11, %w00[%random_index, %random_index] : vector<2xf32>, tensor<?x?xf32>
|
||||
scf.yield %w01, %w11, %w21, %w31 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
scf.yield %1#0, %1#1, %1#2, %1#3 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
return %0#0, %0#1, %0#2, %0#3 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
module attributes {transform.with_named_sequence} {
|
||||
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
||||
%0 = transform.structured.match ops{["func.func"]} in %arg1
|
||||
: (!transform.any_op) -> !transform.any_op
|
||||
transform.structured.hoist_redundant_tensor_subsets %0
|
||||
: (!transform.any_op) -> ()
|
||||
transform.yield
|
||||
}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor_and_slices
|
||||
// CHECK-SAME: %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR3:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR4:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR5:[a-zA-Z0-9]*]]: tensor<?x?xf32>
|
||||
func.func @hoist_vector_transfer_pairs_tensor_and_slices(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>, %tensor2: tensor<?x?xf32>,
|
||||
%tensor3: tensor<?x?xf32>, %tensor4: tensor<?x?xf32>, %tensor5: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>//, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: scf.for %[[I:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR0_ARG:[0-9a-zA-Z]+]] = %[[TENSOR0]],
|
||||
// CHECK-SAME: %[[TENSOR1_ARG:[0-9a-zA-Z]+]] = %[[TENSOR1]],
|
||||
// CHECK-SAME: %[[TENSOR2_ARG:[0-9a-zA-Z]+]] = %[[TENSOR2]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
%0:3 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[ST0:.*]] = tensor.extract_slice %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V0:.*]] = vector.transfer_read %[[ST0]]{{.*}} : tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[R:.*]]:3 = scf.for %[[J:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR1_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR1_ARG]]
|
||||
// CHECK-SAME: %[[TENSOR2_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR2_ARG]]
|
||||
// CHECK-SAME: %[[V0_ARG_L2:[0-9a-zA-Z]+]] = %[[V0]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>
|
||||
%1:3 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
// Hoists.
|
||||
%st0 = tensor.extract_slice %arg6[%i, %i][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r0 = vector.transfer_read %st0[%c0, %c0], %cst: tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[ST1:.*]] = tensor.extract_slice %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V1:.*]] = vector.transfer_read %[[ST1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// Does not hoist (slice depends on %j)
|
||||
%st1 = tensor.extract_slice %arg7[%j, %c0][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r1 = vector.transfer_read %st1[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
|
||||
// CHECK: %[[ST2:.*]] = tensor.extract_slice %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V2:.*]] = vector.transfer_read %[[ST2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// Does not hoist, 2 slice %arg8.
|
||||
%st2 = tensor.extract_slice %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r2 = vector.transfer_read %st2[%c0, %c0], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
|
||||
// CHECK: %[[U0:.*]] = "some_use"(%[[V0_ARG_L2]]) : (vector<1xf32>) -> vector<1xf32>
|
||||
// CHECK: %[[U1:.*]] = "some_use"(%[[V1]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: %[[U2:.*]] = "some_use"(%[[V2]]) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u0 = "some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
|
||||
%u1 = "some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u2 = "some_use"(%r2) : (vector<3xf32>) -> vector<3xf32>
|
||||
|
||||
// Hoists
|
||||
%w0 = vector.transfer_write %u0, %st0[%c0, %c0] : vector<1xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI1:.*]] = vector.transfer_write %[[U1]], %{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// Does not hoist (associated slice depends on %j).
|
||||
%w1 = vector.transfer_write %u1, %st1[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI2:.*]] = vector.transfer_write %[[U2]], %{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%w2 = vector.transfer_write %u2, %st2[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
|
||||
// Hoists.
|
||||
%sti0 = tensor.insert_slice %w0 into %arg6[%i, %i][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI1]] into %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Does not hoist (depends on %j).
|
||||
%sti1 = tensor.insert_slice %w1 into %arg7[%j, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI2]] into %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%sti2 = tensor.insert_slice %w2 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Extract with a different stride to make sure we cannot fold this extract with the above insert.
|
||||
%st22 = tensor.extract_slice %sti2[%i, %c0][%step, %step][2, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%sti22 = tensor.insert_slice %st22 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
scf.yield %sti0, %sti1, %sti22:
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[STI0:.*]] = vector.transfer_write %[[R]]#2, %[[ST0]]{{.*}} : vector<1xf32>, tensor<?x?xf32>
|
||||
// CHECK: tensor.insert_slice %[[STI0]] into %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}} : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
scf.yield %1#0, %1#1, %1#2 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK: }
|
||||
}
|
||||
return %0#0, %0#1, %0#2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
module attributes {transform.with_named_sequence} {
|
||||
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
||||
%0 = transform.structured.match ops{["func.func"]} in %arg1
|
||||
: (!transform.any_op) -> !transform.any_op
|
||||
transform.structured.hoist_redundant_tensor_subsets %0
|
||||
: (!transform.any_op) -> ()
|
||||
transform.yield
|
||||
}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_write_pairs_disjoint_tensor(
|
||||
// CHECK-SAME: %[[T:.*]]: tensor<?x?xf32>,
|
||||
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
|
||||
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
|
||||
// CHECK-DAG: %[[R0:.*]] = vector.transfer_read %[[T]][%[[C0]], %[[C0]]], %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK-DAG: %[[R1:.*]] = vector.transfer_read %[[T]][%[[C0]], %[[C3]]], %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: %[[F:.*]]:2 = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[R3:.*]] = %[[R1:.*]], %[[R2:.*]] = %[[R0]]) -> (vector<2xf32>, vector<2xf32>) {
|
||||
// CHECK: %[[R4:.*]] = "some_use"(%[[R2]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: %[[R5:.*]] = "some_use"(%[[R3]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: scf.yield %[[R5]], %[[R4]] : vector<2xf32>, vector<2xf32>
|
||||
// CHECK: }
|
||||
// CHECK: %[[W0:.*]] = vector.transfer_write %[[F]]#1, %[[T]][%[[C0]], %[[C0]]] : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: %[[W1:.*]] = vector.transfer_write %[[F]]#0, %[[W0]][%[[C0]], %[[C3]]] : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: return %[[W1]] : tensor<?x?xf32>
|
||||
func.func @hoist_vector_transfer_write_pairs_disjoint_tensor(
|
||||
%tensor: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%c1 = arith.constant 1 : index
|
||||
%c3 = arith.constant 3 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
%1 = scf.for %j = %lb to %ub step %step iter_args(%arg5 = %tensor)
|
||||
-> (tensor<?x?xf32>) {
|
||||
%r00 = vector.transfer_read %arg5[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u00 = "some_use"(%r00) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w10 = vector.transfer_write %u00, %arg5[%c0, %c0] : vector<2xf32>, tensor<?x?xf32>
|
||||
|
||||
// Hoist by properly bypassing the disjoint write %w10.
|
||||
%r01 = vector.transfer_read %w10[%c0, %c3], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u01 = "some_use"(%r01) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w11 = vector.transfer_write %u01, %w10[%c0, %c3] : vector<2xf32>, tensor<?x?xf32>
|
||||
scf.yield %w11 : tensor<?x?xf32>
|
||||
}
|
||||
return %1 : tensor<?x?xf32>
|
||||
}
|
||||
|
||||
module attributes {transform.with_named_sequence} {
|
||||
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
||||
%0 = transform.structured.match ops{["func.func"]} in %arg1
|
||||
: (!transform.any_op) -> !transform.any_op
|
||||
transform.structured.hoist_redundant_tensor_subsets %0
|
||||
: (!transform.any_op) -> ()
|
||||
transform.yield
|
||||
}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor_and_slices_static_large_tensor
|
||||
// CHECK-SAME: %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<100x100xf32>,
|
||||
// CHECK-SAME: %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<200x200xf32>,
|
||||
// CHECK-SAME: %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<300x300xf32>
|
||||
func.func @hoist_vector_transfer_pairs_tensor_and_slices_static_large_tensor(
|
||||
%tensor0: tensor<100x100xf32>, %tensor1: tensor<200x200xf32>, %tensor2: tensor<300x300xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(
|
||||
tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: scf.for %[[I:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR0_ARG:[0-9a-zA-Z]+]] = %[[TENSOR0]],
|
||||
// CHECK-SAME: %[[TENSOR1_ARG:[0-9a-zA-Z]+]] = %[[TENSOR1]],
|
||||
// CHECK-SAME: %[[TENSOR2_ARG:[0-9a-zA-Z]+]] = %[[TENSOR2]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
%0:3 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2)
|
||||
-> (tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>) {
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[ST0:.*]] = tensor.extract_slice %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}}: tensor<100x100xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V0:.*]] = vector.transfer_read %[[ST0]]{{.*}} : tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[R:.*]]:3 = scf.for %[[J:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR1_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR1_ARG]]
|
||||
// CHECK-SAME: %[[TENSOR2_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR2_ARG]]
|
||||
// CHECK-SAME: %[[V0_ARG_L2:[0-9a-zA-Z]+]] = %[[V0]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<200x200xf32>, tensor<300x300xf32>, vector<1xf32>
|
||||
%1:3 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2)
|
||||
-> (tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>) {
|
||||
// Hoists.
|
||||
%st0 = tensor.extract_slice %arg6[%i, %i][%step, %step][1, 1] : tensor<100x100xf32> to tensor<?x?xf32>
|
||||
%r0 = vector.transfer_read %st0[%c0, %c0], %cst: tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[ST1:.*]] = tensor.extract_slice %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<200x200xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V1:.*]] = vector.transfer_read %[[ST1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// Does not hoist (slice depends on %j)
|
||||
%st1 = tensor.extract_slice %arg7[%j, %c0][%step, %step][1, 1] : tensor<200x200xf32> to tensor<?x?xf32>
|
||||
%r1 = vector.transfer_read %st1[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
|
||||
// CHECK: %[[ST2:.*]] = tensor.extract_slice %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<300x300xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V2:.*]] = vector.transfer_read %[[ST2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// Does not hoist, 2 slice %arg8.
|
||||
%st2 = tensor.extract_slice %arg8[%i, %c0][%step, %step][1, 1] : tensor<300x300xf32> to tensor<?x?xf32>
|
||||
%r2 = vector.transfer_read %st2[%c0, %c0], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
|
||||
// CHECK: %[[U0:.*]] = "some_use"(%[[V0_ARG_L2]]) : (vector<1xf32>) -> vector<1xf32>
|
||||
// CHECK: %[[U1:.*]] = "some_use"(%[[V1]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: %[[U2:.*]] = "some_use"(%[[V2]]) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u0 = "some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
|
||||
%u1 = "some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u2 = "some_use"(%r2) : (vector<3xf32>) -> vector<3xf32>
|
||||
|
||||
// Hoists
|
||||
%w0 = vector.transfer_write %u0, %st0[%c0, %c0] : vector<1xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI1:.*]] = vector.transfer_write %[[U1]], %{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// Does not hoist (associated slice depends on %j).
|
||||
%w1 = vector.transfer_write %u1, %st1[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI2:.*]] = vector.transfer_write %[[U2]], %{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%w2 = vector.transfer_write %u2, %st2[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
|
||||
// Hoists.
|
||||
%sti0 = tensor.insert_slice %w0 into %arg6[%i, %i][%step, %step][1, 1] : tensor<?x?xf32> into tensor<100x100xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI1]] into %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<?x?xf32> into tensor<200x200xf32>
|
||||
// Does not hoist (depends on %j).
|
||||
%sti1 = tensor.insert_slice %w1 into %arg7[%j, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<200x200xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI2]] into %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<?x?xf32> into tensor<300x300xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%sti2 = tensor.insert_slice %w2 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<300x300xf32>
|
||||
// Extract with a different stride to make sure we cannot fold this extract with the above insert.
|
||||
%st22 = tensor.extract_slice %sti2[%i, %c0][%step, %step][2, 1] : tensor<300x300xf32> to tensor<?x?xf32>
|
||||
%sti22 = tensor.insert_slice %st22 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<300x300xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<200x200xf32>, tensor<300x300xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
scf.yield %sti0, %sti1, %sti22:
|
||||
tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
}
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[STI0:.*]] = vector.transfer_write %[[R]]#2, %[[ST0]]{{.*}} : vector<1xf32>, tensor<?x?xf32>
|
||||
// CHECK: tensor.insert_slice %[[STI0]] into %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}} : tensor<?x?xf32> into tensor<100x100xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
scf.yield %1#0, %1#1, %1#2 :
|
||||
tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
|
||||
// CHECK: }
|
||||
}
|
||||
return %0#0, %0#1, %0#2 : tensor<100x100xf32>, tensor<200x200xf32>, tensor<300x300xf32>
|
||||
}
|
||||
|
||||
module attributes {transform.with_named_sequence} {
|
||||
transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
|
||||
%0 = transform.structured.match ops{["func.func"]} in %arg1
|
||||
: (!transform.any_op) -> !transform.any_op
|
||||
transform.structured.hoist_redundant_tensor_subsets %0
|
||||
: (!transform.any_op) -> ()
|
||||
transform.yield
|
||||
}
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func.func @hoist_vector_transfer_read(
|
||||
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
|
||||
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
|
||||
|
||||
@@ -277,3 +277,321 @@ func.func @nested_hoisting(%arg: tensor<?xf32>) -> tensor<?xf32> {
|
||||
// CHECK: return %[[insert2]]
|
||||
return %0 : tensor<?xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor
|
||||
func.func @hoist_vector_transfer_pairs_tensor(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>, %tensor2: tensor<?x?xf32>,
|
||||
%tensor3: tensor<?x?xf32>, %tensor4: tensor<?x?xf32>, %tensor5: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>) {
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<2xf32>, vector<1xf32>) {
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: "test.some_crippling_use"(%{{.*}}) : (tensor<?x?xf32>) -> ()
|
||||
// CHECK: vector.transfer_read %{{.*}} : tensor<?x?xf32>, vector<5xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (tensor<?x?xf32>) -> vector<3xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<5xf32>) -> vector<5xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<5xf32>, tensor<?x?xf32>
|
||||
// CHECK: "test.some_crippling_use"(%{{.*}}) : (tensor<?x?xf32>) -> ()
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<2xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
// CHECK: vector.transfer_write %{{.*}} : vector<1xf32>, tensor<?x?xf32>
|
||||
%0:6 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2,
|
||||
%arg3 = %tensor3, %arg4 = %tensor4, %arg5 = %tensor5)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%1:6 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2,
|
||||
%arg9 = %arg3, %arg10 = %arg4, %arg11 = %arg5)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%r0 = vector.transfer_read %arg7[%c0, %c0], %cst: tensor<?x?xf32>, vector<1xf32>
|
||||
%r1 = vector.transfer_read %arg6[%i, %i], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r3 = vector.transfer_read %arg9[%c0, %c0], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
"test.some_crippling_use"(%arg10) : (tensor<?x?xf32>) -> ()
|
||||
%r4 = vector.transfer_read %arg10[%c0, %c0], %cst: tensor<?x?xf32>, vector<5xf32>
|
||||
%r5 = vector.transfer_read %arg11[%c0, %c0], %cst: tensor<?x?xf32>, vector<6xf32>
|
||||
"test.some_crippling_use"(%arg11) : (tensor<?x?xf32>) -> ()
|
||||
%u0 = "test.some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
|
||||
%u1 = "test.some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u2 = "test.some_use"(%arg8) : (tensor<?x?xf32>) -> vector<3xf32>
|
||||
%u3 = "test.some_use"(%r3) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u4 = "test.some_use"(%r4) : (vector<5xf32>) -> vector<5xf32>
|
||||
%u5 = "test.some_use"(%r5) : (vector<6xf32>) -> vector<6xf32>
|
||||
%w1 = vector.transfer_write %u0, %arg7[%c0, %c0] : vector<1xf32>, tensor<?x?xf32>
|
||||
%w0 = vector.transfer_write %u1, %arg6[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w2 = vector.transfer_write %u2, %arg8[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w3 = vector.transfer_write %u3, %arg9[%c0, %c0] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w4 = vector.transfer_write %u4, %arg10[%c0, %c0] : vector<5xf32>, tensor<?x?xf32>
|
||||
%w5 = vector.transfer_write %u5, %arg11[%c0, %c0] : vector<6xf32>, tensor<?x?xf32>
|
||||
"test.some_crippling_use"(%w3) : (tensor<?x?xf32>) -> ()
|
||||
scf.yield %w0, %w1, %w2, %w3, %w4, %w5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
scf.yield %1#0, %1#1, %1#2, %1#3, %1#4, %1#5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
return %0#0, %0#1, %0#2, %0#3, %0#4, %0#5 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>,
|
||||
tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_disjoint_tensor(
|
||||
// CHECK-SAME: %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR3:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
func.func @hoist_vector_transfer_pairs_disjoint_tensor(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>,
|
||||
%tensor2: tensor<?x?xf32>, %tensor3: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index,
|
||||
%random_index : index) ->
|
||||
(tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%c1 = arith.constant 1 : index
|
||||
%c3 = arith.constant 3 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: vector.transfer_read %[[TENSOR2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR3]]{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR3]]{{.*}} : tensor<?x?xf32>, vector<4xf32>
|
||||
// CHECK: %[[R:.*]]:8 = scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>) {
|
||||
// CHECK: scf.for {{.*}} iter_args({{.*}}) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>) {
|
||||
// CHECK: vector.transfer_read %[[TENSOR1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: vector.transfer_read %[[TENSOR1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<3xf32>) -> vector<3xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<3xf32>) -> vector<3xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<4xf32>) -> vector<4xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: "test.some_use"(%{{.*}}) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}}, %{{.*}}{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %{{.*}}, %{{.*}}{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>
|
||||
// CHECK: }
|
||||
// CHECK: scf.yield {{.*}} :
|
||||
// CHECK-SAME: tensor<?x?xf32>, tensor<?x?xf32>, vector<3xf32>, vector<3xf32>, vector<4xf32>, vector<4xf32>
|
||||
// CHECK: }
|
||||
// CHECK: %[[TENSOR4:.*]] = vector.transfer_write %[[R]]#7, %[[R]]#3{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %[[R]]#6, %[[TENSOR4]]{{.*}} : vector<4xf32>, tensor<?x?xf32>
|
||||
// CHECK: %[[TENSOR5:.*]] = vector.transfer_write %[[R]]#5, %[[R]]#2{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// CHECK: vector.transfer_write %[[R]]#4, %[[TENSOR5]]{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
%0:4 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2,
|
||||
%arg3 = %tensor3)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%1:4 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg4 = %arg0, %arg5 = %arg1, %arg6 = %arg2,
|
||||
%arg7 = %arg3)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
%r00 = vector.transfer_read %arg5[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r01 = vector.transfer_read %arg5[%c0, %c1], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r20 = vector.transfer_read %arg6[%c0, %c0], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
%r21 = vector.transfer_read %arg6[%c0, %c3], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
%r30 = vector.transfer_read %arg7[%c0, %random_index], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
%r31 = vector.transfer_read %arg7[%c1, %random_index], %cst: tensor<?x?xf32>, vector<4xf32>
|
||||
%r10 = vector.transfer_read %arg4[%i, %i], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%r11 = vector.transfer_read %arg4[%random_index, %random_index], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u00 = "test.some_use"(%r00) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u01 = "test.some_use"(%r01) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u20 = "test.some_use"(%r20) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u21 = "test.some_use"(%r21) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u30 = "test.some_use"(%r30) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u31 = "test.some_use"(%r31) : (vector<4xf32>) -> vector<4xf32>
|
||||
%u10 = "test.some_use"(%r10) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u11 = "test.some_use"(%r11) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w10 = vector.transfer_write %u00, %arg5[%c0, %c0] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w11 = vector.transfer_write %u01, %w10[%c0, %c1] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w20 = vector.transfer_write %u20, %arg6[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w21 = vector.transfer_write %u21, %w20[%c0, %c3] : vector<3xf32>, tensor<?x?xf32>
|
||||
%w30 = vector.transfer_write %u30, %arg7[%c0, %random_index] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w31 = vector.transfer_write %u31, %w30[%c1, %random_index] : vector<4xf32>, tensor<?x?xf32>
|
||||
%w00 = vector.transfer_write %u10, %arg4[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
%w01 = vector.transfer_write %u11, %w00[%random_index, %random_index] : vector<2xf32>, tensor<?x?xf32>
|
||||
scf.yield %w01, %w11, %w21, %w31 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
scf.yield %1#0, %1#1, %1#2, %1#3 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
return %0#0, %0#1, %0#2, %0#3 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_pairs_tensor_and_slices
|
||||
// CHECK-SAME: %[[TENSOR0:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR1:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR2:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR3:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR4:[a-zA-Z0-9]*]]: tensor<?x?xf32>,
|
||||
// CHECK-SAME: %[[TENSOR5:[a-zA-Z0-9]*]]: tensor<?x?xf32>
|
||||
func.func @hoist_vector_transfer_pairs_tensor_and_slices(
|
||||
%tensor0: tensor<?x?xf32>, %tensor1: tensor<?x?xf32>, %tensor2: tensor<?x?xf32>,
|
||||
%tensor3: tensor<?x?xf32>, %tensor4: tensor<?x?xf32>, %tensor5: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>//, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
|
||||
// CHECK: scf.for %[[I:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR0_ARG:[0-9a-zA-Z]+]] = %[[TENSOR0]],
|
||||
// CHECK-SAME: %[[TENSOR1_ARG:[0-9a-zA-Z]+]] = %[[TENSOR1]],
|
||||
// CHECK-SAME: %[[TENSOR2_ARG:[0-9a-zA-Z]+]] = %[[TENSOR2]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
%0:3 = scf.for %i = %lb to %ub step %step
|
||||
iter_args(%arg0 = %tensor0, %arg1 = %tensor1, %arg2 = %tensor2)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[ST0:.*]] = tensor.extract_slice %[[TENSOR0_ARG]][%[[I]], %[[I]]]{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V0:.*]] = vector.transfer_read %[[ST0]]{{.*}} : tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[R:.*]]:5 = scf.for %[[J:.*]] = {{.*}} iter_args(
|
||||
// CHECK-SAME: %[[TENSOR0_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR0_ARG]]
|
||||
// CHECK-SAME: %[[TENSOR1_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR1_ARG]]
|
||||
// CHECK-SAME: %[[TENSOR2_ARG_L2:[0-9a-zA-Z]+]] = %[[TENSOR2_ARG]]
|
||||
// CHECK-SAME: %[[ST0_ARG_L2:[0-9a-zA-Z]+]] = %[[ST0]]
|
||||
// CHECK-SAME: %[[V0_ARG_L2:[0-9a-zA-Z]+]] = %[[V0]]
|
||||
// CHECK-SAME: ) ->
|
||||
// CHECK-SAME: (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>)
|
||||
%1:3 = scf.for %j = %lb to %ub step %step
|
||||
iter_args(%arg6 = %arg0, %arg7 = %arg1, %arg8 = %arg2)
|
||||
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
|
||||
// Hoists.
|
||||
%st0 = tensor.extract_slice %arg6[%i, %i][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r0 = vector.transfer_read %st0[%c0, %c0], %cst: tensor<?x?xf32>, vector<1xf32>
|
||||
|
||||
// CHECK: %[[ST1:.*]] = tensor.extract_slice %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V1:.*]] = vector.transfer_read %[[ST1]]{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// Does not hoist (slice depends on %j)
|
||||
%st1 = tensor.extract_slice %arg7[%j, %c0][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r1 = vector.transfer_read %st1[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
|
||||
// CHECK: %[[ST2:.*]] = tensor.extract_slice %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<?x?xf32> to tensor<?x?xf32>
|
||||
// CHECK: %[[V2:.*]] = vector.transfer_read %[[ST2]]{{.*}} : tensor<?x?xf32>, vector<3xf32>
|
||||
// Does not hoist, 2 slice %arg8.
|
||||
%st2 = tensor.extract_slice %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%r2 = vector.transfer_read %st2[%c0, %c0], %cst: tensor<?x?xf32>, vector<3xf32>
|
||||
|
||||
// CHECK: %[[U0:.*]] = "test.some_use"(%[[V0_ARG_L2]]) : (vector<1xf32>) -> vector<1xf32>
|
||||
// CHECK: %[[U1:.*]] = "test.some_use"(%[[V1]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: %[[U2:.*]] = "test.some_use"(%[[V2]]) : (vector<3xf32>) -> vector<3xf32>
|
||||
%u0 = "test.some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
|
||||
%u1 = "test.some_use"(%r1) : (vector<2xf32>) -> vector<2xf32>
|
||||
%u2 = "test.some_use"(%r2) : (vector<3xf32>) -> vector<3xf32>
|
||||
|
||||
// Hoists
|
||||
%w0 = vector.transfer_write %u0, %st0[%c0, %c0] : vector<1xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI1:.*]] = vector.transfer_write %[[U1]], %{{.*}} : vector<2xf32>, tensor<?x?xf32>
|
||||
// Does not hoist (associated slice depends on %j).
|
||||
%w1 = vector.transfer_write %u1, %st1[%i, %i] : vector<2xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: %[[STI2:.*]] = vector.transfer_write %[[U2]], %{{.*}} : vector<3xf32>, tensor<?x?xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%w2 = vector.transfer_write %u2, %st2[%c0, %c0] : vector<3xf32>, tensor<?x?xf32>
|
||||
|
||||
// Hoists.
|
||||
%sti0 = tensor.insert_slice %w0 into %arg6[%i, %i][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI1]] into %[[TENSOR1_ARG_L2]][%[[J]],{{.*}}: tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Does not hoist (depends on %j).
|
||||
%sti1 = tensor.insert_slice %w1 into %arg7[%j, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK-DAG: tensor.insert_slice %[[STI2]] into %[[TENSOR2_ARG_L2]][%[[I]],{{.*}}: tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Does not hoist, 2 slice / insert_slice for %arg8.
|
||||
%sti2 = tensor.insert_slice %w2 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
// Extract with a different stride to make sure we cannot fold this extract with the above insert.
|
||||
%st22 = tensor.extract_slice %sti2[%i, %c0][%step, %step][2, 1] : tensor<?x?xf32> to tensor<?x?xf32>
|
||||
%sti22 = tensor.insert_slice %st22 into %arg8[%i, %c0][%step, %step][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>, vector<1xf32>
|
||||
// CHECK: }
|
||||
scf.yield %sti0, %sti1, %sti22:
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
// Hoisted
|
||||
// CHECK: %[[STI0:.*]] = vector.transfer_write %[[R]]#4, %[[R]]#3{{.*}} : vector<1xf32>, tensor<?x?xf32>
|
||||
// CHECK: tensor.insert_slice %[[STI0]] into %[[R]]#0[%[[I]], %[[I]]]{{.*}} : tensor<?x?xf32> into tensor<?x?xf32>
|
||||
|
||||
// CHECK: scf.yield {{.*}} : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
scf.yield %1#0, %1#1, %1#2 :
|
||||
tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
|
||||
// CHECK: }
|
||||
}
|
||||
return %0#0, %0#1, %0#2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
|
||||
}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: func @hoist_vector_transfer_write_pairs_disjoint_tensor(
|
||||
// CHECK-SAME: %[[T:.*]]: tensor<?x?xf32>,
|
||||
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
|
||||
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
|
||||
// CHECK-DAG: %[[R0:.*]] = vector.transfer_read %[[T]][%[[C0]], %[[C0]]], %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK-DAG: %[[R1:.*]] = vector.transfer_read %[[T]][%[[C0]], %[[C3]]], %{{.*}} : tensor<?x?xf32>, vector<2xf32>
|
||||
// CHECK: %[[F:.*]]:3 = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[TL:.*]] = %[[T]], %[[R2:.*]] = %[[R0]], %[[R3:.*]] = %[[R1]]) -> (tensor<?x?xf32>, vector<2xf32>, vector<2xf32>) {
|
||||
// CHECK: %[[R4:.*]] = "test.some_use"(%[[R2]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: %[[R5:.*]] = "test.some_use"(%[[R3]]) : (vector<2xf32>) -> vector<2xf32>
|
||||
// CHECK: scf.yield %[[TL]], %[[R4]], %[[R5]] : tensor<?x?xf32>, vector<2xf32>, vector<2xf32>
|
||||
// CHECK: }
|
||||
// CHECK: %[[W0:.*]] = vector.transfer_write %[[F]]#2, %[[F]]#0[%[[C0]], %[[C3]]] : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: %[[W1:.*]] = vector.transfer_write %[[F]]#1, %[[W0]][%[[C0]], %[[C0]]] : vector<2xf32>, tensor<?x?xf32>
|
||||
// CHECK: return %[[W1]] : tensor<?x?xf32>
|
||||
func.func @hoist_vector_transfer_write_pairs_disjoint_tensor(
|
||||
%tensor: tensor<?x?xf32>,
|
||||
%val: index, %lb : index, %ub : index, %step: index) ->
|
||||
(tensor<?x?xf32>) {
|
||||
%c0 = arith.constant 0 : index
|
||||
%c1 = arith.constant 1 : index
|
||||
%c3 = arith.constant 3 : index
|
||||
%cst = arith.constant 0.0 : f32
|
||||
%1 = scf.for %j = %lb to %ub step %step iter_args(%arg5 = %tensor)
|
||||
-> (tensor<?x?xf32>) {
|
||||
%r00 = vector.transfer_read %arg5[%c0, %c0], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u00 = "test.some_use"(%r00) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w10 = vector.transfer_write %u00, %arg5[%c0, %c0] : vector<2xf32>, tensor<?x?xf32>
|
||||
|
||||
// Hoist by properly bypassing the disjoint write %w10.
|
||||
%r01 = vector.transfer_read %w10[%c0, %c3], %cst: tensor<?x?xf32>, vector<2xf32>
|
||||
%u01 = "test.some_use"(%r01) : (vector<2xf32>) -> vector<2xf32>
|
||||
%w11 = vector.transfer_write %u01, %w10[%c0, %c3] : vector<2xf32>, tensor<?x?xf32>
|
||||
scf.yield %w11 : tensor<?x?xf32>
|
||||
}
|
||||
return %1 : tensor<?x?xf32>
|
||||
}
|
||||
|
||||
@@ -4598,6 +4598,7 @@ cc_library(
|
||||
":Pass",
|
||||
":SCFDialect",
|
||||
":SideEffectInterfaces",
|
||||
":SubsetOpInterface",
|
||||
":Support",
|
||||
":TensorDialect",
|
||||
":Transforms",
|
||||
|
||||
Reference in New Issue
Block a user