This is an archive of the discontinued LLVM Phabricator instance.

Differential D135129

[mlir][tensor][NFC] Rename linalg.init_tensor to tensor.empty
ClosedPublic

Authored by springerm on Oct 3 2022, 10:27 PM.

Details

Reviewers
pifon2a
nicolasvasilache
ftynse
herhut
aartbik
mravishankar
dcaballe
Commits
rG81ca5aa45240: [mlir][tensor][NFC] Rename linalg.init_tensor to tensor.empty
Summary

tensor.empty/linalg.init_tensor produces an uninititalized tensor that can be used as a destination operand for destination-style ops (ops that implement DestinationStyleOpInterface).

This change makes it possible to implement TilingInterface for non-destination-style ops without depending on the Linalg dialect.

RFC: https://discourse.llvm.org/t/rfc-add-tensor-from-shape-operation/65101

Diff Detail

Event Timeline

springerm created this revision.Oct 3 2022, 10:27 PM
Herald added a reviewer: aartbik. · View Herald TranscriptOct 3 2022, 10:27 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: mravishankar. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: anlunx, bzcheeseman, mravishankar and 25 others. · View Herald Transcript
springerm requested review of this revision.Oct 3 2022, 10:27 PM
Herald added a reviewer: dcaballe. · View Herald TranscriptOct 3 2022, 10:27 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: limo1996, stephenneuendorffer. · View Herald Transcript
Harbormaster completed remote builds in B190117: Diff 464888.Oct 3 2022, 10:28 PM
springerm updated this revision to Diff 464890.Oct 3 2022, 10:33 PM

add test case

Harbormaster completed remote builds in B190119: Diff 464890.Oct 3 2022, 10:34 PM
springerm added a comment.Oct 3 2022, 10:36 PM

Note: The assembly format of the op has changed so that it matches bufferization.alloc_tensor and memref.alloc. Apart from that, this change is NFC.

springerm added a child revision: D135130: [mlir][bufferization][NFC] Move EmptyTensorToAllocTensorPass.Oct 3 2022, 11:34 PM
ftynse accepted this revision.Oct 4 2022, 12:55 AM
This revision is now accepted and ready to land.Oct 4 2022, 12:55 AM
This revision was landed with ongoing or failed builds.Oct 4 2022, 1:28 AM
Closed by commit rG81ca5aa45240: [mlir][tensor][NFC] Rename linalg.init_tensor to tensor.empty (authored by springerm). · Explain Why
This revision was automatically updated to reflect the committed changes.
springerm added a commit: rG81ca5aa45240: [mlir][tensor][NFC] Rename linalg.init_tensor to tensor.empty.
pifon2a added a comment.Oct 4 2022, 2:39 AM

Nice!

springerm mentioned this in D136179: [mlir][interfaces][NFC] Move DestinationStyleOpInterface to mlir/Interfaces.Oct 18 2022, 8:22 AM
springerm mentioned this in rGcfc9ddaafc5d: [mlir][interfaces][NFC] Move DestinationStyleOpInterface to mlir/Interfaces.Oct 18 2022, 8:40 AM

Revision Contents

PathSize
mlir/
docs/
Dialects/
Linalg/
4 lines
include/
mlir/
Dialect/
Linalg/
IR/
104 lines
4 lines
8 lines
TransformOps/
6 lines
Transforms/
2 lines
8 lines
Tensor/
IR/
1 line
59 lines
Transforms/
2 lines
Utils/
2 lines
lib/
Conversion/
TosaToLinalg/
111 lines
80 lines
Dialect/
Linalg/
IR/
287 lines
Transforms/
12 lines
2 lines
2 lines
22 lines
6 lines
4 lines
13 lines
32 lines
35 lines
3 lines
15 lines
2 lines
SparseTensor/
Transforms/
2 lines
Tensor/
IR/
1 line
1 line
367 lines
4 lines
python/
mlir/
dialects/
38 lines
42 lines
test/
Conversion/
TensorToLinalg/
4 lines
TosaToLinalg/
78 lines
116 lines
Dialect/
Bufferization/
Transforms/
3 lines
Linalg/
8 lines
10 lines
16 lines
198 lines
4 lines
34 lines
12 lines
4 lines
12 lines
4 lines
4 lines
10 lines
4 lines
30 lines
30 lines
78 lines
2 lines
8 lines
4 lines
4 lines
29 lines
2 lines
62 lines
2 lines
8 lines
4 lines
12 lines
26 lines
21 lines
12 lines
8 lines
4 lines
4 lines
2 lines
4 lines
4 lines
8 lines
4 lines
2 lines
42 lines
SparseTensor/
4 lines
4 lines
8 lines
Tensor/
105 lines
14 lines
8 lines
8 lines
Integration/
Dialect/
Linalg/
CPU/
2 lines
SparseTensor/
taco/
tools/
23 lines
Interfaces/
TilingInterface/
32 lines
8 lines
lib/
Dialect/
Tensor/
4 lines
python/
dialects/
linalg/
opdsl/
5 lines
52 lines
34 lines
utils/
bazel/
llvm-project-overlay/
mlir/
1 line

Diff 464890

mlir/docs/Dialects/Linalg/OpDSL.md

Show First 20 Lines • Show All 312 Lines • ▼ Show 20 Lines
The operation sets the elements of the output tensor `O` to `value`. All The operation sets the elements of the output tensor `O` to `value`. All
operands are either scalars or rank zero tensors that are accessed using the operands are either scalars or rank zero tensors that are accessed using the
index `None`. The operation thus performs a scalar computation that trivially index `None`. The operation thus performs a scalar computation that trivially
extends to a multi-dimensional pointwise computation. As a result, we may use extends to a multi-dimensional pointwise computation. As a result, we may use
`fill` with arbitrary ranked output tensors: `fill` with arbitrary ranked output tensors:
```python ```python
tensor_2d = linalg.InitTensorOp([4, 8], f32) tensor_2d = tensor.EmptyOp([4, 8], f32)
tensor_3d = linalg.InitTensorOp([4, 8, 16], f32) tensor_3d = tensor.EmptyOp([4, 8, 16], f32)
fill(value, outs=[tensor_2d]) fill(value, outs=[tensor_2d])
fill(value, outs=[tensor_3d]) fill(value, outs=[tensor_3d])
``` ```

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

Show All 18 Lines
include "mlir/Interfaces/LoopLikeInterface.td" include "mlir/Interfaces/LoopLikeInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td" include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/ViewLikeInterface.td" include "mlir/Interfaces/ViewLikeInterface.td"
// Base class for Linalg dialect ops that do not correspond to library calls. // Base class for Linalg dialect ops that do not correspond to library calls.
class Linalg_Op<string mnemonic, list<Trait> traits = []> : class Linalg_Op<string mnemonic, list<Trait> traits = []> :
Op<Linalg_Dialect, mnemonic, traits>; Op<Linalg_Dialect, mnemonic, traits>;
def Linalg_InitTensorOp : Linalg_Op<"init_tensor",
[NoSideEffect,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
let summary = "operation to define a tensor of particular shape";
let description = [{
`linalg.init_tensor` is an operation that defines a tensor of a particular
shape. The shape could be dynamic or static. The contents of the tensor are
unspecified and the only purpose of the op result is to materialize the
specified shape in IR and make it available to other transformations.
Note: This op can be lowered to a `bufferization.alloc_tensor`, at which
point it turns into an explicit buffer allocation.
}];
let arguments =
(ins Variadic<Index>:$sizes, I64ArrayAttr:$static_sizes);
let results = (outs AnyTensor:$result);
let assemblyFormat = [{
custom<DynamicIndexList>($sizes, $static_sizes,
"ShapedType::kDynamicSize")
attr-dict `:` type($result)
}];
let extraClassDeclaration = [{
static StringRef getStaticSizesAttrStrName() {
return "static_sizes";
}
RankedTensorType getType() {
return getResult().getType().cast<RankedTensorType>(); }
// Infer the shape of the result tensor given the static shapes
// and element type of the result tensor.
static Type inferResultType(ArrayRef<int64_t> staticSizes, Type elementType,
Attribute encoding = {});
// Return true if the size of the tensor is dynamic at `idx`
bool isDynamicSize(unsigned idx) {
APInt v = *(getStaticSizes().getAsValueRange<IntegerAttr>().begin() + idx);
return ShapedType::isDynamic(v.getSExtValue());
}
// Assert that the size of the result tensor is static at `idx`
// and return the shape.
int64_t getStaticSize(unsigned idx) {
assert(!isDynamicSize(idx) && "expected static size");
APInt v = *(getStaticSizes().
template getAsValueRange<IntegerAttr>().begin() + idx);
return v.getSExtValue();
}
// Return the argument position that contains the dynamic size of
// the tensor at dimension `idx`. Asserts that the shape is
// dynamic at that `idx`.
unsigned getIndexOfDynamicSize(unsigned idx) {
assert(isDynamicSize(idx) && "expected dynamic size");
return std::count_if(
getStaticSizes().getValue().begin(),
getStaticSizes().getValue().begin() + idx,
[&](Attribute attr) {
return ShapedType::isDynamic(attr.cast<IntegerAttr>().getInt());
});
}
// Return both static and dynamic sizes as a list of `OpFoldResult`.
SmallVector<OpFoldResult> getMixedSizes();
// Return the Value of the dynamic size of the tensor at dimension
// `idx`. Asserts that the shape is dynamic at that `idx.
Value getDynamicSize(unsigned idx) {
return getOperand(getIndexOfDynamicSize(idx));
}
}];
let builders = [
OpBuilder<(ins "ValueRange":$shape,
"ArrayRef<int64_t>":$staticShape, "Type":$elementType),
[{
build($_builder, $_state,
InitTensorOp::inferResultType(staticShape, elementType),
shape, $_builder.getI64ArrayAttr(staticShape));
}]>,
OpBuilder<(ins "ValueRange":$shape, "Type":$elementType),
[{
SmallVector<int64_t, 4> staticShape(
shape.size(), ShapedType::kDynamicSize);
build($_builder, $_state, shape, staticShape, elementType);
}]>,
OpBuilder<(ins "ArrayRef<int64_t>":$staticShape, "Type":$elementType),
[{
build($_builder, $_state, ValueRange{}, staticShape, elementType);
}]>,
OpBuilder<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
];
let hasCanonicalizer = 1;
let hasCustomAssemblyFormat = 1;
let hasVerifier = 1;
}
def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>, def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
Arguments<(ins Variadic<AnyType>:$values)> { Arguments<(ins Variadic<AnyType>:$values)> {
let summary = "Linalg yield operation"; let summary = "Linalg yield operation";
let description = [{ let description = [{
`linalg.yield` is a special terminator operation for blocks inside regions `linalg.yield` is a special terminator operation for blocks inside regions
in `linalg` generic ops. It returns values to the immediately enclosing in `linalg` generic ops. It returns values to the immediately enclosing
`linalg` generic op. `linalg` generic op.
▲ Show 20 Lines • Show All 56 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/Passes.h

Show First 20 Lines • Show All 55 Lines • ▼ Show 20 Lines
createConvertLinalgToParallelLoopsPass(); createConvertLinalgToParallelLoopsPass();
/// Create a pass to convert Linalg operations to affine.for loops and /// Create a pass to convert Linalg operations to affine.for loops and
/// affine_load/affine_store accesses. /// affine_load/affine_store accesses.
/// Placeholder for now, this is NYI. /// Placeholder for now, this is NYI.
std::unique_ptr<OperationPass<func::FuncOp>> std::unique_ptr<OperationPass<func::FuncOp>>
createConvertLinalgToAffineLoopsPass(); createConvertLinalgToAffineLoopsPass();
/// Create a pass that rewrites init_tensor to alloc_tensor. /// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor.
std::unique_ptr<Pass> createLinalgInitTensorToAllocTensorPass(); std::unique_ptr<Pass> createEmptyTensorToAllocTensorPass();
/// Create a pass to convert Linalg operations which work on tensors to use /// Create a pass to convert Linalg operations which work on tensors to use
/// buffers instead. /// buffers instead.
std::unique_ptr<OperationPass<func::FuncOp>> createLinalgBufferizePass(); std::unique_ptr<OperationPass<func::FuncOp>> createLinalgBufferizePass();
/// Create a pass to convert named Linalg operations to Linalg generic /// Create a pass to convert named Linalg operations to Linalg generic
/// operations. /// operations.
std::unique_ptr<OperationPass<func::FuncOp>> createLinalgGeneralizationPass(); std::unique_ptr<OperationPass<func::FuncOp>> createLinalgGeneralizationPass();
▲ Show 20 Lines • Show All 65 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/Passes.td

Show All 18 Lines let description = [{
This pass only converts ops that operate on ranked tensors. It can be This pass only converts ops that operate on ranked tensors. It can be
run on op which contains linalg ops (most commonly a run on op which contains linalg ops (most commonly a
FunctionOpInterface op). FunctionOpInterface op).
}]; }];
let constructor = "mlir::createConvertElementwiseToLinalgPass()"; let constructor = "mlir::createConvertElementwiseToLinalgPass()";
let dependentDialects = ["linalg::LinalgDialect", "memref::MemRefDialect"]; let dependentDialects = ["linalg::LinalgDialect", "memref::MemRefDialect"];
} }
def LinalgInitTensorToAllocTensor : Pass<"linalg-init-tensor-to-alloc-tensor"> { def EmptyTensorToAllocTensor : Pass<"empty-tensor-to-alloc-tensor"> {
let summary = "Replace all init_tensor ops by alloc_tensor ops."; let summary = "Replace all empty ops by alloc_tensor ops.";
let description = [{ let description = [{
init_tensor ops return a tensor of unspecified contents who's only purpose tensor.empty ops return a tensor of unspecified contents who's only purpose
is to carry the tensor shape. This pass converts such ops to is to carry the tensor shape. This pass converts such ops to
bufferization.alloc_tensor ops, which bufferize to buffer allocations. bufferization.alloc_tensor ops, which bufferize to buffer allocations.
}]; }];
let constructor = "mlir::createLinalgInitTensorToAllocTensorPass()"; let constructor = "mlir::createEmptyTensorToAllocTensorPass()";
} }
def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> { def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> {
let summary = "Remove unit-extent dimension in Linalg ops on tensors"; let summary = "Remove unit-extent dimension in Linalg ops on tensors";
let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()"; let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
let options = [ let options = [
Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool", Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
/*default=*/"false", /*default=*/"false",
▲ Show 20 Lines • Show All 221 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

Show First 20 Lines • Show All 500 Lines • ▼ Show 20 Lines ```
} -> tensor<f32> } -> tensor<f32>
``` ```
is split into: is split into:
``` ```
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32> %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
%1 = linalg.init_tensor [4] : tensor<4xf32> %1 = tensor.empty() : tensor<4xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32>
%3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) { ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
^bb0(%arg3: f32, %arg5: f32): ^bb0(%arg3: f32, %arg5: f32):
%5 = arith.addf %arg3, %arg4 : f32 %5 = arith.addf %arg3, %arg4 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
Show All 34 Lines let description = [{
``` ```
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)> #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)> #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
#map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map5 = affine_map<(d0, d1, d2) -> (d0, d1)> #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
%0 = linalg.init_tensor [16, 32, 64] : tensor<16x32x64xf32> %0 = tensor.empty() : tensor<16x32x64xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) -> %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
tensor<16x32x64xf32> tensor<16x32x64xf32>
%2 = linalg.init_tensor [64, 4] : tensor<64x4xi1> %2 = tensor.empty() : tensor<64x4xi1>
%3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3], %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel", "reduction"]} iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>) ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>)
outs(%1 : tensor<16x32x64xf32>) { outs(%1 : tensor<16x32x64xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
%5 = arith.mulf %arg3, %arg4 : f32 %5 = arith.mulf %arg3, %arg4 : f32
%6 = arith.addf %arg6, %5 : f32 %6 = arith.addf %arg6, %5 : f32
▲ Show 20 Lines • Show All 391 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Linalg/Transforms/HoistPadding.h

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
/// } : tensor<?x?xf32> to tensor<4x8xf32> /// } : tensor<?x?xf32> to tensor<4x8xf32>
/// compute(%0) /// compute(%0)
/// ``` /// ```
/// ///
/// IR resembling the following is produced: /// IR resembling the following is produced:
/// ///
/// ``` /// ```
/// scf.for (%i) { /// scf.for (%i) {
/// %packed_init = linalg.init_tensor range(%j) : tensor<?x4x8xf32> /// %packed_init = tensor.empty range(%j) : tensor<?x4x8xf32>
/// %packed = scf.for (%k) iter_args(%p : %packed_init) { /// %packed = scf.for (%k) iter_args(%p : %packed_init) {
/// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32> /// %st0 = tensor.extract_slice f(%i, %k) : ... to tensor<?x?xf32>
/// %0 = tensor.pad %st0 low[0, 0] high[...] { /// %0 = tensor.pad %st0 low[0, 0] high[...] {
/// ^bb0( ... ): /// ^bb0( ... ):
/// linalg.yield %pad /// linalg.yield %pad
/// } : tensor<?x?xf32> to tensor<4x8xf32> /// } : tensor<?x?xf32> to tensor<4x8xf32>
/// %1 = tensor.insert_slice %0 ... /// %1 = tensor.insert_slice %0 ...
/// : tensor<4x8xf32> to tensor<?x4x8xf32> /// : tensor<4x8xf32> to tensor<?x4x8xf32>
Show All 17 Lines

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

Show First 20 Lines • Show All 1,190 Lines • ▼ Show 20 Lines
rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad, rewriteAsPaddedOp(OpBuilder &b, LinalgOp opToPad,
ArrayRef<int64_t> paddingDimensions, ArrayRef<int64_t> paddingDimensions,
ArrayRef<Attribute> paddingValues, ArrayRef<Attribute> paddingValues,
ArrayRef<bool> packPaddings, LinalgOp &paddedOp); ArrayRef<bool> packPaddings, LinalgOp &paddedOp);
using OptimizeCopyFn = using OptimizeCopyFn =
std::function<LogicalResult(PatternRewriter &, tensor::PadOp, Value)>; std::function<LogicalResult(PatternRewriter &, tensor::PadOp, Value)>;
/// Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
/// InsertSliceOp. For now, only constant padding values are supported. /// InsertSliceOp. For now, only constant padding values are supported.
/// `OptimizeCopyFn` can be used to customize copying step optimization. /// `OptimizeCopyFn` can be used to customize copying step optimization.
struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> { struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
GeneralizePadOpPattern(MLIRContext *context, GeneralizePadOpPattern(MLIRContext *context,
OptimizeCopyFn optimizeCopyFn = nullptr, OptimizeCopyFn optimizeCopyFn = nullptr,
PatternBenefit benefit = 1) PatternBenefit benefit = 1)
: OpRewritePattern<tensor::PadOp>(context, benefit), : OpRewritePattern<tensor::PadOp>(context, benefit),
optimizeCopyFn(std::move(optimizeCopyFn)) {} optimizeCopyFn(std::move(optimizeCopyFn)) {}
▲ Show 20 Lines • Show All 194 Lines • ▼ Show 20 Lines
/// %y = arith.addf %arg1, %arg2 : f32 /// %y = arith.addf %arg1, %arg2 : f32
/// linalg.yield %y : f32 /// linalg.yield %y : f32
/// } -> tensor<f32> /// } -> tensor<f32>
/// ``` /// ```
/// To: /// To:
/// ``` /// ```
/// %cst = arith.constant 0.000000e+00 : f32 /// %cst = arith.constant 0.000000e+00 : f32
/// %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32> /// %0 = tensor.expand_shape %in [[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
/// %1 = linalg.init_tensor [4] : tensor<4xf32> /// %1 = tensor.empty [4] : tensor<4xf32>
/// %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> /// %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32>
/// %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, /// %3 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
/// affine_map<(d0, d1) -> (d0)>], /// affine_map<(d0, d1) -> (d0)>],
/// iterator_types = ["parallel", "reduction"]} /// iterator_types = ["parallel", "reduction"]}
/// ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) { /// ins(%0 : tensor<4x8xf32>) outs(%2 : tensor<4xf32>) {
/// ^bb0(%arg3: f32, %arg5: f32): /// ^bb0(%arg3: f32, %arg5: f32):
/// %5 = arith.addf %arg3, %arg4 : f32 /// %5 = arith.addf %arg3, %arg4 : f32
/// linalg.yield %5 : f32 /// linalg.yield %5 : f32
Show All 40 Lines
/// ///
/// ``` /// ```
/// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)> /// #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2 * 4 + d3)>
/// #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)> /// #map1 = affine_map<(d0, d1, d2, d3) -> (d2 * 4 + d3, d1)>
/// #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)> /// #map2 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
/// #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> /// #map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
/// #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> /// #map4 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
/// #map5 = affine_map<(d0, d1, d2) -> (d0, d1)> /// #map5 = affine_map<(d0, d1, d2) -> (d0, d1)>
/// %0 = linalg.init_tensor [16, 32, 64] : tensor<16x32x64xf32> /// %0 = tensor.empty [16, 32, 64] : tensor<16x32x64xf32>
/// %cst = arith.constant 0.000000e+00 : f32 /// %cst = arith.constant 0.000000e+00 : f32
/// %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) -> /// %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<16x32x64xf32>) ->
/// tensor<16x32x64xf32> /// tensor<16x32x64xf32>
/// %2 = linalg.init_tensor [64, 4] : tensor<64x4xi1> /// %2 = tensor.empty [64, 4] : tensor<64x4xi1>
/// ///
/// %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3], /// %3 = linalg.generic {indexing_maps = [#map0, #map1, #map2, #map3],
/// iterator_types = ["parallel", "parallel", "parallel", "reduction"]} /// iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
/// ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>) /// ins(%A, %B, %2 : tensor<16x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>)
/// outs(%1 : tensor<16x32x64xf32>) { /// outs(%1 : tensor<16x32x64xf32>) {
/// ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32): /// ^bb0(%arg3: f32, %arg4: f32, %arg5: i1, %arg6: f32):
/// %5 = arith.mulf %arg3, %arg4 : f32 /// %5 = arith.mulf %arg3, %arg4 : f32
/// %6 = arith.addf %arg6, %5 : f32 /// %6 = arith.addf %arg6, %5 : f32
Show All 23 Lines

mlir/include/mlir/Dialect/Tensor/IR/TensorBase.td

Show First 20 Lines • Show All 41 Lines • ▼ Show 20 Lines let description = [{
The `tensor` type is actually a builtin type (it lives in the builtin The `tensor` type is actually a builtin type (it lives in the builtin
dialect), and does not live in this dialect. dialect), and does not live in this dialect.
}]; }];
let hasConstantMaterializer = 1; let hasConstantMaterializer = 1;
let dependentDialects = [ let dependentDialects = [
"AffineDialect",
"arith::ArithDialect", "arith::ArithDialect",
"complex::ComplexDialect", "complex::ComplexDialect",
]; ];
} }
#endif // TENSOR_BASE #endif // TENSOR_BASE

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

Show First 20 Lines • Show All 131 Lines • ▼ Show 20 Linesdef Tensor_DimOp : Tensor_Op<"dim", [NoSideEffect, ShapedDimOpInterface]> {
}]; }];
let hasCanonicalizer = 1; let hasCanonicalizer = 1;
let hasFolder = 1; let hasFolder = 1;
let hasVerifier = 1; let hasVerifier = 1;
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// EmptyOp
//===----------------------------------------------------------------------===//
def Tensor_EmptyOp : Tensor_Op<"empty",
[NoSideEffect,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
let summary = "empty tensor operation";
let description = [{
`tensor.empty` is an operation that defines a tensor of a particular shape.
The shape could be dynamic or static. The contents of the tensor are
unspecified and the only purpose of the op result is to materialize the
specified shape in IR and make it available to other transformations.
`tensor.empty` is useful in transformations that expect destination style
ops. I.e., ops that implement `DestinationStyleOpInterface`. Ops that are
not in destination style can be made compatible with such transformations
with a `tensor.empty` destination.
Note: This op can be lowered to a `bufferization.alloc_tensor`, at which
point it turns into an explicit buffer allocation.
}];
let arguments = (ins Variadic<Index>:$dynamicSizes);
let results = (outs AnyRankedTensor:$result);
let assemblyFormat = "`(`$dynamicSizes`)` attr-dict `:` type($result)";
let extraClassDeclaration = [{
RankedTensorType getType() {
return getResult().getType().cast<RankedTensorType>();
}
// Return both static and dynamic sizes as a list of `OpFoldResult`.
SmallVector<OpFoldResult> getMixedSizes();
// Return the Value of the dynamic size of the tensor at dimension `idx`.
// Asserts that the shape is dynamic at that `idx`.
Value getDynamicSize(unsigned idx);
}];
let builders = [
// Build with fully static sizes.
OpBuilder<(ins "ArrayRef<int64_t>":$staticShape, "Type":$elementType)>,
// Build with mixed static/dynamic sizes.
OpBuilder<(ins "ArrayRef<int64_t>":$staticShape, "Type":$elementType,
"ValueRange":$dynamicSizes)>,
// Build with mixed static/dynamic sizes.
OpBuilder<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType)>
];
let hasCanonicalizer = 1;
let hasVerifier = 1;
}
//===----------------------------------------------------------------------===//
// ExtractOp // ExtractOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
def Tensor_ExtractOp : Tensor_Op<"extract", def Tensor_ExtractOp : Tensor_Op<"extract",
[NoSideEffect, [NoSideEffect,
TypesMatchWith<"result type matches element type of tensor", TypesMatchWith<"result type matches element type of tensor",
"tensor", "result", "tensor", "result",
"$_self.cast<ShapedType>().getElementType()">]> { "$_self.cast<ShapedType>().getElementType()">]> {
▲ Show 20 Lines • Show All 1,434 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Tensor/Transforms/TransformUtils.h

Show First 20 Lines • Show All 56 Lines • ▼ Show 20 Lines
/// : tensor<3x?x?x11x?xf32> into tensor<?x?xf32> /// : tensor<3x?x?x11x?xf32> into tensor<?x?xf32>
/// %2 = tensor.extract_slice %1 [%offt0, %offt1][%size0, %size1][1, 1] /// %2 = tensor.extract_slice %1 [%offt0, %offt1][%size0, %size1][1, 1]
/// : tensor<?x?xf32> to tensor<?x?xf32> /// : tensor<?x?xf32> to tensor<?x?xf32>
/// ``` /// ```
/// ///
/// We can construct %2 by generating the following, which only uses `%0`: /// We can construct %2 by generating the following, which only uses `%0`:
/// ///
/// ``` /// ```
/// %dest = linalg.init_tensor [%size0, %size1] : tensor<?x?xf32> /// %dest = tensor.empty(%size0, %size1) : tensor<?x?xf32>
/// %1 = tensor.dim %0, %c1 : tensor<3x?x?x11x?xf32> /// %1 = tensor.dim %0, %c1 : tensor<3x?x?x11x?xf32>
/// %2 = tensor.dim %0, %c2 : tensor<3x?x?x11x?xf32> /// %2 = tensor.dim %0, %c2 : tensor<3x?x?x11x?xf32>
/// %3 = tensor.dim %0, %c4 : tensor<3x?x?x11x?xf32> /// %3 = tensor.dim %0, %c4 : tensor<3x?x?x11x?xf32>
/// ///
/// %result = scf.for %iv0 = %c0 to %arg2 step %c1 iter_args(%arg6 = %dest) -> /// %result = scf.for %iv0 = %c0 to %arg2 step %c1 iter_args(%arg6 = %dest) ->
/// (tensor<?x?xf32>) { /// (tensor<?x?xf32>) {
/// %5 = scf.for %iv1 = %c0 to %arg4 step %c1 iter_args(%arg8 = %arg6) /// %5 = scf.for %iv1 = %c0 to %arg4 step %c1 iter_args(%arg8 = %arg6)
/// -> (tensor<?x?xf32>) { /// -> (tensor<?x?xf32>) {
▲ Show 20 Lines • Show All 137 Lines • Show Last 20 Lines

mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h

Show First 20 Lines • Show All 394 Lines • ▼ Show 20 Lines
// clang-format off // clang-format off
/// ``` /// ```
/// %0 = linalg.generic ... -> tensor<3x7x11x10xf32> /// %0 = linalg.generic ... -> tensor<3x7x11x10xf32>
/// %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32> /// %1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32>
/// %2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32> /// %2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32>
/// ``` /// ```
/// This class helps build the below IR to replace %2: /// This class helps build the below IR to replace %2:
/// ``` /// ```
/// %dest = linalg.init_tensor() : tensor<10x10xf32> /// %dest = tensor.empty() : tensor<10x10xf32>
/// %2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> { /// %2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> {
/// %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv) /// %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv)
/// %3:3 = arith.delinearize_index %iv into (3, 7, 11) /// %3:3 = arith.delinearize_index %iv into (3, 7, 11)
/// ///
/// // This function takes %3 (multiIndices) and the parameters for the slice below. /// // This function takes %3 (multiIndices) and the parameters for the slice below.
/// %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] : /// %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] :
/// tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32> /// tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32>
/// ///
▲ Show 20 Lines • Show All 54 Lines • Show Last 20 Lines

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

Show First 20 Lines • Show All 550 Lines • ▼ Show 20 LineselementwiseMatchAndRewriteHelper(Operation *operation,
// Construct the indexing maps needed for linalg.generic ops. // Construct the indexing maps needed for linalg.generic ops.
SmallVector<Type> bodyArgTypes; SmallVector<Type> bodyArgTypes;
for (Value in : operation->getOperands()) for (Value in : operation->getOperands())
bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType())); bodyArgTypes.emplace_back(getElementTypeOrSelf(in.getType()));
SmallVector<Type> opResultTypes; SmallVector<Type> opResultTypes;
SmallVector<Value> initTensors; SmallVector<Value> emptyTensors;
SmallVector<Value> dynDims; SmallVector<Value> dynDims;
dynDims.resize(results.front().getType().cast<ShapedType>().getRank()); dynDims.resize(results.front().getType().cast<ShapedType>().getRank());
for (auto arg : operation->getOperands()) { for (auto arg : operation->getOperands()) {
auto operandTy = arg.getType().cast<ShapedType>(); auto operandTy = arg.getType().cast<ShapedType>();
for (int i = 0; i < operandTy.getRank(); i++) { for (int i = 0; i < operandTy.getRank(); i++) {
if (operandTy.isDynamicDim(i) && !dynDims[i]) if (operandTy.isDynamicDim(i) && !dynDims[i])
dynDims[i] = rewriter.create<tensor::DimOp>(loc, arg, i); dynDims[i] = rewriter.create<tensor::DimOp>(loc, arg, i);
} }
} }
SmallVector<Value> filteredDims = condenseValues(dynDims); SmallVector<Value> filteredDims = condenseValues(dynDims);
for (auto result : results) { for (auto result : results) {
auto resultTy = result.getType().template cast<ShapedType>(); auto resultTy = result.getType().template cast<ShapedType>();
initTensors.push_back(rewriter.create<linalg::InitTensorOp>( emptyTensors.push_back(rewriter.create<tensor::EmptyOp>(
loc, filteredDims, resultTy.getShape(), resultTy.getElementType())); loc, resultTy.getShape(), resultTy.getElementType(), filteredDims));
opResultTypes.push_back(result.getType()); opResultTypes.push_back(result.getType());
} }
auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range( auto bodyResultTypes = llvm::to_vector<4>(llvm::map_range(
initTensors, [](Value v) { return getElementTypeOrSelf(v); })); emptyTensors, [](Value v) { return getElementTypeOrSelf(v); }));
SmallVector<Value, 2> operands; SmallVector<Value, 2> operands;
SmallVector<AffineMap, 2> indexingMaps; SmallVector<AffineMap, 2> indexingMaps;
indexingMaps.reserve(operation->getNumOperands() + bodyResultTypes.size()); indexingMaps.reserve(operation->getNumOperands() + bodyResultTypes.size());
// Input indexing maps may be broadcasted. // Input indexing maps may be broadcasted.
for (Value operand : operation->getOperands()) { for (Value operand : operation->getOperands()) {
ShapedType type = operand.getType().cast<ShapedType>(); ShapedType type = operand.getType().cast<ShapedType>();
Show All 27 Lines indexingMaps.push_back(AffineMap::get(
rewriter.getContext())); rewriter.getContext()));
} }
indexingMaps.append(operation->getNumResults(), indexingMaps.append(operation->getNumResults(),
rewriter.getMultiDimIdentityMap(rank)); rewriter.getMultiDimIdentityMap(rank));
bool didEncounterError = false; bool didEncounterError = false;
auto linalgOp = rewriter.create<linalg::GenericOp>( auto linalgOp = rewriter.create<linalg::GenericOp>(
loc, opResultTypes, operands, initTensors, indexingMaps, loc, opResultTypes, operands, emptyTensors, indexingMaps,
getNParallelLoopsAttrs(rank), getNParallelLoopsAttrs(rank),
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
Value opResult = createLinalgBodyCalculationForElementwiseOp( Value opResult = createLinalgBodyCalculationForElementwiseOp(
operation, blockArgs.take_front(operation->getNumOperands()), operation, blockArgs.take_front(operation->getNumOperands()),
bodyResultTypes, rewriter); bodyResultTypes, rewriter);
if (!opResult) { if (!opResult) {
didEncounterError = true; didEncounterError = true;
return; return;
▲ Show 20 Lines • Show All 131 Lines • ▼ Show 20 Lines if (axis != i) {
if (inputTy.isDynamicDim(i)) if (inputTy.isDynamicDim(i))
dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i)); dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i));
} }
} }
Type reduceTy = RankedTensorType::get(reduceShape, resultTy.getElementType()); Type reduceTy = RankedTensorType::get(reduceShape, resultTy.getElementType());
// First fill the output buffer with the init value. // First fill the output buffer with the init value.
auto initTensor = rewriter auto emptyTensor =
.create<linalg::InitTensorOp>(loc, dynDims, reduceShape, rewriter
resultTy.getElementType()) .create<tensor::EmptyOp>(loc, reduceShape, resultTy.getElementType(),
dynDims)
.getResult(); .getResult();
auto fillValueAttr = createInitialValueForReduceOp(op, elementTy, rewriter); auto fillValueAttr = createInitialValueForReduceOp(op, elementTy, rewriter);
if (!fillValueAttr) if (!fillValueAttr)
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op, "No initial value found for reduction operation"); op, "No initial value found for reduction operation");
auto fillValue = rewriter.create<arith::ConstantOp>(loc, fillValueAttr); auto fillValue = rewriter.create<arith::ConstantOp>(loc, fillValueAttr);
auto filledTensor = rewriter auto filledTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{fillValue}, .create<linalg::FillOp>(loc, ValueRange{fillValue},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
SmallVector<AffineExpr, 2> srcExprs; SmallVector<AffineExpr, 2> srcExprs;
SmallVector<AffineExpr, 2> dstExprs; SmallVector<AffineExpr, 2> dstExprs;
SmallVector<StringRef, 4> iteratorTypes; SmallVector<StringRef, 4> iteratorTypes;
for (unsigned int i = 0, rank = inputTy.getRank(); i != rank; ++i) { for (unsigned int i = 0, rank = inputTy.getRank(); i != rank; ++i) {
srcExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext())); srcExprs.push_back(mlir::getAffineDimExpr(i, rewriter.getContext()));
▲ Show 20 Lines • Show All 290 Lines • ▼ Show 20 Lines for (const auto &permutation : llvm::enumerate(perms.getValues<APInt>())) {
if (!operandTy.hasRank() || operandTy.isDynamicDim(index)) { if (!operandTy.hasRank() || operandTy.isDynamicDim(index)) {
dynDims[value] = rewriter.create<tensor::DimOp>(loc, input, index); dynDims[value] = rewriter.create<tensor::DimOp>(loc, input, index);
} }
inputExprs[value] = rewriter.getAffineDimExpr(index); inputExprs[value] = rewriter.getAffineDimExpr(index);
} }
SmallVector<Value> filteredDims = condenseValues(dynDims); SmallVector<Value> filteredDims = condenseValues(dynDims);
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, resultTy.getShape(), resultTy.getElementType()); loc, resultTy.getShape(), resultTy.getElementType(), filteredDims);
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
AffineMap::get(resultTy.getRank(), /*symbolCount=*/0, inputExprs, AffineMap::get(resultTy.getRank(), /*symbolCount=*/0, inputExprs,
rewriter.getContext()), rewriter.getContext()),
rewriter.getMultiDimIdentityMap(resultTy.getRank())}; rewriter.getMultiDimIdentityMap(resultTy.getRank())};
rewriter.replaceOpWithNewOp<linalg::GenericOp>( rewriter.replaceOpWithNewOp<linalg::GenericOp>(
op, resultTy, op.getInput1(), ValueRange{initTensor}, affineMaps, op, resultTy, op.getInput1(), ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(resultTy.getRank()), getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
nestedBuilder.create<linalg::YieldOp>(loc, *args.begin()); nestedBuilder.create<linalg::YieldOp>(loc, *args.begin());
}); });
return success(); return success();
} }
}; };
▲ Show 20 Lines • Show All 89 Lines • ▼ Show 20 Lines if (shiftValues.size() == 1) {
rewriter.getContext())); rewriter.getContext()));
shiftArg = indexingMaps.size() - 1; shiftArg = indexingMaps.size() - 1;
} }
// Indexing maps for output values. // Indexing maps for output values.
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(rank)); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(rank));
// Construct the indexing maps needed for linalg.generic ops. // Construct the indexing maps needed for linalg.generic ops.
Value initTensor = rewriter.create<linalg::InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, ArrayRef<Value>({dynDims}), outputTy.getShape(), loc, outputTy.getShape(), outputTy.getElementType(),
outputTy.getElementType()); ArrayRef<Value>({dynDims}));
auto linalgOp = rewriter.create<linalg::GenericOp>( auto linalgOp = rewriter.create<linalg::GenericOp>(
loc, outputTy, genericInputs, ValueRange{initTensor}, indexingMaps, loc, outputTy, genericInputs, ValueRange{emptyTensor}, indexingMaps,
getNParallelLoopsAttrs(rank), getNParallelLoopsAttrs(rank),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange blockArgs) { ValueRange blockArgs) {
Value value = blockArgs[0]; Value value = blockArgs[0];
Type valueTy = value.getType(); Type valueTy = value.getType();
// For now we do all of our math in 64-bit. This is not optimal but // For now we do all of our math in 64-bit. This is not optimal but
// should be correct for now, consider computing correct bit depth // should be correct for now, consider computing correct bit depth
▲ Show 20 Lines • Show All 100 Lines • ▼ Show 20 Lines auto dynamicDimsOr =
checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()}); checkHasDynamicBatchDims(rewriter, op, {input, op.getOutput()});
if (!dynamicDimsOr.has_value()) if (!dynamicDimsOr.has_value())
return failure(); return failure();
SmallVector<Value> dynamicDims = dynamicDimsOr.value(); SmallVector<Value> dynamicDims = dynamicDimsOr.value();
if (op.getMode() != "NEAREST_NEIGHBOR" && op.getMode() != "BILINEAR") if (op.getMode() != "NEAREST_NEIGHBOR" && op.getMode() != "BILINEAR")
return failure(); return failure();
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynamicDims, resultTy.getShape(), resultElementTy); loc, resultTy.getShape(), resultElementTy, dynamicDims);
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
rewriter.getMultiDimIdentityMap(resultTy.getRank())}; rewriter.getMultiDimIdentityMap(resultTy.getRank())};
auto genericOp = rewriter.create<linalg::GenericOp>( auto genericOp = rewriter.create<linalg::GenericOp>(
loc, resultTy, ValueRange({}), ValueRange{initTensor}, affineMaps, loc, resultTy, ValueRange({}), ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(resultTy.getRank())); getNParallelLoopsAttrs(resultTy.getRank()));
rewriter.replaceOp(op, genericOp.getResult(0)); rewriter.replaceOp(op, genericOp.getResult(0));
OpBuilder::InsertionGuard regionGuard(rewriter); OpBuilder::InsertionGuard regionGuard(rewriter);
rewriter.createBlock(&genericOp.getRegion(), genericOp.getRegion().end(), rewriter.createBlock(&genericOp.getRegion(), genericOp.getRegion().end(),
TypeRange({resultElementTy}), loc); TypeRange({resultElementTy}), loc);
Value batch = rewriter.create<linalg::IndexOp>(loc, 0); Value batch = rewriter.create<linalg::IndexOp>(loc, 0);
Value y = rewriter.create<linalg::IndexOp>(loc, 1); Value y = rewriter.create<linalg::IndexOp>(loc, 1);
▲ Show 20 Lines • Show All 282 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::ConcatOp op, OpAdaptor adaptor,
Value resultDimSize = sizes[axis]; Value resultDimSize = sizes[axis];
for (auto arg : adaptor.getOperands().drop_front()) { for (auto arg : adaptor.getOperands().drop_front()) {
auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue); auto size = rewriter.createOrFold<tensor::DimOp>(loc, arg, axisValue);
resultDimSize = resultDimSize =
rewriter.createOrFold<arith::AddIOp>(loc, resultDimSize, size); rewriter.createOrFold<arith::AddIOp>(loc, resultDimSize, size);
} }
sizes[axis] = resultDimSize; sizes[axis] = resultDimSize;
Value init = rewriter.create<linalg::InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynDims, resultType.getShape(), resultType.getElementType()); loc, resultType.getShape(), resultType.getElementType(), dynDims);
Value zeroVal = rewriter.createOrFold<arith::ConstantOp>( Value zeroVal = rewriter.createOrFold<arith::ConstantOp>(
loc, rewriter.getZeroAttr(resultType.getElementType())); loc, rewriter.getZeroAttr(resultType.getElementType()));
Value result = Value result = rewriter
rewriter .create<linalg::FillOp>(loc, ValueRange{zeroVal},
.create<linalg::FillOp>(loc, ValueRange{zeroVal}, ValueRange{init}) ValueRange{emptyTensor})
.result(); .result();
auto toOpFoldResult = [](Value v) -> OpFoldResult { auto toOpFoldResult = [](Value v) -> OpFoldResult {
auto op = v.getDefiningOp<arith::ConstantIndexOp>(); auto op = v.getDefiningOp<arith::ConstantIndexOp>();
if (!op) if (!op)
return v; return v;
return op.getValue(); return op.getValue();
}; };
for (auto arg : adaptor.getOperands()) { for (auto arg : adaptor.getOperands()) {
Show All 28 Lines for (int i = 0; i < inputTy.getRank(); i++) {
if (inputTy.isDynamicDim(i)) { if (inputTy.isDynamicDim(i)) {
dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i)); dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i));
} }
} }
Value axisDimSize = rewriter.create<tensor::DimOp>(loc, input, axis); Value axisDimSize = rewriter.create<tensor::DimOp>(loc, input, axis);
// First fill the output buffer with the init value. // First fill the output buffer with the init value.
auto initTensor = rewriter auto emptyTensor = rewriter
.create<linalg::InitTensorOp>( .create<tensor::EmptyOp>(loc, inputTy.getShape(),
loc, ArrayRef<Value>({dynDims}), inputTy.getElementType(),
inputTy.getShape(), inputTy.getElementType()) ArrayRef<Value>({dynDims}))
.getResult(); .getResult();
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
rewriter.getMultiDimIdentityMap(resultTy.getRank())}; rewriter.getMultiDimIdentityMap(resultTy.getRank())};
rewriter.replaceOpWithNewOp<linalg::GenericOp>( rewriter.replaceOpWithNewOp<linalg::GenericOp>(
op, resultTy, ArrayRef<Value>({}), ValueRange{initTensor}, affineMaps, op, resultTy, ArrayRef<Value>({}), ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(resultTy.getRank()), getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
llvm::SmallVector<Value> indices; llvm::SmallVector<Value> indices;
for (unsigned int i = 0; i < inputTy.getRank(); i++) { for (unsigned int i = 0; i < inputTy.getRank(); i++) {
auto index = auto index =
rewriter.create<linalg::IndexOp>(nestedLoc, i).getResult(); rewriter.create<linalg::IndexOp>(nestedLoc, i).getResult();
if (i == axis) { if (i == axis) {
auto one = rewriter.create<arith::ConstantIndexOp>(nestedLoc, 1); auto one = rewriter.create<arith::ConstantIndexOp>(nestedLoc, 1);
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::TileOp op, OpAdaptor adaptor,
SmallVector<Value> dynDims; SmallVector<Value> dynDims;
for (int i = 0; i < inputTy.getRank(); i++) { for (int i = 0; i < inputTy.getRank(); i++) {
if (inputTy.isDynamicDim(i) || multiples[i] == -1) { if (inputTy.isDynamicDim(i) || multiples[i] == -1) {
dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i)); dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i));
} }
} }
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
op.getLoc(), dynDims, genericShape, elementTy); op.getLoc(), genericShape, elementTy, dynDims);
// We needs to map the input shape to the non-broadcasted dimensions. // We needs to map the input shape to the non-broadcasted dimensions.
SmallVector<AffineExpr, 4> dimExprs; SmallVector<AffineExpr, 4> dimExprs;
dimExprs.reserve(rank); dimExprs.reserve(rank);
for (unsigned i = 0; i < rank; ++i) for (unsigned i = 0; i < rank; ++i)
dimExprs.push_back(rewriter.getAffineDimExpr(i * 2 + 1)); dimExprs.push_back(rewriter.getAffineDimExpr(i * 2 + 1));
auto readAffineMap = auto readAffineMap =
AffineMap::get(/*dimCount=*/rank * 2, /*symbolCount=*/0, dimExprs, AffineMap::get(/*dimCount=*/rank * 2, /*symbolCount=*/0, dimExprs,
rewriter.getContext()); rewriter.getContext());
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
readAffineMap, rewriter.getMultiDimIdentityMap(genericShape.size())}; readAffineMap, rewriter.getMultiDimIdentityMap(genericShape.size())};
auto genericOp = rewriter.create<linalg::GenericOp>( auto genericOp = rewriter.create<linalg::GenericOp>(
loc, RankedTensorType::get(genericShape, elementTy), input, loc, RankedTensorType::get(genericShape, elementTy), input,
ValueRange{initTensor}, affineMaps, ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(genericShape.size()), getNParallelLoopsAttrs(genericShape.size()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) { [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin()); nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
}); });
rewriter.replaceOpWithNewOp<tosa::ReshapeOp>( rewriter.replaceOpWithNewOp<tosa::ReshapeOp>(
op, resultTy, genericOp.getResult(0), op, resultTy, genericOp.getResult(0),
rewriter.getI64ArrayAttr(resultTy.getShape())); rewriter.getI64ArrayAttr(resultTy.getShape()));
▲ Show 20 Lines • Show All 113 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tosa::ArgMaxOp argmaxOp,
SmallVector<Value> dynDims; SmallVector<Value> dynDims;
for (int i = 0; i < inputTy.getRank(); i++) { for (int i = 0; i < inputTy.getRank(); i++) {
if (inputTy.isDynamicDim(i) && i != axis) { if (inputTy.isDynamicDim(i) && i != axis) {
dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i)); dynDims.push_back(rewriter.create<tensor::DimOp>(loc, input, i));
} }
} }
// First fill the output buffer for the index. // First fill the output buffer for the index.
auto initTensorIdx = auto emptyTensorIdx = rewriter
rewriter .create<tensor::EmptyOp>(loc, resultTy.getShape(),
.create<linalg::InitTensorOp>(loc, dynDims, resultTy.getShape(), outElementTy, dynDims)
outElementTy)
.getResult(); .getResult();
auto fillValueIdx = rewriter.create<arith::ConstantOp>( auto fillValueIdx = rewriter.create<arith::ConstantOp>(
loc, rewriter.getIntegerAttr(outElementTy, 0)); loc, rewriter.getIntegerAttr(outElementTy, 0));
auto filledTensorIdx = auto filledTensorIdx =
rewriter rewriter
.create<linalg::FillOp>(loc, ValueRange{fillValueIdx}, .create<linalg::FillOp>(loc, ValueRange{fillValueIdx},
ValueRange{initTensorIdx}) ValueRange{emptyTensorIdx})
.result(); .result();
// Second fill the output buffer for the running max. // Second fill the output buffer for the running max.
auto initTensorMax = rewriter auto emptyTensorMax = rewriter
.create<linalg::InitTensorOp>( .create<tensor::EmptyOp>(loc, resultTy.getShape(),
loc, dynDims, resultTy.getShape(), inElementTy) inElementTy, dynDims)
.getResult(); .getResult();
auto fillValueMaxAttr = auto fillValueMaxAttr =
createInitialValueForReduceOp(argmaxOp, inElementTy, rewriter); createInitialValueForReduceOp(argmaxOp, inElementTy, rewriter);
if (!fillValueMaxAttr) if (!fillValueMaxAttr)
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
argmaxOp, "unsupported tosa.argmax element type"); argmaxOp, "unsupported tosa.argmax element type");
auto fillValueMax = auto fillValueMax =
rewriter.create<arith::ConstantOp>(loc, fillValueMaxAttr); rewriter.create<arith::ConstantOp>(loc, fillValueMaxAttr);
auto filledTensorMax = auto filledTensorMax =
rewriter rewriter
.create<linalg::FillOp>(loc, ValueRange{fillValueMax}, .create<linalg::FillOp>(loc, ValueRange{fillValueMax},
ValueRange{initTensorMax}) ValueRange{emptyTensorMax})
.result(); .result();
// We need to reduce along the arg-max axis, with parallel operations along // We need to reduce along the arg-max axis, with parallel operations along
// the rest. // the rest.
SmallVector<StringRef, 4> iteratorTypes; SmallVector<StringRef, 4> iteratorTypes;
iteratorTypes.resize(inputTy.getRank(), getParallelIteratorTypeName()); iteratorTypes.resize(inputTy.getRank(), getParallelIteratorTypeName());
iteratorTypes[axis] = getReductionIteratorTypeName(); iteratorTypes[axis] = getReductionIteratorTypeName();
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::GatherOp op, OpAdaptor adaptor,
if (!dynamicDimsOr.has_value()) if (!dynamicDimsOr.has_value())
return failure(); return failure();
SmallVector<Value> dynamicDims = dynamicDimsOr.value(); SmallVector<Value> dynamicDims = dynamicDimsOr.value();
auto resultElementTy = resultTy.getElementType(); auto resultElementTy = resultTy.getElementType();
auto loc = op.getLoc(); auto loc = op.getLoc();
auto initTensor = auto emptyTensor =
rewriter rewriter
.create<linalg::InitTensorOp>(loc, dynamicDims, resultTy.getShape(), .create<tensor::EmptyOp>(loc, resultTy.getShape(), resultElementTy,
resultElementTy) dynamicDims)
.getResult(); .getResult();
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
AffineMap::get( AffineMap::get(
/*dimCount=*/resultTy.getRank(), /*symbolCount=*/0, /*dimCount=*/resultTy.getRank(), /*symbolCount=*/0,
{rewriter.getAffineDimExpr(0), rewriter.getAffineDimExpr(1)}, {rewriter.getAffineDimExpr(0), rewriter.getAffineDimExpr(1)},
rewriter.getContext()), rewriter.getContext()),
rewriter.getMultiDimIdentityMap(resultTy.getRank())}; rewriter.getMultiDimIdentityMap(resultTy.getRank())};
auto genericOp = rewriter.create<linalg::GenericOp>( auto genericOp = rewriter.create<linalg::GenericOp>(
loc, ArrayRef<Type>({resultTy}), ValueRange{indices}, loc, ArrayRef<Type>({resultTy}), ValueRange{indices},
ValueRange{initTensor}, affineMaps, ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(resultTy.getRank()), getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &b, Location loc, ValueRange args) { [&](OpBuilder &b, Location loc, ValueRange args) {
auto indexValue = args[0]; auto indexValue = args[0];
auto index0 = rewriter.create<linalg::IndexOp>(loc, 0); auto index0 = rewriter.create<linalg::IndexOp>(loc, 0);
Value index1 = rewriter.create<arith::IndexCastOp>( Value index1 = rewriter.create<arith::IndexCastOp>(
loc, rewriter.getIndexType(), indexValue); loc, rewriter.getIndexType(), indexValue);
auto index2 = rewriter.create<linalg::IndexOp>(loc, 2); auto index2 = rewriter.create<linalg::IndexOp>(loc, 2);
Value extract = rewriter.create<tensor::ExtractOp>( Value extract = rewriter.create<tensor::ExtractOp>(
Show All 28 Lines LogicalResult matchAndRewrite(tosa::TableOp op,
SmallVector<Value> dynDims; SmallVector<Value> dynDims;
for (int i = 0; i < resultTy.getRank(); ++i) { for (int i = 0; i < resultTy.getRank(); ++i) {
if (inputTy.isDynamicDim(i)) { if (inputTy.isDynamicDim(i)) {
dynDims.push_back( dynDims.push_back(
rewriter.create<tensor::DimOp>(loc, op.getOperand(0), i)); rewriter.create<tensor::DimOp>(loc, op.getOperand(0), i));
} }
} }
auto initTensor = auto emptyTensor = rewriter
rewriter .create<tensor::EmptyOp>(loc, resultTy.getShape(),
.create<linalg::InitTensorOp>(loc, dynDims, resultTy.getShape(), resultElementTy, dynDims)
resultElementTy)
.getResult(); .getResult();
SmallVector<AffineMap, 2> affineMaps = { SmallVector<AffineMap, 2> affineMaps = {
rewriter.getMultiDimIdentityMap(resultTy.getRank()), rewriter.getMultiDimIdentityMap(resultTy.getRank()),
rewriter.getMultiDimIdentityMap(resultTy.getRank())}; rewriter.getMultiDimIdentityMap(resultTy.getRank())};
auto genericOp = rewriter.create<linalg::GenericOp>( auto genericOp = rewriter.create<linalg::GenericOp>(
loc, resultTy, ValueRange({input}), ValueRange{initTensor}, affineMaps, loc, resultTy, ValueRange({input}), ValueRange{emptyTensor}, affineMaps,
getNParallelLoopsAttrs(resultTy.getRank())); getNParallelLoopsAttrs(resultTy.getRank()));
rewriter.replaceOp(op, genericOp.getResult(0)); rewriter.replaceOp(op, genericOp.getResult(0));
{ {
OpBuilder::InsertionGuard regionGuard(rewriter); OpBuilder::InsertionGuard regionGuard(rewriter);
Block *block = rewriter.createBlock( Block *block = rewriter.createBlock(
&genericOp.getRegion(), genericOp.getRegion().end(), &genericOp.getRegion(), genericOp.getRegion().end(),
TypeRange({inputElementTy, resultElementTy}), {loc, loc}); TypeRange({inputElementTy, resultElementTy}), {loc, loc});
▲ Show 20 Lines • Show All 141 Lines • Show Last 20 Lines

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

Show First 20 Lines • Show All 235 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::Conv2DOp op, OpAdaptor adaptor,
Value weightPermValue = Value weightPermValue =
rewriter.create<arith::ConstantOp>(loc, weightPermAttr); rewriter.create<arith::ConstantOp>(loc, weightPermAttr);
Type newWeightTy = Type newWeightTy =
RankedTensorType::get(newWeightShape, weightTy.getElementType()); RankedTensorType::get(newWeightShape, weightTy.getElementType());
weight = rewriter.create<tosa::TransposeOp>(loc, newWeightTy, weight, weight = rewriter.create<tosa::TransposeOp>(loc, newWeightTy, weight,
weightPermValue); weightPermValue);
Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy); Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy);
Value initTensor = rewriter.create<linalg::InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, resultTy.getShape(), resultETy); loc, resultTy.getShape(), resultETy, filteredDims);
Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr); Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr);
Value zeroTensor = rewriter Value zeroTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{zero}, .create<linalg::FillOp>(loc, ValueRange{zero},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
// Extract the attributes for convolution. // Extract the attributes for convolution.
llvm::SmallVector<int64_t> stride, dilation; llvm::SmallVector<int64_t> stride, dilation;
getValuesFromIntArrayAttribute(strideTosaAttr, stride); getValuesFromIntArrayAttribute(strideTosaAttr, stride);
getValuesFromIntArrayAttribute(dilationTosaAttr, dilation); getValuesFromIntArrayAttribute(dilationTosaAttr, dilation);
// Create the convolution op. // Create the convolution op.
auto strideAttr = DenseIntElementsAttr::get( auto strideAttr = DenseIntElementsAttr::get(
RankedTensorType::get({2}, rewriter.getI64Type()), stride); RankedTensorType::get({2}, rewriter.getI64Type()), stride);
auto dilationAttr = DenseIntElementsAttr::get( auto dilationAttr = DenseIntElementsAttr::get(
RankedTensorType::get({2}, rewriter.getI64Type()), dilation); RankedTensorType::get({2}, rewriter.getI64Type()), dilation);
// Create maps for the bias broadcasting // Create maps for the bias broadcasting
SmallVector<AffineMap, 4> indexingMaps; SmallVector<AffineMap, 4> indexingMaps;
indexingMaps.push_back(AffineMap::get( indexingMaps.push_back(AffineMap::get(
/*dimCount=*/resultTy.getRank(), /*symbolCount=*/0, /*dimCount=*/resultTy.getRank(), /*symbolCount=*/0,
{rewriter.getAffineDimExpr(3)}, rewriter.getContext())); {rewriter.getAffineDimExpr(3)}, rewriter.getContext()));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank())); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank()));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank())); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultTy.getRank()));
Value biasInitTensor = rewriter.create<linalg::InitTensorOp>( Value biasEmptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, resultTy.getShape(), resultETy); loc, resultTy.getShape(), resultETy, filteredDims);
if (isQuantized) { if (isQuantized) {
auto quantizationInfo = auto quantizationInfo =
op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>(); op->getAttr("quantization_info").cast<tosa::ConvOpQuantizationAttr>();
auto iZp = rewriter.getI32IntegerAttr(quantizationInfo.getInputZp()); auto iZp = rewriter.getI32IntegerAttr(quantizationInfo.getInputZp());
auto kZp = rewriter.getI32IntegerAttr(quantizationInfo.getWeightZp()); auto kZp = rewriter.getI32IntegerAttr(quantizationInfo.getWeightZp());
auto iZpVal = rewriter.create<arith::ConstantOp>(loc, iZp); auto iZpVal = rewriter.create<arith::ConstantOp>(loc, iZp);
auto kZpVal = rewriter.create<arith::ConstantOp>(loc, kZp); auto kZpVal = rewriter.create<arith::ConstantOp>(loc, kZp);
Value conv = Value conv =
rewriter rewriter
.create<linalg::Conv2DNhwcHwcfQOp>( .create<linalg::Conv2DNhwcHwcfQOp>(
loc, resultTy, ValueRange{input, weight, iZpVal, kZpVal}, loc, resultTy, ValueRange{input, weight, iZpVal, kZpVal},
ValueRange{zeroTensor}, strideAttr, dilationAttr) ValueRange{zeroTensor}, strideAttr, dilationAttr)
->getResult(0); ->getResult(0);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, resultTy, ValueRange({bias, conv}), biasInitTensor, loc, resultTy, ValueRange({bias, conv}), biasEmptyTensor,
indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()), indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddIOp>( Value added = nestedBuilder.create<arith::AddIOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
rewriter.replaceOp(op, result); rewriter.replaceOp(op, result);
return success(); return success();
} }
Value conv = rewriter Value conv = rewriter
.create<linalg::Conv2DNhwcHwcfOp>( .create<linalg::Conv2DNhwcHwcfOp>(
loc, resultTy, ValueRange{input, weight}, loc, resultTy, ValueRange{input, weight},
ValueRange{zeroTensor}, strideAttr, dilationAttr) ValueRange{zeroTensor}, strideAttr, dilationAttr)
->getResult(0); ->getResult(0);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, resultTy, ValueRange({bias, conv}), biasInitTensor, loc, resultTy, ValueRange({bias, conv}), biasEmptyTensor,
indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()), indexingMaps, getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddFOp>( Value added = nestedBuilder.create<arith::AddFOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
▲ Show 20 Lines • Show All 98 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::DepthwiseConv2DOp op, OpAdaptor adaptor,
SmallVector<AffineMap, 4> indexingMaps; SmallVector<AffineMap, 4> indexingMaps;
indexingMaps.push_back(AffineMap::get( indexingMaps.push_back(AffineMap::get(
/*dimCount=*/resultRank, /*symbolCount=*/0, /*dimCount=*/resultRank, /*symbolCount=*/0,
{rewriter.getAffineDimExpr(3)}, rewriter.getContext())); {rewriter.getAffineDimExpr(3)}, rewriter.getContext()));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank)); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank)); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(resultRank));
Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy); Attribute resultZeroAttr = rewriter.getZeroAttr(resultETy);
Value initTensor = rewriter.create<linalg::InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, linalgConvTy.getShape(), resultETy); loc, linalgConvTy.getShape(), resultETy, filteredDims);
Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr); Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr);
Value zeroTensor = rewriter Value zeroTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{zero}, .create<linalg::FillOp>(loc, ValueRange{zero},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
Value biasInitTensor = rewriter.create<linalg::InitTensorOp>( Value biasEmptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, resultTy.getShape(), resultETy); loc, resultTy.getShape(), resultETy, filteredDims);
if (!isQuantized) { if (!isQuantized) {
Value conv = rewriter Value conv = rewriter
.create<linalg::DepthwiseConv2DNhwcHwcmOp>( .create<linalg::DepthwiseConv2DNhwcHwcmOp>(
loc, linalgConvTy, ValueRange{input, weight}, loc, linalgConvTy, ValueRange{input, weight},
ValueRange{zeroTensor}, strideAttr, dilationAttr) ValueRange{zeroTensor}, strideAttr, dilationAttr)
.getResult(0); .getResult(0);
SmallVector<ReassociationExprs, 4> reassociationMap; SmallVector<ReassociationExprs, 4> reassociationMap;
createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter); createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter);
Value convReshape = rewriter.create<tensor::CollapseShapeOp>( Value convReshape = rewriter.create<tensor::CollapseShapeOp>(
loc, resultTy, conv, reassociationMap); loc, resultTy, conv, reassociationMap);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, resultTy, ValueRange({bias, convReshape}), loc, resultTy, ValueRange({bias, convReshape}),
biasInitTensor, indexingMaps, biasEmptyTensor, indexingMaps,
getNParallelLoopsAttrs(resultRank), getNParallelLoopsAttrs(resultRank),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddFOp>( Value added = nestedBuilder.create<arith::AddFOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
Show All 10 Lines if (!isQuantized) {
SmallVector<ReassociationExprs, 4> reassociationMap; SmallVector<ReassociationExprs, 4> reassociationMap;
createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter); createDepthwiseConvCollapseMap(resultRank, reassociationMap, rewriter);
Value convReshape = rewriter.create<tensor::CollapseShapeOp>( Value convReshape = rewriter.create<tensor::CollapseShapeOp>(
loc, resultTy, conv, reassociationMap); loc, resultTy, conv, reassociationMap);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, resultTy, ValueRange({bias, convReshape}), loc, resultTy, ValueRange({bias, convReshape}),
biasInitTensor, indexingMaps, biasEmptyTensor, indexingMaps,
getNParallelLoopsAttrs(resultRank), getNParallelLoopsAttrs(resultRank),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddIOp>( Value added = nestedBuilder.create<arith::AddIOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
Show All 31 Lines matchAndRewrite(tosa::MatMulOp op, OpAdaptor adaptor,
if (!secondOperandTy.hasRank() || secondOperandTy.isDynamicDim(2)) { if (!secondOperandTy.hasRank() || secondOperandTy.isDynamicDim(2)) {
dynDims[2] = rewriter.create<tensor::DimOp>(loc, op->getOperand(1), 2); dynDims[2] = rewriter.create<tensor::DimOp>(loc, op->getOperand(1), 2);
} }
SmallVector<Value> filteredDims = condenseValues(dynDims); SmallVector<Value> filteredDims = condenseValues(dynDims);
auto zeroAttr = rewriter.getZeroAttr(outputElementTy); auto zeroAttr = rewriter.getZeroAttr(outputElementTy);
Value zero = rewriter.create<arith::ConstantOp>(loc, zeroAttr); Value zero = rewriter.create<arith::ConstantOp>(loc, zeroAttr);
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, outputTy.getShape(), outputTy.getElementType()); loc, outputTy.getShape(), outputTy.getElementType(), filteredDims);
Value zeroTensor = rewriter Value zeroTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{zero}, .create<linalg::FillOp>(loc, ValueRange{zero},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
if (!op.getQuantizationInfo()) { if (!op.getQuantizationInfo()) {
rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>( rewriter.replaceOpWithNewOp<linalg::BatchMatmulOp>(
op, TypeRange{op.getType()}, op, TypeRange{op.getType()},
ValueRange{adaptor.getA(), adaptor.getB()}, ValueRange{zeroTensor}); ValueRange{adaptor.getA(), adaptor.getB()}, ValueRange{zeroTensor});
return success(); return success();
} }
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines matchAndRewrite(tosa::FullyConnectedOp op, OpAdaptor adaptor,
// Broadcast the bias. // Broadcast the bias.
indexingMaps.push_back(AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0, indexingMaps.push_back(AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0,
{rewriter.getAffineDimExpr(1)}, {rewriter.getAffineDimExpr(1)},
rewriter.getContext())); rewriter.getContext()));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(outputTy.getRank())); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(outputTy.getRank()));
indexingMaps.push_back(rewriter.getMultiDimIdentityMap(outputTy.getRank())); indexingMaps.push_back(rewriter.getMultiDimIdentityMap(outputTy.getRank()));
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, filteredDims, outputTy.getShape(), outputTy.getElementType()); loc, outputTy.getShape(), outputTy.getElementType(), filteredDims);
// When quantized, the input elemeny type is not the same as the output // When quantized, the input elemeny type is not the same as the output
Attribute resultZeroAttr = rewriter.getZeroAttr(outputETy); Attribute resultZeroAttr = rewriter.getZeroAttr(outputETy);
Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr); Value zero = rewriter.create<arith::ConstantOp>(loc, resultZeroAttr);
Value zeroTensor = rewriter Value zeroTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{zero}, .create<linalg::FillOp>(loc, ValueRange{zero},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
SmallVector<int64_t> permutation{1, 0}; SmallVector<int64_t> permutation{1, 0};
auto permutationAttr = DenseIntElementsAttr::get( auto permutationAttr = DenseIntElementsAttr::get(
RankedTensorType::get({2}, rewriter.getI64Type()), permutation); RankedTensorType::get({2}, rewriter.getI64Type()), permutation);
Value permutationValue = Value permutationValue =
rewriter.create<arith::ConstantOp>(loc, permutationAttr); rewriter.create<arith::ConstantOp>(loc, permutationAttr);
SmallVector<int64_t> newWeightShape{weightShape[1], weightShape[0]}; SmallVector<int64_t> newWeightShape{weightShape[1], weightShape[0]};
Type newWeightTy = Type newWeightTy =
RankedTensorType::get(newWeightShape, weightTy.getElementType()); RankedTensorType::get(newWeightShape, weightTy.getElementType());
Value transposedWeight = rewriter.create<tosa::TransposeOp>( Value transposedWeight = rewriter.create<tosa::TransposeOp>(
loc, newWeightTy, weight, permutationValue); loc, newWeightTy, weight, permutationValue);
auto biasInitTensor = auto biasEmptyTensor =
rewriter rewriter
.create<linalg::InitTensorOp>(loc, filteredDims, .create<tensor::EmptyOp>(loc, outputTy.getShape(), outputETy,
outputTy.getShape(), outputETy) filteredDims)
->getResults(); ->getResults();
if (!op.getQuantizationInfo()) { if (!op.getQuantizationInfo()) {
Value matmul = rewriter Value matmul = rewriter
.create<linalg::MatmulOp>( .create<linalg::MatmulOp>(
loc, TypeRange{op.getType()}, loc, TypeRange{op.getType()},
ValueRange{input, transposedWeight}, zeroTensor) ValueRange{input, transposedWeight}, zeroTensor)
->getResult(0); ->getResult(0);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, outputTy, ValueRange({bias, matmul}), biasInitTensor, loc, outputTy, ValueRange({bias, matmul}), biasEmptyTensor,
indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()), indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddFOp>( Value added = nestedBuilder.create<arith::AddFOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
Show All 11 Lines Value matmul =
.create<linalg::QuantizedMatmulOp>( .create<linalg::QuantizedMatmulOp>(
loc, TypeRange{op.getType()}, loc, TypeRange{op.getType()},
ValueRange{input, transposedWeight, inputZp, outputZp}, ValueRange{input, transposedWeight, inputZp, outputZp},
zeroTensor) zeroTensor)
->getResult(0); ->getResult(0);
Value result = Value result =
rewriter rewriter
.create<linalg::GenericOp>( .create<linalg::GenericOp>(
loc, outputTy, ValueRange({bias, matmul}), biasInitTensor, loc, outputTy, ValueRange({bias, matmul}), biasEmptyTensor,
indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()), indexingMaps, getNParallelLoopsAttrs(outputTy.getRank()),
[&](OpBuilder &nestedBuilder, Location nestedLoc, [&](OpBuilder &nestedBuilder, Location nestedLoc,
ValueRange args) { ValueRange args) {
Value added = nestedBuilder.create<arith::AddIOp>( Value added = nestedBuilder.create<arith::AddIOp>(
loc, args[0], args[1]); loc, args[0], args[1]);
nestedBuilder.create<linalg::YieldOp>(nestedLoc, added); nestedBuilder.create<linalg::YieldOp>(nestedLoc, added);
}) })
.getResult(0); .getResult(0);
▲ Show 20 Lines • Show All 50 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
SmallVector<int64_t> kernel, stride; SmallVector<int64_t> kernel, stride;
getValuesFromIntArrayAttribute(op.getKernel(), kernel); getValuesFromIntArrayAttribute(op.getKernel(), kernel);
getValuesFromIntArrayAttribute(op.getStride(), stride); getValuesFromIntArrayAttribute(op.getStride(), stride);
Attribute strideAttr = rewriter.getI64VectorAttr(stride); Attribute strideAttr = rewriter.getI64VectorAttr(stride);
Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1}); Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
// Create the linalg op that performs pooling. // Create the linalg op that performs pooling.
Value initTensor = rewriter.create<linalg::InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynamicDims, resultTy.getShape(), resultTy.getElementType()); loc, resultTy.getShape(), resultTy.getElementType(), dynamicDims);
Value filledInitTensor = Value filledEmptyTensor =
rewriter rewriter
.create<linalg::FillOp>(loc, ValueRange{initialValue}, .create<linalg::FillOp>(loc, ValueRange{initialValue},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
Value fakeWindowDims = Value fakeWindowDims =
rewriter.create<linalg::InitTensorOp>(loc, kernel, resultETy); rewriter.create<tensor::EmptyOp>(loc, kernel, resultETy);
rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxOp>( rewriter.replaceOpWithNewOp<linalg::PoolingNhwcMaxOp>(
op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims}, op, ArrayRef<Type>{resultTy}, ValueRange{paddedInput, fakeWindowDims},
filledInitTensor, strideAttr, dilationAttr); filledEmptyTensor, strideAttr, dilationAttr);
return success(); return success();
} }
}; };
class AvgPool2dConverter : public OpRewritePattern<tosa::AvgPool2dOp> { class AvgPool2dConverter : public OpRewritePattern<tosa::AvgPool2dOp> {
public: public:
using OpRewritePattern<tosa::AvgPool2dOp>::OpRewritePattern; using OpRewritePattern<tosa::AvgPool2dOp>::OpRewritePattern;
Show All 31 Lines LogicalResult matchAndRewrite(tosa::AvgPool2dOp op,
SmallVector<int64_t> kernel, stride; SmallVector<int64_t> kernel, stride;
getValuesFromIntArrayAttribute(op.getKernel(), kernel); getValuesFromIntArrayAttribute(op.getKernel(), kernel);
getValuesFromIntArrayAttribute(op.getStride(), stride); getValuesFromIntArrayAttribute(op.getStride(), stride);
Attribute strideAttr = rewriter.getI64VectorAttr(stride); Attribute strideAttr = rewriter.getI64VectorAttr(stride);
Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1}); Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
// Create the linalg op that performs pooling. // Create the linalg op that performs pooling.
Value poolInitTensor = rewriter.create<linalg::InitTensorOp>( Value poolEmptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynamicDims, accTy.getShape(), accETy); loc, accTy.getShape(), accETy, dynamicDims);
Value filledInitTensor = Value filledEmptyTensor =
rewriter rewriter
.create<linalg::FillOp>(loc, ValueRange{initialValue}, .create<linalg::FillOp>(loc, ValueRange{initialValue},
ValueRange{poolInitTensor}) ValueRange{poolEmptyTensor})
.result(); .result();
Value fakeWindowDims = Value fakeWindowDims =
rewriter.create<linalg::InitTensorOp>(loc, kernel, accETy); rewriter.create<tensor::EmptyOp>(loc, kernel, accETy);
// Sum across the pooled region. // Sum across the pooled region.
Value poolingOp = rewriter Value poolingOp = rewriter
.create<linalg::PoolingNhwcSumOp>( .create<linalg::PoolingNhwcSumOp>(
loc, ArrayRef<Type>{accTy}, loc, ArrayRef<Type>{accTy},
ValueRange{paddedInput, fakeWindowDims}, ValueRange{paddedInput, fakeWindowDims},
filledInitTensor, strideAttr, dilationAttr) filledEmptyTensor, strideAttr, dilationAttr)
.getResult(0); .getResult(0);
// Normalize the summed value by the number of elements grouped in each // Normalize the summed value by the number of elements grouped in each
// pool. // pool.
auto poolingOpTy = poolingOp.getType().cast<ShapedType>(); auto poolingOpTy = poolingOp.getType().cast<ShapedType>();
auto affineMap = rewriter.getMultiDimIdentityMap(resultTy.getRank()); auto affineMap = rewriter.getMultiDimIdentityMap(resultTy.getRank());
Value genericInitTensor = rewriter.create<linalg::InitTensorOp>( Value genericEmptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynamicDims, resultTy.getShape(), resultETy); loc, resultTy.getShape(), resultETy, dynamicDims);
auto genericOp = rewriter.create<linalg::GenericOp>( auto genericOp = rewriter.create<linalg::GenericOp>(
loc, ArrayRef<Type>({resultTy}), ValueRange{poolingOp}, loc, ArrayRef<Type>({resultTy}), ValueRange{poolingOp},
ValueRange{genericInitTensor}, ValueRange{genericEmptyTensor},
ArrayRef<AffineMap>({affineMap, affineMap}), ArrayRef<AffineMap>({affineMap, affineMap}),
getNParallelLoopsAttrs(resultTy.getRank()), getNParallelLoopsAttrs(resultTy.getRank()),
[&](OpBuilder &b, Location loc, ValueRange args) { [&](OpBuilder &b, Location loc, ValueRange args) {
auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0); auto zero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
auto one = rewriter.create<arith::ConstantIndexOp>(loc, 1); auto one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
auto iH = rewriter.create<arith::ConstantIndexOp>( auto iH = rewriter.create<arith::ConstantIndexOp>(
loc, poolingOpTy.getDimSize(1) - 1); loc, poolingOpTy.getDimSize(1) - 1);
auto iW = rewriter.create<arith::ConstantIndexOp>( auto iW = rewriter.create<arith::ConstantIndexOp>(
▲ Show 20 Lines • Show All 142 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

Show First 20 Lines • Show All 539 Lines • ▼ Show 20 Lines ReifyRankedShapedTypeOpInterface interface =
cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation()); cast<ReifyRankedShapedTypeOpInterface>(padOp.getOperation());
if (failed(interface.reifyResultShapes(rewriter, reifiedShape))) if (failed(interface.reifyResultShapes(rewriter, reifiedShape)))
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
padOp, "failed to reify tensor.pad op result shape"); padOp, "failed to reify tensor.pad op result shape");
auto oldResultType = padOp.getResultType(); auto oldResultType = padOp.getResultType();
SmallVector<int64_t, 4> staticShape(oldResultType.getRank(), SmallVector<int64_t, 4> staticShape(oldResultType.getRank(),
ShapedType::kDynamicSize); ShapedType::kDynamicSize);
auto newInitOp = rewriter.create<InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
padOp.getLoc(), reifiedShape.front(), staticShape, padOp.getLoc(), staticShape, oldResultType.getElementType(),
oldResultType.getElementType()); reifiedShape.front());
auto newFillOp = rewriter.create<FillOp>( auto newFillOp = rewriter.create<FillOp>(
fillOp.getLoc(), ValueRange{padValue}, ValueRange{newInitOp}); fillOp.getLoc(), ValueRange{padValue}, ValueRange{emptyTensor});
rewriter.replaceOpWithNewOp<tensor::CastOp>(padOp, oldResultType, rewriter.replaceOpWithNewOp<tensor::CastOp>(padOp, oldResultType,
newFillOp.result()); newFillOp.result());
return success(); return success();
} }
}; };
/// Fold tensor.insert_slice(tensor.pad(<input>), linalg.fill) into /// Fold tensor.insert_slice(tensor.pad(<input>), linalg.fill) into
▲ Show 20 Lines • Show All 821 Lines • ▼ Show 20 Lines if (outputElementType != bbArg.getType())
<< "output element type " << outputElementType << "output element type " << outputElementType
<< " does not match corresponding block argument type " << " does not match corresponding block argument type "
<< bbArg.getType(); << bbArg.getType();
} }
return success(); return success();
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// InitTensorOp
//===----------------------------------------------------------------------===//
void InitTensorOp::build(OpBuilder &b, OperationState &result,
ArrayRef<OpFoldResult> sizes, Type elementType,
ArrayRef<NamedAttribute> attrs) {
SmallVector<Value, 4> dynamicSizes;
SmallVector<int64_t, 4> staticSizes;
dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
ShapedType::kDynamicSize);
auto resultType = RankedTensorType ::get(staticSizes, elementType);
build(b, result, resultType, dynamicSizes, b.getI64ArrayAttr(staticSizes));
result.addAttributes(attrs);
}
LogicalResult InitTensorOp::verify() {
RankedTensorType resultType = getType();
SmallVector<int64_t, 4> staticSizes = llvm::to_vector<4>(llvm::map_range(
getStaticSizes().cast<ArrayAttr>(),
[](Attribute a) -> int64_t { return a.cast<IntegerAttr>().getInt(); }));
if (failed(verifyListOfOperandsOrIntegers(
*this, "sizes", resultType.getRank(), getStaticSizes(), getSizes(),
ShapedType::isDynamic)))
return failure();
if (getStaticSizes().size() != static_cast<unsigned>(resultType.getRank()))
return emitError("expected ") << resultType.getRank() << " sizes values";
Type expectedType = InitTensorOp::inferResultType(
staticSizes, resultType.getElementType(), resultType.getEncoding());
if (resultType != expectedType) {
return emitError("specified type ")
<< resultType << " does not match the inferred type "
<< expectedType;
}
return success();
}
Type InitTensorOp::inferResultType(ArrayRef<int64_t> staticSizes,
Type elementType, Attribute encoding) {
return RankedTensorType::get(staticSizes, elementType, encoding);
}
SmallVector<OpFoldResult> InitTensorOp::getMixedSizes() {
SmallVector<OpFoldResult> mixedSizes;
mixedSizes.reserve(getType().getRank());
unsigned dynamicValIndex = 0;
for (Attribute attr : getStaticSizes()) {
auto intAttr = attr.cast<IntegerAttr>();
if (!ShapedType::isDynamic(intAttr.getInt())) {
mixedSizes.push_back(intAttr);
continue;
}
mixedSizes.push_back(getSizes()[dynamicValIndex++]);
}
return mixedSizes;
}
namespace {
/// Change the type of the result of a `linalg.init_tensor` by making the result
/// type statically sized along dimension that in the original operation where
/// defined as dynamic, but the size was defined using a `constant` op. For
/// example
///
/// %c5 = arith.constant 5: index
/// %0 = linalg.init_tensor [%arg0, %c5] : tensor<?x?xf32>
///
/// to
///
/// %0 = linalg.init_tensor [%arg0, 5] : tensor<?x5xf32>
struct ReplaceStaticShapeDims : OpRewritePattern<InitTensorOp> {
using OpRewritePattern<InitTensorOp>::OpRewritePattern;
LogicalResult matchAndRewrite(InitTensorOp op,
PatternRewriter &rewriter) const override {
SmallVector<Value, 4> dynamicSizes;
SmallVector<int64_t, 4> staticSizes;
for (unsigned i = 0, e = op.getType().getRank(); i != e; ++i) {
// If the size is already static, nothing to do.
if (!op.isDynamicSize(i)) {
staticSizes.push_back(op.getStaticSize(i));
continue;
}
// If the size is dynamic but defined using a `constant` op, get the
// constant value to find the static size to use.
unsigned operandNum = op.getIndexOfDynamicSize(i);
Value sizeOperand = op.getOperand(operandNum);
if (auto constantIndexOp =
sizeOperand.getDefiningOp<arith::ConstantIndexOp>()) {
staticSizes.push_back(constantIndexOp.value());
continue;
}
// Fallback case. Keep the size dynamic.
dynamicSizes.push_back(sizeOperand);
staticSizes.push_back(ShapedType::kDynamicSize);
}
RankedTensorType newType =
RankedTensorType::get(staticSizes, op.getType().getElementType());
if (newType == op.getType())
return failure();
auto newOp =
rewriter.create<InitTensorOp>(op.getLoc(), newType, dynamicSizes,
rewriter.getI64ArrayAttr(staticSizes));
rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
return success();
}
};
} // namespace
namespace {
/// Since `init_tensor` operation creates a tensor needed only for its shape, a
/// slice of this is also needed only for its shape. The result can be
/// replaced by a new init_tensor operation of the same size as the extract
/// slice op.
struct FoldInitTensorWithExtractSliceOp
: public OpRewritePattern<tensor::ExtractSliceOp> {
using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
PatternRewriter &rewriter) const override {
if (!sliceOp.getSource().getDefiningOp<linalg::InitTensorOp>())
return failure();
// ExtractSliceOp may be rank-reducing; its dynamic sizes must be preserved
// as well as its result type.
rewriter.replaceOpWithNewOp<linalg::InitTensorOp>(
sliceOp, sliceOp.getSizes(),
sliceOp.getResult().getType().cast<RankedTensorType>().getShape(),
sliceOp.getSourceType().getElementType());
return success();
}
};
template <typename TensorReshapeOp>
struct FoldInitTensorWithTensorReshapeOp
: public OpRewritePattern<TensorReshapeOp> {
using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
PatternRewriter &rewriter) const override {
if (!reshapeOp.getSrc().template getDefiningOp<InitTensorOp>())
return failure();
Location loc = reshapeOp.getLoc();
ReifiedRankedShapedTypeDims resultShapes;
ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
cast<ReifyRankedShapedTypeOpInterface>(reshapeOp.getOperation());
if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
resultShapes)) ||
!llvm::hasSingleElement(resultShapes))
return failure();
Value initTensor = rewriter.create<InitTensorOp>(
loc, getAsOpFoldResult(resultShapes[0]),
reshapeOp.getResultType().getElementType());
if (initTensor.getType() != reshapeOp.getResultType()) {
rewriter.replaceOpWithNewOp<tensor::CastOp>(
reshapeOp, reshapeOp.getResultType(), initTensor);
} else {
rewriter.replaceOp(reshapeOp, initTensor);
}
return success();
}
};
struct FoldInitTensorWithDimOp : public OpRewritePattern<tensor::DimOp> {
using OpRewritePattern<tensor::DimOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tensor::DimOp dimOp,
PatternRewriter &rewriter) const override {
Optional<int64_t> maybeConstantIndex = dimOp.getConstantIndex();
auto initTensorOp = dimOp.getSource().getDefiningOp<linalg::InitTensorOp>();
if (!initTensorOp || !maybeConstantIndex)
return failure();
if (!initTensorOp.isDynamicSize(*maybeConstantIndex))
return failure();
rewriter.replaceOp(dimOp, initTensorOp.getDynamicSize(*maybeConstantIndex));
return success();
}
};
/// Canonicalize
///
/// ```mlir
/// %0 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
/// %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x?xf32>
/// ```
///
/// into
///
/// ```mlir
/// %0 = linalg.init_tensor [4, %d1] : tensor<4x?xf32>
/// ```
///
/// This assumes the input program is correct in terms of its shape. So it
/// is safe to assume that `%d0` is in fact 4. If that was not the case, the
/// input program is wrong to begin with, so its undefined behavior anyway (i.e.
/// this optimization can still triggering without violating program semantics).
struct FoldInitTensorWithTensorCastOp
: public OpRewritePattern<tensor::CastOp> {
using OpRewritePattern<tensor::CastOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tensor::CastOp castOp,
PatternRewriter &rewriter) const override {
if (!canFoldIntoProducerOp(castOp))
return failure();
auto producer = castOp.getSource().getDefiningOp<InitTensorOp>();
if (!producer)
return failure();
auto resultType = castOp->getResult(0).getType().cast<RankedTensorType>();
ArrayRef<int64_t> resultShape = resultType.getShape();
SmallVector<OpFoldResult> currMixedSizes = producer.getMixedSizes();
SmallVector<OpFoldResult> newMixedSizes;
newMixedSizes.reserve(currMixedSizes.size());
assert(resultShape.size() == currMixedSizes.size() &&
"mismatch in result shape and sizes of init_tensor op");
for (auto it : llvm::zip(resultShape, currMixedSizes)) {
int64_t newDim = std::get<0>(it);
OpFoldResult currDim = std::get<1>(it);
// Case 1: The init tensor dim is static. Check that the tensor cast
// result dim matches.
if (auto attr = currDim.dyn_cast<Attribute>()) {
if (ShapedType::isDynamic(newDim) ||
newDim != attr.cast<IntegerAttr>().getInt()) {
// Something is off, the cast result shape cannot be more dynamic than
// the init tensor result shape (enforced by `canFoldIntoProducer`).
// Abort for now.
return rewriter.notifyMatchFailure(
producer, "mismatch in static value of shape of init "
"tensor result and cast result");
}
newMixedSizes.push_back(attr);
continue;
}
// Case 2 : The tensor cast shape is static, but init tensor result shape
// is dynamic.
if (!ShapedType::isDynamic(newDim)) {
newMixedSizes.push_back(rewriter.getIndexAttr(newDim));
continue;
}
// Case 3 : The tensor cast shape is dynamic and init tensor result shape
// is dynamic. Use the dynamic value from the init tensor op.
newMixedSizes.push_back(currDim);
}
rewriter.replaceOpWithNewOp<InitTensorOp>(castOp, newMixedSizes,
resultType.getElementType());
return success();
}
};
} // namespace
void InitTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<FoldInitTensorWithTensorCastOp, FoldInitTensorWithDimOp,
FoldInitTensorWithExtractSliceOp,
FoldInitTensorWithTensorReshapeOp<tensor::ExpandShapeOp>,
FoldInitTensorWithTensorReshapeOp<tensor::CollapseShapeOp>,
ReplaceStaticShapeDims>(context);
}
LogicalResult InitTensorOp::reifyResultShapes(
OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
auto shapes = llvm::to_vector<4>(llvm::map_range(
llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
if (isDynamicSize(dim))
return getDynamicSize(dim);
return builder.create<arith::ConstantIndexOp>(getLoc(),
getStaticSize(dim));
}));
reifiedReturnShapes.emplace_back(std::move(shapes));
return success();
}
//===----------------------------------------------------------------------===//
// YieldOp // YieldOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
void linalg::YieldOp::print(OpAsmPrinter &p) { void linalg::YieldOp::print(OpAsmPrinter &p) {
if (getNumOperands() > 0) if (getNumOperands() > 0)
p << ' ' << getOperands(); p << ' ' << getOperands();
p.printOptionalAttrDict((*this)->getAttrs()); p.printOptionalAttrDict((*this)->getAttrs());
if (getNumOperands() > 0) if (getNumOperands() > 0)
▲ Show 20 Lines • Show All 451 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp

Show All 40 Lines
/// linalg.yield %0, %1 : ... /// linalg.yield %0, %1 : ...
/// } -> (..., ...) /// } -> (..., ...)
/// return %result#0, %result#1 /// return %result#0, %result#1
/// ``` /// ```
/// ///
/// gets split into /// gets split into
/// ///
/// ```mlir /// ```mlir
/// %init = linalg.init_tensor ... /// %init = tensor.empty ...
/// %op0:3 = linalg.generic ... ins(%arg0, %arg1, %arg2 : ...) /// %op0:3 = linalg.generic ... ins(%arg0, %arg1, %arg2 : ...)
/// outs(%init0, %init1, %init : ...) /// outs(%init0, %init1, %init : ...)
/// ^bb0(%b0: ... , %b1: ... , %b2: ... , %b3: ..., %b4: ..., %b5: ...): /// ^bb0(%b0: ... , %b1: ... , %b2: ... , %b3: ..., %b4: ..., %b5: ...):
/// %0 = <s0> %b0, %b1 : ... /// %0 = <s0> %b0, %b1 : ...
/// linalg.yield %0, %..., %0 : ... /// linalg.yield %0, %..., %0 : ...
/// } -> (..., ..., ...) /// } -> (..., ..., ...)
/// %op1:2 = linalg.generic ... ins(%arg0, %arg1, %arg2, %op0#2 : ...) /// %op1:2 = linalg.generic ... ins(%arg0, %arg1, %arg2, %op0#2 : ...)
/// outs(%init0, %init1 : ...) { /// outs(%init0, %init1 : ...) {
/// ^bb0(%b0: ... , %b1: ... , %b2: ... , %b3: ..., %b4: ..., %b5: ...): /// ^bb0(%b0: ... , %b1: ... , %b2: ... , %b3: ..., %b4: ..., %b5: ...):
/// %1 = <s1> %b3, %b2 : ... /// %1 = <s1> %b3, %b2 : ...
/// linalg.yield %..., %1 : ... /// linalg.yield %..., %1 : ...
/// } -> (..., ...) /// } -> (..., ...)
/// return %op0#0, %op1#1 /// return %op0#0, %op1#1
/// ``` /// ```
/// ///
/// After canonicalization this is expected to be /// After canonicalization this is expected to be
/// ///
/// ```mlir /// ```mlir
/// %init = linalg.init_tensor ... /// %init = tensor.empty ...
/// %op0 = linalg.generic ... ins(%arg0, %arg1, : ...) /// %op0 = linalg.generic ... ins(%arg0, %arg1, : ...)
/// outs(%init : ...) /// outs(%init : ...)
/// ^bb0(%b0: ... , %b1: ... , %b2: ...): /// ^bb0(%b0: ... , %b1: ... , %b2: ...):
/// %0 = <s0> %b0, %b1 : ... /// %0 = <s0> %b0, %b1 : ...
/// linalg.yield %0 : ... /// linalg.yield %0 : ...
/// } -> ... /// } -> ...
/// %op1 = linalg.generic ... ins(%arg2, %op0#2 : ...) /// %op1 = linalg.generic ... ins(%arg2, %op0#2 : ...)
/// outs(%init1 : ...) { /// outs(%init1 : ...) {
▲ Show 20 Lines • Show All 104 Lines • ▼ Show 20 Lines if (resultNumber) {
newResultTypes.push_back(result.getType()); newResultTypes.push_back(result.getType());
peeledGenericOpIndexingMaps.push_back( peeledGenericOpIndexingMaps.push_back(
genericOp.getTiedIndexingMapForResult(result)); genericOp.getTiedIndexingMapForResult(result));
continue; continue;
} }
// Fall back path, use an `init_tensor` and identity indexing map. // Fall back path, use an `init_tensor` and identity indexing map.
AffineMap indexingMap = rewriter.getMultiDimIdentityMap(domain.size()); AffineMap indexingMap = rewriter.getMultiDimIdentityMap(domain.size());
Value initTensor = rewriter.create<linalg::InitTensorOp>( Value emptyTensor =
loc, domain, scalarOpResult.getType()); rewriter.create<tensor::EmptyOp>(loc, domain, scalarOpResult.getType());
newInitValues.push_back(initTensor); newInitValues.push_back(emptyTensor);
newResultTypes.push_back(initTensor.getType()); newResultTypes.push_back(emptyTensor.getType());
peeledGenericOpIndexingMaps.push_back(indexingMap); peeledGenericOpIndexingMaps.push_back(indexingMap);
} }
/// Create the peeled generic op with an empty body. /// Create the peeled generic op with an empty body.
SmallVector<Value> outsOperands = genericOp.getOutputOperands(); SmallVector<Value> outsOperands = genericOp.getOutputOperands();
outsOperands.append(newInitValues.begin(), newInitValues.end()); outsOperands.append(newInitValues.begin(), newInitValues.end());
SmallVector<Type> resultTypes = llvm::to_vector(genericOp.getResultTypes()); SmallVector<Type> resultTypes = llvm::to_vector(genericOp.getResultTypes());
resultTypes.append(newResultTypes.begin(), newResultTypes.end()); resultTypes.append(newResultTypes.begin(), newResultTypes.end());
▲ Show 20 Lines • Show All 181 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp

Show First 20 Lines • Show All 182 Lines • ▼ Show 20 Lines public:
/// BBs), we need to maintain the sub-set of arguments that should be /// BBs), we need to maintain the sub-set of arguments that should be
/// detensored (i.e. converted by typeConverter) for each affected BB. /// detensored (i.e. converted by typeConverter) for each affected BB.
/// ///
/// Example: /// Example:
/// ///
/// For the following snippet: /// For the following snippet:
/// ... /// ...
/// ^bb1(%6: tensor<i32>, %9: tensor<i32>): /// ^bb1(%6: tensor<i32>, %9: tensor<i32>):
/// %7 = linalg.init_tensor [] : tensor<i32> /// %7 = tensor.empty() : tensor<i32>
/// %8 = linalg.generic #attrs /// %8 = linalg.generic #attrs
/// ins(%6, %6 : tensor<i32>, tensor<i32>) /// ins(%6, %6 : tensor<i32>, tensor<i32>)
/// outs(%7 : tensor<i32>) { /// outs(%7 : tensor<i32>) {
/// ^bb0(%arg0: i32, %arg1: i32, %arg2: i32): /// ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
/// %9 = arith.addi %arg0, %arg1 : i32 /// %9 = arith.addi %arg0, %arg1 : i32
/// linalg.yield %9 : i32 /// linalg.yield %9 : i32
/// } -> tensor<i32> /// } -> tensor<i32>
/// %10 = "some.op"(%9) /// %10 = "some.op"(%9)
▲ Show 20 Lines • Show All 364 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp

Show First 20 Lines • Show All 534 Lines • ▼ Show 20 Lines
void mlir::linalg::populateFoldUnitExtentDimsPatterns( void mlir::linalg::populateFoldUnitExtentDimsPatterns(
RewritePatternSet &patterns) { RewritePatternSet &patterns) {
auto *context = patterns.getContext(); auto *context = patterns.getContext();
patterns.add<FoldUnitDimLoops, ReplaceUnitExtents, RankReducedExtractSliceOp, patterns.add<FoldUnitDimLoops, ReplaceUnitExtents, RankReducedExtractSliceOp,
RankReducedInsertSliceOp<tensor::InsertSliceOp>, RankReducedInsertSliceOp<tensor::InsertSliceOp>,
RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>( RankReducedInsertSliceOp<tensor::ParallelInsertSliceOp>>(
context); context);
linalg::FillOp::getCanonicalizationPatterns(patterns, context); linalg::FillOp::getCanonicalizationPatterns(patterns, context);
linalg::InitTensorOp::getCanonicalizationPatterns(patterns, context);
tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context); tensor::CollapseShapeOp::getCanonicalizationPatterns(patterns, context);
tensor::EmptyOp::getCanonicalizationPatterns(patterns, context);
tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context); tensor::ExpandShapeOp::getCanonicalizationPatterns(patterns, context);
} }
namespace { namespace {
/// Pass that removes unit-extent dims within generic ops. /// Pass that removes unit-extent dims within generic ops.
struct LinalgFoldUnitExtentDimsPass struct LinalgFoldUnitExtentDimsPass
: public impl::LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> { : public impl::LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
void runOnOperation() override { void runOnOperation() override {
Show All 15 Lines

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

Show First 20 Lines • Show All 998 Lines • ▼ Show 20 Lines
// operation. If all dimensions created by expansion can be collapsed in the // operation. If all dimensions created by expansion can be collapsed in the
// iteration space then the reshape is defunct. // iteration space then the reshape is defunct.
// //
// Example: // Example:
// //
// ```mlir // ```mlir
// #map = affine_map<(d0, d1) -> (d0, d1)> // #map = affine_map<(d0, d1) -> (d0, d1)>
// %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32> // %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>
// %2 = linalg.init_tensor [..] : tensor<?x4xf32> // %2 = tensor.empty [..] : tensor<?x4xf32>
// %3 = linalg.generic { // %3 = linalg.generic {
// indexing_maps = [#map, #map], // indexing_maps = [#map, #map],
// iterator_types = ["parallel" ,"parallel"]} // iterator_types = ["parallel" ,"parallel"]}
// ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<?x4xf32>) {.. } // ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<?x4xf32>) {.. }
// ``` // ```
// //
// can be fused by collapsing the dimensions of the iteration space. // can be fused by collapsing the dimensions of the iteration space.
// //
// ```mlir // ```mlir
// #map = affine_map<(d0) -> (d0)> // #map = affine_map<(d0) -> (d0)>
// %2 = linalg.init_tensor [..] : tensor<?xf32> // %2 = tensor.empty [..] : tensor<?xf32>
// %3 = linalg.generic { // %3 = linalg.generic {
// indexing_maps = [#map, #map], // indexing_maps = [#map, #map],
// iterator_types = ["parallel"]} // iterator_types = ["parallel"]}
// ins(%1 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {.. } // ins(%1 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {.. }
// %4 = tensor.expand_shape %3 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32> // %4 = tensor.expand_shape %3 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>
// ``` // ```
// //
// In the following example, // In the following example,
// //
// ```mlir // ```mlir
// #map0 = affine_map<(d0, d1) -> (d0, d1)> // #map0 = affine_map<(d0, d1) -> (d0, d1)>
// #map1 = affine_map<(d0, d1) -> (d1, d0)> // #map1 = affine_map<(d0, d1) -> (d1, d0)>
// %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32> // %1 = tensor.expand_shape %0 [[0, 1]] : tensor<?xf32> into tensor<?x4xf32>
// %2 = linalg.init_tensor [..] : tensor<4x?xf32> // %2 = tensor.empty [..] : tensor<4x?xf32>
// %2 = linalg.generic { // %2 = linalg.generic {
// indexing_maps = [#map0, #map1], // indexing_maps = [#map0, #map1],
// iterator_types = ["parallel" ,"parallel"]} // iterator_types = ["parallel" ,"parallel"]}
// ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<4x?xf32>) {.. } // ins(%1 : tensor<?x4xf32>) outs(%2 : tensor<4x?xf32>) {.. }
// ``` // ```
// //
// the reshape cannot be fused with the generic op by collapsing the op // the reshape cannot be fused with the generic op by collapsing the op
// dimensions since the indexing maps will have to contain mods and divs // dimensions since the indexing maps will have to contain mods and divs
▲ Show 20 Lines • Show All 599 Lines • ▼ Show 20 Lines
} // namespace } // namespace
//===---------------------------------------------------------------------===// //===---------------------------------------------------------------------===//
// Miscellaneous patterns that help fusion. // Miscellaneous patterns that help fusion.
//===---------------------------------------------------------------------===// //===---------------------------------------------------------------------===//
namespace { namespace {
/// Forces `outs` operands of linalg operations to use `linalg.init_tensor` if /// Forces `outs` operands of linalg operations to use `tensor.empty` if the
/// the value of the `outs` operand is not used within the op. This is only /// value of the `outs` operand is not used within the op. This is only
/// implemented for `linalg.generic` operations for now, but should hold for all /// implemented for `linalg.generic` operations for now, but should hold for all
/// linalg structured ops. /// linalg structured ops.
struct RemoveOutsDependency : public OpRewritePattern<GenericOp> { struct RemoveOutsDependency : public OpRewritePattern<GenericOp> {
using OpRewritePattern<GenericOp>::OpRewritePattern; using OpRewritePattern<GenericOp>::OpRewritePattern;
LogicalResult matchAndRewrite(GenericOp op, LogicalResult matchAndRewrite(GenericOp op,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
rewriter.startRootUpdate(op); rewriter.startRootUpdate(op);
bool modifiedOutput = false; bool modifiedOutput = false;
Location loc = op.getLoc(); Location loc = op.getLoc();
for (OpOperand *opOperand : op.getOutputOperands()) { for (OpOperand *opOperand : op.getOutputOperands()) {
if (!op.payloadUsesValueFromOperand(opOperand)) { if (!op.payloadUsesValueFromOperand(opOperand)) {
Value operandVal = opOperand->get(); Value operandVal = opOperand->get();
auto operandType = operandVal.getType().dyn_cast<RankedTensorType>(); auto operandType = operandVal.getType().dyn_cast<RankedTensorType>();
if (!operandType) if (!operandType)
continue; continue;
// If outs is sparse, leave it to the sparse compiler. // If outs is sparse, leave it to the sparse compiler.
if (sparse_tensor::getSparseTensorEncoding(operandVal.getType())) if (sparse_tensor::getSparseTensorEncoding(operandVal.getType()))
continue; continue;
// If outs is already an `init_tensor` operation, nothing to do. // If outs is already an `empty` operation, nothing to do.
auto definingOp = operandVal.getDefiningOp<InitTensorOp>(); auto definingOp = operandVal.getDefiningOp<tensor::EmptyOp>();
if (definingOp) if (definingOp)
continue; continue;
modifiedOutput = true; modifiedOutput = true;
SmallVector<Value> dynamicDims; SmallVector<Value> dynamicDims;
for (const auto &dim : llvm::enumerate(operandType.getShape())) { for (const auto &dim : llvm::enumerate(operandType.getShape())) {
if (dim.value() != ShapedType::kDynamicSize) if (dim.value() != ShapedType::kDynamicSize)
continue; continue;
dynamicDims.push_back(rewriter.createOrFold<tensor::DimOp>( dynamicDims.push_back(rewriter.createOrFold<tensor::DimOp>(
loc, operandVal, dim.index())); loc, operandVal, dim.index()));
} }
Value initTensor = rewriter.create<InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, dynamicDims, operandType.getShape(), loc, operandType.getShape(), operandType.getElementType(),
operandType.getElementType()); dynamicDims);
op->setOperand(opOperand->getOperandNumber(), initTensor); op->setOperand(opOperand->getOperandNumber(), emptyTensor);
} }
} }
if (!modifiedOutput) { if (!modifiedOutput) {
rewriter.cancelRootUpdate(op); rewriter.cancelRootUpdate(op);
return failure(); return failure();
} }
rewriter.finalizeRootUpdate(op); rewriter.finalizeRootUpdate(op);
return success(); return success();
▲ Show 20 Lines • Show All 107 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp

Show All 31 Lines
} }
/// Given `op` assumed `isElementwiseMappableOpOnRankedTensors`, iterate over /// Given `op` assumed `isElementwiseMappableOpOnRankedTensors`, iterate over
/// the result types and return a list of values such that, for each result type /// the result types and return a list of values such that, for each result type
/// `t` and value `v` at the same index `idx`: /// `t` and value `v` at the same index `idx`:
/// 1. `v.getType() == t` /// 1. `v.getType() == t`
/// 2. If an operand of `op` has type `t`, let `operand_first` be the first /// 2. If an operand of `op` has type `t`, let `operand_first` be the first
/// such operand. Then`v == operand_first`. /// such operand. Then`v == operand_first`.
/// 3. Otherwise, v is a newly created `linalg::InitTensorOp` with: /// 3. Otherwise, v is a newly created `tensor::EmptyOp` with:
/// a. Static and dynamic dims extracted from the first operand of `op`. /// a. Static and dynamic dims extracted from the first operand of `op`.
/// b. Elemental type equal to the elemental type of `t`. /// b. Elemental type equal to the elemental type of `t`.
/// ///
/// This is sufficient because ElementwiseMappable guarantees that "The static /// This is sufficient because ElementwiseMappable guarantees that "The static
/// types of all vector (resp. tensor) operands and results must have the same /// types of all vector (resp. tensor) operands and results must have the same
/// shape". /// shape".
static SmallVector<Value, 4> static SmallVector<Value, 4>
getOrCreateOperandsMatchingResultTypes(OpBuilder &b, Operation *op) { getOrCreateOperandsMatchingResultTypes(OpBuilder &b, Operation *op) {
Show All 17 Lines if (found)
continue; continue;
// Extract static / dynamic shape mix from the first operand. // Extract static / dynamic shape mix from the first operand.
Value firstOperand = operands.front(); Value firstOperand = operands.front();
auto rankedTensorType = t.cast<RankedTensorType>(); auto rankedTensorType = t.cast<RankedTensorType>();
auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape()); auto staticShape = llvm::to_vector<4>(rankedTensorType.getShape());
auto dynamicShape = linalg::getDynOperands(loc, firstOperand, b); auto dynamicShape = linalg::getDynOperands(loc, firstOperand, b);
res.push_back(b.create<linalg::InitTensorOp>( res.push_back(b.create<tensor::EmptyOp>(
loc, dynamicShape, staticShape, rankedTensorType.getElementType())); loc, staticShape, rankedTensorType.getElementType(), dynamicShape));
} }
return res; return res;
} }
namespace { namespace {
struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern { struct ConvertAnyElementwiseMappableOpOnRankedTensors : public RewritePattern {
ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context) ConvertAnyElementwiseMappableOpOnRankedTensors(MLIRContext *context)
: RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {} : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}
▲ Show 20 Lines • Show All 67 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp

Show First 20 Lines • Show All 70 Lines • ▼ Show 20 Lines if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
padOp, "failed to get shape of pad op result"); padOp, "failed to get shape of pad op result");
} }
Location loc = padOp.getLoc(); Location loc = padOp.getLoc();
// Create the tensor of same size as output of the pad op. // Create the tensor of same size as output of the pad op.
RankedTensorType padResultType = padOp.getResultType(); RankedTensorType padResultType = padOp.getResultType();
auto resultSizes = getAsOpFoldResult(resultShape[0]); auto resultSizes = getAsOpFoldResult(resultShape[0]);
auto initTensor = rewriter.create<linalg::InitTensorOp>( auto emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, resultSizes, padResultType.getElementType()); loc, resultSizes, padResultType.getElementType());
// Fill the tensor with the pad value. // Fill the tensor with the pad value.
// TODO: There is an option to fill only the boundaries. For now just // TODO: There is an option to fill only the boundaries. For now just
// filling the whole tensor. // filling the whole tensor.
auto fillTensor = auto fillTensor =
rewriter.create<linalg::FillOp>(loc, padValue, initTensor.getResult()); rewriter.create<linalg::FillOp>(loc, padValue, emptyTensor.getResult());
// Construct a slice of the fill result that is to be replaced with the // Construct a slice of the fill result that is to be replaced with the
// result of the generic op. The low pad values are the offsets, the size of // result of the generic op. The low pad values are the offsets, the size of
// the source is the size of the slice. // the source is the size of the slice.
// TODO: This insert/extract could be potentially made a utility method. // TODO: This insert/extract could be potentially made a utility method.
unsigned resultNumber = source.cast<OpResult>().getResultNumber(); unsigned resultNumber = source.cast<OpResult>().getResultNumber();
SmallVector<OpFoldResult> offsets = padOp.getMixedLowPad(); SmallVector<OpFoldResult> offsets = padOp.getMixedLowPad();
SmallVector<OpFoldResult> sizes; SmallVector<OpFoldResult> sizes;
Show All 33 Lines

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

Show First 20 Lines • Show All 425 Lines • ▼ Show 20 LinesFailureOr<Value> mlir::linalg::hoistPaddingOnTensors(
// Create the packed tensor<?x?x..?xtransposedShape> into which we amortize // Create the packed tensor<?x?x..?xtransposedShape> into which we amortize
// padding. // padding.
SmallVector<int64_t> packedShape(nPackedLoops, ShapedType::kDynamicSize); SmallVector<int64_t> packedShape(nPackedLoops, ShapedType::kDynamicSize);
// TODO: go grab dims when necessary, for now tensor::PadOp returns a static // TODO: go grab dims when necessary, for now tensor::PadOp returns a static
// tensor. // tensor.
llvm::append_range(packedShape, transposedTensorType->getShape()); llvm::append_range(packedShape, transposedTensorType->getShape());
auto packedTensorType = RankedTensorType::get( auto packedTensorType = RankedTensorType::get(
packedShape, transposedTensorType->getElementType()); packedShape, transposedTensorType->getElementType());
Value packedTensor = b.create<linalg::InitTensorOp>( Value packedTensor = b.create<tensor::EmptyOp>(
loc, dynamicTensorSizes, packedTensorType.getShape(), loc, packedTensorType.getShape(), packedTensorType.getElementType(),
packedTensorType.getElementType()); dynamicTensorSizes);
// Clone the operations involved in the backward slice, iteratively stepping // Clone the operations involved in the backward slice, iteratively stepping
// into the loops that we encounter. // into the loops that we encounter.
// The implementation proceeds in a stack-like fashion: // The implementation proceeds in a stack-like fashion:
// 1. Iteratively clone and step into the loops, pushing the `packedTensor` // 1. Iteratively clone and step into the loops, pushing the `packedTensor`
// deeper in the stack. // deeper in the stack.
// 2. Create a GenericOp if `transposeVector` is non-empty. // 2. Create a GenericOp if `transposeVector` is non-empty.
// 3. Create a InsertSliceOp at the top of the stack. // 3. Create a InsertSliceOp at the top of the stack.
▲ Show 20 Lines • Show All 93 Lines • ▼ Show 20 LinesFailureOr<Value> mlir::linalg::hoistPaddingOnTensors(
// strides = [1 .. 1] (defined above) // strides = [1 .. 1] (defined above)
packedTensor = packedTensor =
scf::getForInductionVarOwner(clonedLoopIvs.front())->getResult(0); scf::getForInductionVarOwner(clonedLoopIvs.front())->getResult(0);
Value newResult = b.create<tensor::ExtractSliceOp>( Value newResult = b.create<tensor::ExtractSliceOp>(
loc, *transposedTensorType, packedTensor, offsets, sizes, strides); loc, *transposedTensorType, packedTensor, offsets, sizes, strides);
// Transpose the packed tensor back to the original storage order. // Transpose the packed tensor back to the original storage order.
if (!transposeVector.empty()) { if (!transposeVector.empty()) {
Value initTensor = Value emptyTensor = b.create<tensor::EmptyOp>(
b.create<InitTensorOp>(loc, ValueRange{}, paddedTensorType.getShape(), loc, paddedTensorType.getShape(), paddedTensorType.getElementType());
paddedTensorType.getElementType());
transposeOps.push_back( transposeOps.push_back(
makeTransposeOp(b, loc, newResult, initTensor, transposeVector)); makeTransposeOp(b, loc, newResult, emptyTensor, transposeVector));
newResult = transposeOps.back()->getResult(0); newResult = transposeOps.back()->getResult(0);
} }
// Make the newly cloned `opToHoist` available to the caller. // Make the newly cloned `opToHoist` available to the caller.
hoistedOp = hoistedOp =
cast<tensor::PadOp>(bvm.lookup(opToHoist.getResult()).getDefiningOp()); cast<tensor::PadOp>(bvm.lookup(opToHoist.getResult()).getDefiningOp());
return newResult; return newResult;
} }

mlir/lib/Dialect/Linalg/Transforms/InitTensorToAllocTensor.cpp

//===- InitTensorToAllocTensor.cpp - Lower init_tensor to alloc_tensor ----===// //===- InitTensorToAllocTensor.cpp - Lower init_tensor to alloc_tensor ----===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Pass/Pass.h" #include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
namespace mlir { namespace mlir {
#define GEN_PASS_DEF_LINALGINITTENSORTOALLOCTENSOR #define GEN_PASS_DEF_EMPTYTENSORTOALLOCTENSOR
#include "mlir/Dialect/Linalg/Passes.h.inc" #include "mlir/Dialect/Linalg/Passes.h.inc"
} // namespace mlir } // namespace mlir
using namespace mlir; using namespace mlir;
using namespace mlir::bufferization; using namespace mlir::bufferization;
using namespace mlir::linalg; using namespace mlir::tensor;
namespace { namespace {
struct InitTensorLoweringPattern : public OpRewritePattern<InitTensorOp> { struct EmptyTensorLoweringPattern : public OpRewritePattern<tensor::EmptyOp> {
using OpRewritePattern<InitTensorOp>::OpRewritePattern; using OpRewritePattern<tensor::EmptyOp>::OpRewritePattern;
LogicalResult matchAndRewrite(InitTensorOp op, LogicalResult matchAndRewrite(tensor::EmptyOp op,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(op, op.getType(), rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
op.getSizes()); op, op.getType(), op.getDynamicSizes());
return success(); return success();
} }
}; };
struct LinalgInitTensorToAllocTensor struct EmptyTensorToAllocTensor
: public impl::LinalgInitTensorToAllocTensorBase< : public impl::EmptyTensorToAllocTensorBase<EmptyTensorToAllocTensor> {
LinalgInitTensorToAllocTensor> { EmptyTensorToAllocTensor() = default;
LinalgInitTensorToAllocTensor() = default;
void runOnOperation() override; void runOnOperation() override;
void getDependentDialects(DialectRegistry &registry) const override { void getDependentDialects(DialectRegistry &registry) const override {
registry registry
.insert<linalg::LinalgDialect, bufferization::BufferizationDialect>(); .insert<tensor::TensorDialect, bufferization::BufferizationDialect>();
} }
}; };
} // namespace } // namespace
void LinalgInitTensorToAllocTensor::runOnOperation() { void EmptyTensorToAllocTensor::runOnOperation() {
Operation *op = getOperation(); Operation *op = getOperation();
RewritePatternSet patterns(op->getContext()); RewritePatternSet patterns(op->getContext());
patterns.insert<InitTensorLoweringPattern>(op->getContext()); patterns.insert<EmptyTensorLoweringPattern>(op->getContext());
if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns)))) if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
signalPassFailure(); signalPassFailure();
} }
std::unique_ptr<Pass> mlir::createLinalgInitTensorToAllocTensorPass() { std::unique_ptr<Pass> mlir::createEmptyTensorToAllocTensorPass() {
return std::make_unique<LinalgInitTensorToAllocTensor>(); return std::make_unique<EmptyTensorToAllocTensor>();
} }

mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp

Show First 20 Lines • Show All 190 Lines • ▼ Show 20 Lines for (unsigned idx :
if (control.innerParallel) { if (control.innerParallel) {
outputExpr.push_back( outputExpr.push_back(
b.getAffineDimExpr(dim <= reductionDim ? dim : dim + 1)); b.getAffineDimExpr(dim <= reductionDim ? dim : dim + 1));
} else { } else {
outputExpr.push_back( outputExpr.push_back(
b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1)); b.getAffineDimExpr(dim < insertSplitDimension ? dim : dim + 1));
} }
} }
Value initOrAllocTensor; Value emptyOrAllocTensor;
if (useAlloc) { if (useAlloc) {
initOrAllocTensor = b.create<bufferization::AllocTensorOp>( emptyOrAllocTensor = b.create<bufferization::AllocTensorOp>(
loc, loc,
RankedTensorType::get(newOutputShape, RankedTensorType::get(newOutputShape,
op.getRegionOutputArgs()[0].getType()), op.getRegionOutputArgs()[0].getType()),
ValueRange{}); ValueRange{});
} else { } else {
initOrAllocTensor = b.create<linalg::InitTensorOp>( emptyOrAllocTensor = b.create<tensor::EmptyOp>(
loc, newOutputShape, op.getRegionOutputArgs()[0].getType()); loc, newOutputShape, op.getRegionOutputArgs()[0].getType());
} }
Value constantOp = b.create<arith::ConstantOp>(loc, identity); Value constantOp = b.create<arith::ConstantOp>(loc, identity);
Value identityTensor = Value identityTensor =
b.create<linalg::FillOp>(op->getLoc(), constantOp, initOrAllocTensor) b.create<linalg::FillOp>(op->getLoc(), constantOp, emptyOrAllocTensor)
.getResult(0); .getResult(0);
newMaps.push_back(AffineMap::get(oldOutputMap.getNumDims() + 1, 0, outputExpr, newMaps.push_back(AffineMap::get(oldOutputMap.getNumDims() + 1, 0, outputExpr,
op.getContext())); op.getContext()));
SmallVector<StringRef> newIteratorTypes; SmallVector<StringRef> newIteratorTypes;
for (auto &it : llvm::enumerate(op.iterator_types())) { for (auto &it : llvm::enumerate(op.iterator_types())) {
if (insertSplitDimension == it.index() && !control.innerParallel) if (insertSplitDimension == it.index() && !control.innerParallel)
newIteratorTypes.push_back(getParallelIteratorTypeName()); newIteratorTypes.push_back(getParallelIteratorTypeName());
newIteratorTypes.push_back(it.value().cast<StringAttr>().getValue()); newIteratorTypes.push_back(it.value().cast<StringAttr>().getValue());
if (insertSplitDimension == it.index() && control.innerParallel) if (insertSplitDimension == it.index() && control.innerParallel)
newIteratorTypes.push_back(getParallelIteratorTypeName()); newIteratorTypes.push_back(getParallelIteratorTypeName());
} }
// Create the new op matching the original op with an extra parallel // Create the new op matching the original op with an extra parallel
// dimension. // dimension.
GenericOp genericOp = b.create<GenericOp>( GenericOp genericOp = b.create<GenericOp>(
loc, TypeRange({initOrAllocTensor.getType()}), newInputs, loc, TypeRange({emptyOrAllocTensor.getType()}), newInputs,
ValueRange({identityTensor}), newMaps, newIteratorTypes); ValueRange({identityTensor}), newMaps, newIteratorTypes);
b.inlineRegionBefore(op->getRegion(0), genericOp.getRegion(), b.inlineRegionBefore(op->getRegion(0), genericOp.getRegion(),
genericOp.getRegion().begin()); genericOp.getRegion().begin());
// Then create a new reduction that only reduce the newly added dimension // Then create a new reduction that only reduce the newly added dimension
// from the previous op. // from the previous op.
unsigned intermRank = newOutputShape.size(); unsigned intermRank = newOutputShape.size();
AffineMap inputMap = b.getMultiDimIdentityMap(intermRank); AffineMap inputMap = b.getMultiDimIdentityMap(intermRank);
Show All 17 Lines auto reduction = b.create<GenericOp>(
[reductionOp](OpBuilder &b, Location loc, ValueRange inputs) { [reductionOp](OpBuilder &b, Location loc, ValueRange inputs) {
Operation *clonedReductionOp = b.clone(*reductionOp); Operation *clonedReductionOp = b.clone(*reductionOp);
clonedReductionOp->setOperand(0, inputs[0]); clonedReductionOp->setOperand(0, inputs[0]);
clonedReductionOp->setOperand(1, inputs[1]); clonedReductionOp->setOperand(1, inputs[1]);
b.create<linalg::YieldOp>(loc, clonedReductionOp->getResult(0)); b.create<linalg::YieldOp>(loc, clonedReductionOp->getResult(0));
}); });
b.replaceOp(op, reduction.getResults()); b.replaceOp(op, reduction.getResults());
return SplitReductionResult{ return SplitReductionResult{emptyOrAllocTensor.getDefiningOp(),
initOrAllocTensor.getDefiningOp(), identityTensor.getDefiningOp<FillOp>(), identityTensor.getDefiningOp<FillOp>(),
cast<LinalgOp>(genericOp.getOperation()), reduction}; cast<LinalgOp>(genericOp.getOperation()),
reduction};
} }
/// Rewrite f(i, j, k, ...) into f(i, j, k * ratio + kk, ...) /// Rewrite f(i, j, k, ...) into f(i, j, k * ratio + kk, ...)
/// TODO: Additional pattern to rewrite f(i, j, k * ratio + kk, ...) into /// TODO: Additional pattern to rewrite f(i, j, k * ratio + kk, ...) into
/// f(i, j, k, kk, ...) with a proper ExpandShapeOp. This is probably better /// f(i, j, k, kk, ...) with a proper ExpandShapeOp. This is probably better
/// done as a transform to enable better vectorization. /// done as a transform to enable better vectorization.
static AffineMap scaleReductionDim(LinalgOp op, OpOperand &opOperand, static AffineMap scaleReductionDim(LinalgOp op, OpOperand &opOperand,
unsigned reductionDimPos, unsigned reductionDimPos,
▲ Show 20 Lines • Show All 79 Lines • ▼ Show 20 LinesFailureOr<SplitReductionResult> mlir::linalg::splitReductionByScaling(
// b. O(i, j) += O_i(kk, i, j) // b. O(i, j) += O_i(kk, i, j)
// The intermediate tensor O_i is of shape (128/16)x3x5 == 8x3x5. // The intermediate tensor O_i is of shape (128/16)x3x5 == 8x3x5.
Location loc = op->getLoc(); Location loc = op->getLoc();
MLIRContext *context = op.getContext(); MLIRContext *context = op.getContext();
// For now assume outputs are 1-1 with reduction neutralElements. // For now assume outputs are 1-1 with reduction neutralElements.
// TODO: generalize when multi-reduction support is available. // TODO: generalize when multi-reduction support is available.
SmallVector<Value> newOutputs; SmallVector<Value> newOutputs;
newOutputs.reserve(op.getNumOutputs()); newOutputs.reserve(op.getNumOutputs());
SmallVector<Operation *> initOrAllocTensorOps; SmallVector<Operation *> emptyOrAllocTensorOps;
SmallVector<linalg::FillOp> fillOps; SmallVector<linalg::FillOp> fillOps;
fillOps.reserve(op.getNumOutputs()); fillOps.reserve(op.getNumOutputs());
for (auto it : llvm::zip(op.outputs(), neutralElements)) { for (auto it : llvm::zip(op.outputs(), neutralElements)) {
Value rankedTensor = std::get<0>(it); Value rankedTensor = std::get<0>(it);
auto t = rankedTensor.getType().cast<RankedTensorType>(); auto t = rankedTensor.getType().cast<RankedTensorType>();
RankedTensorType newT = RankedTensorType::Builder(t).insertDim( RankedTensorType newT = RankedTensorType::Builder(t).insertDim(
reductionDimSize / splitFactor, insertSplitDimension); reductionDimSize / splitFactor, insertSplitDimension);
SmallVector<Value> dims = SmallVector<Value> dims =
tensor::createDynamicDimValues(b, loc, rankedTensor); tensor::createDynamicDimValues(b, loc, rankedTensor);
Value initOrAllocTensor; Value emptyOrAllocTensor;
if (useAlloc) { if (useAlloc) {
initOrAllocTensor = emptyOrAllocTensor =
b.create<bufferization::AllocTensorOp>(loc, newT, dims); b.create<bufferization::AllocTensorOp>(loc, newT, dims);
} else { } else {
initOrAllocTensor = b.create<linalg::InitTensorOp>( emptyOrAllocTensor = b.create<tensor::EmptyOp>(loc, newT.getShape(),
loc, dims, newT.getShape(), t.getElementType()); t.getElementType(), dims);
} }
Value constantOp = b.create<arith::ConstantOp>(loc, std::get<1>(it)); Value constantOp = b.create<arith::ConstantOp>(loc, std::get<1>(it));
fillOps.push_back( fillOps.push_back(
b.create<linalg::FillOp>(op->getLoc(), constantOp, initOrAllocTensor)); b.create<linalg::FillOp>(op->getLoc(), constantOp, emptyOrAllocTensor));
newOutputs.push_back(fillOps.back().getResult(0)); newOutputs.push_back(fillOps.back().getResult(0));
initOrAllocTensorOps.push_back(initOrAllocTensor.getDefiningOp()); emptyOrAllocTensorOps.push_back(emptyOrAllocTensor.getDefiningOp());
} }
// Step 2. Reindex / expand indexing maps. // Step 2. Reindex / expand indexing maps.
// Reindex existing input indexings: k -> k * splitFactor + k'. // Reindex existing input indexings: k -> k * splitFactor + k'.
SmallVector<AffineMap> newMaps; SmallVector<AffineMap> newMaps;
newMaps.reserve(op.getNumInputsAndOutputs() + 1); newMaps.reserve(op.getNumInputsAndOutputs() + 1);
for (OpOperand *o : op.getInputOperands()) for (OpOperand *o : op.getInputOperands())
newMaps.push_back(scaleReductionDim(op, *o, reductionDimPos, splitFactor)); newMaps.push_back(scaleReductionDim(op, *o, reductionDimPos, splitFactor));
Show All 10 Lines for (OpOperand *o : op.getOutputOperands())
newMaps.push_back(insertParallelDim(op, *o, reductionDimPos, newMaps.push_back(insertParallelDim(op, *o, reductionDimPos,
reductionDimSize / splitFactor)); reductionDimSize / splitFactor));
// Step 3. Handle operands. // Step 3. Handle operands.
// Compute the new input tensors. // Compute the new input tensors.
auto newInputs = llvm::to_vector<4>(op.inputs()); auto newInputs = llvm::to_vector<4>(op.inputs());
// Add a single shape-only tensor to carry the dimensions without resorting to // Add a single shape-only tensor to carry the dimensions without resorting to
// more complex inversions. // more complex inversions.
newInputs.push_back(b.create<linalg::InitTensorOp>( newInputs.push_back(b.create<tensor::EmptyOp>(
loc, ArrayRef<int64_t>{reductionDimSize / splitFactor, splitFactor}, loc, ArrayRef<int64_t>{reductionDimSize / splitFactor, splitFactor},
b.getIntegerType(1))); b.getIntegerType(1)));
// Output tensors are already good to go. // Output tensors are already good to go.
// Step 4. Create the new op matching the original op with an extra parallel // Step 4. Create the new op matching the original op with an extra parallel
// dimension. // dimension.
SmallVector<StringRef> iteratorTypes = SmallVector<StringRef> iteratorTypes =
llvm::to_vector<4>(op.getIteratorTypes().getAsValueRange<StringAttr>()); llvm::to_vector<4>(op.getIteratorTypes().getAsValueRange<StringAttr>());
▲ Show 20 Lines • Show All 47 Lines • ▼ Show 20 Lines for (auto it :
// clang-format on // clang-format on
results.push_back(reductionOp); results.push_back(reductionOp);
} }
// TODO: extend when multi-reduction support is available. // TODO: extend when multi-reduction support is available.
assert(fillOps.size() == results.size() && results.size() == 1); assert(fillOps.size() == results.size() && results.size() == 1);
b.replaceOp(op, results.front()->getResults()); b.replaceOp(op, results.front()->getResults());
return SplitReductionResult{initOrAllocTensorOps.front(), fillOps.front(), return SplitReductionResult{emptyOrAllocTensorOps.front(), fillOps.front(),
cast<LinalgOp>(genericOp.getOperation()), cast<LinalgOp>(genericOp.getOperation()),
results.front()}; results.front()};
} }
namespace { namespace {
struct LinalgSplitReduction : public OpInterfaceRewritePattern<LinalgOp> { struct LinalgSplitReduction : public OpInterfaceRewritePattern<LinalgOp> {
/// Construct a generic pattern applied to all LinalgOp that verify `filter`. /// Construct a generic pattern applied to all LinalgOp that verify `filter`.
Show All 29 Lines

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

Show First 20 Lines • Show All 715 Lines • ▼ Show 20 Linesvoid mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx); memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
memref::ViewOp::getCanonicalizationPatterns(patterns, ctx); memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
scf::ForOp::getCanonicalizationPatterns(patterns, ctx); scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx); scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
tensor::CastOp::getCanonicalizationPatterns(patterns, ctx); tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
tensor::EmptyOp::getCanonicalizationPatterns(patterns, ctx);
tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx); tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx); tensor::InsertSliceOp::getCanonicalizationPatterns(patterns, ctx);
InitTensorOp::getCanonicalizationPatterns(patterns, ctx);
tensor::PadOp::getCanonicalizationPatterns(patterns, ctx); tensor::PadOp::getCanonicalizationPatterns(patterns, ctx);
ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns); ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
CanonicalizationPatternList< CanonicalizationPatternList<
#define GET_OP_LIST #define GET_OP_LIST
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc" #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
>::insert(patterns); >::insert(patterns);
} }
▲ Show 20 Lines • Show All 75 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

Show First 20 Lines • Show All 624 Lines • ▼ Show 20 LinesLogicalResult mlir::linalg::applyStagedPatterns(
} }
return success(); return success();
} }
static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) { static SmallVector<StringRef> getNParallelLoopsAttrs(unsigned nParallelLoops) {
return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName()); return SmallVector<StringRef>(nParallelLoops, getParallelIteratorTypeName());
} }
/// Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp (to /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp (to
/// initialize with pad_val) and GenericOp (to copy contents). /// initialize with pad_val) and GenericOp (to copy contents).
LogicalResult LogicalResult
PadOpTransformationPattern::matchAndRewrite(tensor::PadOp padOp, PadOpTransformationPattern::matchAndRewrite(tensor::PadOp padOp,
PatternRewriter &rewriter) const { PatternRewriter &rewriter) const {
auto inputShapedType = padOp.getSource().getType().cast<ShapedType>(); auto inputShapedType = padOp.getSource().getType().cast<ShapedType>();
auto resultShapedType = padOp.getResult().getType().cast<ShapedType>(); auto resultShapedType = padOp.getResult().getType().cast<ShapedType>();
Show All 14 Lines if (definingOp && definingOp->getBlock() == &block)
return failure(); return failure();
if (!definingOp && padValue.cast<BlockArgument>().getOwner() == &block) if (!definingOp && padValue.cast<BlockArgument>().getOwner() == &block)
return failure(); return failure();
// Create tensor with the padded shape // Create tensor with the padded shape
Location loc = padOp.getLoc(); Location loc = padOp.getLoc();
SmallVector<Value> indices(resultShapedType.getRank(), SmallVector<Value> indices(resultShapedType.getRank(),
rewriter.create<arith::ConstantIndexOp>(loc, 0)); rewriter.create<arith::ConstantIndexOp>(loc, 0));
Value initTensor = rewriter.create<InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
loc, resultShapedType.getShape(), resultShapedType.getElementType()); loc, resultShapedType.getShape(), resultShapedType.getElementType());
// Initialize tensor with the pad value // Initialize tensor with the pad value
Value tmpTensor = rewriter Value tmpTensor = rewriter
.create<linalg::FillOp>(loc, ValueRange{padValue}, .create<linalg::FillOp>(loc, ValueRange{padValue},
ValueRange{initTensor}) ValueRange{emptyTensor})
.result(); .result();
// Copy original contents into new tensor // Copy original contents into new tensor
// Uses linalg.generic, but could be done with tensor.insert_slice // Uses linalg.generic, but could be done with tensor.insert_slice
SmallVector<AffineExpr, 4> outputExprs; SmallVector<AffineExpr, 4> outputExprs;
for (unsigned i = 0; i < resultShapedType.getRank(); ++i) { for (unsigned i = 0; i < resultShapedType.getRank(); ++i) {
outputExprs.push_back(getAffineDimExpr(i, rewriter.getContext()) + outputExprs.push_back(getAffineDimExpr(i, rewriter.getContext()) +
padOp.getStaticLow()[i].cast<IntegerAttr>().getInt()); padOp.getStaticLow()[i].cast<IntegerAttr>().getInt());
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines if (auto val = ofr.dyn_cast<Value>())
return val; return val;
return rewriter return rewriter
.create<arith::ConstantIndexOp>( .create<arith::ConstantIndexOp>(
padOp.getLoc(), ofr.get<Attribute>().cast<IntegerAttr>().getInt()) padOp.getLoc(), ofr.get<Attribute>().cast<IntegerAttr>().getInt())
.getResult(); .getResult();
}; };
auto resultType = padOp.getResultType(); auto resultType = padOp.getResultType();
// Compute size of InitTensorOp. Any combination of static/dynamic is // Compute size of EmptyOp. Any combination of static/dynamic is supported.
// supported.
SmallVector<Value> dynSizes; SmallVector<Value> dynSizes;
SmallVector<int64_t> staticSizes; SmallVector<int64_t> staticSizes;
for (unsigned dim = 0; dim < resultType.getRank(); ++dim) { for (unsigned dim = 0; dim < resultType.getRank(); ++dim) {
if (resultType.isDynamicDim(dim)) { if (resultType.isDynamicDim(dim)) {
auto srcSize = rewriter.createOrFold<tensor::DimOp>( auto srcSize = rewriter.createOrFold<tensor::DimOp>(
padOp.getLoc(), padOp.getSource(), dim); padOp.getLoc(), padOp.getSource(), dim);
// Add low and high padding value. // Add low and high padding value.
auto plusLow = rewriter.createOrFold<arith::AddIOp>( auto plusLow = rewriter.createOrFold<arith::AddIOp>(
padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim])); padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim]));
auto plusHigh = rewriter.createOrFold<arith::AddIOp>( auto plusHigh = rewriter.createOrFold<arith::AddIOp>(
padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim])); padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim]));
dynSizes.push_back(plusHigh); dynSizes.push_back(plusHigh);
} }
staticSizes.push_back(resultType.getDimSize(dim)); staticSizes.push_back(resultType.getDimSize(dim));
} }
// Init tensor and fill it with padding. // Init tensor and fill it with padding.
Value init = rewriter.create<InitTensorOp>( Value emptyTensor = rewriter.create<tensor::EmptyOp>(
padOp.getLoc(), dynSizes, staticSizes, resultType.getElementType()); padOp.getLoc(), staticSizes, resultType.getElementType(), dynSizes);
Value fill = createFillOrGenerateOp(rewriter, padOp, init, dynSizes); Value fill = createFillOrGenerateOp(rewriter, padOp, emptyTensor, dynSizes);
// Try optimize the copy of source. // Try optimize the copy of source.
if (optimizeCopyFn && optimizeCopyFn(rewriter, padOp, fill).succeeded()) if (optimizeCopyFn && optimizeCopyFn(rewriter, padOp, fill).succeeded())
return success(); return success();
// tensor::PadOps cannot be optimized. Generate a InsertSliceOp instead // tensor::PadOps cannot be optimized. Generate a InsertSliceOp instead
// for copying the PadOp source. // for copying the PadOp source.
auto sourceType = padOp.getSourceType(); auto sourceType = padOp.getSourceType();
▲ Show 20 Lines • Show All 232 Lines • Show Last 20 Lines

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Show First 20 Lines • Show All 680 Lines • ▼ Show 20 Lines for (auto o : ofrs) {
} else { } else {
result.push_back(builder.create<arith::ConstantIndexOp>( result.push_back(builder.create<arith::ConstantIndexOp>(
loc, getIntFromAttr(o.template get<Attribute>()))); loc, getIntFromAttr(o.template get<Attribute>())));
} }
} }
return result; return result;
} }
/// Rewrite a tensor::PadOp into a sequence of InitTensorOp, FillOp and /// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
/// InsertSliceOp. For now, only constant padding values are supported. /// InsertSliceOp. For now, only constant padding values are supported.
/// If there is enough static type information, TransferReadOps and /// If there is enough static type information, TransferReadOps and
/// TransferWriteOps may be generated instead of InsertSliceOps. /// TransferWriteOps may be generated instead of InsertSliceOps.
struct GenericPadOpVectorizationPattern : public GeneralizePadOpPattern { struct GenericPadOpVectorizationPattern : public GeneralizePadOpPattern {
GenericPadOpVectorizationPattern(MLIRContext *context, GenericPadOpVectorizationPattern(MLIRContext *context,
PatternBenefit benefit = 1) PatternBenefit benefit = 1)
: GeneralizePadOpPattern(context, tryVectorizeCopy, benefit) {} : GeneralizePadOpPattern(context, tryVectorizeCopy, benefit) {}
/// Vectorize the copying of a tensor::PadOp's source. This is possible if /// Vectorize the copying of a tensor::PadOp's source. This is possible if
▲ Show 20 Lines • Show All 1,094 Lines • Show Last 20 Lines

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

Show First 20 Lines • Show All 332 Lines • ▼ Show 20 Linesstatic bool computeIterationGraph(Merger &merger, linalg::GenericOp op,
// Report failure for a cyclic iteration graph. // Report failure for a cyclic iteration graph.
topSort.clear(); topSort.clear();
topSort.reserve(n); topSort.reserve(n);
return topSortOptimal(n, iteratorTypes, topSort, inDegree, adjM); return topSortOptimal(n, iteratorTypes, topSort, inDegree, adjM);
} }
/// Returns true if tensor materializes uninitialized into the computation. /// Returns true if tensor materializes uninitialized into the computation.
static bool isMaterializing(Value val) { static bool isMaterializing(Value val) {
return val.getDefiningOp<linalg::InitTensorOp>() || return val.getDefiningOp<tensor::EmptyOp>() ||
val.getDefiningOp<bufferization::AllocTensorOp>(); val.getDefiningOp<bufferization::AllocTensorOp>();
} }
/// Returns true when the tensor expression is admissible for codegen. /// Returns true when the tensor expression is admissible for codegen.
/// Since all sparse input tensors are admissible, we just need to check /// Since all sparse input tensors are admissible, we just need to check
/// whether the out tensor in the tensor expression codegen is admissible. /// whether the out tensor in the tensor expression codegen is admissible.
/// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective /// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective
/// nesting depth when a "truly dynamic" sparse tensor output occurs. /// nesting depth when a "truly dynamic" sparse tensor output occurs.
▲ Show 20 Lines • Show All 1,600 Lines • Show Last 20 Lines

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

Show All 13 Linesadd_mlir_dialect_library(MLIRTensorDialect
DEPENDS DEPENDS
MLIRTensorOpsIncGen MLIRTensorOpsIncGen
LINK_COMPONENTS LINK_COMPONENTS
Core Core
LINK_LIBS PUBLIC LINK_LIBS PUBLIC
MLIRAffineDialect
MLIRArithDialect MLIRArithDialect
MLIRArithUtils MLIRArithUtils
MLIRCastInterfaces MLIRCastInterfaces
MLIRComplexDialect MLIRComplexDialect
MLIRDialectUtils MLIRDialectUtils
MLIRIR MLIRIR
MLIRInferTypeOpInterface MLIRInferTypeOpInterface
MLIRParallelCombiningOpInterface MLIRParallelCombiningOpInterface
Show All 35 Lines

mlir/lib/Dialect/Tensor/IR/TensorDialect.cpp

//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// //
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information. // See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// //
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Complex/IR/Complex.h" #include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/InliningUtils.h" #include "mlir/Transforms/InliningUtils.h"
using namespace mlir; using namespace mlir;
using namespace mlir::tensor; using namespace mlir::tensor;
Show All 31 Lines

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

Show First 20 Lines • Show All 411 Lines • ▼ Show 20 Lines
} // namespace } // namespace
void DimOp::getCanonicalizationPatterns(RewritePatternSet &results, void DimOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) { MLIRContext *context) {
results.add<DimOfCastOp>(context); results.add<DimOfCastOp>(context);
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// EmptyOp
//===----------------------------------------------------------------------===//
void EmptyOp::build(OpBuilder &builder, OperationState &result,
ArrayRef<int64_t> staticShape, Type elementType) {
assert(all_of(staticShape,
[](int64_t sz) { return !ShapedType::isDynamic(sz); }) &&
"expected only static sizes");
build(builder, result, staticShape, elementType, {});
}
void EmptyOp::build(OpBuilder &builder, OperationState &result,
ArrayRef<int64_t> staticShape, Type elementType,
ValueRange dynamicSizes) {
auto tensorType = RankedTensorType::get(staticShape, elementType);
build(builder, result, tensorType, dynamicSizes);
}
void EmptyOp::build(OpBuilder &builder, OperationState &result,
ArrayRef<OpFoldResult> sizes, Type elementType) {
SmallVector<int64_t> staticShape;
SmallVector<Value> dynamicSizes;
dispatchIndexOpFoldResults(sizes, dynamicSizes, staticShape,
ShapedType::kDynamicSize);
build(builder, result, staticShape, elementType, dynamicSizes);
}
LogicalResult EmptyOp::verify() {
if (getType().getNumDynamicDims() != getDynamicSizes().size())
return emitOpError("incorrect number of dynamic sizes, has ")
<< getDynamicSizes().size() << ", expected "
<< getType().getNumDynamicDims();
return success();
}
LogicalResult
EmptyOp::reifyResultShapes(OpBuilder &builder,
ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
unsigned ctr = 0;
for (int64_t i = 0; i < getType().getRank(); ++i) {
if (getType().isDynamicDim(i)) {
reifiedReturnShapes[0][i] = getDynamicSizes()[ctr++];
} else {
reifiedReturnShapes[0][i] =
builder.create<arith::ConstantIndexOp>(getLoc(), i);
}
}
return success();
}
Value EmptyOp::getDynamicSize(unsigned idx) {
assert(getType().isDynamicDim(idx) && "expected dynamic dim");
unsigned ctr = 0;
for (int64_t i = 0; i < static_cast<int64_t>(idx); ++i)
if (getType().isDynamicDim(i))
++ctr;
return getDynamicSizes()[ctr];
}
SmallVector<OpFoldResult> EmptyOp::getMixedSizes() {
SmallVector<OpFoldResult> result;
unsigned ctr = 0;
OpBuilder b(getContext());
for (int64_t i = 0; i < getType().getRank(); ++i) {
if (getType().isDynamicDim(i)) {
result.push_back(getDynamicSizes()[ctr++]);
} else {
result.push_back(b.getIndexAttr(getType().getShape()[i]));
}
}
return result;
}
namespace {
/// Change the type of the result of a `tensor.empty` by making the result
/// type statically sized along dimensions that in the original operation were
/// defined as dynamic, but the size was defined using a `constant` op. For
/// example
///
/// %c5 = arith.constant 5: index
/// %0 = tensor.empty(%arg0, %c5) : tensor<?x?xf32>
///
/// to
///
/// %0 = tensor.empty(%arg0) : tensor<?x5xf32>
struct ReplaceEmptyTensorStaticShapeDims : OpRewritePattern<EmptyOp> {
using OpRewritePattern<EmptyOp>::OpRewritePattern;
LogicalResult matchAndRewrite(EmptyOp op,
PatternRewriter &rewriter) const override {
SmallVector<int64_t> staticShape(op.getType().getShape().begin(),
op.getType().getShape().end());
SmallVector<Value> dynamicSizes;
// Compute new static and dynamic sizes.
unsigned ctr = 0;
bool changedType = false;
for (int64_t i = 0; i < op.getType().getRank(); ++i) {
if (op.getType().isDynamicDim(i)) {
Value dynamicSize = op.getDynamicSizes()[ctr++];
Optional<int64_t> cst = getConstantIntValue(dynamicSize);
if (cst.hasValue()) {
staticShape[i] = *cst;
changedType = true;
} else {
dynamicSizes.push_back(dynamicSize);
}
}
}
// Stop here if no dynamic size was promoted to static.
if (!changedType)
return failure();
auto tensorType = RankedTensorType::get(
staticShape, op.getType().getElementType(), op.getType().getEncoding());
auto newOp =
rewriter.create<EmptyOp>(op.getLoc(), tensorType, dynamicSizes);
rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
return success();
}
};
/// `tensor.empty` does not define any tensor contents, so a slice of a
/// `tensor.empty` can be canonicalized to a smaller `tensor.empty`.
struct FoldEmptyTensorWithExtractSliceOp
: public OpRewritePattern<ExtractSliceOp> {
using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ExtractSliceOp sliceOp,
PatternRewriter &rewriter) const override {
if (!sliceOp.getSource().getDefiningOp<EmptyOp>())
return failure();
// ExtractSliceOp may be rank-reducing; its dynamic sizes must be
// preserved as well as its result type.
auto tensorType = RankedTensorType::get(sliceOp.getType().getShape(),
sliceOp.getType().getElementType(),
sliceOp.getType().getEncoding());
rewriter.replaceOpWithNewOp<EmptyOp>(sliceOp, tensorType,
sliceOp.getSizes());
return success();
}
};
template <typename ReshapeOp>
struct FoldEmptyTensorWithReshapeOp : public OpRewritePattern<ReshapeOp> {
using OpRewritePattern<ReshapeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(ReshapeOp reshapeOp,
PatternRewriter &rewriter) const override {
if (!reshapeOp.getSrc().template getDefiningOp<EmptyOp>())
return failure();
Location loc = reshapeOp.getLoc();
ReifiedRankedShapedTypeDims resultShapes;
ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
cast<ReifyRankedShapedTypeOpInterface>(reshapeOp.getOperation());
if (failed(reifyShapedTypeInterface.reifyResultShapes(rewriter,
resultShapes)) ||
!llvm::hasSingleElement(resultShapes))
return failure();
// TODO: Do not drop tensor type encoding.
Value emptyTensor =
rewriter.create<EmptyOp>(loc, getAsOpFoldResult(resultShapes[0]),
reshapeOp.getResultType().getElementType());
if (emptyTensor.getType() != reshapeOp.getResultType()) {
rewriter.replaceOpWithNewOp<tensor::CastOp>(
reshapeOp, reshapeOp.getResultType(), emptyTensor);
} else {
rewriter.replaceOp(reshapeOp, emptyTensor);
}
return success();
}
};
struct FoldEmptyTensorWithDimOp : public OpRewritePattern<DimOp> {
using OpRewritePattern<DimOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tensor::DimOp dimOp,
PatternRewriter &rewriter) const override {
Optional<int64_t> maybeConstantIndex = dimOp.getConstantIndex();
auto emptyTensorOp = dimOp.getSource().getDefiningOp<EmptyOp>();
if (!emptyTensorOp || !maybeConstantIndex)
return failure();
if (!emptyTensorOp.getType().isDynamicDim(*maybeConstantIndex))
return failure();
rewriter.replaceOp(dimOp,
emptyTensorOp.getDynamicSize(*maybeConstantIndex));
return success();
}
};
/// Canonicalize
///
/// ```mlir
/// %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
/// %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x?xf32>
/// ```
///
/// into
///
/// ```mlir
/// %0 = tensor.empty(%d1) : tensor<4x?xf32>
/// ```
///
/// This assumes the input program is correct in terms of its shape. So it is
/// safe to assume that `%d0` is in fact 4.
struct FoldEmptyTensorWithCastOp : public OpRewritePattern<CastOp> {
using OpRewritePattern<CastOp>::OpRewritePattern;
LogicalResult matchAndRewrite(CastOp castOp,
PatternRewriter &rewriter) const override {
if (!canFoldIntoProducerOp(castOp))
return failure();
auto producer = castOp.getSource().getDefiningOp<EmptyOp>();
if (!producer)
return failure();
auto resultType = castOp->getResult(0).getType().cast<RankedTensorType>();
ArrayRef<int64_t> resultShape = resultType.getShape();
SmallVector<OpFoldResult> currMixedSizes = producer.getMixedSizes();
SmallVector<OpFoldResult> newMixedSizes;
newMixedSizes.reserve(currMixedSizes.size());
assert(resultShape.size() == currMixedSizes.size() &&
"mismatch in result shape and sizes of empty op");
for (auto it : llvm::zip(resultShape, currMixedSizes)) {
int64_t newDim = std::get<0>(it);
OpFoldResult currDim = std::get<1>(it);
// Case 1: The empty tensor dim is static. Check that the tensor cast
// result dim matches.
if (auto attr = currDim.dyn_cast<Attribute>()) {
if (ShapedType::isDynamic(newDim) ||
newDim != attr.cast<IntegerAttr>().getInt()) {
// Something is off, the cast result shape cannot be more dynamic
// than the empty tensor result shape (enforced by
// `canFoldIntoProducer`). Abort for now.
return rewriter.notifyMatchFailure(
producer, "mismatch in static value of shape of empty tensor "
"result and cast result");
}
newMixedSizes.push_back(attr);
continue;
}
// Case 2 : The tensor cast shape is static, but empty tensor result
// shape is dynamic.
if (!ShapedType::isDynamic(newDim)) {
newMixedSizes.push_back(rewriter.getIndexAttr(newDim));
continue;
}
// Case 3 : The tensor cast shape is dynamic and empty tensor result
// shape is dynamic. Use the dynamic value from the empty tensor op.
newMixedSizes.push_back(currDim);
}
// TODO: Do not drop tensor encoding.
rewriter.replaceOpWithNewOp<EmptyOp>(castOp, newMixedSizes,
resultType.getElementType());
return success();
}
};
} // namespace
void EmptyOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<FoldEmptyTensorWithCastOp, FoldEmptyTensorWithDimOp,
FoldEmptyTensorWithExtractSliceOp,
FoldEmptyTensorWithReshapeOp<tensor::ExpandShapeOp>,
FoldEmptyTensorWithReshapeOp<tensor::CollapseShapeOp>,
ReplaceEmptyTensorStaticShapeDims>(context);
}
//===----------------------------------------------------------------------===//
// ExtractOp // ExtractOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
LogicalResult ExtractOp::verify() { LogicalResult ExtractOp::verify() {
// Verify the # indices match if we have a ranked type. // Verify the # indices match if we have a ranked type.
if (auto tensorType = getTensor().getType().dyn_cast<RankedTensorType>()) if (auto tensorType = getTensor().getType().dyn_cast<RankedTensorType>())
if (tensorType.getRank() != static_cast<int64_t>(getIndices().size())) if (tensorType.getRank() != static_cast<int64_t>(getIndices().size()))
return emitOpError("incorrect number of indices for extract_element"); return emitOpError("incorrect number of indices for extract_element");
return success(); return success();
} }
OpFoldResult ExtractOp::fold(ArrayRef<Attribute> operands) { OpFoldResult ExtractOp::fold(ArrayRef<Attribute> operands) {
// If this is a splat elements attribute, simply return the value. All of the // If this is a splat elements attribute, simply return the value. All of
// elements of a splat attribute are the same. // the elements of a splat attribute are the same.
if (Attribute tensor = operands.front()) if (Attribute tensor = operands.front())
if (auto splatTensor = tensor.dyn_cast<SplatElementsAttr>()) if (auto splatTensor = tensor.dyn_cast<SplatElementsAttr>())
return splatTensor.getSplatValue<Attribute>(); return splatTensor.getSplatValue<Attribute>();
// Collect the constant indices into the tensor. // Collect the constant indices into the tensor.
SmallVector<uint64_t, 8> indices; SmallVector<uint64_t, 8> indices;
for (Attribute indice : llvm::drop_begin(operands, 1)) { for (Attribute indice : llvm::drop_begin(operands, 1)) {
if (!indice || !indice.isa<IntegerAttr>()) if (!indice || !indice.isa<IntegerAttr>())
Show All 9 Lines assert(static_cast<int64_t>(indices.size()) == tensorType.getRank() &&
"rank mismatch"); "rank mismatch");
int flatIndex = 0; int flatIndex = 0;
int stride = 1; int stride = 1;
for (int i = rank - 1; i >= 0; --i) { for (int i = rank - 1; i >= 0; --i) {
if (i < rank - 1) if (i < rank - 1)
stride *= tensorType.getDimSize(i); stride *= tensorType.getDimSize(i);
flatIndex += indices[i] * stride; flatIndex += indices[i] * stride;
} }
// Prevent out of bounds accesses. This can happen in invalid code that will // Prevent out of bounds accesses. This can happen in invalid code that
// never execute. // will never execute.
if (static_cast<int>(fromElementsOp.getElements().size()) <= flatIndex || if (static_cast<int>(fromElementsOp.getElements().size()) <= flatIndex ||
flatIndex < 0) flatIndex < 0)
return {}; return {};
return fromElementsOp.getElements()[flatIndex]; return fromElementsOp.getElements()[flatIndex];
} }
// If this is an elements attribute, query the value at the given indices. // If this is an elements attribute, query the value at the given indices.
if (Attribute tensor = operands.front()) { if (Attribute tensor = operands.front()) {
Show All 40 Lines
// //
// to the following: // to the following:
// //
// %extract = tensor.extract %tensor[%index] : tensor<1xindex> // %extract = tensor.extract %tensor[%index] : tensor<1xindex>
// %cast = arith.index_cast %extract : i32 to index // %cast = arith.index_cast %extract : i32 to index
// //
// to just %element. // to just %element.
// //
// Consider expanding this to a template and handle all tensor cast operations. // Consider expanding this to a template and handle all tensor cast
// operations.
struct ExtractElementFromIndexCast struct ExtractElementFromIndexCast
: public OpRewritePattern<tensor::ExtractOp> { : public OpRewritePattern<tensor::ExtractOp> {
using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern; using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
LogicalResult matchAndRewrite(tensor::ExtractOp extract, LogicalResult matchAndRewrite(tensor::ExtractOp extract,
PatternRewriter &rewriter) const final { PatternRewriter &rewriter) const final {
Location loc = extract.getLoc(); Location loc = extract.getLoc();
auto indexCast = extract.getTensor().getDefiningOp<arith::IndexCastOp>(); auto indexCast = extract.getTensor().getDefiningOp<arith::IndexCastOp>();
▲ Show 20 Lines • Show All 139 Lines • ▼ Show 20 Lines if (getType().isDynamicDim(dim)) {
reifiedReturnShapes[0][dim] = builder.create<arith::ConstantIndexOp>( reifiedReturnShapes[0][dim] = builder.create<arith::ConstantIndexOp>(
getLoc(), getType().getDimSize(dim)); getLoc(), getType().getDimSize(dim));
} }
} }
return success(); return success();
} }
LogicalResult GenerateOp::verify() { LogicalResult GenerateOp::verify() {
// Ensure that the tensor type has as many dynamic dimensions as are specified // Ensure that the tensor type has as many dynamic dimensions as are
// by the operands. // specified by the operands.
RankedTensorType resultTy = getType().cast<RankedTensorType>(); RankedTensorType resultTy = getType().cast<RankedTensorType>();
if (getNumOperands() != resultTy.getNumDynamicDims()) if (getNumOperands() != resultTy.getNumDynamicDims())
return emitError("must have as many index operands as dynamic extents " return emitError("must have as many index operands as dynamic extents "
"in the result type"); "in the result type");
return success(); return success();
} }
▲ Show 20 Lines • Show All 219 Lines • ▼ Show 20 Lines
SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() { SmallVector<AffineMap, 4> ExpandShapeOp::getReassociationMaps() {
return getSymbolLessAffineMaps(getReassociationExprs()); return getSymbolLessAffineMaps(getReassociationExprs());
} }
SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() { SmallVector<ReassociationExprs, 4> ExpandShapeOp::getReassociationExprs() {
return convertReassociationIndicesToExprs(getContext(), return convertReassociationIndicesToExprs(getContext(),
getReassociationIndices()); getReassociationIndices());
} }
/// Compute the RankedTensorType obtained by applying `reassociation` to `type`. /// Compute the RankedTensorType obtained by applying `reassociation` to
/// `type`.
static RankedTensorType static RankedTensorType
computeTensorReshapeCollapsedType(RankedTensorType type, computeTensorReshapeCollapsedType(RankedTensorType type,
ArrayRef<AffineMap> reassociation) { ArrayRef<AffineMap> reassociation) {
auto shape = type.getShape(); auto shape = type.getShape();
SmallVector<int64_t, 4> newShape; SmallVector<int64_t, 4> newShape;
newShape.reserve(reassociation.size()); newShape.reserve(reassociation.size());
// Use the fact that reassociation is valid to simplify the logic: only use // Use the fact that reassociation is valid to simplify the logic: only use
▲ Show 20 Lines • Show All 81 Lines • ▼ Show 20 Lines if (!attr || !attr.isSplat())
return failure(); return failure();
DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer( DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
reshapeOp.getResultType(), attr.getRawData()); reshapeOp.getResultType(), attr.getRawData());
rewriter.replaceOpWithNewOp<arith::ConstantOp>(reshapeOp, newAttr); rewriter.replaceOpWithNewOp<arith::ConstantOp>(reshapeOp, newAttr);
return success(); return success();
} }
}; };
/// Reshape of a FromElements can be replaced with a FromElements of the result /// Reshape of a FromElements can be replaced with a FromElements of the
/// type /// result type
template <typename TensorReshapeOp> template <typename TensorReshapeOp>
struct FoldReshapeWithFromElements : OpRewritePattern<TensorReshapeOp> { struct FoldReshapeWithFromElements : OpRewritePattern<TensorReshapeOp> {
using OpRewritePattern<TensorReshapeOp>::OpRewritePattern; using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp, LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
PatternRewriter &rewriter) const override { PatternRewriter &rewriter) const override {
auto fromElements = auto fromElements =
reshapeOp.getSrc().template getDefiningOp<FromElementsOp>(); reshapeOp.getSrc().template getDefiningOp<FromElementsOp>();
if (!fromElements) if (!fromElements)
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Lines
/// An extract_slice result type can be inferred, when it is not /// An extract_slice result type can be inferred, when it is not
/// rank-reduced, from the source type and the static representation of /// rank-reduced, from the source type and the static representation of
/// offsets, sizes and strides. Special sentinels encode the dynamic case. /// offsets, sizes and strides. Special sentinels encode the dynamic case.
RankedTensorType ExtractSliceOp::inferResultType( RankedTensorType ExtractSliceOp::inferResultType(
ShapedType sourceShapedTensorType, ArrayRef<int64_t> staticOffsets, ShapedType sourceShapedTensorType, ArrayRef<int64_t> staticOffsets,
ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides) { ArrayRef<int64_t> staticSizes, ArrayRef<int64_t> staticStrides) {
// An extract_slice op may specify only a leading subset of offset/sizes/ // An extract_slice op may specify only a leading subset of offset/sizes/
// strides in which case we complete with offset=0, sizes from memref type and // strides in which case we complete with offset=0, sizes from memref type
// strides=1. // and strides=1.
assert(static_cast<int64_t>(staticSizes.size()) == assert(static_cast<int64_t>(staticSizes.size()) ==
sourceShapedTensorType.getRank() && sourceShapedTensorType.getRank() &&
"unexpected staticSizes not equal to rank of source"); "unexpected staticSizes not equal to rank of source");
return RankedTensorType::get(staticSizes, return RankedTensorType::get(staticSizes,
sourceShapedTensorType.getElementType()); sourceShapedTensorType.getElementType());
} }
RankedTensorType ExtractSliceOp::inferResultType( RankedTensorType ExtractSliceOp::inferResultType(
ShapedType sourceShapedTensorType, ArrayRef<OpFoldResult> offsets, ShapedType sourceShapedTensorType, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) { ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> strides) {
SmallVector<int64_t> staticOffsets, staticSizes, staticStrides; SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides; SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets, dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets,
ShapedType::kDynamicStrideOrOffset); ShapedType::kDynamicStrideOrOffset);
dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes, dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
ShapedType::kDynamicSize); ShapedType::kDynamicSize);
dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
ShapedType::kDynamicStrideOrOffset); ShapedType::kDynamicStrideOrOffset);
return ExtractSliceOp::inferResultType(sourceShapedTensorType, staticOffsets, return ExtractSliceOp::inferResultType(sourceShapedTensorType, staticOffsets,
staticSizes, staticStrides); staticSizes, staticStrides);
} }
/// If the rank is reduced (i.e. the desiredResultRank is smaller than the /// If the rank is reduced (i.e. the desiredResultRank is smaller than the
/// number of sizes), drop as many size 1 as needed to produce an inferred type /// number of sizes), drop as many size 1 as needed to produce an inferred
/// with the desired rank. /// type with the desired rank.
/// ///
/// Note that there may be multiple ways to compute this rank-reduced type: /// Note that there may be multiple ways to compute this rank-reduced type:
/// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors. /// e.g. 1x6x1 can rank-reduce to either 1x6 or 6x1 2-D tensors.
/// ///
/// To disambiguate, this function always drops the first 1 sizes occurrences. /// To disambiguate, this function always drops the first 1 sizes occurrences.
RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType( RankedTensorType ExtractSliceOp::inferCanonicalRankReducedResultType(
unsigned desiredResultRank, RankedTensorType sourceRankedTensorType, unsigned desiredResultRank, RankedTensorType sourceRankedTensorType,
ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes, ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes,
▲ Show 20 Lines • Show All 70 Lines • ▼ Show 20 Lines
void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source, void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, ArrayRef<OpFoldResult> sizes,
ArrayRef<OpFoldResult> strides, ArrayRef<OpFoldResult> strides,
ArrayRef<NamedAttribute> attrs) { ArrayRef<NamedAttribute> attrs) {
build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs);
} }
/// Build an ExtractSliceOp with mixed static and dynamic entries packed into a /// Build an ExtractSliceOp with mixed static and dynamic entries packed into
/// Range vector. /// a Range vector.
void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source, void ExtractSliceOp::build(OpBuilder &b, OperationState &result, Value source,
ArrayRef<Range> ranges, ArrayRef<Range> ranges,
ArrayRef<NamedAttribute> attrs) { ArrayRef<NamedAttribute> attrs) {
auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges);
build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs); build(b, result, RankedTensorType(), source, offsets, sizes, strides, attrs);
} }
/// Build an ExtractSliceOp with dynamic entries and custom result type. If the /// Build an ExtractSliceOp with dynamic entries and custom result type. If
/// type passed is nullptr, it is inferred. /// the type passed is nullptr, it is inferred.
void ExtractSliceOp::build(OpBuilder &b, OperationState &result, void ExtractSliceOp::build(OpBuilder &b, OperationState &result,
RankedTensorType resultType, Value source, RankedTensorType resultType, Value source,
ValueRange offsets, ValueRange sizes, ValueRange offsets, ValueRange sizes,
ValueRange strides, ArrayRef<NamedAttribute> attrs) { ValueRange strides, ArrayRef<NamedAttribute> attrs) {
SmallVector<OpFoldResult> offsetValues = llvm::to_vector<4>( SmallVector<OpFoldResult> offsetValues = llvm::to_vector<4>(
llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; })); llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; }));
SmallVector<OpFoldResult> sizeValues = llvm::to_vector<4>( SmallVector<OpFoldResult> sizeValues = llvm::to_vector<4>(
llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; })); llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; }));
▲ Show 20 Lines • Show All 156 Lines • ▼ Show 20 Linesstatic void sliceElements(IterTy values, ArrayRef<int64_t> counts,
for (int64_t i = 0; i < size; ++i, offset += stride) { for (int64_t i = 0; i < size; ++i, offset += stride) {
auto begin = values + offset * counts.front(); auto begin = values + offset * counts.front();
sliceElements<IterTy, ElemTy>(begin, counts.drop_front(), sliceElements<IterTy, ElemTy>(begin, counts.drop_front(),
offsets.drop_front(), sizes.drop_front(), offsets.drop_front(), sizes.drop_front(),
strides.drop_front(), outValues); strides.drop_front(), outValues);
} }
} }
/// Fold arith.constant and tensor.extract_slice into arith.constant. The folded /// Fold arith.constant and tensor.extract_slice into arith.constant. The
/// operation might introduce more constant data; Users can control their /// folded operation might introduce more constant data; Users can control
/// heuristics by the control function. /// their heuristics by the control function.
class ConstantOpExtractSliceFolder final class ConstantOpExtractSliceFolder final
: public OpRewritePattern<ExtractSliceOp> { : public OpRewritePattern<ExtractSliceOp> {
public: public:
using OpRewritePattern<ExtractSliceOp>::OpRewritePattern; using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
ConstantOpExtractSliceFolder(MLIRContext *context, ConstantOpExtractSliceFolder(MLIRContext *context,
ControlConstantExtractSliceFusionFn controlFn) ControlConstantExtractSliceFusionFn controlFn)
: OpRewritePattern<ExtractSliceOp>(context), : OpRewritePattern<ExtractSliceOp>(context),
▲ Show 20 Lines • Show All 117 Lines • ▼ Show 20 Lines
// //
static LogicalResult static LogicalResult
foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op, foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op,
ShapedType shapedType) { ShapedType shapedType) {
OpBuilder b(op.getContext()); OpBuilder b(op.getContext());
for (OpFoldResult ofr : op.getMixedOffsets()) for (OpFoldResult ofr : op.getMixedOffsets())
if (getConstantIntValue(ofr) != static_cast<int64_t>(0)) if (getConstantIntValue(ofr) != static_cast<int64_t>(0))
return failure(); return failure();
// Rank-reducing noops only need to inspect the leading dimensions: llvm::zip // Rank-reducing noops only need to inspect the leading dimensions:
// is appropriate. // llvm::zip is appropriate.
auto shape = shapedType.getShape(); auto shape = shapedType.getShape();
for (auto it : llvm::zip(op.getMixedSizes(), shape)) for (auto it : llvm::zip(op.getMixedSizes(), shape))
if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it)) if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it))
return failure(); return failure();
for (OpFoldResult ofr : op.getMixedStrides()) for (OpFoldResult ofr : op.getMixedStrides())
if (getConstantIntValue(ofr) != static_cast<int64_t>(1)) if (getConstantIntValue(ofr) != static_cast<int64_t>(1))
return failure(); return failure();
return success(); return success();
} }
/// If we have an ExtractSliceOp consuming an InsertSliceOp with the same slice, /// If we have an ExtractSliceOp consuming an InsertSliceOp with the same
/// we can return the InsertSliceOp's source directly. /// slice, we can return the InsertSliceOp's source directly.
// TODO: This only checks the immediate producer; extend to go up the // TODO: This only checks the immediate producer; extend to go up the
// insert/extract chain if the slices are disjoint. // insert/extract chain if the slices are disjoint.
static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) { static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) {
auto insertOp = extractOp.getSource().getDefiningOp<InsertSliceOp>(); auto insertOp = extractOp.getSource().getDefiningOp<InsertSliceOp>();
auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; }; auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; };
if (insertOp && insertOp.getSource().getType() == extractOp.getType() && if (insertOp && insertOp.getSource().getType() == extractOp.getType() &&
insertOp.isSameAs(extractOp, isSame)) insertOp.isSameAs(extractOp, isSame))
▲ Show 20 Lines • Show All 76 Lines • ▼ Show 20 Lines
/// Rank-reducing type verification for both InsertSliceOp and /// Rank-reducing type verification for both InsertSliceOp and
/// ParallelInsertSliceOp. /// ParallelInsertSliceOp.
static SliceVerificationResult static SliceVerificationResult
verifyInsertSliceOp(ShapedType srcType, ShapedType dstType, verifyInsertSliceOp(ShapedType srcType, ShapedType dstType,
ArrayAttr staticOffsets, ArrayAttr staticSizes, ArrayAttr staticOffsets, ArrayAttr staticSizes,
ArrayAttr staticStrides, ArrayAttr staticStrides,
ShapedType *expectedType = nullptr) { ShapedType *expectedType = nullptr) {
// insert_slice is the inverse of extract_slice, use the same type inference. // insert_slice is the inverse of extract_slice, use the same type
// inference.
RankedTensorType expected = ExtractSliceOp::inferResultType( RankedTensorType expected = ExtractSliceOp::inferResultType(
dstType, extractFromI64ArrayAttr(staticOffsets), dstType, extractFromI64ArrayAttr(staticOffsets),
extractFromI64ArrayAttr(staticSizes), extractFromI64ArrayAttr(staticSizes),
extractFromI64ArrayAttr(staticStrides)); extractFromI64ArrayAttr(staticStrides));
if (expectedType) if (expectedType)
*expectedType = expected; *expectedType = expected;
return isRankReducedType(expected, srcType); return isRankReducedType(expected, srcType);
} }
▲ Show 20 Lines • Show All 107 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
// Create the new op in canonical form. // Create the new op in canonical form.
auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType( auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType(
insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(), insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(),
mixedOffsets, mixedSizes, mixedStrides); mixedOffsets, mixedSizes, mixedStrides);
Value toInsert = insertSliceOp.getSource(); Value toInsert = insertSliceOp.getSource();
if (sourceType != insertSliceOp.getSourceType()) { if (sourceType != insertSliceOp.getSourceType()) {
OpBuilder::InsertionGuard g(rewriter); OpBuilder::InsertionGuard g(rewriter);
// The only difference between InsertSliceOp and ParallelInsertSliceOp is // The only difference between InsertSliceOp and ParallelInsertSliceOp
// the the insertion point is just before the ParallelCombiningOp in the // is the the insertion point is just before the ParallelCombiningOp in
// parallel case. // the parallel case.
if (std::is_same<InsertOpTy, ParallelInsertSliceOp>::value) if (std::is_same<InsertOpTy, ParallelInsertSliceOp>::value)
rewriter.setInsertionPoint(insertSliceOp->getParentOp()); rewriter.setInsertionPoint(insertSliceOp->getParentOp());
toInsert = rewriter.create<tensor::CastOp>(insertSliceOp.getLoc(), toInsert = rewriter.create<tensor::CastOp>(insertSliceOp.getLoc(),
sourceType, toInsert); sourceType, toInsert);
} }
rewriter.replaceOpWithNewOp<InsertOpTy>( rewriter.replaceOpWithNewOp<InsertOpTy>(
insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets, insertSliceOp, toInsert, insertSliceOp.getDest(), mixedOffsets,
mixedSizes, mixedStrides); mixedSizes, mixedStrides);
return success(); return success();
} }
}; };
/// Fold tensor_casts with insert_slice operations. If the source or destination /// Fold tensor_casts with insert_slice operations. If the source or
/// tensor is a tensor_cast that removes static type information, the cast is /// destination tensor is a tensor_cast that removes static type information,
/// folded into the insert_slice operation. E.g.: /// the cast is folded into the insert_slice operation. E.g.:
/// ///
/// ```mlir /// ```mlir
/// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32> /// %1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32>
/// %2 = tensor.insert_slice %1 into ... : tensor<?x?xf32> into ... /// %2 = tensor.insert_slice %1 into ... : tensor<?x?xf32> into ...
/// ``` /// ```
/// ///
/// folds into: /// folds into:
/// ///
▲ Show 20 Lines • Show All 382 Lines • ▼ Show 20 Lines
/// hold: /// hold:
/// 1) the tensor::ExtractSliceOps are not rank-reducing, /// 1) the tensor::ExtractSliceOps are not rank-reducing,
/// 2) the tensor::ExtractSliceOps have only unit-strides, /// 2) the tensor::ExtractSliceOps have only unit-strides,
/// 3) the tensor::PadOps perform only high-padding, /// 3) the tensor::PadOps perform only high-padding,
/// 4) the tensor::PadOps have the same constant padding value, /// 4) the tensor::PadOps have the same constant padding value,
/// 5) the tensor::PadOps do not have common padding dimensions, /// 5) the tensor::PadOps do not have common padding dimensions,
/// 6) one tensor::ExtractSliceOp, tensor::PadOp pair has zero-padding and /// 6) one tensor::ExtractSliceOp, tensor::PadOp pair has zero-padding and
/// zero-offset for every dimension. /// zero-offset for every dimension.
/// 7) the tensor::ExtractSliceOp sizes match the source tensor sizes for the /// 7) the tensor::ExtractSliceOp sizes match the source tensor sizes for
/// the
/// padded source dimensions. /// padded source dimensions.
/// ///
/// Example: /// Example:
/// ///
/// ```mlir /// ```mlir
/// %0 = tensor.extract_slice %input[16, 0] [%sz0, 64] [1, 1] /// %0 = tensor.extract_slice %input[16, 0] [%sz0, 64] [1, 1]
/// : tensor<64x64xf32> to tensor<?x64xf32> /// : tensor<64x64xf32> to tensor<?x64xf32>
/// %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] { ... /// %1 = tensor.pad %0 low[0, 0] high[%pw0, 0] { ...
▲ Show 20 Lines • Show All 84 Lines • ▼ Show 20 Lines for (auto &en : enumerate(newOffsets)) {
(getConstantIntValue(outerOffset) == static_cast<int64_t>(0))) { (getConstantIntValue(outerOffset) == static_cast<int64_t>(0))) {
en.value() = innerOffset; en.value() = innerOffset;
continue; continue;
} }
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
padOp, "cannot find zero-offset and zero-padding pair"); padOp, "cannot find zero-offset and zero-padding pair");
} }
// 7) Combine the sizes of the two tensor::ExtractSliceOps. Take the size of // 7) Combine the sizes of the two tensor::ExtractSliceOps. Take the size
// the outer tensor::ExtractSliceOp for the dimensions padded by the outer // of the outer tensor::ExtractSliceOp for the dimensions padded by the
// tensor::PadOp and fail if the size of the inner tensor::ExtractSliceOp // outer tensor::PadOp and fail if the size of the inner
// does not match the size of the padded dimension. Otherwise, take the size // tensor::ExtractSliceOp does not match the size of the padded dimension.
// of the inner tensor::ExtractSliceOp. // Otherwise, take the size of the inner tensor::ExtractSliceOp.
SmallVector<OpFoldResult> newSizes = innerSliceOp.getMixedSizes(); SmallVector<OpFoldResult> newSizes = innerSliceOp.getMixedSizes();
for (auto &en : enumerate(newSizes)) { for (auto &en : enumerate(newSizes)) {
if (!outerDims.test(en.index())) if (!outerDims.test(en.index()))
continue; continue;
OpFoldResult sliceSize = innerSliceOp.getMixedSizes()[en.index()]; OpFoldResult sliceSize = innerSliceOp.getMixedSizes()[en.index()];
int64_t sourceSize = innerSliceOp.getSourceType().getShape()[en.index()]; int64_t sourceSize = innerSliceOp.getSourceType().getShape()[en.index()];
assert(!ShapedType::isDynamic(sourceSize) && assert(!ShapedType::isDynamic(sourceSize) &&
"expected padded dimension to have a static size"); "expected padded dimension to have a static size");
Show All 9 Lines LogicalResult matchAndRewrite(PadOp padOp,
SmallVector<OpFoldResult> newHighPad(rank, rewriter.getIndexAttr(0)); SmallVector<OpFoldResult> newHighPad(rank, rewriter.getIndexAttr(0));
for (auto &en : enumerate(newHighPad)) { for (auto &en : enumerate(newHighPad)) {
if (innerDims.test(en.index())) if (innerDims.test(en.index()))
newHighPad[en.index()] = padOp.getMixedHighPad()[en.index()]; newHighPad[en.index()] = padOp.getMixedHighPad()[en.index()];
if (outerDims.test(en.index())) if (outerDims.test(en.index()))
newHighPad[en.index()] = outerPadOp.getMixedHighPad()[en.index()]; newHighPad[en.index()] = outerPadOp.getMixedHighPad()[en.index()];
} }
// Create a new tensor::ExtractSliceOp, tensor::PadOp pair that performs the // Create a new tensor::ExtractSliceOp, tensor::PadOp pair that performs
// two paddings in one step. // the two paddings in one step.
auto newSliceOp = rewriter.create<ExtractSliceOp>( auto newSliceOp = rewriter.create<ExtractSliceOp>(
padOp.getLoc(), outerSliceOp.getSource(), newOffsets, newSizes, padOp.getLoc(), outerSliceOp.getSource(), newOffsets, newSizes,
innerSliceOp.getMixedStrides()); innerSliceOp.getMixedStrides());
auto newPadOp = rewriter.create<PadOp>( auto newPadOp = rewriter.create<PadOp>(
padOp.getLoc(), padOp.getResultType(), newSliceOp.getResult(), padOp.getLoc(), padOp.getResultType(), newSliceOp.getResult(),
padOp.getMixedLowPad(), newHighPad, padOp.getNofold()); padOp.getMixedLowPad(), newHighPad, padOp.getNofold());
rewriter.inlineRegionBefore(padOp.getRegion(), newPadOp.getRegion(), rewriter.inlineRegionBefore(padOp.getRegion(), newPadOp.getRegion(),
newPadOp.getRegion().begin()); newPadOp.getRegion().begin());
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Linesvoid ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides, dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
ShapedType::kDynamicStrideOrOffset); ShapedType::kDynamicStrideOrOffset);
build(b, result, {}, source, dest, dynamicOffsets, dynamicSizes, build(b, result, {}, source, dest, dynamicOffsets, dynamicSizes,
dynamicStrides, b.getI64ArrayAttr(staticOffsets), dynamicStrides, b.getI64ArrayAttr(staticOffsets),
b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides)); b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides));
result.addAttributes(attrs); result.addAttributes(attrs);
} }
/// Build an ParallelInsertSliceOp with mixed static and dynamic entries packed /// Build an ParallelInsertSliceOp with mixed static and dynamic entries
/// into a Range vector. /// packed into a Range vector.
void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result, void ParallelInsertSliceOp::build(OpBuilder &b, OperationState &result,
Value source, Value dest, Value source, Value dest,
ArrayRef<Range> ranges, ArrayRef<Range> ranges,
ArrayRef<NamedAttribute> attrs) { ArrayRef<NamedAttribute> attrs) {
auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges); auto [offsets, sizes, strides] = getOffsetsSizesAndStrides(ranges);
build(b, result, source, dest, offsets, sizes, strides, attrs); build(b, result, source, dest, offsets, sizes, strides, attrs);
} }
▲ Show 20 Lines • Show All 70 Lines • ▼ Show 20 Lines
// SplatOp // SplatOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
OpFoldResult SplatOp::fold(ArrayRef<Attribute> operands) { OpFoldResult SplatOp::fold(ArrayRef<Attribute> operands) {
auto constOperand = operands.front(); auto constOperand = operands.front();
if (!constOperand.isa_and_nonnull<IntegerAttr, FloatAttr>()) if (!constOperand.isa_and_nonnull<IntegerAttr, FloatAttr>())
return {}; return {};
// SplatElementsAttr::get treats single value for second arg as being a splat. // SplatElementsAttr::get treats single value for second arg as being a
// splat.
return SplatElementsAttr::get(getType(), {constOperand}); return SplatElementsAttr::get(getType(), {constOperand});
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// TableGen'd op method definitions // TableGen'd op method definitions
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
#define GET_OP_CLASSES #define GET_OP_CLASSES
#include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc" #include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

Show All 25 Lines SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
b.setInsertionPoint(op); b.setInsertionPoint(op);
ReifiedRankedShapedTypeDims reifiedShapes; ReifiedRankedShapedTypeDims reifiedShapes;
ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface = ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
dyn_cast<ReifyRankedShapedTypeOpInterface>(op); dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
(void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes); (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
auto padOp = cast<PadOp>(op); auto padOp = cast<PadOp>(op);
SmallVector<OpFoldResult> mixedSizes = getAsOpFoldResult(reifiedShapes[0]); SmallVector<OpFoldResult> mixedSizes = getAsOpFoldResult(reifiedShapes[0]);
Value initTensor = b.create<linalg::InitTensorOp>( Value emptyTensor = b.create<EmptyOp>(
op->getLoc(), mixedSizes, padOp.getResultType().getElementType()); op->getLoc(), mixedSizes, padOp.getResultType().getElementType());
return {initTensor}; return {emptyTensor};
} }
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const { SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto padOp = cast<PadOp>(op); auto padOp = cast<PadOp>(op);
SmallVector<utils::IteratorType> iteratorTypes( SmallVector<utils::IteratorType> iteratorTypes(
padOp.getResultType().getRank(), utils::IteratorType::parallel); padOp.getResultType().getRank(), utils::IteratorType::parallel);
return iteratorTypes; return iteratorTypes;
} }
▲ Show 20 Lines • Show All 257 Lines • Show Last 20 Lines

mlir/python/mlir/dialects/_linalg_ops_ext.py

Show All 14 Lines
def isa(cls: Type, ty: Type): def isa(cls: Type, ty: Type):
try: try:
cls(ty) cls(ty)
return True return True
except ValueError: except ValueError:
return False return False
class InitTensorOp:
"""Extends the linalg.init_tensor op."""
def __init__(self,
sizes: Union[Sequence[int], Sequence[Value]],
element_type: Type,
*,
loc=None,
ip=None):
"""Constructs an `init_tensor` with either static or dynamic sizes."""
context = get_default_loc_context(loc)
operands = []
attributes = {}
# TODO: Refactor the InitTensorOp to take an element type attribute and
# then use normal result type inference, unifying the Python and C++ side
# with a standard mechanism (versus stashing that in builders).
if sizes and isinstance(sizes[0], Value):
# Dynamic sizes.
operands.extend(sizes)
static_size_ints = [-1] * len(sizes)
result_type = RankedTensorType.get(static_size_ints, element_type)
else:
# Static sizes.
result_type = RankedTensorType.get(sizes, element_type)
static_size_ints = sizes
i64_type = IntegerType.get_signless(64)
attributes["static_sizes"] = ArrayAttr.get(
[IntegerAttr.get(i64_type, s) for s in static_size_ints],
context=context)
op = self.build_generic(results=[result_type],
operands=operands,
attributes=attributes,
loc=loc,
ip=ip)
OpView.__init__(self, op)
class StructuredOpMixin: class StructuredOpMixin:
"""All structured ops use the same mixin class.""" """All structured ops use the same mixin class."""
def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None): def __init__(self, inputs, outputs=(), results=(), loc=None, ip=None):
super().__init__( super().__init__(
self.build_generic(results=list(results), self.build_generic(results=list(results),
operands=[list(inputs), list(outputs)], operands=[list(inputs), list(outputs)],
loc=loc, loc=loc,
Show All 11 Lines

mlir/python/mlir/dialects/_tensor_ops_ext.py

  • This file was added.
# 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
try:
from ..ir import *
except ImportError as e:
raise RuntimeError("Error loading imports from extension module") from e
from typing import Any, Optional, Sequence, Union
from ._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
class EmptyOp:
"""Extends the tensor.empty op."""
def __init__(self,
sizes: Sequence[Union[int, Value]],
element_type: Type,
*,
loc=None,
ip=None):
"""Constructs an `empty` with mixed static/dynamic sizes."""
# TODO: Refactor the EmptyOp to take an element type attribute and
# then use normal result type inference, unifying the Python and C++ side
# with a standard mechanism (versus stashing that in builders).
dynamic_sizes = []
static_sizes = []
for s in sizes:
if isinstance(s, int):
static_sizes.append(s)
else:
static_sizes.append(-1)
dynamic_sizes.append(s)
result_type = RankedTensorType.get(static_sizes, element_type)
op = self.build_generic(
results=[result_type],
operands=dynamic_sizes,
attributes={},
loc=loc,
ip=ip)
OpView.__init__(self, op)

mlir/test/Conversion/TensorToLinalg/tensor-ops-to-linalg.mlir

// RUN: mlir-opt -split-input-file -convert-tensor-to-linalg -cse -verify-diagnostics %s | FileCheck %s // RUN: mlir-opt -split-input-file -convert-tensor-to-linalg -cse -verify-diagnostics %s | FileCheck %s
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// tensor.pad // tensor.pad
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// CHECK-LABEL: func @generalize_pad_tensor_static_shape( // CHECK-LABEL: func @generalize_pad_tensor_static_shape(
// CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> { // CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[INIT:.*]] = linalg.init_tensor [1, 32, 32, 1] : tensor<1x32x32x1xf32> // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<1x32x32x1xf32>
// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32> // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32>
// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32> // CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32>
// CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32> // CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32>
func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> { func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] { %0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32> } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
return %0 : tensor<1x32x32x1xf32> return %0 : tensor<1x32x32x1xf32>
} }
// CHECK-LABEL: func @generalize_pad_tensor_dynamic_shape( // CHECK-LABEL: func @generalize_pad_tensor_dynamic_shape(
// CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>, // CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>,
// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> { // CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[C2]], %[[OFFSET]] : index // CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[C2]], %[[OFFSET]] : index
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>
// CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index // CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index
// CHECK: %[[INIT:.*]] = linalg.init_tensor [4, %[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]] : tensor<4x?x?x?xf32> // CHECK: %[[INIT:.*]] = tensor.empty(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : tensor<4x?x?x?xf32>
// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<4x?x?x?xf32>) -> tensor<4x?x?x?xf32> // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<4x?x?x?xf32>) -> tensor<4x?x?x?xf32>
// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]]{{\[}}%[[C0]], %[[C0]], %[[OFFSET]], %[[C0]]] [4, %[[DIM1]], 2, %[[DIM3]]] [1, 1, 1, 1] : tensor<4x?x2x?xf32> into tensor<4x?x?x?xf32> // CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]]{{\[}}%[[C0]], %[[C0]], %[[OFFSET]], %[[C0]]] [4, %[[DIM1]], 2, %[[DIM3]]] [1, 1, 1, 1] : tensor<4x?x2x?xf32> into tensor<4x?x?x?xf32>
// CHECK: return %[[PADDED]] : tensor<4x?x?x?xf32> // CHECK: return %[[PADDED]] : tensor<4x?x?x?xf32>
// CHECK: } // CHECK: }
func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> { func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] { %out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] {
^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index): ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32> } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
return %out : tensor<4x?x?x?xf32> return %out : tensor<4x?x?x?xf32>
} }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-named.mlir

// RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg-named)" %s -verify-diagnostics -o -| FileCheck %s // RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg-named)" %s -verify-diagnostics -o -| FileCheck %s
// CHECK-LABEL: @matmul // CHECK-LABEL: @matmul
func.func @matmul(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x6xf32>) -> (tensor<1x5x6xf32>) { func.func @matmul(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x6xf32>) -> (tensor<1x5x6xf32>) {
// CHECK: [[C0:%.+]] = arith.constant 0 // CHECK: [[C0:%.+]] = arith.constant 0
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 6] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) outs([[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) outs([[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
// CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x3xf32>, tensor<1x3x6xf32>) outs([[FILLED]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x3xf32>, tensor<1x3x6xf32>) outs([[FILLED]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
%0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x3xf32>, tensor<1x3x6xf32>) -> (tensor<1x5x6xf32>) %0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x3xf32>, tensor<1x3x6xf32>) -> (tensor<1x5x6xf32>)
return %0 : tensor<1x5x6xf32> return %0 : tensor<1x5x6xf32>
} }
// ----- // -----
// CHECK-LABEL: @matmul_quantized // CHECK-LABEL: @matmul_quantized
func.func @matmul_quantized(%arg0: tensor<1x5x3xi8>, %arg1: tensor<1x3x6xi8>) -> (tensor<1x5x6xi32>) { func.func @matmul_quantized(%arg0: tensor<1x5x3xi8>, %arg1: tensor<1x3x6xi8>) -> (tensor<1x5x6xi32>) {
// CHECK: [[C0:%.+]] = arith.constant 0 // CHECK: [[C0:%.+]] = arith.constant 0
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 6] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : i32) outs([[INIT]] : tensor<1x5x6xi32>) -> tensor<1x5x6xi32> // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : i32) outs([[INIT]] : tensor<1x5x6xi32>) -> tensor<1x5x6xi32>
// CHECK: [[ONE:%.+]] = arith.constant 1 // CHECK: [[ONE:%.+]] = arith.constant 1
// CHECK: [[TWO:%.+]] = arith.constant 2 // CHECK: [[TWO:%.+]] = arith.constant 2
// CHECK: linalg.quantized_batch_matmul ins(%arg0, %arg1, [[ONE]], [[TWO]] : tensor<1x5x3xi8>, tensor<1x3x6xi8>, i32, i32) outs([[FILLED]] : tensor<1x5x6xi32>) -> tensor<1x5x6xi32> // CHECK: linalg.quantized_batch_matmul ins(%arg0, %arg1, [[ONE]], [[TWO]] : tensor<1x5x3xi8>, tensor<1x3x6xi8>, i32, i32) outs([[FILLED]] : tensor<1x5x6xi32>) -> tensor<1x5x6xi32>
%0 = "tosa.matmul"(%arg0, %arg1) {quantization_info = #tosa.matmul_quant<a_zp = 1, b_zp = 2>} : (tensor<1x5x3xi8>, tensor<1x3x6xi8>) -> (tensor<1x5x6xi32>) %0 = "tosa.matmul"(%arg0, %arg1) {quantization_info = #tosa.matmul_quant<a_zp = 1, b_zp = 2>} : (tensor<1x5x3xi8>, tensor<1x3x6xi8>) -> (tensor<1x5x6xi32>)
return %0 : tensor<1x5x6xi32> return %0 : tensor<1x5x6xi32>
} }
// ----- // -----
// CHECK-LABEL: @matmul_dyn_batch // CHECK-LABEL: @matmul_dyn_batch
func.func @matmul_dyn_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x3x6xf32>) -> (tensor<?x5x6xf32>) { func.func @matmul_dyn_batch(%arg0: tensor<?x5x3xf32>, %arg1: tensor<?x3x6xf32>) -> (tensor<?x5x6xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[C0_0:.+]] = arith.constant 0 // CHECK: %[[C0_0:.+]] = arith.constant 0
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM]], 5, 6] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0_0]] : f32) outs(%[[INIT]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32> // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0_0]] : f32) outs(%[[INIT]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
// CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<?x5x3xf32>, tensor<?x3x6xf32>) outs(%[[FILLED]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32> // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<?x5x3xf32>, tensor<?x3x6xf32>) outs(%[[FILLED]] : tensor<?x5x6xf32>) -> tensor<?x5x6xf32>
%0 = "tosa.matmul"(%arg0, %arg1) : (tensor<?x5x3xf32>, tensor<?x3x6xf32>) -> (tensor<?x5x6xf32>) %0 = "tosa.matmul"(%arg0, %arg1) : (tensor<?x5x3xf32>, tensor<?x3x6xf32>) -> (tensor<?x5x6xf32>)
return %0 : tensor<?x5x6xf32> return %0 : tensor<?x5x6xf32>
} }
// ----- // -----
// CHECK-LABEL: @matmul_dyn_independent_dim // CHECK-LABEL: @matmul_dyn_independent_dim
func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>) { func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x3xf32>, %arg1: tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>) {
// CHECK: %[[C2:.+]] = arith.constant 2 // CHECK: %[[C2:.+]] = arith.constant 2
// CHECK: %[[DIM:.+]] = tensor.dim %arg1, %[[C2]] // CHECK: %[[DIM:.+]] = tensor.dim %arg1, %[[C2]]
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 5, %[[DIM]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x?xf32>) -> tensor<1x5x?xf32> // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x?xf32>) -> tensor<1x5x?xf32>
// CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x3xf32>, tensor<1x3x?xf32>) outs(%[[FILLED]] : tensor<1x5x?xf32>) -> tensor<1x5x?xf32> // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x3xf32>, tensor<1x3x?xf32>) outs(%[[FILLED]] : tensor<1x5x?xf32>) -> tensor<1x5x?xf32>
%0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x3xf32>, tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>) %0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x3xf32>, tensor<1x3x?xf32>) -> (tensor<1x5x?xf32>)
return %0 : tensor<1x5x?xf32> return %0 : tensor<1x5x?xf32>
} }
// ----- // -----
// CHECK-LABEL: @matmul_dyn_independent_dim // CHECK-LABEL: @matmul_dyn_independent_dim
func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x?xf32>, %arg1: tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>) { func.func @matmul_dyn_independent_dim(%arg0: tensor<1x5x?xf32>, %arg1: tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 5, 6] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
// CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x?xf32>, tensor<1x?x6xf32>) outs(%[[FILLED]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x5x?xf32>, tensor<1x?x6xf32>) outs(%[[FILLED]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32>
%0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x?xf32>, tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>) %0 = "tosa.matmul"(%arg0, %arg1) : (tensor<1x5x?xf32>, tensor<1x?x6xf32>) -> (tensor<1x5x6xf32>)
return %0 : tensor<1x5x6xf32> return %0 : tensor<1x5x6xf32>
} }
// ----- // -----
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @fully_connected // CHECK-LABEL: @fully_connected
func.func @fully_connected(%arg0: tensor<5x3xf32>, %arg1: tensor<6x3xf32>, %arg2: tensor<6xf32>) -> (tensor<5x6xf32>) { func.func @fully_connected(%arg0: tensor<5x3xf32>, %arg1: tensor<6x3xf32>, %arg2: tensor<6xf32>) -> (tensor<5x6xf32>) {
// CHECK: [[INITT:%.+]] = linalg.init_tensor [5, 6] // CHECK: [[INITT:%.+]] = tensor.empty()
// CHECK: [[ZERO:%.+]] = arith.constant 0 // CHECK: [[ZERO:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[ZERO]]{{.*}}outs([[INITT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[ZERO]]{{.*}}outs([[INITT]]
// CHECK: [[PERM:%.+]] = arith.constant dense<[1, 0]> // CHECK: [[PERM:%.+]] = arith.constant dense<[1, 0]>
// CHECK: [[TRANSPOSE:%.+]] = "tosa.transpose"(%arg1, [[PERM]]) // CHECK: [[TRANSPOSE:%.+]] = "tosa.transpose"(%arg1, [[PERM]])
// CHECK: [[INITB:%.+]] = linalg.init_tensor [5, 6] // CHECK: [[INITB:%.+]] = tensor.empty()
// CHECK: [[MATMUL:%.+]] = linalg.matmul ins(%arg0, [[TRANSPOSE]] : tensor<5x3xf32>, tensor<3x6xf32>) outs([[FILL]] : tensor<5x6xf32>) -> tensor<5x6xf32> // CHECK: [[MATMUL:%.+]] = linalg.matmul ins(%arg0, [[TRANSPOSE]] : tensor<5x3xf32>, tensor<3x6xf32>) outs([[FILL]] : tensor<5x6xf32>) -> tensor<5x6xf32>
// CHECK: [[ADDED:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, [[MATMUL]] : tensor<6xf32>, tensor<5x6xf32>) outs([[INITB]] : tensor<5x6xf32>) { // CHECK: [[ADDED:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, [[MATMUL]] : tensor<6xf32>, tensor<5x6xf32>) outs([[INITB]] : tensor<5x6xf32>) {
// CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
// CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32 // CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
// CHECK: linalg.yield [[ADD]] : f32 // CHECK: linalg.yield [[ADD]] : f32
%0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) : (tensor<5x3xf32>, tensor<6x3xf32>, tensor<6xf32>) -> (tensor<5x6xf32>) %0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) : (tensor<5x3xf32>, tensor<6x3xf32>, tensor<6xf32>) -> (tensor<5x6xf32>)
return %0 : tensor<5x6xf32> return %0 : tensor<5x6xf32>
} }
// ----- // -----
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @quantized_fully_connected // CHECK-LABEL: @quantized_fully_connected
func.func @quantized_fully_connected(%arg0: tensor<5x3xi8>, %arg1: tensor<6x3xi8>, %arg2: tensor<6xi32>) -> (tensor<5x6xi32>) { func.func @quantized_fully_connected(%arg0: tensor<5x3xi8>, %arg1: tensor<6x3xi8>, %arg2: tensor<6xi32>) -> (tensor<5x6xi32>) {
// CHECK: [[INITT:%.+]] = linalg.init_tensor [5, 6] // CHECK: [[INITT:%.+]] = tensor.empty()
// CHECK: [[ZERO:%.+]] = arith.constant 0 // CHECK: [[ZERO:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[ZERO]]{{.*}}outs([[INITT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[ZERO]]{{.*}}outs([[INITT]]
// CHECK: [[PERM:%.+]] = arith.constant dense<[1, 0]> // CHECK: [[PERM:%.+]] = arith.constant dense<[1, 0]>
// CHECK: [[TRANSPOSE:%.+]] = "tosa.transpose"(%arg1, [[PERM]]) // CHECK: [[TRANSPOSE:%.+]] = "tosa.transpose"(%arg1, [[PERM]])
// CHECK: [[INITB:%.+]] = linalg.init_tensor [5, 6] // CHECK: [[INITB:%.+]] = tensor.empty()
// CHECK: [[ONE:%.+]] = arith.constant 1 // CHECK: [[ONE:%.+]] = arith.constant 1
// CHECK: [[TWO:%.+]] = arith.constant 2 // CHECK: [[TWO:%.+]] = arith.constant 2
// CHECK: [[MATMUL:%.+]] = linalg.quantized_matmul ins(%arg0, [[TRANSPOSE]], [[ONE]], [[TWO]] : tensor<5x3xi8>, tensor<3x6xi8>, i32, i32) outs([[FILL]] : tensor<5x6xi32>) -> tensor<5x6xi32> // CHECK: [[MATMUL:%.+]] = linalg.quantized_matmul ins(%arg0, [[TRANSPOSE]], [[ONE]], [[TWO]] : tensor<5x3xi8>, tensor<3x6xi8>, i32, i32) outs([[FILL]] : tensor<5x6xi32>) -> tensor<5x6xi32>
// CHECK: [[ADDED:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, [[MATMUL]] : tensor<6xi32>, tensor<5x6xi32>) outs([[INITB]] // CHECK: [[ADDED:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, [[MATMUL]] : tensor<6xi32>, tensor<5x6xi32>) outs([[INITB]]
// CHECK: ^bb0([[IN1:%.+]]: i32, [[IN2:%.+]]: i32, [[UNUSED:%.+]]: i32): // CHECK: ^bb0([[IN1:%.+]]: i32, [[IN2:%.+]]: i32, [[UNUSED:%.+]]: i32):
// CHECK: [[ADD:%.+]] = arith.addi // CHECK: [[ADD:%.+]] = arith.addi
// CHECK: linalg.yield [[ADD]] : i32 // CHECK: linalg.yield [[ADD]] : i32
%0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) {quantization_info = #tosa.conv_quant<input_zp = 1, weight_zp = 2>} : (tensor<5x3xi8>, tensor<6x3xi8>, tensor<6xi32>) -> (tensor<5x6xi32>) %0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) {quantization_info = #tosa.conv_quant<input_zp = 1, weight_zp = 2>} : (tensor<5x3xi8>, tensor<6x3xi8>, tensor<6xi32>) -> (tensor<5x6xi32>)
return %0 : tensor<5x6xi32> return %0 : tensor<5x6xi32>
} }
// ----- // -----
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @fully_connected_dyn // CHECK-LABEL: @fully_connected_dyn
func.func @fully_connected_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<6x3xf32>, %arg2: tensor<6xf32>) -> (tensor<?x6xf32>) { func.func @fully_connected_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<6x3xf32>, %arg2: tensor<6xf32>) -> (tensor<?x6xf32>) {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INITT:.+]] = linalg.init_tensor [%[[DIM]], 6] // CHECK: %[[INITT:.+]] = tensor.empty(%[[DIM]])
// CHECK: %[[ZERO:.+]] = arith.constant 0 // CHECK: %[[ZERO:.+]] = arith.constant 0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]]{{.*}}outs(%[[INITT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ZERO]]{{.*}}outs(%[[INITT]]
// CHECK: %[[PERM:.+]] = arith.constant dense<[1, 0]> // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 0]>
// CHECK: %[[TRANSPOSE:.+]] = "tosa.transpose"(%arg1, %[[PERM]]) // CHECK: %[[TRANSPOSE:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
// CHECK: %[[INITB:.+]] = linalg.init_tensor [%[[DIM]], 6] // CHECK: %[[INITB:.+]] = tensor.empty(%[[DIM]])
// CHECK: %[[MATMUL:.+]] = linalg.matmul ins(%arg0, %[[TRANSPOSE]] : tensor<?x3xf32>, tensor<3x6xf32>) outs(%[[FILL]] : tensor<?x6xf32>) -> tensor<?x6xf32> // CHECK: %[[MATMUL:.+]] = linalg.matmul ins(%arg0, %[[TRANSPOSE]] : tensor<?x3xf32>, tensor<3x6xf32>) outs(%[[FILL]] : tensor<?x6xf32>) -> tensor<?x6xf32>
// CHECK: %[[ADDED:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, %[[MATMUL]] : tensor<6xf32>, tensor<?x6xf32>) outs(%[[INITB]] : tensor<?x6xf32>) { // CHECK: %[[ADDED:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel"]} ins(%arg2, %[[MATMUL]] : tensor<6xf32>, tensor<?x6xf32>) outs(%[[INITB]] : tensor<?x6xf32>) {
// CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
// CHECK: %[[ADD:.+]] = arith.addf %arg3, %arg4 : f32 // CHECK: %[[ADD:.+]] = arith.addf %arg3, %arg4 : f32
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
%0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) : (tensor<?x3xf32>, tensor<6x3xf32>, tensor<6xf32>) -> (tensor<?x6xf32>) %0 = "tosa.fully_connected"(%arg0, %arg1, %arg2) : (tensor<?x3xf32>, tensor<6x3xf32>, tensor<6xf32>) -> (tensor<?x6xf32>)
return %0 : tensor<?x6xf32> return %0 : tensor<?x6xf32>
} }
// ----- // -----
// CHECK-LABEL: @max_pool // CHECK-LABEL: @max_pool
func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () { func.func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
// CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 // CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38
// CHECK-DAG: [[INIT:%.+]] = linalg.init_tensor [1, 4, 32, 62] // CHECK-DAG: [[INIT:%.+]] = tensor.empty()
// CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[INIT]] // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[INIT]]
// CHECK-DAG: [[KERNEL:%.+]] = linalg.init_tensor [3, 3] // CHECK-DAG: [[KERNEL:%.+]] = tensor.empty()
// CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, [[KERNEL]] : tensor<1x6x34x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x32x62xf32>) // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, [[KERNEL]] : tensor<1x6x34x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x32x62xf32>)
%0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>) -> (tensor<1x4x32x62xf32>) %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>) -> (tensor<1x4x32x62xf32>)
return return
} }
// CHECK-LABEL: @max_pool_padded // CHECK-LABEL: @max_pool_padded
func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () { func.func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
// CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 : f32 // CHECK-DAG: [[CONST:%.+]] = arith.constant -3.40282347E+38 : f32
// CHECK-DAG: [[PAD:%.+]] = tensor.pad %arg0 low[0, 0, 0, 0] high[0, 0, 1, 0] // CHECK-DAG: [[PAD:%.+]] = tensor.pad %arg0 low[0, 0, 0, 0] high[0, 0, 1, 0]
// CHECK-DAG: tensor.yield [[CONST]] // CHECK-DAG: tensor.yield [[CONST]]
// CHECK-DAG: [[INITVAL:%.+]] = arith.constant -3.40282347E+38 : f32 // CHECK-DAG: [[INITVAL:%.+]] = arith.constant -3.40282347E+38 : f32
// CHECK-DAG: [[INIT:%.+]] = linalg.init_tensor [1, 4, 33, 62] // CHECK-DAG: [[INIT:%.+]] = tensor.empty()
// CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INITVAL]]{{.*}}outs([[INIT]] // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INITVAL]]{{.*}}outs([[INIT]]
// CHECK-DAG: [[KERNEL:%.+]] = linalg.init_tensor [3, 3] // CHECK-DAG: [[KERNEL:%.+]] = tensor.empty()
// CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x6x35x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x33x62xf32>) // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x6x35x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x33x62xf32>)
%0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 1], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>) -> (tensor<1x4x33x62xf32>) %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 1], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>) -> (tensor<1x4x33x62xf32>)
return return
} }
// CHECK-LABEL: @max_pool_dyn // CHECK-LABEL: @max_pool_dyn
func.func @max_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> () { func.func @max_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[CONST:.+]] = arith.constant -3.40282347E+38 // CHECK: %[[CONST:.+]] = arith.constant -3.40282347E+38
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 4, 32, 62] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}outs(%[[INIT]]
// CHECK: %[[KERNEL:.+]] = linalg.init_tensor [3, 3] // CHECK: %[[KERNEL:.+]] = tensor.empty()
// CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x6x34x62xf32>, tensor<3x3xf32>) outs(%[[FILL]] : tensor<?x4x32x62xf32>) // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x6x34x62xf32>, tensor<3x3xf32>) outs(%[[FILL]] : tensor<?x4x32x62xf32>)
%0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<?x6x34x62xf32>) -> (tensor<?x4x32x62xf32>) %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<?x6x34x62xf32>) -> (tensor<?x4x32x62xf32>)
return return
} }
// CHECK-LABEL: @max_pool_i8 // CHECK-LABEL: @max_pool_i8
func.func @max_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> () { func.func @max_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> () {
// CHECK: arith.constant -128 // CHECK: arith.constant -128
Show All 20 Lines
// ----- // -----
// CHECK-LABEL: @avg_pool // CHECK-LABEL: @avg_pool
func.func @avg_pool(%arg0: tensor<1x6x34x62xf32>) -> (tensor<1x5x33x62xf32>) { func.func @avg_pool(%arg0: tensor<1x6x34x62xf32>) -> (tensor<1x5x33x62xf32>) {
// Initial piece computes the sum of the pooling region, with appropriate padding. // Initial piece computes the sum of the pooling region, with appropriate padding.
// CHECK: [[CONST:%.+]] = arith.constant 0 // CHECK: [[CONST:%.+]] = arith.constant 0
// CHECK: [[PAD:%.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] // CHECK: [[PAD:%.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
// CHECK: [[CONST:%.+]] = arith.constant 0 // CHECK: [[CONST:%.+]] = arith.constant 0
// CHECK: [[POOLINIT:%.+]] = linalg.init_tensor [1, 5, 33, 62] // CHECK: [[POOLINIT:%.+]] = tensor.empty()
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[POOLINIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[POOLINIT]]
// CHECK: [[KERNEL:%.+]] = linalg.init_tensor [4, 4] // CHECK: [[KERNEL:%.+]] = tensor.empty()
// CHECK: [[POOL:%.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x8x36x62xf32>, tensor<4x4xf32>) outs([[FILL]] : tensor<1x5x33x62xf32>) // CHECK: [[POOL:%.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x8x36x62xf32>, tensor<4x4xf32>) outs([[FILL]] : tensor<1x5x33x62xf32>)
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 33, 62] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins([[POOL]] : tensor<1x5x33x62xf32>) outs([[INIT]] : tensor<1x5x33x62xf32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins([[POOL]] : tensor<1x5x33x62xf32>) outs([[INIT]] : tensor<1x5x33x62xf32>)
// CHECK: [[ZERO:%.0]] = arith.constant 0 // CHECK: [[ZERO:%.0]] = arith.constant 0
// CHECK: [[ONE:%.+]] = arith.constant 1 // CHECK: [[ONE:%.+]] = arith.constant 1
// CHECK: [[HEIGHT:%.+]] = arith.constant 4 // CHECK: [[HEIGHT:%.+]] = arith.constant 4
// CHECK: [[WIDTH:%.+]] = arith.constant 32 // CHECK: [[WIDTH:%.+]] = arith.constant 32
// CHECK: [[IDX1:%.+]] = linalg.index 1 // CHECK: [[IDX1:%.+]] = linalg.index 1
// CHECK: [[IDX2:%.+]] = linalg.index 2 // CHECK: [[IDX2:%.+]] = linalg.index 2
Show All 40 Lines
// ----- // -----
// CHECK-LABEL: @avg_pool_dyn // CHECK-LABEL: @avg_pool_dyn
func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>) { func.func @avg_pool_dyn(%arg0: tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>) {
// The calculations remain the same as above, only testing for dyn behavior // The calculations remain the same as above, only testing for dyn behavior
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[PAD:.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] // CHECK: %[[PAD:.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
// CHECK: %[[POOLINIT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 33, 62] // CHECK: %[[POOLINIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[KERNEL:.+]] = linalg.init_tensor [4, 4] // CHECK: %[[KERNEL:.+]] = tensor.empty()
// CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%[[PAD]], %[[KERNEL]] : tensor<?x8x36x62xf32>, tensor<4x4xf32>) outs(%[[FILL]] : tensor<?x5x33x62xf32>) // CHECK: %[[POOL:.+]] = linalg.pooling_nhwc_sum {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%[[PAD]], %[[KERNEL]] : tensor<?x8x36x62xf32>, tensor<4x4xf32>) outs(%[[FILL]] : tensor<?x5x33x62xf32>)
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 33, 62] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[POOL]] : tensor<?x5x33x62xf32>) outs(%[[INIT]] : tensor<?x5x33x62xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[POOL]] : tensor<?x5x33x62xf32>) outs(%[[INIT]] : tensor<?x5x33x62xf32>)
%0 = "tosa.avg_pool2d"(%arg0) {pad = [1, 1, 1, 1], kernel = [4, 4], stride = [1, 1]} : (tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>) %0 = "tosa.avg_pool2d"(%arg0) {pad = [1, 1, 1, 1], kernel = [4, 4], stride = [1, 1]} : (tensor<?x6x34x62xf32>) -> (tensor<?x5x33x62xf32>)
return %0 : tensor<?x5x33x62xf32> return %0 : tensor<?x5x33x62xf32>
} }
// ----- // -----
// CHECK-LABEL: @avg_pool_i8 // CHECK-LABEL: @avg_pool_i8
▲ Show 20 Lines • Show All 56 Lines • ▼ Show 20 Lines
// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @conv2d_f32 // CHECK-LABEL: @conv2d_f32
func.func @conv2d_f32(%input: tensor<1x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () { func.func @conv2d_f32(%input: tensor<1x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
// CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]> // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]>
// CHECK: %[[W:.+]] = "tosa.transpose"(%arg1, %[[PERM]]) // CHECK: %[[W:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
// CHECK: %[[M_IN:.+]] = linalg.init_tensor [1, 45, 40, 28] // CHECK: %[[M_IN:.+]] = tensor.empty()
// CHECK: %[[CST:.+]] = arith.constant 0 // CHECK: %[[CST:.+]] = arith.constant 0
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[B_IN:.+]] = linalg.init_tensor [1, 45, 40, 28] // CHECK: %[[B_IN:.+]] = tensor.empty()
// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[W]] : tensor<1x49x42x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<1x45x40x28xf32>) // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[W]] : tensor<1x49x42x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<1x45x40x28xf32>)
// CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<1x45x40x28xf32>) outs(%[[B_IN]] : tensor<1x45x40x28xf32>) // CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<1x45x40x28xf32>) outs(%[[B_IN]] : tensor<1x45x40x28xf32>)
// CHECK: arith.addf // CHECK: arith.addf
// CHECK: linalg.yield // CHECK: linalg.yield
%0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<1x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<1x45x40x28xf32>) %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<1x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<1x45x40x28xf32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @conv2d_dyn // CHECK-LABEL: @conv2d_dyn
func.func @conv2d_dyn(%input: tensor<?x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () { func.func @conv2d_dyn(%input: tensor<?x49x42x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]> // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]>
// CHECK: %[[W:.+]] = "tosa.transpose"(%arg1, %[[PERM]]) // CHECK: %[[W:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
// CHECK: %[[M_IN:.+]] = linalg.init_tensor [%[[BATCH]], 45, 40, 28] // CHECK: %[[M_IN:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[CST:.+]] = arith.constant 0 // CHECK: %[[CST:.+]] = arith.constant 0
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[B_IN:.+]] = linalg.init_tensor [%[[BATCH]], 45, 40, 28] // CHECK: %[[B_IN:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[W]] : tensor<?x49x42x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<?x45x40x28xf32>) // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[W]] : tensor<?x49x42x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<?x45x40x28xf32>)
// CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<?x45x40x28xf32>) outs(%[[B_IN]] : tensor<?x45x40x28xf32>) // CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<?x45x40x28xf32>) outs(%[[B_IN]] : tensor<?x45x40x28xf32>)
// CHECK: %[[ADD:.+]] = arith.addf // CHECK: %[[ADD:.+]] = arith.addf
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
%0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<?x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<?x45x40x28xf32>) %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<?x49x42x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<?x45x40x28xf32>)
return return
} }
Show All 40 Linesfunc.func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
// CHECK: %[[SUBTRACTED_0:.+]] = arith.subi %[[ADD_PAD_3]], %[[ADD_ONE_0]] : index // CHECK: %[[SUBTRACTED_0:.+]] = arith.subi %[[ADD_PAD_3]], %[[ADD_ONE_0]] : index
// CHECK: %[[STRIDE_W:.+]] = arith.constant 1 : index // CHECK: %[[STRIDE_W:.+]] = arith.constant 1 : index
// CHECK: %[[DIVIDED_0:.+]] = arith.divui %[[SUBTRACTED_0]], %[[STRIDE_W]] : index // CHECK: %[[DIVIDED_0:.+]] = arith.divui %[[SUBTRACTED_0]], %[[STRIDE_W]] : index
// CHECK: %[[W_OUT:.+]] = arith.addi %[[DIVIDED_0]], %[[ONE_0]] : index // CHECK: %[[W_OUT:.+]] = arith.addi %[[DIVIDED_0]], %[[ONE_0]] : index
// Running convolution // Running convolution
// CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]> // CHECK: %[[PERM:.+]] = arith.constant dense<[1, 2, 3, 0]>
// CHECK: %[[WEIGHT:.+]] = "tosa.transpose"(%arg1, %[[PERM]]) // CHECK: %[[WEIGHT:.+]] = "tosa.transpose"(%arg1, %[[PERM]])
// CHECK: %[[M_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28] // CHECK: %[[M_IN:.+]] = tensor.empty(%[[H_OUT]], %[[W_OUT]])
// CHECK: %[[CST:.+]] = arith.constant 0 // CHECK: %[[CST:.+]] = arith.constant 0
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[B_IN:.+]] = linalg.init_tensor [1, %[[H_OUT]], %[[W_OUT]], 28] // CHECK: %[[B_IN:.+]] = tensor.empty(%[[H_OUT]], %[[W_OUT]])
// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[WEIGHT]] : tensor<1x?x?x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<1x?x?x28xf32>) // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %[[WEIGHT]] : tensor<1x?x?x27xf32>, tensor<3x3x27x28xf32>) outs(%[[FILL]] : tensor<1x?x?x28xf32>)
// CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<1x?x?x28xf32>) outs(%[[B_IN]] : tensor<1x?x?x28xf32>) // CHECK: %[[B:.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[CONV]] : tensor<28xf32>, tensor<1x?x?x28xf32>) outs(%[[B_IN]] : tensor<1x?x?x28xf32>)
// CHECK: %[[ADD:.+]] = arith.addf // CHECK: %[[ADD:.+]] = arith.addf
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
%0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<1x?x?x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<1x?x?x28xf32>) %0 = "tosa.conv2d"(%input, %weights, %bias) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [2, 1]} : (tensor<1x?x?x27xf32>, tensor<28x3x3x27xf32>, tensor<28xf32>) -> (tensor<1x?x?x28xf32>)
return return
} }
Show All 23 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @depthwise_conv // CHECK-LABEL: @depthwise_conv
func.func @depthwise_conv(%arg0 : tensor<1x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () { func.func @depthwise_conv(%arg0 : tensor<1x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 5, 3, 11] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 5, 5, 33] // CHECK: [[OUT:%.+]] = tensor.empty()
// CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>) // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>)
// CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
// CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) { // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) {
// CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
// CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32 // CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
// CHECK: linalg.yield [[ADD]] : f32 // CHECK: linalg.yield [[ADD]] : f32
// CHECK: } -> tensor<1x5x5x33xf32> // CHECK: } -> tensor<1x5x5x33xf32>
%2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<1x5x5x33xf32>) %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<1x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<1x5x5x33xf32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @depthwise_conv_dyn // CHECK-LABEL: @depthwise_conv_dyn
func.func @depthwise_conv_dyn(%arg0 : tensor<?x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () { func.func @depthwise_conv_dyn(%arg0 : tensor<?x7x5x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 5, 3, 11] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[CST0:.+]] = arith.constant 0 // CHECK: %[[CST0:.+]] = arith.constant 0
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[OUT:.+]] = linalg.init_tensor [%[[BATCH]], 5, 5, 33] // CHECK: %[[OUT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[DEPTH:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>) outs(%[[FILL]] : tensor<?x5x5x3x11xf32>) // CHECK: %[[DEPTH:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>) outs(%[[FILL]] : tensor<?x5x5x3x11xf32>)
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[DEPTH]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
// CHECK: %[[BIAS:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[COLLAPSED]] : tensor<33xf32>, tensor<?x5x5x33xf32>) outs(%[[OUT]] : tensor<?x5x5x33xf32>) { // CHECK: %[[BIAS:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, %[[COLLAPSED]] : tensor<33xf32>, tensor<?x5x5x33xf32>) outs(%[[OUT]] : tensor<?x5x5x33xf32>) {
// CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
// CHECK: %[[ADD:.+]] = arith.addf %arg3, %arg4 : f32 // CHECK: %[[ADD:.+]] = arith.addf %arg3, %arg4 : f32
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
// CHECK: } -> tensor<?x5x5x33xf32> // CHECK: } -> tensor<?x5x5x33xf32>
%2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<?x5x5x33xf32>) %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1] } : (tensor<?x7x5x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<?x5x5x33xf32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @depthwise_conv_strides // CHECK-LABEL: @depthwise_conv_strides
func.func @depthwise_conv_strides(%arg0 : tensor<1x11x9x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () { func.func @depthwise_conv_strides(%arg0 : tensor<1x11x9x3xf32>, %arg1 : tensor<3x1x3x11xf32>, %arg2 : tensor<33xf32>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 5, 5, 3, 11] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 5, 5, 33] // CHECK: [[OUT:%.+]] = tensor.empty()
// CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>) // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>) outs([[FILL]] : tensor<1x5x5x3x11xf32>)
// CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
// CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) { // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<33xf32>, tensor<1x5x5x33xf32>) outs([[OUT]] : tensor<1x5x5x33xf32>) {
// CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32): // CHECK: ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
// CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32 // CHECK: [[ADD:%.+]] = arith.addf %arg3, %arg4 : f32
// CHECK: linalg.yield [[ADD]] : f32 // CHECK: linalg.yield [[ADD]] : f32
// CHECK: } -> tensor<1x5x5x33xf32> // CHECK: } -> tensor<1x5x5x33xf32>
%2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [2, 2], dilation = [1, 1] } : (tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<1x5x5x33xf32>) %2 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) { pad = [0, 0, 0, 0], stride = [2, 2], dilation = [1, 1] } : (tensor<1x11x9x3xf32>, tensor<3x1x3x11xf32>, tensor<33xf32>) -> (tensor<1x5x5x33xf32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @depthwise_conv_quant // CHECK-LABEL: @depthwise_conv_quant
func.func @depthwise_conv_quant(%arg0 : tensor<1x12x12x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () { func.func @depthwise_conv_quant(%arg0 : tensor<1x12x12x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () {
// CHECK: [[PADV:%.+]] = arith.constant -128 // CHECK: [[PADV:%.+]] = arith.constant -128
// CHECK: [[PAD:%.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0] // CHECK: [[PAD:%.+]] = tensor.pad %arg0 low[0, 1, 1, 0] high[0, 1, 1, 0]
// CHECK: tensor.yield [[PADV]] // CHECK: tensor.yield [[PADV]]
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 12, 12, 4, 128] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 12, 12, 512] // CHECK: [[OUT:%.+]] = tensor.empty()
// CHECK: [[C128:%.+]] = arith.constant -128 // CHECK: [[C128:%.+]] = arith.constant -128
// CHECK: [[C42:%.+]] = arith.constant 42 // CHECK: [[C42:%.+]] = arith.constant 42
// CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins([[PAD]], %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x12x12x4x128xi32>) // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins([[PAD]], %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x12x12x4x128xi32>)
// CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
// CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x12x12x512xi32>) outs([[OUT]] : tensor<1x12x12x512xi32>) { // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x12x12x512xi32>) outs([[OUT]] : tensor<1x12x12x512xi32>) {
// CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32): // CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32):
// CHECK: [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32 // CHECK: [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32
// CHECK: linalg.yield [[ADD]] : i32 // CHECK: linalg.yield [[ADD]] : i32
// CHECK: } -> tensor<1x12x12x512xi32> // CHECK: } -> tensor<1x12x12x512xi32>
%0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 1, 1, 1], quantization_info = #tosa.conv_quant<input_zp = -128, weight_zp = 42>, stride = [1, 1], dilation = [1, 1] } : (tensor<1x12x12x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>) -> tensor<1x12x12x512xi32> %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 1, 1, 1], quantization_info = #tosa.conv_quant<input_zp = -128, weight_zp = 42>, stride = [1, 1], dilation = [1, 1] } : (tensor<1x12x12x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>) -> tensor<1x12x12x512xi32>
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d3)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @depthwise_conv_quant_dilations // CHECK-LABEL: @depthwise_conv_quant_dilations
func.func @depthwise_conv_quant_dilations(%arg0 : tensor<1x14x14x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () { func.func @depthwise_conv_quant_dilations(%arg0 : tensor<1x14x14x4xi8>, %arg1 : tensor<3x3x4x128xi8>, %arg2 : tensor<512xi32>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 10, 10, 4, 128] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[OUT:%.+]] = linalg.init_tensor [1, 10, 10, 512] // CHECK: [[OUT:%.+]] = tensor.empty()
// CHECK: [[C128:%.+]] = arith.constant -128 // CHECK: [[C128:%.+]] = arith.constant -128
// CHECK: [[C42:%.+]] = arith.constant 42 // CHECK: [[C42:%.+]] = arith.constant 42
// CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x10x10x4x128xi32>) // CHECK: [[DEPTH:%.+]] = linalg.depthwise_conv_2d_nhwc_hwcm_q {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1, [[C128]], [[C42]] : tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, i32, i32) outs([[FILL]] : tensor<1x10x10x4x128xi32>)
// CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: [[COLLAPSED:%.+]] = tensor.collapse_shape [[DEPTH]] {{\[}}[0], [1], [2], [3, 4]]
// CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x10x10x512xi32>) outs([[OUT]] : tensor<1x10x10x512xi32>) { // CHECK: [[BIAS:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2, [[COLLAPSED]] : tensor<512xi32>, tensor<1x10x10x512xi32>) outs([[OUT]] : tensor<1x10x10x512xi32>) {
// CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32): // CHECK: ^bb0(%arg3: i32, %arg4: i32, %arg5: i32):
// CHECK: [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32 // CHECK: [[ADD:%.+]] = arith.addi %arg3, %arg4 : i32
// CHECK: linalg.yield [[ADD]] : i32 // CHECK: linalg.yield [[ADD]] : i32
Show All 20 Lines

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

// RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg)" %s -verify-diagnostics -o -| FileCheck %s // RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg)" %s -verify-diagnostics -o -| FileCheck %s
// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()> // CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
// CHECK-LABEL: @test_abs // CHECK-LABEL: @test_abs
func.func @test_abs(%arg0: tensor<f32>) -> tensor<f32> { func.func @test_abs(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [] : tensor<f32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<f32>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = []} ins(%arg0 : tensor<f32>) outs([[INIT]] : tensor<f32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = []} ins(%arg0 : tensor<f32>) outs([[INIT]] : tensor<f32>) {
// CHECK: ^bb0(%arg1: f32, %arg2: f32): // CHECK: ^bb0(%arg1: f32, %arg2: f32):
// CHECK: [[ELEMENT:%.+]] = math.absf %arg1 // CHECK: [[ELEMENT:%.+]] = math.absf %arg1
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<f32> // CHECK: } -> tensor<f32>
%0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32> %0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
// CHECK: return [[GENERIC]] // CHECK: return [[GENERIC]]
return %0 : tensor<f32> return %0 : tensor<f32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_abs // CHECK-LABEL: @test_abs
func.func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> { func.func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2xf32>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
// CHECK: ^bb0(%arg1: f32, %arg2: f32): // CHECK: ^bb0(%arg1: f32, %arg2: f32):
// CHECK: [[ELEMENT:%.+]] = math.absf %arg1 // CHECK: [[ELEMENT:%.+]] = math.absf %arg1
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<2xf32> // CHECK: } -> tensor<2xf32>
%0 = "tosa.abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32> %0 = "tosa.abs"(%arg0) : (tensor<2xf32>) -> tensor<2xf32>
// CHECK: return [[GENERIC]] // CHECK: return [[GENERIC]]
return %0 : tensor<2xf32> return %0 : tensor<2xf32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @test_abs // CHECK-LABEL: @test_abs
func.func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { func.func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2x3xf32>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<2x3xf32>) outs([[INIT]] : tensor<2x3xf32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<2x3xf32>) outs([[INIT]] : tensor<2x3xf32>) {
// CHECK: ^bb0(%arg1: f32, %arg2: f32): // CHECK: ^bb0(%arg1: f32, %arg2: f32):
// CHECK: [[ELEMENT:%.+]] = math.absf %arg1 // CHECK: [[ELEMENT:%.+]] = math.absf %arg1
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<2x3xf32> // CHECK: } -> tensor<2x3xf32>
%0 = "tosa.abs"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32> %0 = "tosa.abs"(%arg0) : (tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK: return [[GENERIC]] // CHECK: return [[GENERIC]]
return %0 : tensor<2x3xf32> return %0 : tensor<2x3xf32>
} }
// ----- // -----
// CHECK-LABEL: @test_abs // CHECK-LABEL: @test_abs
func.func @test_abs(%arg0: tensor<?xf32>) -> tensor<?xf32> { func.func @test_abs(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK: math.absf // CHECK: math.absf
%0 = "tosa.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32> %0 = "tosa.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @test_abs_dyn // CHECK-LABEL: @test_abs_dyn
func.func @test_abs_dyn(%arg0: tensor<2x?xf32>) -> tensor<2x?xf32> { func.func @test_abs_dyn(%arg0: tensor<2x?xf32>) -> tensor<2x?xf32> {
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [2, %[[DIM]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]])
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK: math.absf // CHECK: math.absf
%0 = "tosa.abs"(%arg0) : (tensor<2x?xf32>) -> tensor<2x?xf32> %0 = "tosa.abs"(%arg0) : (tensor<2x?xf32>) -> tensor<2x?xf32>
return %0 : tensor<2x?xf32> return %0 : tensor<2x?xf32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> ()> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> ()>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_broadcast // CHECK-LABEL: @test_broadcast
func.func @test_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> { func.func @test_broadcast(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2xf32>
// CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0 // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg0
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins([[RESHAPE]], %arg1 : tensor<f32>, tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins([[RESHAPE]], %arg1 : tensor<f32>, tensor<2xf32>) outs([[INIT]] : tensor<2xf32>) {
// CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
// CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32 // CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<2xf32> // CHECK: } -> tensor<2xf32>
%0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32> %0 = "tosa.add"(%arg0, %arg1) : (tensor<1xf32>, tensor<2xf32>) -> tensor<2xf32>
return %0 : tensor<2xf32> return %0 : tensor<2xf32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> ()> // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> ()>
// CHECK-LABEL: @test_broadcast_swapped_args // CHECK-LABEL: @test_broadcast_swapped_args
func.func @test_broadcast_swapped_args(%arg0: tensor<2xf32>, %arg1: tensor<1xf32>) -> tensor<2xf32> { func.func @test_broadcast_swapped_args(%arg0: tensor<2xf32>, %arg1: tensor<1xf32>) -> tensor<2xf32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] : tensor<2xf32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2xf32>
// CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg1 // CHECK: [[RESHAPE:%.+]] = tensor.collapse_shape %arg1
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[RESHAPE]] : tensor<2xf32>, tensor<f32>) outs([[INIT]] : tensor<2xf32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[RESHAPE]] : tensor<2xf32>, tensor<f32>) outs([[INIT]] : tensor<2xf32>) {
// CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
// CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32 // CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<2xf32> // CHECK: } -> tensor<2xf32>
%0 = "tosa.add"(%arg0, %arg1) : (tensor<2xf32>, tensor<1xf32>) -> tensor<2xf32> %0 = "tosa.add"(%arg0, %arg1) : (tensor<2xf32>, tensor<1xf32>) -> tensor<2xf32>
return %0 : tensor<2xf32> return %0 : tensor<2xf32>
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: @test_multibroadcast // CHECK-LABEL: @test_multibroadcast
func.func @test_multibroadcast(%arg0: tensor<1x3xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x3xf32> { func.func @test_multibroadcast(%arg0: tensor<1x3xf32>, %arg1: tensor<2x1xf32>) -> tensor<2x3xf32> {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3] : tensor<2x3xf32> // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2x3xf32>
// CHECK: [[RESHAPE1:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]] // CHECK: [[RESHAPE1:%.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1]]
// CHECK: [[RESHAPE2:%.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1]] // CHECK: [[RESHAPE2:%.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins([[RESHAPE1]], [[RESHAPE2]] : tensor<3xf32>, tensor<2xf32>) outs([[INIT]] : tensor<2x3xf32>) { // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins([[RESHAPE1]], [[RESHAPE2]] : tensor<3xf32>, tensor<2xf32>) outs([[INIT]] : tensor<2x3xf32>) {
// CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32): // CHECK: ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
// CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32 // CHECK: [[ELEMENT:%.+]] = arith.addf %arg2, %arg3 : f32
// CHECK: linalg.yield [[ELEMENT]] : f32 // CHECK: linalg.yield [[ELEMENT]] : f32
// CHECK: } -> tensor<2x3xf32> // CHECK: } -> tensor<2x3xf32>
%0 = "tosa.add"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32> %0 = "tosa.add"(%arg0, %arg1) : (tensor<1x3xf32>, tensor<2x1xf32>) -> tensor<2x3xf32>
▲ Show 20 Lines • Show All 488 Lines • ▼ Show 20 Lines
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: @test_transpose // CHECK-LABEL: @test_transpose
// CHECK-SAME: ([[ARG0:%.+]]: tensor<1x2x3xi32>) // CHECK-SAME: ([[ARG0:%.+]]: tensor<1x2x3xi32>)
func.func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () { func.func @test_transpose(%arg0: tensor<1x2x3xi32>) -> () {
%0 = arith.constant dense<[1, 2, 0]> : tensor<3xi32> %0 = arith.constant dense<[1, 2, 0]> : tensor<3xi32>
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 3, 1] // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<2x3x1xi32>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel"]} ins([[ARG0]] : tensor<1x2x3xi32>) outs([[OUT:%.+]] : tensor<2x3x1xi32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel"]} ins([[ARG0]] : tensor<1x2x3xi32>) outs([[OUT:%.+]] : tensor<2x3x1xi32>)
// CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32) // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
// CHECK: linalg.yield [[ARG1]] // CHECK: linalg.yield [[ARG1]]
// CHECK: } // CHECK: }
%1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>) %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x2x3xi32>, tensor<3xi32>) -> (tensor<2x3x1xi32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @test_transpose_dyn // CHECK-LABEL: @test_transpose_dyn
// CHECK-SAME: (%[[ARG0:.+]]: tensor<1x?x3x4xi32>) // CHECK-SAME: (%[[ARG0:.+]]: tensor<1x?x3x4xi32>)
func.func @test_transpose_dyn(%arg0: tensor<1x?x3x4xi32>) -> () { func.func @test_transpose_dyn(%arg0: tensor<1x?x3x4xi32>) -> () {
%0 = arith.constant dense<[1, 3, 0, 2]> : tensor<4xi32> %0 = arith.constant dense<[1, 3, 0, 2]> : tensor<4xi32>
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM]], 4, 1, 3] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM]]) : tensor<?x4x1x3xi32>
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[ARG0]] : tensor<1x?x3x4xi32>) outs([[OUT:%.+]] : tensor<?x4x1x3xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%[[ARG0]] : tensor<1x?x3x4xi32>) outs([[OUT:%.+]] : tensor<?x4x1x3xi32>)
// CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32) // CHECK: ^bb0([[ARG1:%.+]]: i32, [[ARG2:%.+]]: i32)
// CHECK: linalg.yield [[ARG1]] // CHECK: linalg.yield [[ARG1]]
// CHECK: } // CHECK: }
%1 = "tosa.transpose"(%arg0, %0) : (tensor<1x?x3x4xi32>, tensor<4xi32>) -> (tensor<?x4x1x3xi32>) %1 = "tosa.transpose"(%arg0, %0) : (tensor<1x?x3x4xi32>, tensor<4xi32>) -> (tensor<?x4x1x3xi32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @test_transpose_dyn // CHECK-LABEL: @test_transpose_dyn
// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?xf32>) // CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?xf32>)
func.func @test_transpose_dyn_multiple(%arg0: tensor<?x?xf32>) -> () { func.func @test_transpose_dyn_multiple(%arg0: tensor<?x?xf32>) -> () {
%0 = arith.constant dense<[1, 0]> : tensor<2xi32> %0 = arith.constant dense<[1, 0]> : tensor<2xi32>
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DIM0:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DIM1]], %[[DIM0]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM1]], %[[DIM0]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%[[ARG0]] : tensor<?x?xf32>) outs([[OUT:%.+]] : tensor<?x?xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%[[ARG0]] : tensor<?x?xf32>) outs([[OUT:%.+]] : tensor<?x?xf32>)
// CHECK: ^bb0([[ARG1:%.+]]: f32, [[ARG2:%.+]]: f32) // CHECK: ^bb0([[ARG1:%.+]]: f32, [[ARG2:%.+]]: f32)
// CHECK: linalg.yield [[ARG1]] // CHECK: linalg.yield [[ARG1]]
// CHECK: } // CHECK: }
%1 = "tosa.transpose"(%arg0, %0) : (tensor<?x?xf32>, tensor<2xi32>) -> (tensor<?x?xf32>) %1 = "tosa.transpose"(%arg0, %0) : (tensor<?x?xf32>, tensor<2xi32>) -> (tensor<?x?xf32>)
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: @reduce_float // CHECK-LABEL: @reduce_float
// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xf32> // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xf32>
func.func @reduce_float(%arg0: tensor<5x4xf32>) -> () { func.func @reduce_float(%arg0: tensor<5x4xf32>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [4] // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<4xf32>
// CHECK: [[CST0:%.+]] = arith.constant 0.0 // CHECK: [[CST0:%.+]] = arith.constant 0.0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<4xf32>)
// CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: ^bb0(%arg1: f32, %arg2: f32)
// CHECK: [[RES:%.+]] = arith.addf %arg1, %arg2 : f32 // CHECK: [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
// CHECK: linalg.yield [[RES]] : f32 // CHECK: linalg.yield [[RES]] : f32
// CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xf32> into tensor<1x4xf32> // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xf32> into tensor<1x4xf32>
%0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32> %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xf32>) -> tensor<1x4xf32>
// CHECK: [[INIT:%.+]] = linalg.init_tensor [5] // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<5xf32>
// CHECK: [[CST0:%.+]] = arith.constant 0.0 // CHECK: [[CST0:%.+]] = arith.constant 0.0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xf32>) outs([[FILL]] : tensor<5xf32>)
// CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: ^bb0(%arg1: f32, %arg2: f32)
// CHECK: [[RES:%.+]] = arith.addf %arg1, %arg2 : f32 // CHECK: [[RES:%.+]] = arith.addf %arg1, %arg2 : f32
// CHECK: linalg.yield [[RES]] : f32 // CHECK: linalg.yield [[RES]] : f32
// CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xf32> into tensor<5x1xf32> // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
%1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5x1xf32> %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xf32>) -> tensor<5x1xf32>
Show All 22 Lines
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-LABEL: @reduce_float_dyn // CHECK-LABEL: @reduce_float_dyn
func.func @reduce_float_dyn(%arg0: tensor<?x5x4xf32>) -> () { func.func @reduce_float_dyn(%arg0: tensor<?x5x4xf32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DYN]], 4] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<?x4xf32>
// CHECK: %[[CST0:.+]] = arith.constant 0.0 // CHECK: %[[CST0:.+]] = arith.constant 0.0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]]{{.*}}outs(%[[INIT]]
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "reduction", "parallel"]} ins(%arg0 : tensor<?x5x4xf32>) outs(%[[FILL]] : tensor<?x4xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "reduction", "parallel"]} ins(%arg0 : tensor<?x5x4xf32>) outs(%[[FILL]] : tensor<?x4xf32>)
// CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: ^bb0(%arg1: f32, %arg2: f32)
// CHECK: %[[RES:.+]] = arith.addf %arg1, %arg2 : f32 // CHECK: %[[RES:.+]] = arith.addf %arg1, %arg2 : f32
// CHECK: linalg.yield %[[RES]] : f32 // CHECK: linalg.yield %[[RES]] : f32
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<?x4xf32> into tensor<?xf32> // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<?x4xf32> into tensor<?xf32>
// CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1, 2]] : tensor<?xf32> into tensor<?x1x4xf32> // CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1, 2]] : tensor<?xf32> into tensor<?x1x4xf32>
%0 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<?x5x4xf32>) -> tensor<?x1x4xf32> %0 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<?x5x4xf32>) -> tensor<?x1x4xf32>
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: @reduce_float_dyn_nonzero_batch // CHECK-LABEL: @reduce_float_dyn_nonzero_batch
func.func @reduce_float_dyn_nonzero_batch(%arg0: tensor<5x?x4xf32>) -> () { func.func @reduce_float_dyn_nonzero_batch(%arg0: tensor<5x?x4xf32>) -> () {
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [5, %[[DYN]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<5x?xf32>
// CHECK: %[[CST1:.+]] = arith.constant 1.0 // CHECK: %[[CST1:.+]] = arith.constant 1.0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST1]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST1]]{{.*}}outs(%[[INIT]]
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<5x?x4xf32>) outs(%[[FILL]] : tensor<5x?xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<5x?x4xf32>) outs(%[[FILL]] : tensor<5x?xf32>)
// CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: ^bb0(%arg1: f32, %arg2: f32)
// CHECK: %[[RES:.+]] = arith.mulf %arg1, %arg2 : f32 // CHECK: %[[RES:.+]] = arith.mulf %arg1, %arg2 : f32
// CHECK: linalg.yield %[[RES]] : f32 // CHECK: linalg.yield %[[RES]] : f32
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<5x?xf32> into tensor<?xf32> // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<5x?xf32> into tensor<?xf32>
// CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1, 2]] : tensor<?xf32> into tensor<5x?x1xf32> // CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1, 2]] : tensor<?xf32> into tensor<5x?x1xf32>
%0 = "tosa.reduce_prod"(%arg0) {axis = 2 : i64} : (tensor<5x?x4xf32>) -> tensor<5x?x1xf32> %0 = "tosa.reduce_prod"(%arg0) {axis = 2 : i64} : (tensor<5x?x4xf32>) -> tensor<5x?x1xf32>
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: @reduce_float_dyn_multiple // CHECK-LABEL: @reduce_float_dyn_multiple
func.func @reduce_float_dyn_multiple(%arg0: tensor<?x?xf32>) -> () { func.func @reduce_float_dyn_multiple(%arg0: tensor<?x?xf32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DYN]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[CMIN:.+]] = arith.constant -3.40282347E+38 // CHECK: %[[CMIN:.+]] = arith.constant -3.40282347E+38
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CMIN]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CMIN]]{{.*}}outs(%[[INIT]]
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<?x?xf32>) outs(%[[FILL]] : tensor<?xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<?x?xf32>) outs(%[[FILL]] : tensor<?xf32>)
// CHECK: ^bb0(%arg1: f32, %arg2: f32) // CHECK: ^bb0(%arg1: f32, %arg2: f32)
// CHECK: %[[MAX:.+]] = arith.maxf %arg1, %arg2 : f32 // CHECK: %[[MAX:.+]] = arith.maxf %arg1, %arg2 : f32
// CHECK: linalg.yield %[[MAX]] : f32 // CHECK: linalg.yield %[[MAX]] : f32
// CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<?xf32> into tensor<?x1xf32> // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}[0, 1]] : tensor<?xf32> into tensor<?x1xf32>
%0 = "tosa.reduce_max"(%arg0) {axis = 1 : i64} : (tensor<?x?xf32>) -> tensor<?x1xf32> %0 = "tosa.reduce_max"(%arg0) {axis = 1 : i64} : (tensor<?x?xf32>) -> tensor<?x1xf32>
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: @reduce_int // CHECK-LABEL: @reduce_int
// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi32> // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi32>
func.func @reduce_int(%arg0: tensor<5x4xi32>) -> () { func.func @reduce_int(%arg0: tensor<5x4xi32>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [4] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<4xi32>)
// CHECK: ^bb0(%arg1: i32, %arg2: i32) // CHECK: ^bb0(%arg1: i32, %arg2: i32)
// CHECK: [[RES:%.+]] = arith.addi %arg1, %arg2 : i32 // CHECK: [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
// CHECK: linalg.yield [[RES]] : i32 // CHECK: linalg.yield [[RES]] : i32
// CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi32> into tensor<1x4xi32> // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi32> into tensor<1x4xi32>
%0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32> %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<1x4xi32>
// CHECK: [[INIT:%.+]] = linalg.init_tensor [5] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant 0 // CHECK: [[CST0:%.+]] = arith.constant 0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP2]]], iterator_types = ["parallel", "reduction"]} ins([[ARG0]] : tensor<5x4xi32>) outs([[FILL]] : tensor<5xi32>)
// CHECK: ^bb0(%arg1: i32, %arg2: i32) // CHECK: ^bb0(%arg1: i32, %arg2: i32)
// CHECK: [[RES:%.+]] = arith.addi %arg1, %arg2 : i32 // CHECK: [[RES:%.+]] = arith.addi %arg1, %arg2 : i32
// CHECK: linalg.yield [[RES]] : i32 // CHECK: linalg.yield [[RES]] : i32
// CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xi32> into tensor<5x1xi32> // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<5xi32> into tensor<5x1xi32>
%1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x1xi32> %1 = "tosa.reduce_sum"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x1xi32>
Show All 23 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-LABEL: @reduce_bool // CHECK-LABEL: @reduce_bool
// CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi1> // CHECK-SAME: [[ARG0:%.+]]: tensor<5x4xi1>
func.func @reduce_bool(%arg0: tensor<5x4xi1>) -> () { func.func @reduce_bool(%arg0: tensor<5x4xi1>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [4] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[CST0:%.+]] = arith.constant true // CHECK: [[CST0:%.+]] = arith.constant true
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST0]]{{.*}}outs([[INIT]]
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi1>) outs([[FILL]] : tensor<4xi1>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins([[ARG0]] : tensor<5x4xi1>) outs([[FILL]] : tensor<4xi1>)
// CHECK: ^bb0(%arg1: i1, %arg2: i1) // CHECK: ^bb0(%arg1: i1, %arg2: i1)
// CHECK: [[RES:%.+]] = arith.andi %arg1, %arg2 : i1 // CHECK: [[RES:%.+]] = arith.andi %arg1, %arg2 : i1
// CHECK: linalg.yield [[RES]] : i1 // CHECK: linalg.yield [[RES]] : i1
// CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi1> into tensor<1x4xi1> // CHECK: tensor.expand_shape [[GENERIC]] {{\[}}[0, 1]] : tensor<4xi1> into tensor<1x4xi1>
%0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1> %0 = "tosa.reduce_all"(%arg0) {axis = 0 : i64} : (tensor<5x4xi1>) -> tensor<1x4xi1>
Show All 11 Lines
// CHECK-LABEL: @concat // CHECK-LABEL: @concat
func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () { func.func @concat(%arg0: tensor<5x1xf32>, %arg1: tensor<6x1xf32>) -> () {
// CHECK: [[AXIS:%.+]] = arith.constant 0 // CHECK: [[AXIS:%.+]] = arith.constant 0
// CHECK: [[STRIDE:%.+]] = arith.constant 1 // CHECK: [[STRIDE:%.+]] = arith.constant 1
// CHECK: [[OFFSET:%.+]] = arith.constant 0 : index // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
// CHECK: [[IDX0:%.+]] = arith.constant 0 : index // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
// CHECK: [[IDX1:%.+]] = arith.constant 1 : index // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
// CHECK: [[INIT:%.+]] = linalg.init_tensor [11, 1] // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<11x1xf32>
// CHECK: [[CST:%.+]] = arith.constant 0.0 // CHECK: [[CST:%.+]] = arith.constant 0.0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST]]{{.*}}outs([[INIT]]
// CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1] // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1]
// CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg1 into [[INSERT0]][5, 0] [6, 1] [1, 1] // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg1 into [[INSERT0]][5, 0] [6, 1] [1, 1]
%0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x1xf32>, tensor<6x1xf32>) -> (tensor<11x1xf32>) %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x1xf32>, tensor<6x1xf32>) -> (tensor<11x1xf32>)
// CHECK: [[AXIS:%.+]] = arith.constant 1 // CHECK: [[AXIS:%.+]] = arith.constant 1
// CHECK: [[STRIDE:%.+]] = arith.constant 1 // CHECK: [[STRIDE:%.+]] = arith.constant 1
// CHECK: [[OFFSET:%.+]] = arith.constant 0 : index // CHECK: [[OFFSET:%.+]] = arith.constant 0 : index
// CHECK: [[IDX0:%.+]] = arith.constant 0 : index // CHECK: [[IDX0:%.+]] = arith.constant 0 : index
// CHECK: [[IDX1:%.+]] = arith.constant 1 : index // CHECK: [[IDX1:%.+]] = arith.constant 1 : index
// CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 2] // CHECK: [[INIT:%.+]] = tensor.empty() : tensor<5x2xf32>
// CHECK: [[CST:%.+]] = arith.constant 0.0 // CHECK: [[CST:%.+]] = arith.constant 0.0
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CST]]{{.*}}outs([[INIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CST]]{{.*}}outs([[INIT]]
// CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1] // CHECK: [[INSERT0:%.+]] = tensor.insert_slice %arg0 into [[FILL]][0, 0] [5, 1] [1, 1]
// CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg0 into [[INSERT0]][0, 1] [5, 1] [1, 1] // CHECK: [[INSERT1:%.+]] = tensor.insert_slice %arg0 into [[INSERT0]][0, 1] [5, 1] [1, 1]
%1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i64} : (tensor<5x1xf32>, tensor<5x1xf32>) -> (tensor<5x2xf32>) %1 = "tosa.concat"(%arg0, %arg0) { axis = 1 : i64} : (tensor<5x1xf32>, tensor<5x1xf32>) -> (tensor<5x2xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @concat_non_axis_dyn // CHECK-LABEL: @concat_non_axis_dyn
func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () { func.func @concat_non_axis_dyn(%arg0: tensor<5x?xf32>, %arg1: tensor<6x?xf32>) -> () {
// CHECK: %[[AXIS:.+]] = arith.constant 0 // CHECK: %[[AXIS:.+]] = arith.constant 0
// CHECK: %[[STRIDE:.+]] = arith.constant 1 // CHECK: %[[STRIDE:.+]] = arith.constant 1
// CHECK: %[[OFFSET:.+]] = arith.constant 0 : index // CHECK: %[[OFFSET:.+]] = arith.constant 0 : index
// CHECK: %[[IDX0:.+]] = arith.constant 0 : index // CHECK: %[[IDX0:.+]] = arith.constant 0 : index
// CHECK: %[[IDX1:.+]] = arith.constant 1 : index // CHECK: %[[IDX1:.+]] = arith.constant 1 : index
// CHECK: %[[SIZE:.+]] = tensor.dim %arg0, %[[IDX1]] // CHECK: %[[SIZE:.+]] = tensor.dim %arg0, %[[IDX1]]
// CHECK: %[[IDX1_2:.+]] = arith.constant 1 : index // CHECK: %[[IDX1_2:.+]] = arith.constant 1 : index
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[IDX1_2]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[IDX1_2]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [11, %[[DYN]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<11x?xf32>
// CHECK: %[[CST:.+]] = arith.constant 0.0 // CHECK: %[[CST:.+]] = arith.constant 0.0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]]{{.*}}outs(%[[INIT]]
// CHECK: %[[INSERT0:.+]] = tensor.insert_slice %arg0 into %[[FILL]][0, 0] [5, %[[SIZE]]] [1, 1] // CHECK: %[[INSERT0:.+]] = tensor.insert_slice %arg0 into %[[FILL]][0, 0] [5, %[[SIZE]]] [1, 1]
// CHECK: %[[INSERT1:.+]] = tensor.insert_slice %arg1 into %[[INSERT0]][5, 0] [6, %[[SIZE]]] [1, 1] // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %arg1 into %[[INSERT0]][5, 0] [6, %[[SIZE]]] [1, 1]
%0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x?xf32>, tensor<6x?xf32>) -> (tensor<11x?xf32>) %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<5x?xf32>, tensor<6x?xf32>) -> (tensor<11x?xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @concat_axis_dyn // CHECK-LABEL: @concat_axis_dyn
func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () { func.func @concat_axis_dyn(%arg0: tensor<?x3xf32>, %arg1: tensor<?x3xf32>) -> () {
// CHECK: %[[AXIS:.+]] = arith.constant 0 // CHECK: %[[AXIS:.+]] = arith.constant 0
// CHECK: %[[STRIDE:.+]] = arith.constant 1 // CHECK: %[[STRIDE:.+]] = arith.constant 1
// CHECK: %[[OFFSET:.+]] = arith.constant 0 : index // CHECK: %[[OFFSET:.+]] = arith.constant 0 : index
// CHECK: %[[IDX0:.+]] = arith.constant 0 : index // CHECK: %[[IDX0:.+]] = arith.constant 0 : index
// CHECK: %[[SIZE:.+]] = tensor.dim %arg0, %[[IDX0]] // CHECK: %[[SIZE:.+]] = tensor.dim %arg0, %[[IDX0]]
// CHECK: %[[IDX0_2:.+]] = arith.constant 0 : index // CHECK: %[[IDX0_2:.+]] = arith.constant 0 : index
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[IDX0_2]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[IDX0_2]]
// CHECK: %[[IDX1:.+]] = arith.constant 1 : index // CHECK: %[[IDX1:.+]] = arith.constant 1 : index
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DYN]], 3] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]]) : tensor<?x3xf32>
// CHECK: %[[CST:.+]] = arith.constant 0.0 // CHECK: %[[CST:.+]] = arith.constant 0.0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST]]{{.*}}outs(%[[INIT]]
// CHECK: %[[DYN1:.+]] = tensor.dim %arg0, %[[AXIS]] // CHECK: %[[DYN1:.+]] = tensor.dim %arg0, %[[AXIS]]
// CHECK: %[[INSERT0:.+]] = tensor.insert_slice %arg0 into %[[FILL]][0, 0] [%[[DYN1]], 3] [1, 1] // CHECK: %[[INSERT0:.+]] = tensor.insert_slice %arg0 into %[[FILL]][0, 0] [%[[DYN1]], 3] [1, 1]
// CHECK: %[[SUM:.+]] = arith.addi %[[OFFSET]], %[[DYN1]] // CHECK: %[[SUM:.+]] = arith.addi %[[OFFSET]], %[[DYN1]]
// CHECK: %[[DYN2:.+]] = tensor.dim %arg1, %[[AXIS]] // CHECK: %[[DYN2:.+]] = tensor.dim %arg1, %[[AXIS]]
// CHECK: %[[INSERT1:.+]] = tensor.insert_slice %arg1 into %[[INSERT0]][%[[SUM]], 0] [%[[DYN2]], 3] [1, 1] // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %arg1 into %[[INSERT0]][%[[SUM]], 0] [%[[DYN2]], 3] [1, 1]
%0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<?x3xf32>, tensor<?x3xf32>) -> (tensor<?x3xf32>) %0 = "tosa.concat"(%arg0, %arg1) { axis = 0 : i64} : (tensor<?x3xf32>, tensor<?x3xf32>) -> (tensor<?x3xf32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @rescale_i8 // CHECK-LABEL: @rescale_i8
func.func @rescale_i8(%arg0 : tensor<2xi8>) -> () { func.func @rescale_i8(%arg0 : tensor<2xi8>) -> () {
// CHECK: [[C0:%.+]] = arith.constant 19689 // CHECK: [[C0:%.+]] = arith.constant 19689
// CHECK: [[C1:%.+]] = arith.constant 15 // CHECK: [[C1:%.+]] = arith.constant 15
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xi8>)
// CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: i8): // CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: i8):
// CHECK: [[C17:%.+]] = arith.constant 17 // CHECK: [[C17:%.+]] = arith.constant 17
// CHECK: [[C22:%.+]] = arith.constant 22 // CHECK: [[C22:%.+]] = arith.constant 22
// CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]] // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
// CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]] // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
// CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false} // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
// CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]] // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]]
// CHECK-DAG: [[CMIN:%.+]] = arith.constant -128 // CHECK-DAG: [[CMIN:%.+]] = arith.constant -128
// CHECK-DAG: [[CMAX:%.+]] = arith.constant 127 // CHECK-DAG: [[CMAX:%.+]] = arith.constant 127
// CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]] // CHECK-DAG: [[MINLT:%.+]] = arith.cmpi slt, [[SCALED_ZEROED]], [[CMIN]]
// CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]] // CHECK-DAG: [[MAXLT:%.+]] = arith.cmpi slt, [[CMAX]], [[SCALED_ZEROED]]
// CHECK-DAG: [[LOWER:%.+]] = arith.select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]] // CHECK-DAG: [[LOWER:%.+]] = arith.select [[MINLT]], [[CMIN]], [[SCALED_ZEROED]]
// CHECK-DAG: [[BOUNDED:%.+]] = arith.select [[MAXLT]], [[CMAX]], [[LOWER]] // CHECK-DAG: [[BOUNDED:%.+]] = arith.select [[MAXLT]], [[CMAX]], [[LOWER]]
// CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]] // CHECK-DAG: [[TRUNC:%.+]] = arith.trunci [[BOUNDED]]
// CHECK-DAG: linalg.yield [[TRUNC]] // CHECK-DAG: linalg.yield [[TRUNC]]
%0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xi8>) -> (tensor<2xi8>) %0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<2xi8>) -> (tensor<2xi8>)
// CHECK: [[C0:%.+]] = arith.constant 19689 // CHECK: [[C0:%.+]] = arith.constant 19689
// CHECK: [[C1:%.+]] = arith.constant 15 // CHECK: [[C1:%.+]] = arith.constant 15
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xui8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xi8>) outs([[INIT]] : tensor<2xui8>)
// CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: ui8): // CHECK: ^bb0([[IN:%.+]]: i8, [[UNUSED:%.+]]: ui8):
// CHECK: [[C17:%.+]] = arith.constant 17 // CHECK: [[C17:%.+]] = arith.constant 17
// CHECK: [[C22:%.+]] = arith.constant 22 // CHECK: [[C22:%.+]] = arith.constant 22
// CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]] // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
// CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]] // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
// CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false} // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
// CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]] // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C22]]
Show All 15 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @rescale_i8_dyn_batch // CHECK-LABEL: @rescale_i8_dyn_batch
func.func @rescale_i8_dyn_batch(%arg0 : tensor<?x2xi8>) -> () { func.func @rescale_i8_dyn_batch(%arg0 : tensor<?x2xi8>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 2] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x2xi8>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<?x2xi8>) outs(%[[INIT]] : tensor<?x2xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<?x2xi8>) outs(%[[INIT]] : tensor<?x2xi8>)
%0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<?x2xi8>) -> (tensor<?x2xi8>) %0 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<?x2xi8>) -> (tensor<?x2xi8>)
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 2] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x2xui8>
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<?x2xi8>) outs(%[[INIT]] : tensor<?x2xui8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<?x2xi8>) outs(%[[INIT]] : tensor<?x2xui8>)
%1 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<?x2xi8>) -> (tensor<?x2xui8>) %1 = "tosa.rescale"(%arg0) {input_zp = 17 : i32, output_zp = 22 : i32, multiplier = [19689 : i32], shift = [15 : i32], scale32 = false, double_round = false, per_channel = false} : (tensor<?x2xi8>) -> (tensor<?x2xui8>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @rescale_dyn // CHECK-LABEL: @rescale_dyn
func.func @rescale_dyn(%arg0 : tensor<1x?x?x32xi32>) -> () { func.func @rescale_dyn(%arg0 : tensor<1x?x?x32xi32>) -> () {
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[DIM1:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[C2:.+]] = arith.constant 2 // CHECK: %[[C2:.+]] = arith.constant 2
// CHECK: %[[DIM2:.+]] = tensor.dim %arg0, %[[C2]] // CHECK: %[[DIM2:.+]] = tensor.dim %arg0, %[[C2]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, %[[DIM1]], %[[DIM2]], 32] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM1]], %[[DIM2]])
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<1x?x?x32xi32>) outs(%[[INIT]] : tensor<1x?x?x32xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<1x?x?x32xi32>) outs(%[[INIT]] : tensor<1x?x?x32xi8>)
%0 = "tosa.rescale"(%arg0) {double_round = true, input_zp = 0 : i32, multiplier = [1376784203 : i32], output_zp = 0 : i32, per_channel = false, scale32 = true, shift = [38 : i32]} : (tensor<1x?x?x32xi32>) -> tensor<1x?x?x32xi8> %0 = "tosa.rescale"(%arg0) {double_round = true, input_zp = 0 : i32, multiplier = [1376784203 : i32], output_zp = 0 : i32, per_channel = false, scale32 = true, shift = [38 : i32]} : (tensor<1x?x?x32xi32>) -> tensor<1x?x?x32xi8>
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @rescale_ui8 // CHECK-LABEL: @rescale_ui8
func.func @rescale_ui8(%arg0 : tensor<2xui8>) -> () { func.func @rescale_ui8(%arg0 : tensor<2xui8>) -> () {
// CHECK: [[C0:%.+]] = arith.constant 19689 // CHECK: [[C0:%.+]] = arith.constant 19689
// CHECK: [[C1:%.+]] = arith.constant 15 // CHECK: [[C1:%.+]] = arith.constant 15
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xui8>) outs([[INIT]] : tensor<2xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0 : tensor<2xui8>) outs([[INIT]] : tensor<2xi8>)
// CHECK: ^bb0([[IN:%.+]]: ui8, [[UNUSED:%.+]]: i8): // CHECK: ^bb0([[IN:%.+]]: ui8, [[UNUSED:%.+]]: i8):
// CHECK: [[C17:%.+]] = arith.constant 17 // CHECK: [[C17:%.+]] = arith.constant 17
// CHECK: [[C22:%.+]] = arith.constant 22 // CHECK: [[C22:%.+]] = arith.constant 22
// CHECK-DAG: [[CAST:%.+]] = builtin.unrealized_conversion_cast [[IN]] : ui8 to i8 // CHECK-DAG: [[CAST:%.+]] = builtin.unrealized_conversion_cast [[IN]] : ui8 to i8
// CHECK-DAG: [[IN32:%.+]] = arith.extui [[CAST]] // CHECK-DAG: [[IN32:%.+]] = arith.extui [[CAST]]
// CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]] // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C17]]
// CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false} // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[C0]], [[C1]]) {double_round = false}
Show All 14 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @rescale_per_channel // CHECK-LABEL: @rescale_per_channel
func.func @rescale_per_channel(%arg0 : tensor<3xi8>) -> (tensor<3xi8>) { func.func @rescale_per_channel(%arg0 : tensor<3xi8>) -> (tensor<3xi8>) {
// CHECK: [[MULTIPLIERS:%.+]] = arith.constant dense<[42, 43, 0]> // CHECK: [[MULTIPLIERS:%.+]] = arith.constant dense<[42, 43, 0]>
// CHECK: [[SHIFTS:%.+]] = arith.constant dense<[14, 15, 0]> // CHECK: [[SHIFTS:%.+]] = arith.constant dense<[14, 15, 0]>
// CHECK: [[INIT:%.+]] = linalg.init_tensor [3] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[MULTIPLIERS]], [[SHIFTS]] : tensor<3xi8>, tensor<3xi32>, tensor<3xi8>) outs([[INIT]] : tensor<3xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%arg0, [[MULTIPLIERS]], [[SHIFTS]] : tensor<3xi8>, tensor<3xi32>, tensor<3xi8>) outs([[INIT]] : tensor<3xi8>)
// CHECK: ^bb0([[IN:%.+]]: i8, [[MULTIPLIER:%.+]]: i32, [[SHIFT:%.+]]: i8, [[UNUSED:%.+]]: i8): // CHECK: ^bb0([[IN:%.+]]: i8, [[MULTIPLIER:%.+]]: i32, [[SHIFT:%.+]]: i8, [[UNUSED:%.+]]: i8):
// CHECK: [[C243:%.+]] = arith.constant 243 // CHECK: [[C243:%.+]] = arith.constant 243
// CHECK: [[C252:%.+]] = arith.constant 252 // CHECK: [[C252:%.+]] = arith.constant 252
// CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]] // CHECK-DAG: [[IN32:%.+]] = arith.extsi [[IN]]
// CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C243]] // CHECK-DAG: [[IN_ZEROED:%.+]] = arith.subi [[IN32]], [[C243]]
// CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[MULTIPLIER]], [[SHIFT]]) {double_round = false} // CHECK-DAG: [[SCALED:%.+]] = "tosa.apply_scale"([[IN_ZEROED]], [[MULTIPLIER]], [[SHIFT]]) {double_round = false}
Show All 35 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @reverse // CHECK-LABEL: @reverse
func.func @reverse(%arg0: tensor<5x4xi32>) -> () { func.func @reverse(%arg0: tensor<5x4xi32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [5, 4] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>)
// CHECK-DAG: %[[I0:.+]] = linalg.index 0 // CHECK-DAG: %[[I0:.+]] = linalg.index 0
// CHECK-DAG: %[[I1:.+]] = linalg.index 1 // CHECK-DAG: %[[I1:.+]] = linalg.index 1
// CHECK-DAG: %[[SUB1:.+]] = arith.constant 1 // CHECK-DAG: %[[SUB1:.+]] = arith.constant 1
// CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]] // CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]]
// CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]] // CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]]
// CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]], %[[I1]]] : tensor<5x4xi32> // CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]], %[[I1]]] : tensor<5x4xi32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32> %0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C1]] // CHECK: %[[RDIM:.+]] = tensor.dim %arg0, %[[C1]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [5, 4] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel", "parallel"]} outs(%[[INIT]] : tensor<5x4xi32>)
// CHECK-DAG: %[[I0:.+]] = linalg.index 0 // CHECK-DAG: %[[I0:.+]] = linalg.index 0
// CHECK-DAG: %[[I1:.+]] = linalg.index 1 // CHECK-DAG: %[[I1:.+]] = linalg.index 1
// CHECK-DAG: %[[SUB1:.+]] = arith.constant 1 // CHECK-DAG: %[[SUB1:.+]] = arith.constant 1
// CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]] // CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[RDIM]], %[[SUB1]]
// CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I1]] // CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I1]]
// CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[I0]], %[[READ_DIM]]] : tensor<5x4xi32> // CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[I0]], %[[READ_DIM]]] : tensor<5x4xi32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%1 = "tosa.reverse"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32> %1 = "tosa.reverse"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @reverse_dyn // CHECK-LABEL: @reverse_dyn
func.func @reverse_dyn(%arg0: tensor<?xi32>) -> () { func.func @reverse_dyn(%arg0: tensor<?xi32>) -> () {
// CHECK: %[[C0_1:.+]] = arith.constant 0 // CHECK: %[[C0_1:.+]] = arith.constant 0
// CHECK: %[[D0_1:.+]] = tensor.dim %arg0, %[[C0_1]] // CHECK: %[[D0_1:.+]] = tensor.dim %arg0, %[[C0_1]]
// CHECK: %[[C0_2:.+]] = arith.constant 0 // CHECK: %[[C0_2:.+]] = arith.constant 0
// CHECK: %[[D0_2:.+]] = tensor.dim %arg0, %[[C0_2]] // CHECK: %[[D0_2:.+]] = tensor.dim %arg0, %[[C0_2]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[D0_1]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[D0_1]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel"]} outs(%[[INIT]] : tensor<?xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]]], iterator_types = ["parallel"]} outs(%[[INIT]] : tensor<?xi32>)
// CHECK-DAG: %[[I0:.+]] = linalg.index 0 // CHECK-DAG: %[[I0:.+]] = linalg.index 0
// CHECK-DAG: %[[SUB1:.+]] = arith.constant 1 // CHECK-DAG: %[[SUB1:.+]] = arith.constant 1
// CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[D0_2]], %[[SUB1]] // CHECK-DAG: %[[RDIM_MINUS_C1:.+]] = arith.subi %[[D0_2]], %[[SUB1]]
// CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]] // CHECK-DAG: %[[READ_DIM:.+]] = arith.subi %[[RDIM_MINUS_C1]], %[[I0]]
// CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]]] : tensor<?xi32> // CHECK-DAG: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[READ_DIM]]] : tensor<?xi32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<?xi32>) -> tensor<?xi32> %0 = "tosa.reverse"(%arg0) {axis = 0 : i64} : (tensor<?xi32>) -> tensor<?xi32>
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @tile // CHECK-LABEL: @tile
func.func @tile(%arg0 : tensor<2x3xi8>) -> () { func.func @tile(%arg0 : tensor<2x3xi8>) -> () {
// CHECK: [[INIT:%.+]] = linalg.init_tensor [2, 2, 1, 3] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<2x2x1x3xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<2x2x1x3xi8>)
// CHECK: linalg.yield %arg1 : i8 // CHECK: linalg.yield %arg1 : i8
// CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1, 2], [3]] // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1, 2], [3]]
%0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<2x3xi8>) -> (tensor<4x3xi8>) %0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<2x3xi8>) -> (tensor<4x3xi8>)
// CHECK: [[INIT:%.+]] = linalg.init_tensor [1, 2, 2, 3] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<1x2x2x3xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<1x2x2x3xi8>)
// CHECK: linalg.yield %arg1 : i8 // CHECK: linalg.yield %arg1 : i8
// CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]] // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
%1 = "tosa.tile"(%arg0) {multiples = [1, 2]} : (tensor<2x3xi8>) -> (tensor<2x6xi8>) %1 = "tosa.tile"(%arg0) {multiples = [1, 2]} : (tensor<2x3xi8>) -> (tensor<2x6xi8>)
// CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 2, 7, 3] // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<5x2x7x3xi8>) // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs([[INIT]] : tensor<5x2x7x3xi8>)
// CHECK: linalg.yield %arg1 : i8 // CHECK: linalg.yield %arg1 : i8
// CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]] // CHECK: tensor.collapse_shape [[GENERIC]] {{\[}}[0, 1], [2, 3]]
%2 = "tosa.tile"(%arg0) {multiples = [5, 7]} : (tensor<2x3xi8>) -> (tensor<10x21xi8>) %2 = "tosa.tile"(%arg0) {multiples = [5, 7]} : (tensor<2x3xi8>) -> (tensor<10x21xi8>)
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @tile_dyn_input // CHECK-LABEL: @tile_dyn_input
func.func @tile_dyn_input(%arg0 : tensor<?x3xi8>) -> () { func.func @tile_dyn_input(%arg0 : tensor<?x3xi8>) -> () {
// CHECK: %[[CST0:.+]] = arith.constant 0 // CHECK: %[[CST0:.+]] = arith.constant 0
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST0]] : tensor<?x3xi8> // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST0]] : tensor<?x3xi8>
// CHECK: %[[INIT:.+]] = linalg.init_tensor [2, %[[DYN]], 1, 3] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<?x3xi8>) outs(%[[INIT]] : tensor<2x?x1x3xi8>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<?x3xi8>) outs(%[[INIT]] : tensor<2x?x1x3xi8>)
// CHECK: linalg.yield %arg1 : i8 // CHECK: linalg.yield %arg1 : i8
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1, 2, 3]] // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1, 2, 3]]
// CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1]] // CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1]]
%0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<?x3xi8>) -> (tensor<?x3xi8>) %0 = "tosa.tile"(%arg0) {multiples = [2, 1]} : (tensor<?x3xi8>) -> (tensor<?x3xi8>)
return return
} }
// ----- // -----
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK-LABEL: @tile_dyn_multiples // CHECK-LABEL: @tile_dyn_multiples
func.func @tile_dyn_multiples(%arg0 : tensor<2x3xi8>) -> () { func.func @tile_dyn_multiples(%arg0 : tensor<2x3xi8>) -> () {
// CHECK: %[[CST1:.+]] = arith.constant 1 // CHECK: %[[CST1:.+]] = arith.constant 1
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST1]] : tensor<2x3xi8> // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST1]] : tensor<2x3xi8>
// CHECK: %[[INIT:.+]] = linalg.init_tensor [2, 2, %[[DYN]], 3] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs(%[[INIT]] : tensor<2x2x?x3xi8>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<2x3xi8>) outs(%[[INIT]] : tensor<2x2x?x3xi8>)
// CHECK: linalg.yield %arg1 : i8 // CHECK: linalg.yield %arg1 : i8
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1, 2, 3]] // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1, 2, 3]]
// CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1]] // CHECK: tensor.expand_shape %[[COLLAPSED]] {{\[}}[0, 1]]
%0 = "tosa.tile"(%arg0) {multiples = [2, -1]} : (tensor<2x3xi8>) -> (tensor<2x?xi8>) %0 = "tosa.tile"(%arg0) {multiples = [2, -1]} : (tensor<2x3xi8>) -> (tensor<2x?xi8>)
return return
} }
▲ Show 20 Lines • Show All 91 Lines • ▼ Show 20 Lines
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP4:.*]] = affine_map<(d0) -> ()> // CHECK: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () { func.func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
// CHECK: [[IDX_INIT:%.+]] = linalg.init_tensor [2] // CHECK: [[IDX_INIT:%.+]] = tensor.empty()
// CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32 // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
// CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]] // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]]
// CHECK: [[VAL_INIT:%.+]] = linalg.init_tensor [2] // CHECK: [[VAL_INIT:%.+]] = tensor.empty()
// CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648 // CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648
// CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]]{{.*}}outs([[VAL_INIT]] // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]]{{.*}}outs([[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<2xi32>, tensor<2xi32>) // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<2xi32>, tensor<2xi32>)
// CHECK: [[IDX:%.+]] = linalg.index 0 // CHECK: [[IDX:%.+]] = linalg.index 0
// CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]]
// CHECK: [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4 // CHECK: [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4
// CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %arg2, %arg4 // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %arg2, %arg4
// CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %arg3 // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %arg3
// CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]] // CHECK: linalg.yield [[SELECT_IDX]], [[SELECT_VAL]]
%0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x2xi32>) -> (tensor<2xi32>) %0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x2xi32>) -> (tensor<2xi32>)
// CHECK: [[IDX_INIT:%.+]] = linalg.init_tensor [3] // CHECK: [[IDX_INIT:%.+]] = tensor.empty()
// CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32 // CHECK: [[IDX_MIN:%.+]] = arith.constant 0 : i32
// CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]] // CHECK: [[IDX_FILL:%.+]] = linalg.fill ins([[IDX_MIN]]{{.*}}outs([[IDX_INIT]]
// CHECK: [[VAL_INIT:%.+]] = linalg.init_tensor [3] // CHECK: [[VAL_INIT:%.+]] = tensor.empty()
// CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648 // CHECK: [[VAL_MIN:%.+]] = arith.constant -2147483648
// CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]]{{.*}}outs([[VAL_INIT]] // CHECK: [[VAL_FILL:%.+]] = linalg.fill ins([[VAL_MIN]]{{.*}}outs([[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#map0, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>) // CHECK: linalg.generic {indexing_maps = [#map0, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
// CHECK: [[IDX:%.+]] = linalg.index 1 // CHECK: [[IDX:%.+]] = linalg.index 1
// CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]] // CHECK: [[CAST:%.+]] = arith.index_cast [[IDX]]
// CHECK: [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4 // CHECK: [[CMP:%.+]] = arith.cmpi sgt, %arg2, %arg4
// CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %arg2, %arg4 // CHECK: [[SELECT_VAL:%.+]] = arith.select [[CMP]], %arg2, %arg4
// CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %arg3 // CHECK: [[SELECT_IDX:%.+]] = arith.select [[CMP]], [[CAST]], %arg3
Show All 15 Lines
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>
func.func @argmax_dyn_non_axis(%arg0 : tensor<3x?xi32>) -> () { func.func @argmax_dyn_non_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: %[[CST1:.+]] = arith.constant 1 // CHECK: %[[CST1:.+]] = arith.constant 1
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST1]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST1]]
// CHECK: %[[IDX_INIT:.+]] = linalg.init_tensor [%[[DYN]]] // CHECK: %[[IDX_INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32 // CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32
// CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]] // CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]]
// CHECK: %[[VAL_INIT:.+]] = linalg.init_tensor [%[[DYN]]] // CHECK: %[[VAL_INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[VAL_MIN:.+]] = arith.constant -2147483648 // CHECK: %[[VAL_MIN:.+]] = arith.constant -2147483648
// CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MIN]]{{.*}}outs(%[[VAL_INIT]] // CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MIN]]{{.*}}outs(%[[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<?xi32>, tensor<?xi32>) // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["reduction", "parallel"]} ins(%arg0 : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<?xi32>, tensor<?xi32>)
// CHECK: %[[IDX:.+]] = linalg.index 0 // CHECK: %[[IDX:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]]
// CHECK: %[[CMP:.+]] = arith.cmpi sgt, %arg1, %arg3 // CHECK: %[[CMP:.+]] = arith.cmpi sgt, %arg1, %arg3
// CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %arg1, %arg3 // CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %arg1, %arg3
// CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %arg2 // CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %arg2
// CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]] // CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]]
%0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x?xi32>) -> (tensor<?xi32>) %0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x?xi32>) -> (tensor<?xi32>)
return return
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () { func.func @argmax_dyn_axis(%arg0 : tensor<3x?xi32>) -> () {
// CHECK: %[[IDX_INIT:.+]] = linalg.init_tensor [3] // CHECK: %[[IDX_INIT:.+]] = tensor.empty()
// CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32 // CHECK: %[[IDX_MIN:.+]] = arith.constant 0 : i32
// CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]] // CHECK: %[[IDX_FILL:.+]] = linalg.fill ins(%[[IDX_MIN]]{{.*}}outs(%[[IDX_INIT]]
// CHECK: %[[VAL_INIT:.+]] = linalg.init_tensor [3] // CHECK: %[[VAL_INIT:.+]] = tensor.empty()
// CHECK: %[[VAL_MIN:.+]] = arith.constant -2147483648 // CHECK: %[[VAL_MIN:.+]] = arith.constant -2147483648
// CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MIN]]{{.*}}outs(%[[VAL_INIT]] // CHECK: %[[VAL_FILL:.+]] = linalg.fill ins(%[[VAL_MIN]]{{.*}}outs(%[[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>) // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<3x?xi32>) outs(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
// CHECK: %[[IDX:.+]] = linalg.index 1 // CHECK: %[[IDX:.+]] = linalg.index 1
// CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[IDX]]
// CHECK: %[[CMP:.+]] = arith.cmpi sgt, %arg1, %arg3 // CHECK: %[[CMP:.+]] = arith.cmpi sgt, %arg1, %arg3
// CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %arg1, %arg3 // CHECK: %[[SELECT_VAL:.+]] = arith.select %[[CMP]], %arg1, %arg3
// CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %arg2 // CHECK: %[[SELECT_IDX:.+]] = arith.select %[[CMP]], %[[CAST]], %arg2
// CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]] // CHECK: linalg.yield %[[SELECT_IDX]], %[[SELECT_VAL]]
%0 = "tosa.argmax"(%arg0) { axis = 1 : i64} : (tensor<3x?xi32>) -> (tensor<3xi32>) %0 = "tosa.argmax"(%arg0) { axis = 1 : i64} : (tensor<3x?xi32>) -> (tensor<3xi32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @gather_float // CHECK-LABEL: @gather_float
func.func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () { func.func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [2, 3, 2] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xf32>)
// CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: f32) // CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]]
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xf32> // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xf32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.gather"(%arg0, %arg1) : (tensor<2x3x2xf32>, tensor<2x3xi32>) -> (tensor<2x3x2xf32>) %0 = "tosa.gather"(%arg0, %arg1) : (tensor<2x3x2xf32>, tensor<2x3xi32>) -> (tensor<2x3x2xf32>)
return return
} }
// CHECK-LABEL: @gather_float_dyn // CHECK-LABEL: @gather_float_dyn
func.func @gather_float_dyn(%arg0: tensor<?x3x2xf32>, %arg1: tensor<?x3xi32>) -> () { func.func @gather_float_dyn(%arg0: tensor<?x3x2xf32>, %arg1: tensor<?x3xi32>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 3, 2] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<?x3xi32>) outs(%[[INIT]] : tensor<?x3x2xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<?x3xi32>) outs(%[[INIT]] : tensor<?x3x2xf32>)
// CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: f32) // CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]]
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x3x2xf32> // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<?x3x2xf32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.gather"(%arg0, %arg1) : (tensor<?x3x2xf32>, tensor<?x3xi32>) -> (tensor<?x3x2xf32>) %0 = "tosa.gather"(%arg0, %arg1) : (tensor<?x3x2xf32>, tensor<?x3xi32>) -> (tensor<?x3x2xf32>)
return return
} }
// CHECK-LABEL: @gather_int // CHECK-LABEL: @gather_int
func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () { func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [2, 3, 2] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map0, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg1 : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>)
// CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: i32) // CHECK: ^bb0(%[[ARG0:.+]]: i32, %[[ARG1:.+]]: i32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG0]]
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xi32> // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg0[%[[IDX0]], %[[CAST]], %[[IDX2]]] : tensor<2x3x2xi32>
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.gather"(%arg0, %arg1) : (tensor<2x3x2xi32>, tensor<2x3xi32>) -> (tensor<2x3x2xi32>) %0 = "tosa.gather"(%arg0, %arg1) : (tensor<2x3x2xi32>, tensor<2x3xi32>) -> (tensor<2x3x2xi32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @table8 // CHECK-LABEL: @table8
func.func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () { func.func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [6] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>)
// CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8) // CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8)
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]]
// CHECK: %[[OFFSET:.+]] = arith.constant 128 // CHECK: %[[OFFSET:.+]] = arith.constant 128
// CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]] // CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]]
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]] // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]]
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<512xi8>) -> (tensor<6xi8>) %0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<512xi8>) -> (tensor<6xi8>)
return return
} }
// ----- // -----
// CHECK-LABEL: @table16 // CHECK-LABEL: @table16
func.func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () { func.func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [6] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi16>) outs(%[[INIT]] : tensor<6xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi16>) outs(%[[INIT]] : tensor<6xi32>)
// CHECK: ^bb0(%arg2: i16, %arg3: i32) // CHECK: ^bb0(%arg2: i16, %arg3: i32)
// CHECK: %[[EXT_IN:.+]] = arith.extsi %arg2 // CHECK: %[[EXT_IN:.+]] = arith.extsi %arg2
// CHECK: %[[C32768:.+]] = arith.constant 32768 // CHECK: %[[C32768:.+]] = arith.constant 32768
// CHECK: %[[C7:.+]] = arith.constant 7 // CHECK: %[[C7:.+]] = arith.constant 7
// CHECK: %[[C1:.+]] = arith.constant 1 // CHECK: %[[C1:.+]] = arith.constant 1
// CHECK: %[[C127:.+]] = arith.constant 127 // CHECK: %[[C127:.+]] = arith.constant 127
// CHECK: %[[INADD:.+]] = arith.addi %[[EXT_IN]], %[[C32768]] // CHECK: %[[INADD:.+]] = arith.addi %[[EXT_IN]], %[[C32768]]
Show All 16 Lines
} }
// ----- // -----
// CHECK-LABEL: @table8_dyn // CHECK-LABEL: @table8_dyn
func.func @table8_dyn(%arg0: tensor<?xi8>, %arg1: tensor<512xi8>) -> () { func.func @table8_dyn(%arg0: tensor<?xi8>, %arg1: tensor<512xi8>) -> () {
// CHECK: %[[CST0:.+]] = arith.constant 0 // CHECK: %[[CST0:.+]] = arith.constant 0
// CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST0]] // CHECK: %[[DYN:.+]] = tensor.dim %arg0, %[[CST0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[DYN]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[DYN]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<?xi8>) outs(%[[INIT]] : tensor<?xi8>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<?xi8>) outs(%[[INIT]] : tensor<?xi8>)
// CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8) // CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8)
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]]
// CHECK: %[[OFFSET:.+]] = arith.constant 128 // CHECK: %[[OFFSET:.+]] = arith.constant 128
// CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]] // CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]]
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]] // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]]
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.table"(%arg0, %arg1) : (tensor<?xi8>, tensor<512xi8>) -> (tensor<?xi8>) %0 = "tosa.table"(%arg0, %arg1) : (tensor<?xi8>, tensor<512xi8>) -> (tensor<?xi8>)
return return
} }
// ----- // -----
// CHECK-LABEL: @table8_dyn_table // CHECK-LABEL: @table8_dyn_table
func.func @table8_dyn_table(%arg0: tensor<6xi8>, %arg1: tensor<?xi8>) -> () { func.func @table8_dyn_table(%arg0: tensor<6xi8>, %arg1: tensor<?xi8>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [6] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>)
// CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8) // CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8)
// CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[ARG_IN]]
// CHECK: %[[OFFSET:.+]] = arith.constant 128 // CHECK: %[[OFFSET:.+]] = arith.constant 128
// CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]] // CHECK: %[[ADD:.+]] = arith.addi %[[CAST]], %[[OFFSET]]
// CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]] // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]]
// CHECK: linalg.yield %[[EXTRACT]] // CHECK: linalg.yield %[[EXTRACT]]
%0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<?xi8>) -> (tensor<6xi8>) %0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<?xi8>) -> (tensor<6xi8>)
return return
} }
// ----- // -----
// CHECK-LABEL: @resize_nearest // CHECK-LABEL: @resize_nearest
func.func @resize_nearest(%input: tensor<1x2x2x1xf32>) -> () { func.func @resize_nearest(%input: tensor<1x2x2x1xf32>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[IDX1:.+]] = linalg.index 1 // CHECK: %[[IDX1:.+]] = linalg.index 1
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[IDX3:.+]] = linalg.index 3 // CHECK: %[[IDX3:.+]] = linalg.index 3
// CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0 // CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0
// CHECK-DAG: %[[YMAX:.+]] = arith.constant 1 // CHECK-DAG: %[[YMAX:.+]] = arith.constant 1
// CHECK-DAG: %[[XMAX:.+]] = arith.constant 1 // CHECK-DAG: %[[XMAX:.+]] = arith.constant 1
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Linesfunc.func @resize_nearest(%input: tensor<1x2x2x1xf32>) -> () {
return return
} }
// ----- // -----
// CHECK-LABEL: @resize_bilinear // CHECK-LABEL: @resize_bilinear
func.func @resize_bilinear(%input: tensor<1x2x2x1xf32>) -> () { func.func @resize_bilinear(%input: tensor<1x2x2x1xf32>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[IDX1:.+]] = linalg.index 1 // CHECK: %[[IDX1:.+]] = linalg.index 1
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[IDX3:.+]] = linalg.index 3 // CHECK: %[[IDX3:.+]] = linalg.index 3
// CHECK: %[[XYMIN:.+]] = arith.constant 0 // CHECK: %[[XYMIN:.+]] = arith.constant 0
// CHECK: %[[YMAX:.+]] = arith.constant 1 // CHECK: %[[YMAX:.+]] = arith.constant 1
// CHECK: %[[XMAX:.+]] = arith.constant 1 // CHECK: %[[XMAX:.+]] = arith.constant 1
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Linesfunc.func @resize_bilinear(%input: tensor<1x2x2x1xf32>) -> () {
%output = "tosa.resize"(%input) { output_size = [4, 4], stride = [0, 0], offset = [0, 0], stride_fp = [0.5 : f32, 0.5 : f32], offset_fp = [0.1 : f32, 0.2 : f32], shift = 0 : i32, mode = "BILINEAR" } : (tensor<1x2x2x1xf32>) -> (tensor<1x4x4x1xf32>) %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [0, 0], offset = [0, 0], stride_fp = [0.5 : f32, 0.5 : f32], offset_fp = [0.1 : f32, 0.2 : f32], shift = 0 : i32, mode = "BILINEAR" } : (tensor<1x2x2x1xf32>) -> (tensor<1x4x4x1xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @resize_nearest_int // CHECK-LABEL: @resize_nearest_int
func.func @resize_nearest_int(%input: tensor<1x2x2x1xi32>) -> () { func.func @resize_nearest_int(%input: tensor<1x2x2x1xi32>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[IDX1:.+]] = linalg.index 1 // CHECK: %[[IDX1:.+]] = linalg.index 1
// CHECK: %[[IDX2:.+]] = linalg.index 2 // CHECK: %[[IDX2:.+]] = linalg.index 2
// CHECK: %[[IDX3:.+]] = linalg.index 3 // CHECK: %[[IDX3:.+]] = linalg.index 3
// CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0 // CHECK-DAG: %[[XYMIN:.+]] = arith.constant 0
// CHECK-DAG: %[[YMAX:.+]] = arith.constant 1 // CHECK-DAG: %[[YMAX:.+]] = arith.constant 1
// CHECK-DAG: %[[XMAX:.+]] = arith.constant 1 // CHECK-DAG: %[[XMAX:.+]] = arith.constant 1
▲ Show 20 Lines • Show All 51 Lines • ▼ Show 20 Linesfunc.func @resize_nearest_int(%input: tensor<1x2x2x1xi32>) -> () {
%output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "NEAREST_NEIGHBOR" } : (tensor<1x2x2x1xi32>) -> (tensor<1x4x4x1xi32>) %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "NEAREST_NEIGHBOR" } : (tensor<1x2x2x1xi32>) -> (tensor<1x4x4x1xi32>)
return return
} }
// ----- // -----
// CHECK-LABEL: @resize_bilinear_int // CHECK-LABEL: @resize_bilinear_int
func.func @resize_bilinear_int(%input: tensor<1x2x2x1xi8>) -> () { func.func @resize_bilinear_int(%input: tensor<1x2x2x1xi8>) -> () {
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 4, 4, 1] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[IDX3:.+]] = linalg.index 3 // CHECK: %[[IDX3:.+]] = linalg.index 3
// CHECK: %[[XYMIN:.+]] = arith.constant 0 // CHECK: %[[XYMIN:.+]] = arith.constant 0
// CHECK: %[[YMAX:.+]] = arith.constant 1 // CHECK: %[[YMAX:.+]] = arith.constant 1
// CHECK: %[[XMAX:.+]] = arith.constant 1 // CHECK: %[[XMAX:.+]] = arith.constant 1
▲ Show 20 Lines • Show All 70 Lines • ▼ Show 20 Lines
} }
// ----- // -----
// CHECK-LABEL: @resize_dyn // CHECK-LABEL: @resize_dyn
func.func @resize_dyn(%input: tensor<?x2x2x1xi8>) -> () { func.func @resize_dyn(%input: tensor<?x2x2x1xi8>) -> () {
// CHECK: %[[C0:.+]] = arith.constant 0 // CHECK: %[[C0:.+]] = arith.constant 0
// CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]] // CHECK: %[[BATCH:.+]] = tensor.dim %arg0, %[[C0]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[BATCH]], 4, 4, 1] // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
%output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "BILINEAR" } : (tensor<?x2x2x1xi8>) -> (tensor<?x4x4x1xi32>) %output = "tosa.resize"(%input) { output_size = [4, 4], stride = [128, 128], offset = [1, 2], stride_fp = [0. : f32, 0. : f32], offset_fp = [0. : f32, 0. : f32], shift = 8 : i32, mode = "BILINEAR" } : (tensor<?x2x2x1xi8>) -> (tensor<?x4x4x1xi32>)
return return
} }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-alloc-tensor-elimination.mlir

Show First 20 Lines • Show All 41 Lines • ▼ Show 20 Lines
func.func @buffer_forwarding_no_conflict( func.func @buffer_forwarding_no_conflict(
%t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true}, %t: tensor<?xf32> {bufferization.buffer_layout = affine_map<(d0) -> (d0)>, bufferization.writable = true},
%sz: index) %sz: index)
-> (tensor<?xf32>) -> (tensor<?xf32>)
{ {
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
// alloc_tensor itself does not alloc but forwards to the insert_slice. // alloc_tensor itself does not alloc but forwards to the insert_slice.
// InitTensorOp replaces the alloc_tensor with an inplace extract_slice. // AllocTensorOpElimination replaces the alloc_tensor with an inplace
// extract_slice.
// CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1] // CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
%a = bufferization.alloc_tensor(%sz) : tensor<?xf32> %a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
// CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW]] : memref<?xf32 // CHECK: linalg.fill ins({{.*}} : f32) outs(%[[T_SUBVIEW]] : memref<?xf32
%f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32> %f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32>
// Self-copy canonicalizes away later. // Self-copy canonicalizes away later.
%r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32> %r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
▲ Show 20 Lines • Show All 80 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/bubble-up-extract-slice-op.mlir

Show First 20 Lines • Show All 129 Lines • ▼ Show 20 Linesfunc.func @conv_slice(%input: tensor<1x225x225x3xf32>, %filter: tensor<3x3x3x32xf32>) -> tensor<1x32x32x16xf32> {
%c112 = arith.constant 112 : index %c112 = arith.constant 112 : index
%c32 = arith.constant 32 : index %c32 = arith.constant 32 : index
%c16 = arith.constant 16 : index %c16 = arith.constant 16 : index
%c8 = arith.constant 8 : index %c8 = arith.constant 8 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%init = linalg.init_tensor [1, 112, 112, 32] : tensor<1x112x112x32xf32> %init = tensor.empty() : tensor<1x112x112x32xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32>
%conv = linalg.conv_2d_nhwc_hwcf %conv = linalg.conv_2d_nhwc_hwcf
{dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
ins(%input, %filter : tensor<1x225x225x3xf32>, tensor<3x3x3x32xf32>) ins(%input, %filter : tensor<1x225x225x3xf32>, tensor<3x3x3x32xf32>)
outs(%fill : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> outs(%fill : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32>
%slice = tensor.extract_slice %conv [0, 64, 64, 16] [1, 32, 32, 16] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x32x32x16xf32> %slice = tensor.extract_slice %conv [0, 64, 64, 16] [1, 32, 32, 16] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x32x32x16xf32>
return %slice : tensor<1x32x32x16xf32> return %slice : tensor<1x32x32x16xf32>
} }
// CHECK: func @conv_slice // CHECK: func @conv_slice
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 112, 112, 32] : tensor<1x112x112x32xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x112x112x32xf32>
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[0, 128, 128, 0] [1, 65, 65, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x65x65x3xf32> // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[0, 128, 128, 0] [1, 65, 65, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x65x65x3xf32>
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[0, 0, 0, 16] [3, 3, 3, 16] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x16xf32> // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[0, 0, 0, 16] [3, 3, 3, 16] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x16xf32>
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[INIT]][0, 64, 64, 16] [1, 32, 32, 16] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x32x32x16xf32> // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[INIT]][0, 64, 64, 16] [1, 32, 32, 16] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x32x32x16xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST:.+]] : f32) outs(%[[SLICE2]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32> // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST:.+]] : f32) outs(%[[SLICE2]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32>
// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%[[SLICE0]], %[[SLICE1]] : tensor<1x65x65x3xf32>, tensor<3x3x3x16xf32>) outs(%[[FILL]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32> // CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%[[SLICE0]], %[[SLICE1]] : tensor<1x65x65x3xf32>, tensor<3x3x3x16xf32>) outs(%[[FILL]] : tensor<1x32x32x16xf32>) -> tensor<1x32x32x16xf32>
// CHECK: return %[[CONV]] : tensor<1x32x32x16xf32> // CHECK: return %[[CONV]] : tensor<1x32x32x16xf32>
//----- //-----
// The slice is not supposed to be bubbled up when it is rank-reducing. // The slice is not supposed to be bubbled up when it is rank-reducing.
func.func @rank_reducing_slice(%width : index) -> tensor<1x1x1x?xf32> { func.func @rank_reducing_slice(%width : index) -> tensor<1x1x1x?xf32> {
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%init = linalg.init_tensor [1, %width] : tensor<1x?xf32> %init = tensor.empty(%width) : tensor<1x?xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x?xf32>) -> tensor<1x?xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x?xf32>) -> tensor<1x?xf32>
%slice = tensor.extract_slice %fill[0, 0] [1, %width] [1, 1] : tensor<1x?xf32> to tensor<?xf32> %slice = tensor.extract_slice %fill[0, 0] [1, %width] [1, 1] : tensor<1x?xf32> to tensor<?xf32>
%expand = tensor.expand_shape %slice [[0, 1, 2, 3]] : tensor<?xf32> into tensor<1x1x1x?xf32> %expand = tensor.expand_shape %slice [[0, 1, 2, 3]] : tensor<?xf32> into tensor<1x1x1x?xf32>
return %expand : tensor<1x1x1x?xf32> return %expand : tensor<1x1x1x?xf32>
} }
// CHECK: func @rank_reducing_slice // CHECK: func @rank_reducing_slice
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[FILL:.+]] = linalg.fill ins // CHECK: %[[FILL:.+]] = linalg.fill ins
// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]] // CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]]
// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]] // CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[SLICE]]
// CHECK: return %[[EXPAND]] // CHECK: return %[[EXPAND]]

mlir/test/Dialect/Linalg/bufferize.mlir

Show All 35 Linesfunc.func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> {
return %0 : tensor<4xf32> return %0 : tensor<4xf32>
} }
// ----- // -----
#map0 = affine_map<(d0) -> (d0)> #map0 = affine_map<(d0) -> (d0)>
// Same as above but with linalg.init_tensor op. // Same as above but with tensor.empty op.
// CHECK: #map = affine_map<(d0) -> (d0)> // CHECK: #map = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @init_tensor( // CHECK-LABEL: func @empty_tensor(
// CHECK-SAME: %[[IN:.*]]: tensor<?xf32>, %[[SIZE:.*]]: index) // CHECK-SAME: %[[IN:.*]]: tensor<?xf32>, %[[SIZE:.*]]: index)
// CHECK-DAG: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<?xf32> // CHECK-DAG: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<?xf32>
// CHECK-DAG: %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) {{.*}} : memref<?xf32> // CHECK-DAG: %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) {{.*}} : memref<?xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[MEMREF]] : memref<?xf32>) // CHECK-SAME: ins(%[[MEMREF]] : memref<?xf32>)
// CHECK-SAME: outs(%[[OUT_BUF]] : memref<?xf32>) { // CHECK-SAME: outs(%[[OUT_BUF]] : memref<?xf32>) {
func.func @init_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> { func.func @empty_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> {
%init = linalg.init_tensor [%size] : tensor<?xf32> %init = tensor.empty(%size) : tensor<?xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map0, #map0], indexing_maps = [#map0, #map0],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%in : tensor<?xf32>) } ins(%in : tensor<?xf32>)
outs(%init : tensor<?xf32>) { outs(%init : tensor<?xf32>) {
^bb0(%gen_arg1: f32, %out: f32): ^bb0(%gen_arg1: f32, %out: f32):
%tmp1 = math.exp %gen_arg1 : f32 %tmp1 = math.exp %gen_arg1 : f32
linalg.yield %tmp1 : f32 linalg.yield %tmp1 : f32
▲ Show 20 Lines • Show All 139 Lines • ▼ Show 20 Lines
// CHECK: %[[collapse_m:.*]] = bufferization.to_memref %[[collapse]] // CHECK: %[[collapse_m:.*]] = bufferization.to_memref %[[collapse]]
// CHECK: %[[alloc:.*]] = memref.alloc() // CHECK: %[[alloc:.*]] = memref.alloc()
// CHECK: linalg.generic {{.*}} ins(%[[collapse_m]] : memref<6xi1>) outs(%[[alloc]] : memref<6xi64>) // CHECK: linalg.generic {{.*}} ins(%[[collapse_m]] : memref<6xi1>) outs(%[[alloc]] : memref<6xi64>)
// CHECK: %[[generic_t:.*]] = bufferization.to_tensor %[[alloc]] // CHECK: %[[generic_t:.*]] = bufferization.to_tensor %[[alloc]]
// CHECK: %[[call:.*]] = call @csum(%[[generic_t]]) // CHECK: %[[call:.*]] = call @csum(%[[generic_t]])
// CHECK: return %[[call]] // CHECK: return %[[call]]
func.func public @main(%arg0: tensor<2x3xi1>) -> tensor<6xi64> { func.func public @main(%arg0: tensor<2x3xi1>) -> tensor<6xi64> {
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<2x3xi1> into tensor<6xi1> %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<2x3xi1> into tensor<6xi1>
%1 = linalg.init_tensor [6] : tensor<6xi64> %1 = tensor.empty() : tensor<6xi64>
%2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%0 : tensor<6xi1>) outs(%1 : tensor<6xi64>) { %2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%0 : tensor<6xi1>) outs(%1 : tensor<6xi64>) {
^bb0(%arg1: i1, %arg2: i64): ^bb0(%arg1: i1, %arg2: i64):
%4 = arith.extui %arg1 : i1 to i64 %4 = arith.extui %arg1 : i1 to i64
linalg.yield %4 : i64 linalg.yield %4 : i64
} -> tensor<6xi64> } -> tensor<6xi64>
%3 = func.call @csum(%2) : (tensor<6xi64>) -> tensor<6xi64> %3 = func.call @csum(%2) : (tensor<6xi64>) -> tensor<6xi64>
return %3 : tensor<6xi64> return %3 : tensor<6xi64>
} }

mlir/test/Dialect/Linalg/canonicalize-duplicate-inputs.mlir

Show First 20 Lines • Show All 125 Lines • ▼ Show 20 Lines
// ----- // -----
// Current argmax lowering to `linalg.generic`. Cannot drop the // Current argmax lowering to `linalg.generic`. Cannot drop the
// first return even though it isnt used since it has an internal // first return even though it isnt used since it has an internal
// use. // use.
#map0 = affine_map<(d0) -> (d0)> #map0 = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> ()> #map1 = affine_map<(d0) -> ()>
func.func @argmax_lowering(%arg0 : tensor<?xf32>) -> tensor<i32> { func.func @argmax_lowering(%arg0 : tensor<?xf32>) -> tensor<i32> {
%init0 = linalg.init_tensor [] : tensor<f32> %init0 = tensor.empty() : tensor<f32>
%init1 = linalg.init_tensor [] : tensor<i32> %init1 = tensor.empty() : tensor<i32>
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [#map0, #map1, #map1], indexing_maps = [#map0, #map1, #map1],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0 : tensor<?xf32>) ins(%arg0 : tensor<?xf32>)
outs(%init0, %init1 : tensor<f32>, tensor<i32>) { outs(%init0, %init1 : tensor<f32>, tensor<i32>) {
^bb0(%b0: f32, %b1: f32, %b2: i32): ^bb0(%b0: f32, %b1: f32, %b2: i32):
%8 = linalg.index 0 : index %8 = linalg.index 0 : index
%9 = arith.index_cast %8 : index to i32 %9 = arith.index_cast %8 : index to i32
%10 = arith.cmpf oge, %b0, %b1 : f32 %10 = arith.cmpf oge, %b0, %b1 : f32
%11 = arith.select %10, %b0, %b1 : f32 %11 = arith.select %10, %b0, %b1 : f32
%12 = arith.cmpf oeq, %b0, %b1 : f32 %12 = arith.cmpf oeq, %b0, %b1 : f32
%13 = arith.minsi %9, %b2 : i32 %13 = arith.minsi %9, %b2 : i32
%14 = arith.select %10, %9, %b2 : i32 %14 = arith.select %10, %9, %b2 : i32
%15 = arith.select %12, %13, %14 : i32 %15 = arith.select %12, %13, %14 : i32
linalg.yield %11, %15 : f32, i32 linalg.yield %11, %15 : f32, i32
} -> (tensor<f32>, tensor<i32>) } -> (tensor<f32>, tensor<i32>)
return %0#1 : tensor<i32> return %0#1 : tensor<i32>
} }
// CHECK: func @argmax_lowering( // CHECK: func @argmax_lowering(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?xf32>
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [] : tensor<f32> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<f32>
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [] : tensor<i32> // CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<i32>
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] :
// CHECK: return %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#1
// ----- // -----
// Do not remove operand needed for loop dim. // Do not remove operand needed for loop dim.
func.func @loop_dim_operand(%arg0 : tensor<?xf32>) -> tensor<i32> { func.func @loop_dim_operand(%arg0 : tensor<?xf32>) -> tensor<i32> {
%cst = arith.constant 0 : i32 %cst = arith.constant 0 : i32
%init = linalg.init_tensor [] : tensor<i32> %init = tensor.empty() : tensor<i32>
%fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32> %fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>], indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0 : tensor<?xf32>) outs(%fill : tensor<i32>) { ins(%arg0 : tensor<?xf32>) outs(%fill : tensor<i32>) {
^bb0(%b0: f32, %b1: i32): ^bb0(%b0: f32, %b1: i32):
%1 = linalg.index 0 : index %1 = linalg.index 0 : index
%2 = arith.index_cast %1 : index to i32 %2 = arith.index_cast %1 : index to i32
%3 = arith.addi %b1, %2 : i32 %3 = arith.addi %b1, %2 : i32
linalg.yield %3 : i32 linalg.yield %3 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
// CHECK: func @loop_dim_operand( // CHECK: func @loop_dim_operand(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0]] : // CHECK-SAME: ins(%[[ARG0]] :
// ----- // -----
// Do not remove outs operand needed for loop bound computation. // Do not remove outs operand needed for loop bound computation.
func.func @loop_dim_outs_operand(%arg0 : index) -> tensor<i32> { func.func @loop_dim_outs_operand(%arg0 : index) -> tensor<i32> {
%cst = arith.constant 0 : i32 %cst = arith.constant 0 : i32
%init1 = linalg.init_tensor [%arg0] : tensor<?xi32> %init1 = tensor.empty(%arg0) : tensor<?xi32>
%init = linalg.init_tensor [] : tensor<i32> %init = tensor.empty() : tensor<i32>
%fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32> %fill = linalg.fill ins(%cst : i32) outs(%init : tensor<i32>) -> tensor<i32>
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>], indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%init1, %fill : tensor<?xi32>, tensor<i32>) { outs(%init1, %fill : tensor<?xi32>, tensor<i32>) {
^bb0(%b0: i32, %b1: i32): ^bb0(%b0: i32, %b1: i32):
%1 = linalg.index 0 : index %1 = linalg.index 0 : index
%2 = arith.index_cast %1 : index to i32 %2 = arith.index_cast %1 : index to i32
%3 = arith.addi %b1, %2 : i32 %3 = arith.addi %b1, %2 : i32
linalg.yield %2, %3 : i32, i32 linalg.yield %2, %3 : i32, i32
} -> (tensor<?xi32>, tensor<i32>) } -> (tensor<?xi32>, tensor<i32>)
return %0#1 : tensor<i32> return %0#1 : tensor<i32>
} }
// CHECK: func @loop_dim_outs_operand( // CHECK: func @loop_dim_outs_operand(
// CHECK-SAME: %[[ARG0:.+]]: index // CHECK-SAME: %[[ARG0:.+]]: index
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[ARG0]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[ARG0]])
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: outs(%[[INIT]] // CHECK-SAME: outs(%[[INIT]]
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1, d0)> #map1 = affine_map<(d0, d1) -> (d1, d0)>
#map2 = affine_map<(d0, d1) -> (d0)> #map2 = affine_map<(d0, d1) -> (d0)>
▲ Show 20 Lines • Show All 72 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/canonicalize.mlir

Show First 20 Lines • Show All 79 Lines • ▼ Show 20 Linesfunc.func @linalg_effects(%a : tensor<?x?xf32>, %b : memref<?x?xf32>, %c : tensor<?x?xf32>) {
// CHECK: linalg.matmul // CHECK: linalg.matmul
linalg.matmul ins(%a, %c : tensor<?x?xf32>, tensor<?x?xf32>) linalg.matmul ins(%a, %c : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%b : memref<?x?xf32>) outs(%b : memref<?x?xf32>)
return return
} }
// ----- // -----
func.func @init_tensor_canonicalize() -> (tensor<4x5x?xf32>) {
%c6 = arith.constant 6 : index
%0 = linalg.init_tensor [4, 5, %c6] : tensor<4x5x?xf32>
return %0 : tensor<4x5x?xf32>
}
// CHECK: func @init_tensor_canonicalize
// CHECK: %[[T0:.+]] = linalg.init_tensor [4, 5, 6] : tensor<4x5x6xf32>
// CHECK: %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>
// CHECK: return %[[T1]]
// -----
func.func @init_tensor_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
%0 = linalg.init_tensor [6, 5, %arg0] : tensor<6x5x?xf32>
%1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
: tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
return %1 : tensor<2x3x5x4x?x7xf32>
}
// CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
// CHECK: func @init_tensor_reshape_expansion
// CHECK-SAME: %[[ARG0:.+]]: index
// CHECK-NEXT: %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
// CHECK-NEXT: %[[INIT:.+]] = linalg.init_tensor [2, 3, 5, 4, %[[D]], 7]
// CHECK-NEXT: return %[[INIT]]
// -----
func.func @init_tensor_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
%0 = linalg.init_tensor [2, 3, 5, 4, %arg0, 7] : tensor<2x3x5x4x?x7xf32>
%1 = tensor.collapse_shape %0 [[0, 1], [2], [3, 4, 5]]
: tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
return %1 : tensor<6x5x?xf32>
}
// CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 * 28)>
// CHECK: func @init_tensor_reshape_collapse
// CHECK-SAME: %[[ARG0:.+]]: index
// CHECK-NEXT: %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
// CHECK-NEXT: %[[INIT:.+]] = linalg.init_tensor [6, 5, %[[D]]]
// CHECK-NEXT: return %[[INIT]]
// -----
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @remove_no_op(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) func.func @remove_no_op(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>)
-> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) { -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4, %5 = linalg.generic { %4, %5 = linalg.generic {
indexing_maps = [#map, #map, #map, #map], indexing_maps = [#map, #map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) } ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
outs(%3, %3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) { outs(%3, %3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg3, %arg2 : f32, f32 linalg.yield %arg3, %arg2 : f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %4, %5 : tensor<?x?x?xf32>, tensor<?x?x?xf32> return %4, %5 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
} }
// CHECK-LABEL: func @remove_no_op // CHECK-LABEL: func @remove_no_op
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK: return %[[ARG1]], %[[ARG0]] // CHECK: return %[[ARG1]], %[[ARG0]]
// ----- // -----
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @remove_no_op_mismatched_types(%arg0 : tensor<?x?x?xf32>) func.func @remove_no_op_mismatched_types(%arg0 : tensor<?x?x?xf32>)
-> tensor<1x2x3xf32> { -> tensor<1x2x3xf32> {
%out = linalg.init_tensor [1, 2, 3] : tensor<1x2x3xf32> %out = tensor.empty() : tensor<1x2x3xf32>
%g = linalg.generic { %g = linalg.generic {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0 : tensor<?x?x?xf32>) } ins(%arg0 : tensor<?x?x?xf32>)
outs(%out : tensor<1x2x3xf32>) { outs(%out : tensor<1x2x3xf32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> (tensor<1x2x3xf32>) } -> (tensor<1x2x3xf32>)
return %g : tensor<1x2x3xf32> return %g : tensor<1x2x3xf32>
} }
// CHECK-LABEL: func @remove_no_op_mismatched_types // CHECK-LABEL: func @remove_no_op_mismatched_types
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK: %[[CAST:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<1x2x3xf32> // CHECK: %[[CAST:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<1x2x3xf32>
// CHECK: return %[[CAST]] // CHECK: return %[[CAST]]
// ----- // -----
#map = affine_map<() -> ()> #map = affine_map<() -> ()>
func.func @cant_fold_to_tensor_cast(%arg0 : f32) -> tensor<f32> { func.func @cant_fold_to_tensor_cast(%arg0 : f32) -> tensor<f32> {
%out = linalg.init_tensor [] : tensor<f32> %out = tensor.empty() : tensor<f32>
%g = linalg.generic { %g = linalg.generic {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = [] iterator_types = []
} ins(%arg0 : f32) } ins(%arg0 : f32)
outs(%out : tensor<f32>) { outs(%out : tensor<f32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> (tensor<f32>) } -> (tensor<f32>)
return %g : tensor<f32> return %g : tensor<f32>
} }
// CHECK-LABEL: func @cant_fold_to_tensor_cast // CHECK-LABEL: func @cant_fold_to_tensor_cast
// CHECK: linalg.generic // CHECK: linalg.generic
// ----- // -----
#map = affine_map<(d0, d1) -> (d0, d1)> #map = affine_map<(d0, d1) -> (d0, d1)>
func.func @keep_not_noop(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @keep_not_noop(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
cf.br ^bb1(%cst : f32) cf.br ^bb1(%cst : f32)
^bb1(%arg1 : f32): ^bb1(%arg1 : f32):
%3 = linalg.generic %3 = linalg.generic
{indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?x?xf32>) outs(%2 : tensor<?x?xf32>) { ins(%arg0 : tensor<?x?xf32>) outs(%2 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3 : f32): ^bb0(%arg2: f32, %arg3 : f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
Show All 9 Lines
#map = affine_map<(d0, d1) -> (d0, d1)> #map = affine_map<(d0, d1) -> (d0, d1)>
func.func @keep_not_noop(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) func.func @keep_not_noop(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>) { -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
cf.br ^bb1(%cst : f32) cf.br ^bb1(%cst : f32)
^bb1(%arg2 : f32): ^bb1(%arg2 : f32):
%3:2 = linalg.generic %3:2 = linalg.generic
{indexing_maps = [#map, #map, #map, #map], {indexing_maps = [#map, #map, #map, #map],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>) { outs(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4 : f32, %arg5 : f32, %arg6 : f32): ^bb0(%arg3: f32, %arg4 : f32, %arg5 : f32, %arg6 : f32):
linalg.yield %arg2, %arg4 : f32, f32 linalg.yield %arg2, %arg4 : f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>)
return %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32> return %3#0, %3#1 : tensor<?x?xf32>, tensor<?x?xf32>
} }
// CHECK-LABEL: func @keep_not_noop // CHECK-LABEL: func @keep_not_noop
// CHECK: %[[RESULT:.+]]:2 = linalg.generic // CHECK: %[[RESULT:.+]]:2 = linalg.generic
// CHECK: return %[[RESULT]]#0, %[[RESULT]]#1 // CHECK: return %[[RESULT]]#0, %[[RESULT]]#1
// ----- // -----
func.func @fold_init_tensor_with_slice
(%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
{
%0 = linalg.init_tensor[%arg0, 10, 40] : tensor<?x10x40xf32>
%1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
: tensor<?x10x40xf32> to tensor<5x?x20xf32>
return %1 : tensor<5x?x20xf32>
}
// CHECK: func @fold_init_tensor_with_slice
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK: %[[T0:.+]] = linalg.init_tensor [5, %[[ARG1]], 20]
// CHECK: return %[[T0]]
// -----
func.func @fold_init_tensor_with_cast(%arg0 : index) -> tensor<1x12xf32> {
%0 = linalg.init_tensor [%arg0, 12] : tensor<?x12xf32>
%1 = tensor.cast %0 : tensor<?x12xf32> to tensor<1x12xf32>
return %1 : tensor<1x12xf32>
}
// CHECK: func @fold_init_tensor_with_cast(%[[ARG0:.+]]: index)
// CHECK: %[[T0:.+]] = linalg.init_tensor [1, 12] : tensor<1x12xf32>
// CHECK: return %[[T0]] : tensor<1x12xf32>
// -----
#accesses = [ #accesses = [
affine_map<(i, j) -> (i, j)> affine_map<(i, j) -> (i, j)>
] ]
#trait = { #trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
Show All 26 Lines
// ----- // -----
func.func @propogate_casts(%arg0 : tensor<?x?xf32>, %arg1 : f32, %arg2 : index, func.func @propogate_casts(%arg0 : tensor<?x?xf32>, %arg1 : f32, %arg2 : index,
%arg3 : index) -> tensor<?x?xf32> { %arg3 : index) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c21 = arith.constant 21 : index %c21 = arith.constant 21 : index
%c42 = arith.constant 42 : index %c42 = arith.constant 42 : index
%0 = linalg.init_tensor [%c21, %c42] : tensor<?x?xf32> %0 = tensor.empty(%c21, %c42) : tensor<?x?xf32>
%1 = linalg.fill ins(%arg1 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32> %1 = linalg.fill ins(%arg1 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%2 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %2 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%3 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %3 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%4 = tensor.insert_slice %arg0 into %1[%arg2, %arg3] [%2, %3] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %4 = tensor.insert_slice %arg0 into %1[%arg2, %arg3] [%2, %3] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
return %4 : tensor<?x?xf32> return %4 : tensor<?x?xf32>
} }
// CHECK-LABEL: func @propogate_casts // CHECK-LABEL: func @propogate_casts
// CHECK: %[[INIT:.+]] = linalg.init_tensor [21, 42] // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]
// CHECK: %[[INSERTED:.+]] = tensor.insert_slice %{{.+}} into %[[FILL]] // CHECK: %[[INSERTED:.+]] = tensor.insert_slice %{{.+}} into %[[FILL]]
// CHECK: %[[RESULT:.+]] = tensor.cast %[[INSERTED]] // CHECK: %[[RESULT:.+]] = tensor.cast %[[INSERTED]]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
// CHECK-LABEL: @self_copy // CHECK-LABEL: @self_copy
Show All 13 Lines
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>) // CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)
// CHECK-SAME: outs({{.*}} : tensor<?xf32>) { // CHECK-SAME: outs({{.*}} : tensor<?xf32>) {
#map0 = affine_map<(d0) -> (d0)> #map0 = affine_map<(d0) -> (d0)>
func.func @remove_deadargs_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) { func.func @remove_deadargs_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 7.0 : f32 %cst = arith.constant 7.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
%1 = linalg.init_tensor [%0] : tensor<?xf32> %1 = tensor.empty(%0) : tensor<?xf32>
%2 = linalg.init_tensor [%0] : tensor<?xf32> %2 = tensor.empty(%0) : tensor<?xf32>
%3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %1 : tensor<?xf32>, tensor<?xf32>) outs (%2:tensor<?xf32>) { %3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %1 : tensor<?xf32>, tensor<?xf32>) outs (%2:tensor<?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%4 = arith.addf %arg1, %cst : f32 %4 = arith.addf %arg1, %cst : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %3 : tensor<?xf32> return %3 : tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @remove_deadargs_generic_mixedaccess // CHECK-LABEL: func @remove_deadargs_generic_mixedaccess
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-NOT: ins // CHECK-NOT: ins
// CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) { // CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1, d0)> #map1 = affine_map<(d0, d1) -> (d1, d0)>
func.func @remove_deadargs_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) { func.func @remove_deadargs_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 0 : index %c1 = arith.constant 0 : index
%cst1 = arith.constant 7.0 : f32 %cst1 = arith.constant 7.0 : f32
%cst2 = arith.constant 6.0 : f32 %cst2 = arith.constant 6.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.init_tensor [%1, %0] : tensor<?x?xf32> %3 = tensor.empty(%1, %0) : tensor<?x?xf32>
%4 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %4 = tensor.empty(%0, %1) : tensor<?x?xf32>
%5 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%2, %3 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%4:tensor<?x?xf32>) { %5 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%2, %3 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%4:tensor<?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%6 = arith.divf %cst1, %cst2 : f32 %6 = arith.divf %cst1, %cst2 : f32
linalg.yield %6 : f32 linalg.yield %6 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %5 : tensor<?x?xf32> return %5 : tensor<?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fold_fill_reshape() // CHECK-LABEL: func @fold_fill_reshape()
func.func @fold_fill_reshape() -> tensor<6x4xf32> { func.func @fold_fill_reshape() -> tensor<6x4xf32> {
%zero = arith.constant 0.0 : f32 %zero = arith.constant 0.0 : f32
// CHECK: %[[INIT:.+]] = linalg.init_tensor [6, 4] : tensor<6x4xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<6x4xf32>
%init = linalg.init_tensor [1, 2, 3, 4] : tensor<1x2x3x4xf32> %empty = tensor.empty() : tensor<1x2x3x4xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<6x4xf32>) -> tensor<6x4xf32> // CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<6x4xf32>) -> tensor<6x4xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> %fill = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32>
%reshape = tensor.collapse_shape %fill [[0, 1, 2], [3]] %reshape = tensor.collapse_shape %fill [[0, 1, 2], [3]]
: tensor<1x2x3x4xf32> into tensor<6x4xf32> : tensor<1x2x3x4xf32> into tensor<6x4xf32>
// CHECK: return %[[FILL]] : tensor<6x4xf32> // CHECK: return %[[FILL]] : tensor<6x4xf32>
return %reshape : tensor<6x4xf32> return %reshape : tensor<6x4xf32>
} }
// ----- // -----
// CHECK: func @fold_fill_reshape_dynamic // CHECK: func @fold_fill_reshape_dynamic
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?x?xf32>
func.func @fold_fill_reshape_dynamic(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> { func.func @fold_fill_reshape_dynamic(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32> {
%zero = arith.constant 0.0 : f32 %zero = arith.constant 0.0 : f32
// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] // CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
%0 = linalg.fill ins(%zero : f32) outs(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> %0 = linalg.fill ins(%zero : f32) outs(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
// CHECK: %[[RESULT:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[RESHAPE]] // CHECK: %[[RESULT:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[RESHAPE]]
%1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]] %1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]]
: tensor<?x?x?x?x?xf32> into tensor<?x?xf32> : tensor<?x?x?x?x?xf32> into tensor<?x?xf32>
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
// -----
func.func private @some_use(%i : index, %j : index)
// CHECK-LABEL: func @init_canonicalize
// CHECK-SAME: %[[I:.*]]: index
func.func @init_canonicalize(%i : index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
// CHECK-NOT: init_tensor
%0 = linalg.init_tensor [%i, 42] : tensor<?x42xf32>
// CHECK-NOT: tensor.dim
%1 = tensor.dim %0, %c0: tensor<?x42xf32>
%2 = tensor.dim %0, %c1: tensor<?x42xf32>
// CHECK: %[[c42:.*]] = arith.constant 42 : index
// CHECK: call @some_use(%[[I]], %[[c42]])
call @some_use(%1, %2) : (index, index) -> ()
return
}
// -----
// CHECK-LABEL: func @rank_reducing_init_extract
func.func @rank_reducing_init_extract(%sz : index, %idx : index) -> tensor<2xf32> {
// CHECK: linalg.init_tensor [2] : tensor<2xf32>
%a = linalg.init_tensor [%sz, 2] : tensor<?x2xf32>
// CHECK-NOT: extract
%r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
return %r: tensor<2xf32>
}
// ----- // -----
// CHECK: func @fold_self_copy // CHECK: func @fold_self_copy
func.func @fold_self_copy(%0 : memref<4x16xf32>) { func.func @fold_self_copy(%0 : memref<4x16xf32>) {
// CHECK-NEXT: return // CHECK-NEXT: return
linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%0 : memref<4x16xf32>) ins(%0 : memref<4x16xf32>)
outs(%0 : memref<4x16xf32>) { outs(%0 : memref<4x16xf32>) {
^bb0(%arg4: f32, %arg5: f32): ^bb0(%arg4: f32, %arg5: f32):
linalg.yield %arg4 : f32 linalg.yield %arg4 : f32
} }
return return
} }
// ----- // -----
// CHECK-LABEL: func @fold_static_pad_fill // CHECK-LABEL: func @fold_static_pad_fill
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[INIT:.+]] = linalg.init_tensor [412, 276] : tensor<412x276xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<412x276xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]]
// CHECK: return %[[FILL]] // CHECK: return %[[FILL]]
func.func @fold_static_pad_fill() -> tensor<412x276xf32> { func.func @fold_static_pad_fill() -> tensor<412x276xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%init = linalg.init_tensor [400, 273] : tensor<400x273xf32> %empty = tensor.empty() : tensor<400x273xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<400x273xf32>) -> tensor<400x273xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<400x273xf32>) -> tensor<400x273xf32>
%pad = tensor.pad %fill low[4, 1] high[8, 2] { %pad = tensor.pad %fill low[4, 1] high[8, 2] {
^bb0(%arg1: index, %arg2: index): ^bb0(%arg1: index, %arg2: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<400x273xf32> to tensor<412x276xf32> } : tensor<400x273xf32> to tensor<412x276xf32>
return %pad : tensor<412x276xf32> return %pad : tensor<412x276xf32>
} }
// ----- // -----
Show All 9 Lines
// CHECK-DAG: %[[I1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[I1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[OF:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[SRC]] : tensor<8x?x16x32xf32>) // CHECK: %[[OF:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[SRC]] : tensor<8x?x16x32xf32>)
// CHECK: %[[S0:.+]] = affine.apply #[[MAP0]]()[%[[LOW0]]] // CHECK: %[[S0:.+]] = affine.apply #[[MAP0]]()[%[[LOW0]]]
// CHECK: %[[DIM1:.+]] = tensor.dim %[[OF]], %[[I1]] : tensor<8x?x16x32xf32> // CHECK: %[[DIM1:.+]] = tensor.dim %[[OF]], %[[I1]] : tensor<8x?x16x32xf32>
// CHECK: %[[S1:.+]] = affine.apply #[[MAP1]]()[%[[DIM1]]] // CHECK: %[[S1:.+]] = affine.apply #[[MAP1]]()[%[[DIM1]]]
// CHECK: %[[S2:.+]] = affine.apply #[[MAP2]]()[%[[HIGH2]]] // CHECK: %[[S2:.+]] = affine.apply #[[MAP2]]()[%[[HIGH2]]]
// CHECK: %[[S3:.+]] = affine.apply #[[MAP3]]()[%[[LOW3]], %[[HIGH3]]] // CHECK: %[[S3:.+]] = affine.apply #[[MAP3]]()[%[[LOW3]], %[[HIGH3]]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[S0]], %[[S1]], %[[S2]], %[[S3]]] : tensor<?x?x?x?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%[[S0]], %[[S1]], %[[S2]], %[[S3]]) : tensor<?x?x?x?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]]
// CHECK: return %[[FILL]] // CHECK: return %[[FILL]]
func.func @fold_dynamic_pad_fill(%init: tensor<8x?x16x32xf32>, %low0: index, %low3: index, %high2: index, %high3: index) -> tensor<?x?x?x?xf32> { func.func @fold_dynamic_pad_fill(%empty: tensor<8x?x16x32xf32>, %low0: index, %low3: index, %high2: index, %high3: index) -> tensor<?x?x?x?xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x?x16x32xf32>) -> tensor<8x?x16x32xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x?x16x32xf32>) -> tensor<8x?x16x32xf32>
%pad = tensor.pad %fill low[%low0, 8, 7, %low3] high[1, 2, %high2, %high3] { %pad = tensor.pad %fill low[%low0, 8, 7, %low3] high[1, 2, %high2, %high3] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<8x?x16x32xf32> to tensor<?x?x?x?xf32> } : tensor<8x?x16x32xf32> to tensor<?x?x?x?xf32>
return %pad : tensor<?x?x?x?xf32> return %pad : tensor<?x?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @no_fold_pad_fill_value_mismatch // CHECK-LABEL: func @no_fold_pad_fill_value_mismatch
func.func @no_fold_pad_fill_value_mismatch() -> tensor<412x276xf32> { func.func @no_fold_pad_fill_value_mismatch() -> tensor<412x276xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%f1 = arith.constant 1.0 : f32 %f1 = arith.constant 1.0 : f32
%init = linalg.init_tensor [400, 273] : tensor<400x273xf32> %empty = tensor.empty() : tensor<400x273xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<400x273xf32>) -> tensor<400x273xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<400x273xf32>) -> tensor<400x273xf32>
// CHECK: tensor.pad // CHECK: tensor.pad
%pad = tensor.pad %fill low[4, 1] high[8, 2] { %pad = tensor.pad %fill low[4, 1] high[8, 2] {
^bb0(%arg1: index, %arg2: index): ^bb0(%arg1: index, %arg2: index):
tensor.yield %f1 : f32 tensor.yield %f1 : f32
} : tensor<400x273xf32> to tensor<412x276xf32> } : tensor<400x273xf32> to tensor<412x276xf32>
return %pad : tensor<412x276xf32> return %pad : tensor<412x276xf32>
} }
// ----- // -----
// Tests below verify whether static information is propagated through all the operands of generic op. // Tests below verify whether static information is propagated through all the operands of generic op.
// 1. If one of the inputs of generic op has static info and it has no cast source. // 1. If one of the inputs of generic op has static info and it has no cast source.
// 2. If one of the inputs of generic op has static info and it is coming from tensr.cast operation. // 2. If one of the inputs of generic op has static info and it is coming from tensr.cast operation.
// 3. If one of the outputs of generic op has static info and it is coming from tenso.cast operation. // 3. If one of the outputs of generic op has static info and it is coming from tenso.cast operation.
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @static_input_without_cast // CHECK-LABEL: func @static_input_without_cast
// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
func.func @static_input_without_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> { func.func @static_input_without_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32> %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
%1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32> %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
%2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32> %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4 = linalg.generic { %4 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<?x?x?xf32>) } ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<?x?x?xf32>)
outs(%3 : tensor<?x?x?xf32>) { outs(%3 : tensor<?x?x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32): ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
%9 = arith.addf %arg2, %arg3 : f32 %9 = arith.addf %arg2, %arg3 : f32
linalg.yield %9 : f32 linalg.yield %9 : f32
Show All 13 Lines
// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
func.func @static_input_with_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> { func.func @static_input_with_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32> %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
%1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32> %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
%2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32> %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32> %4 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
%5 = linalg.generic { %5 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0, %4 : tensor<2x3x4xf32>, tensor<2x?x?xf32>) } ins(%arg0, %4 : tensor<2x3x4xf32>, tensor<2x?x?xf32>)
outs(%3 : tensor<?x?x?xf32>) { outs(%3 : tensor<?x?x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32): ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
%9 = arith.addf %arg2, %arg3 : f32 %9 = arith.addf %arg2, %arg3 : f32
Show All 14 Lines
// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>, %[[ARG2:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>, %[[ARG2:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
func.func @static_output_with_cast(%arg0 : tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>, %arg2: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { func.func @static_output_with_cast(%arg0 : tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>, %arg2: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg2, %c0 : tensor<2x3x4xf32> %0 = tensor.dim %arg2, %c0 : tensor<2x3x4xf32>
%1 = tensor.dim %arg2, %c1 : tensor<2x3x4xf32> %1 = tensor.dim %arg2, %c1 : tensor<2x3x4xf32>
%2 = tensor.dim %arg2, %c2 : tensor<2x3x4xf32> %2 = tensor.dim %arg2, %c2 : tensor<2x3x4xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4 = tensor.cast %3 : tensor<?x?x?xf32> to tensor<2x3x4xf32> %4 = tensor.cast %3 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
%5 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32> %5 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
%6 = linalg.generic { %6 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0, %5 : tensor<?x?x?xf32>, tensor<2x?x?xf32>) } ins(%arg0, %5 : tensor<?x?x?xf32>, tensor<2x?x?xf32>)
outs(%4 : tensor<2x3x4xf32>) { outs(%4 : tensor<2x3x4xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
Show All 17 Lines
// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
func.func @cast_source(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> { func.func @cast_source(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32> %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
%1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32> %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
%2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32> %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4 = tensor.cast %arg0 : tensor<2x3x4xf32> to tensor<2x?x?xf32> %4 = tensor.cast %arg0 : tensor<2x3x4xf32> to tensor<2x?x?xf32>
%5 = tensor.cast %arg1 : tensor<2x3x4xf32> to tensor<2x?x?xf32> %5 = tensor.cast %arg1 : tensor<2x3x4xf32> to tensor<2x?x?xf32>
%6 = linalg.generic { %6 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%4, %5 : tensor<2x?x?xf32>, tensor<2x?x?xf32>) } ins(%4, %5 : tensor<2x?x?xf32>, tensor<2x?x?xf32>)
outs(%3 : tensor<?x?x?xf32>) { outs(%3 : tensor<?x?x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32): ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
%9 = arith.addf %arg2, %arg3 : f32 %9 = arith.addf %arg2, %arg3 : f32
linalg.yield %9 : f32 linalg.yield %9 : f32
} -> (tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>)
%7 = tensor.cast %6 : tensor<?x?x?xf32> to tensor<2x3x4xf32> %7 = tensor.cast %6 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
return %7: tensor<2x3x4xf32> return %7: tensor<2x3x4xf32>
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>) // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
// CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>) // CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
} }
// ----- // -----
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @cast_dest // CHECK-LABEL: func @cast_dest
// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<1x?x?xf32>, // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<1x?x?xf32>,
func.func @cast_dest(%arg0: tensor<?x?x?xf32>, %arg1: tensor<1x?x?xf32>, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?x?xf32> { func.func @cast_dest(%arg0: tensor<?x?x?xf32>, %arg1: tensor<1x?x?xf32>, %arg2: index, %arg3: index, %arg4: index) -> tensor<?x?x?xf32> {
%0 = linalg.init_tensor [%arg2, %arg3, %arg4] : tensor<?x?x?xf32> %0 = tensor.empty(%arg2, %arg3, %arg4) : tensor<?x?x?xf32>
%1 = tensor.cast %arg1 : tensor<1x?x?xf32> to tensor<?x?x?xf32> %1 = tensor.cast %arg1 : tensor<1x?x?xf32> to tensor<?x?x?xf32>
%2 = linalg.generic { %2 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<1x?x?xf32>) } ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<1x?x?xf32>)
outs(%0 : tensor<?x?x?xf32>) { outs(%0 : tensor<?x?x?xf32>) {
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
%3 = arith.subf %arg5, %arg6 : f32 %3 = arith.subf %arg5, %arg6 : f32
Show All 10 Lines
// CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 + 1)> // CHECK: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 + 1)>
// CHECK-LABEL: func @insert_pad_into_fill // CHECK-LABEL: func @insert_pad_into_fill
// CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x?xf32>, %[[LOW0:.+]]: index, %[[LOW1:.+]]: index, %{{.+}}: index, %{{.+}}: index) // CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x?xf32>, %[[LOW0:.+]]: index, %[[LOW1:.+]]: index, %{{.+}}: index, %{{.+}}: index)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[INIT:.+]] = linalg.init_tensor [8, 384, 384] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]]
// CHECK: %[[OFFSET1:.+]] = affine.apply #[[$MAP]]()[%[[LOW1]]] // CHECK: %[[OFFSET1:.+]] = affine.apply #[[$MAP]]()[%[[LOW1]]]
// CHECK: %[[D0:.+]] = tensor.dim %[[INPUT]], %[[C0]] : tensor<?x?x?xf32> // CHECK: %[[D0:.+]] = tensor.dim %[[INPUT]], %[[C0]] : tensor<?x?x?xf32>
// CHECK: %[[D1:.+]] = tensor.dim %[[INPUT]], %[[C1]] : tensor<?x?x?xf32> // CHECK: %[[D1:.+]] = tensor.dim %[[INPUT]], %[[C1]] : tensor<?x?x?xf32>
// CHECK: %[[D2:.+]] = tensor.dim %[[INPUT]], %[[C2]] : tensor<?x?x?xf32> // CHECK: %[[D2:.+]] = tensor.dim %[[INPUT]], %[[C2]] : tensor<?x?x?xf32>
// CHECK: tensor.insert_slice %[[INPUT]] into %[[FILL]][%[[LOW0]], %[[OFFSET1]], 2] [%[[D0]], %[[D1]], %[[D2]]] [1, 1, 1] // CHECK: tensor.insert_slice %[[INPUT]] into %[[FILL]][%[[LOW0]], %[[OFFSET1]], 2] [%[[D0]], %[[D1]], %[[D2]]] [1, 1, 1]
func.func @insert_pad_into_fill(%input: tensor<?x?x?xf32>, %low0: index, %low1: index, %high1: index, %high2: index) -> tensor<8x384x384xf32> { func.func @insert_pad_into_fill(%input: tensor<?x?x?xf32>, %low0: index, %low1: index, %high1: index, %high2: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%pad = tensor.pad %input low[%low0, %low1, %c0] high[%c0, %high1, %high2] { %pad = tensor.pad %input low[%low0, %low1, %c0] high[%c0, %high1, %high2] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<?x?x?xf32> to tensor<8x128x128xf32> } : tensor<?x?x?xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
%0 = tensor.insert_slice %pad into %fill[0, 1, 2] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %pad into %fill[0, 1, 2] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %0: tensor<8x384x384xf32> return %0: tensor<8x384x384xf32>
} }
// ----- // -----
// CHECK-LABEL: func @multi_insert_pad_into_fill // CHECK-LABEL: func @multi_insert_pad_into_fill
// CHECK-SAME: (%[[INPUT:.+]]: tensor<7x123x124xf32>, %[[A:.+]]: tensor<8x128x128xf32>, %[[OFFSET:.+]]: index) // CHECK-SAME: (%[[INPUT:.+]]: tensor<7x123x124xf32>, %[[A:.+]]: tensor<8x128x128xf32>, %[[OFFSET:.+]]: index)
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[INSERT0:.+]] = tensor.insert_slice %[[A]] into %[[FILL]][%[[OFFSET]], 0, 0] [8, 128, 128] [1, 1, 1] // CHECK: %[[INSERT0:.+]] = tensor.insert_slice %[[A]] into %[[FILL]][%[[OFFSET]], 0, 0] [8, 128, 128] [1, 1, 1]
// CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[A]] into %[[INSERT0]][0, 128, %[[OFFSET]]] [8, 128, 128] [1, 1, 1] // CHECK: %[[INSERT1:.+]] = tensor.insert_slice %[[A]] into %[[INSERT0]][0, 128, %[[OFFSET]]] [8, 128, 128] [1, 1, 1]
// CHECK: tensor.insert_slice %[[INPUT]] into %[[INSERT1]][1, 2, 256] [7, 123, 124] [1, 1, 1] // CHECK: tensor.insert_slice %[[INPUT]] into %[[INSERT1]][1, 2, 256] [7, 123, 124] [1, 1, 1]
func.func @multi_insert_pad_into_fill(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> { func.func @multi_insert_pad_into_fill(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] { %pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<7x123x124xf32> to tensor<8x128x128xf32> } : tensor<7x123x124xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
%0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%1 = tensor.insert_slice %a into %0 [0, 128, %offset][8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %1 = tensor.insert_slice %a into %0 [0, 128, %offset][8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %2: tensor<8x384x384xf32> return %2: tensor<8x384x384xf32>
} }
// ----- // -----
// CHECK-LABEL: func @multi_insert_pad_into_fill_overlap // CHECK-LABEL: func @multi_insert_pad_into_fill_overlap
func.func @multi_insert_pad_into_fill_overlap(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> { func.func @multi_insert_pad_into_fill_overlap(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
// CHECK: tensor.pad // CHECK: tensor.pad
%pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] { %pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<7x123x124xf32> to tensor<8x128x128xf32> } : tensor<7x123x124xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
%0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%1 = tensor.insert_slice %a into %0 [0, 0, 129] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %1 = tensor.insert_slice %a into %0 [0, 0, 129] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
// Range overlap with %1 at dim#3 // Range overlap with %1 at dim#3
%2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %2: tensor<8x384x384xf32> return %2: tensor<8x384x384xf32>
} }
// ----- // -----
// CHECK-LABEL: func @multi_insert_pad_into_fill_overlap // CHECK-LABEL: func @multi_insert_pad_into_fill_overlap
func.func @multi_insert_pad_into_fill_overlap(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> { func.func @multi_insert_pad_into_fill_overlap(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
// CHECK: tensor.pad // CHECK: tensor.pad
%pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] { %pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<7x123x124xf32> to tensor<8x128x128xf32> } : tensor<7x123x124xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
%0 = tensor.insert_slice %a into %fill[0, 0, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %a into %fill[0, 0, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%1 = tensor.insert_slice %a into %0 [0, 128, 255] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %1 = tensor.insert_slice %a into %0 [0, 128, 255] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
// Range overlap with %0 at dim#3 // Range overlap with %0 at dim#3
%2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %2: tensor<8x384x384xf32> return %2: tensor<8x384x384xf32>
} }
// ----- // -----
// CHECK-LABEL: func @multi_insert_pad_into_fill // CHECK-LABEL: func @multi_insert_pad_into_fill
func.func @multi_insert_pad_into_fill(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> { func.func @multi_insert_pad_into_fill(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
// CHECK-NOT: tensor.pad // CHECK-NOT: tensor.pad
%pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] { %pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<7x123x124xf32> to tensor<8x128x128xf32> } : tensor<7x123x124xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
// Overlap btween %0 and %1 is fine but not with %2 is fine. // Overlap btween %0 and %1 is fine but not with %2 is fine.
// CHECK-COUNT-3: tensor.insert_slice // CHECK-COUNT-3: tensor.insert_slice
%0 = tensor.insert_slice %a into %fill[0, 0, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %a into %fill[0, 0, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%1 = tensor.insert_slice %a into %0 [0, 1, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %1 = tensor.insert_slice %a into %0 [0, 1, %offset] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%2 = tensor.insert_slice %pad into %1 [0, 256, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %2 = tensor.insert_slice %pad into %1 [0, 256, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %2: tensor<8x384x384xf32> return %2: tensor<8x384x384xf32>
} }
// ----- // -----
// CHECK-LABEL: func @multi_insert_pad_into_fill_mismatch // CHECK-LABEL: func @multi_insert_pad_into_fill_mismatch
func.func @multi_insert_pad_into_fill_mismatch(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> { func.func @multi_insert_pad_into_fill_mismatch(%input: tensor<7x123x124xf32>, %a: tensor<8x128x128xf32>, %offset: index) -> tensor<8x384x384xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%f1 = arith.constant 1.0 : f32 %f1 = arith.constant 1.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
// CHECK: tensor.pad // CHECK: tensor.pad
%pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] { %pad = tensor.pad %input low[1, 2, 0] high[0, 3, 4] {
^bb0(%arg3: index, %arg4: index, %arg5: index): ^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %f0 : f32 tensor.yield %f0 : f32
} : tensor<7x123x124xf32> to tensor<8x128x128xf32> } : tensor<7x123x124xf32> to tensor<8x128x128xf32>
%init = linalg.init_tensor [8, 384, 384] : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
// Different filling value than padding value. // Different filling value than padding value.
%fill = linalg.fill ins(%f1 : f32) outs(%init : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f1 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32>
%0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %0 = tensor.insert_slice %a into %fill[%offset, 0, 0] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%1 = tensor.insert_slice %a into %0 [0, 128, %offset][8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %1 = tensor.insert_slice %a into %0 [0, 128, %offset][8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
%2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32> %2 = tensor.insert_slice %pad into %1 [0, 0, 256] [8, 128, 128] [1, 1, 1] : tensor<8x128x128xf32> into tensor<8x384x384xf32>
return %2: tensor<8x384x384xf32> return %2: tensor<8x384x384xf32>
} }
// ----- // -----
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines
func.func @fold_multi_use_generic_op_with_consumer(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<2x3x4xf32>) { func.func @fold_multi_use_generic_op_with_consumer(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<2x3x4xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init1 = linalg.init_tensor [%d1, %d2, %d0] : tensor<?x?x?xf32> %empty1 = tensor.empty(%d1, %d2, %d0) : tensor<?x?x?xf32>
%init2 = linalg.init_tensor [%d2, %d1, %d0] : tensor<?x?x?xf32> %empty2 = tensor.empty(%d2, %d1, %d0) : tensor<?x?x?xf32>
%0:2 = linalg.generic { %0:2 = linalg.generic {
iterator_types = ["parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel"],
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d1, d2, d0)>, affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
affine_map<(d0, d1, d2) -> (d2, d1, d0)>]} affine_map<(d0, d1, d2) -> (d2, d1, d0)>]}
ins(%arg0 : tensor<?x?x?xf32>) outs(%init1, %init2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) { ins(%arg0 : tensor<?x?x?xf32>) outs(%empty1, %empty2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32) :
linalg.yield %b0, %b0 : f32, f32 linalg.yield %b0, %b0 : f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
%1 = tensor.cast %0#1 : tensor<?x?x?xf32> to tensor<2x3x4xf32> %1 = tensor.cast %0#1 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
return %0#0, %1 : tensor<?x?x?xf32>, tensor<2x3x4xf32> return %0#0, %1 : tensor<?x?x?xf32>, tensor<2x3x4xf32>
} }
// CHECK: func @fold_multi_use_generic_op_with_consumer // CHECK: func @fold_multi_use_generic_op_with_consumer
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32> // CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<2x3x4xf32>
// CHECK-DAG: %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x3x2xf32> // CHECK-DAG: %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x3x2xf32>
// CHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [3, 2, 4] : tensor<3x2x4xf32> // CHECK-DAG: %[[INIT2:.+]] = tensor.empty() : tensor<3x2x4xf32>
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: ins(%[[CAST]] : // CHECK-SAME: ins(%[[CAST]] :
// CHECK-SAME: outs(%[[INIT2]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT2]], %[[INIT1]] :
// CHECK: %[[RETURN_CAST:.+]] = tensor.cast %[[GENERIC]]#0 : tensor<3x2x4xf32> to tensor<?x?x?xf32> // CHECK: %[[RETURN_CAST:.+]] = tensor.cast %[[GENERIC]]#0 : tensor<3x2x4xf32> to tensor<?x?x?xf32>
// CHECK: return %[[RETURN_CAST]], %[[GENERIC]]#1 // CHECK: return %[[RETURN_CAST]], %[[GENERIC]]#1

mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir

Show First 20 Lines • Show All 69 Lines • ▼ Show 20 Lines
// ----- // -----
// Spot-check an op that requires copying attributes properly to the created scalar op. // Spot-check an op that requires copying attributes properly to the created scalar op.
// Also checks proper init_tensor usage. // Also checks proper init_tensor usage.
// CHECK-LABEL: func @cmpf( // CHECK-LABEL: func @cmpf(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<f32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<f32>
func.func @cmpf(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> { func.func @cmpf(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
// CHECK: %[[INIT:.*]] = linalg.init_tensor [] : tensor<i1> // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<i1>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]
// CHECK-SAME: outs(%[[INIT]] // CHECK-SAME: outs(%[[INIT]]
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1): // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1):
// CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f32 // CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f32
%0 = arith.cmpf olt, %arg0, %arg1 : tensor<f32> %0 = arith.cmpf olt, %arg0, %arg1 : tensor<f32>
return %0 : tensor<i1> return %0 : tensor<i1>
} }
// ----- // -----
// Check proper init_tensor usage in a mixed case. // Check proper init_tensor usage in a mixed case.
// CHECK-LABEL: func @cmpf( // CHECK-LABEL: func @cmpf(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<4x?x?x8x2x?xf32>
func.func @cmpf(%arg0: tensor<4x?x?x8x2x?xf32>, %arg1: tensor<4x?x?x8x2x?xf32>) -> tensor<4x?x?x8x2x?xi1> { func.func @cmpf(%arg0: tensor<4x?x?x8x2x?xf32>, %arg1: tensor<4x?x?x8x2x?xf32>) -> tensor<4x?x?x8x2x?xi1> {
// CHECK: %[[C1:.*]] = arith.constant 1 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<4x?x?x8x2x?xf32> // CHECK: %[[D1:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<4x?x?x8x2x?xf32>
// CHECK: %[[C2:.*]] = arith.constant 2 : index // CHECK: %[[C2:.*]] = arith.constant 2 : index
// CHECK: %[[D2:.*]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<4x?x?x8x2x?xf32> // CHECK: %[[D2:.*]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<4x?x?x8x2x?xf32>
// CHECK: %[[C5:.*]] = arith.constant 5 : index // CHECK: %[[C5:.*]] = arith.constant 5 : index
// CHECK: %[[D5:.*]] = tensor.dim %[[ARG0]], %[[C5]] : tensor<4x?x?x8x2x?xf32> // CHECK: %[[D5:.*]] = tensor.dim %[[ARG0]], %[[C5]] : tensor<4x?x?x8x2x?xf32>
// CHECK: %[[INIT:.*]] = linalg.init_tensor [4, %[[D1]], %[[D2]], 8, 2, %[[D5]]] : tensor<4x?x?x8x2x?xi1> // CHECK: %[[INIT:.*]] = tensor.empty(%[[D1]], %[[D2]], %[[D5]]) : tensor<4x?x?x8x2x?xi1>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]
// CHECK-SAME: outs(%[[INIT]] // CHECK-SAME: outs(%[[INIT]]
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1): // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: i1):
// CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f32 // CHECK: arith.cmpf olt, %{{.*}}, %{{.*}} : f32
%0 = arith.cmpf olt, %arg0, %arg1 : tensor<4x?x?x8x2x?xf32> %0 = arith.cmpf olt, %arg0, %arg1 : tensor<4x?x?x8x2x?xf32>
return %0 : tensor<4x?x?x8x2x?xi1> return %0 : tensor<4x?x?x8x2x?xi1>
} }

mlir/test/Dialect/Linalg/decompose-ops.mlir

// RUN: mlir-opt -test-linalg-decompose-ops -cse -split-input-file %s | FileCheck %s // RUN: mlir-opt -test-linalg-decompose-ops -cse -split-input-file %s | FileCheck %s
// RUN: mlir-opt -test-linalg-decompose-ops -cse -canonicalize -split-input-file %s | FileCheck %s --check-prefix=CANONICALIZECHECK // RUN: mlir-opt -test-linalg-decompose-ops -cse -canonicalize -split-input-file %s | FileCheck %s --check-prefix=CANONICALIZECHECK
func.func @simple_op(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) func.func @simple_op(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>) { -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%init1 = linalg.init_tensor [%d1, %d0] : tensor<?x?xf32> %init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32>
%init2 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init2 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%result:2 = linalg.generic { %result:2 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>)
outs(%init1, %init2 : tensor<?x?xf32>, tensor<?x?xf32>) { outs(%init1, %init2 : tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
Show All 10 Lines
// CHECK: func @simple_op( // CHECK: func @simple_op(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [%[[D1]], %[[D0]]] // CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK-DAG: %[[GENERIC1:.+]]:3 = linalg.generic // CHECK-DAG: %[[GENERIC1:.+]]:3 = linalg.generic
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP0]], #[[MAP3]]] // CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP0]], #[[MAP3]]]
// CHECK-SAME: ["parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] :
// CHECK-SAME: outs(%[[INIT1]], %[[INIT2]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT1]], %[[INIT2]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32
// CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32
Show All 28 Lines
// CANONICALIZECHECK: func @simple_op( // CANONICALIZECHECK: func @simple_op(
// CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CANONICALIZECHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32> // CANONICALIZECHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32>
// CANONICALIZECHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32> // CANONICALIZECHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>
// CANONICALIZECHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CANONICALIZECHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CANONICALIZECHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CANONICALIZECHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CANONICALIZECHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CANONICALIZECHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CANONICALIZECHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]] // CANONICALIZECHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CANONICALIZECHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [%[[D1]], %[[D0]]] // CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CANONICALIZECHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CANONICALIZECHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CANONICALIZECHECK-DAG: %[[GENERIC1:.+]] = linalg.generic // CANONICALIZECHECK-DAG: %[[GENERIC1:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CANONICALIZECHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CANONICALIZECHECK-SAME: ["parallel", "parallel"] // CANONICALIZECHECK-SAME: ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CANONICALIZECHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT1]] : // CANONICALIZECHECK-SAME: outs(%[[INIT1]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32 // CANONICALIZECHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32
// CANONICALIZECHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32 // CANONICALIZECHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32
Show All 17 Lines
// ----- // -----
func.func @simple_op_permuted_outputs(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) func.func @simple_op_permuted_outputs(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) { -> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%init1 = linalg.init_tensor [%d1, %d0] : tensor<?x?xf32> %init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32>
%init2 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init2 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%result:3 = linalg.generic { %result:3 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1, d0)>,
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?xf32>, tensor<?xf32>)
outs(%init1, %init2, %init2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) { outs(%init1, %init2, %init2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32, %b5 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32, %b5 : f32) :
Show All 10 Lines
// CHECK: func @simple_op_permuted_outputs( // CHECK: func @simple_op_permuted_outputs(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [%[[D1]], %[[D0]]] // CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK-DAG: %[[GENERIC1:.+]]:4 = linalg.generic // CHECK-DAG: %[[GENERIC1:.+]]:4 = linalg.generic
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP0]], #[[MAP0]], #[[MAP3]]] // CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP0]], #[[MAP0]], #[[MAP3]]]
// CHECK-SAME: ["parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] :
// CHECK-SAME: outs(%[[INIT1]], %[[INIT2]], %[[INIT2]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT1]], %[[INIT2]], %[[INIT2]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32
// CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32
Show All 27 Lines
// CANONICALIZECHECK: func @simple_op_permuted_outputs( // CANONICALIZECHECK: func @simple_op_permuted_outputs(
// CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CANONICALIZECHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32> // CANONICALIZECHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?xf32>
// CANONICALIZECHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32> // CANONICALIZECHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>
// CANONICALIZECHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CANONICALIZECHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CANONICALIZECHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CANONICALIZECHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CANONICALIZECHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CANONICALIZECHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CANONICALIZECHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]] // CANONICALIZECHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CANONICALIZECHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [%[[D1]], %[[D0]]] // CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CANONICALIZECHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CANONICALIZECHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CANONICALIZECHECK-DAG: %[[GENERIC1:.+]]:2 = linalg.generic // CANONICALIZECHECK-DAG: %[[GENERIC1:.+]]:2 = linalg.generic
// CANONICALIZECHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]] // CANONICALIZECHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]]
// CANONICALIZECHECK-SAME: ["parallel", "parallel"] // CANONICALIZECHECK-SAME: ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CANONICALIZECHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT1]], %[[INIT2]] : // CANONICALIZECHECK-SAME: outs(%[[INIT1]], %[[INIT2]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32 // CANONICALIZECHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32
// CANONICALIZECHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32 // CANONICALIZECHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32
Show All 14 Lines
// CANONICALIZECHECK: return %[[GENERIC1]]#0, %[[GENERIC2]], %[[GENERIC1]]#1 // CANONICALIZECHECK: return %[[GENERIC1]]#0, %[[GENERIC2]], %[[GENERIC1]]#1
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0)> #map1 = affine_map<(d0, d1) -> (d0)>
#map2 = affine_map<(d0, d1) -> (d1, d0)> #map2 = affine_map<(d0, d1) -> (d1, d0)>
func.func @multi_statement(%arg0 : tensor<10x20xf32>, %arg1 : tensor<10xi32>) -> tensor<20x10xf64> { func.func @multi_statement(%arg0 : tensor<10x20xf32>, %arg1 : tensor<10xi32>) -> tensor<20x10xf64> {
%init = linalg.init_tensor [20, 10] : tensor<20x10xf64> %init = tensor.empty() : tensor<20x10xf64>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map0, #map1, #map2], indexing_maps = [#map0, #map1, #map2],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<10x20xf32>, tensor<10xi32>) ins(%arg0, %arg1 : tensor<10x20xf32>, tensor<10xi32>)
outs(%init : tensor<20x10xf64>) { outs(%init : tensor<20x10xf64>) {
^bb0(%b0 : f32, %b1 : i32, %b2 : f64): ^bb0(%b0 : f32, %b1 : i32, %b2 : f64):
%1 = arith.sitofp %b1 : i32 to f64 %1 = arith.sitofp %b1 : i32 to f64
%2 = arith.extf %b0 : f32 to f64 %2 = arith.extf %b0 : f32 to f64
%3 = arith.addf %1, %2 : f64 %3 = arith.addf %1, %2 : f64
linalg.yield %3 : f64 linalg.yield %3 : f64
} -> tensor<20x10xf64> } -> tensor<20x10xf64>
return %0 : tensor<20x10xf64> return %0 : tensor<20x10xf64>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK: func @multi_statement( // CHECK: func @multi_statement(
// CHECK-SAME: %[[ARG0:.+]]: tensor<10x20xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<10x20xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<10xi32>) // CHECK-SAME: %[[ARG1:.+]]: tensor<10xi32>)
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [20, 10] : tensor<20x10xf64> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<20x10xf64>
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [10, 20] : tensor<10x20xf64> // CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<10x20xf64>
// CHECK: %[[GENERIC0:.+]]:2 = linalg.generic // CHECK: %[[GENERIC0:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B0:.+]]: f32 // CHECK-SAME: %[[B0:.+]]: f32
// CHECK-SAME: %[[B1:.+]]: i32 // CHECK-SAME: %[[B1:.+]]: i32
Show All 30 Lines
// CHECK: return %[[GENERIC2]] // CHECK: return %[[GENERIC2]]
// CANONICALIZECHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)> // CANONICALIZECHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)>
// CANONICALIZECHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CANONICALIZECHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CANONICALIZECHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1, d0)> // CANONICALIZECHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CANONICALIZECHECK: func @multi_statement( // CANONICALIZECHECK: func @multi_statement(
// CANONICALIZECHECK-SAME: %[[ARG0:.+]]: tensor<10x20xf32> // CANONICALIZECHECK-SAME: %[[ARG0:.+]]: tensor<10x20xf32>
// CANONICALIZECHECK-SAME: %[[ARG1:.+]]: tensor<10xi32>) // CANONICALIZECHECK-SAME: %[[ARG1:.+]]: tensor<10xi32>)
// CANONICALIZECHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [20, 10] : tensor<20x10xf64> // CANONICALIZECHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<20x10xf64>
// CANONICALIZECHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [10, 20] : tensor<10x20xf64> // CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<10x20xf64>
// CANONICALIZECHECK: %[[GENERIC0:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC0:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] // CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
// CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"] // CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG1]] : // CANONICALIZECHECK-SAME: ins(%[[ARG1]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT1]] : // CANONICALIZECHECK-SAME: outs(%[[INIT1]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B0:.+]]: i32 // CANONICALIZECHECK-SAME: %[[B0:.+]]: i32
// CANONICALIZECHECK-SAME: %[[B1:.+]]: f64 // CANONICALIZECHECK-SAME: %[[B1:.+]]: f64
▲ Show 20 Lines • Show All 116 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/detensorize_0d.mlir

// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s // RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s
#map = affine_map<() -> ()> #map = affine_map<() -> ()>
func.func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { func.func @detensor_simple(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { outs(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<f32> } -> tensor<f32>
return %1: tensor<f32> return %1: tensor<f32>
} }
// CHECK-LABEL: func @detensor_simple // CHECK-LABEL: func @detensor_simple
// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) // CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] // CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]]
// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] // CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]]
// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] // CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]] // CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]]
// CHECK: return %[[new_tensor_res]] // CHECK: return %[[new_tensor_res]]
func.func @detensor_op_sequence(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { func.func @detensor_op_sequence(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { outs(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<f32> } -> tensor<f32>
%3 = linalg.init_tensor [] : tensor<f32> %3 = tensor.empty() : tensor<f32>
%4 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %4 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%arg1, %1 : tensor<f32>, tensor<f32>) ins(%arg1, %1 : tensor<f32>, tensor<f32>)
outs(%3 : tensor<f32>) { outs(%3 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%5 = arith.mulf %arg3, %arg4 : f32 %5 = arith.mulf %arg3, %arg4 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<f32> } -> tensor<f32>
%6 = linalg.init_tensor [] : tensor<f32> %6 = tensor.empty() : tensor<f32>
%7 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %7 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%1, %4 : tensor<f32>, tensor<f32>) ins(%1, %4 : tensor<f32>, tensor<f32>)
outs(%6 : tensor<f32>) { outs(%6 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%5 = arith.divf %arg3, %arg4 : f32 %5 = arith.divf %arg3, %arg4 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<f32> } -> tensor<f32>
return %7: tensor<f32> return %7: tensor<f32>
} }
// CHECK-LABEL: func @detensor_op_sequence // CHECK-LABEL: func @detensor_op_sequence
// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) // CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] // CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]]
// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] // CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]]
// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] // CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
// CHECK: %[[detensored_res2:.*]] = arith.mulf %[[arg1_val]], %[[detensored_res]] // CHECK: %[[detensored_res2:.*]] = arith.mulf %[[arg1_val]], %[[detensored_res]]
// CHECK: %[[detensored_res3:.*]] = arith.divf %[[detensored_res]], %[[detensored_res2]] // CHECK: %[[detensored_res3:.*]] = arith.divf %[[detensored_res]], %[[detensored_res2]]
// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res3]] // CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res3]]
// CHECK: return %[[new_tensor_res]] // CHECK: return %[[new_tensor_res]]
func.func @detensor_multiple_ops(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { func.func @detensor_multiple_ops(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { outs(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
%3 = arith.mulf %2, %arg4 : f32 %3 = arith.mulf %2, %arg4 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<f32> } -> tensor<f32>
return %1: tensor<f32> return %1: tensor<f32>
} }
// CHECK-LABEL: func @detensor_multiple_ops // CHECK-LABEL: func @detensor_multiple_ops
// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) // CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] // CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]]
// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] // CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]]
// CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]] // CHECK: %[[detensored_res:.*]] = arith.addf %[[arg1_val]], %[[arg2_val]]
// CHECK: %[[detensored_res2:.*]] = arith.mulf %[[detensored_res]], %[[arg2_val]] // CHECK: %[[detensored_res2:.*]] = arith.mulf %[[detensored_res]], %[[arg2_val]]
// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res2]] // CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res2]]
// CHECK: return %[[new_tensor_res]] // CHECK: return %[[new_tensor_res]]
func.func @detensor_foreign_op(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} { func.func @detensor_foreign_op(%arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> attributes {iree.module.export} {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []} %1 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%arg1, %arg2 : tensor<f32>, tensor<f32>) ins(%arg1, %arg2 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { outs(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = "foreign.do_something"(%arg3, %arg4) {} : (f32, f32) -> f32 %2 = "foreign.do_something"(%arg3, %arg4) {} : (f32, f32) -> f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<f32> } -> tensor<f32>
return %1: tensor<f32> return %1: tensor<f32>
} }
// CHECK-LABEL: func @detensor_foreign_op // CHECK-LABEL: func @detensor_foreign_op
// CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>) // CHECK-SAME: (%[[arg1:.*]]: tensor<f32>, %[[arg2:.*]]: tensor<f32>)
// CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]] // CHECK-DAG: %[[arg1_val:.*]] = tensor.extract %[[arg1]]
// CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]] // CHECK-DAG: %[[arg2_val:.*]] = tensor.extract %[[arg2]]
// CHECK: %[[detensored_res:.*]] = "foreign.do_something"(%[[arg1_val]], %[[arg2_val]]) // CHECK: %[[detensored_res:.*]] = "foreign.do_something"(%[[arg1_val]], %[[arg2_val]])
// CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]] // CHECK: %[[new_tensor_res:.*]] = tensor.from_elements %[[detensored_res]]
// CHECK: return %[[new_tensor_res]] // CHECK: return %[[new_tensor_res]]

mlir/test/Dialect/Linalg/detensorize_br_operands.mlir

// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s // RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s
// TODO: Detensoring breaks if %arg0 or %arg1 are passed directly as tensors. Fix that. // TODO: Detensoring breaks if %arg0 or %arg1 are passed directly as tensors. Fix that.
func.func @if_true_test(%arg0: i1, %arg1: i32) -> tensor<i32> attributes {} { func.func @if_true_test(%arg0: i1, %arg1: i32) -> tensor<i32> attributes {} {
%arg0_t = tensor.from_elements %arg0 : tensor<i1> %arg0_t = tensor.from_elements %arg0 : tensor<i1>
%arg1_t = tensor.from_elements %arg1 : tensor<i32> %arg1_t = tensor.from_elements %arg1 : tensor<i32>
%cst = arith.constant dense<10> : tensor<i32> %cst = arith.constant dense<10> : tensor<i32>
%2 = linalg.init_tensor [] : tensor<i8> %2 = tensor.empty() : tensor<i8>
%3 = linalg.generic %3 = linalg.generic
{indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
ins(%arg0_t : tensor<i1>) ins(%arg0_t : tensor<i1>)
outs(%2 : tensor<i8>) { outs(%2 : tensor<i8>) {
^bb0(%arg2: i1, %arg3: i8): ^bb0(%arg2: i1, %arg3: i8):
%10 = arith.extui %arg2 : i1 to i8 %10 = arith.extui %arg2 : i1 to i8
linalg.yield %10 : i8 linalg.yield %10 : i8
} -> tensor<i8> } -> tensor<i8>
%4 = tensor.extract %3[] : tensor<i8> %4 = tensor.extract %3[] : tensor<i8>
%5 = arith.trunci %4 : i8 to i1 %5 = arith.trunci %4 : i8 to i1
cf.cond_br %5, ^bb1, ^bb2(%arg1_t : tensor<i32>) cf.cond_br %5, ^bb1, ^bb2(%arg1_t : tensor<i32>)
^bb1: ^bb1:
%6 = linalg.init_tensor [] : tensor<i32> %6 = tensor.empty() : tensor<i32>
%7 = linalg.generic %7 = linalg.generic
{indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []} {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>, affine_map<() -> ()>], iterator_types = []}
ins(%arg1_t, %cst : tensor<i32>, tensor<i32>) ins(%arg1_t, %cst : tensor<i32>, tensor<i32>)
outs(%6 : tensor<i32>) { outs(%6 : tensor<i32>) {
^bb0(%arg2: i32, %arg3: i32, %arg4: i32): ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):
%10 = arith.addi %arg2, %arg3 : i32 %10 = arith.addi %arg2, %arg3 : i32
linalg.yield %10 : i32 linalg.yield %10 : i32
} -> tensor<i32> } -> tensor<i32>
Show All 15 Lines

mlir/test/Dialect/Linalg/detensorize_if.mlir

Show All 9 Lines
func.func @main() -> (tensor<i32>) attributes {} { func.func @main() -> (tensor<i32>) attributes {} {
%c0 = arith.constant 0 : i32 %c0 = arith.constant 0 : i32
%0 = tensor.from_elements %c0 : tensor<i32> %0 = tensor.from_elements %c0 : tensor<i32>
%c10 = arith.constant 10 : i32 %c10 = arith.constant 10 : i32
%1 = tensor.from_elements %c10 : tensor<i32> %1 = tensor.from_elements %c10 : tensor<i32>
cf.br ^bb1(%0 : tensor<i32>) cf.br ^bb1(%0 : tensor<i32>)
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%2, %1 : tensor<i32>, tensor<i32>) ins(%2, %1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1> %5 = tensor.extract %4[] : tensor<i1>
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>) cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
^bb2(%6: tensor<i32>): // pred: ^bb1 ^bb2(%6: tensor<i32>): // pred: ^bb1
%7 = linalg.init_tensor [] : tensor<i32> %7 = tensor.empty() : tensor<i32>
%8 = linalg.generic #attrs %8 = linalg.generic #attrs
ins(%6, %6 : tensor<i32>, tensor<i32>) ins(%6, %6 : tensor<i32>, tensor<i32>)
outs(%7 : tensor<i32>) { outs(%7 : tensor<i32>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i32): ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
%9 = arith.addi %arg0, %arg1 : i32 %9 = arith.addi %arg0, %arg1 : i32
linalg.yield %9 : i32 linalg.yield %9 : i32
} -> tensor<i32> } -> tensor<i32>
cf.br ^bb3(%8 : tensor<i32>) cf.br ^bb3(%8 : tensor<i32>)
Show All 32 Lines
func.func @main() -> (tensor<i32>) attributes {} { func.func @main() -> (tensor<i32>) attributes {} {
%c0 = arith.constant 0 : i32 %c0 = arith.constant 0 : i32
%0 = tensor.from_elements %c0 : tensor<i32> %0 = tensor.from_elements %c0 : tensor<i32>
%c10 = arith.constant 10 : i32 %c10 = arith.constant 10 : i32
%1 = tensor.from_elements %c10 : tensor<i32> %1 = tensor.from_elements %c10 : tensor<i32>
cf.br ^bb1(%0 : tensor<i32>) cf.br ^bb1(%0 : tensor<i32>)
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%2, %1 : tensor<i32>, tensor<i32>) ins(%2, %1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1> %5 = tensor.extract %4[] : tensor<i1>
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>) cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
^bb2(%6: tensor<i32>): // pred: ^bb1 ^bb2(%6: tensor<i32>): // pred: ^bb1
%7 = linalg.init_tensor [] : tensor<i32> %7 = tensor.empty() : tensor<i32>
%8 = linalg.generic #attrs %8 = linalg.generic #attrs
ins(%6, %6 : tensor<i32>, tensor<i32>) ins(%6, %6 : tensor<i32>, tensor<i32>)
outs(%7 : tensor<i32>) { outs(%7 : tensor<i32>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i32): ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
%9 = arith.addi %arg0, %arg1 : i32 %9 = arith.addi %arg0, %arg1 : i32
linalg.yield %9 : i32 linalg.yield %9 : i32
} -> tensor<i32> } -> tensor<i32>
cf.br ^bb3(%8 : tensor<i32>) cf.br ^bb3(%8 : tensor<i32>)
Show All 34 Lines
func.func @main() -> (tensor<i32>) attributes {} { func.func @main() -> (tensor<i32>) attributes {} {
%c0 = arith.constant 0 : i32 %c0 = arith.constant 0 : i32
%0 = tensor.from_elements %c0 : tensor<i32> %0 = tensor.from_elements %c0 : tensor<i32>
%c10 = arith.constant 10 : i32 %c10 = arith.constant 10 : i32
%1 = tensor.from_elements %c10 : tensor<i32> %1 = tensor.from_elements %c10 : tensor<i32>
cf.br ^bb1(%0 : tensor<i32>) cf.br ^bb1(%0 : tensor<i32>)
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%2, %1 : tensor<i32>, tensor<i32>) ins(%2, %1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1> %5 = tensor.extract %4[] : tensor<i1>
// This cf.cond_br intentionally has bb2 as it's target for both branches. This // This cf.cond_br intentionally has bb2 as it's target for both branches. This
// is to make sure that the "forward phase" of the cost-model correctly adds // is to make sure that the "forward phase" of the cost-model correctly adds
// the users of a block argument (in this case bb2's argument) to the work // the users of a block argument (in this case bb2's argument) to the work
// list. // list.
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>) cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb2(%2 : tensor<i32>)
^bb2(%6: tensor<i32>): // pred: ^bb1 ^bb2(%6: tensor<i32>): // pred: ^bb1
%12 = tensor.from_elements %c10 : tensor<i32> %12 = tensor.from_elements %c10 : tensor<i32>
%7 = linalg.init_tensor [] : tensor<i32> %7 = tensor.empty() : tensor<i32>
%8 = linalg.generic #attrs %8 = linalg.generic #attrs
ins(%6, %12 : tensor<i32>, tensor<i32>) ins(%6, %12 : tensor<i32>, tensor<i32>)
outs(%7 : tensor<i32>) { outs(%7 : tensor<i32>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i32): ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
%9 = arith.addi %arg0, %arg1 : i32 %9 = arith.addi %arg0, %arg1 : i32
linalg.yield %9 : i32 linalg.yield %9 : i32
} -> tensor<i32> } -> tensor<i32>
cf.br ^bb3(%8 : tensor<i32>) cf.br ^bb3(%8 : tensor<i32>)
Show All 19 Lines

mlir/test/Dialect/Linalg/detensorize_trivial.mlir

// RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL // RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF // RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()> #map0 = affine_map<() -> ()>
#attrs = { #attrs = {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = [] iterator_types = []
} }
func.func @main(%farg0 : tensor<i32>) -> (tensor<i1>) attributes {} { func.func @main(%farg0 : tensor<i32>) -> (tensor<i1>) attributes {} {
%c10 = arith.constant 10 : i32 %c10 = arith.constant 10 : i32
%1 = tensor.from_elements %c10 : tensor<i32> %1 = tensor.from_elements %c10 : tensor<i32>
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%farg0, %1 : tensor<i32>, tensor<i32>) ins(%farg0, %1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
return %4 : tensor<i1> return %4 : tensor<i1>
} }
// DET-ALL-LABEL: func @main(%{{.*}}: tensor<i32>) // DET-ALL-LABEL: func @main(%{{.*}}: tensor<i32>)
// DET-ALL-NEXT: arith.constant 10 // DET-ALL-NEXT: arith.constant 10
// DET-ALL-NEXT: tensor.extract %{{.*}}[] // DET-ALL-NEXT: tensor.extract %{{.*}}[]
// DET-ALL-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}} // DET-ALL-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}}
// DET-ALL-NEXT: tensor.from_elements %{{.*}} // DET-ALL-NEXT: tensor.from_elements %{{.*}}
// DET-ALL-NEXT: return %{{.*}} : tensor<i1> // DET-ALL-NEXT: return %{{.*}} : tensor<i1>
// DET-ALL-NEXT: } // DET-ALL-NEXT: }
// DET-CF-LABEL: func @main(%{{.*}}: tensor<i32>) // DET-CF-LABEL: func @main(%{{.*}}: tensor<i32>)
// DET-CF-NEXT: arith.constant dense<10> : tensor<i32> // DET-CF-NEXT: arith.constant dense<10> : tensor<i32>
// DET-CF-NEXT: linalg.init_tensor [] : tensor<i1> // DET-CF-NEXT: tensor.empty() : tensor<i1>
// DET-CF-NEXT: linalg.generic // DET-CF-NEXT: linalg.generic
// DET-CF-NEXT: ^{{.*}}(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i1) // DET-CF-NEXT: ^{{.*}}(%{{.*}}: i32, %{{.*}}: i32, %{{.*}}: i1)
// DET-CF-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}} // DET-CF-NEXT: arith.cmpi slt, %{{.*}}, %{{.*}}
// DET-CF-NEXT: linalg.yield %{{.*}} // DET-CF-NEXT: linalg.yield %{{.*}}
// DET-CF-NEXT: } -> tensor<i1> // DET-CF-NEXT: } -> tensor<i1>
// DET-CF-NEXT: return %{{.*}} // DET-CF-NEXT: return %{{.*}}
// DET-CF-NEXT: } // DET-CF-NEXT: }

mlir/test/Dialect/Linalg/detensorize_while.mlir

// RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL // RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF // RUN: mlir-opt %s -pass-pipeline="func.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()> #map0 = affine_map<() -> ()>
#attrs = { #attrs = {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = [] iterator_types = []
} }
func.func @main(%farg0: tensor<i32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} { func.func @main(%farg0: tensor<i32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} {
cf.br ^bb1(%farg0 : tensor<i32>) cf.br ^bb1(%farg0 : tensor<i32>)
^bb1(%0: tensor<i32>): // 2 preds: ^bb0, ^bb2 ^bb1(%0: tensor<i32>): // 2 preds: ^bb0, ^bb2
%1 = linalg.init_tensor [] : tensor<i1> %1 = tensor.empty() : tensor<i1>
%2 = linalg.generic #attrs %2 = linalg.generic #attrs
ins(%0, %farg1 : tensor<i32>, tensor<i32>) ins(%0, %farg1 : tensor<i32>, tensor<i32>)
outs(%1 : tensor<i1>) { outs(%1 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%3 = tensor.extract %2[] : tensor<i1> %3 = tensor.extract %2[] : tensor<i1>
cf.cond_br %3, ^bb2(%0 : tensor<i32>), ^bb3(%0 : tensor<i32>) cf.cond_br %3, ^bb2(%0 : tensor<i32>), ^bb3(%0 : tensor<i32>)
^bb2(%4: tensor<i32>): // pred: ^bb1 ^bb2(%4: tensor<i32>): // pred: ^bb1
%5 = linalg.init_tensor [] : tensor<i32> %5 = tensor.empty() : tensor<i32>
%6 = linalg.generic #attrs %6 = linalg.generic #attrs
ins(%4, %4 : tensor<i32>, tensor<i32>) ins(%4, %4 : tensor<i32>, tensor<i32>)
outs(%5 : tensor<i32>) { outs(%5 : tensor<i32>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i32): ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
%8 = arith.addi %arg0, %arg1 : i32 %8 = arith.addi %arg0, %arg1 : i32
linalg.yield %8 : i32 linalg.yield %8 : i32
} -> tensor<i32> } -> tensor<i32>
cf.br ^bb1(%6 : tensor<i32>) cf.br ^bb1(%6 : tensor<i32>)
Show All 36 Lines

mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir

Show All 19 Lines#broadcast_attrs = {
indexing_maps = [#map1, #map2], indexing_maps = [#map1, #map2],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
func.func @main(%farg0: tensor<10xi32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} { func.func @main(%farg0: tensor<10xi32>, %farg1: tensor<i32>) -> tensor<i32> attributes {} {
cf.br ^bb1(%farg0 : tensor<10xi32>) cf.br ^bb1(%farg0 : tensor<10xi32>)
^bb1(%0: tensor<10xi32>): // 2 preds: ^bb0, ^bb2 ^bb1(%0: tensor<10xi32>): // 2 preds: ^bb0, ^bb2
%1 = linalg.init_tensor [] : tensor<i32> %1 = tensor.empty() : tensor<i32>
%2 = linalg.generic #sum_reduction_attrs %2 = linalg.generic #sum_reduction_attrs
ins(%0: tensor<10xi32>) ins(%0: tensor<10xi32>)
outs(%1: tensor<i32>) { outs(%1: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%b = arith.addi %x, %a : i32 %b = arith.addi %x, %a : i32
linalg.yield %b : i32 linalg.yield %b : i32
} -> tensor<i32> } -> tensor<i32>
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%2, %farg1 : tensor<i32>, tensor<i32>) ins(%2, %farg1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1> %5 = tensor.extract %4[] : tensor<i1>
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>) cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3(%2 : tensor<i32>)
^bb2(%6: tensor<i32>): // pred: ^bb1 ^bb2(%6: tensor<i32>): // pred: ^bb1
%7 = linalg.init_tensor [10] : tensor<10xi32> %7 = tensor.empty() : tensor<10xi32>
%9 = linalg.generic #broadcast_attrs %9 = linalg.generic #broadcast_attrs
ins(%6: tensor<i32>) ins(%6: tensor<i32>)
outs(%7: tensor<10xi32>) { outs(%7: tensor<10xi32>) {
^bb(%a: i32, %b: i32) : ^bb(%a: i32, %b: i32) :
linalg.yield %a : i32 linalg.yield %a : i32
} -> tensor<10xi32> } -> tensor<10xi32>
cf.br ^bb1(%9 : tensor<10xi32>) cf.br ^bb1(%9 : tensor<10xi32>)
^bb3(%10: tensor<i32>): // pred: ^bb1 ^bb3(%10: tensor<i32>): // pred: ^bb1
return %10 : tensor<i32> return %10 : tensor<i32>
} }
// Test aggresively detensoring all detensorable ops. // Test aggresively detensoring all detensorable ops.
// //
// DET-ALL-LABEL: func @main // DET-ALL-LABEL: func @main
// DET-ALL-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>) // DET-ALL-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>)
// DET-ALL: cf.br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>) // DET-ALL: cf.br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>)
// DET-ALL: ^[[bb1]](%{{.*}}: tensor<10xi32>) // DET-ALL: ^[[bb1]](%{{.*}}: tensor<10xi32>)
// DET-ALL: linalg.init_tensor [] : tensor<i32> // DET-ALL: tensor.empty() : tensor<i32>
// DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) { // DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) {
// DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32): // DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32):
// DET-ALL: %{{.*}} = arith.addi %{{.*}}, %{{.*}} // DET-ALL: %{{.*}} = arith.addi %{{.*}}, %{{.*}}
// DET-ALL: linalg.yield %{{.*}} : i32 // DET-ALL: linalg.yield %{{.*}} : i32
// DET-ALL: } -> tensor<i32> // DET-ALL: } -> tensor<i32>
// DET-ALL: tensor.extract %{{.*}}[] : tensor<i32> // DET-ALL: tensor.extract %{{.*}}[] : tensor<i32>
// DET-ALL: cmpi slt, %{{.*}}, %{{.*}} : i32 // DET-ALL: cmpi slt, %{{.*}}, %{{.*}} : i32
// DET-ALL: cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32) // DET-ALL: cf.cond_br %{{.*}}, ^[[bb2:.*]](%{{.*}} : i32), ^[[bb3:.*]](%{{.*}} : i32)
// DET-ALL: ^[[bb2]](%{{.*}}: i32) // DET-ALL: ^[[bb2]](%{{.*}}: i32)
// DET-ALL: tensor.from_elements %{{.*}} : tensor<i32> // DET-ALL: tensor.from_elements %{{.*}} : tensor<i32>
// DET-ALL: linalg.init_tensor [10] : tensor<10xi32> // DET-ALL: tensor.empty() : tensor<10xi32>
// DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) { // DET-ALL: linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) {
// DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32): // DET-ALL: ^bb0(%{{.*}}: i32, %{{.*}}: i32):
// DET-ALL: linalg.yield %{{.*}} : i32 // DET-ALL: linalg.yield %{{.*}} : i32
// DET-ALL: } -> tensor<10xi32> // DET-ALL: } -> tensor<10xi32>
// DET-ALL: cf.br ^[[bb1]](%{{.*}} : tensor<10xi32>) // DET-ALL: cf.br ^[[bb1]](%{{.*}} : tensor<10xi32>)
// DET-ALL: ^[[bb3]](%{{.*}}: i32) // DET-ALL: ^[[bb3]](%{{.*}}: i32)
// DET-ALL: tensor.from_elements %{{.*}} : tensor<i32> // DET-ALL: tensor.from_elements %{{.*}} : tensor<i32>
// DET-ALL: return %{{.*}} : tensor<i32> // DET-ALL: return %{{.*}} : tensor<i32>
Show All 16 Lines

mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir

Show All 11 Linesfunc.func @main() -> () attributes {} {
%0 = tensor.from_elements %c0 : tensor<1xi32> %0 = tensor.from_elements %c0 : tensor<1xi32>
%reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32> %reshaped0 = tensor.collapse_shape %0 [] : tensor<1xi32> into tensor<i32>
%c10 = arith.constant 10 : i32 %c10 = arith.constant 10 : i32
%1 = tensor.from_elements %c10 : tensor<1xi32> %1 = tensor.from_elements %c10 : tensor<1xi32>
%reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32> %reshaped1 = tensor.collapse_shape %1 [] : tensor<1xi32> into tensor<i32>
cf.br ^bb1(%reshaped0 : tensor<i32>) cf.br ^bb1(%reshaped0 : tensor<i32>)
^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2 ^bb1(%2: tensor<i32>): // 2 preds: ^bb0, ^bb2
%3 = linalg.init_tensor [] : tensor<i1> %3 = tensor.empty() : tensor<i1>
%4 = linalg.generic #attrs %4 = linalg.generic #attrs
ins(%2, %reshaped1 : tensor<i32>, tensor<i32>) ins(%2, %reshaped1 : tensor<i32>, tensor<i32>)
outs(%3 : tensor<i1>) { outs(%3 : tensor<i1>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i1): ^bb0(%arg0: i32, %arg1: i32, %arg2: i1):
%8 = arith.cmpi slt, %arg0, %arg1 : i32 %8 = arith.cmpi slt, %arg0, %arg1 : i32
linalg.yield %8 : i1 linalg.yield %8 : i1
} -> tensor<i1> } -> tensor<i1>
%5 = tensor.extract %4[] : tensor<i1> %5 = tensor.extract %4[] : tensor<i1>
cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3 cf.cond_br %5, ^bb2(%2 : tensor<i32>), ^bb3
^bb2(%6: tensor<i32>): // pred: ^bb1 ^bb2(%6: tensor<i32>): // pred: ^bb1
%7 = linalg.init_tensor [] : tensor<i32> %7 = tensor.empty() : tensor<i32>
%8 = linalg.generic #attrs %8 = linalg.generic #attrs
ins(%6, %6 : tensor<i32>, tensor<i32>) ins(%6, %6 : tensor<i32>, tensor<i32>)
outs(%7 : tensor<i32>) { outs(%7 : tensor<i32>) {
^bb0(%arg0: i32, %arg1: i32, %arg2: i32): ^bb0(%arg0: i32, %arg1: i32, %arg2: i32):
%9 = arith.addi %arg0, %arg1 : i32 %9 = arith.addi %arg0, %arg1 : i32
linalg.yield %9 : i32 linalg.yield %9 : i32
} -> tensor<i32> } -> tensor<i32>
cf.br ^bb1(%8 : tensor<i32>) cf.br ^bb1(%8 : tensor<i32>)
Show All 18 Lines

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

Show First 20 Lines • Show All 240 Lines • ▼ Show 20 Lines
// CHECK-SAME: %[[A]] // CHECK-SAME: %[[A]]
// ----- // -----
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2)> #map1 = affine_map<(d0, d1, d2) -> (d2)>
func.func @fold_unit_dim_tensor_reshape_op(%arg0 : tensor<5xf32>) -> tensor<2x5xf32> func.func @fold_unit_dim_tensor_reshape_op(%arg0 : tensor<5xf32>) -> tensor<2x5xf32>
{ {
%1 = linalg.init_tensor [1, 2, 5] : tensor<1x2x5xf32> %1 = tensor.empty() : tensor<1x2x5xf32>
%2 = linalg.generic {i64, indexing_maps = [#map1, #map0], %2 = linalg.generic {i64, indexing_maps = [#map1, #map0],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<5xf32>) outs(%1 : tensor<1x2x5xf32>) { ins(%arg0 : tensor<5xf32>) outs(%1 : tensor<1x2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<1x2x5xf32> } -> tensor<1x2x5xf32>
%3 = tensor.collapse_shape %2 [[0, 1], [2]] %3 = tensor.collapse_shape %2 [[0, 1], [2]]
: tensor<1x2x5xf32> into tensor<2x5xf32> : tensor<1x2x5xf32> into tensor<2x5xf32>
return %3 : tensor<2x5xf32> return %3 : tensor<2x5xf32>
} }
// CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op // CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
func.func @fold_unit_dim_for_init_tensor(%input: tensor<1x1000xf32>) -> tensor<1xf32> { func.func @fold_unit_dim_for_empty_tensor(%input: tensor<1x1000xf32>) -> tensor<1xf32> {
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%init = linalg.init_tensor [1] : tensor<1xf32> %init = tensor.empty() : tensor<1xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1xf32>) -> tensor<1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1xf32>) -> tensor<1xf32>
%add = linalg.generic { %add = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%input : tensor<1x1000xf32>)outs(%fill : tensor<1xf32>) { ins(%input : tensor<1x1000xf32>)outs(%fill : tensor<1xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%1823 = arith.addf %arg1, %arg2 : f32 %1823 = arith.addf %arg1, %arg2 : f32
linalg.yield %1823 : f32 linalg.yield %1823 : f32
} -> tensor<1xf32> } -> tensor<1xf32>
return %add : tensor<1xf32> return %add : tensor<1xf32>
} }
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>
// CHECK: func @fold_unit_dim_for_init_tensor // CHECK: func @fold_unit_dim_for_empty_tensor
// CHECK: %[[INPUT_RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x1000xf32> into tensor<1000xf32> // CHECK: %[[INPUT_RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x1000xf32> into tensor<1000xf32>
// CHECK: %[[INIT:.+]] = linalg.init_tensor [] : tensor<f32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<f32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<f32>) -> tensor<f32> // CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<f32>) -> tensor<f32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["reduction"] // CHECK-SAME: iterator_types = ["reduction"]
// CHECK-SAME: ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>) // CHECK-SAME: ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>)
// CHECK-SAME: outs(%[[FILL]] : tensor<f32>) // CHECK-SAME: outs(%[[FILL]] : tensor<f32>)
// CHECK: %[[GENERIC_RESHAPE:.+]] = tensor.expand_shape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32> // CHECK: %[[GENERIC_RESHAPE:.+]] = tensor.expand_shape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32>
// CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32> // CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32>
Show All 27 Lines
// CHECK: return %[[RESULT1]], %[[RESULT2]] // CHECK: return %[[RESULT1]], %[[RESULT2]]
// ----- // -----
func.func @unit_dim_for_reduction(%arg0: tensor<1x?x1x?xf32>) -> tensor<1x?xf32> { func.func @unit_dim_for_reduction(%arg0: tensor<1x?x1x?xf32>) -> tensor<1x?xf32> {
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%0 = tensor.dim %arg0, %c3 : tensor<1x?x1x?xf32> %0 = tensor.dim %arg0, %c3 : tensor<1x?x1x?xf32>
%1 = linalg.init_tensor [1, %0] : tensor<1x?xf32> %1 = tensor.empty(%0) : tensor<1x?xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x?xf32>) -> tensor<1x?xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x?xf32>) -> tensor<1x?xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1)>], affine_map<(d0, d1, d2, d3) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction", "reduction"]} iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
ins(%arg0 : tensor<1x?x1x?xf32>) ins(%arg0 : tensor<1x?x1x?xf32>)
outs(%2 : tensor<1x?xf32>) { outs(%2 : tensor<1x?xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%4 = arith.addf %arg1, %arg2 : f32 %4 = arith.addf %arg1, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<1x?xf32> } -> tensor<1x?xf32>
return %3 : tensor<1x?xf32> return %3 : tensor<1x?xf32>
} }
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
// CHECK: func @unit_dim_for_reduction // CHECK: func @unit_dim_for_reduction
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x?xf32>
// CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]] // CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%{{.+}}] : tensor<?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "reduction"]
// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>) // CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>)
// CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>) // CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>)
// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] // CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
// ----- // -----
func.func @unit_dim_for_both_reduction(%arg0: tensor<1x?x1x1xf32>) -> tensor<1x1xf32> { func.func @unit_dim_for_both_reduction(%arg0: tensor<1x?x1x1xf32>) -> tensor<1x1xf32> {
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%1 = linalg.init_tensor [1, 1] : tensor<1x1xf32> %1 = tensor.empty() : tensor<1x1xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1)>], affine_map<(d0, d1, d2, d3) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction", "reduction"]} iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
ins(%arg0 : tensor<1x?x1x1xf32>) ins(%arg0 : tensor<1x?x1x1xf32>)
outs(%2 : tensor<1x1xf32>) { outs(%2 : tensor<1x1xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%4 = arith.addf %arg1, %arg2 : f32 %4 = arith.addf %arg1, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<1x1xf32> } -> tensor<1x1xf32>
return %3 : tensor<1x1xf32> return %3 : tensor<1x1xf32>
} }
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0)>
// CHECK: func @unit_dim_for_both_reduction // CHECK: func @unit_dim_for_both_reduction
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x1xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x1xf32>
// CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3] // CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1] : tensor<1xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel"] // CHECK-SAME: iterator_types = ["parallel"]
// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?xf32>) // CHECK-SAME: ins(%[[RESHAPE]] : tensor<?xf32>)
// CHECK-SAME: outs(%[[FILL]] : tensor<1xf32>) // CHECK-SAME: outs(%[[FILL]] : tensor<1xf32>)
// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] // CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
// ----- // -----
func.func @unit_dim_for_reduction_inner(%arg0: tensor<?x1x?x1xf32>) -> tensor<?x1xf32> { func.func @unit_dim_for_reduction_inner(%arg0: tensor<?x1x?x1xf32>) -> tensor<?x1xf32> {
%cst = arith.constant 1.000000e+00 : f32 %cst = arith.constant 1.000000e+00 : f32
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c2 : tensor<?x1x?x1xf32> %0 = tensor.dim %arg0, %c2 : tensor<?x1x?x1xf32>
%1 = linalg.init_tensor [%0, 1] : tensor<?x1xf32> %1 = tensor.empty(%0) : tensor<?x1xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?x1xf32>) -> tensor<?x1xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?x1xf32>) -> tensor<?x1xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1)>], affine_map<(d0, d1, d2, d3) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction", "reduction"]} iterator_types = ["parallel", "parallel", "reduction", "reduction"]}
ins(%arg0 : tensor<?x1x?x1xf32>) ins(%arg0 : tensor<?x1x?x1xf32>)
outs(%2 : tensor<?x1xf32>) { outs(%2 : tensor<?x1xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%4 = arith.addf %arg1, %arg2 : f32 %4 = arith.addf %arg1, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x1xf32> } -> tensor<?x1xf32>
return %3 : tensor<?x1xf32> return %3 : tensor<?x1xf32>
} }
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0)>
// CHECK: func @unit_dim_for_reduction_inner // CHECK: func @unit_dim_for_reduction_inner
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x1xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x1xf32>
// CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]] // CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%{{.+}}] : tensor<?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "reduction"]
// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>) // CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>)
// CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>) // CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>)
// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] // CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]]
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
▲ Show 20 Lines • Show All 375 Lines • ▼ Show 20 Lines indexing_maps = [
affine_map<(i,j) -> (i)> // x (out) affine_map<(i,j) -> (i)> // x (out)
], ],
iterator_types = ["parallel", "reduction"] iterator_types = ["parallel", "reduction"]
} }
#CSR = #sparse_tensor.encoding<{ dimLevelType = ["dense", "compressed"] }> #CSR = #sparse_tensor.encoding<{ dimLevelType = ["dense", "compressed"] }>
func.func @sparse_case(%arg0: tensor<8x8xf32, #CSR>, %arg1: tensor<8xf32>) -> tensor<8xf32> { func.func @sparse_case(%arg0: tensor<8x8xf32, #CSR>, %arg1: tensor<8xf32>) -> tensor<8xf32> {
%0 = linalg.init_tensor [8] : tensor<8xf32> %0 = tensor.empty() : tensor<8xf32>
%1 = linalg.generic #matvec %1 = linalg.generic #matvec
ins(%arg0, %arg1: tensor<8x8xf32, #CSR>, tensor<8xf32>) ins(%arg0, %arg1: tensor<8x8xf32, #CSR>, tensor<8xf32>)
outs(%0: tensor<8xf32>) { outs(%0: tensor<8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%m = arith.mulf %a, %b : f32 %m = arith.mulf %a, %b : f32
%add = arith.addf %x, %m : f32 %add = arith.addf %x, %m : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<8xf32> } -> tensor<8xf32>
return %1: tensor<8xf32> return %1: tensor<8xf32>
} }
// CHECK-LABEL: func @sparse_case // CHECK-LABEL: func @sparse_case
// CHECK-NEXT: linalg.init_tensor // CHECK-NEXT: tensor.empty
// CHECK-NEXT: linalg.generic // CHECK-NEXT: linalg.generic
// ----- // -----
func.func @reduce_dispatch_0() -> tensor<4x2xf32> { func.func @reduce_dispatch_0() -> tensor<4x2xf32> {
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%0 = linalg.init_tensor [4, 2] : tensor<4x2xf32> %0 = tensor.empty() : tensor<4x2xf32>
%res = scf.foreach_thread (%arg0, %arg1) in (%c4, %c2) shared_outs(%o = %0) -> (tensor<4x2xf32>) { %res = scf.foreach_thread (%arg0, %arg1) in (%c4, %c2) shared_outs(%o = %0) -> (tensor<4x2xf32>) {
%1 = linalg.init_tensor [1, 1] : tensor<1x1xf32> %1 = tensor.empty() : tensor<1x1xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
// CHECK: tensor.parallel_insert_slice %{{[0-9a-z]*}} into %{{[0-9a-z]*}} // CHECK: tensor.parallel_insert_slice %{{[0-9a-z]*}} into %{{[0-9a-z]*}}
// CHECK-SAME: [%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<4x2xf32> // CHECK-SAME: [%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<4x2xf32>
tensor.parallel_insert_slice %2 into %o[%arg0, %arg1] [1, 1] [1, 1] : tensor.parallel_insert_slice %2 into %o[%arg0, %arg1] [1, 1] [1, 1] :
tensor<1x1xf32> into tensor<4x2xf32> tensor<1x1xf32> into tensor<4x2xf32>
} }
} }
Show All 28 Lines

mlir/test/Dialect/Linalg/fuse-with-reshape-by-collapsing.mlir

// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing -split-input-file | FileCheck %s
// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing-control -split-input-file | FileCheck %s --check-prefix=CONTROL // RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing-control -split-input-file | FileCheck %s --check-prefix=CONTROL
// Static problem sizes. Checks all aspects of fusion by collapsing. Rest of the // Static problem sizes. Checks all aspects of fusion by collapsing. Rest of the
// tests only check a subset of conditions. // tests only check a subset of conditions.
#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> #map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)> #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>
#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)> #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>
#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> #map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
func.func @fuse_by_collapsing(%arg0 : tensor<2x12x5x336x9xi32>, func.func @fuse_by_collapsing(%arg0 : tensor<2x12x5x336x9xi32>,
%arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> { %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> {
%expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]]
: tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32> : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>
%init = linalg.init_tensor [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x3x4x5x6x7x8x9xi32> %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>
%generic = linalg.generic { %generic = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>) ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)
outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) { outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {
^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32): ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):
%t0 = arith.addi %b0, %b1 : i32 %t0 = arith.addi %b0, %b1 : i32
%t1 = arith.addi %t0, %b2 : i32 %t1 = arith.addi %t0, %b2 : i32
linalg.yield %t1 : i32 linalg.yield %t1 : i32
} -> tensor<2x3x4x5x6x7x8x9xi32> } -> tensor<2x3x4x5x6x7x8x9xi32>
return %generic : tensor<2x3x4x5x6x7x8x9xi32> return %generic : tensor<2x3x4x5x6x7x8x9xi32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>
// CHECK: func @fuse_by_collapsing( // CHECK: func @fuse_by_collapsing(
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32> // CHECK-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32> // CHECK-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32>
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [2, 3, 4, 5, 6, 7, 8, 9] // CHECK-DAG: %[[INIT:.+]] = tensor.empty()
// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1, 2]{{\]}} // CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1, 2]{{\]}}
// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}} // CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}
// CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}} // CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}
// CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic // CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :
// CHECK-SAME: outs(%[[INIT_RESHAPE]] : // CHECK-SAME: outs(%[[INIT_RESHAPE]] :
Show All 14 Lines
#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> #map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)> #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>
#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)> #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>
#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> #map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
func.func @fuse_by_collapsing_indexing_op(%arg0 : tensor<2x12x5x336x9xi32>, func.func @fuse_by_collapsing_indexing_op(%arg0 : tensor<2x12x5x336x9xi32>,
%arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> { %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> {
%expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]]
: tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32> : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>
%init = linalg.init_tensor [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x3x4x5x6x7x8x9xi32> %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>
%generic = linalg.generic { %generic = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>) ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)
outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) { outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {
^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32): ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):
%iv0 = linalg.index 0: index %iv0 = linalg.index 0: index
%iv1 = linalg.index 1: index %iv1 = linalg.index 1: index
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d7, d5, d6, d0, d1, d2, d3, d4)> #map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d7, d5, d6, d0, d1, d2, d3, d4)>
#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d5, d6, d0)> #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d5, d6, d0)>
#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d1, d2, d3)> #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d1, d2, d3)>
#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)> #map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
func.func @fuse_by_collapsing_change_reshape_order(%arg0 : tensor<9x56x2x60x6xi32>, func.func @fuse_by_collapsing_change_reshape_order(%arg0 : tensor<9x56x2x60x6xi32>,
%arg1 : tensor<7x8x2xi32>, %arg2 : tensor<6x3x4x5xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> { %arg1 : tensor<7x8x2xi32>, %arg2 : tensor<6x3x4x5xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> {
%expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]]
: tensor<9x56x2x60x6xi32> into tensor<9x7x8x2x3x4x5x6xi32> : tensor<9x56x2x60x6xi32> into tensor<9x7x8x2x3x4x5x6xi32>
%init = linalg.init_tensor [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x3x4x5x6x7x8x9xi32> %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>
%generic = linalg.generic { %generic = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
ins(%expand, %arg1, %arg2 : tensor<9x7x8x2x3x4x5x6xi32>, tensor<7x8x2xi32>, tensor<6x3x4x5xi32>) ins(%expand, %arg1, %arg2 : tensor<9x7x8x2x3x4x5x6xi32>, tensor<7x8x2xi32>, tensor<6x3x4x5xi32>)
outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) { outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {
^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32): ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):
%t0 = arith.addi %b0, %b1 : i32 %t0 = arith.addi %b0, %b1 : i32
%t1 = arith.addi %t0, %b2 : i32 %t1 = arith.addi %t0, %b2 : i32
linalg.yield %t1 : i32 linalg.yield %t1 : i32
} -> tensor<2x3x4x5x6x7x8x9xi32> } -> tensor<2x3x4x5x6x7x8x9xi32>
return %generic : tensor<2x3x4x5x6x7x8x9xi32> return %generic : tensor<2x3x4x5x6x7x8x9xi32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d4, d3, d0, d1, d2)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d4, d3, d0, d1, d2)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d0)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
// CHECK: func @fuse_by_collapsing_change_reshape_order( // CHECK: func @fuse_by_collapsing_change_reshape_order(
// CHECK-SAME: %[[ARG0:.+]]: tensor<9x56x2x60x6xi32> // CHECK-SAME: %[[ARG0:.+]]: tensor<9x56x2x60x6xi32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<7x8x2xi32> // CHECK-SAME: %[[ARG1:.+]]: tensor<7x8x2xi32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<6x3x4x5xi32> // CHECK-SAME: %[[ARG2:.+]]: tensor<6x3x4x5xi32>
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [2, 3, 4, 5, 6, 7, 8, 9] // CHECK-DAG: %[[INIT:.+]] = tensor.empty()
// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1], [2]{{\]}} // CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1], [2]{{\]}}
// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}} // CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}
// CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}} // CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}}
// CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic // CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :
// CHECK-SAME: outs(%[[INIT_RESHAPE]] : // CHECK-SAME: outs(%[[INIT_RESHAPE]] :
Show All 14 Linesfunc.func @fuse_by_collapsing_dynamic(%arg0 : tensor<?x?x?x?x?xi32>,
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]]
: tensor<?x?x?x?x?xi32> into tensor<?x7x?x?x3x?x5x?xi32> : tensor<?x?x?x?x?xi32> into tensor<?x7x?x?x3x?x5x?xi32>
%d0 = tensor.dim %arg1, %c2 : tensor<?x?x?xi32> %d0 = tensor.dim %arg1, %c2 : tensor<?x?x?xi32>
%d2 = tensor.dim %arg2, %c2 : tensor<?x?x?x?xi32> %d2 = tensor.dim %arg2, %c2 : tensor<?x?x?x?xi32>
%d4 = tensor.dim %arg2, %c0 : tensor<?x?x?x?xi32> %d4 = tensor.dim %arg2, %c0 : tensor<?x?x?x?xi32>
%d6 = tensor.dim %arg1, %c1 : tensor<?x?x?xi32> %d6 = tensor.dim %arg1, %c1 : tensor<?x?x?xi32>
%d7 = tensor.dim %arg0, %c0 : tensor<?x?x?x?x?xi32> %d7 = tensor.dim %arg0, %c0 : tensor<?x?x?x?x?xi32>
%init = linalg.init_tensor [%d0, 3, %d2, 5, %d4, 7, %d6, %d7] : tensor<?x3x?x5x?x7x?x?xi32> %init = tensor.empty(%d0, %d2, %d4, %d6, %d7) : tensor<?x3x?x5x?x7x?x?xi32>
%generic = linalg.generic { %generic = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}
ins(%expand, %arg1, %arg2 : tensor<?x7x?x?x3x?x5x?xi32>, tensor<?x?x?xi32>, tensor<?x?x?x?xi32>) ins(%expand, %arg1, %arg2 : tensor<?x7x?x?x3x?x5x?xi32>, tensor<?x?x?xi32>, tensor<?x?x?x?xi32>)
outs(%init : tensor<?x3x?x5x?x7x?x?xi32>) { outs(%init : tensor<?x3x?x5x?x7x?x?xi32>) {
^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32): ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):
%iv0 = linalg.index 0: index %iv0 = linalg.index 0: index
%iv1 = linalg.index 1: index %iv1 = linalg.index 1: index
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines
// ----- // -----
// Test no fusion because the folded dimensions are not all preserved. // Test no fusion because the folded dimensions are not all preserved.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
func.func @no_fuse_unpreserved_folding(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x3x4x5xf32> { func.func @no_fuse_unpreserved_folding(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x3x4x5xf32> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32> %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>
%init = linalg.init_tensor [2, 3, 4, 5] : tensor<2x3x4x5xf32> %init = tensor.empty(): tensor<2x3x4x5xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map0], indexing_maps = [#map0, #map1, #map0],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x3x4x5xf32>) { ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x3x4x5xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%2 = arith.addf %b0, %b1 : f32 %2 = arith.addf %b0, %b1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<2x3x4x5xf32> } -> tensor<2x3x4x5xf32>
Show All 10 Lines
// ----- // -----
// Test no fusion because the folded dimensions are not all preserved. // Test no fusion because the folded dimensions are not all preserved.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)> #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>
func.func @no_fuse_unpreserved_folding_transpose(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2xf32>) -> tensor<2x4x3x5xf32> { func.func @no_fuse_unpreserved_folding_transpose(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2xf32>) -> tensor<2x4x3x5xf32> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32> %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>
%init = linalg.init_tensor [2, 4, 3, 5] : tensor<2x4x3x5xf32> %init = tensor.empty() : tensor<2x4x3x5xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map2], indexing_maps = [#map0, #map1, #map2],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2xf32>) outs(%init : tensor<2x4x3x5xf32>) { ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2xf32>) outs(%init : tensor<2x4x3x5xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%2 = arith.addf %b0, %b1 : f32 %2 = arith.addf %b0, %b1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<2x4x3x5xf32> } -> tensor<2x4x3x5xf32>
Show All 10 Lines
// ----- // -----
// Test no fusion because the iterator types of folded dims are not preserved. // Test no fusion because the iterator types of folded dims are not preserved.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3)> #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3)>
func.func @no_fuse_mismatched_iterator_types(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x5xf32> { func.func @no_fuse_mismatched_iterator_types(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x5xf32> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32> %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>
%init = linalg.init_tensor [2, 5] : tensor<2x5xf32> %init = tensor.empty() : tensor<2x5xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map2], indexing_maps = [#map0, #map1, #map2],
iterator_types = ["parallel", "reduction", "parallel", "parallel"]} iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x5xf32>) { ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x5xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%2 = arith.addf %b0, %b1 : f32 %2 = arith.addf %b0, %b1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<2x5xf32> } -> tensor<2x5xf32>
Show All 12 Lines
// Test control of fusion using control function // Test control of fusion using control function
// Test no fusion because the folded dimensions are not all preserved. // Test no fusion because the folded dimensions are not all preserved.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d3)> #map1 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> #map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
func.func @control_fusion(%arg0 : tensor<6xf32>, %arg1 : tensor<20xf32>) -> tensor<2x3x4x5xf32> { func.func @control_fusion(%arg0 : tensor<6xf32>, %arg1 : tensor<20xf32>) -> tensor<2x3x4x5xf32> {
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<6xf32> into tensor<2x3xf32> %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<6xf32> into tensor<2x3xf32>
%1 = tensor.expand_shape %arg1 [[0, 1]] : tensor<20xf32> into tensor<4x5xf32> %1 = tensor.expand_shape %arg1 [[0, 1]] : tensor<20xf32> into tensor<4x5xf32>
%init = linalg.init_tensor [2, 3, 4, 5] : tensor<2x3x4x5xf32> %init = tensor.empty() : tensor<2x3x4x5xf32>
%2 = linalg.generic { %2 = linalg.generic {
indexing_maps = [#map0, #map1, #map2], indexing_maps = [#map0, #map1, #map2],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%0, %1 : tensor<2x3xf32>, tensor<4x5xf32>) outs(%init : tensor<2x3x4x5xf32>) { ins(%0, %1 : tensor<2x3xf32>, tensor<4x5xf32>) outs(%init : tensor<2x3x4x5xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%3 = arith.addf %b0, %b1 : f32 %3 = arith.addf %b0, %b1 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<2x3x4x5xf32> } -> tensor<2x3x4x5xf32>
Show All 16 Lines
// CONTROL-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CONTROL-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CONTROL-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)> // CONTROL-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)>
// CONTROL-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CONTROL-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CONTROL: func @control_fusion( // CONTROL: func @control_fusion(
// CONTROL-SAME: %[[ARG0:.+]]: tensor<6xf32> // CONTROL-SAME: %[[ARG0:.+]]: tensor<6xf32>
// CONTROL-SAME: %[[ARG1:.+]]: tensor<20xf32> // CONTROL-SAME: %[[ARG1:.+]]: tensor<20xf32>
// CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] // CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]
// CONTROL: %[[INIT:.+]] = linalg.init_tensor [2, 3, 4, 5] // CONTROL: %[[INIT:.+]] = tensor.empty()
// CONTROL: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]{{\]}} // CONTROL: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]{{\]}}
// CONTROL: %[[GENERIC:.+]] = linalg.generic // CONTROL: %[[GENERIC:.+]] = linalg.generic
// CONTROL-SAME: ins(%[[EXPAND]], %[[ARG1]] : // CONTROL-SAME: ins(%[[EXPAND]], %[[ARG1]] :
// CONTROL-SAME: outs(%[[INIT_RESHAPE]] : // CONTROL-SAME: outs(%[[INIT_RESHAPE]] :
// CONTROL: %[[RESULT:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1], [2, 3]{{\]}} // CONTROL: %[[RESULT:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1], [2, 3]{{\]}}
// ----- // -----
// Corner case that isnt handled currently. // Corner case that isnt handled currently.
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
func.func @zero_D_test(%arg0: tensor<f32>) -> tensor<1xf32> { func.func @zero_D_test(%arg0: tensor<f32>) -> tensor<1xf32> {
%0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32> %0 = tensor.expand_shape %arg0 [] : tensor<f32> into tensor<1xf32>
%init = linalg.init_tensor [1] : tensor<1xf32> %init = tensor.empty() : tensor<1xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%0: tensor<1xf32>) outs(%init : tensor<1xf32>) { ins(%0: tensor<1xf32>) outs(%init : tensor<1xf32>) {
^bb0(%b0 : f32, %b1 : f32): ^bb0(%b0 : f32, %b1 : f32):
linalg.yield %b0: f32 linalg.yield %b0: f32
} -> tensor<1xf32> } -> tensor<1xf32>
return %1 : tensor<1xf32> return %1 : tensor<1xf32>
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d1, d0, d2)> #map1 = affine_map<(d0, d1, d2, d3) -> (d3, d1, d0, d2)>
func.func @fold_non_consecutive_dims(%arg0 : tensor<?x?xi32>) -> tensor<?x8x?x4xi32> { func.func @fold_non_consecutive_dims(%arg0 : tensor<?x?xi32>) -> tensor<?x8x?x4xi32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xi32> into tensor<?x4x?x8xi32> %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xi32> into tensor<?x4x?x8xi32>
%d0 = tensor.dim %0, %c0 : tensor<?x4x?x8xi32> %d0 = tensor.dim %0, %c0 : tensor<?x4x?x8xi32>
%d1 = tensor.dim %0, %c2 : tensor<?x4x?x8xi32> %d1 = tensor.dim %0, %c2 : tensor<?x4x?x8xi32>
%init = linalg.init_tensor [%d1, 8, %d0, 4] : tensor<?x8x?x4xi32> %init = tensor.empty(%d1, %d0) : tensor<?x8x?x4xi32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1], indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<?x8x?x4xi32>) { ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<?x8x?x4xi32>) {
^bb0(%b0 : i32, %b1 : i32): ^bb0(%b0 : i32, %b1 : i32):
%2 = linalg.index 0 : index %2 = linalg.index 0 : index
%3 = linalg.index 1 : index %3 = linalg.index 1 : index
%4 = linalg.index 2 : index %4 = linalg.index 2 : index
%5 = linalg.index 3 : index %5 = linalg.index 3 : index
%6 = arith.addi %2, %3 : index %6 = arith.addi %2, %3 : index
%7 = arith.addi %6, %4 : index %7 = arith.addi %6, %4 : index
%8 = arith.addi %7, %5 : index %8 = arith.addi %7, %5 : index
%9 = arith.index_cast %8 : index to i32 %9 = arith.index_cast %8 : index to i32
linalg.yield %9: i32 linalg.yield %9: i32
} -> tensor<?x8x?x4xi32> } -> tensor<?x8x?x4xi32>
return %1 : tensor<?x8x?x4xi32> return %1 : tensor<?x8x?x4xi32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK: func @fold_non_consecutive_dims( // CHECK: func @fold_non_consecutive_dims(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>)
// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index // CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index // CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[COLLAPSE_INIT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2, 3]{{\]}} // CHECK: %[[COLLAPSE_INIT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2, 3]{{\]}}
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]] : // CHECK-SAME: ins(%[[ARG0]] :
// CHECK-SAME: outs(%[[COLLAPSE_INIT]] : // CHECK-SAME: outs(%[[COLLAPSE_INIT]] :
// CHECK-NEXT: ^bb{{[0-9]}} // CHECK-NEXT: ^bb{{[0-9]}}
// CHECK: %[[ID0:.+]] = linalg.index 0 // CHECK: %[[ID0:.+]] = linalg.index 0
Show All 15 Lines
// None of the folded iteration space dims are contiguous reduction dimensions. // None of the folded iteration space dims are contiguous reduction dimensions.
// So no change in the code. // So no change in the code.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>
#map1 = affine_map<(d0, d1, d2, d3) -> ()> #map1 = affine_map<(d0, d1, d2, d3) -> ()>
func.func @no_fold_non_consecutive_reduction_dims(%arg0 : tensor<?x?xi32>) -> tensor<i32> { func.func @no_fold_non_consecutive_reduction_dims(%arg0 : tensor<?x?xi32>) -> tensor<i32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xi32> into tensor<?x4x?x8xi32> %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xi32> into tensor<?x4x?x8xi32>
%init = linalg.init_tensor [] : tensor<i32> %init = tensor.empty() : tensor<i32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1], indexing_maps = [#map0, #map1],
iterator_types = ["reduction", "reduction", "reduction", "reduction"]} iterator_types = ["reduction", "reduction", "reduction", "reduction"]}
ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<i32>) { ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<i32>) {
^bb0(%b0 : i32, %b1 : i32): ^bb0(%b0 : i32, %b1 : i32):
%2 = linalg.index 0 : index %2 = linalg.index 0 : index
%3 = linalg.index 1 : index %3 = linalg.index 1 : index
%4 = linalg.index 2 : index %4 = linalg.index 2 : index
Show All 15 Lines

mlir/test/Dialect/Linalg/fusion-elementwise-ops.mlir

// RUN: mlir-opt %s -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s // RUN: mlir-opt %s -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s
// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @add_mul_fusion // CHECK-LABEL: @add_mul_fusion
func.func @add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>) { outs(%2 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%4 = arith.addf %arg3, %arg4 : f32 %4 = arith.addf %arg3, %arg4 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
Show All 25 Lines
// CHECK-LABEL: @scalar_add_mul_fusion // CHECK-LABEL: @scalar_add_mul_fusion
func.func @scalar_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : f32, %arg2 : f32) -> tensor<?x?xf32> func.func @scalar_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : f32, %arg2 : f32) -> tensor<?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, f32) ins(%arg0, %arg1 : tensor<?x?xf32>, f32)
outs(%2 : tensor<?x?xf32>) { outs(%2 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%4 = arith.addf %arg3, %arg4 : f32 %4 = arith.addf %arg3, %arg4 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
Show All 25 Lines
// CHECK-LABEL: @transpose_add_mul_fusion // CHECK-LABEL: @transpose_add_mul_fusion
func.func @transpose_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @transpose_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>) { outs(%2 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%4 = arith.addf %arg3, %arg4 : f32 %4 = arith.addf %arg3, %arg4 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
Show All 17 Lines
// CHECK-LABEL: @add_transpose_mul_fusion // CHECK-LABEL: @add_transpose_mul_fusion
func.func @add_transpose_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @add_transpose_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>) { outs(%2 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%4 = arith.addf %arg3, %arg4 : f32 %4 = arith.addf %arg3, %arg4 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
Show All 17 Lines
#map2 = affine_map<(d0) -> (d0)> #map2 = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @add_broadcast_mul_fusion // CHECK-LABEL: @add_broadcast_mul_fusion
func.func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
%1 = linalg.init_tensor [%0] : tensor<?xf32> %1 = tensor.empty(%0) : tensor<?xf32>
%2 = linalg.generic {indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]} %2 = linalg.generic {indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%1 : tensor<?xf32>) { outs(%1 : tensor<?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%3 = arith.addf %arg3, %arg4 : f32 %3 = arith.addf %arg3, %arg4 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP1]], [[$MAP0]], [[$MAP0]] // CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP1]], [[$MAP0]], [[$MAP0]]
%3 = tensor.dim %arg2, %c1 : tensor<?x?xf32> %3 = tensor.dim %arg2, %c1 : tensor<?x?xf32>
%4 = linalg.init_tensor [%0, %3] : tensor<?x?xf32> %4 = tensor.empty(%0, %3) : tensor<?x?xf32>
%5 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %5 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%2, %arg2 : tensor<?xf32>, tensor<?x?xf32>) ins(%2, %arg2 : tensor<?xf32>, tensor<?x?xf32>)
outs(%4 : tensor<?x?xf32>){ outs(%4 : tensor<?x?xf32>){
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
%6 = arith.mulf %arg5, %arg6 : f32 %6 = arith.mulf %arg5, %arg6 : f32
linalg.yield %6 : f32 linalg.yield %6 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %5 : tensor<?x?xf32> return %5 : tensor<?x?xf32>
} }
// ----- // -----
// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()> // CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
#map0 = affine_map<() -> ()> #map0 = affine_map<() -> ()>
// CHECK-LABEL: @add_mul_scalar_fusion // CHECK-LABEL: @add_mul_scalar_fusion
func.func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> func.func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32>
{ {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []} %1 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []}
ins(%arg0, %arg1 : tensor<f32>, tensor<f32>) ins(%arg0, %arg1 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { outs(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<f32> } -> tensor<f32>
// CHECK: linalg.generic { // CHECK: linalg.generic {
Show All 17 Lines
func.func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32> func.func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%cst = arith.constant dense<42.0> : tensor<5xf32> %cst = arith.constant dense<42.0> : tensor<5xf32>
%0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32> %0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32>
%1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32> %1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32>
%2 = linalg.init_tensor [5, %0, %1] : tensor<5x?x?xf32> %2 = tensor.empty(%0, %1) : tensor<5x?x?xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map0, #map1, #map1], indexing_maps = [#map0, #map1, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%cst, %arg0 : tensor<5xf32>, tensor<5x?x?xf32>) ins(%cst, %arg0 : tensor<5xf32>, tensor<5x?x?xf32>)
outs(%2 : tensor<5x?x?xf32>) { outs(%2 : tensor<5x?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%4 = arith.mulf %arg1, %arg2 : f32 %4 = arith.mulf %arg1, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
Show All 15 Linesfunc.func @generic_op_zero_dim_constant_fusion(%arg0 : tensor<5x?x?xf32>)
-> tensor<5x?x?xf32> -> tensor<5x?x?xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%cst = arith.constant dense<42.0> : tensor<f32> %cst = arith.constant dense<42.0> : tensor<f32>
%0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32> %0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32>
%1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32> %1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32>
%2 = linalg.init_tensor [5, %0, %1] : tensor<5x?x?xf32> %2 = tensor.empty(%0, %1) : tensor<5x?x?xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map0, #map1, #map1], indexing_maps = [#map0, #map1, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%cst, %arg0 : tensor<f32>, tensor<5x?x?xf32>) ins(%cst, %arg0 : tensor<f32>, tensor<5x?x?xf32>)
outs(%2 : tensor<5x?x?xf32>) { outs(%2 : tensor<5x?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%4 = arith.mulf %arg1, %arg2 : f32 %4 = arith.mulf %arg1, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
Show All 11 Lines
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
func.func @producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>, func.func @producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>,
%arg1: tensor<?x?xi32>) -> tensor<?x?xi32> { %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xi32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xi32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xi32> %2 = tensor.empty(%0, %1) : tensor<?x?xi32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>) ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { outs(%2 : tensor<?x?xi32>) {
^bb0(%arg2: i32, %arg3: i32, %arg4: i32): ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):
%10 = arith.addi %arg2, %arg3 : i32 %10 = arith.addi %arg2, %arg3 : i32
linalg.yield %10 : i32 linalg.yield %10 : i32
Show All 34 Lines
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
func.func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> { func.func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xi32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xi32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xi32> %2 = tensor.empty(%0, %1) : tensor<?x?xi32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map0, #map0], indexing_maps = [#map0, #map0],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%arg0 : tensor<?x?xi32>) ins(%arg0 : tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { outs(%2 : tensor<?x?xi32>) {
^bb0(%arg4: i32, %arg5: i32): ^bb0(%arg4: i32, %arg5: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
Show All 37 Lines
#map0 = affine_map<(d0, d1) -> (d1, d0)> #map0 = affine_map<(d0, d1) -> (d1, d0)>
#map1 = affine_map<(d0, d1) -> (d0, d1)> #map1 = affine_map<(d0, d1) -> (d0, d1)>
func.func @indexed_producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>) func.func @indexed_producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>)
-> tensor<?x?xi32> { -> tensor<?x?xi32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?xi32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xi32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xi32> %2 = tensor.empty(%0, %1) : tensor<?x?xi32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map0, #map0], indexing_maps = [#map0, #map0],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%arg0 : tensor<?x?xi32>) ins(%arg0 : tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { outs(%2 : tensor<?x?xi32>) {
^bb0(%arg2: i32, %arg3: i32): ^bb0(%arg2: i32, %arg3: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
#map1 = affine_map<(d0) -> (d0)> #map1 = affine_map<(d0) -> (d0)>
#map2 = affine_map<(d0, d1) -> (d0, d1)> #map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1) -> (d1)> #map3 = affine_map<(d0, d1) -> (d1)>
func.func @one_dim_indexed_producer_consumer_fusion(%arg0 : tensor<?xi32>, func.func @one_dim_indexed_producer_consumer_fusion(%arg0 : tensor<?xi32>,
%arg1 : tensor<?x?xi32>) -> tensor<?x?xi32> { %arg1 : tensor<?x?xi32>) -> tensor<?x?xi32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?xi32> %d0 = tensor.dim %arg0, %c0 : tensor<?xi32>
%0 = linalg.init_tensor [%d0] : tensor<?xi32> %0 = tensor.empty(%d0) : tensor<?xi32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [#map1, #map1], {indexing_maps = [#map1, #map1],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0 : tensor<?xi32>) outs(%0 : tensor<?xi32>) { ins(%arg0 : tensor<?xi32>) outs(%0 : tensor<?xi32>) {
^bb0(%arg2 : i32, %arg3 : i32): ^bb0(%arg2 : i32, %arg3 : i32):
%2 = linalg.index 0 : index %2 = linalg.index 0 : index
%3 = arith.index_cast %2 : index to i32 %3 = arith.index_cast %2 : index to i32
%4 = arith.addi %arg2, %3 : i32 %4 = arith.addi %arg2, %3 : i32
linalg.yield %4 : i32 linalg.yield %4 : i32
} -> tensor<?xi32> } -> tensor<?xi32>
%2 = tensor.dim %arg1, %c0 : tensor<?x?xi32> %2 = tensor.dim %arg1, %c0 : tensor<?x?xi32>
%3 = tensor.dim %arg1, %c1 : tensor<?x?xi32> %3 = tensor.dim %arg1, %c1 : tensor<?x?xi32>
%4 = linalg.init_tensor [%2, %3] : tensor<?x?xi32> %4 = tensor.empty(%2, %3) : tensor<?x?xi32>
%5 = linalg.generic %5 = linalg.generic
{indexing_maps = [#map2, #map3, #map2], {indexing_maps = [#map2, #map3, #map2],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg1, %1 : tensor<?x?xi32>, tensor<?xi32>) ins(%arg1, %1 : tensor<?x?xi32>, tensor<?xi32>)
outs(%4 : tensor<?x?xi32>) { outs(%4 : tensor<?x?xi32>) {
^bb0(%arg2 : i32, %arg3 : i32, %arg4: i32): ^bb0(%arg2 : i32, %arg3 : i32, %arg4: i32):
%6 = arith.addi %arg2, %arg3 : i32 %6 = arith.addi %arg2, %arg3 : i32
linalg.yield %6 : i32 linalg.yield %6 : i32
Show All 16 Lines
// ----- // -----
func.func @scalar_generic_fusion func.func @scalar_generic_fusion
(%arg0: tensor<5x1x1xf32>, %arg1 : tensor<i32>) -> tensor<10xf32> (%arg0: tensor<5x1x1xf32>, %arg1 : tensor<i32>) -> tensor<10xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant dense<1.000000e+00> : tensor<10xf32> %cst = arith.constant dense<1.000000e+00> : tensor<10xf32>
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>], {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>],
iterator_types = []} iterator_types = []}
ins(%arg1 : tensor<i32>) outs(%0 : tensor<f32>) { ins(%arg1 : tensor<i32>) outs(%0 : tensor<f32>) {
^bb0(%arg2: i32, %arg3: f32): ^bb0(%arg2: i32, %arg3: f32):
%3 = arith.index_cast %arg2 : i32 to index %3 = arith.index_cast %arg2 : i32 to index
%4 = tensor.extract %arg0[%3, %c0, %c0] : tensor<5x1x1xf32> %4 = tensor.extract %arg0[%3, %c0, %c0] : tensor<5x1x1xf32>
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<f32> } -> tensor<f32>
%2 = linalg.init_tensor [10] : tensor<10xf32> %2 = tensor.empty() : tensor<10xf32>
%3 = linalg.generic %3 = linalg.generic
{indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>, {indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>], affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%1, %cst : tensor<f32>, tensor<10xf32>) outs(%2 : tensor<10xf32>) { ins(%1, %cst : tensor<f32>, tensor<10xf32>) outs(%2 : tensor<10xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
%4 = arith.mulf %arg2, %arg3 : f32 %4 = arith.mulf %arg2, %arg3 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
Show All 12 Lines
// CHECK: tensor.extract %[[ARG0]] // CHECK: tensor.extract %[[ARG0]]
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK return %[[T0]] // CHECK return %[[T0]]
// ----- // -----
func.func @constant_fusion(%arg0 : tensor<4xf32>) -> (tensor<4xf32>) { func.func @constant_fusion(%arg0 : tensor<4xf32>) -> (tensor<4xf32>) {
%cst = arith.constant dense<1.0> : tensor<4xf32> %cst = arith.constant dense<1.0> : tensor<4xf32>
%1 = linalg.init_tensor [4] : tensor<4xf32> %1 = tensor.empty() : tensor<4xf32>
%2 = linalg.generic %2 = linalg.generic
{indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>], affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins (%arg0, %cst : tensor<4xf32>, tensor<4xf32>) ins (%arg0, %cst : tensor<4xf32>, tensor<4xf32>)
outs (%1 : tensor<4xf32>) { outs (%1 : tensor<4xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%3 = arith.addf %arg1, %arg2 : f32 %3 = arith.addf %arg1, %arg2 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %2 : tensor<4xf32> return %2 : tensor<4xf32>
} }
// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[MAP:.+]] = affine_map<(d0) -> (d0)>
// CHECK: func @constant_fusion(%[[ARG0:.+]]: tensor<4xf32>) // CHECK: func @constant_fusion(%[[ARG0:.+]]: tensor<4xf32>)
// CHECK-DAG: %[[CST:.+]] = arith.constant 1.000000e+00 : f32 // CHECK-DAG: %[[CST:.+]] = arith.constant 1.000000e+00 : f32
// CHECK-DAG: %[[T0:.+]] = linalg.init_tensor [4] : tensor<4xf32> // CHECK-DAG: %[[T0:.+]] = tensor.empty() : tensor<4xf32>
// CHECK: %[[T1:.+]] = linalg.generic // CHECK: %[[T1:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]] // CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]
// CHECK-SAME: ins(%[[ARG0]] : tensor<4xf32>) // CHECK-SAME: ins(%[[ARG0]] : tensor<4xf32>)
// CHECK-SAME: outs(%[[T0]] : tensor<4xf32>) // CHECK-SAME: outs(%[[T0]] : tensor<4xf32>)
// CHECK: ^{{.+}}( // CHECK: ^{{.+}}(
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: f32, %[[ARG2:[a-zA-Z0-9_]+]]: f32) // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: f32, %[[ARG2:[a-zA-Z0-9_]+]]: f32)
// CHECK: %[[T2:.+]] = arith.addf %[[ARG1]], %[[CST]] // CHECK: %[[T2:.+]] = arith.addf %[[ARG1]], %[[CST]]
// CHECK: linalg.yield %[[T2]] // CHECK: linalg.yield %[[T2]]
// CHECK: return %[[T1]] // CHECK: return %[[T1]]
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0) -> (0, d0)> #map1 = affine_map<(d0) -> (0, d0)>
#map2 = affine_map<(d0) -> (0)> #map2 = affine_map<(d0) -> (0)>
func.func @consumer_with_reduction(%arg0: tensor<1x10xf32>, func.func @consumer_with_reduction(%arg0: tensor<1x10xf32>,
%arg1: tensor<1x10xf32>, %arg1: tensor<1x10xf32>,
%arg2: tensor<1xf32>) -> tensor<1xf32> { %arg2: tensor<1xf32>) -> tensor<1xf32> {
%init = linalg.init_tensor [1, 10] : tensor<1x10xf32> %init = tensor.empty() : tensor<1x10xf32>
%0 = linalg.generic %0 = linalg.generic
{indexing_maps = [#map0, #map0, #map0], {indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<1x10xf32>, tensor<1x10xf32>) ins(%arg0, %arg1 : tensor<1x10xf32>, tensor<1x10xf32>)
outs(%init : tensor<1x10xf32>) { outs(%init : tensor<1x10xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
Show All 29 Lines
// CHECK-NOT: linalg.generic // CHECK-NOT: linalg.generic
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
func.func @sigmoid_dynamic_dim(%0: tensor<?x1xf32>) -> tensor<?x1xf32> { func.func @sigmoid_dynamic_dim(%0: tensor<?x1xf32>) -> tensor<?x1xf32> {
%cp5 = arith.constant 5.000000e-01 : f32 %cp5 = arith.constant 5.000000e-01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%shape = shape.shape_of %0 : tensor<?x1xf32> -> tensor<?xindex> %shape = shape.shape_of %0 : tensor<?x1xf32> -> tensor<?xindex>
%extend = shape.to_extent_tensor %shape : tensor<?xindex> -> tensor<2xindex> %extend = shape.to_extent_tensor %shape : tensor<?xindex> -> tensor<2xindex>
%extracted = tensor.extract %extend[%c0] : tensor<2xindex> %extracted = tensor.extract %extend[%c0] : tensor<2xindex>
%init0 = linalg.init_tensor [%extracted, 1] : tensor<?x1xf32> %init0 = tensor.empty(%extracted) : tensor<?x1xf32>
%1 = linalg.generic {indexing_maps = [ %1 = linalg.generic {indexing_maps = [
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
outs(%init0 : tensor<?x1xf32>) { outs(%init0 : tensor<?x1xf32>) {
^bb0(%a: f32): ^bb0(%a: f32):
linalg.yield %cp5 : f32 linalg.yield %cp5 : f32
} -> tensor<?x1xf32> } -> tensor<?x1xf32>
%d0 = tensor.dim %0, %c0 : tensor<?x1xf32> %d0 = tensor.dim %0, %c0 : tensor<?x1xf32>
%init1 = linalg.init_tensor [%d0, 1] : tensor<?x1xf32> %init1 = tensor.empty(%d0) : tensor<?x1xf32>
%2 = linalg.generic {indexing_maps = [ %2 = linalg.generic {indexing_maps = [
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%0, %1 : tensor<?x1xf32>, tensor<?x1xf32>) ins(%0, %1 : tensor<?x1xf32>, tensor<?x1xf32>)
outs(%init1 : tensor<?x1xf32>) { outs(%init1 : tensor<?x1xf32>) {
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
// CHECK-LABEL: func @no_fuse_constant_with_reduction // CHECK-LABEL: func @no_fuse_constant_with_reduction
func.func @no_fuse_constant_with_reduction() -> tensor<3xf32> func.func @no_fuse_constant_with_reduction() -> tensor<3xf32>
{ {
// CHECK: %[[CONST:.+]] = arith.constant {{.+}} : tensor<3x2xf32> // CHECK: %[[CONST:.+]] = arith.constant {{.+}} : tensor<3x2xf32>
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: ins(%[[CONST]] : tensor<3x2xf32>) // CHECK-SAME: ins(%[[CONST]] : tensor<3x2xf32>)
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
%three = arith.constant dense<3.0> : tensor<3x2xf32> %three = arith.constant dense<3.0> : tensor<3x2xf32>
%init = linalg.init_tensor [3] : tensor<3xf32> %init = tensor.empty() : tensor<3xf32>
%result = linalg.generic { %result = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%three : tensor<3x2xf32>) outs(%init : tensor<3xf32>) { ins(%three : tensor<3x2xf32>) outs(%init : tensor<3xf32>) {
^bb0(%arg0 : f32, %arg1 : f32): ^bb0(%arg0 : f32, %arg1 : f32):
%0 = arith.addf %arg0, %arg1 : f32 %0 = arith.addf %arg0, %arg1 : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
Show All 23 Linesfunc.func @break_outs_dependency(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %2 : tensor<?x?xf32> return %2 : tensor<?x?xf32>
} }
// CHECK: func @break_outs_dependency( // CHECK: func @break_outs_dependency(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK: %[[GENERIC1:.+]] = linalg.generic // CHECK: %[[GENERIC1:.+]] = linalg.generic
// CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>) // CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>)
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[GENERIC1]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[GENERIC1]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[GENERIC1]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[GENERIC1]], %[[C1]]
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>) // CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>)
// ----- // -----
func.func @fuse_scalar_constant(%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi32>) { func.func @fuse_scalar_constant(%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi32>) {
%cst = arith.constant 4.0 : f32 %cst = arith.constant 4.0 : f32
%c42 = arith.constant 42 : i32 %c42 = arith.constant 42 : i32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%0 = linalg.init_tensor[%d0, %d1] : tensor<?x?xf32> %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%1 = linalg.init_tensor[%d0, %d1] : tensor<?x?xi32> %1 = tensor.empty(%d0, %d1) : tensor<?x?xi32>
%2:2 = linalg.generic { %2:2 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> ()>,
affine_map<(d0, d1) -> ()>, affine_map<(d0, d1) -> ()>,
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %cst, %c42 : tensor<?x?xf32>, f32, i32) ins(%arg0, %cst, %c42 : tensor<?x?xf32>, f32, i32)
▲ Show 20 Lines • Show All 157 Lines • ▼ Show 20 Lines
#map2 = affine_map<(d0, d1) -> (d1, d0)> #map2 = affine_map<(d0, d1) -> (d1, d0)>
#map3 = affine_map<(d0, d1) -> (d0)> #map3 = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: @no_fusion_missing_reduction_shape // CHECK-LABEL: @no_fusion_missing_reduction_shape
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK: linalg.generic // CHECK: linalg.generic
func.func @no_fusion_missing_reduction_shape(%arg0: tensor<f32>, %arg1: index) -> tensor<?xf32> { func.func @no_fusion_missing_reduction_shape(%arg0: tensor<f32>, %arg1: index) -> tensor<?xf32> {
%cst = arith.constant 0xFF800000 : f32 %cst = arith.constant 0xFF800000 : f32
%4 = linalg.init_tensor [%arg1, %arg1] : tensor<?x?xf32> %4 = tensor.empty(%arg1, %arg1) : tensor<?x?xf32>
%5 = linalg.generic { %5 = linalg.generic {
indexing_maps = [#map0, #map1], indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0 : tensor<f32>) outs(%4 : tensor<?x?xf32>) { } ins(%arg0 : tensor<f32>) outs(%4 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%6 = linalg.init_tensor [%arg1] : tensor<?xf32> %6 = tensor.empty(%arg1) : tensor<?xf32>
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32> %7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>
%8 = linalg.generic { %8 = linalg.generic {
indexing_maps = [#map2, #map3], indexing_maps = [#map2, #map3],
iterator_types = ["parallel", "reduction"] iterator_types = ["parallel", "reduction"]
} ins(%5 : tensor<?x?xf32>) outs(%7 : tensor<?xf32>) { } ins(%5 : tensor<?x?xf32>) outs(%7 : tensor<?xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
%9 = arith.maxf %arg2, %arg3 : f32 %9 = arith.maxf %arg2, %arg3 : f32
linalg.yield %9 : f32 linalg.yield %9 : f32
Show All 9 Lines %0 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>], indexing_maps = [affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%arg0 : tensor<5000xi64>) { outs(%arg0 : tensor<5000xi64>) {
^bb0(%arg3: i64): // no predecessors ^bb0(%arg3: i64): // no predecessors
%22 = linalg.index 0 : index %22 = linalg.index 0 : index
%23 = arith.index_cast %22 : index to i64 %23 = arith.index_cast %22 : index to i64
linalg.yield %23 : i64 linalg.yield %23 : i64
} -> tensor<5000xi64> } -> tensor<5000xi64>
%1 = linalg.init_tensor [5000] : tensor<5000xi32> %1 = tensor.empty() : tensor<5000xi32>
%2 = linalg.generic { %2 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%0 : tensor<5000xi64>) outs(%1 : tensor<5000xi32>) { ins(%0 : tensor<5000xi64>) outs(%1 : tensor<5000xi32>) {
^bb0(%arg3: i64, %arg5: i32): // no predecessors ^bb0(%arg3: i64, %arg5: i32): // no predecessors
%22 = arith.index_cast %arg3 : i64 to index %22 = arith.index_cast %arg3 : i64 to index
%23 = tensor.extract %arg1[%22] : tensor<5000xi32> %23 = tensor.extract %arg1[%22] : tensor<5000xi32>
linalg.yield %23 : i32 linalg.yield %23 : i32
} -> tensor<5000xi32> } -> tensor<5000xi32>
return %2 : tensor<5000xi32> return %2 : tensor<5000xi32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)>
// CHECK: func @fusion_different_axes( // CHECK: func @fusion_different_axes(
// CHECK-SAME: %[[ARG0:.+]]: tensor<5000xi64> // CHECK-SAME: %[[ARG0:.+]]: tensor<5000xi64>
// CHECK-SAME: %[[ARG1:.+]]: tensor<5000xi32> // CHECK-SAME: %[[ARG1:.+]]: tensor<5000xi32>
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [5000] : tensor<5000xi64> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<5000xi64>
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [5000] : tensor<5000xi32> // CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<5000xi32>
// CHECK: %[[RESULT:.+]]:2 = linalg.generic // CHECK: %[[RESULT:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]
// CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] : // CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B0:.+]]: i64 // CHECK-SAME: %[[B0:.+]]: i64
// CHECK-SAME: %[[B1:.+]]: i32 // CHECK-SAME: %[[B1:.+]]: i32
// CHECK-DAG: %[[T0:.+]] = linalg.index 0 // CHECK-DAG: %[[T0:.+]] = linalg.index 0
// CHECK-DAG: %[[CAST1:.+]] = arith.index_cast %[[T0]] : index to i64 // CHECK-DAG: %[[CAST1:.+]] = arith.index_cast %[[T0]] : index to i64
Show All 10 Lines
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>) // CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)
// CHECK-SAME: outs({{.*}} : tensor<?xf32>) { // CHECK-SAME: outs({{.*}} : tensor<?xf32>) {
#map0 = affine_map<(d0) -> (d0)> #map0 = affine_map<(d0) -> (d0)>
func.func @fold_fill_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) { func.func @fold_fill_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 7.0 : f32 %cst = arith.constant 7.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?xf32>
%1 = linalg.init_tensor [%0] : tensor<?xf32> %1 = tensor.empty(%0) : tensor<?xf32>
%2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32> %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32>
%3 = linalg.init_tensor [%0] : tensor<?xf32> %3 = tensor.empty(%0) : tensor<?xf32>
%4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf32>, tensor<?xf32>) outs (%3:tensor<?xf32>) { %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf32>, tensor<?xf32>) outs (%3:tensor<?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%5 = arith.addf %arg1, %arg2 : f32 %5 = arith.addf %arg1, %arg2 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %4 : tensor<?xf32> return %4 : tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fold_fill_generic_mixedaccess // CHECK-LABEL: func @fold_fill_generic_mixedaccess
// CHECK-NOT: linalg.fill // CHECK-NOT: linalg.fill
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-NOT: ins // CHECK-NOT: ins
// CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) { // CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1, d0)> #map1 = affine_map<(d0, d1) -> (d1, d0)>
func.func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) { func.func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 0 : index %c1 = arith.constant 0 : index
%cst1 = arith.constant 7.0 : f32 %cst1 = arith.constant 7.0 : f32
%cst2 = arith.constant 6.0 : f32 %cst2 = arith.constant 6.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.fill ins(%cst1 : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32> %3 = linalg.fill ins(%cst1 : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%4 = linalg.init_tensor [%1, %0] : tensor<?x?xf32> %4 = tensor.empty(%1, %0) : tensor<?x?xf32>
%5 = linalg.fill ins(%cst2 : f32) outs(%4 : tensor<?x?xf32>) -> tensor<?x?xf32> %5 = linalg.fill ins(%cst2 : f32) outs(%4 : tensor<?x?xf32>) -> tensor<?x?xf32>
%6 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %6 = tensor.empty(%0, %1) : tensor<?x?xf32>
%7 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%3, %5 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%6:tensor<?x?xf32>) { %7 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%3, %5 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%6:tensor<?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%8 = arith.divf %arg1, %arg2 : f32 %8 = arith.divf %arg1, %arg2 : f32
linalg.yield %8 : f32 linalg.yield %8 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %7 : tensor<?x?xf32> return %7 : tensor<?x?xf32>
} }
// ----- // -----
#map = affine_map<() -> ()> #map = affine_map<() -> ()>
module { module {
func.func @fuse_multi_result_producer(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<f32>) -> tensor<f32> { func.func @fuse_multi_result_producer(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<f32>) -> tensor<f32> {
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.init_tensor [] : tensor<f32> %1 = tensor.empty() : tensor<f32>
%2:2 = linalg.generic { %2:2 = linalg.generic {
indexing_maps = [#map, #map, #map, #map, #map], iterator_types = []} indexing_maps = [#map, #map, #map, #map, #map], iterator_types = []}
ins(%arg0, %arg1, %arg1 : tensor<f32>, tensor<f32>, tensor<f32>) outs(%0, %1 : tensor<f32>, tensor<f32>) { ins(%arg0, %arg1, %arg1 : tensor<f32>, tensor<f32>, tensor<f32>) outs(%0, %1 : tensor<f32>, tensor<f32>) {
^bb0(%arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, %arg9: f32):
%4 = arith.addf %arg5, %arg6 : f32 %4 = arith.addf %arg5, %arg6 : f32
%5 = arith.addf %4, %arg7 : f32 %5 = arith.addf %4, %arg7 : f32
linalg.yield %4, %5 : f32, f32 linalg.yield %4, %5 : f32, f32
} -> (tensor<f32>, tensor<f32>) } -> (tensor<f32>, tensor<f32>)
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map, #map, #map], iterator_types = []} indexing_maps = [#map, #map, #map], iterator_types = []}
ins(%2#1, %arg1 : tensor<f32>, tensor<f32>) outs(%arg4 : tensor<f32>) { ins(%2#1, %arg1 : tensor<f32>, tensor<f32>) outs(%arg4 : tensor<f32>) {
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
%4 = arith.addf %arg5, %arg6 : f32 %4 = arith.addf %arg5, %arg6 : f32
%5 = arith.addf %4, %arg6 : f32 %5 = arith.addf %4, %arg6 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<f32> } -> tensor<f32>
return %3 : tensor<f32> return %3 : tensor<f32>
} }
} }
// CHECK-LABEL: func.func @fuse_multi_result_producer // CHECK-LABEL: func.func @fuse_multi_result_producer
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<f32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<f32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<f32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<f32>
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CHECK-SAME: outs(%[[INIT]] : // CHECK-SAME: outs(%[[INIT]] :
// CHECK-NEXT: ^bb0 // CHECK-NEXT: ^bb0
// CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B0:[a-zA-Z0-9]+]]: f32
// CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[B1:[a-zA-Z0-9]+]]: f32
// CHECK-DAG: %[[T0:.+]] = arith.addf %[[B0]], %[[B1]] // CHECK-DAG: %[[T0:.+]] = arith.addf %[[B0]], %[[B1]]
// CHECK-DAG: %[[T1:.+]] = arith.addf %[[T0]], %[[B1]] // CHECK-DAG: %[[T1:.+]] = arith.addf %[[T0]], %[[B1]]
// CHECK-DAG: %[[T2:.+]] = arith.addf %[[T1]], %[[B1]] // CHECK-DAG: %[[T2:.+]] = arith.addf %[[T1]], %[[B1]]
// CHECK-DAG: %[[T3:.+]] = arith.addf %[[T2]], %[[B1]] // CHECK-DAG: %[[T3:.+]] = arith.addf %[[T2]], %[[B1]]
// CHECK: linalg.yield %[[T3]] : f32 // CHECK: linalg.yield %[[T3]] : f32
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]

mlir/test/Dialect/Linalg/fusion-elementwise-options.mlir

Show All 12 Lines
func.func @test_fusion_limit( func.func @test_fusion_limit(
%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>, %arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>,
%arg3 : tensor<?x?xf32>, %arg4 : tensor<?x?xf32>, %arg5 : tensor<?x?xf32>) %arg3 : tensor<?x?xf32>, %arg4 : tensor<?x?xf32>, %arg5 : tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%0 = linalg.generic #binary2Dpointwise %0 = linalg.generic #binary2Dpointwise
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%init : tensor<?x?xf32>) { outs(%init : tensor<?x?xf32>) {
^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32): ^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32):
%1 = arith.mulf %arg6, %arg7 : f32 %1 = arith.mulf %arg6, %arg7 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%2 = linalg.generic #binary2Dpointwise %2 = linalg.generic #binary2Dpointwise
▲ Show 20 Lines • Show All 46 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/fusion-push-reshape.mlir

Show All 28 Lines
// ----- // -----
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-LABEL: func @reshape_multiple // CHECK-LABEL: func @reshape_multiple
// CHECK-SAME: (%[[A:.*]]: tensor<12544x16xf32>, %[[B:.*]]: tensor<12544x16xf32>, %[[C:.*]]: tensor<16xf32>) // CHECK-SAME: (%[[A:.*]]: tensor<12544x16xf32>, %[[B:.*]]: tensor<12544x16xf32>, %[[C:.*]]: tensor<16xf32>)
// CHECK: %[[I:.*]] = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32> // CHECK: %[[I:.*]] = tensor.empty() : tensor<112x112x16xf32>
// CHECK: %[[RI:.*]] = tensor.collapse_shape %[[I]] {{\[}}[0, 1], [2]] : tensor<112x112x16xf32> into tensor<12544x16xf32> // CHECK: %[[RI:.*]] = tensor.collapse_shape %[[I]] {{\[}}[0, 1], [2]] : tensor<112x112x16xf32> into tensor<12544x16xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP2]], #[[$MAP3]], #[[$MAP2]]], // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[A]], %[[B]], %[[C]] : tensor<12544x16xf32>, tensor<12544x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<12544x16xf32>) // CHECK-SAME: ins(%[[A]], %[[B]], %[[C]] : tensor<12544x16xf32>, tensor<12544x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<12544x16xf32>)
// CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<12544x16xf32> into tensor<112x112x16xf32> // CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<12544x16xf32> into tensor<112x112x16xf32>
// CHECK: return %[[RR]] : tensor<112x112x16xf32> // CHECK: return %[[RR]] : tensor<112x112x16xf32>
func.func @reshape_multiple(%A: tensor<12544x16xf32>, %B: tensor<12544x16xf32>, func.func @reshape_multiple(%A: tensor<12544x16xf32>, %B: tensor<12544x16xf32>,
%C: tensor<16xf32>) -> tensor<112x112x16xf32> { %C: tensor<16xf32>) -> tensor<112x112x16xf32> {
%0 = tensor.expand_shape %A [[0, 1], [2]] %0 = tensor.expand_shape %A [[0, 1], [2]]
: tensor<12544x16xf32> into tensor<112x112x16xf32> : tensor<12544x16xf32> into tensor<112x112x16xf32>
%1 = tensor.expand_shape %B [[0, 1], [2]] %1 = tensor.expand_shape %B [[0, 1], [2]]
: tensor<12544x16xf32> into tensor<112x112x16xf32> : tensor<12544x16xf32> into tensor<112x112x16xf32>
%2 = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32> %2 = tensor.empty() : tensor<112x112x16xf32>
%3 = linalg.generic {indexing_maps = [ %3 = linalg.generic {indexing_maps = [
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>], affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%0, %1, %C : tensor<112x112x16xf32>, tensor<112x112x16xf32>, tensor<16xf32>) ins(%0, %1, %C : tensor<112x112x16xf32>, tensor<112x112x16xf32>, tensor<16xf32>)
outs(%2 : tensor<112x112x16xf32>) { outs(%2 : tensor<112x112x16xf32>) {
Show All 11 Lines
// d0 and d1 dimensions // d0 and d1 dimensions
// CHECK-LABEL: func @reshape_negative // CHECK-LABEL: func @reshape_negative
// CHECK: tensor.expand_shape {{.*}} : tensor<12544x16xf32> into tensor<112x112x16xf32> // CHECK: tensor.expand_shape {{.*}} : tensor<12544x16xf32> into tensor<112x112x16xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK: } -> tensor<112x112x16xf32> // CHECK: } -> tensor<112x112x16xf32>
func.func @reshape_negative(%A: tensor<12544x16xf32>, %B: tensor<112xf32>) -> tensor<112x112x16xf32> { func.func @reshape_negative(%A: tensor<12544x16xf32>, %B: tensor<112xf32>) -> tensor<112x112x16xf32> {
%20 = tensor.expand_shape %A [[0, 1], [2]] %20 = tensor.expand_shape %A [[0, 1], [2]]
: tensor<12544x16xf32> into tensor<112x112x16xf32> : tensor<12544x16xf32> into tensor<112x112x16xf32>
%21 = linalg.init_tensor [112, 112, 16] : tensor<112x112x16xf32> %21 = tensor.empty() : tensor<112x112x16xf32>
%22 = linalg.generic {indexing_maps = [ %22 = linalg.generic {indexing_maps = [
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d1)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d1)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>], affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>) ins(%20, %B : tensor<112x112x16xf32>, tensor<112xf32>)
outs(%21 : tensor<112x112x16xf32>) { outs(%21 : tensor<112x112x16xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%s = arith.subf %arg1, %arg2 : f32 %s = arith.subf %arg1, %arg2 : f32
linalg.yield %s : f32 linalg.yield %s : f32
} -> tensor<112x112x16xf32> } -> tensor<112x112x16xf32>
return %22 : tensor<112x112x16xf32> return %22 : tensor<112x112x16xf32>
} }
// ----- // -----
func.func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>, func.func @type_correctness(%arg0 : tensor<6x5xi32>, %arg1 : tensor<5xf32>,
%arg2 : tensor<5xf32>) -> tensor<2x3x5xf32> { %arg2 : tensor<5xf32>) -> tensor<2x3x5xf32> {
%cst_6 = arith.constant 1.000000e+00 : f32 %cst_6 = arith.constant 1.000000e+00 : f32
%cst_7 = arith.constant 7.000000e+00 : f32 %cst_7 = arith.constant 7.000000e+00 : f32
%cst_8 = arith.constant 1.1920929E-7 : f32 %cst_8 = arith.constant 1.1920929E-7 : f32
%25 = tensor.expand_shape %arg0 [[0, 1], [2]] %25 = tensor.expand_shape %arg0 [[0, 1], [2]]
: tensor<6x5xi32> into tensor<2x3x5xi32> : tensor<6x5xi32> into tensor<2x3x5xi32>
%26 = linalg.init_tensor [2, 3, 5] : tensor<2x3x5xf32> %26 = tensor.empty() : tensor<2x3x5xf32>
%28 = linalg.generic { %28 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2)>,
affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d2)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>], affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%25, %arg1, %arg2 : tensor<2x3x5xi32>, tensor<5xf32>, tensor<5xf32>) ins(%25, %arg1, %arg2 : tensor<2x3x5xi32>, tensor<5xf32>, tensor<5xf32>)
outs(%26 : tensor<2x3x5xf32>) { outs(%26 : tensor<2x3x5xf32>) {
Show All 17 Lines

mlir/test/Dialect/Linalg/generalize-pad-tensor.mlir

// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-pad-tensor" %s | FileCheck --check-prefix=CHECK %s // RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-pad-tensor" %s | FileCheck --check-prefix=CHECK %s
// CHECK-LABEL: func @generalize_pad_tensor_static_shape( // CHECK-LABEL: func @generalize_pad_tensor_static_shape(
// CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> { // CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[INIT:.*]] = linalg.init_tensor [1, 32, 32, 1] : tensor<1x32x32x1xf32> // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<1x32x32x1xf32>
// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32> // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32>
// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32> // CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32>
// CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32> // CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32>
func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> { func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] { %0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32> } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
return %0 : tensor<1x32x32x1xf32> return %0 : tensor<1x32x32x1xf32>
} }
// CHECK-LABEL: func @generalize_pad_tensor_dynamic_shape( // CHECK-LABEL: func @generalize_pad_tensor_dynamic_shape(
// CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>, // CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>,
// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> { // CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>
// CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[OFFSET]], %[[C2]] : index // CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[OFFSET]], %[[C2]] : index
// CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>
// CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index // CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index
// CHECK: %[[INIT:.*]] = linalg.init_tensor [4, %[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]] : tensor<4x?x?x?xf32> // CHECK: %[[INIT:.*]] = tensor.empty(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : tensor<4x?x?x?xf32>
// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<4x?x?x?xf32>) -> tensor<4x?x?x?xf32> // CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<4x?x?x?xf32>) -> tensor<4x?x?x?xf32>
// CHECK: %[[DIM1_1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM1_1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>
// CHECK: %[[DIM3_1:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32> // CHECK: %[[DIM3_1:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>
// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]]{{\[}}%[[C0]], %[[C0]], %[[OFFSET]], %[[C0]]] [4, %[[DIM1_1]], 2, %[[DIM3_1]]] [1, 1, 1, 1] : tensor<4x?x2x?xf32> into tensor<4x?x?x?xf32> // CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]]{{\[}}%[[C0]], %[[C0]], %[[OFFSET]], %[[C0]]] [4, %[[DIM1_1]], 2, %[[DIM3_1]]] [1, 1, 1, 1] : tensor<4x?x2x?xf32> into tensor<4x?x?x?xf32>
// CHECK: return %[[PADDED]] : tensor<4x?x?x?xf32> // CHECK: return %[[PADDED]] : tensor<4x?x?x?xf32>
// CHECK: } // CHECK: }
func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> { func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] { %out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] {
^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index): ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32> } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>
return %out : tensor<4x?x?x?xf32> return %out : tensor<4x?x?x?xf32>
} }

mlir/test/Dialect/Linalg/inline-scalar-operands.mlir

// RUN: mlir-opt %s -linalg-inline-scalar-operands -split-input-file | FileCheck %s // RUN: mlir-opt %s -linalg-inline-scalar-operands -split-input-file | FileCheck %s
// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>
#map2 = affine_map<(d0) -> (d0)> #map2 = affine_map<(d0) -> (d0)>
#map3 = affine_map<(d0) -> ()> #map3 = affine_map<(d0) -> ()>
// CHECK: func @inline_zerod(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<f32>) // CHECK: func @inline_zerod(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<f32>)
func.func @inline_zerod(%arg0: tensor<4xf32>, %scalar: tensor<f32>) -> tensor<4xf32> { func.func @inline_zerod(%arg0: tensor<4xf32>, %scalar: tensor<f32>) -> tensor<4xf32> {
%0 = linalg.init_tensor [4] : tensor<4xf32> %0 = tensor.empty() : tensor<4xf32>
// CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]], // CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],
// CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>) // CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)
%1 = linalg.generic {indexing_maps = [#map2, #map3, #map2], %1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0, %scalar : tensor<4xf32>, tensor<f32>) ins(%arg0, %scalar : tensor<4xf32>, tensor<f32>)
outs(%0 : tensor<4xf32>) { outs(%0 : tensor<4xf32>) {
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32) // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
// CHECK: tensor.extract %[[SCALAR]][] // CHECK: tensor.extract %[[SCALAR]][]
%2 = arith.divf %arg1, %arg2 : f32 %2 = arith.divf %arg1, %arg2 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %1 : tensor<4xf32> return %1 : tensor<4xf32>
} }
// ----- // -----
// CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[MAP:.*]] = affine_map<(d0) -> (d0)>
#map2 = affine_map<(d0) -> (d0)> #map2 = affine_map<(d0) -> (d0)>
#map3 = affine_map<(d0) -> (0)> #map3 = affine_map<(d0) -> (0)>
// CHECK: func @inline_oned(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<1xf32>) // CHECK: func @inline_oned(%[[ARG:.*]]: tensor<4xf32>, %[[SCALAR:.*]]: tensor<1xf32>)
func.func @inline_oned(%arg0: tensor<4xf32>, %scalar: tensor<1xf32>) -> tensor<4xf32> { func.func @inline_oned(%arg0: tensor<4xf32>, %scalar: tensor<1xf32>) -> tensor<4xf32> {
// CHECK: %[[ZERO:.*]] = arith.constant 0 : index // CHECK: %[[ZERO:.*]] = arith.constant 0 : index
%0 = linalg.init_tensor [4] : tensor<4xf32> %0 = tensor.empty() : tensor<4xf32>
// CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]], // CHECK: linalg.generic {indexing_maps = [#[[MAP]], #[[MAP]]],
// CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>) // CHECK-SAME: iterator_types = ["parallel"]} ins(%[[ARG]] : tensor<4xf32>)
%1 = linalg.generic {indexing_maps = [#map2, #map3, #map2], %1 = linalg.generic {indexing_maps = [#map2, #map3, #map2],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0, %scalar : tensor<4xf32>, tensor<1xf32>) ins(%arg0, %scalar : tensor<4xf32>, tensor<1xf32>)
outs(%0 : tensor<4xf32>) { outs(%0 : tensor<4xf32>) {
// CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32) // CHECK: ^bb0(%{{.*}}: f32, %{{.*}}: f32)
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
// CHECK: tensor.extract %[[SCALAR]][%[[ZERO]]] // CHECK: tensor.extract %[[SCALAR]][%[[ZERO]]]
%2 = arith.divf %arg1, %arg2 : f32 %2 = arith.divf %arg1, %arg2 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %1 : tensor<4xf32> return %1 : tensor<4xf32>
} }

mlir/test/Dialect/Linalg/invalid.mlir

Show First 20 Lines • Show All 317 Lines • ▼ Show 20 Linesfunc.func @matching_inits(%m: memref<?x?xf32>, %t: tensor<?x?xf32>) {
%res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>) %res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>)
outs(%t : tensor<?x?xf32>) outs(%t : tensor<?x?xf32>)
-> tensor<?xf32> -> tensor<?xf32>
return return
} }
// ----- // -----
func.func @init_tensor_err(%arg0 : index, %arg1 : index)
{
// expected-error @+1 {{specified type 'tensor<4x?x?x5xf32>' does not match the inferred type 'tensor<4x5x?x?xf32>'}}
%1 = linalg.init_tensor [4, 5, %arg0, %arg1] : tensor<4x?x?x5xf32>
return
}
// -----
func.func @init_tensor_err(%arg0 : index)
{
// expected-error @+1 {{expected 4 sizes values}}
%1 = linalg.init_tensor [4, 5, %arg0] : tensor<4x?x?x5xf32>
return
}
// -----
func.func @init_tensor_err(%arg0 : index)
{
// expected-error @+1 {{expected 2 dynamic sizes values}}
%1 = "linalg.init_tensor"(%arg0) {static_sizes = [4, -1, -1, 5]} : (index) -> tensor<4x?x?x5xf32>
return
}
// -----
func.func @illegal_fill_tensor_no_return(%arg0 : index, %arg1 : index, %arg2 : f32) func.func @illegal_fill_tensor_no_return(%arg0 : index, %arg1 : index, %arg2 : f32)
{ {
%0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32> %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
// expected-error @+1 {{expected the number of results (0) to be equal to the number of output tensors (1)}} // expected-error @+1 {{expected the number of results (0) to be equal to the number of output tensors (1)}}
linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>) linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>)
} }
// ----- // -----
func.func @illegal_fill_memref_with_tensor_return func.func @illegal_fill_memref_with_tensor_return
(%arg0 : memref<?x?xf32>, %arg1 : f32) -> tensor<?x?xf32> (%arg0 : memref<?x?xf32>, %arg1 : f32) -> tensor<?x?xf32>
▲ Show 20 Lines • Show All 226 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/lower-pad-tensor.mlir

Show First 20 Lines • Show All 46 Lines • ▼ Show 20 Linesfunc.func @pad_tensor_detailed(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32> } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
return %0 : tensor<1x32x32x1xf32> return %0 : tensor<1x32x32x1xf32>
} }
// CHECK: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> // CHECK: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32>
// CHECK: %[[CTE:.+]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[CTE:.+]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[TMP:.+]] = linalg.init_tensor [1, 32, 32, 1] : tensor<1x32x32x1xf32> // CHECK: %[[TMP:.+]] = tensor.empty() : tensor<1x32x32x1xf32>
// CHECK: %[[R1c:.+]] = linalg.fill // CHECK: %[[R1c:.+]] = linalg.fill
// CHECK: %[[R2c:.+]] = linalg.generic // CHECK: %[[R2c:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP4]], #[[$MAP5]]] // CHECK-SAME: indexing_maps = [#[[$MAP4]], #[[$MAP5]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]
// CHECK: ins(%arg0 : tensor<1x28x28x1xf32>) outs(%1 : tensor<1x32x32x1xf32>) // CHECK: ins(%arg0 : tensor<1x28x28x1xf32>) outs(%1 : tensor<1x32x32x1xf32>)
// CHECK: ^bb0(%[[VAL:.+]]: f32, %arg2: f32) // CHECK: ^bb0(%[[VAL:.+]]: f32, %arg2: f32)
// CHECK: linalg.yield %[[VAL]] : f32 // CHECK: linalg.yield %[[VAL]] : f32
// CHECK: return %[[R2c:.+]] // CHECK: return %[[R2c:.+]]

mlir/test/Dialect/Linalg/named-ops.mlir

// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s // RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
// CHECK-LABEL: func @depthwise_conv_1d_nwc_wcm // CHECK-LABEL: func @depthwise_conv_1d_nwc_wcm
func.func @depthwise_conv_1d_nwc_wcm(%input: tensor<1x12x8xf32>, %filter: tensor<3x8x8xf32>) -> tensor<1x10x8x8xf32> { func.func @depthwise_conv_1d_nwc_wcm(%input: tensor<1x12x8xf32>, %filter: tensor<3x8x8xf32>) -> tensor<1x10x8x8xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [1, 10, 8, 8] : tensor<1x10x8x8xf32> %init = tensor.empty() : tensor<1x10x8x8xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32>
// CHECK: depthwise_conv_1d_nwc_wcm // CHECK: depthwise_conv_1d_nwc_wcm
%0 = linalg.depthwise_conv_1d_nwc_wcm {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} %0 = linalg.depthwise_conv_1d_nwc_wcm {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8x8xf32>) ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8x8xf32>)
outs(%fill : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32> outs(%fill : tensor<1x10x8x8xf32>) -> tensor<1x10x8x8xf32>
return %0 : tensor<1x10x8x8xf32> return %0 : tensor<1x10x8x8xf32>
} }
// ----- // -----
// CHECK-LABEL: func @depthwise_conv_1d_nwc_wc // CHECK-LABEL: func @depthwise_conv_1d_nwc_wc
func.func @depthwise_conv_1d_nwc_wc(%input: tensor<1x12x8xf32>, %filter: tensor<3x8xf32>) -> tensor<1x10x8xf32> { func.func @depthwise_conv_1d_nwc_wc(%input: tensor<1x12x8xf32>, %filter: tensor<3x8xf32>) -> tensor<1x10x8xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [1, 10, 8] : tensor<1x10x8xf32> %init = tensor.empty() : tensor<1x10x8xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8xf32>) -> tensor<1x10x8xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<1x10x8xf32>) -> tensor<1x10x8xf32>
// CHECK: depthwise_conv_1d_nwc_wc // CHECK: depthwise_conv_1d_nwc_wc
%0 = linalg.depthwise_conv_1d_nwc_wc {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} %0 = linalg.depthwise_conv_1d_nwc_wc {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8xf32>) ins(%input, %filter : tensor<1x12x8xf32>, tensor<3x8xf32>)
outs(%fill : tensor<1x10x8xf32>) -> tensor<1x10x8xf32> outs(%fill : tensor<1x10x8xf32>) -> tensor<1x10x8xf32>
return %0 : tensor<1x10x8xf32> return %0 : tensor<1x10x8xf32>
} }
// ----- // -----
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_tensor // CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_tensor
func.func @depthwise_conv_2d_nhwc_hwcm_tensor(%input: tensor<2x4x5x2xf32>, %filter: tensor<2x2x2x3xf32>) -> tensor<2x3x4x2x3xf32> { func.func @depthwise_conv_2d_nhwc_hwcm_tensor(%input: tensor<2x4x5x2xf32>, %filter: tensor<2x2x2x3xf32>) -> tensor<2x3x4x2x3xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [2, 3, 4, 2, 3] : tensor<2x3x4x2x3xf32> %init = tensor.empty() : tensor<2x3x4x2x3xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32>
// CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwcm // CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwcm
// CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} // CHECK-SAME: {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<2x3x4x2x3xf32>) // CHECK-SAME: outs(%{{.+}} : tensor<2x3x4x2x3xf32>)
%0 = linalg.depthwise_conv_2d_nhwc_hwcm %0 = linalg.depthwise_conv_2d_nhwc_hwcm
{ dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
ins(%input, %filter : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>) ins(%input, %filter : tensor<2x4x5x2xf32>, tensor<2x2x2x3xf32>)
Show All 11 Lines linalg.depthwise_conv_2d_nhwc_hwcm
{ dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>) ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
outs(%output : memref<2x3x4x2x3xf32>) outs(%output : memref<2x3x4x2x3xf32>)
return return
} }
// CHECK-LABEL: func @depthwise_conv_1d_nw_tensor // CHECK-LABEL: func @depthwise_conv_1d_nw_tensor
func.func @depthwise_conv_1d_nw_tensor(%input: tensor<1x113x96xf32>, %filter: tensor<3x96xf32>) -> tensor<1x56x96xf32> { func.func @depthwise_conv_1d_nw_tensor(%input: tensor<1x113x96xf32>, %filter: tensor<3x96xf32>) -> tensor<1x56x96xf32> {
%init = linalg.init_tensor [1, 56, 96] : tensor<1x56x96xf32> %init = tensor.empty() : tensor<1x56x96xf32>
// CHECK: %{{.+}} = linalg.depthwise_conv_1d_nw // CHECK: %{{.+}} = linalg.depthwise_conv_1d_nw
// CHECK-SAME: {dilations = dense<1> : vector<1xi64>, strides = dense<2> : vector<1xi64>} // CHECK-SAME: {dilations = dense<1> : vector<1xi64>, strides = dense<2> : vector<1xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x113x96xf32>, tensor<3x96xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x113x96xf32>, tensor<3x96xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x56x96xf32>) -> tensor<1x56x96xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x56x96xf32>) -> tensor<1x56x96xf32>
%0 = linalg.depthwise_conv_1d_nwc_wc {dilations = dense<1> : vector<1xi64>, strides = dense<2> : vector<1xi64>} %0 = linalg.depthwise_conv_1d_nwc_wc {dilations = dense<1> : vector<1xi64>, strides = dense<2> : vector<1xi64>}
ins(%input, %filter: tensor<1x113x96xf32>, tensor<3x96xf32>) ins(%input, %filter: tensor<1x113x96xf32>, tensor<3x96xf32>)
outs(%init: tensor<1x56x96xf32>) -> tensor<1x56x96xf32> outs(%init: tensor<1x56x96xf32>) -> tensor<1x56x96xf32>
return %0: tensor<1x56x96xf32> return %0: tensor<1x56x96xf32>
} }
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_tensor // CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_tensor
func.func @depthwise_conv_2d_nhwc_hwc_tensor(%input: tensor<1x113x113x96xf32>, %filter: tensor<3x3x96xf32>) -> tensor<1x56x56x96xf32> { func.func @depthwise_conv_2d_nhwc_hwc_tensor(%input: tensor<1x113x113x96xf32>, %filter: tensor<3x3x96xf32>) -> tensor<1x56x56x96xf32> {
%init = linalg.init_tensor [1, 56, 56, 96] : tensor<1x56x56x96xf32> %init = tensor.empty() : tensor<1x56x56x96xf32>
// CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwc // CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwc
// CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} // CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x113x113x96xf32>, tensor<3x3x96xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x113x113x96xf32>, tensor<3x3x96xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
%0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} %0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: tensor<1x113x113x96xf32>, tensor<3x3x96xf32>) ins(%input, %filter: tensor<1x113x113x96xf32>, tensor<3x3x96xf32>)
outs(%init: tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32> outs(%init: tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
return %0: tensor<1x56x56x96xf32> return %0: tensor<1x56x56x96xf32>
} }
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_memref // CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_memref
func.func @depthwise_conv_2d_nhwc_hwc_memref(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) { func.func @depthwise_conv_2d_nhwc_hwc_memref(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
// CHECK: linalg.depthwise_conv_2d_nhwc_hwc // CHECK: linalg.depthwise_conv_2d_nhwc_hwc
// CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} // CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x56x56x96xf32>) // CHECK-SAME: outs(%{{.+}} : memref<1x56x56x96xf32>)
linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>) ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
outs(%output: memref<1x56x56x96xf32>) outs(%output: memref<1x56x56x96xf32>)
return return
} }
// CHECK-LABEL: func @depthwise_conv_2d_nchw_chw_tensor // CHECK-LABEL: func @depthwise_conv_2d_nchw_chw_tensor
func.func @depthwise_conv_2d_nchw_chw_tensor(%input: tensor<1x96x113x113xf32>, %filter: tensor<96x3x3xf32>) -> tensor<1x96x56x56xf32> { func.func @depthwise_conv_2d_nchw_chw_tensor(%input: tensor<1x96x113x113xf32>, %filter: tensor<96x3x3xf32>) -> tensor<1x96x56x56xf32> {
%init = linalg.init_tensor [1, 96, 56, 56] : tensor<1x96x56x56xf32> %init = tensor.empty() : tensor<1x96x56x56xf32>
// CHECK: %{{.+}} = linalg.depthwise_conv_2d_nchw_chw // CHECK: %{{.+}} = linalg.depthwise_conv_2d_nchw_chw
// CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} // CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x96x113x113xf32>, tensor<96x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x96x113x113xf32>, tensor<96x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x96x56x56xf32>) -> tensor<1x96x56x56xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x96x56x56xf32>) -> tensor<1x96x56x56xf32>
%0 = linalg.depthwise_conv_2d_nchw_chw {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} %0 = linalg.depthwise_conv_2d_nchw_chw {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: tensor<1x96x113x113xf32>, tensor<96x3x3xf32>) ins(%input, %filter: tensor<1x96x113x113xf32>, tensor<96x3x3xf32>)
outs(%init: tensor<1x96x56x56xf32>) -> tensor<1x96x56x56xf32> outs(%init: tensor<1x96x56x56xf32>) -> tensor<1x96x56x56xf32>
return %0: tensor<1x96x56x56xf32> return %0: tensor<1x96x56x56xf32>
} }
// CHECK-LABEL: func @depthwise_conv_2d_nchw_chw_memref // CHECK-LABEL: func @depthwise_conv_2d_nchw_chw_memref
func.func @depthwise_conv_2d_nchw_chw_memref(%input: memref<1x96x113x113xf32>, %filter: memref<96x3x3xf32>, %output: memref<1x96x56x56xf32>) { func.func @depthwise_conv_2d_nchw_chw_memref(%input: memref<1x96x113x113xf32>, %filter: memref<96x3x3xf32>, %output: memref<1x96x56x56xf32>) {
// CHECK: linalg.depthwise_conv_2d_nchw_chw // CHECK: linalg.depthwise_conv_2d_nchw_chw
// CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} // CHECK-SAME: {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x96x113x113xf32>, memref<96x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x96x113x113xf32>, memref<96x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x96x56x56xf32>) // CHECK-SAME: outs(%{{.+}} : memref<1x96x56x56xf32>)
linalg.depthwise_conv_2d_nchw_chw {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} linalg.depthwise_conv_2d_nchw_chw {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: memref<1x96x113x113xf32>, memref<96x3x3xf32>) ins(%input, %filter: memref<1x96x113x113xf32>, memref<96x3x3xf32>)
outs(%output: memref<1x96x56x56xf32>) outs(%output: memref<1x96x56x56xf32>)
return return
} }
func.func @depthwise_conv_2d_nhwc_hwcm_tensor_dilated(%input: tensor<2x8x9x2xf32>, %filter: tensor<2x2x2x3xf32>) -> tensor<2x6x7x2x3xf32> { func.func @depthwise_conv_2d_nhwc_hwcm_tensor_dilated(%input: tensor<2x8x9x2xf32>, %filter: tensor<2x2x2x3xf32>) -> tensor<2x6x7x2x3xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [2, 6, 7, 2, 3] : tensor<2x6x7x2x3xf32> %init = tensor.empty() : tensor<2x6x7x2x3xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x6x7x2x3xf32>) -> tensor<2x6x7x2x3xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x6x7x2x3xf32>) -> tensor<2x6x7x2x3xf32>
// CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwcm // CHECK: %{{.+}} = linalg.depthwise_conv_2d_nhwc_hwcm
// CHECK-SAME: {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} // CHECK-SAME: {dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x8x9x2xf32>, tensor<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<2x8x9x2xf32>, tensor<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<2x6x7x2x3xf32>) // CHECK-SAME: outs(%{{.+}} : tensor<2x6x7x2x3xf32>)
%0 = linalg.depthwise_conv_2d_nhwc_hwcm %0 = linalg.depthwise_conv_2d_nhwc_hwcm
{ dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> } { dilations = dense<2> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
ins(%input, %filter : tensor<2x8x9x2xf32>, tensor<2x2x2x3xf32>) ins(%input, %filter : tensor<2x8x9x2xf32>, tensor<2x2x2x3xf32>)
▲ Show 20 Lines • Show All 47 Lines • ▼ Show 20 Linesfunc.func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_size(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
return return
} }
// ----- // -----
// CHECK-LABEL: func @depthwise_conv_3d_ndhwc_dhwcm // CHECK-LABEL: func @depthwise_conv_3d_ndhwc_dhwcm
func.func @depthwise_conv_3d_ndhwc_dhwcm(%input: tensor<2x6x13x12x6xf32>, %filter: tensor<2x1x3x6x6xf32>) -> tensor<2x3x13x4x6x6xf32> { func.func @depthwise_conv_3d_ndhwc_dhwcm(%input: tensor<2x6x13x12x6xf32>, %filter: tensor<2x1x3x6x6xf32>) -> tensor<2x3x13x4x6x6xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [2, 3, 13, 4, 6, 6] : tensor<2x3x13x4x6x6xf32> %init = tensor.empty() : tensor<2x3x13x4x6x6xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x13x4x6x6xf32>) -> tensor<2x3x13x4x6x6xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x13x4x6x6xf32>) -> tensor<2x3x13x4x6x6xf32>
// CHECK: depthwise_conv_3d_ndhwc_dhwcm // CHECK: depthwise_conv_3d_ndhwc_dhwcm
%0 = linalg.depthwise_conv_3d_ndhwc_dhwcm {dilations = dense<1> : tensor<3xi64>, strides = dense<[2, 1, 3]> : tensor<3xi64>} %0 = linalg.depthwise_conv_3d_ndhwc_dhwcm {dilations = dense<1> : tensor<3xi64>, strides = dense<[2, 1, 3]> : tensor<3xi64>}
ins(%input, %filter : tensor<2x6x13x12x6xf32>, tensor<2x1x3x6x6xf32>) ins(%input, %filter : tensor<2x6x13x12x6xf32>, tensor<2x1x3x6x6xf32>)
outs(%fill : tensor<2x3x13x4x6x6xf32>) -> tensor<2x3x13x4x6x6xf32> outs(%fill : tensor<2x3x13x4x6x6xf32>) -> tensor<2x3x13x4x6x6xf32>
return %0 : tensor<2x3x13x4x6x6xf32> return %0 : tensor<2x3x13x4x6x6xf32>
} }
// ----- // -----
// CHECK-LABEL: func @depthwise_conv_3d_ndhwc_dhwc // CHECK-LABEL: func @depthwise_conv_3d_ndhwc_dhwc
func.func @depthwise_conv_3d_ndhwc_dhwc(%input: tensor<2x6x13x12x6xf32>, %filter: tensor<2x1x3x6xf32>) -> tensor<2x3x13x4x6xf32> { func.func @depthwise_conv_3d_ndhwc_dhwc(%input: tensor<2x6x13x12x6xf32>, %filter: tensor<2x1x3x6xf32>) -> tensor<2x3x13x4x6xf32> {
%zero = arith.constant 0.000000e+00 : f32 %zero = arith.constant 0.000000e+00 : f32
%init = linalg.init_tensor [2, 3, 13, 4, 6] : tensor<2x3x13x4x6xf32> %init = tensor.empty() : tensor<2x3x13x4x6xf32>
%fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x13x4x6xf32>) -> tensor<2x3x13x4x6xf32> %fill = linalg.fill ins(%zero : f32) outs(%init : tensor<2x3x13x4x6xf32>) -> tensor<2x3x13x4x6xf32>
// CHECK: depthwise_conv_3d_ndhwc_dhwc // CHECK: depthwise_conv_3d_ndhwc_dhwc
%0 = linalg.depthwise_conv_3d_ndhwc_dhwc {dilations = dense<1> : tensor<3xi64>, strides = dense<[2, 1, 3]> : tensor<3xi64>} %0 = linalg.depthwise_conv_3d_ndhwc_dhwc {dilations = dense<1> : tensor<3xi64>, strides = dense<[2, 1, 3]> : tensor<3xi64>}
ins(%input, %filter : tensor<2x6x13x12x6xf32>, tensor<2x1x3x6xf32>) ins(%input, %filter : tensor<2x6x13x12x6xf32>, tensor<2x1x3x6xf32>)
outs(%fill : tensor<2x3x13x4x6xf32>) -> tensor<2x3x13x4x6xf32> outs(%fill : tensor<2x3x13x4x6xf32>) -> tensor<2x3x13x4x6xf32>
return %0 : tensor<2x3x13x4x6xf32> return %0 : tensor<2x3x13x4x6xf32>
} }
▲ Show 20 Lines • Show All 193 Lines • ▼ Show 20 Lines
// CHECK-LABEL: func @pooling_nhwc_sum_tensor // CHECK-LABEL: func @pooling_nhwc_sum_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_sum // CHECK: %{{.+}} = linalg.pooling_nhwc_sum
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func.func @pooling_nhwc_sum_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> { func.func @pooling_nhwc_sum_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>) ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32> return %res : tensor<1x2x2x1xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_nchw_sum_tensor // CHECK-LABEL: func @pooling_nchw_sum_tensor
// CHECK: %{{.+}} = linalg.pooling_nchw_sum // CHECK: %{{.+}} = linalg.pooling_nchw_sum
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x1x4x4xf32>, tensor<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x1x4x4xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
func.func @pooling_nchw_sum_tensor(%input: tensor<1x1x4x4xf32>) -> tensor<1x1x2x2xf32> { func.func @pooling_nchw_sum_tensor(%input: tensor<1x1x4x4xf32>) -> tensor<1x1x2x2xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = linalg.init_tensor [1, 1, 2, 2] : tensor<1x1x2x2xf32> %init = tensor.empty() : tensor<1x1x2x2xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
%res = linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x1x4x4xf32>, tensor<3x3xf32>) ins(%input, %fake: tensor<1x1x4x4xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> outs(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
return %res : tensor<1x1x2x2xf32> return %res : tensor<1x1x2x2xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_nhwc_max_tensor // CHECK-LABEL: func @pooling_nhwc_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max // CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func.func @pooling_nhwc_max_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> { func.func @pooling_nhwc_max_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>) ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32> return %res : tensor<1x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nchw_max_tensor // CHECK-LABEL: func @pooling_nchw_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nchw_max // CHECK: %{{.+}} = linalg.pooling_nchw_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x1x4x4xf32>, tensor<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x1x4x4xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
func.func @pooling_nchw_max_tensor(%input: tensor<1x1x4x4xf32>) -> tensor<1x1x2x2xf32> { func.func @pooling_nchw_max_tensor(%input: tensor<1x1x4x4xf32>) -> tensor<1x1x2x2xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = linalg.init_tensor [1, 1, 2, 2] : tensor<1x1x2x2xf32> %init = tensor.empty() : tensor<1x1x2x2xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
%res = linalg.pooling_nchw_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nchw_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x1x4x4xf32>, tensor<3x3xf32>) ins(%input, %fake: tensor<1x1x4x4xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32> outs(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
return %res : tensor<1x1x2x2xf32> return %res : tensor<1x1x2x2xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_nhwc_i8_max_tensor // CHECK-LABEL: func @pooling_nhwc_i8_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max // CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi8>, tensor<3x3xi8>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi8>, tensor<3x3xi8>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
func.func @pooling_nhwc_i8_max_tensor(%input: tensor<1x4x4x1xi8>) -> tensor<1x2x2x1xi8> { func.func @pooling_nhwc_i8_max_tensor(%input: tensor<1x4x4x1xi8>) -> tensor<1x2x2x1xi8> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xi8> %fake = tensor.empty() : tensor<3x3xi8>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi8> %init = tensor.empty() : tensor<1x2x2x1xi8>
%cst = arith.constant 0 : i8 %cst = arith.constant 0 : i8
%fill = linalg.fill ins(%cst : i8) outs(%init : tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8> %fill = linalg.fill ins(%cst : i8) outs(%init : tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
%res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xi8>, tensor<3x3xi8>) ins(%input, %fake: tensor<1x4x4x1xi8>, tensor<3x3xi8>)
outs(%fill: tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8> outs(%fill: tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
return %res : tensor<1x2x2x1xi8> return %res : tensor<1x2x2x1xi8>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_nhwc_i16_max_tensor // CHECK-LABEL: func @pooling_nhwc_i16_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max // CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi16>, tensor<3x3xi16>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi16>, tensor<3x3xi16>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
func.func @pooling_nhwc_i16_max_tensor(%input: tensor<1x4x4x1xi16>) -> tensor<1x2x2x1xi16> { func.func @pooling_nhwc_i16_max_tensor(%input: tensor<1x4x4x1xi16>) -> tensor<1x2x2x1xi16> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xi16> %fake = tensor.empty() : tensor<3x3xi16>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi16> %init = tensor.empty() : tensor<1x2x2x1xi16>
%cst = arith.constant 0 : i16 %cst = arith.constant 0 : i16
%fill = linalg.fill ins(%cst : i16) outs(%init : tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16> %fill = linalg.fill ins(%cst : i16) outs(%init : tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
%res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xi16>, tensor<3x3xi16>) ins(%input, %fake: tensor<1x4x4x1xi16>, tensor<3x3xi16>)
outs(%fill: tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16> outs(%fill: tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
return %res : tensor<1x2x2x1xi16> return %res : tensor<1x2x2x1xi16>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_nhwc_i32_max_tensor // CHECK-LABEL: func @pooling_nhwc_i32_max_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_max // CHECK: %{{.+}} = linalg.pooling_nhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi32>, tensor<3x3xi32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi32>, tensor<3x3xi32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
func.func @pooling_nhwc_i32_max_tensor(%input: tensor<1x4x4x1xi32>) -> tensor<1x2x2x1xi32> { func.func @pooling_nhwc_i32_max_tensor(%input: tensor<1x4x4x1xi32>) -> tensor<1x2x2x1xi32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xi32> %fake = tensor.empty() : tensor<3x3xi32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi32> %init = tensor.empty() : tensor<1x2x2x1xi32>
%cst = arith.constant 0 : i32 %cst = arith.constant 0 : i32
%fill = linalg.fill ins(%cst : i32) outs(%init : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32> %fill = linalg.fill ins(%cst : i32) outs(%init : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
%res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xi32>, tensor<3x3xi32>) ins(%input, %fake: tensor<1x4x4x1xi32>, tensor<3x3xi32>)
outs(%fill: tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32> outs(%fill: tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
return %res : tensor<1x2x2x1xi32> return %res : tensor<1x2x2x1xi32>
} }
Show All 17 Lines
// CHECK-LABEL: func @pooling_nhwc_min_tensor // CHECK-LABEL: func @pooling_nhwc_min_tensor
// CHECK: %{{.+}} = linalg.pooling_nhwc_min // CHECK: %{{.+}} = linalg.pooling_nhwc_min
// CHECK-SAME: dilations = dense<1> : tensor<2xi64> // CHECK-SAME: dilations = dense<1> : tensor<2xi64>
// CHECK-SAME: strides = dense<1> : tensor<2xi64> // CHECK-SAME: strides = dense<1> : tensor<2xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xf32>, tensor<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
func.func @pooling_nhwc_min_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> { func.func @pooling_nhwc_min_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3] : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
%res = linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} %res = linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>) ins(%input, %fake: tensor<1x4x4x1xf32>, tensor<3x3xf32>)
outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32> outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32> return %res : tensor<1x2x2x1xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_ndhwc_sum_tensor // CHECK-LABEL: func @pooling_ndhwc_sum_tensor
// CHECK: %{{.+}} = linalg.pooling_ndhwc_sum // CHECK: %{{.+}} = linalg.pooling_ndhwc_sum
// CHECK-SAME: dilations = dense<1> : tensor<3xi64> // CHECK-SAME: dilations = dense<1> : tensor<3xi64>
// CHECK-SAME: strides = dense<1> : tensor<3xi64> // CHECK-SAME: strides = dense<1> : tensor<3xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
func.func @pooling_ndhwc_sum_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> { func.func @pooling_ndhwc_sum_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3, 3] : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 2, 1] : tensor<1x2x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
%res = linalg.pooling_ndhwc_sum {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} %res = linalg.pooling_ndhwc_sum {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_ndhwc_max_tensor // CHECK-LABEL: func @pooling_ndhwc_max_tensor
// CHECK: %{{.+}} = linalg.pooling_ndhwc_max // CHECK: %{{.+}} = linalg.pooling_ndhwc_max
// CHECK-SAME: dilations = dense<1> : tensor<3xi64> // CHECK-SAME: dilations = dense<1> : tensor<3xi64>
// CHECK-SAME: strides = dense<1> : tensor<3xi64> // CHECK-SAME: strides = dense<1> : tensor<3xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
func.func @pooling_ndhwc_max_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> { func.func @pooling_ndhwc_max_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3, 3] : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 2, 1] : tensor<1x2x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
%res = linalg.pooling_ndhwc_max {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} %res = linalg.pooling_ndhwc_max {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
Show All 16 Lines
// CHECK-LABEL: func @pooling_ndhwc_min_tensor // CHECK-LABEL: func @pooling_ndhwc_min_tensor
// CHECK: %{{.+}} = linalg.pooling_ndhwc_min // CHECK: %{{.+}} = linalg.pooling_ndhwc_min
// CHECK-SAME: dilations = dense<1> : tensor<3xi64> // CHECK-SAME: dilations = dense<1> : tensor<3xi64>
// CHECK-SAME: strides = dense<1> : tensor<3xi64> // CHECK-SAME: strides = dense<1> : tensor<3xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> // CHECK-SAME: outs(%{{.+}} : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
func.func @pooling_ndhwc_min_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> { func.func @pooling_ndhwc_min_tensor(%input: tensor<1x4x4x4x1xf32>) -> tensor<1x2x2x2x1xf32> {
%fake = linalg.init_tensor [3, 3, 3] : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = linalg.init_tensor [1, 2, 2, 2, 1] : tensor<1x2x2x2x1xf32> %init = tensor.empty() : tensor<1x2x2x2x1xf32>
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
%res = linalg.pooling_ndhwc_min {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} %res = linalg.pooling_ndhwc_min {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>) ins(%input, %fake: tensor<1x4x4x4x1xf32>, tensor<3x3x3xf32>)
outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32> outs(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
▲ Show 20 Lines • Show All 68 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/one-shot-bufferize-analysis-init-tensor-elimination.mlir

// RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries test-analysis-only allow-return-allocs" -split-input-file | FileCheck %s // RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries test-analysis-only allow-return-allocs" -split-input-file | FileCheck %s
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// InitTensorOp elimination // AllocTensorOp elimination
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// CHECK-LABEL: func @buffer_forwarding_conflict // CHECK-LABEL: func @buffer_forwarding_conflict
func.func @buffer_forwarding_conflict(%arg0: tensor<?xf32> {bufferization.writable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) { func.func @buffer_forwarding_conflict(%arg0: tensor<?xf32> {bufferization.writable = true}, %arg1: index) -> (tensor<?xf32>, tensor<?xf32>) {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["false", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["false", "none"]
// Instead of allocating, share buffer with some inplace bufferization? // Instead of allocating, share buffer with some inplace bufferization?
▲ Show 20 Lines • Show All 41 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/pad_fusion.mlir

// RUN: mlir-opt -test-linalg-pad-fusion -split-input-file %s | FileCheck %s // RUN: mlir-opt -test-linalg-pad-fusion -split-input-file %s | FileCheck %s
func.func @dynamic_pad_fusion(%arg0 : tensor<?x?xf32>, %arg1 : index, %arg2 : index, func.func @dynamic_pad_fusion(%arg0 : tensor<?x?xf32>, %arg1 : index, %arg2 : index,
%arg3 : index, %arg4 : index, %arg5 : f32) -> tensor<?x?xf32> { %arg3 : index, %arg4 : index, %arg5 : f32) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?x?xf32>) outs(%init : tensor<?x?xf32>) { ins(%arg0 : tensor<?x?xf32>) outs(%init : tensor<?x?xf32>) {
^bb0(%arg6 : f32, %arg7 : f32): ^bb0(%arg6 : f32, %arg7 : f32):
%1 = arith.mulf %arg6, %arg6 : f32 %1 = arith.mulf %arg6, %arg6 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
Show All 14 Lines
// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: f32
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[SOURCE:.+]] = linalg.generic // CHECK-DAG: %[[SOURCE:.+]] = linalg.generic
// CHECK-DAG: %[[SOURCE_D0:.+]] = tensor.dim %[[SOURCE]], %[[C0]] // CHECK-DAG: %[[SOURCE_D0:.+]] = tensor.dim %[[SOURCE]], %[[C0]]
// CHECK-DAG: %[[TARGET_D0:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG3]], %[[SOURCE_D0]]] // CHECK-DAG: %[[TARGET_D0:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG3]], %[[SOURCE_D0]]]
// CHECK-DAG: %[[SOURCE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]] // CHECK-DAG: %[[SOURCE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]]
// CHECK-DAG: %[[TARGET_D1:.+]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG4]], %[[SOURCE_D1]]] // CHECK-DAG: %[[TARGET_D1:.+]] = affine.apply #[[MAP]]()[%[[ARG2]], %[[ARG4]], %[[SOURCE_D1]]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[TARGET_D0]], %[[TARGET_D1]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[TARGET_D0]], %[[TARGET_D1]])
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG5]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG5]]{{.*}}outs(%[[INIT]]
// CHECK-DAG: %[[SIZE_D0:.+]] = tensor.dim %[[SOURCE]], %[[C0]] // CHECK-DAG: %[[SIZE_D0:.+]] = tensor.dim %[[SOURCE]], %[[C0]]
// CHECK-DAG: %[[SIZE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]] // CHECK-DAG: %[[SIZE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]]
// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]] // CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]]
// CHECK-SAME: [%[[ARG1]], %[[ARG2]]] [%[[SIZE_D0]], %[[SIZE_D1]]] [1, 1] // CHECK-SAME: [%[[ARG1]], %[[ARG2]]] [%[[SIZE_D0]], %[[SIZE_D1]]] [1, 1]
// CHECK: %[[SOURCE:.+]] = linalg.generic // CHECK: %[[SOURCE:.+]] = linalg.generic
// CHECK-SAME: outs(%[[SLICE]] : tensor<?x?xf32>) // CHECK-SAME: outs(%[[SLICE]] : tensor<?x?xf32>)
// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SOURCE]] into %[[FILL]] // CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SOURCE]] into %[[FILL]]
// CHECK-SAME: [%[[ARG1]], %[[ARG2]]] [%[[SIZE_D0]], %[[SIZE_D1]]] [1, 1] // CHECK-SAME: [%[[ARG1]], %[[ARG2]]] [%[[SIZE_D0]], %[[SIZE_D1]]] [1, 1]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
func.func @mixed_pad_fusion(%arg0 : tensor<?x42xf32>, %arg1 : index, %arg2 : index, func.func @mixed_pad_fusion(%arg0 : tensor<?x42xf32>, %arg1 : index, %arg2 : index,
%arg3 : f32) -> tensor<49x?xf32> { %arg3 : f32) -> tensor<49x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x42xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x42xf32>
%init = linalg.init_tensor [42, %d0] : tensor<42x?xf32> %init = tensor.empty(%d0) : tensor<42x?xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?x42xf32>) outs(%init : tensor<42x?xf32>) { ins(%arg0 : tensor<?x42xf32>) outs(%init : tensor<42x?xf32>) {
^bb0(%arg4 : f32, %arg5 : f32): ^bb0(%arg4 : f32, %arg5 : f32):
%1 = arith.mulf %arg4, %arg4 : f32 %1 = arith.mulf %arg4, %arg4 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<42x?xf32> } -> tensor<42x?xf32>
Show All 9 Lines
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: f32 // CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: f32
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[SOURCE:.+]] = linalg.generic // CHECK-DAG: %[[SOURCE:.+]] = linalg.generic
// CHECK-DAG: %[[SOURCE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]] // CHECK-DAG: %[[SOURCE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]]
// CHECK-DAG: %[[TARGET_D1:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG2]], %[[SOURCE_D1]]] // CHECK-DAG: %[[TARGET_D1:.+]] = affine.apply #[[MAP]]()[%[[ARG1]], %[[ARG2]], %[[SOURCE_D1]]]
// CHECK: %[[INIT:.+]] = linalg.init_tensor [49, %[[TARGET_D1]]] // CHECK: %[[INIT:.+]] = tensor.empty(%[[TARGET_D1]])
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG3]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG3]]{{.*}}outs(%[[INIT]]
// CHECK-DAG: %[[SIZE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]] // CHECK-DAG: %[[SIZE_D1:.+]] = tensor.dim %[[SOURCE]], %[[C1]]
// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]] // CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[FILL]]
// CHECK-SAME: [3, %[[ARG1]]] [42, %[[SIZE_D1]]] [1, 1] // CHECK-SAME: [3, %[[ARG1]]] [42, %[[SIZE_D1]]] [1, 1]
// CHECK: %[[SOURCE:.+]] = linalg.generic // CHECK: %[[SOURCE:.+]] = linalg.generic
// CHECK-SAME: outs(%[[SLICE]] : tensor<42x?xf32>) // CHECK-SAME: outs(%[[SLICE]] : tensor<42x?xf32>)
// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SOURCE]] into %[[FILL]] // CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[SOURCE]] into %[[FILL]]
// CHECK-SAME: [3, %[[ARG1]]] [42, %[[SIZE_D1]]] [1, 1] // CHECK-SAME: [3, %[[ARG1]]] [42, %[[SIZE_D1]]] [1, 1]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]

mlir/test/Dialect/Linalg/reshape_control_fusion.mlir

// RUN: mlir-opt -test-linalg-elementwise-fusion-patterns=control-fusion-by-expansion %s -split-input-file | FileCheck %s // RUN: mlir-opt -test-linalg-elementwise-fusion-patterns=control-fusion-by-expansion %s -split-input-file | FileCheck %s
func.func @control_producer_reshape_fusion(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?xf32>) -> tensor<?x?xf32> { func.func @control_producer_reshape_fusion(%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32> %0 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<?x?x?xf32> into tensor<?x?xf32>
%d0 = tensor.dim %0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %0, %c1 : tensor<?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%0, %arg1 : tensor<?x?xf32>, tensor<?xf32>) ins(%0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
outs(%init : tensor<?x?xf32>) { outs(%init : tensor<?x?xf32>) {
^bb0(%arg2 : f32, %arg3:f32, %arg4 : f32): ^bb0(%arg2 : f32, %arg3:f32, %arg4 : f32):
%2 = arith.addf %arg2, %arg3 : f32 %2 = arith.addf %arg2, %arg3 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
Show All 17 Lines
// ----- // -----
func.func @control_consumer_reshape_fusion(%arg0 : tensor<1x?x?xf32>, %arg1 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32> { func.func @control_consumer_reshape_fusion(%arg0 : tensor<1x?x?xf32>, %arg1 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%d0 = tensor.dim %arg0, %c1 : tensor<1x?x?xf32> %d0 = tensor.dim %arg0, %c1 : tensor<1x?x?xf32>
%d1 = tensor.dim %arg1, %c2 : tensor<1x?x?xf32> %d1 = tensor.dim %arg1, %c2 : tensor<1x?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%fill = linalg.generic { %fill = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%init : tensor<?x?xf32>) { outs(%init : tensor<?x?xf32>) {
^bb0(%arg2: f32): ^bb0(%arg2: f32):
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%0 = tensor.expand_shape %fill [[0, 1], [2]] : tensor<?x?xf32> into tensor<1x?x?xf32> %0 = tensor.expand_shape %fill [[0, 1], [2]] : tensor<?x?xf32> into tensor<1x?x?xf32>
Show All 11 Lines

mlir/test/Dialect/Linalg/reshape_fusion.mlir

Show First 20 Lines • Show All 136 Lines • ▼ Show 20 Lines
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1)> #map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d2)> #map2 = affine_map<(d0, d1, d2) -> (d2)>
func.func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>) func.func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>)
-> tensor<8x33x4xf32> { -> tensor<8x33x4xf32> {
%cst = arith.constant dense<2.000000e+00> : tensor<264x4xf32> %cst = arith.constant dense<2.000000e+00> : tensor<264x4xf32>
%0 = linalg.init_tensor [264, 4] : tensor<264x4xf32> %0 = tensor.empty() : tensor<264x4xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>) ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>)
outs(%0 : tensor<264x4xf32>) { outs(%0 : tensor<264x4xf32>) {
^bb0(%arg1: f32, %arg2: f32, %s: f32): ^bb0(%arg1: f32, %arg2: f32, %s: f32):
%2 = arith.mulf %arg1, %arg2 : f32 %2 = arith.mulf %arg1, %arg2 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<264x4xf32> } -> tensor<264x4xf32>
%2 = tensor.expand_shape %1 [[0, 1], [2]] : %2 = tensor.expand_shape %1 [[0, 1], [2]] :
tensor<264x4xf32> into tensor<8x33x4xf32> tensor<264x4xf32> into tensor<8x33x4xf32>
return %2 : tensor<8x33x4xf32> return %2 : tensor<8x33x4xf32>
} }
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK: func @generic_op_reshape_consumer_static // CHECK: func @generic_op_reshape_consumer_static
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32>
// CHECK-DAG: %[[CST:.+]] = arith.constant // CHECK-DAG: %[[CST:.+]] = arith.constant
// CHECK-SAME: : tensor<8x33x4xf32> // CHECK-SAME: : tensor<8x33x4xf32>
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [264, 4] // CHECK-DAG: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] // CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK-SAME: [0, 1], [2] // CHECK-SAME: [0, 1], [2]
// CHECK-SAME: tensor<264x4xf32> into tensor<8x33x4xf32> // CHECK-SAME: tensor<264x4xf32> into tensor<8x33x4xf32>
// CHECK: %[[T1:.+]] = tensor.expand_shape %[[INIT]] // CHECK: %[[T1:.+]] = tensor.expand_shape %[[INIT]]
// CHECK-SAME: [0, 1], [2] // CHECK-SAME: [0, 1], [2]
// CHECK-SAME: : tensor<264x4xf32> into tensor<8x33x4xf32> // CHECK-SAME: : tensor<264x4xf32> into tensor<8x33x4xf32>
// CHECK: %[[T2:.+]] = linalg.generic // CHECK: %[[T2:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]
▲ Show 20 Lines • Show All 102 Lines • ▼ Show 20 Lines
// CHECK: %[[T8:.+]] = arith.addi %[[T6]], %[[T7]] // CHECK: %[[T8:.+]] = arith.addi %[[T6]], %[[T7]]
// CHECK: linalg.yield %[[T8]] // CHECK: linalg.yield %[[T8]]
// ----- // -----
func.func @reshape_as_consumer_permutation func.func @reshape_as_consumer_permutation
(%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>) (%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>)
-> tensor<2x3x4x5x6x7xi32> { -> tensor<2x3x4x5x6x7xi32> {
%shape = linalg.init_tensor [6, 4, 210] : tensor<6x4x210xi32> %shape = tensor.empty() : tensor<6x4x210xi32>
%c = linalg.generic { %c = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,
affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>,
affine_map<(d0, d1, d2) -> (d0, d2, d1)>], affine_map<(d0, d1, d2) -> (d0, d2, d1)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%a, %b : tensor<210x6x4xi32>, tensor<210x4xi32>) ins(%a, %b : tensor<210x6x4xi32>, tensor<210x4xi32>)
outs(%shape : tensor<6x4x210xi32>) { outs(%shape : tensor<6x4x210xi32>) {
^bb0(%arg3 : i32, %arg4: i32, %s: i32): ^bb0(%arg3 : i32, %arg4: i32, %s: i32):
Show All 20 Lines
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 3)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 3)>
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 6)> // CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 6)>
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 7)> // CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1) -> (d0 + d1 * 7)>
// CHECK: func @reshape_as_consumer_permutation // CHECK: func @reshape_as_consumer_permutation
// CHECK-SAME: %[[ARG0:.+]]: tensor<210x6x4xi32> // CHECK-SAME: %[[ARG0:.+]]: tensor<210x6x4xi32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<210x4xi32> // CHECK-SAME: %[[ARG1:.+]]: tensor<210x4xi32>
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor [6, 4, 210] // CHECK-DAG: %[[INIT:.+]] = tensor.empty()
// CHECK-DAG: %[[T1:.+]] = tensor.expand_shape %[[ARG0]] // CHECK-DAG: %[[T1:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK-SAME: [0, 1, 2], [3, 4], [5] // CHECK-SAME: [0, 1, 2], [3, 4], [5]
// CHECK-DAG: %[[T2:.+]] = tensor.expand_shape %[[ARG1]] // CHECK-DAG: %[[T2:.+]] = tensor.expand_shape %[[ARG1]]
// CHECK-SAME: [0, 1, 2], [3] // CHECK-SAME: [0, 1, 2], [3]
// CHECK-DAG: %[[T3:.+]] = tensor.expand_shape %[[INIT]] // CHECK-DAG: %[[T3:.+]] = tensor.expand_shape %[[INIT]]
// CHECK-SAME: [0, 1], [2], [3, 4, 5] // CHECK-SAME: [0, 1], [2], [3, 4, 5]
// CHECK-SAME: : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32> // CHECK-SAME: : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>
// CHECK: %[[T4:.+]] = linalg.generic // CHECK: %[[T4:.+]] = linalg.generic
▲ Show 20 Lines • Show All 120 Lines • ▼ Show 20 Lines
// CHECK: return %[[T3]] : tensor<?x?x4x5xf32> // CHECK: return %[[T3]] : tensor<?x?x4x5xf32>
// ----- // -----
func.func @no_fuse_dynamic_dims(%arg0: tensor<?x?xf32>) -> tensor<?xf32> { func.func @no_fuse_dynamic_dims(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32> %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>
%1 = tensor.dim %0, %c0 : tensor<?xf32> %1 = tensor.dim %0, %c0 : tensor<?xf32>
%2 = linalg.init_tensor [%1] : tensor<?xf32> %2 = tensor.empty(%1) : tensor<?xf32>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%0 : tensor<?xf32>) outs(%2 : tensor<?xf32>) { ins(%0 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {
^bb0(%arg1 : f32, %arg2: f32): ^bb0(%arg1 : f32, %arg2: f32):
%4 = arith.addf %arg1, %arg1 : f32 %4 = arith.addf %arg1, %arg1 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %3 : tensor<?xf32> return %3 : tensor<?xf32>
} }
// CHECK: func @no_fuse_dynamic_dims // CHECK: func @no_fuse_dynamic_dims
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] // CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?xf32>) // CHECK-SAME: ins(%[[RESHAPE]] : tensor<?xf32>)
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
// ----- // -----
func.func @no_fuse_mismatched_dynamism(%arg0: tensor<2x1xi64>, %arg1: tensor<?xi64>) -> tensor<2xi64> { func.func @no_fuse_mismatched_dynamism(%arg0: tensor<2x1xi64>, %arg1: tensor<?xi64>) -> tensor<2xi64> {
%0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<2x1xi64> into tensor<2xi64> %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<2x1xi64> into tensor<2xi64>
%1 = linalg.init_tensor [2] : tensor<2xi64> %1 = tensor.empty() : tensor<2xi64>
%2 = linalg.generic %2 = linalg.generic
{indexing_maps = [affine_map<(d0) -> (d0)>, {indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (d0)>], affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%0, %arg1 : tensor<2xi64>, tensor<?xi64>) ins(%0, %arg1 : tensor<2xi64>, tensor<?xi64>)
outs(%1 : tensor<2xi64>) { outs(%1 : tensor<2xi64>) {
^bb0(%arg4: i64, %arg5: i64, %arg6: i64): ^bb0(%arg4: i64, %arg5: i64, %arg6: i64):
▲ Show 20 Lines • Show All 87 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/resolve-shaped-type-result-dims.mlir

// RUN: mlir-opt -resolve-shaped-type-result-dims -split-input-file %s | FileCheck %s // RUN: mlir-opt -resolve-shaped-type-result-dims -split-input-file %s | FileCheck %s
func.func @init_tensor_static_dim() -> (index, index) { func.func @empty_tensor_static_dim() -> (index, index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c6 = arith.constant 6 : index %c6 = arith.constant 6 : index
%0 = linalg.init_tensor [4, 5, %c6] : tensor<4x5x?xf32> %0 = tensor.empty(%c6) : tensor<4x5x?xf32>
%1 = tensor.dim %0, %c2 : tensor<4x5x?xf32> %1 = tensor.dim %0, %c2 : tensor<4x5x?xf32>
%2 = tensor.dim %0, %c0 : tensor<4x5x?xf32> %2 = tensor.dim %0, %c0 : tensor<4x5x?xf32>
return %1, %2 : index, index return %1, %2 : index, index
} }
// CHECK: func @init_tensor_static_dim // CHECK: func @empty_tensor_static_dim
// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index // CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index
// CHECK-DAG: %[[C6:.+]] = arith.constant 6 : index // CHECK-DAG: %[[C6:.+]] = arith.constant 6 : index
// CHECK: return %[[C6]], %[[C4]] // CHECK: return %[[C6]], %[[C4]]
// ----- // -----
func.func @init_tensor_dynamic_dim(%arg0 : index) -> (index) { func.func @empty_tensor_dynamic_dim(%arg0 : index) -> (index) {
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = linalg.init_tensor [4, 5, %arg0] : tensor<4x5x?xf32> %0 = tensor.empty(%arg0) : tensor<4x5x?xf32>
%1 = tensor.dim %0, %c2 : tensor<4x5x?xf32> %1 = tensor.dim %0, %c2 : tensor<4x5x?xf32>
return %1 : index return %1 : index
} }
// CHECK: func @init_tensor_dynamic_dim // CHECK: func @empty_tensor_dynamic_dim
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index
// CHECK: return %[[ARG0]] // CHECK: return %[[ARG0]]
// ----- // -----
func.func @init_tensor_dynamic_dim2(%arg0 : index, %arg1 : index) -> (index, index) { func.func @empty_tensor_dynamic_dim2(%arg0 : index, %arg1 : index) -> (index, index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32> %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
%1 = tensor.dim %0, %c0 : tensor<?x?xf32> %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
%2 = tensor.dim %0, %c1 : tensor<?x?xf32> %2 = tensor.dim %0, %c1 : tensor<?x?xf32>
return %1, %2 : index, index return %1, %2 : index, index
} }
// CHECK: func @init_tensor_dynamic_dim2 // CHECK: func @empty_tensor_dynamic_dim2
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK: return %[[ARG0]], %[[ARG1]] // CHECK: return %[[ARG0]], %[[ARG1]]
// ----- // -----
func.func @remove_dim_result_uses func.func @remove_dim_result_uses
(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, (%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
Show All 34 Lines
// ----- // -----
func.func @remove_dim_result_uses_outs func.func @remove_dim_result_uses_outs
(%arg0 : tensor<?xf32>, %arg1 : index) -> (index) { (%arg0 : tensor<?xf32>, %arg1 : index) -> (index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?xf32>
%0 = linalg.init_tensor [%d0, %arg1] : tensor<?x?xf32> %0 = tensor.empty(%d0, %arg1) : tensor<?x?xf32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [affine_map<(d0, d1) -> (d0)>, {indexing_maps = [affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?xf32>) outs(%0 : tensor<?x?xf32>) { ins(%arg0 : tensor<?xf32>) outs(%0 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32) : ^bb0(%arg2: f32, %arg3: f32) :
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
// ----- // -----
func.func @keep_result_dim_uses_sequence2 func.func @keep_result_dim_uses_sequence2
(%arg0 : tensor<?xf32>, %arg1 : index) -> (index, index) { (%arg0 : tensor<?xf32>, %arg1 : index) -> (index, index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?xf32>
%0 = linalg.init_tensor [%d0, %arg1] : tensor<?x?xf32> %0 = tensor.empty(%d0, %arg1) : tensor<?x?xf32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [affine_map<(d0, d1) -> (d0)>, {indexing_maps = [affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?xf32>) outs(%0 : tensor<?x?xf32>) { ins(%arg0 : tensor<?xf32>) outs(%0 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3 : f32): ^bb0(%arg2: f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%2 = tensor.dim %1, %c0 : tensor<?x?xf32> %2 = tensor.dim %1, %c0 : tensor<?x?xf32>
%3 = tensor.dim %1, %c1 : tensor<?x?xf32> %3 = tensor.dim %1, %c1 : tensor<?x?xf32>
return %2, %3 : index, index return %2, %3 : index, index
} }
// CHECK: func @keep_result_dim_uses_sequence2 // CHECK: func @keep_result_dim_uses_sequence2
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[T0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CHECK-DAG: %[[T0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK: return %[[T0]], %[[ARG1]] // CHECK: return %[[T0]], %[[ARG1]]
// ----- // -----
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
func.func @init_tensor_dim_of_linalg_result(%arg_0 : tensor<?xf32>, func.func @empty_tensor_dim_of_linalg_result(%arg_0 : tensor<?xf32>,
%arg_1: tensor<?xf32>) -> (index, index) { %arg_1: tensor<?xf32>) -> (index, index) {
%0, %1 = linalg.generic { %0, %1 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%arg_0 : tensor<?xf32>) } ins(%arg_0 : tensor<?xf32>)
outs(%arg_0, %arg_1 : tensor<?xf32>, tensor<?xf32>) { outs(%arg_0, %arg_1 : tensor<?xf32>, tensor<?xf32>) {
^bb0(%in: f32, %out_0: f32, %out_1: f32): ^bb0(%in: f32, %out_0: f32, %out_1: f32):
linalg.yield %in, %in : f32, f32 linalg.yield %in, %in : f32, f32
} -> (tensor<?xf32>, tensor<?xf32>) } -> (tensor<?xf32>, tensor<?xf32>)
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%num_elem_0 = tensor.dim %0, %c0 : tensor<?xf32> %num_elem_0 = tensor.dim %0, %c0 : tensor<?xf32>
%num_elem_1 = tensor.dim %1, %c0 : tensor<?xf32> %num_elem_1 = tensor.dim %1, %c0 : tensor<?xf32>
return %num_elem_0, %num_elem_1 : index, index return %num_elem_0, %num_elem_1 : index, index
} }
// CHECK: func @init_tensor_dim_of_linalg_result( // CHECK: func @empty_tensor_dim_of_linalg_result(
// CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG_0:[a-zA-Z0-9_]+]]: tensor<?xf32>
// CHECK-SAME: %[[ARG_1:[a-zA-Z0-9_]+]]: tensor<?xf32>) // CHECK-SAME: %[[ARG_1:[a-zA-Z0-9_]+]]: tensor<?xf32>)
// CHECK: %[[R0:.+]] = tensor.dim %[[ARG_0]] // CHECK: %[[R0:.+]] = tensor.dim %[[ARG_0]]
// CHECK: %[[R1:.+]] = tensor.dim %[[ARG_0]] // CHECK: %[[R1:.+]] = tensor.dim %[[ARG_0]]
// CHECK: return %[[R0]], %[[R1]] // CHECK: return %[[R0]], %[[R1]]
// ----- // -----
▲ Show 20 Lines • Show All 77 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/roundtrip.mlir

Show First 20 Lines • Show All 196 Lines • ▼ Show 20 Lines
// CHECK: return {{.*}} : tensor<?x?x?xf32> // CHECK: return {{.*}} : tensor<?x?x?xf32>
// ----- // -----
func.func @generic_with_multiple_tensor_outputs( func.func @generic_with_multiple_tensor_outputs(
%arg0: tensor<?xi32>, %arg1: tensor<?xi32>, %arg2: i32) %arg0: tensor<?xi32>, %arg1: tensor<?xi32>, %arg2: i32)
-> (tensor<i32>, tensor<i32>) { -> (tensor<i32>, tensor<i32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = linalg.init_tensor [] : tensor<i32> %0 = tensor.empty() : tensor<i32>
%1 = linalg.fill ins(%arg2 : i32) outs(%0 : tensor<i32>) -> tensor<i32> %1 = linalg.fill ins(%arg2 : i32) outs(%0 : tensor<i32>) -> tensor<i32>
%2 = linalg.init_tensor [] : tensor<i32> %2 = tensor.empty() : tensor<i32>
%3 = linalg.fill ins(%arg2 : i32) outs(%2 : tensor<i32>) -> tensor<i32> %3 = linalg.fill ins(%arg2 : i32) outs(%2 : tensor<i32>) -> tensor<i32>
%4:2 = linalg.generic { %4:2 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0, %arg1 : tensor<?xi32>, tensor<?xi32>) ins(%arg0, %arg1 : tensor<?xi32>, tensor<?xi32>)
outs(%1, %3 : tensor<i32>, tensor<i32>) { outs(%1, %3 : tensor<i32>, tensor<i32>) {
^bb0(%arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32): ^bb0(%arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32):
%5 = arith.cmpi sge, %arg3, %arg5 : i32 %5 = arith.cmpi sge, %arg3, %arg5 : i32
▲ Show 20 Lines • Show All 103 Lines • ▼ Show 20 Lines
// CHECK-LABEL: func @named_ops // CHECK-LABEL: func @named_ops
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// ----- // -----
#attr = {"foo"}
func.func @init_tensor(%arg0 : index, %arg1 : index)
{
%0 = linalg.init_tensor [3, 42] : tensor<3x42xf32>
%1 = linalg.init_tensor [4, %arg0, %arg1, 5] : tensor<4x?x?x5xf32>
%2 = linalg.init_tensor [2, 2] : tensor<2x2xf32, #attr>
return
}
// CHECK-LABEL: func @init_tensor
// CHECK: linalg.init_tensor [3, 42] : tensor<3x42xf32>
// CHECK: linalg.init_tensor [4, %{{.*}}, %{{.*}}, 5] : tensor<4x?x?x5xf32>
// CHECK: linalg.init_tensor [2, 2] : tensor<2x2xf32, {foo}>
// -----
func.func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> { func.func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {
%0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32> %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
%1 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32> %1 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
// CHECK: %{{.+}} = linalg.fill ins(%{{.+}} : f32) outs(%{{.+}} : tensor<?x?xf32>) -> tensor<?x?xf32> // CHECK: %{{.+}} = linalg.fill ins(%{{.+}} : f32) outs(%{{.+}} : tensor<?x?xf32>) -> tensor<?x?xf32>
// ----- // -----
func.func @mixed_parallel_reduced_results(%arg0 : tensor<?x?x?xf32>, func.func @mixed_parallel_reduced_results(%arg0 : tensor<?x?x?xf32>,
▲ Show 20 Lines • Show All 53 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/split_reduction.mlir

Show All 10 Lines
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: @matmul_split // CHECK-LABEL: @matmul_split
// CHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x4x64xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x4x64xf32>
// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<4x64x32xf32> // CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<4x64x32xf32>
// CHECK-DAG: %[[INI:.*]] = linalg.init_tensor [16, 32, 4] : tensor<16x32x4xf32> // CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<16x32x4xf32>
// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32> // CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32>
// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]] // CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
// CHECK-SAME: , iterator_types = ["parallel", "parallel", "parallel", "reduction"]} // CHECK-SAME: , iterator_types = ["parallel", "parallel", "parallel", "reduction"]}
// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x4x64xf32>, tensor<4x64x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) { // CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x4x64xf32>, tensor<4x64x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) {
// CHECK: arith.mulf // CHECK: arith.mulf
// CHECK: arith.addf // CHECK: arith.addf
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK: } -> tensor<16x32x4xf32> // CHECK: } -> tensor<16x32x4xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) { // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} ins(%[[G]] : tensor<16x32x4xf32>) outs(%{{.*}} : tensor<16x32xf32>) {
// CHECK: arith.addf // CHECK: arith.addf
// CHECK: linalg.yield %{{.*}} : f32 // CHECK: linalg.yield %{{.*}} : f32
// CHECK: } -> tensor<16x32xf32> // CHECK: } -> tensor<16x32xf32>
// CHECK: return %[[R]] : tensor<16x32xf32> // CHECK: return %[[R]] : tensor<16x32xf32>
// INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> // INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)> // INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d1)>
// INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> // INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// INNERPARALLELCHECK-LABEL: @matmul_split // INNERPARALLELCHECK-LABEL: @matmul_split
// INNERPARALLELCHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32 // INNERPARALLELCHECK-DAG: %[[ID:.*]] = arith.constant 0.000000e+00 : f32
// INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x64x4xf32> // INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x64x4xf32>
// INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<64x4x32xf32> // INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<64x4x32xf32>
// INNERPARALLELCHECK-DAG: %[[INI:.*]] = linalg.init_tensor [16, 32, 4] : tensor<16x32x4xf32> // INNERPARALLELCHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<16x32x4xf32>
// INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32> // INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<16x32x4xf32>) -> tensor<16x32x4xf32>
// INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]] // INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
// INNERPARALLELCHECK-SAME: , iterator_types = ["parallel", "parallel", "reduction", "parallel"]} // INNERPARALLELCHECK-SAME: , iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
// INNERPARALLELCHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) { // INNERPARALLELCHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) {
// INNERPARALLELCHECK: arith.mulf // INNERPARALLELCHECK: arith.mulf
// INNERPARALLELCHECK: arith.addf // INNERPARALLELCHECK: arith.addf
// INNERPARALLELCHECK: linalg.yield // INNERPARALLELCHECK: linalg.yield
// INNERPARALLELCHECK: } -> tensor<16x32x4xf32> // INNERPARALLELCHECK: } -> tensor<16x32x4xf32>
Show All 25 Lines
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()> // CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
//CHECK-LABEL: @generic_split_1d //CHECK-LABEL: @generic_split_1d
// CHECK: %[[ID:.*]] = arith.constant 1.000000e+00 : f32 // CHECK: %[[ID:.*]] = arith.constant 1.000000e+00 : f32
// CHECK: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<4x8xf32> // CHECK: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
// CHECK: %[[INI:.*]] = linalg.init_tensor [4] : tensor<4xf32> // CHECK: %[[INI:.*]] = tensor.empty() : tensor<4xf32>
// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32> // CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>
// CHECK: %[[G:.*]] = linalg.generic // CHECK: %[[G:.*]] = linalg.generic
// CHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], // CHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],
// CHECK: iterator_types = ["parallel", "reduction"]} ins(%[[I1]], %{{.*}} : tensor<4x8xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) { // CHECK: iterator_types = ["parallel", "reduction"]} ins(%[[I1]], %{{.*}} : tensor<4x8xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) {
// CHECK: arith.subf // CHECK: arith.subf
// CHECK: math.exp // CHECK: math.exp
// CHECK: arith.mulf // CHECK: arith.mulf
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK: } -> tensor<4xf32> // CHECK: } -> tensor<4xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) { // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["reduction"]} ins(%[[G]] : tensor<4xf32>) outs(%{{.*}} : tensor<f32>) {
// CHECK: arith.mulf // CHECK: arith.mulf
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK: } -> tensor<f32> // CHECK: } -> tensor<f32>
// CHECK: return %[[R]] : tensor<f32> // CHECK: return %[[R]] : tensor<f32>
// INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()> // INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> ()>
// INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d1)> // INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d1)>
// INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)> // INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0) -> (d0)>
// INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()> // INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0) -> ()>
//INNERPARALLELCHECK-LABEL: @generic_split_1d //INNERPARALLELCHECK-LABEL: @generic_split_1d
// INNERPARALLELCHECK: %[[ID:.*]] = arith.constant 1.000000e+00 : f32 // INNERPARALLELCHECK: %[[ID:.*]] = arith.constant 1.000000e+00 : f32
// INNERPARALLELCHECK: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<8x4xf32> // INNERPARALLELCHECK: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<8x4xf32>
// INNERPARALLELCHECK: %[[INI:.*]] = linalg.init_tensor [4] : tensor<4xf32> // INNERPARALLELCHECK: %[[INI:.*]] = tensor.empty() : tensor<4xf32>
// INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32> // INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<4xf32>) -> tensor<4xf32>
// INNERPARALLELCHECK: %[[G:.*]] = linalg.generic // INNERPARALLELCHECK: %[[G:.*]] = linalg.generic
// INNERPARALLELCHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], // INNERPARALLELCHECK: {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]],
// INNERPARALLELCHECK: iterator_types = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) { // INNERPARALLELCHECK: iterator_types = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) {
// INNERPARALLELCHECK: arith.subf // INNERPARALLELCHECK: arith.subf
// INNERPARALLELCHECK: math.exp // INNERPARALLELCHECK: math.exp
// INNERPARALLELCHECK: arith.mulf // INNERPARALLELCHECK: arith.mulf
// INNERPARALLELCHECK: linalg.yield // INNERPARALLELCHECK: linalg.yield
Show All 29 Lines
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d2, d1)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>
// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: func @generic_split_3d // CHECK-LABEL: func @generic_split_3d
// CHECK: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32 // CHECK: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32
// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<4x8x2xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<4x8x2xf32>
// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x4x8xf32> // CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x4x8xf32>
// CHECK: %[[INI:.*]] = linalg.init_tensor [5, 2, 4] : tensor<5x2x4xf32> // CHECK: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>
// CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32> // CHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>
// CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]} // CHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<4x8x2xf32>, tensor<5x4x8xf32>) outs(%[[F]] : tensor<5x2x4xf32>) { // CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<4x8x2xf32>, tensor<5x4x8xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {
// CHECK: arith.addf // CHECK: arith.addf
// CHECK: arith.maxf // CHECK: arith.maxf
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK: } -> tensor<5x2x4xf32> // CHECK: } -> tensor<5x2x4xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]} // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
// CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) { // CHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
// CHECK: arith.maxf // CHECK: arith.maxf
// CHECK: linalg.yield // CHECK: linalg.yield
// CHECK: } -> tensor<5x2xf32> // CHECK: } -> tensor<5x2xf32>
// CHECK: return %[[R]] : tensor<5x2xf32> // CHECK: return %[[R]] : tensor<5x2xf32>
// INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)> // INNERPARALLELCHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>
// INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)> // INNERPARALLELCHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>
// INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)> // INNERPARALLELCHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d2)>
// INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // INNERPARALLELCHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)> // INNERPARALLELCHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// INNERPARALLELCHECK-LABEL: func @generic_split_3d // INNERPARALLELCHECK-LABEL: func @generic_split_3d
// INNERPARALLELCHECK: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32 // INNERPARALLELCHECK: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32
// INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32> // INNERPARALLELCHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32>
// INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32> // INNERPARALLELCHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32>
// INNERPARALLELCHECK: %[[INI:.*]] = linalg.init_tensor [5, 2, 4] : tensor<5x2x4xf32> // INNERPARALLELCHECK: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32>
// INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32> // INNERPARALLELCHECK: %[[F:.*]] = linalg.fill ins(%[[ID]] : f32) outs(%[[INI]] : tensor<5x2x4xf32>) -> tensor<5x2x4xf32>
// INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]} // INNERPARALLELCHECK: %[[G:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "reduction", "parallel", "parallel"]}
// INNERPARALLELCHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) { // INNERPARALLELCHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) {
// INNERPARALLELCHECK: arith.addf // INNERPARALLELCHECK: arith.addf
// INNERPARALLELCHECK: arith.maxf // INNERPARALLELCHECK: arith.maxf
// INNERPARALLELCHECK: linalg.yield // INNERPARALLELCHECK: linalg.yield
// INNERPARALLELCHECK: } -> tensor<5x2x4xf32> // INNERPARALLELCHECK: } -> tensor<5x2x4xf32>
// INNERPARALLELCHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]} // INNERPARALLELCHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP3]], #[[$MAP4]]], iterator_types = ["parallel", "parallel", "reduction"]}
// INNERPARALLELCHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) { // INNERPARALLELCHECK-SAME: ins(%[[G]] : tensor<5x2x4xf32>) outs(%{{.*}} : tensor<5x2xf32>) {
// INNERPARALLELCHECK: arith.maxf // INNERPARALLELCHECK: arith.maxf
// INNERPARALLELCHECK: linalg.yield // INNERPARALLELCHECK: linalg.yield
// INNERPARALLELCHECK: } -> tensor<5x2xf32> // INNERPARALLELCHECK: } -> tensor<5x2xf32>
// INNERPARALLELCHECK: return %[[R]] : tensor<5x2xf32> // INNERPARALLELCHECK: return %[[R]] : tensor<5x2xf32>

mlir/test/Dialect/Linalg/tile-and-fuse-tensors.mlir

Show First 20 Lines • Show All 59 Lines • ▼ Show 20 Linesfunc.func @conv_tensors_static(%input: tensor<1x225x225x3xf32>, %filter: tensor<3x3x3x32xf32>, %elementwise: tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> {
%c112 = arith.constant 112 : index %c112 = arith.constant 112 : index
%c32 = arith.constant 32 : index %c32 = arith.constant 32 : index
%c16 = arith.constant 16 : index %c16 = arith.constant 16 : index
%c8 = arith.constant 8 : index %c8 = arith.constant 8 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%init = linalg.init_tensor [1, 112, 112, 32] : tensor<1x112x112x32xf32> %init = tensor.empty() : tensor<1x112x112x32xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32>
%conv = linalg.conv_2d_nhwc_hwcf %conv = linalg.conv_2d_nhwc_hwcf
{dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
ins(%input, %filter : tensor<1x225x225x3xf32>, tensor<3x3x3x32xf32>) ins(%input, %filter : tensor<1x225x225x3xf32>, tensor<3x3x3x32xf32>)
outs(%fill : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> outs(%fill : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32>
%for0 = scf.for %iv0 = %c0 to %c112 step %c8 iter_args(%arg0 = %fill) -> tensor<1x112x112x32xf32> { %for0 = scf.for %iv0 = %c0 to %c112 step %c8 iter_args(%arg0 = %fill) -> tensor<1x112x112x32xf32> {
Show All 27 Lines
} }
// CHECK: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 2)> // CHECK: #[[MAP0:.+]] = affine_map<(d0) -> (d0 * 2)>
// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> // CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
// CHECK: func @conv_tensors_static // CHECK: func @conv_tensors_static
// CHECK-SAME: (%[[INPUT:.+]]: tensor<1x225x225x3xf32>, %[[FILTER:.+]]: tensor<3x3x3x32xf32>, %[[ELEM:.+]]: tensor<1x112x112x32xf32>) // CHECK-SAME: (%[[INPUT:.+]]: tensor<1x225x225x3xf32>, %[[FILTER:.+]]: tensor<3x3x3x32xf32>, %[[ELEM:.+]]: tensor<1x112x112x32xf32>)
// CHECK: %[[INIT:.+]] = linalg.init_tensor [1, 112, 112, 32] : tensor<1x112x112x32xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x112x112x32xf32>
// CHECK-NEXT: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> // CHECK-NEXT: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32>
// CHECK-NEXT: scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG0:.+]] = %[[FILL]]) // CHECK-NEXT: scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG0:.+]] = %[[FILL]])
// CHECK-NEXT: %[[OFFSET_H:.+]] = affine.apply #[[MAP0]](%[[IV0]]) // CHECK-NEXT: %[[OFFSET_H:.+]] = affine.apply #[[MAP0]](%[[IV0]])
// CHECK-NEXT: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG1:.+]] = %[[ARG0]]) // CHECK-NEXT: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG1:.+]] = %[[ARG0]])
// CHECK-NEXT: %[[OFFSET_W:.+]] = affine.apply #[[MAP0]](%[[IV1]]) // CHECK-NEXT: %[[OFFSET_W:.+]] = affine.apply #[[MAP0]](%[[IV1]])
// CHECK-NEXT: %[[ST_INPUT:.+]] = tensor.extract_slice %arg0[0, %[[OFFSET_H]], %[[OFFSET_W]], 0] [1, 17, 33, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x17x33x3xf32> // CHECK-NEXT: %[[ST_INPUT:.+]] = tensor.extract_slice %arg0[0, %[[OFFSET_H]], %[[OFFSET_W]], 0] [1, 17, 33, 3] [1, 1, 1, 1] : tensor<1x225x225x3xf32> to tensor<1x17x33x3xf32>
// CHECK-NEXT: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG2:.+]] = %[[ARG1]]) // CHECK-NEXT: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[ARG2:.+]] = %[[ARG1]])
Show All 21 Linesfunc.func @conv_tensors_dynamic(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %elementwise: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
%c8 = arith.constant 8 : index %c8 = arith.constant 8 : index
%c16 = arith.constant 16 : index %c16 = arith.constant 16 : index
%n = tensor.dim %elementwise, %c0 : tensor<?x?x?x?xf32> %n = tensor.dim %elementwise, %c0 : tensor<?x?x?x?xf32>
%oh = tensor.dim %elementwise, %c1 : tensor<?x?x?x?xf32> %oh = tensor.dim %elementwise, %c1 : tensor<?x?x?x?xf32>
%ow = tensor.dim %elementwise, %c2 : tensor<?x?x?x?xf32> %ow = tensor.dim %elementwise, %c2 : tensor<?x?x?x?xf32>
%oc = tensor.dim %elementwise, %c3 : tensor<?x?x?x?xf32> %oc = tensor.dim %elementwise, %c3 : tensor<?x?x?x?xf32>
%init = linalg.init_tensor [%n, %oh, %ow, %oc] : tensor<?x?x?x?xf32> %init = tensor.empty(%n, %oh, %ow, %oc) : tensor<?x?x?x?xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
%conv = linalg.conv_2d_nhwc_hwcf %conv = linalg.conv_2d_nhwc_hwcf
{dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
ins(%input, %filter : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) ins(%input, %filter : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
outs(%fill : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> outs(%fill : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
%for0 = scf.for %iv0 = %c0 to %n step %c8 iter_args(%arg0 = %fill) -> tensor<?x?x?x?xf32> { %for0 = scf.for %iv0 = %c0 to %n step %c8 iter_args(%arg0 = %fill) -> tensor<?x?x?x?xf32> {
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Lines
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index // CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index
// CHECK-DAG: %[[ELEM_N:.+]] = tensor.dim %[[ELEM]], %[[C0]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[ELEM_N:.+]] = tensor.dim %[[ELEM]], %[[C0]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[ELEM_OH:.+]] = tensor.dim %[[ELEM]], %[[C1]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[ELEM_OH:.+]] = tensor.dim %[[ELEM]], %[[C1]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[ELEM_OW:.+]] = tensor.dim %[[ELEM]], %[[C2]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[ELEM_OW:.+]] = tensor.dim %[[ELEM]], %[[C2]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[ELEM_OC:.+]] = tensor.dim %[[ELEM]], %[[C3]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[ELEM_OC:.+]] = tensor.dim %[[ELEM]], %[[C3]] : tensor<?x?x?x?xf32>
// CHECK: %[[INIT:.+]] = linalg.init_tensor [%[[ELEM_N]], %[[ELEM_OH]], %[[ELEM_OW]], %[[ELEM_OC]]] : tensor<?x?x?x?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%[[ELEM_N]], %[[ELEM_OH]], %[[ELEM_OW]], %[[ELEM_OC]]) : tensor<?x?x?x?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> // CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
// CHECK-DAG: %[[FILTER_H:.+]] = tensor.dim %[[FILTER]], %[[C0]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[FILTER_H:.+]] = tensor.dim %[[FILTER]], %[[C0]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[FILTER_W:.+]] = tensor.dim %[[FILTER]], %[[C1]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[FILTER_W:.+]] = tensor.dim %[[FILTER]], %[[C1]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[FILTER_IC:.+]] = tensor.dim %[[FILTER]], %[[C2]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[FILTER_IC:.+]] = tensor.dim %[[FILTER]], %[[C2]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[FILTER_OC:.+]] = tensor.dim %[[FILTER]], %[[C3]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[FILTER_OC:.+]] = tensor.dim %[[FILTER]], %[[C3]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[INPUT_N:.+]] = tensor.dim %[[INPUT]], %[[C0]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[INPUT_N:.+]] = tensor.dim %[[INPUT]], %[[C0]] : tensor<?x?x?x?xf32>
// CHECK-DAG: %[[INPUT_C:.+]] = tensor.dim %[[INPUT]], %[[C3]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[INPUT_C:.+]] = tensor.dim %[[INPUT]], %[[C3]] : tensor<?x?x?x?xf32>
▲ Show 20 Lines • Show All 99 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/tile-and-peel-tensors.mlir

Show First 20 Lines • Show All 41 Lines • ▼ Show 20 Lines
// CHECK-PEEL-12: linalg.matmul ins({{.*}} : tensor<?x64xf32>, tensor<64x128xf32>) outs({{.*}} : tensor<?x128xf32>) // CHECK-PEEL-12: linalg.matmul ins({{.*}} : tensor<?x64xf32>, tensor<64x128xf32>) outs({{.*}} : tensor<?x128xf32>)
// CHECK-PEEL-12: } // CHECK-PEEL-12: }
// CHECK-PEEL-12: scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} { // CHECK-PEEL-12: scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} {
// CHECK-PEEL-12: linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x36xf32>) outs({{.*}} : tensor<?x36xf32>) // CHECK-PEEL-12: linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x36xf32>) outs({{.*}} : tensor<?x36xf32>)
// CHECK-PEEL-12: } // CHECK-PEEL-12: }
// CHECK-PEEL-12: } // CHECK-PEEL-12: }
func.func @matmul_static_tensor(%arg0: tensor<1500x1600xf32>, %arg1: tensor<1600x1700xf32>) func.func @matmul_static_tensor(%arg0: tensor<1500x1600xf32>, %arg1: tensor<1600x1700xf32>)
-> tensor<1500x1700xf32> { -> tensor<1500x1700xf32> {
%out = linalg.init_tensor [1500, 1700] : tensor<1500x1700xf32> %out = tensor.empty() : tensor<1500x1700xf32>
%r = linalg.matmul {__internal_linalg_transform__ = "tile"} %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
ins(%arg0, %arg1: tensor<1500x1600xf32>, tensor<1600x1700xf32>) ins(%arg0, %arg1: tensor<1500x1600xf32>, tensor<1600x1700xf32>)
outs(%out: tensor<1500x1700xf32>) -> tensor<1500x1700xf32> outs(%out: tensor<1500x1700xf32>) -> tensor<1500x1700xf32>
return %r : tensor<1500x1700xf32> return %r : tensor<1500x1700xf32>
} }
// ----- // -----
Show All 38 Lines
// CHECK-PEEL-12: } // CHECK-PEEL-12: }
// CHECK-PEEL-12: } // CHECK-PEEL-12: }
func.func @matmul_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) func.func @matmul_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%out = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %out = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%r = linalg.matmul {__internal_linalg_transform__ = "tile"} %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32> outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32>
return %r : tensor<?x?xf32> return %r : tensor<?x?xf32>
} }

mlir/test/Dialect/Linalg/tile-fuse-and-distribute.mlir

// RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-fuse-and-distribute-options -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-fuse-and-distribute-options -split-input-file | FileCheck %s
// CHECK: #[[MULMAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)> // CHECK: #[[MULMAP:.+]] = affine_map<()[s0, s1] -> (s0 * s1)>
// CHECK: #[[ADDMAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)> // CHECK: #[[ADDMAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1)>
// CHECK: func @fill_matmul_tensors( // CHECK: func @fill_matmul_tensors(
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32> // CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { // CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
func.func @fill_matmul_tensors( func.func @fill_matmul_tensors(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y // CHECK-DAG: %[[BIDY:.*]] = gpu.block_id y
// CHECK-DAG: %[[NBLOCKSY:.*]] = gpu.grid_dim y // CHECK-DAG: %[[NBLOCKSY:.*]] = gpu.grid_dim y
// CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x // CHECK-DAG: %[[BIDX:.*]] = gpu.block_id x
// CHECK-DAG: %[[NBLOCKSX:.*]] = gpu.grid_dim x // CHECK-DAG: %[[NBLOCKSX:.*]] = gpu.grid_dim x
// CHECK-DAG: %[[INIT:.+]] = linalg.init_tensor // CHECK-DAG: %[[INIT:.+]] = tensor.empty
// CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDY]], %[[C8]]] // CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDY]], %[[C8]]]
// CHECK: %[[LBY:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]] // CHECK: %[[LBY:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]]
// CHECK: %[[STEPY:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSY]], %[[C8]]] // CHECK: %[[STEPY:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSY]], %[[C8]]]
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[INIT]]) -> (tensor<?x?xf32>) { // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[INIT]]) -> (tensor<?x?xf32>) {
// CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDX]], %[[C8]]] // CHECK: %[[MUL:.+]] = affine.apply #[[MULMAP]]()[%[[BIDX]], %[[C8]]]
// CHECK: %[[LBX:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]] // CHECK: %[[LBX:.+]] = affine.apply #[[ADDMAP]]()[%[[MUL]], %[[C0]]]
// CHECK: %[[STEPX:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSX]], %[[C8]]] // CHECK: %[[STEPX:.+]] = affine.apply #[[MULMAP]]()[%[[NBLOCKSX]], %[[C8]]]
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) {
Show All 12 Lines
// CHECK: %[[TD2:.*]] = tensor.insert_slice %[[sTD2]] into %[[TC1]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32> // CHECK: %[[TD2:.*]] = tensor.insert_slice %[[sTD2]] into %[[TC1]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32> // CHECK: scf.yield %[[TD2]] : tensor<?x?xf32>
// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32> // CHECK: scf.yield %[[TD1]] : tensor<?x?xf32>
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%2 = linalg.init_tensor [%0, %1] : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32> %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%4 = linalg.matmul {__internal_linalg_transform__ = "tensors_fuse_distribute1"} %4 = linalg.matmul {__internal_linalg_transform__ = "tensors_fuse_distribute1"}
ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%3: tensor<?x?xf32>) outs(%3: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: return %[[TD0]] : tensor<?x?xf32> // CHECK: return %[[TD0]] : tensor<?x?xf32>
return %4 : tensor<?x?xf32> return %4 : tensor<?x?xf32>
} }

mlir/test/Dialect/Linalg/tile-scalarize-dynamic-dims.mlir

Show All 13 Lines
// CHECK: %[[S2:.*]] = tensor.extract_slice %[[ARG1]][%[[IV1]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32> // CHECK: %[[S2:.*]] = tensor.extract_slice %[[ARG1]][%[[IV1]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32>
// CHECK: %[[S3:.*]] = tensor.extract_slice %{{.*}}[%[[IV0]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32> // CHECK: %[[S3:.*]] = tensor.extract_slice %{{.*}}[%[[IV0]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32>
// CHECK: linalg.matmul ins(%[[S1]], %[[S2]] : tensor<1x1xf32>, tensor<1x2000xf32>) outs(%[[S3]] : tensor<1x2000xf32>) -> tensor<1x2000xf32> // CHECK: linalg.matmul ins(%[[S1]], %[[S2]] : tensor<1x1xf32>, tensor<1x2000xf32>) outs(%[[S3]] : tensor<1x2000xf32>) -> tensor<1x2000xf32>
func.func @matmul_partly_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x2000xf32>) func.func @matmul_partly_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x2000xf32>)
-> tensor<?x2000xf32> { -> tensor<?x2000xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%out = linalg.init_tensor [%d0, 2000] : tensor<?x2000xf32> %out = tensor.empty(%d0) : tensor<?x2000xf32>
%r = linalg.matmul {__internal_linalg_transform__ = "tile"} %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x2000xf32>) ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x2000xf32>)
outs(%out: tensor<?x2000xf32>) -> tensor<?x2000xf32> outs(%out: tensor<?x2000xf32>) -> tensor<?x2000xf32>
return %r : tensor<?x2000xf32> return %r : tensor<?x2000xf32>
} }
// ----- // -----
▲ Show 20 Lines • Show All 44 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/tile-tensors.mlir

Show All 31 Lines
func.func @generic_op_tensors( func.func @generic_op_tensors(
%arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32> %3 = tensor.empty(%0, %1, %2) : tensor<?x?x?xf32>
%4 = linalg.generic %4 = linalg.generic
{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d0, d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d2, d1)>,
affine_map<(d0, d1, d2) -> (d2, d1, d0)>], affine_map<(d0, d1, d2) -> (d2, d1, d0)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
outs(%3 : tensor<?x?x?xf32>) { outs(%3 : tensor<?x?x?xf32>) {
^bb0(%arg2 : f32, %arg3: f32, %arg4: f32): ^bb0(%arg2 : f32, %arg3: f32, %arg4: f32):
%5 = arith.addf %arg2, %arg3 : f32 %5 = arith.addf %arg2, %arg3 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
return %4 : tensor<?x?x?xf32> return %4 : tensor<?x?x?xf32>
} }
// CHECK-LABEL: func @generic_op_tensors // CHECK-LABEL: func @generic_op_tensors
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) { // CHECK: %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) { // CHECK: %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) { // CHECK: %[[TD2:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC2:.+]] = %[[TC1]]) -> (tensor<?x?x?xf32>) {
// CHECK: %[[STARG0:.+]] = tensor.extract_slice %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> // CHECK: %[[STARG0:.+]] = tensor.extract_slice %[[ARG0]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[STARG1:.+]] = tensor.extract_slice %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> // CHECK: %[[STARG1:.+]] = tensor.extract_slice %[[ARG1]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[STARG2:.+]] = tensor.extract_slice %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32> // CHECK: %[[STARG2:.+]] = tensor.extract_slice %[[TC2]][{{.+}}] : tensor<?x?x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[STRETURN:.+]] = linalg.generic // CHECK: %[[STRETURN:.+]] = linalg.generic
// CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>) // CHECK-SAME: ins(%[[STARG0]], %[[STARG1]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
▲ Show 20 Lines • Show All 53 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/transform-op-decompose.mlir

Show All 34 Linesfunc.func @conv_2d_nchw_fchw(%input: tensor<?x?x1x?xf32>, %filter: tensor<?x?x1x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
return %0 : tensor<?x?x1x?xf32> return %0 : tensor<?x?x1x?xf32>
} }
// CHECK-LABEL: @depthwise_conv_2d_nhwc_hwc // CHECK-LABEL: @depthwise_conv_2d_nhwc_hwc
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x113x96xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x113x96xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x3x96xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<1x3x96xf32>
func.func @depthwise_conv_2d_nhwc_hwc(%input: tensor<1x1x113x96xf32>, %filter: tensor<1x3x96xf32>) -> tensor<1x1x56x96xf32> { func.func @depthwise_conv_2d_nhwc_hwc(%input: tensor<1x1x113x96xf32>, %filter: tensor<1x3x96xf32>) -> tensor<1x1x56x96xf32> {
// CHECK: %[[RES:.+]] = linalg.init_tensor // CHECK: %[[RES:.+]] = tensor.empty
%init = linalg.init_tensor [1, 1, 56, 96] : tensor<1x1x56x96xf32> %init = tensor.empty() : tensor<1x1x56x96xf32>
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]] // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]] // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
// CHECK: %[[SLICERES:.+]] = tensor.extract_slice %[[RES]] // CHECK: %[[SLICERES:.+]] = tensor.extract_slice %[[RES]]
// CHECK: %[[OPRES:.+]] = linalg.depthwise_conv_1d_nwc_wc // CHECK: %[[OPRES:.+]] = linalg.depthwise_conv_1d_nwc_wc
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]]
// CHECK-SAME: outs(%[[SLICERES]] // CHECK-SAME: outs(%[[SLICERES]]
// CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[OPRES]] into %[[RES]] // CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[OPRES]] into %[[RES]]
%0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>} %0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
Show All 14 Lines

mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir

Show First 20 Lines • Show All 60 Lines • ▼ Show 20 Lines
#map2 = affine_map<(d0)[s0] -> (-(d0 * s0) + 64, s0)> #map2 = affine_map<(d0)[s0] -> (-(d0 * s0) + 64, s0)>
module { module {
// CHECK-LABEL: func.func @fuse_untileable_op // CHECK-LABEL: func.func @fuse_untileable_op
// CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
// CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<64xf32> // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<64xf32>
// CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<64xf32> // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<64xf32>
func.func @fuse_untileable_op(%arg0: index, %arg1: tensor<64xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> { func.func @fuse_untileable_op(%arg0: index, %arg1: tensor<64xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {
%0 = linalg.init_tensor [%arg0] : tensor<?xf32> %0 = tensor.empty(%arg0) : tensor<?xf32>
%1 = affine.apply #map0()[%arg0] %1 = affine.apply #map0()[%arg0]
// CHECK: scf.foreach_thread {{.*}} { // CHECK: scf.foreach_thread {{.*}} {
%2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<64xf32>) { %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<64xf32>) {
// CHECK: %[[INIT_TENSOR:.*]] = linalg.init_tensor // CHECK: %[[INIT_TENSOR:.*]] = tensor.empty
%3 = affine.apply #map1(%arg3)[%arg0] %3 = affine.apply #map1(%arg3)[%arg0]
%4 = affine.min #map2(%arg3)[%arg0] %4 = affine.min #map2(%arg3)[%arg0]
%5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<64xf32> to tensor<?xf32> %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<64xf32> to tensor<?xf32>
// CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[INIT_TENSOR]] // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[INIT_TENSOR]]
%7 = linalg.elemwise_unary ins(%0 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32> %7 = linalg.elemwise_unary ins(%0 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<64xf32> tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<64xf32>
} }
} }
// CHECK: } // CHECK: }
func.return %2 : tensor<64xf32> func.return %2 : tensor<64xf32>
} }
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
transform.sequence %arg0 failures(propagate) { transform.sequence %arg0 failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.init_tensor"]} in %arg1 %0 = transform.structured.match ops{["tensor.empty"]} in %arg1
%1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1 %1 = transform.structured.match ops{["scf.foreach_thread"]} in %arg1
// linalg.init_tensor is not tileable. The op is cloned and fused. // tensor.empty is not tileable. The op is cloned and fused.
transform.structured.fuse_into_containing_op %0 into %1 transform.structured.fuse_into_containing_op %0 into %1
} }
} }
} }
// ----- // -----
#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)> #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
▲ Show 20 Lines • Show All 44 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/transform-op-fuse.mlir

Show First 20 Lines • Show All 53 Lines • ▼ Show 20 Lines
} }
// ----- // -----
// CHECK-LABEL: func.func @interchange_reduction // CHECK-LABEL: func.func @interchange_reduction
// CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>) // CHECK-SAME: (%[[INPUT:.+]]: tensor<12x7x25xf32>)
func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> { func.func @interchange_reduction(%input: tensor<12x7x25xf32>) -> tensor<12x25xf32> {
%five = arith.constant 5.0 : f32 %five = arith.constant 5.0 : f32
%init = linalg.init_tensor [12, 25] : tensor<12x25xf32> %init = tensor.empty() : tensor<12x25xf32>
// CHECK: %[[INIT:.+]] = linalg.init_tensor [12, 25] // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index // CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index
// CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index // CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index
// CHECK: scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %[[C5]] iter_args(%[[FOR_ARG0:.+]] = %[[INIT]]) // CHECK: scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %[[C5]] iter_args(%[[FOR_ARG0:.+]] = %[[INIT]])
// CHECK: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %[[C7]] iter_args(%[[FOR_ARG1:.+]] = %[[FOR_ARG0]]) // CHECK: scf.for %[[IV1:.+]] = %{{.+}} to %{{.+}} step %[[C7]] iter_args(%[[FOR_ARG1:.+]] = %[[FOR_ARG0]])
// CHECK: %[[OUT_SLICE0:.+]] = tensor.extract_slice %[[INPUT]][%[[IV0]], 0, %[[IV1]]] // CHECK: %[[OUT_SLICE0:.+]] = tensor.extract_slice %[[INPUT]][%[[IV0]], 0, %[[IV1]]]
// CHECK: %[[OUT_SLICE1:.+]] = tensor.extract_slice %[[FOR_ARG1]][%[[IV0]], %[[IV1]]] // CHECK: %[[OUT_SLICE1:.+]] = tensor.extract_slice %[[FOR_ARG1]][%[[IV0]], %[[IV1]]]
// CHECK: %[[FILL:.+]] = linalg.fill {{.+}} outs(%[[OUT_SLICE1]] : tensor<?x?xf32>) // CHECK: %[[FILL:.+]] = linalg.fill {{.+}} outs(%[[OUT_SLICE1]] : tensor<?x?xf32>)
// CHECK: %[[C4:.+]] = arith.constant 4 : index // CHECK: %[[C4:.+]] = arith.constant 4 : index
Show All 25 Lines

mlir/test/Dialect/Linalg/transform-op-match.mlir

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
} }
// ----- // -----
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
func.func @match_complex_attribute(%arg0: tensor<12x128x32xf32>) func.func @match_complex_attribute(%arg0: tensor<12x128x32xf32>)
-> tensor<128x12x32xf32> { -> tensor<128x12x32xf32> {
%0 = linalg.init_tensor [128, 12, 32] : tensor<128x12x32xf32> %0 = tensor.empty() : tensor<128x12x32xf32>
// expected-remark @below {{matched complex attr}} // expected-remark @below {{matched complex attr}}
%1 = linalg.generic {indexing_maps = [#map0, #map1], %1 = linalg.generic {indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<12x128x32xf32>) ins(%arg0 : tensor<12x128x32xf32>)
outs(%0 : tensor<128x12x32xf32>) { outs(%0 : tensor<128x12x32xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<128x12x32xf32> } -> tensor<128x12x32xf32>
Show All 21 Lines

mlir/test/Dialect/Linalg/vectorization.mlir

Show First 20 Lines • Show All 176 Lines • ▼ Show 20 Lines
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map1 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)> // CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
// CHECK: func @generic_interchanged_transpose // CHECK: func @generic_interchanged_transpose
func.func @generic_interchanged_transpose(%arg0: tensor<12x128x32xf32>) -> tensor<128x12x32xf32> { func.func @generic_interchanged_transpose(%arg0: tensor<12x128x32xf32>) -> tensor<128x12x32xf32> {
// CHECK: %[[IN:.+]] = vector.transfer_read // CHECK: %[[IN:.+]] = vector.transfer_read
// CHECK: vector.transfer_write %[[IN]], {{.+}} permutation_map = #[[MAP]] // CHECK: vector.transfer_write %[[IN]], {{.+}} permutation_map = #[[MAP]]
%0 = linalg.init_tensor [128, 12, 32] : tensor<128x12x32xf32> %0 = tensor.empty() : tensor<128x12x32xf32>
%1 = linalg.generic {indexing_maps = [#map0, #map1], %1 = linalg.generic {indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<12x128x32xf32>) ins(%arg0 : tensor<12x128x32xf32>)
outs(%0 : tensor<128x12x32xf32>) { outs(%0 : tensor<128x12x32xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<128x12x32xf32> } -> tensor<128x12x32xf32>
return %1 : tensor<128x12x32xf32> return %1 : tensor<128x12x32xf32>
▲ Show 20 Lines • Show All 587 Lines • ▼ Show 20 Lines
// ----- // -----
// CHECK-LABEL: func @pad_static( // CHECK-LABEL: func @pad_static(
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x?x2xf32>, %[[PAD:.*]]: f32 // CHECK-SAME: %[[ARG0:.*]]: tensor<2x?x2xf32>, %[[PAD:.*]]: f32
// CHECK-NOT: tensor.pad // CHECK-NOT: tensor.pad
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[INIT:.*]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32> // CHECK-DAG: %[[INIT:.*]] = tensor.empty() : tensor<2x3x4xf32>
// CHECK-DAG: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x3x4xf32> // CHECK-DAG: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x3x4xf32>
// CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]]{{.*}} : vector<2x3x4xf32>, tensor<2x3x4xf32> // CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]]{{.*}} : vector<2x3x4xf32>, tensor<2x3x4xf32>
// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, false, true]} : tensor<2x?x2xf32>, vector<2x3x2xf32> // CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, false, true]} : tensor<2x?x2xf32>, vector<2x3x2xf32>
// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x3x2xf32>, tensor<2x3x4xf32> // CHECK: %[[RESULT:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x3x2xf32>, tensor<2x3x4xf32>
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
func.func @pad_static(%arg0: tensor<2x?x2xf32>, %pad_value: f32) -> tensor<2x3x4xf32> { func.func @pad_static(%arg0: tensor<2x?x2xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {
%0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] { %0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] {
^bb0(%arg1: index, %arg2: index, %arg3: index): ^bb0(%arg1: index, %arg2: index, %arg3: index):
Show All 15 Lines
// ----- // -----
// CHECK-LABEL: func @pad_static_source( // CHECK-LABEL: func @pad_static_source(
// CHECK-SAME: %[[ARG0:.*]]: tensor<2x5x2xf32>, %[[PAD:.*]]: f32 // CHECK-SAME: %[[ARG0:.*]]: tensor<2x5x2xf32>, %[[PAD:.*]]: f32
// CHECK-NOT: tensor.pad // CHECK-NOT: tensor.pad
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
// CHECK: %[[INIT:.*]] = linalg.init_tensor [2, 6, 4] : tensor<2x6x4xf32> // CHECK: %[[INIT:.*]] = tensor.empty() : tensor<2x6x4xf32>
// CHECK: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x6x4xf32> // CHECK: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x6x4xf32>
// CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<2x6x4xf32>, tensor<2x6x4xf32> // CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<2x6x4xf32>, tensor<2x6x4xf32>
// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : tensor<2x5x2xf32>, vector<2x5x2xf32> // CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : tensor<2x5x2xf32>, vector<2x5x2xf32>
// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x5x2xf32>, tensor<2x6x4xf32> // CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x5x2xf32>, tensor<2x6x4xf32>
// CHECK: return %[[WRITE]] // CHECK: return %[[WRITE]]
func.func @pad_static_source(%arg0: tensor<2x5x2xf32>, %pad_value: f32) -> tensor<2x6x4xf32> { func.func @pad_static_source(%arg0: tensor<2x5x2xf32>, %pad_value: f32) -> tensor<2x6x4xf32> {
%0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] { %0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] {
^bb0(%arg1: index, %arg2: index, %arg3: index): ^bb0(%arg1: index, %arg2: index, %arg3: index):
Show All 23 Lines
// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index // CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index
// CHECK: %[[V0:.*]] = arith.addi %[[LOW]], %[[C2]] : index // CHECK: %[[V0:.*]] = arith.addi %[[LOW]], %[[C2]] : index
// CHECK: %[[V1:.*]] = arith.addi %[[V0]], %[[C3]] : index // CHECK: %[[V1:.*]] = arith.addi %[[V0]], %[[C3]] : index
// CHECK: %[[V2:.*]] = arith.addi %[[HIGH]], %[[C5]] : index // CHECK: %[[V2:.*]] = arith.addi %[[HIGH]], %[[C5]] : index
// CHECK: %[[DIM3:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32> // CHECK: %[[DIM3:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
// CHECK: %[[V4:.*]] = arith.addi %[[DIM3]], %[[C3]] : index // CHECK: %[[V4:.*]] = arith.addi %[[DIM3]], %[[C3]] : index
// CHECK: %[[V5:.*]] = arith.addi %[[V4]], %[[C2]] : index // CHECK: %[[V5:.*]] = arith.addi %[[V4]], %[[C2]] : index
// CHECK: %[[INIT:.*]] = linalg.init_tensor [6, %[[V1]], %[[V2]], %[[V5]]] : tensor<6x?x?x?xf32> // CHECK: %[[INIT:.*]] = tensor.empty(%[[V1]], %[[V2]], %[[V5]]) : tensor<6x?x?x?xf32>
// CHECK: %[[FILL:.*]] = linalg.fill ins(%{{.*}} : f32) outs(%[[INIT]] : tensor<6x?x?x?xf32>) -> tensor<6x?x?x?xf32> // CHECK: %[[FILL:.*]] = linalg.fill ins(%{{.*}} : f32) outs(%[[INIT]] : tensor<6x?x?x?xf32>) -> tensor<6x?x?x?xf32>
// CHECK: %[[SRCDIM:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32> // CHECK: %[[SRCDIM:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
// CHECK: %[[RESULT:.*]] = tensor.insert_slice %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32> // CHECK: %[[RESULT:.*]] = tensor.insert_slice %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32>
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
func.func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, func.func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
%pad_value: f32) -> tensor<6x?x?x?xf32> { %pad_value: f32) -> tensor<6x?x?x?xf32> {
%0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] { %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
▲ Show 20 Lines • Show All 317 Lines • ▼ Show 20 Lines^bb0(%arg0: !pdl.operation):
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_max_2d( // CHECK-LABEL: func @red_max_2d(
func.func @red_max_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { func.func @red_max_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> {
// CHECK: %[[CMINF:.+]] = arith.constant dense<-3.402820e+38> : vector<4xf32> // CHECK: %[[CMINF:.+]] = arith.constant dense<-3.402820e+38> : vector<4xf32>
// CHECK: linalg.init_tensor [4] : tensor<4xf32> // CHECK: tensor.empty() : tensor<4xf32>
// CHECK: vector.multi_reduction <maxf>, {{.*}}, %[[CMINF]] [1] : vector<4x4xf32> to vector<4xf32> // CHECK: vector.multi_reduction <maxf>, {{.*}}, %[[CMINF]] [1] : vector<4x4xf32> to vector<4xf32>
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32>
%ident = arith.constant -3.40282e+38 : f32 %ident = arith.constant -3.40282e+38 : f32
%init = linalg.init_tensor [4] : tensor<4xf32> %init = tensor.empty() : tensor<4xf32>
%fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> %fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%max = arith.maxf %in0, %out0 : f32 %max = arith.maxf %in0, %out0 : f32
linalg.yield %max : f32 linalg.yield %max : f32
Show All 12 Lines^bb0(%arg0: !pdl.operation):
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_min_2d( // CHECK-LABEL: func @red_min_2d(
func.func @red_min_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { func.func @red_min_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> {
// CHECK: %[[CMAXF:.+]] = arith.constant dense<3.402820e+38> : vector<4xf32> // CHECK: %[[CMAXF:.+]] = arith.constant dense<3.402820e+38> : vector<4xf32>
// CHECK: linalg.init_tensor [4] : tensor<4xf32> // CHECK: tensor.empty() : tensor<4xf32>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.multi_reduction <minf>, {{.*}}, %[[CMAXF]] [1] : vector<4x4xf32> to vector<4xf32> // CHECK: vector.multi_reduction <minf>, {{.*}}, %[[CMAXF]] [1] : vector<4x4xf32> to vector<4xf32>
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32>
%maxf32 = arith.constant 3.40282e+38 : f32 %maxf32 = arith.constant 3.40282e+38 : f32
%init = linalg.init_tensor [4] : tensor<4xf32> %init = tensor.empty() : tensor<4xf32>
%fill = linalg.fill ins(%maxf32 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> %fill = linalg.fill ins(%maxf32 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%min = arith.minf %out0, %in0 : f32 %min = arith.minf %out0, %in0 : f32
linalg.yield %min : f32 linalg.yield %min : f32
Show All 11 Lines ^bb1(%arg1: !pdl.operation):
%5 = transform.structured.vectorize %4 %5 = transform.structured.vectorize %4
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_mul_2d( // CHECK-LABEL: func @red_mul_2d(
func.func @red_mul_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { func.func @red_mul_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> {
// CHECK: linalg.init_tensor [4] : tensor<4xf32> // CHECK: tensor.empty() : tensor<4xf32>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.multi_reduction <mul>, {{.*}}, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> // CHECK: vector.multi_reduction <mul>, {{.*}}, {{.*}} [1] : vector<4x4xf32> to vector<4xf32>
// CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32>
%ident = arith.constant 1.0 : f32 %ident = arith.constant 1.0 : f32
%init = linalg.init_tensor [4] : tensor<4xf32> %init = tensor.empty() : tensor<4xf32>
%fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> %fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%mul = arith.mulf %in0, %out0 : f32 %mul = arith.mulf %in0, %out0 : f32
linalg.yield %mul : f32 linalg.yield %mul : f32
Show All 11 Lines ^bb1(%arg1: !pdl.operation):
%5 = transform.structured.vectorize %4 %5 = transform.structured.vectorize %4
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_or_2d( // CHECK-LABEL: func @red_or_2d(
func.func @red_or_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { func.func @red_or_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> {
// CHECK: linalg.init_tensor [4] : tensor<4xi1> // CHECK: tensor.empty() : tensor<4xi1>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1>
// CHECK: vector.multi_reduction <or>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> // CHECK: vector.multi_reduction <or>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1>
// CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1>
%ident = arith.constant false %ident = arith.constant false
%init = linalg.init_tensor [4] : tensor<4xi1> %init = tensor.empty() : tensor<4xi1>
%fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) {
^bb0(%in0: i1, %out0: i1): ^bb0(%in0: i1, %out0: i1):
%or = arith.ori %in0, %out0 : i1 %or = arith.ori %in0, %out0 : i1
linalg.yield %or : i1 linalg.yield %or : i1
Show All 11 Lines ^bb1(%arg1: !pdl.operation):
%5 = transform.structured.vectorize %4 %5 = transform.structured.vectorize %4
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_and_2d( // CHECK-LABEL: func @red_and_2d(
func.func @red_and_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { func.func @red_and_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> {
// CHECK: linalg.init_tensor [4] : tensor<4xi1> // CHECK: tensor.empty() : tensor<4xi1>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1>
// CHECK: vector.multi_reduction <and>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> // CHECK: vector.multi_reduction <and>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1>
// CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1>
%ident = arith.constant true %ident = arith.constant true
%init = linalg.init_tensor [4] : tensor<4xi1> %init = tensor.empty() : tensor<4xi1>
%fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) {
^bb0(%in0: i1, %out0: i1): ^bb0(%in0: i1, %out0: i1):
%and = arith.andi %in0, %out0 : i1 %and = arith.andi %in0, %out0 : i1
linalg.yield %and : i1 linalg.yield %and : i1
Show All 11 Lines ^bb1(%arg1: !pdl.operation):
%5 = transform.structured.vectorize %4 %5 = transform.structured.vectorize %4
} }
} }
// ----- // -----
// CHECK-LABEL: func @red_xor_2d( // CHECK-LABEL: func @red_xor_2d(
func.func @red_xor_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { func.func @red_xor_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> {
// CHECK: linalg.init_tensor [4] : tensor<4xi1> // CHECK: tensor.empty() : tensor<4xi1>
// CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1>
// CHECK: vector.multi_reduction <xor>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> // CHECK: vector.multi_reduction <xor>, {{.*}}, {{.*}} [1] : vector<4x4xi1> to vector<4xi1>
// CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1>
%ident = arith.constant false %ident = arith.constant false
%init = linalg.init_tensor [4] : tensor<4xi1> %init = tensor.empty() : tensor<4xi1>
%fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) {
^bb0(%in0: i1, %out0: i1): ^bb0(%in0: i1, %out0: i1):
%xor = arith.xori %in0, %out0 : i1 %xor = arith.xori %in0, %out0 : i1
linalg.yield %xor : i1 linalg.yield %xor : i1
Show All 18 Lines
// CHECK-LABEL: func @explicit_broadcast( // CHECK-LABEL: func @explicit_broadcast(
func.func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> { func.func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> {
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32>
// CHECK: subf {{.*}} : vector<4x4xf32> // CHECK: subf {{.*}} : vector<4x4xf32>
// CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32> // CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32>
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
%init = linalg.init_tensor [4, 4] : tensor<4x4xf32> %init = tensor.empty() : tensor<4x4xf32>
%fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, 0)>, affine_map<(d0, d1) -> (d0, 0)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>) ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>)
outs(%fill : tensor<4x4xf32>) { outs(%fill : tensor<4x4xf32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
Show All 22 Lines
func.func @fused_broadcast_red_2d(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4xf32> { func.func @fused_broadcast_red_2d(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4xf32> {
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32>
// CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M6]]} : tensor<4x1xf32>, vector<4x4xf32> // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M6]]} : tensor<4x1xf32>, vector<4x4xf32>
// CHECK: subf {{.*}} : vector<4x4xf32> // CHECK: subf {{.*}} : vector<4x4xf32>
// CHECK: math.exp {{.*}} : vector<4x4xf32> // CHECK: math.exp {{.*}} : vector<4x4xf32>
// CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} : vector<4x4xf32> to vector<4xf32> // CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} : vector<4x4xf32> to vector<4xf32>
// CHECK: vector.transfer_write {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<4xf32> // CHECK: vector.transfer_write {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<4xf32>
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
%init = linalg.init_tensor [4] : tensor<4xf32> %init = tensor.empty() : tensor<4xf32>
%fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32>
%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, 0)>, affine_map<(d0, d1) -> (d0, 0)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>) ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>)
outs(%fill : tensor<4xf32>) { outs(%fill : tensor<4xf32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
Show All 25 Lines
// CHECK-LABEL: func @reduce_1d( // CHECK-LABEL: func @reduce_1d(
// CHECK-SAME: %[[A:.*]]: tensor<32xf32> // CHECK-SAME: %[[A:.*]]: tensor<32xf32>
func.func @reduce_1d(%arg0: tensor<32xf32>) -> tensor<f32> { func.func @reduce_1d(%arg0: tensor<32xf32>) -> tensor<f32> {
// CHECK-DAG: %[[vF0:.*]] = arith.constant dense<0.000000e+00> : vector<f32> // CHECK-DAG: %[[vF0:.*]] = arith.constant dense<0.000000e+00> : vector<f32>
// CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
%f0 = arith.constant 0.000000e+00 : f32 %f0 = arith.constant 0.000000e+00 : f32
// CHECK: %[[init:.*]] = linalg.init_tensor [] : tensor<f32> // CHECK: %[[init:.*]] = tensor.empty() : tensor<f32>
%0 = linalg.init_tensor [] : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = linalg.fill ins(%f0 : f32) outs(%0 : tensor<f32>) -> tensor<f32> %1 = linalg.fill ins(%f0 : f32) outs(%0 : tensor<f32>) -> tensor<f32>
// CHECK: %[[r:.*]] = vector.transfer_read %[[A]][%[[C0]]] // CHECK: %[[r:.*]] = vector.transfer_read %[[A]][%[[C0]]]
// CHECK-SAME: : tensor<32xf32>, vector<32xf32> // CHECK-SAME: : tensor<32xf32>, vector<32xf32>
// CHECK: %[[f0:.*]] = vector.extractelement %[[vF0]][] : vector<f32> // CHECK: %[[f0:.*]] = vector.extractelement %[[vF0]][] : vector<f32>
// CHECK: %[[red:.*]] = vector.multi_reduction <add>, %[[r]], %[[f0]] [0] // CHECK: %[[red:.*]] = vector.multi_reduction <add>, %[[r]], %[[f0]] [0]
// CHECK-SAME: : vector<32xf32> to f32 // CHECK-SAME: : vector<32xf32> to f32
// CHECK: %[[red_v1:.*]] = vector.broadcast %[[red]] : f32 to vector<f32> // CHECK: %[[red_v1:.*]] = vector.broadcast %[[red]] : f32 to vector<f32>
Show All 28 Lines
// This test checks that vectorization does not occur when an input indexing map // This test checks that vectorization does not occur when an input indexing map
// is not a projected permutation. In the future, this can be converted to a // is not a projected permutation. In the future, this can be converted to a
// positive test when support is added. // positive test when support is added.
// CHECK-LABEL: func @not_projected_permutation // CHECK-LABEL: func @not_projected_permutation
func.func @not_projected_permutation(%arg0: tensor<8x8xf32>) -> tensor<6x6x3x3xf32> { func.func @not_projected_permutation(%arg0: tensor<8x8xf32>) -> tensor<6x6x3x3xf32> {
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
%init = linalg.init_tensor [6, 6, 3, 3] : tensor<6x6x3x3xf32> %init = tensor.empty() : tensor<6x6x3x3xf32>
%fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<6x6x3x3xf32>) -> tensor<6x6x3x3xf32> %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<6x6x3x3xf32>) -> tensor<6x6x3x3xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
%result = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>, %result = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0 + d2, d1 + d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<8x8xf32>) ins(%arg0 : tensor<8x8xf32>)
outs(%fill : tensor<6x6x3x3xf32>) { outs(%fill : tensor<6x6x3x3xf32>) {
^bb0(%arg7: f32, %arg9: f32): ^bb0(%arg7: f32, %arg9: f32):
▲ Show 20 Lines • Show All 56 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/sparse_1d.mlir

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK-LABEL: func @add_d_init( // CHECK-LABEL: func @add_d_init(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: f32) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_1:.*]]: f32) -> tensor<32xf32> {
// CHECK: %[[VAL_2:.*]] = arith.constant 32 : index // CHECK: %[[VAL_2:.*]] = arith.constant 32 : index
// CHECK: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_INITTENSOR:.*]] = linalg.init_tensor [32] : tensor<32xf32> // CHECK: %[[VAL_INITTENSOR:.*]] = tensor.empty() : tensor<32xf32>
// CHECK: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_INITTENSOR]] : memref<32xf32> // CHECK: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_INITTENSOR]] : memref<32xf32>
// CHECK: linalg.fill ins(%[[VAL_3]] : f32) outs(%[[VAL_7]] : memref<32xf32>) // CHECK: linalg.fill ins(%[[VAL_3]] : f32) outs(%[[VAL_7]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_8:.*]] = %[[VAL_4]] to %[[VAL_2]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_8:.*]] = %[[VAL_4]] to %[[VAL_2]] step %[[VAL_5]] {
// CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_8]]] : memref<?xf32> // CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_8]]] : memref<?xf32>
// CHECK: %[[VAL_10:.*]] = arith.addf %[[VAL_9]], %[[VAL_1]] : f32 // CHECK: %[[VAL_10:.*]] = arith.addf %[[VAL_9]], %[[VAL_1]] : f32
// CHECK: memref.store %[[VAL_10]], %[[VAL_7]]{{\[}}%[[VAL_8]]] : memref<32xf32> // CHECK: memref.store %[[VAL_10]], %[[VAL_7]]{{\[}}%[[VAL_8]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_11:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf32> // CHECK: %[[VAL_11:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf32>
// CHECK: return %[[VAL_11]] : tensor<32xf32> // CHECK: return %[[VAL_11]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_d_init(%arga: tensor<32xf32, #DV>, %argb: f32) -> tensor<32xf32> { func.func @add_d_init(%arga: tensor<32xf32, #DV>, %argb: f32) -> tensor<32xf32> {
%u = linalg.init_tensor [32] : tensor<32xf32> %u = tensor.empty() : tensor<32xf32>
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #DV>) ins(%arga: tensor<32xf32, #DV>)
outs(%u: tensor<32xf32>) { outs(%u: tensor<32xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %a, %argb : f32 %0 = arith.addf %a, %argb : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
▲ Show 20 Lines • Show All 1,497 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

Show All 21 Lines
// CHECK-LABEL: func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> { // CHECK-LABEL: func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
// CHECK: %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf64> // CHECK: %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<1024x1024xf64>
// CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<1024x1024xf64> // CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<1024x1024xf64>
// CHECK: return %[[VAL_1]] : tensor<1024x1024xf64> // CHECK: return %[[VAL_1]] : tensor<1024x1024xf64>
// CHECK: } // CHECK: }
func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> { func.func @fold_yield_arg_zero() -> tensor<1024x1024xf64> {
%cst = arith.constant 0.000000e+00 : f64 %cst = arith.constant 0.000000e+00 : f64
%0 = linalg.init_tensor [1024, 1024] : tensor<1024x1024xf64> %0 = tensor.empty() : tensor<1024x1024xf64>
%1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>, %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%cst : f64) ins(%cst : f64)
outs(%0 : tensor<1024x1024xf64>) { outs(%0 : tensor<1024x1024xf64>) {
^bb0(%a: f64, %x: f64): ^bb0(%a: f64, %x: f64):
linalg.yield %a : f64 linalg.yield %a : f64
} -> tensor<1024x1024xf64> } -> tensor<1024x1024xf64>
return %1 : tensor<1024x1024xf64> return %1 : tensor<1024x1024xf64>
} }
// CHECK-LABEL: func.func @fold_yield_direct_zero() -> tensor<32xf64> { // CHECK-LABEL: func.func @fold_yield_direct_zero() -> tensor<32xf64> {
// CHECK: %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<32xf64> // CHECK: %[[VAL_0:.*]] = arith.constant dense<0.000000e+00> : tensor<32xf64>
// CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<32xf64> // CHECK: %[[VAL_1:.*]] = bufferization.alloc_tensor() copy(%[[VAL_0]]) {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<32xf64>
// CHECK: return %[[VAL_1]] : tensor<32xf64> // CHECK: return %[[VAL_1]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @fold_yield_direct_zero() -> tensor<32xf64> { func.func @fold_yield_direct_zero() -> tensor<32xf64> {
%cst = arith.constant 0.000000e+00 : f64 %cst = arith.constant 0.000000e+00 : f64
%0 = linalg.init_tensor [32] : tensor<32xf64> %0 = tensor.empty() : tensor<32xf64>
%1 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>], %1 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%0 : tensor<32xf64>) { outs(%0 : tensor<32xf64>) {
^bb0(%x: f64): ^bb0(%x: f64):
linalg.yield %cst : f64 linalg.yield %cst : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %1 : tensor<32xf64> return %1 : tensor<32xf64>
} }
▲ Show 20 Lines • Show All 149 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir

Show All 26 Lines
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<0> : vector<8xindex> // CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<0> : vector<8xindex>
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : i64
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
// CHECK-DAG: %[[VAL_10a:.*]] = linalg.init_tensor [8] : tensor<8xi64> // CHECK-DAG: %[[VAL_10a:.*]] = tensor.empty() : tensor<8xi64>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_10a]] : memref<8xi64> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_10a]] : memref<8xi64>
// CHECK-DAG: linalg.fill ins(%[[VAL_5]] : i64) outs(%[[VAL_10]] : memref<8xi64>) // CHECK-DAG: linalg.fill ins(%[[VAL_5]] : i64) outs(%[[VAL_10]] : memref<8xi64>)
// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_3]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_3]] {
// CHECK: %[[VAL_14:.*]] = affine.min #map0(%[[VAL_13]]){{\[}}%[[VAL_12]]] // CHECK: %[[VAL_14:.*]] = affine.min #map0(%[[VAL_13]]){{\[}}%[[VAL_12]]]
// CHECK: %[[VAL_15:.*]] = vector.create_mask %[[VAL_14]] : vector<8xi1> // CHECK: %[[VAL_15:.*]] = vector.create_mask %[[VAL_14]] : vector<8xi1>
// CHECK: %[[VAL_16:.*]] = vector.maskedload %[[VAL_8]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_2]] : memref<?xindex>, vector<8xi1>, vector<8xindex> into vector<8xindex> // CHECK: %[[VAL_16:.*]] = vector.maskedload %[[VAL_8]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_2]] : memref<?xindex>, vector<8xi1>, vector<8xindex> into vector<8xindex>
// CHECK: %[[VAL_17:.*]] = vector.maskedload %[[VAL_9]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_1]] : memref<?xi64>, vector<8xi1>, vector<8xi64> into vector<8xi64> // CHECK: %[[VAL_17:.*]] = vector.maskedload %[[VAL_9]]{{\[}}%[[VAL_13]]], %[[VAL_15]], %[[VAL_1]] : memref<?xi64>, vector<8xi1>, vector<8xi64> into vector<8xi64>
// CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_16]] : vector<8xindex> to vector<8xi64> // CHECK: %[[VAL_18:.*]] = arith.index_cast %[[VAL_16]] : vector<8xindex> to vector<8xi64>
// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_18]] : vector<8xi64> // CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_17]], %[[VAL_18]] : vector<8xi64>
// CHECK: vector.scatter %[[VAL_10]]{{\[}}%[[VAL_6]]] {{\[}}%[[VAL_16]]], %[[VAL_15]], %[[VAL_19]] : memref<8xi64>, vector<8xindex>, vector<8xi1>, vector<8xi64> // CHECK: vector.scatter %[[VAL_10]]{{\[}}%[[VAL_6]]] {{\[}}%[[VAL_16]]], %[[VAL_15]], %[[VAL_19]] : memref<8xi64>, vector<8xindex>, vector<8xi1>, vector<8xi64>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<8xi64> // CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<8xi64>
// CHECK: return %[[VAL_20]] : tensor<8xi64> // CHECK: return %[[VAL_20]] : tensor<8xi64>
// CHECK: } // CHECK: }
func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> { func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> {
%init = linalg.init_tensor [8] : tensor<8xi64> %init = tensor.empty() : tensor<8xi64>
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%init: tensor<8xi64>) { outs(%init: tensor<8xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%m1 = arith.muli %a, %ii : i64 %m1 = arith.muli %a, %ii : i64
linalg.yield %m1 : i64 linalg.yield %m1 : i64
} -> tensor<8xi64> } -> tensor<8xi64>
return %r : tensor<8xi64> return %r : tensor<8xi64>
} }
// CHECK-LABEL: func @sparse_index_1d_disj( // CHECK-LABEL: func @sparse_index_1d_disj(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<8xi64> { // CHECK-SAME: %[[VAL_0:.*]]: tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<8xi64> {
// CHECK-DAG: %[[VAL_1:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex> // CHECK-DAG: %[[VAL_1:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : i64
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<8xi64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xi64>
// CHECK-DAG: %[[VAL_9a:.*]] = linalg.init_tensor [8] : tensor<8xi64> // CHECK-DAG: %[[VAL_9a:.*]] = tensor.empty() : tensor<8xi64>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_9a]] : memref<8xi64> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_9a]] : memref<8xi64>
// CHECK-DAG: linalg.fill ins(%[[VAL_3]] : i64) outs(%[[VAL_9]] : memref<8xi64>) // CHECK-DAG: linalg.fill ins(%[[VAL_3]] : i64) outs(%[[VAL_9]] : memref<8xi64>)
// CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_2]]] : memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_2]]] : memref<?xindex>
// CHECK: %[[VAL_12:.*]]:2 = scf.while (%[[VAL_13:.*]] = %[[VAL_10]], %[[VAL_14:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_12:.*]]:2 = scf.while (%[[VAL_13:.*]] = %[[VAL_10]], %[[VAL_14:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_15:.*]] = arith.cmpi ult, %[[VAL_13]], %[[VAL_11]] : index // CHECK: %[[VAL_15:.*]] = arith.cmpi ult, %[[VAL_13]], %[[VAL_11]] : index
// CHECK: scf.condition(%[[VAL_15]]) %[[VAL_13]], %[[VAL_14]] : index, index // CHECK: scf.condition(%[[VAL_15]]) %[[VAL_13]], %[[VAL_14]] : index, index
// CHECK: } do { // CHECK: } do {
Show All 22 Lines
// CHECK: %[[VAL_33:.*]] = arith.addi %[[VAL_32]], %[[VAL_1]] : vector<8xindex> // CHECK: %[[VAL_33:.*]] = arith.addi %[[VAL_32]], %[[VAL_1]] : vector<8xindex>
// CHECK: %[[VAL_34:.*]] = arith.index_cast %[[VAL_33]] : vector<8xindex> to vector<8xi64> // CHECK: %[[VAL_34:.*]] = arith.index_cast %[[VAL_33]] : vector<8xindex> to vector<8xi64>
// CHECK: vector.maskedstore %[[VAL_9]]{{\[}}%[[VAL_28]]], %[[VAL_31]], %[[VAL_34]] : memref<8xi64>, vector<8xi1>, vector<8xi64> // CHECK: vector.maskedstore %[[VAL_9]]{{\[}}%[[VAL_28]]], %[[VAL_31]], %[[VAL_34]] : memref<8xi64>, vector<8xi1>, vector<8xi64>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_35:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<8xi64> // CHECK: %[[VAL_35:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<8xi64>
// CHECK: return %[[VAL_35]] : tensor<8xi64> // CHECK: return %[[VAL_35]] : tensor<8xi64>
// CHECK: } // CHECK: }
func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> { func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>) -> tensor<8xi64> {
%init = linalg.init_tensor [8] : tensor<8xi64> %init = tensor.empty() : tensor<8xi64>
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%init: tensor<8xi64>) { outs(%init: tensor<8xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%m1 = arith.addi %a, %ii : i64 %m1 = arith.addi %a, %ii : i64
linalg.yield %m1 : i64 linalg.yield %m1 : i64
} -> tensor<8xi64> } -> tensor<8xi64>
return %r : tensor<8xi64> return %r : tensor<8xi64>
} }

mlir/test/Dialect/Tensor/canonicalize.mlir

Show First 20 Lines • Show All 1,517 Lines • ▼ Show 20 Lines
func.func @dont_fold_mismatched_parameters(%input: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> { func.func @dont_fold_mismatched_parameters(%input: tensor<1x2x2x4xf32>) -> tensor<1x2x2x4xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
%0 = tensor.extract_slice %input[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32> %0 = tensor.extract_slice %input[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x2x4xf32> to tensor<1x2x4xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
%1 = tensor.insert_slice %0 into %input[%c0, 1, %c0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32> %1 = tensor.insert_slice %0 into %input[%c0, 1, %c0, 0] [1, 1, 2, 4] [1, 1, 1, 1] : tensor<1x2x4xf32> into tensor<1x2x2x4xf32>
return %1: tensor<1x2x2x4xf32> return %1: tensor<1x2x2x4xf32>
} }
// -----
func.func @empty_canonicalize() -> (tensor<4x5x?xf32>) {
%c6 = arith.constant 6 : index
%0 = tensor.empty(%c6) : tensor<4x5x?xf32>
return %0 : tensor<4x5x?xf32>
}
// CHECK: func @empty_canonicalize
// CHECK: %[[T0:.+]] = tensor.empty() : tensor<4x5x6xf32>
// CHECK: %[[T1:.+]] = tensor.cast %[[T0]] : tensor<4x5x6xf32> to tensor<4x5x?xf32>
// CHECK: return %[[T1]]
// -----
func.func @empty_reshape_expansion(%arg0 : index) -> tensor<2x3x5x4x?x7xf32> {
%0 = tensor.empty(%arg0) : tensor<6x5x?xf32>
%1 = tensor.expand_shape %0 [[0, 1], [2], [3, 4, 5]]
: tensor<6x5x?xf32> into tensor<2x3x5x4x?x7xf32>
return %1 : tensor<2x3x5x4x?x7xf32>
}
// CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 floordiv 28)>
// CHECK: func @empty_reshape_expansion
// CHECK-SAME: %[[ARG0:.+]]: index
// CHECK-NEXT: %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
// CHECK-NEXT: %[[INIT:.+]] = tensor.empty(%[[D]])
// CHECK-NEXT: return %[[INIT]]
// -----
func.func @empty_reshape_collapse(%arg0 : index) -> tensor<6x5x?xf32> {
%0 = tensor.empty(%arg0) : tensor<2x3x5x4x?x7xf32>
%1 = tensor.collapse_shape %0 [[0, 1], [2], [3, 4, 5]]
: tensor<2x3x5x4x?x7xf32> into tensor<6x5x?xf32>
return %1 : tensor<6x5x?xf32>
}
// CHECK: #[[MAP:.+]] = affine_map<()[s0] -> (s0 * 28)>
// CHECK: func @empty_reshape_collapse
// CHECK-SAME: %[[ARG0:.+]]: index
// CHECK-NEXT: %[[D:.+]] = affine.apply #[[MAP]]()[%[[ARG0]]]
// CHECK-NEXT: %[[INIT:.+]] = tensor.empty(%[[D]])
// CHECK-NEXT: return %[[INIT]]
// -----
func.func @fold_empty_tensor_with_slice
(%arg0 : index, %arg1 : index) -> tensor<5x?x20xf32>
{
%0 = tensor.empty(%arg0) : tensor<?x10x40xf32>
%1 = tensor.extract_slice %0[0, 0, 0] [5, %arg1, 20] [1, 1, 1]
: tensor<?x10x40xf32> to tensor<5x?x20xf32>
return %1 : tensor<5x?x20xf32>
}
// CHECK: func @fold_empty_tensor_with_slice
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: index
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK: %[[T0:.+]] = tensor.empty(%[[ARG1]])
// CHECK: return %[[T0]]
// -----
func.func @fold_empty_tensor_with_cast(%arg0 : index) -> tensor<1x12xf32> {
%0 = tensor.empty(%arg0) : tensor<?x12xf32>
%1 = tensor.cast %0 : tensor<?x12xf32> to tensor<1x12xf32>
return %1 : tensor<1x12xf32>
}
// CHECK: func @fold_empty_tensor_with_cast(%[[ARG0:.+]]: index)
// CHECK: %[[T0:.+]] = tensor.empty() : tensor<1x12xf32>
// CHECK: return %[[T0]] : tensor<1x12xf32>
// -----
func.func private @some_use(%i : index, %j : index)
// CHECK-LABEL: func @empty_tensor_canonicalize
// CHECK-SAME: %[[I:.*]]: index
func.func @empty_tensor_canonicalize(%i : index) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
// CHECK-NOT: tensor.empty
%0 = tensor.empty(%i) : tensor<?x42xf32>
// CHECK-NOT: tensor.dim
%1 = tensor.dim %0, %c0: tensor<?x42xf32>
%2 = tensor.dim %0, %c1: tensor<?x42xf32>
// CHECK: %[[c42:.*]] = arith.constant 42 : index
// CHECK: call @some_use(%[[I]], %[[c42]])
call @some_use(%1, %2) : (index, index) -> ()
return
}
// -----
// CHECK-LABEL: func @rank_reducing_empty_tensor_extract
func.func @rank_reducing_empty_tensor_extract(%sz : index, %idx : index) -> tensor<2xf32> {
// CHECK: tensor.empty() : tensor<2xf32>
%a = tensor.empty(%sz) : tensor<?x2xf32>
// CHECK-NOT: extract
%r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
return %r: tensor<2xf32>
}

mlir/test/Dialect/Tensor/extract-slice-from-collapse-shape.mlir

// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-rewrite-extract-slice-from-collapse-shape %s | FileCheck %s // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-rewrite-extract-slice-from-collapse-shape %s | FileCheck %s
// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-rewrite-extract-slice-from-collapse-shape use-foreach" %s | FileCheck %s --check-prefix=FOREACH // RUN: mlir-opt -split-input-file -test-tensor-transform-patterns="test-rewrite-extract-slice-from-collapse-shape use-foreach" %s | FileCheck %s --check-prefix=FOREACH
func.func @extract_slice_static(%input: tensor<3x5x7x11xf32>) -> tensor<20x11xf32> { func.func @extract_slice_static(%input: tensor<3x5x7x11xf32>) -> tensor<20x11xf32> {
%collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32> %collapsed = tensor.collapse_shape %input [[0, 1, 2], [3]] : tensor<3x5x7x11xf32> into tensor<105x11xf32>
%slice = tensor.extract_slice %collapsed [0, 0] [20, 11] [1, 1] : tensor<105x11xf32> to tensor<20x11xf32> %slice = tensor.extract_slice %collapsed [0, 0] [20, 11] [1, 1] : tensor<105x11xf32> to tensor<20x11xf32>
return %slice : tensor<20x11xf32> return %slice : tensor<20x11xf32>
} }
// CHECK: func.func @extract_slice_static(%[[arg0:.+]]: // CHECK: func.func @extract_slice_static(%[[arg0:.+]]:
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[c20:.+]] = arith.constant 20 : index // CHECK-DAG: %[[c20:.+]] = arith.constant 20 : index
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index // CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
// CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index // CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index
// CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index // CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index
// CHECK-DAG: %[[init:.+]] = linalg.init_tensor [20, 11] : // CHECK-DAG: %[[init:.+]] = tensor.empty() : tensor<20x11xf32>
// CHECK-DAG: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c20]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]]) // CHECK-DAG: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c20]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])
// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (%[[c3]], %[[c5]], %[[c7]] // CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (%[[c3]], %[[c5]], %[[c7]]
// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] : // CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] :
// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : // CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :
// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 11] [1, 1] : // CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 11] [1, 1] :
// CHECK: scf.yield %[[update]] : // CHECK: scf.yield %[[update]] :
// CHECK: return %[[tile]] // CHECK: return %[[tile]]
// FOREACH: func.func @extract_slice_static(%[[arg0:.+]]: // FOREACH: func.func @extract_slice_static(%[[arg0:.+]]:
// FOREACH-DAG: %[[c20:.+]] = arith.constant 20 : index // FOREACH-DAG: %[[c20:.+]] = arith.constant 20 : index
// FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index // FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index
// FOREACH-DAG: %[[c5:.+]] = arith.constant 5 : index // FOREACH-DAG: %[[c5:.+]] = arith.constant 5 : index
// FOREACH-DAG: %[[c7:.+]] = arith.constant 7 : index // FOREACH-DAG: %[[c7:.+]] = arith.constant 7 : index
// FOREACH-DAG: %[[init:.+]] = linalg.init_tensor [20, 11] : // FOREACH-DAG: %[[init:.+]] = tensor.empty() : tensor<20x11xf32>
// FOREACH: %[[tile:.+]] = scf.foreach_thread (%[[iv:.+]]) in (%[[c20]]) shared_outs(%[[dest:.+]] = %[[init]]) // FOREACH: %[[tile:.+]] = scf.foreach_thread (%[[iv:.+]]) in (%[[c20]]) shared_outs(%[[dest:.+]] = %[[init]])
// FOREACH: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (%[[c3]], %[[c5]], %[[c7]] // FOREACH: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[iv]] into (%[[c3]], %[[c5]], %[[c7]]
// FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] : // FOREACH: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 11] [1, 1, 1, 1] :
// FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : // FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :
// FOREACH: perform_concurrently // FOREACH: perform_concurrently
// FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[iv]], 0] [1, 11] [1, 1] : // FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[iv]], 0] [1, 11] [1, 1] :
// FOREACH: return %[[tile]] // FOREACH: return %[[tile]]
Show All 9 Lines
// CHECK: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2 + 13)> // CHECK: #[[$map0:.+]] = affine_map<(d0) -> (d0 * 2 + 13)>
// CHECK: func.func @extract_slice_static_strided(%[[arg0:.+]]: // CHECK: func.func @extract_slice_static_strided(%[[arg0:.+]]:
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[c10:.+]] = arith.constant 10 : index // CHECK-DAG: %[[c10:.+]] = arith.constant 10 : index
// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index // CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
// CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index // CHECK-DAG: %[[c5:.+]] = arith.constant 5 : index
// CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index // CHECK-DAG: %[[c7:.+]] = arith.constant 7 : index
// CHECK: %[[init:.+]] = linalg.init_tensor [10, 5] : // CHECK: %[[init:.+]] = tensor.empty() : tensor<10x5xf32>
// CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c10]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]]) // CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[c10]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])
// CHECK: %[[inputIv:.+]] = affine.apply #[[$map0]](%[[iv]]) // CHECK: %[[inputIv:.+]] = affine.apply #[[$map0]](%[[iv]])
// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (%[[c3]], %[[c5]], %[[c7]] // CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (%[[c3]], %[[c5]], %[[c7]]
// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] : // CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] :
// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : // CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :
// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] : // CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] :
// CHECK: scf.yield %[[update]] : // CHECK: scf.yield %[[update]] :
// CHECK: return %[[tile]] // CHECK: return %[[tile]]
Show All 9 Lines
} }
// CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)> // CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
// CHECK: func.func @extract_slice_dynamic(%[[arg0:.+]]: tensor<{{.*}}>, %[[lb:.+]]: index, %[[sz:.+]]: index) // CHECK: func.func @extract_slice_dynamic(%[[arg0:.+]]: tensor<{{.*}}>, %[[lb:.+]]: index, %[[sz:.+]]: index)
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index // CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index // CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
// CHECK: %[[init:.+]] = linalg.init_tensor [%[[sz]], 5] : tensor<?x5xf32> // CHECK: %[[init:.+]] = tensor.empty(%[[sz]]) : tensor<?x5xf32>
// CHECK-DAG: %[[d1:.+]] = tensor.dim %arg0, %[[c1]] : tensor<3x?x?x11xf32> // CHECK-DAG: %[[d1:.+]] = tensor.dim %arg0, %[[c1]] : tensor<3x?x?x11xf32>
// CHECK-DAG: %[[d2:.+]] = tensor.dim %arg0, %[[c2]] : tensor<3x?x?x11xf32> // CHECK-DAG: %[[d2:.+]] = tensor.dim %arg0, %[[c2]] : tensor<3x?x?x11xf32>
// CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[sz]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]]) // CHECK: %[[tile:.+]] = scf.for %[[iv:.+]] = %[[c0]] to %[[sz]] step %[[c1]] iter_args(%[[iterArg:.+]] = %[[init]])
// CHECK: %[[inputIv:.+]] = affine.apply #[[map0]](%[[iv]])[%[[lb]]] // CHECK: %[[inputIv:.+]] = affine.apply #[[map0]](%[[iv]])[%[[lb]]]
// CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (%[[c3]], %[[d1]], %[[d2]]) : // CHECK: %[[multiIndex:.+]]:3 = affine.delinearize_index %[[inputIv]] into (%[[c3]], %[[d1]], %[[d2]]) :
// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] : // CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex]]#0, %[[multiIndex]]#1, %[[multiIndex]]#2, 0] [1, 1, 1, 5] [1, 1, 1, 2] :
// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} : // CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3]{{\]}} :
// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] : // CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg]][%[[iv]], 0] [1, 5] [1, 1] :
Show All 12 Lines
// CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)> // CHECK: #[[map0:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// CHECK: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index) // CHECK: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)
// CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[c1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index // CHECK-DAG: %[[c2:.+]] = arith.constant 2 : index
// CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index // CHECK-DAG: %[[c3:.+]] = arith.constant 3 : index
// CHECK-DAG: %[[c4:.+]] = arith.constant 4 : index // CHECK-DAG: %[[c4:.+]] = arith.constant 4 : index
// CHECK-DAG: %[[c11:.+]] = arith.constant 11 : index // CHECK-DAG: %[[c11:.+]] = arith.constant 11 : index
// CHECK: %[[init:.+]] = linalg.init_tensor [%[[sz1]], %[[sz2]]] : tensor<?x?xf32> // CHECK: %[[init:.+]] = tensor.empty(%[[sz1]], %[[sz2]]) : tensor<?x?xf32>
// CHECK-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : // CHECK-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] :
// CHECK-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : // CHECK-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] :
// CHECK-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] : // CHECK-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :
// CHECK: %[[tile1:.+]] = scf.for %[[iv1:.+]] = %[[c0]] to %[[sz1]] step %[[c1]] iter_args(%[[iterArg1:.+]] = %[[init]]) // CHECK: %[[tile1:.+]] = scf.for %[[iv1:.+]] = %[[c0]] to %[[sz1]] step %[[c1]] iter_args(%[[iterArg1:.+]] = %[[init]])
// CHECK: %[[tile2:.+]] = scf.for %[[iv2:.+]] = %[[c0]] to %[[sz2]] step %[[c1]] iter_args(%[[iterArg2:.+]] = %[[iterArg1]]) // CHECK: %[[tile2:.+]] = scf.for %[[iv2:.+]] = %[[c0]] to %[[sz2]] step %[[c1]] iter_args(%[[iterArg2:.+]] = %[[iterArg1]])
// CHECK: %[[inputIv1:.+]] = affine.apply #[[map0:.+]](%[[iv1]])[%[[lb1]]] // CHECK: %[[inputIv1:.+]] = affine.apply #[[map0:.+]](%[[iv1]])[%[[lb1]]]
// CHECK: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[inputIv1]] into (%[[c3]], %[[d1]], %[[d2]]) : // CHECK: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[inputIv1]] into (%[[c3]], %[[d1]], %[[d2]]) :
// CHECK: %[[inputIv2:.+]] = affine.apply #[[map0:.+]](%[[iv2]])[%[[lb2]]] // CHECK: %[[inputIv2:.+]] = affine.apply #[[map0:.+]](%[[iv2]])[%[[lb2]]]
// CHECK: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[inputIv2]] into (%[[c11]], %[[d4]]) : // CHECK: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[inputIv2]] into (%[[c11]], %[[d4]]) :
// CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : // CHECK: %[[slice:.+]] = tensor.extract_slice %[[arg0]][%[[multiIndex1]]#0, %[[multiIndex1]]#1, %[[multiIndex1]]#2, %[[multiIndex2]]#0, %[[multiIndex2]]#1] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] :
// CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} : // CHECK: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} :
// CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg2]][%[[iv1]], %[[iv2]]] [1, 1] [1, 1] : // CHECK: %[[update:.+]] = tensor.insert_slice %[[sliceFlat]] into %[[iterArg2]][%[[iv1]], %[[iv2]]] [1, 1] [1, 1] :
// CHECK: scf.yield %[[update]] : // CHECK: scf.yield %[[update]] :
// CHECK: scf.yield %[[tile2]] : // CHECK: scf.yield %[[tile2]] :
// CHECK: return %[[tile1]] : // CHECK: return %[[tile1]] :
// FOREACH: #[[map1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)> // FOREACH: #[[map1:.+]] = affine_map<(d0)[s0] -> (d0 + s0)>
// FOREACH: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index) // FOREACH: func.func @extract_slice_dynamic_multidim(%[[arg0:.+]]: tensor<3x?x?x11x?xf32>, %[[lb1:.+]]: index, %[[sz1:.+]]: index, %[[lb2:.+]]: index, %[[sz2:.+]]: index)
// FOREACH-DAG: %[[c1:.+]] = arith.constant 1 : index // FOREACH-DAG: %[[c1:.+]] = arith.constant 1 : index
// FOREACH-DAG: %[[c2:.+]] = arith.constant 2 : index // FOREACH-DAG: %[[c2:.+]] = arith.constant 2 : index
// FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index // FOREACH-DAG: %[[c3:.+]] = arith.constant 3 : index
// FOREACH-DAG: %[[c4:.+]] = arith.constant 4 : index // FOREACH-DAG: %[[c4:.+]] = arith.constant 4 : index
// FOREACH-DAG: %[[c11:.+]] = arith.constant 11 : index // FOREACH-DAG: %[[c11:.+]] = arith.constant 11 : index
// FOREACH: %[[init:.+]] = linalg.init_tensor [%[[sz1]], %[[sz2]]] : tensor<?x?xf32> // FOREACH: %[[init:.+]] = tensor.empty(%[[sz1]], %[[sz2]]) : tensor<?x?xf32>
// FOREACH-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] : // FOREACH-DAG: %[[d1:.+]] = tensor.dim %[[arg0]], %[[c1]] :
// FOREACH-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] : // FOREACH-DAG: %[[d2:.+]] = tensor.dim %[[arg0]], %[[c2]] :
// FOREACH-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] : // FOREACH-DAG: %[[d4:.+]] = tensor.dim %[[arg0]], %[[c4]] :
// FOREACH: %[[tile1:.+]] = scf.foreach_thread (%[[tid1:.+]], %[[tid2:.+]]) in (%[[sz1]], %[[sz2]]) shared_outs(%[[dest:.+]] = %[[init]]) // FOREACH: %[[tile1:.+]] = scf.foreach_thread (%[[tid1:.+]], %[[tid2:.+]]) in (%[[sz1]], %[[sz2]]) shared_outs(%[[dest:.+]] = %[[init]])
// FOREACH-DAG: %[[iv1:.+]] = affine.apply #[[map1]](%[[tid1]])[%[[lb1]]] // FOREACH-DAG: %[[iv1:.+]] = affine.apply #[[map1]](%[[tid1]])[%[[lb1]]]
// FOREACH: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[iv1]] into (%[[c3]], %[[d1]], %[[d2]]) : // FOREACH: %[[multiIndex1:.+]]:3 = affine.delinearize_index %[[iv1]] into (%[[c3]], %[[d1]], %[[d2]]) :
// FOREACH-DAG: %[[iv2:.+]] = affine.apply #[[map1]](%[[tid2]])[%[[lb2]]] // FOREACH-DAG: %[[iv2:.+]] = affine.apply #[[map1]](%[[tid2]])[%[[lb2]]]
// FOREACH: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[iv2]] into (%[[c11]], %[[d4]]) : // FOREACH: %[[multiIndex2:.+]]:2 = affine.delinearize_index %[[iv2]] into (%[[c11]], %[[d4]]) :
Show All 20 Lines
// ----- // -----
// CHECK: @no_sliced_linearized_dims(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index // CHECK: @no_sliced_linearized_dims(%[[arg0:.+]]: tensor<{{.*}}>, %[[arg1:.+]]: index, %[[arg2:.+]]: index
func.func @no_sliced_linearized_dims(%input: tensor<30x11x100xf32>, %offt: index, %size: index) -> tensor<330x?xf32> { func.func @no_sliced_linearized_dims(%input: tensor<30x11x100xf32>, %offt: index, %size: index) -> tensor<330x?xf32> {
%collapsed = tensor.collapse_shape %input [[0, 1], [2]] : tensor<30x11x100xf32> into tensor<330x100xf32> %collapsed = tensor.collapse_shape %input [[0, 1], [2]] : tensor<30x11x100xf32> into tensor<330x100xf32>
%slice = tensor.extract_slice %collapsed [0, %offt] [330, %size] [1, 1] : tensor<330x100xf32> to tensor<330x?xf32> %slice = tensor.extract_slice %collapsed [0, %offt] [330, %size] [1, 1] : tensor<330x100xf32> to tensor<330x?xf32>
// CHECK-NOT: scf.for // CHECK-NOT: scf.for
// CHECK: %[[init:.+]] = linalg.init_tensor [330, %[[arg2]]] // CHECK: %[[init:.+]] = tensor.empty(%[[arg2]])
// CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, %[[arg1]]] [30, 11, %[[arg2]]] [1, 1, 1] // CHECK: %[[e:.+]] = tensor.extract_slice %[[arg0]][0, 0, %[[arg1]]] [30, 11, %[[arg2]]] [1, 1, 1]
// CHECK: %[[c:.+]] = tensor.collapse_shape %[[e]] {{\[}}[0, 1], [2]] // CHECK: %[[c:.+]] = tensor.collapse_shape %[[e]] {{\[}}[0, 1], [2]]
// CHECK: %[[res:.+]] = tensor.insert_slice %[[c]] into %[[init]] // CHECK: %[[res:.+]] = tensor.insert_slice %[[c]] into %[[init]]
// CHECK: return %[[res]] // CHECK: return %[[res]]
return %slice : tensor<330x?xf32> return %slice : tensor<330x?xf32>
} }

mlir/test/Dialect/Tensor/invalid.mlir

Show First 20 Lines • Show All 508 Lines • ▼ Show 20 Lines
func.func @scatter_wrong_result_type( func.func @scatter_wrong_result_type(
%source : tensor<f32>, %source : tensor<f32>,
%dest : tensor<4x5x6xf32>, %indices: tensor<1x2x3xindex>) { %dest : tensor<4x5x6xf32>, %indices: tensor<1x2x3xindex>) {
// expected-error@+1 {{source type mismatch: expected 'tensor<1x2x1x5x1xf32>' or its rank-reduced variant 'tensor<1x2x5xf32>' (got: 'tensor<f32>')}} // expected-error@+1 {{source type mismatch: expected 'tensor<1x2x1x5x1xf32>' or its rank-reduced variant 'tensor<1x2x5xf32>' (got: 'tensor<f32>')}}
%out = tensor.scatter %source into %dest[%indices] scatter_dims([0, 2]) unique: %out = tensor.scatter %source into %dest[%indices] scatter_dims([0, 2]) unique:
(tensor<f32>, tensor<4x5x6xf32>, tensor<1x2x3xindex>) -> tensor<1x2x1xf32> (tensor<f32>, tensor<4x5x6xf32>, tensor<1x2x3xindex>) -> tensor<1x2x1xf32>
return return
} }
// -----
func.func @empty_wrong_number_of_operands(%sz : index) {
// expected-error@+1 {{incorrect number of dynamic sizes, has 1, expected 2}}
%out = tensor.empty(%sz) : tensor<2x?x?x5xf32>
return
}

mlir/test/Dialect/Tensor/ops.mlir

// RUN: mlir-opt <%s | mlir-opt | FileCheck %s // RUN: mlir-opt <%s | mlir-opt | FileCheck %s
// CHECK-LABEL: func @cast( // CHECK-LABEL: func @cast(
func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor<?x?xf32>) { func.func @cast(%arg0: tensor<*xf32>, %arg1 : tensor<4x4xf32>, %arg2: tensor<?x?xf32>) {
// CHECK: tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32> // CHECK: tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32>
%0 = tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32> %0 = tensor.cast %arg0 : tensor<*xf32> to tensor<?x?xf32>
// CHECK: tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32> // CHECK: tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32>
%1 = tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32> %1 = tensor.cast %arg1 : tensor<4x4xf32> to tensor<*xf32>
// CHECK: tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32> // CHECK: tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32>
%2 = tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32> %2 = tensor.cast %arg2 : tensor<?x?xf32> to tensor<4x?xf32>
// CHECK: tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32> // CHECK: tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32>
%3 = tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32> %3 = tensor.cast %2 : tensor<4x?xf32> to tensor<?x?xf32>
return return
} }
// CHECK-LABEL: func @empty(
// CHECK-SAME: %[[sz:.*]]: index
func.func @empty(%sz: index) -> tensor<5x?x6xf32> {
// CHECK: tensor.empty(%[[sz]]) : tensor<5x?x6xf32>
%0 = tensor.empty(%sz) : tensor<5x?x6xf32>
return %0 : tensor<5x?x6xf32>
}
// CHECK-LABEL: func @extract( // CHECK-LABEL: func @extract(
// CHECK-SAME: %[[TENSOR:.*]]: tensor<?x?x?xf32>, // CHECK-SAME: %[[TENSOR:.*]]: tensor<?x?x?xf32>,
// CHECK-SAME: %[[INDEX:.*]]: index) { // CHECK-SAME: %[[INDEX:.*]]: index) {
func.func @extract(%arg0: tensor<?x?x?xf32>, %arg1: index) { func.func @extract(%arg0: tensor<?x?x?xf32>, %arg1: index) {
// CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32> // CHECK: tensor.extract %[[TENSOR]][%[[INDEX]], %[[INDEX]], %[[INDEX]]] : tensor<?x?x?xf32>
%0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor<?x?x?xf32> %0 = tensor.extract %arg0[%arg1, %arg1, %arg1] : tensor<?x?x?xf32>
return return
} }
▲ Show 20 Lines • Show All 258 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-padtensor.mlir

// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns \ // RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns \
// RUN: -linalg-init-tensor-to-alloc-tensor -linalg-bufferize -arith-bufferize \ // RUN: -empty-tensor-to-alloc-tensor -linalg-bufferize -arith-bufferize \
// RUN: -bufferization-bufferize -tensor-bufferize -func-bufferize \ // RUN: -bufferization-bufferize -tensor-bufferize -func-bufferize \
// RUN: -finalizing-bufferize -buffer-deallocation \ // RUN: -finalizing-bufferize -buffer-deallocation \
// RUN: -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \
// RUN: mlir-cpu-runner -e main -entry-point-result=void \ // RUN: mlir-cpu-runner -e main -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_lib_dir/libmlir_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
Show All 24 Lines

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

Show All 33 Lines
from mlir import ir from mlir import ir
from mlir import runtime from mlir import runtime
from mlir.dialects import arith from mlir.dialects import arith
from mlir.dialects import bufferization from mlir.dialects import bufferization
from mlir.dialects import builtin from mlir.dialects import builtin
from mlir.dialects import func from mlir.dialects import func
from mlir.dialects import linalg from mlir.dialects import linalg
from mlir.dialects import sparse_tensor from mlir.dialects import sparse_tensor
from mlir.dialects import tensor
from mlir.dialects.linalg.opdsl import lang from mlir.dialects.linalg.opdsl import lang
from . import mlir_pytaco_utils as utils from . import mlir_pytaco_utils as utils
# TACO naming prefixes. # TACO naming prefixes.
_TACO_INDEX_PREFIX = "i" _TACO_INDEX_PREFIX = "i"
_TACO_TENSOR_PREFIX = "A" _TACO_TENSOR_PREFIX = "A"
▲ Show 20 Lines • Show All 844 Lines • ▼ Show 20 Lines def __post_init__(self) -> None:
"""Verifies the integrity of the attribute values.""" """Verifies the integrity of the attribute values."""
assert len(self.dst_indices) == len(self.dst_dims) assert len(self.dst_indices) == len(self.dst_dims)
def emit_tensor_init(self) -> ir.RankedTensorType: def emit_tensor_init(self) -> ir.RankedTensorType:
"""Returns an initialization for the destination tensor.""" """Returns an initialization for the destination tensor."""
if self.dst_format is None or self.dst_format.rank() == 0: if self.dst_format is None or self.dst_format.rank() == 0:
# Initialize the dense tensor. # Initialize the dense tensor.
ir_type = _mlir_type_from_taco_type(self.dst_dtype) ir_type = _mlir_type_from_taco_type(self.dst_dtype)
tensor = linalg.InitTensorOp(self.dst_dims, ir_type).result empty = tensor.EmptyOp(self.dst_dims, ir_type).result
zero = arith.ConstantOp(ir_type, 0.0) zero = arith.ConstantOp(ir_type, 0.0)
return linalg.fill(zero, outs=[tensor]) return linalg.fill(zero, outs=[empty])
# Initialize the sparse tensor. # Initialize the sparse tensor.
mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims, mlir_type = _mlir_tensor_type(self.dst_dtype, self.dst_dims,
self.dst_format.mlir_tensor_attr()) self.dst_format.mlir_tensor_attr())
index_type = ir.IndexType.get() index_type = ir.IndexType.get()
return bufferization.AllocTensorOp(mlir_type, [], None, None) return bufferization.AllocTensorOp(mlir_type, [], None, None)
▲ Show 20 Lines • Show All 276 Lines • ▼ Show 20 Lines def from_array(array: np.ndarray) -> "Tensor":
Returns: Returns:
A Tensor object. A Tensor object.
Raises: Raises:
ValueError if the data type of the numpy array is not supported. ValueError if the data type of the numpy array is not supported.
""" """
if array.dtype != np.float32 and array.dtype != np.float64: if array.dtype != np.float32 and array.dtype != np.float64:
raise ValueError(f"Expected floating point value type: {array.dtype}.") raise ValueError(f"Expected floating point value type: {array.dtype}.")
tensor = Tensor( t = Tensor(
array.shape, array.shape,
dtype=_nptype_to_taco_type(array.dtype.type), dtype=_nptype_to_taco_type(array.dtype.type),
is_dense=True) is_dense=True)
tensor._dense_storage = np.copy(array) t._dense_storage = np.copy(array)
return tensor return t
@staticmethod @staticmethod
def from_coo( def from_coo(
coordinates: List[Tuple[int, ...]], coordinates: List[Tuple[int, ...]],
values: List[_AnyRuntimeType], values: List[_AnyRuntimeType],
fmt: Format, fmt: Format,
dtype: DType, dtype: DType,
) -> "Tensor": ) -> "Tensor":
Show All 18 Lines ) -> "Tensor":
rank = fmt.rank() rank = fmt.rank()
assert all(len(c) == rank and _all_instance_of(c, int) for c in coordinates) assert all(len(c) == rank and _all_instance_of(c, int) for c in coordinates)
# Find the maximum coordinate value for each dimension. # Find the maximum coordinate value for each dimension.
max_coordinate = list(map(max, zip(*coordinates))) max_coordinate = list(map(max, zip(*coordinates)))
# The size of each dimension is one more that such a maximum coordinate # The size of each dimension is one more that such a maximum coordinate
# value. # value.
shape = [c + 1 for c in max_coordinate] shape = [c + 1 for c in max_coordinate]
tensor = Tensor(shape, fmt, dtype=dtype) t = Tensor(shape, fmt, dtype=dtype)
tensor._coords = coordinates t._coords = coordinates
tensor._values = values t._values = values
return tensor return tensor
@staticmethod @staticmethod
def from_file( def from_file(
filename: str, filename: str,
fmt: Format, fmt: Format,
dtype: DType, dtype: DType,
) -> "Tensor": ) -> "Tensor":
"""Constructs a sparse tensor using the COO-flavored values from a file. """Constructs a sparse tensor using the COO-flavored values from a file.
Args: Args:
filename: A string for the name of the file that contains the sparse filename: A string for the name of the file that contains the sparse
tensor data. tensor data.
fmt: The tensor storage format. fmt: The tensor storage format.
dtype: The tensor element data type. dtype: The tensor element data type.
Returns: Returns:
A tensor with the given non-zero values and storage format. The tensor A tensor with the given non-zero values and storage format. The tensor
value is stored as an MLIR sparse tensor. value is stored as an MLIR sparse tensor.
""" """
sparse_tensor, shape = utils.create_sparse_tensor(filename, sparse_tensor, shape = utils.create_sparse_tensor(filename,
fmt.format_pack.formats, fmt.format_pack.formats,
_dtype_to_mlir_str(dtype)) _dtype_to_mlir_str(dtype))
tensor = Tensor(shape.tolist(), fmt, dtype=dtype) t = Tensor(shape.tolist(), fmt, dtype=dtype)
tensor._set_packed_sparse_tensor(sparse_tensor) t._set_packed_sparse_tensor(sparse_tensor)
return tensor return t
def to_file(self, filename: str) -> None: def to_file(self, filename: str) -> None:
"""Output the tensor value to a file. """Output the tensor value to a file.
This method evaluates any pending assignment to the tensor and outputs the This method evaluates any pending assignment to the tensor and outputs the
tensor value. tensor value.
Args: Args:
▲ Show 20 Lines • Show All 920 Lines • Show Last 20 Lines

mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-interface.mlir

// RUN: mlir-opt -test-tiling-interface=tile-consumer-and-fuse-producer-using-scf-for -split-input-file %s | FileCheck %s // RUN: mlir-opt -test-tiling-interface=tile-consumer-and-fuse-producer-using-scf-for -split-input-file %s | FileCheck %s
func.func @gemm_fill_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @gemm_fill_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
%gemm = linalg.matmul {__internal_linalg_transform__ = "fusion"} %gemm = linalg.matmul {__internal_linalg_transform__ = "fusion"}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
return %gemm : tensor<?x?xf32> return %gemm : tensor<?x?xf32>
} }
// CHECK: func.func @gemm_fill_fusion( // CHECK: func.func @gemm_fill_fusion(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>) // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]]) // CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]])
// CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]]) // CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
// CHECK: %[[FILL_TILE:.+]] = linalg.fill // CHECK: %[[FILL_TILE:.+]] = linalg.fill
// CHECK-SAME: outs(%[[INIT_TILE]] : // CHECK-SAME: outs(%[[INIT_TILE]] :
// CHECK: %[[GEMM_TILE:.+]] = linalg.matmul // CHECK: %[[GEMM_TILE:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS_TILE]], %[[RHS_TILE]] : // CHECK-SAME: ins(%[[LHS_TILE]], %[[RHS_TILE]] :
// CHECK-SAME: outs(%[[FILL_TILE]] : // CHECK-SAME: outs(%[[FILL_TILE]] :
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]] // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]]
// CHECK: scf.yield %[[INSERT]] // CHECK: scf.yield %[[INSERT]]
// ----- // -----
func.func @gemm_generic_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, func.func @gemm_generic_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?xf32>) -> tensor<?x?xf32> { %arg2 : tensor<?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%init = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32> %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<?x?xf32>) -> tensor<?x?xf32>
%gemm = linalg.matmul %gemm = linalg.matmul
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
%generic = linalg.generic { %generic = linalg.generic {
__internal_linalg_transform__ = "fusion", __internal_linalg_transform__ = "fusion",
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%gemm, %arg2 : tensor<?x?xf32>, tensor<?xf32>) outs(%init : tensor<?x?xf32>) { ins(%gemm, %arg2 : tensor<?x?xf32>, tensor<?xf32>) outs(%init : tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%add = arith.addf %b0, %b1 : f32 %add = arith.addf %b0, %b1 : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %generic : tensor<?x?xf32> return %generic : tensor<?x?xf32>
} }
// CHECK: func.func @gemm_generic_fusion( // CHECK: func.func @gemm_generic_fusion(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>, // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>,
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>) // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?xf32>)
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]]) // CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]])
// CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]]) // CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV0]], %[[IV1]]]
// CHECK: %[[FILL_TILE:.+]] = linalg.fill // CHECK: %[[FILL_TILE:.+]] = linalg.fill
Show All 12 Lines
// ----- // -----
func.func @gemm_gemm_fusion(%lhs0 : tensor<?x?xf32>, %rhs0 : tensor<?x?xf32>, %rhs1 : tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @gemm_gemm_fusion(%lhs0 : tensor<?x?xf32>, %rhs0 : tensor<?x?xf32>, %rhs1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%d0 = tensor.dim %lhs0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %lhs0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %rhs0, %c1 : tensor<?x?xf32> %d1 = tensor.dim %rhs0, %c1 : tensor<?x?xf32>
%init0 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%fill0 = linalg.fill ins(%cst : f32) outs(%init0 : tensor<?x?xf32>) -> tensor<?x?xf32> %fill0 = linalg.fill ins(%cst : f32) outs(%init0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%gemm0 = linalg.matmul %gemm0 = linalg.matmul
ins(%lhs0, %rhs0 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%fill0 : tensor<?x?xf32>) -> tensor<?x?xf32> ins(%lhs0, %rhs0 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%fill0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%d2 = tensor.dim %rhs1, %c1 : tensor<?x?xf32> %d2 = tensor.dim %rhs1, %c1 : tensor<?x?xf32>
%init1 = linalg.init_tensor [%d0, %d2] : tensor<?x?xf32> %init1 = tensor.empty(%d0, %d2) : tensor<?x?xf32>
%fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<?x?xf32>) -> tensor<?x?xf32> %fill1 = linalg.fill ins(%cst : f32) outs(%init1 : tensor<?x?xf32>) -> tensor<?x?xf32>
%gemm1 = linalg.matmul {__internal_linalg_transform__ = "gemm_fusion"} %gemm1 = linalg.matmul {__internal_linalg_transform__ = "gemm_fusion"}
ins(%gemm0, %rhs1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%fill1 : tensor<?x?xf32>) -> tensor<?x?xf32> ins(%gemm0, %rhs1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%fill1 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %gemm1 : tensor<?x?xf32> return %gemm1 : tensor<?x?xf32>
} }
// CHECK: func.func @gemm_gemm_fusion( // CHECK: func.func @gemm_gemm_fusion(
// CHECK-SAME: %[[LHS0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[LHS0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[RHS0:[a-zA-Z0-9]+]]: tensor<?x?xf32>, // CHECK-SAME: %[[RHS0:[a-zA-Z0-9]+]]: tensor<?x?xf32>,
// CHECK-SAME: %[[RHS1:[a-zA-Z0-9]+]]: tensor<?x?xf32>) // CHECK-SAME: %[[RHS1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[LHS0]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[LHS0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[RHS0]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[RHS0]], %[[C1]]
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT0:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK-DAG: %[[D2:.+]] = tensor.dim %[[RHS1]], %[[C1]] // CHECK-DAG: %[[D2:.+]] = tensor.dim %[[RHS1]], %[[C1]]
// CHECK: %[[INIT1:.+]] = linalg.init_tensor [%[[D0]], %[[D2]]] // CHECK: %[[INIT1:.+]] = tensor.empty(%[[D0]], %[[D2]])
// CHECK: scf.for %[[IV:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG:.+]] = %[[INIT1]]) // CHECK-SAME: iter_args(%[[ITERARG:.+]] = %[[INIT1]])
// CHECK-DAG: %[[LHS0_TILE:.+]] = tensor.extract_slice %[[LHS0]][%[[IV]], 0] // CHECK-DAG: %[[LHS0_TILE:.+]] = tensor.extract_slice %[[LHS0]][%[[IV]], 0]
// CHECK-DAG: %[[RHS0_TILE:.+]] = tensor.extract_slice %[[RHS0]][0, 0] // CHECK-DAG: %[[RHS0_TILE:.+]] = tensor.extract_slice %[[RHS0]][0, 0]
// CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[INIT0]][%[[IV]], 0] // CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[INIT0]][%[[IV]], 0]
// CHECK: %[[FILL0_TILE:.+]] = linalg.fill // CHECK: %[[FILL0_TILE:.+]] = linalg.fill
// CHECK-SAME: outs(%[[INIT0_TILE]] : // CHECK-SAME: outs(%[[INIT0_TILE]] :
// CHECK: %[[GEMM0_TILE:.+]] = linalg.matmul // CHECK: %[[GEMM0_TILE:.+]] = linalg.matmul
Show All 12 Lines
// ----- // -----
func.func @gemm_transpose_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @gemm_transpose_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%init0 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %init0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%fill = linalg.fill ins(%cst : f32) outs(%init0 : tensor<?x?xf32>) -> tensor<?x?xf32> %fill = linalg.fill ins(%cst : f32) outs(%init0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%gemm = linalg.matmul %gemm = linalg.matmul
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
%init1 = linalg.init_tensor [%d1, %d0] : tensor<?x?xf32> %init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32>
%transpose = linalg.generic { %transpose = linalg.generic {
__internal_linalg_transform__ = "fusion", __internal_linalg_transform__ = "fusion",
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%gemm : tensor<?x?xf32>) outs(%init1 : tensor<?x?xf32>) { ins(%gemm : tensor<?x?xf32>) outs(%init1 : tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32): ^bb0(%b0 : f32, %b1 : f32):
linalg.yield %b0 : f32 linalg.yield %b0 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %transpose : tensor<?x?xf32> return %transpose : tensor<?x?xf32>
} }
// CHECK: func.func @gemm_transpose_fusion( // CHECK: func.func @gemm_transpose_fusion(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>) // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]] // CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]
// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C1]] // CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG1]], %[[C1]]
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [%[[D0]], %[[D1]]] // CHECK-DAG: %[[INIT0:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [%[[D1]], %[[D0]]] // CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT1]]) // CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT1]])
// CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]]) // CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[INIT0]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[INIT0]][%[[IV0]], %[[IV1]]]
// CHECK: %[[FILL_TILE:.+]] = linalg.fill // CHECK: %[[FILL_TILE:.+]] = linalg.fill
Show All 11 Lines
// ----- // -----
func.func @interchange_matmul_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @interchange_matmul_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
%d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%0 = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32> %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32> %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%2 = linalg.matmul %2 = linalg.matmul
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%1 : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%1 : tensor<?x?xf32>) -> tensor<?x?xf32>
%3 = linalg.generic { %3 = linalg.generic {
__internal_linalg_transform__ = "gemm_interchange_fusion", __internal_linalg_transform__ = "gemm_interchange_fusion",
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%2 : tensor<?x?xf32>) outs(%0 : tensor<?x?xf32>) { ins(%2 : tensor<?x?xf32>) outs(%0 : tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32): ^bb0(%b0 : f32, %b1 : f32):
%4 = arith.addf %b0, %b0 : f32 %4 = arith.addf %b0, %b0 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %3 : tensor<?x?xf32> return %3 : tensor<?x?xf32>
} }
// CHECK: func.func @interchange_matmul_fusion( // CHECK: func.func @interchange_matmul_fusion(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>) // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
// CHECK: %[[INIT:.+]] = linalg.init_tensor // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV0:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]]) // CHECK-SAME: iter_args(%[[ITERARG0:.+]] = %[[INIT]])
// CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] = // CHECK: scf.for %[[IV1:[a-zA-Z0-9]+]] =
// CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]]) // CHECK-SAME: iter_args(%[[ITERARG1:.+]] = %[[ITERARG0]])
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV0]]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV0]]]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV1]], %[[IV0]]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT]][%[[IV1]], %[[IV0]]]
// CHECK: %[[FILL_TILE:.+]] = linalg.fill // CHECK: %[[FILL_TILE:.+]] = linalg.fill
Show All 15 Linesfunc.func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg2, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg2, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg2, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg2, %c1 : tensor<?x?xf32>
%2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%3 = tensor.dim %2, %c0 : tensor<?x?xf32> %3 = tensor.dim %2, %c0 : tensor<?x?xf32>
%4 = tensor.dim %2, %c1 : tensor<?x?xf32> %4 = tensor.dim %2, %c1 : tensor<?x?xf32>
%5 = linalg.init_tensor [%3, %4] : tensor<?x?xf32> %5 = tensor.empty(%3, %4) : tensor<?x?xf32>
%6 = linalg.generic %6 = linalg.generic
{indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
__internal_linalg_transform__ = "gemm_plus_gemm_fusion"} __internal_linalg_transform__ = "gemm_plus_gemm_fusion"}
ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%5 : tensor<?x?xf32>) { outs(%5 : tensor<?x?xf32>) {
▲ Show 20 Lines • Show All 42 Lines • ▼ Show 20 Linesfunc.func @matmul_plus_transpose_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %arg2, %c0 : tensor<?x?xf32> %0 = tensor.dim %arg2, %c0 : tensor<?x?xf32>
%1 = tensor.dim %arg2, %c1 : tensor<?x?xf32> %1 = tensor.dim %arg2, %c1 : tensor<?x?xf32>
%2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %2 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%3 = tensor.dim %2, %c0 : tensor<?x?xf32> %3 = tensor.dim %2, %c0 : tensor<?x?xf32>
%4 = tensor.dim %2, %c1 : tensor<?x?xf32> %4 = tensor.dim %2, %c1 : tensor<?x?xf32>
%5 = linalg.init_tensor [%3, %4] : tensor<?x?xf32> %5 = tensor.empty(%3, %4) : tensor<?x?xf32>
%6 = linalg.generic %6 = linalg.generic
{indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1, d0)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
__internal_linalg_transform__ = "gemm_plus_gemm_fusion"} __internal_linalg_transform__ = "gemm_plus_gemm_fusion"}
ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%2, %2 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%5 : tensor<?x?xf32>) { outs(%5 : tensor<?x?xf32>) {
▲ Show 20 Lines • Show All 90 Lines • Show Last 20 Lines

mlir/test/Interfaces/TilingInterface/tile-using-interface.mlir

Show First 20 Lines • Show All 81 Lines • ▼ Show 20 Lines
// CHECK-SAME: outs(%[[OUT_TILE]] : // CHECK-SAME: outs(%[[OUT_TILE]] :
// ----- // -----
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d2, d1)> #map1 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map2 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
func.func @multi_result(%arg0 : tensor<128x200x300xf32>) -> (tensor<128x300x200xf32>, tensor<300x128x200xf32>) { func.func @multi_result(%arg0 : tensor<128x200x300xf32>) -> (tensor<128x300x200xf32>, tensor<300x128x200xf32>) {
%init0 = linalg.init_tensor [128, 300, 200] : tensor<128x300x200xf32> %init0 = tensor.empty() : tensor<128x300x200xf32>
%init1 = linalg.init_tensor [300, 128, 200] : tensor<300x128x200xf32> %init1 = tensor.empty() : tensor<300x128x200xf32>
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [#map0, #map1, #map2], indexing_maps = [#map0, #map1, #map2],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
{__internal_linalg_transform__ = "parallel_generic_transpose"} {__internal_linalg_transform__ = "parallel_generic_transpose"}
ins(%arg0 : tensor<128x200x300xf32>) ins(%arg0 : tensor<128x200x300xf32>)
outs(%init0, %init1 : tensor<128x300x200xf32>, tensor<300x128x200xf32>) { outs(%init0, %init1 : tensor<128x300x200xf32>, tensor<300x128x200xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
linalg.yield %b0, %b0 : f32, f32 linalg.yield %b0, %b0 : f32, f32
} -> (tensor<128x300x200xf32>, tensor<300x128x200xf32>) } -> (tensor<128x300x200xf32>, tensor<300x128x200xf32>)
return %0#0, %0#1 : tensor<128x300x200xf32>, tensor<300x128x200xf32> return %0#0, %0#1 : tensor<128x300x200xf32>, tensor<300x128x200xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK: func.func @multi_result( // CHECK: func.func @multi_result(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<128x200x300xf32>) // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<128x200x300xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[C10:.+]] = arith.constant 10 : index // CHECK-DAG: %[[C10:.+]] = arith.constant 10 : index
// CHECK-DAG: %[[C20:.+]] = arith.constant 20 : index // CHECK-DAG: %[[C20:.+]] = arith.constant 20 : index
// CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index // CHECK-DAG: %[[C128:.+]] = arith.constant 128 : index
// CHECK-DAG: %[[C300:.+]] = arith.constant 300 : index // CHECK-DAG: %[[C300:.+]] = arith.constant 300 : index
// CHECK-DAG: %[[INIT0:.+]] = linalg.init_tensor [128, 300, 200] // CHECK-DAG: %[[INIT0:.+]] = tensor.empty()
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [300, 128, 200] // CHECK-DAG: %[[INIT1:.+]] = tensor.empty()
// CHECK: %[[OUTER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C128]] step %[[C10]] // CHECK: %[[OUTER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV0:[a-zA-Z0-9]+]] = %[[C0]] to %[[C128]] step %[[C10]]
// CHECK-SAME: iter_args(%[[ARG1:[a-zA-Z0-9]+]] = %[[INIT0]], %[[ARG2:[a-zA-Z0-9]+]] = %[[INIT1]]) // CHECK-SAME: iter_args(%[[ARG1:[a-zA-Z0-9]+]] = %[[INIT0]], %[[ARG2:[a-zA-Z0-9]+]] = %[[INIT1]])
// CHECK: %[[TS_Y:.+]] = affine.min #[[MAP0]](%[[IV0]])[%[[C10]], %[[C128]]] // CHECK: %[[TS_Y:.+]] = affine.min #[[MAP0]](%[[IV0]])[%[[C10]], %[[C128]]]
// CHECK: %[[INNER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV1:[a-zA-Z0-9]+]] = %[[C0]] to %[[C300]] step %[[C20]] // CHECK: %[[INNER:[a-zA-Z0-9]+]]:2 = scf.for %[[IV1:[a-zA-Z0-9]+]] = %[[C0]] to %[[C300]] step %[[C20]]
// CHECK-SAME: iter_args(%[[ARG3:[a-zA-Z0-9]+]] = %[[ARG1]], %[[ARG4:[a-zA-Z0-9]+]] = %[[ARG2]]) // CHECK-SAME: iter_args(%[[ARG3:[a-zA-Z0-9]+]] = %[[ARG1]], %[[ARG4:[a-zA-Z0-9]+]] = %[[ARG2]])
// CHECK-DAG: %[[ARG_TILE:.+]] = tensor.extract_slice %[[ARG0]] // CHECK-DAG: %[[ARG_TILE:.+]] = tensor.extract_slice %[[ARG0]]
// CHECK-SAME: [%[[IV0]], 0, %[[IV1]]] [%[[TS_Y]], 200, 20] [1, 1, 1] // CHECK-SAME: [%[[IV0]], 0, %[[IV1]]] [%[[TS_Y]], 200, 20] [1, 1, 1]
// CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[ARG3]] // CHECK-DAG: %[[INIT0_TILE:.+]] = tensor.extract_slice %[[ARG3]]
▲ Show 20 Lines • Show All 155 Lines • Show Last 20 Lines

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

Show First 20 Lines • Show All 130 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tensor::ExtractSliceOp op,
// Materialize the output shape values of the slice operation.a // Materialize the output shape values of the slice operation.a
ReifiedRankedShapedTypeDims reifiedShapes; ReifiedRankedShapedTypeDims reifiedShapes;
if (failed(op.reifyResultShapes(rewriter, reifiedShapes))) if (failed(op.reifyResultShapes(rewriter, reifiedShapes)))
return rewriter.notifyMatchFailure(op, "failed to reify result shapes"); return rewriter.notifyMatchFailure(op, "failed to reify result shapes");
// Create the destination tensor using the above values. // Create the destination tensor using the above values.
Type elementType = op.getSourceType().getElementType(); Type elementType = op.getSourceType().getElementType();
SmallVector<OpFoldResult> outputShape = getAsOpFoldResult(reifiedShapes[0]); SmallVector<OpFoldResult> outputShape = getAsOpFoldResult(reifiedShapes[0]);
Value dest = rewriter.create<linalg::InitTensorOp>( Value dest = rewriter.create<tensor::EmptyOp>(op->getLoc(), outputShape,
op->getLoc(), outputShape, elementType); elementType);
// Calculate the parameters for the tile loop nest. // Calculate the parameters for the tile loop nest.
FailureOr<tensor::ExtractSliceFromCollapseHelper> params = FailureOr<tensor::ExtractSliceFromCollapseHelper> params =
tensor::ExtractSliceFromCollapseHelper::create(rewriter, collapseOp, tensor::ExtractSliceFromCollapseHelper::create(rewriter, collapseOp,
op); op);
if (failed(params)) if (failed(params))
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op, "could not calculate tiling parameters"); op, "could not calculate tiling parameters");
▲ Show 20 Lines • Show All 115 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/opdsl/emit_matmul.py

# RUN: %PYTHON %s | FileCheck %s # RUN: %PYTHON %s | FileCheck %s
from mlir.ir import * from mlir.ir import *
from mlir.dialects import builtin from mlir.dialects import builtin
from mlir.dialects import func from mlir.dialects import func
from mlir.dialects import linalg from mlir.dialects import linalg
from mlir.dialects import tensor
from mlir.dialects.linalg.opdsl.lang import * from mlir.dialects.linalg.opdsl.lang import *
T1 = TV.T1 T1 = TV.T1
T2 = TV.T2 T2 = TV.T2
@linalg_structured_op @linalg_structured_op
Show All 30 Lines with InsertionPoint(module.body):
# generation behavior. # generation behavior.
# CHECK: #[[$MUL_MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> # CHECK: #[[$MUL_MAP_A:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
# CHECK: #[[$MUL_MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> # CHECK: #[[$MUL_MAP_B:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
# CHECK: #[[$MUL_MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> # CHECK: #[[$MUL_MAP_C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
# CHECK-LABEL: func @test_matmul_mono # CHECK-LABEL: func @test_matmul_mono
# CHECK-SAME: %[[A:.+]]: tensor<4x16xf32> # CHECK-SAME: %[[A:.+]]: tensor<4x16xf32>
# CHECK-SAME: %[[B:.+]]: tensor<16x8xf32> # CHECK-SAME: %[[B:.+]]: tensor<16x8xf32>
# CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32> # CHECK: %[[INITC:.+]] = tensor.empty() : tensor<4x8xf32>
# CHECK: linalg.generic # CHECK: linalg.generic
# CHECK-SAME: indexing_maps = [#[[$MUL_MAP_A]], #[[$MUL_MAP_B]], #[[$MUL_MAP_C]]] # CHECK-SAME: indexing_maps = [#[[$MUL_MAP_A]], #[[$MUL_MAP_B]], #[[$MUL_MAP_C]]]
# CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] # CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
# CHECK-SAME: ins(%[[A]], %[[B]] # CHECK-SAME: ins(%[[A]], %[[B]]
# CHECK-SAME: outs(%[[INITC]] # CHECK-SAME: outs(%[[INITC]]
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32)) RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32))
def test_matmul_mono(lhs, rhs): def test_matmul_mono(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = tensor.EmptyOp([4, 8], f32)
return matmul_mono(lhs, rhs, outs=[init_result.result]) return matmul_mono(lhs, rhs, outs=[init_result.result])
# CHECK-LABEL: @test_i8i8i32_matmul # CHECK-LABEL: @test_i8i8i32_matmul
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32) # CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)
# CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i32 # CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i8 to i32
# CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i8 to i32 # CHECK-NEXT: %[[B_CAST:.+]] = arith.extsi %[[B_ARG]] : i8 to i32
# CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32 # CHECK-NEXT: %[[MUL:.+]] = arith.muli %[[A_CAST]], %[[B_CAST]] : i32
# CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32 # CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[C_ARG]], %[[MUL]] : i32
▲ Show 20 Lines • Show All 100 Lines • Show Last 20 Lines

mlir/test/python/dialects/linalg/ops.py

# RUN: %PYTHON %s | FileCheck %s # RUN: %PYTHON %s | FileCheck %s
from mlir.dialects import arith, builtin, func, linalg from mlir.dialects import arith, builtin, func, linalg, tensor
from mlir.dialects.linalg.opdsl.lang import * from mlir.dialects.linalg.opdsl.lang import *
from mlir.ir import * from mlir.ir import *
def run(f): def run(f):
print("\nTEST:", f.__name__) print("\nTEST:", f.__name__)
f() f()
return f return f
# CHECK-LABEL: TEST: testInitTensor
@run
def testInitTensor():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
# CHECK-LABEL: func @static_sizes
# CHECK: %0 = linalg.init_tensor [3, 4] : tensor<3x4xf32>
@func.FuncOp.from_py_func()
def static_sizes():
return linalg.InitTensorOp([3, 4], f32)
# CHECK-LABEL: func @dynamic_sizes
# CHECK: %0 = linalg.init_tensor [%arg0, %arg1] : tensor<?x?xf32>
@func.FuncOp.from_py_func(IndexType.get(), IndexType.get())
def dynamic_sizes(d0, d1):
return linalg.InitTensorOp([d0, d1], f32)
# CHECK-LABEL: func @zero_d
# CHECK: %0 = linalg.init_tensor [] : tensor<f32>
@func.FuncOp.from_py_func()
def zero_d():
return linalg.InitTensorOp([], f32)
print(module)
# CHECK-LABEL: TEST: testInitTensorStaticSizesAttribute
@run
def testInitTensorStaticSizesAttribute():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
op = linalg.InitTensorOp([3, 4], f32)
# CHECK: [3, 4]
print(op.attributes["static_sizes"])
# CHECK-LABEL: TEST: testFill # CHECK-LABEL: TEST: testFill
@run @run
def testFill(): def testFill():
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
# CHECK-LABEL: func @fill_tensor # CHECK-LABEL: func @fill_tensor
Show All 25 Linesdef testNamedStructuredOpCustomForm():
with Context() as ctx, Location.unknown(): with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((4, 8), f32), RankedTensorType.get((4, 8), f32)) RankedTensorType.get((4, 8), f32), RankedTensorType.get((4, 8), f32))
def named_form(lhs, rhs): def named_form(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = tensor.EmptyOp([4, 8], f32)
# Check for the named form with custom format # Check for the named form with custom format
# CHECK: linalg.elemwise_unary # CHECK: linalg.elemwise_unary
# CHECK-SAME: cast = #linalg.type_fn<cast_signed> # CHECK-SAME: cast = #linalg.type_fn<cast_signed>
# CHECK-SAME: fun = #linalg.unary_fn<exp> # CHECK-SAME: fun = #linalg.unary_fn<exp>
# CHECK-SAME: ins(%{{.*}} : tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>) # CHECK-SAME: ins(%{{.*}} : tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>)
unary_result = linalg.elemwise_unary(lhs, outs=[init_result.result]) unary_result = linalg.elemwise_unary(lhs, outs=[init_result.result])
# CHECK: linalg.elemwise_binary # CHECK: linalg.elemwise_binary
# CHECK-SAME: cast = #linalg.type_fn<cast_unsigned> # CHECK-SAME: cast = #linalg.type_fn<cast_unsigned>
Show All 18 Lines with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
f32)) f32))
def named_form(lhs, rhs): def named_form(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = tensor.EmptyOp([4, 8], f32)
# CHECK: "linalg.matmul"(%{{.*}}) # CHECK: "linalg.matmul"(%{{.*}})
# CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): # CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
# CHECK-NEXT: arith.mulf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: arith.mulf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: arith.addf{{.*}} (f32, f32) -> f32 # CHECK-NEXT: arith.addf{{.*}} (f32, f32) -> f32
# CHECK-NEXT: linalg.yield{{.*}} (f32) -> () # CHECK-NEXT: linalg.yield{{.*}} (f32) -> ()
# CHECK-NEXT: cast = #linalg.type_fn<cast_signed> # CHECK-NEXT: cast = #linalg.type_fn<cast_signed>
# CHECK-SAME: operand_segment_sizes = array<i32: 2, 1> # CHECK-SAME: operand_segment_sizes = array<i32: 2, 1>
# CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32> # CHECK-SAME: (tensor<4x16xf32>, tensor<16x8xf32>, tensor<4x8xf32>) -> tensor<4x8xf32>
Show All 9 Lines with Context() as ctx, Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
f32)) f32))
def generic_form(lhs, rhs): def generic_form(lhs, rhs):
init_result = linalg.InitTensorOp([4, 8], f32) init_result = tensor.EmptyOp([4, 8], f32)
# CHECK: linalg.generic # CHECK: linalg.generic
return linalg.matmul( return linalg.matmul(
lhs, rhs, outs=[init_result.result], emit_generic=True) lhs, rhs, outs=[init_result.result], emit_generic=True)
print(module) print(module)
# CHECK-LABEL: TEST: testOpResultFromOtherOp # CHECK-LABEL: TEST: testOpResultFromOtherOp
@run @run
def testOpResultFromOtherOp(): def testOpResultFromOtherOp():
with Context(), Location.unknown(): with Context(), Location.unknown():
module = Module.create() module = Module.create()
f32 = F32Type.get() f32 = F32Type.get()
with InsertionPoint(module.body): with InsertionPoint(module.body):
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8),
f32)) f32))
def pass_an_op_directly(arg0, arg1): def pass_an_op_directly(arg0, arg1):
one = arith.ConstantOp(F32Type.get(), 1.0) one = arith.ConstantOp(F32Type.get(), 1.0)
# CHECK: %[[LHS:.*]] = linalg.fill # CHECK: %[[LHS:.*]] = linalg.fill
lhs = linalg.fill(one, outs=[arg0]) lhs = linalg.fill(one, outs=[arg0])
# CHECK: %[[RHS:.*]] = linalg.fill # CHECK: %[[RHS:.*]] = linalg.fill
rhs = linalg.fill(one, outs=[arg1]) rhs = linalg.fill(one, outs=[arg1])
# CHECK: %[[INIT:.*]] = linalg.init_tensor # CHECK: %[[INIT:.*]] = tensor.empty
init = linalg.InitTensorOp([4, 8], f32) init = tensor.EmptyOp([4, 8], f32)
# CHECK: linalg.matmul # CHECK: linalg.matmul
# CHECK: ins(%[[LHS]], %[[RHS]] # CHECK: ins(%[[LHS]], %[[RHS]]
# CHECK: outs(%[[INIT]] # CHECK: outs(%[[INIT]]
return linalg.matmul(lhs, rhs, outs=init) return linalg.matmul(lhs, rhs, outs=init)
print(module) print(module)

mlir/test/python/dialects/tensor.py

Show All 31 Lines with InsertionPoint(module.body):
def tensor_static_dim(t): def tensor_static_dim(t):
c0 = arith.ConstantOp(indexType, 0) c0 = arith.ConstantOp(indexType, 0)
c1 = arith.ConstantOp(indexType, 1) c1 = arith.ConstantOp(indexType, 1)
d0 = tensor.DimOp(t, c0) d0 = tensor.DimOp(t, c0)
d1 = tensor.DimOp(t, c1) d1 = tensor.DimOp(t, c1)
return [d0.result, d1.result] return [d0.result, d1.result]
print(module) print(module)
# CHECK-LABEL: TEST: testEmptyOp
@run
def testEmptyOp():
with Context() as ctx, Location.unknown():
module = Module.create()
f32 = F32Type.get()
with InsertionPoint(module.body):
# CHECK-LABEL: func @static_sizes
# CHECK: %0 = tensor.empty() : tensor<3x4xf32>
@func.FuncOp.from_py_func()
def static_sizes():
return tensor.EmptyOp([3, 4], f32)
# CHECK-LABEL: func @dynamic_sizes
# CHECK: %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>
@func.FuncOp.from_py_func(IndexType.get(), IndexType.get())
def dynamic_sizes(d0, d1):
return tensor.EmptyOp([d0, d1], f32)
# CHECK-LABEL: func @mixed_static_dynamic_sizes
# CHECK: %0 = tensor.empty(%arg0) : tensor<?x4xf32>
@func.FuncOp.from_py_func(IndexType.get())
def mixed_static_dynamic_sizes(d0):
return tensor.EmptyOp([d0, 4], f32)
# CHECK-LABEL: func @zero_d
# CHECK: %0 = tensor.empty() : tensor<f32>
@func.FuncOp.from_py_func()
def zero_d():
return tensor.EmptyOp([], f32)
print(module)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

Show First 20 Lines • Show All 5,026 Lines • ▼ Show 20 Linescc_library(
srcs = [ srcs = [
"include/mlir/Transforms/InliningUtils.h", "include/mlir/Transforms/InliningUtils.h",
"lib/Dialect/Tensor/IR/TensorDialect.cpp", "lib/Dialect/Tensor/IR/TensorDialect.cpp",
"lib/Dialect/Tensor/IR/TensorOps.cpp", "lib/Dialect/Tensor/IR/TensorOps.cpp",
], ],
hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"], hdrs = ["include/mlir/Dialect/Tensor/IR/Tensor.h"],
includes = ["include"], includes = ["include"],
deps = [ deps = [
":AffineDialect",
":ArithDialect", ":ArithDialect",
":ArithUtils", ":ArithUtils",
":CastOpInterfaces", ":CastOpInterfaces",
":ComplexDialect", ":ComplexDialect",
":ControlFlowInterfaces", ":ControlFlowInterfaces",
":DialectUtils", ":DialectUtils",
":IR", ":IR",
":InferTypeOpInterface", ":InferTypeOpInterface",
▲ Show 20 Lines • Show All 4,625 Lines • Show Last 20 Lines