This is an archive of the discontinued LLVM Phabricator instance.

Differential D136875

[mlir] Rename `outs` to `inits` for linalg ops.
AbandonedPublic

Authored by pifon2a on Oct 27 2022, 12:06 PM.

Details

Reviewers
nicolasvasilache
springerm
antiagainst
aartbik
ThomasRaoux
herhut
dcaballe
ftynse
Summary

https://discourse.llvm.org/t/rfc-interface-for-destination-style-ops/64056/15?u=pifon2a

Diff Detail

Event Timeline

pifon2a created this revision.Oct 27 2022, 12:06 PM
Herald added a reviewer: antiagainst. · View Herald TranscriptOct 27 2022, 12:06 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: ThomasRaoux. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: zero9178, anlunx, bzcheeseman and 29 others. · View Herald Transcript
pifon2a requested review of this revision.Oct 27 2022, 12:06 PM
Herald added a reviewer: herhut. · View Herald TranscriptOct 27 2022, 12:06 PM
Herald added a reviewer: dcaballe. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: • pcwang-thead, limo1996, stephenneuendorffer. · View Herald Transcript
aartbik added a comment.Oct 27 2022, 1:12 PM

LGTM for the sparse part, since it is a mechanical change. But it will have impact on many out-of-tree projects using the linalg syntax.
Was there a discourse posting with a heads up on this?

cota added a comment.Oct 27 2022, 1:14 PM

Can you please provide a rationale for the rename in the commit message? Or a link to a previous discussion on this topic. Thanks!

Harbormaster completed remote builds in B194728: Diff 471254.Oct 27 2022, 1:46 PM
nicolasvasilache added a comment.Oct 27 2022, 2:39 PM

This may be contentious and not worth the churn.. not all outs are technically inits and vice versa.

readonly, writeonly, readwrite is the verbose spelling for ins, outs, inits (we even used to have getter/setter for the 3 categories at some point).

Not opposed to the change but if you have a better way to look at this that made you conclude that "init" is unambiguously better thatn "out", I'd love to understand it.

mravishankar added a comment.Oct 27 2022, 2:44 PM
In D136875#3889917, @nicolasvasilache wrote:

This may be contentious and not worth the churn.. not all outs are technically inits and vice versa.

readonly, writeonly, readwrite is the verbose spelling for ins, outs, inits (we even used to have getter/setter for the 3 categories at some point).

Not opposed to the change but if you have a better way to look at this that made you conclude that "init" is unambiguously better thatn "out", I'd love to understand it.

Second this. I have seen lots of people confused about outs and somehow saying init clarifies things when looking at Linalg on tensors.... but this might be fixable with documentation. The churn is not desirable though. It will just be a nuisance for downstream users.

nicolasvasilache mentioned this in D136838: [mlir][linalg] Add nicer builders for `map` and `reduce`..Oct 28 2022, 1:37 AM
pifon2a added a comment.EditedOct 28 2022, 2:03 AM
In D136875#3889917, @nicolasvasilache wrote:

This may be contentious and not worth the churn.. not all outs are technically inits and vice versa.

readonly, writeonly, readwrite is the verbose spelling for ins, outs, inits (we even used to have getter/setter for the 3 categories at some point).

Not opposed to the change but if you have a better way to look at this that made you conclude that "init" is unambiguously better thatn "out", I'd love to understand it.

When we are in tensor world, isn't it always an "init" technically? even if it is a "tensor.empty".

@cota, @aartbik It was discussed in this RFC https://discourse.llvm.org/t/rfc-interface-for-destination-style-ops/64056/15?u=pifon2a

Actually it is a very simple sed command that replaces outs to inits in .mlir files. A much harder change is to rename the methods in DPS interface, which I have to do anyway because of name collisions: after mlir migrated to camel-case getters, there are conflicts between getInputs generated by tablegen and getInputs in the DPS interface. Since I am currently at it, I suggest we rename it to inits.

nicolasvasilache added a comment.Oct 28 2022, 3:07 AM

@cota, @aartbik It was discussed in this RFC https://discourse.llvm.org/t/rfc-interface-for-destination-style-ops/64056/15?u=pifon2a

Oh cool, I lost track of the fact consensus had been reached, fine with me then!

aartbik added a comment.Oct 28 2022, 11:35 AM
In D136875#3891174, @pifon2a wrote:

@cota, @aartbik It was discussed in this RFC https://discourse.llvm.org/t/rfc-interface-for-destination-style-ops/64056/15?u=pifon2a

Thanks. Please add this link in the description on this revision. Also probably a good idea to "reply to" the discourse discussion with a note that this is happening
(last post was a while back ;-). But other than that LGTM as before.

Herald added a subscriber: Moerafaat. · View Herald TranscriptOct 28 2022, 11:35 AM
pifon2a edited the summary of this revision. (Show Details)Oct 28 2022, 12:10 PM
pifon2a updated this revision to Diff 474238.Nov 9 2022, 5:35 AM

Rebase

Harbormaster completed remote builds in B196884: Diff 474238.Nov 9 2022, 6:16 AM
dcaballe added a comment.Nov 9 2022, 10:01 PM

Thanks for proposing this change, Alex! I agree that we should pay attention to using descriptive and precise names. Unfortunately, I've been digesting the getDPSInits() changes and I'm not sure this is going in the right direction. We now have code snippets that are somewhat confusing. For instance:

Process operation *outputs* doing X and Y.
... = getDPSInits();

I have spent some time scratching my head thinking... "hmm, where are the outputs in this code?". With this I'm not suggesting that we go and change the comments but that the name seems to be a bit counterintuitive. I haven't been following the discussion about why this change is needed but hopefully we can reconsider a name that is more intuitive like outputs, inouts, etc. Same for getDPSInits. It's true that in DPS outputs have an initial value but the output aspect of it has much more weight in what they represent, IMO.

Hopefully that helps!

pifon2a added a comment.Nov 10 2022, 2:10 AM

@dcaballe

Hi Diego. The whole reason for this change is the discussion in https://discourse.llvm.org/t/rfc-interface-for-destination-style-ops/64056/15?u=pifon2a where people found "outputs" confusing :). I don't care about naming at all. I just want the getters to be consistent with the IR that we print.

pifon2a abandoned this revision.Jan 13 2023, 5:34 AM
Herald added subscribers: hanchung, jsetoain. · View Herald TranscriptJan 13 2023, 5:34 AM
mehdi_amini added a comment.Mar 14 2023, 5:13 AM
In D136875#3918674, @dcaballe wrote:

we can reconsider a name that is more intuitive like outputs

For me it is highly confusing to read "outputs" and associate this to immutable objects (tensors). Seeing how people try to form a mental model around this for the last 2 years, I really want an alternative spelling that won't bring this confusing and respect the immutability of tensors.
Seems that inits is where we converged on Discourse, but if you want to propose other spelling (that don't imply "writing to a tensor") please do!

Herald added a reviewer: ftynse. · View Herald TranscriptMar 14 2023, 5:13 AM
Herald added a subscriber: thopre. · View Herald Transcript

Revision Contents

PathSize
mlir/
include/
mlir/
Dialect/
Linalg/
IR/
12 lines
lib/
Dialect/
Linalg/
IR/
8 lines
test/
Analysis/
6 lines
Conversion/
LinalgToSPIRV/
8 lines
TensorToLinalg/
4 lines
TosaToLinalg/
102 lines
10 lines
130 lines
Dialect/
Bufferization/
Transforms/
18 lines
4 lines
4 lines
62 lines
2 lines
22 lines
8 lines
GPU/
2 lines
Linalg/
4 lines
30 lines
32 lines
50 lines
118 lines
8 lines
20 lines
12 lines
44 lines
12 lines
4 lines
12 lines
2 lines
4 lines
14 lines
4 lines
82 lines
4 lines
8 lines
6 lines
40 lines
4 lines
122 lines
8 lines
12 lines
14 lines
74 lines
44 lines
60 lines
4 lines
4 lines
2 lines
86 lines
4 lines
48 lines
2 lines
4 lines
224 lines
8 lines
144 lines
4 lines
50 lines
12 lines
6 lines
12 lines
4 lines
10 lines
42 lines
12 lines
74 lines
4 lines
6 lines
36 lines
4 lines
4 lines
10 lines
18 lines
8 lines
20 lines
16 lines
2 lines
4 lines
2 lines
4 lines
10 lines
2 lines
6 lines
48 lines
30 lines
8 lines
8 lines
20 lines
20 lines
12 lines
92 lines
16 lines
SCF/
8 lines
6 lines
6 lines
20 lines
4 lines
26 lines
2 lines
SparseTensor/
14 lines
12 lines
14 lines
6 lines
6 lines
2 lines
8 lines
76 lines
60 lines
74 lines
8 lines
6 lines
18 lines
8 lines
20 lines
4 lines
22 lines
10 lines
2 lines
2 lines
2 lines
4 lines
10 lines
10 lines
6 lines
8 lines
2 lines
2 lines
12 lines
4 lines
2 lines
10 lines
Tensor/
8 lines
4 lines
10 lines
Transform/
14 lines
Vector/
4 lines
IR/
4 lines
Integration/
Dialect/
Async/
CPU/
6 lines
4 lines
Linalg/
CPU/
12 lines
4 lines
4 lines
4 lines
4 lines
4 lines
4 lines
10 lines
2 lines
SparseTensor/
CPU/
2 lines
2 lines
2 lines
4 lines
22 lines
20 lines
4 lines
4 lines
14 lines
2 lines
2 lines
4 lines
2 lines
10 lines
8 lines
6 lines
8 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
4 lines
8 lines
14 lines
2 lines
12 lines
2 lines
4 lines
4 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
2 lines
4 lines
10 lines
12 lines
python/
2 lines
Interfaces/
TilingInterface/
14 lines
92 lines
22 lines
lib/
Dialect/
Test/
4 lines
mlir-cpu-runner/
2 lines
6 lines
10 lines
4 lines
mlir-opt/
2 lines
python/
dialects/
linalg/
opdsl/
2 lines
10 lines
integration/
dialects/
linalg/
28 lines

Diff 471254

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

Show First 20 Lines • Show All 65 Lines • ▼ Show 20 Lines let description = [{
Generic Linalg op form where the key properties of the computation are Generic Linalg op form where the key properties of the computation are
specified as attributes. In pretty form, a `linalg.generic` op is written specified as attributes. In pretty form, a `linalg.generic` op is written
as: as:
```mlir ```mlir
linalg.generic #trait_attribute linalg.generic #trait_attribute
ins(%A, %B : memref<?x?xf32, stride_specification>, ins(%A, %B : memref<?x?xf32, stride_specification>,
memref<?x?xf32, stride_specification>) memref<?x?xf32, stride_specification>)
outs(%C : memref<?x?xf32, stride_specification>) inits(%C : memref<?x?xf32, stride_specification>)
attrs = {other-optional-attributes} attrs = {other-optional-attributes}
{region} {region}
``` ```
Where #trait_attributes is an alias of a dictionary attribute containing: Where #trait_attributes is an alias of a dictionary attribute containing:
- doc [optional]: a documentation string - doc [optional]: a documentation string
- indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input - indexing_maps: a list of AffineMapAttr, one AffineMapAttr per each input
and output view. Such AffineMapAttr specifies the mapping between the and output view. Such AffineMapAttr specifies the mapping between the
Show All 24 Lines Defining a #matmul_trait attribute in MLIR can be done as follows:
} }
``` ```
And can be reused in multiple places as: And can be reused in multiple places as:
```mlir ```mlir
linalg.generic #matmul_trait linalg.generic #matmul_trait
ins(%A, %B : memref<?x?xf32, stride_specification>, ins(%A, %B : memref<?x?xf32, stride_specification>,
memref<?x?xf32, stride_specification>) memref<?x?xf32, stride_specification>)
outs(%C : memref<?x?xf32, stride_specification>) inits(%C : memref<?x?xf32, stride_specification>)
{other-optional-attributes} { {other-optional-attributes} {
^bb0(%a: f32, %b: f32, %c: f32) : ^bb0(%a: f32, %b: f32, %c: f32) :
%d = arith.mulf %a, %b: f32 %d = arith.mulf %a, %b: f32
%e = arith.addf %c, %d: f32 %e = arith.addf %c, %d: f32
linalg.yield %e : f32 linalg.yield %e : f32
} }
``` ```
Show All 24 Lines let description = [{
To allow progressive lowering from the value world (a.k.a tensor values) to To allow progressive lowering from the value world (a.k.a tensor values) to
the buffer world (a.k.a memref values), a `linalg.generic` op allows mixing the buffer world (a.k.a memref values), a `linalg.generic` op allows mixing
tensors and buffers operands and tensor results. tensors and buffers operands and tensor results.
```mlir ```mlir
%C = linalg.generic #trait_attribute %C = linalg.generic #trait_attribute
ins(%A, %B : tensor<?x?xf32>, memref<?x?xf32, stride_specification>) ins(%A, %B : tensor<?x?xf32>, memref<?x?xf32, stride_specification>)
outs(%C : tensor<?x?xf32>) inits(%C : tensor<?x?xf32>)
{other-optional-attributes} {other-optional-attributes}
{region} {region}
-> (tensor<?x?xf32>) -> (tensor<?x?xf32>)
``` ```
}]; }];
let arguments = (ins Variadic<AnyType>:$inputs, let arguments = (ins Variadic<AnyType>:$inputs,
Variadic<AnyShaped>:$outputs, Variadic<AnyShaped>:$outputs,
▲ Show 20 Lines • Show All 79 Lines • ▼ Show 20 Linesdef MapOp : LinalgStructuredBase_Op<"map", [
let description = [{ let description = [{
Models elementwise operations on tensors in terms of arithmetic operations Models elementwise operations on tensors in terms of arithmetic operations
on the corresponding elements. on the corresponding elements.
Example: Example:
``` ```
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
outs(%init: tensor<64xf32>) inits(%init: tensor<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
``` ```
}]; }];
let arguments = (ins let arguments = (ins
▲ Show 20 Lines • Show All 54 Lines • ▼ Show 20 Lines let description = [{
Executes `combiner` on the `dimensions` of `inputs` and returns the Executes `combiner` on the `dimensions` of `inputs` and returns the
reduced result. The `dimensions` attribute needs to list the reduction reduced result. The `dimensions` attribute needs to list the reduction
dimensions in increasing order. dimensions in increasing order.
Example: Example:
``` ```
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
``` ```
}]; }];
▲ Show 20 Lines • Show All 48 Lines • ▼ Show 20 Lines let description = [{
This op actually moves data, unlike `memref.transpose` which is a metadata This op actually moves data, unlike `memref.transpose` which is a metadata
operation only that produces a transposed "view". operation only that produces a transposed "view".
Example: Example:
``` ```
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x64xf32>) ins(%input:tensor<16x64xf32>)
outs(%init:tensor<64x16xf32>) inits(%init:tensor<64x16xf32>)
permutation = [1, 0] permutation = [1, 0]
``` ```
}]; }];
let arguments = (ins let arguments = (ins
// Input arg // Input arg
TensorOrMemref:$input, TensorOrMemref:$input,
// Output arg // Output arg
▲ Show 20 Lines • Show All 50 Lines • Show Last 20 Lines

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

Show First 20 Lines • Show All 137 Lines • ▼ Show 20 Lines if (parser.parseLParen())
return failure(); return failure();
inputsOperandsLoc = parser.getCurrentLocation(); inputsOperandsLoc = parser.getCurrentLocation();
if (parser.parseOperandList(inputsOperands) || if (parser.parseOperandList(inputsOperands) ||
parser.parseColonTypeList(inputTypes) || parser.parseRParen()) parser.parseColonTypeList(inputTypes) || parser.parseRParen())
return failure(); return failure();
} }
if (succeeded(parser.parseOptionalKeyword("outs"))) { if (succeeded(parser.parseOptionalKeyword("inits"))) {
outputsOperandsLoc = parser.getCurrentLocation(); outputsOperandsLoc = parser.getCurrentLocation();
if (parser.parseLParen() || parser.parseOperandList(outputsOperands) || if (parser.parseLParen() || parser.parseOperandList(outputsOperands) ||
parser.parseColonTypeList(outputTypes) || parser.parseRParen()) parser.parseColonTypeList(outputTypes) || parser.parseRParen())
return failure(); return failure();
} }
if (parser.resolveOperands(inputsOperands, inputTypes, inputsOperandsLoc, if (parser.resolveOperands(inputsOperands, inputTypes, inputsOperandsLoc,
result.operands) || result.operands) ||
Show All 10 LinesparseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
return success(); return success();
} }
static void printCommonStructuredOpParts(OpAsmPrinter &p, ValueRange inputs, static void printCommonStructuredOpParts(OpAsmPrinter &p, ValueRange inputs,
ValueRange outputs) { ValueRange outputs) {
if (!inputs.empty()) if (!inputs.empty())
p << " ins(" << inputs << " : " << inputs.getTypes() << ")"; p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
if (!outputs.empty()) if (!outputs.empty())
p << " outs(" << outputs << " : " << outputs.getTypes() << ")"; p << " inits(" << outputs << " : " << outputs.getTypes() << ")";
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Specific parsing and printing for named structured ops created by ods-gen. // Specific parsing and printing for named structured ops created by ods-gen.
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
static ParseResult parseNamedStructuredOpRegion( static ParseResult parseNamedStructuredOpRegion(
OpAsmParser &parser, Region &region, unsigned numRegionArgs, OpAsmParser &parser, Region &region, unsigned numRegionArgs,
▲ Show 20 Lines • Show All 467 Lines • ▼ Show 20 Lines
// GenericOp // GenericOp
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
void GenericOp::getAsmBlockArgumentNames(Region &region, void GenericOp::getAsmBlockArgumentNames(Region &region,
OpAsmSetValueNameFn setNameFn) { OpAsmSetValueNameFn setNameFn) {
for (Value v : getRegionInputArgs()) for (Value v : getRegionInputArgs())
setNameFn(v, "in"); setNameFn(v, "in");
for (Value v : getRegionOutputArgs()) for (Value v : getRegionOutputArgs())
setNameFn(v, "out"); setNameFn(v, "init");
} }
void GenericOp::build( void GenericOp::build(
OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes, OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps, ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps,
ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall, ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall,
function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild, function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
ArrayRef<NamedAttribute> attributes) { ArrayRef<NamedAttribute> attributes) {
▲ Show 20 Lines • Show All 1,314 Lines • ▼ Show 20 Lines LogicalResult matchAndRewrite(tensor::CastOp castOp,
OpBuilder::InsertionGuard guard(rewriter); OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPoint(linalgOp); rewriter.setInsertionPoint(linalgOp);
Location loc = linalgOp.getLoc(); Location loc = linalgOp.getLoc();
OpResult resultValue = castOp.getSource().cast<OpResult>(); OpResult resultValue = castOp.getSource().cast<OpResult>();
unsigned resultNumber = resultValue.getResultNumber(); unsigned resultNumber = resultValue.getResultNumber();
auto resultType = castOp->getResult(0).getType().cast<RankedTensorType>(); auto resultType = castOp->getResult(0).getType().cast<RankedTensorType>();
// Replace the `outs` for the result with a `tensor.cast`. This cast is now // Replace the `inits` for the result with a `tensor.cast`. This cast is now
// going from a more dynamic shape to a less dynamic shape. If the producer // going from a more dynamic shape to a less dynamic shape. If the producer
// for this cast, i.e. producer of the out operand, is also an operation // for this cast, i.e. producer of the out operand, is also an operation
// that folds with tensor.cast consumer (like this pattern), the cast will // that folds with tensor.cast consumer (like this pattern), the cast will
// continue to propagate as far up the stack as it can go. // continue to propagate as far up the stack as it can go.
OpOperand *outOperand = linalgOp.getOutputOperand(resultNumber); OpOperand *outOperand = linalgOp.getOutputOperand(resultNumber);
Value newOperand = Value newOperand =
rewriter.create<tensor::CastOp>(loc, resultType, outOperand->get()); rewriter.create<tensor::CastOp>(loc, resultType, outOperand->get());
SmallVector<Value> newOperands{linalgOp.getInputOperands()}; SmallVector<Value> newOperands{linalgOp.getInputOperands()};
▲ Show 20 Lines • Show All 197 Lines • Show Last 20 Lines

mlir/test/Analysis/test-match-reduction.mlir

// RUN: mlir-opt %s -pass-pipeline="func.func(test-match-reduction)" -verify-diagnostics -split-input-file // RUN: mlir-opt %s -pass-pipeline="func.func(test-match-reduction)" -verify-diagnostics -split-input-file
// Verify that the generic reduction detection utility works on different // Verify that the generic reduction detection utility works on different
// dialects. // dialects.
// expected-remark@below {{Testing function}} // expected-remark@below {{Testing function}}
func.func @linalg_red_add(%in0t : tensor<?xf32>, %out0t : tensor<1xf32>) { func.func @linalg_red_add(%in0t : tensor<?xf32>, %out0t : tensor<1xf32>) {
// expected-remark@below {{Reduction found in output #0!}} // expected-remark@below {{Reduction found in output #0!}}
// expected-remark@below {{Reduced Value: <block argument> of type 'f32' at index: 0}} // expected-remark@below {{Reduced Value: <block argument> of type 'f32' at index: 0}}
// expected-remark@below {{Combiner Op: %1 = arith.addf }} // expected-remark@below {{Combiner Op: %1 = arith.addf }}
%red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> (0)>], affine_map<(d0) -> (0)>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%in0t : tensor<?xf32>) ins(%in0t : tensor<?xf32>)
outs(%out0t : tensor<1xf32>) { inits(%out0t : tensor<1xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%add = arith.addf %in0, %out0 : f32 %add = arith.addf %in0, %out0 : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<1xf32> } -> tensor<1xf32>
return return
} }
// ----- // -----
Show All 20 Lines
// TODO: Iteration-carried values with multiple uses are not supported yet. // TODO: Iteration-carried values with multiple uses are not supported yet.
// expected-remark@below {{Testing function}} // expected-remark@below {{Testing function}}
func.func @linalg_red_max(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) { func.func @linalg_red_max(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
// expected-remark@below {{Reduction NOT found in output #0!}} // expected-remark@below {{Reduction NOT found in output #0!}}
%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(%in0t : tensor<4x4xf32>) ins(%in0t : tensor<4x4xf32>)
outs(%out0t : tensor<4xf32>) { inits(%out0t : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%cmp = arith.cmpf ogt, %in0, %out0 : f32 %cmp = arith.cmpf ogt, %in0, %out0 : f32
%sel = arith.select %cmp, %in0, %out0 : f32 %sel = arith.select %cmp, %in0, %out0 : f32
linalg.yield %sel : f32 linalg.yield %sel : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return return
} }
// ----- // -----
// expected-remark@below {{Testing function}} // expected-remark@below {{Testing function}}
func.func @linalg_fused_red_add(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) { func.func @linalg_fused_red_add(%in0t: tensor<4x4xf32>, %out0t: tensor<4xf32>) {
// expected-remark@below {{Reduction found in output #0!}} // expected-remark@below {{Reduction found in output #0!}}
// expected-remark@below {{Reduced Value: %2 = arith.subf}} // expected-remark@below {{Reduced Value: %2 = arith.subf}}
// expected-remark@below {{Combiner Op: %3 = arith.addf}} // expected-remark@below {{Combiner Op: %3 = arith.addf}}
%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(%in0t : tensor<4x4xf32>) ins(%in0t : tensor<4x4xf32>)
outs(%out0t : tensor<4xf32>) { inits(%out0t : tensor<4xf32>) {
^bb0(%in0: f32, %out0: f32): ^bb0(%in0: f32, %out0: f32):
%mul = arith.mulf %in0, %in0 : f32 %mul = arith.mulf %in0, %in0 : f32
%sub = arith.subf %mul, %in0 : f32 %sub = arith.subf %mul, %in0 : f32
%add = arith.addf %sub, %out0 : f32 %add = arith.addf %sub, %out0 : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return return
} }
Show All 34 Lines

mlir/test/Conversion/LinalgToSPIRV/linalg-to-spirv.mlir

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: spirv.Return // CHECK: spirv.Return
func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes { func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes {
spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>> spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 1, 1]>: vector<3xi32>>
} { } {
linalg.generic #single_workgroup_reduction_trait linalg.generic #single_workgroup_reduction_trait
ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>) ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>)
outs(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) { inits(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) {
^bb(%in: i32, %out: i32): ^bb(%in: i32, %out: i32):
%sum = arith.addi %in, %out : i32 %sum = arith.addi %in, %out : i32
linalg.yield %sum : i32 linalg.yield %sum : i32
} }
spirv.Return spirv.Return
} }
} }
Show All 12 Lines
module attributes { module attributes {
spirv.target_env = #spirv.target_env< spirv.target_env = #spirv.target_env<
#spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>> #spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>>
} { } {
func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) { func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) {
// expected-error @+1 {{failed to legalize operation 'linalg.generic'}} // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
linalg.generic #single_workgroup_reduction_trait linalg.generic #single_workgroup_reduction_trait
ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>) ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>)
outs(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) { inits(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) {
^bb(%in: i32, %out: i32): ^bb(%in: i32, %out: i32):
%sum = arith.addi %in, %out : i32 %sum = arith.addi %in, %out : i32
linalg.yield %sum : i32 linalg.yield %sum : i32
} }
return return
} }
} }
Show All 14 Lines spirv.target_env = #spirv.target_env<
#spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>> #spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>>
} { } {
func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes { func.func @single_workgroup_reduction(%input: memref<16xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<1xi32, #spirv.storage_class<StorageBuffer>>) attributes {
spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 1, 1]>: vector<3xi32>> spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[32, 1, 1]>: vector<3xi32>>
} { } {
// expected-error @+1 {{failed to legalize operation 'linalg.generic'}} // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
linalg.generic #single_workgroup_reduction_trait linalg.generic #single_workgroup_reduction_trait
ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>) ins(%input : memref<16xi32, #spirv.storage_class<StorageBuffer>>)
outs(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) { inits(%output : memref<1xi32, #spirv.storage_class<StorageBuffer>>) {
^bb(%in: i32, %out: i32): ^bb(%in: i32, %out: i32):
%sum = arith.addi %in, %out : i32 %sum = arith.addi %in, %out : i32
linalg.yield %sum : i32 linalg.yield %sum : i32
} }
spirv.Return spirv.Return
} }
} }
Show All 14 Lines spirv.target_env = #spirv.target_env<
#spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>> #spirv.vce<v1.3, [Shader, GroupNonUniformArithmetic], []>, #spirv.resource_limits<>>
} { } {
func.func @single_workgroup_reduction(%input: memref<16x8xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<16xi32, #spirv.storage_class<StorageBuffer>>) attributes { func.func @single_workgroup_reduction(%input: memref<16x8xi32, #spirv.storage_class<StorageBuffer>>, %output: memref<16xi32, #spirv.storage_class<StorageBuffer>>) attributes {
spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 8, 1]>: vector<3xi32>> spirv.entry_point_abi = #spirv.entry_point_abi<local_size = dense<[16, 8, 1]>: vector<3xi32>>
} { } {
// expected-error @+1 {{failed to legalize operation 'linalg.generic'}} // expected-error @+1 {{failed to legalize operation 'linalg.generic'}}
linalg.generic #single_workgroup_reduction_trait linalg.generic #single_workgroup_reduction_trait
ins(%input : memref<16x8xi32, #spirv.storage_class<StorageBuffer>>) ins(%input : memref<16x8xi32, #spirv.storage_class<StorageBuffer>>)
outs(%output : memref<16xi32, #spirv.storage_class<StorageBuffer>>) { inits(%output : memref<16xi32, #spirv.storage_class<StorageBuffer>>) {
^bb(%in: i32, %out: i32): ^bb(%in: i32, %out: i32):
%sum = arith.addi %in, %out : i32 %sum = arith.addi %in, %out : i32
linalg.yield %sum : i32 linalg.yield %sum : i32
} }
spirv.Return spirv.Return
} }
} }

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:.*]] = tensor.empty() : 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) inits(%[[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:.*]] = tensor.empty(%[[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) inits(%[[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:%.+]] = tensor.empty() // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) outs([[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : f32) inits([[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>) inits([[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:%.+]] = tensor.empty() // CHECK: [[INIT:%.+]] = tensor.empty()
// CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : i32) outs([[INIT]] : tensor<1x5x6xi32>) -> tensor<1x5x6xi32> // CHECK: [[FILLED:%.+]] = linalg.fill ins([[C0]] : i32) inits([[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) inits([[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:.+]] = tensor.empty(%[[DIM]]) // 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) inits(%[[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>) inits(%[[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:.+]] = tensor.empty(%[[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) inits(%[[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>) inits(%[[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:.+]] = tensor.empty() // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x5x6xf32>) -> tensor<1x5x6xf32> // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) inits(%[[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>) inits(%[[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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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:%.+]] = tensor.empty() // 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>) inits([[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>) inits([[INITB]] : tensor<5x6xf32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: f32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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:%.+]] = tensor.empty() // 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) inits([[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>) inits([[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:.+]] = tensor.empty(%[[DIM]]) // 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]]{{.*}}inits(%[[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:.+]] = tensor.empty(%[[DIM]]) // 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>) inits(%[[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>) inits(%[[INITB]] : tensor<?x6xf32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: f32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:%.+]] = tensor.empty() // CHECK-DAG: [[INIT:%.+]] = tensor.empty()
// CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[INIT]] // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}inits([[INIT]]
// CHECK-DAG: [[KERNEL:%.+]] = tensor.empty() // 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>) inits([[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:%.+]] = tensor.empty() // CHECK-DAG: [[INIT:%.+]] = tensor.empty()
// CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INITVAL]]{{.*}}outs([[INIT]] // CHECK-DAG: [[FILL:%.+]] = linalg.fill ins([[INITVAL]]{{.*}}inits([[INIT]]
// CHECK-DAG: [[KERNEL:%.+]] = tensor.empty() // 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>) inits([[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:.+]] = tensor.empty(%[[BATCH]]) // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CONST]]{{.*}}inits(%[[INIT]]
// CHECK: %[[KERNEL:.+]] = tensor.empty() // 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>) inits(%[[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
// CHECK: linalg.pooling_nhwc_max // CHECK: linalg.pooling_nhwc_max
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:%.+]] = tensor.empty() // CHECK: [[POOLINIT:%.+]] = tensor.empty()
// CHECK: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}outs([[POOLINIT]] // CHECK: [[FILL:%.+]] = linalg.fill ins([[CONST]]{{.*}}inits([[POOLINIT]]
// CHECK: [[KERNEL:%.+]] = tensor.empty() // 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>) inits([[FILL]] : tensor<1x5x33x62xf32>)
// CHECK: [[INIT:%.+]] = tensor.empty() // 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>) inits([[INIT]] : tensor<1x5x33x62xf32>)
// CHECK: ^bb0(%[[BBARG1:[a-zA-Z0-9_]+]]: f32, // CHECK: ^bb0(%[[BBARG1:[a-zA-Z0-9_]+]]: f32,
// 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 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
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:.+]] = tensor.empty(%[[BATCH]]) // CHECK: %[[POOLINIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[FILL:.+]] = linalg.fill // CHECK: %[[FILL:.+]] = linalg.fill
// CHECK: %[[KERNEL:.+]] = tensor.empty() // 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>) inits(%[[FILL]] : tensor<?x5x33x62xf32>)
// CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]]) // 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>) inits(%[[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
func.func @avg_pool_i8(%arg0 : tensor<1x128x128x2xi8>) -> () { func.func @avg_pool_i8(%arg0 : tensor<1x128x128x2xi8>) -> () {
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines
// 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:.+]] = tensor.empty() // 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:.+]] = tensor.empty() // 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>) inits(%[[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>) inits(%[[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:.+]] = tensor.empty(%[[BATCH]]) // 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:.+]] = tensor.empty(%[[BATCH]]) // 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>) inits(%[[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>) inits(%[[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 20 Lines • Show All 42 Lines • ▼ Show 20 Linesfunc.func @conv2d_dyn_w_h(%input: tensor<1x?x?x27xf32>, %weights: tensor<28x3x3x27xf32>, %bias: tensor<28xf32>) -> () {
// 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:.+]] = tensor.empty(%[[H_OUT]], %[[W_OUT]]) // 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:.+]] = tensor.empty(%[[H_OUT]], %[[W_OUT]]) // 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>) inits(%[[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>) inits(%[[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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[INIT]]
// CHECK: [[OUT:%.+]] = tensor.empty() // 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>) inits([[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>) inits([[OUT]] : tensor<1x5x5x33xf32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: f32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:.+]] = tensor.empty(%[[BATCH]]) // 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:.+]] = tensor.empty(%[[BATCH]]) // 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>) inits(%[[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>) inits(%[[OUT]] : tensor<?x5x5x33xf32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: f32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[INIT]]
// CHECK: [[OUT:%.+]] = tensor.empty() // 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>) inits([[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>) inits([[OUT]] : tensor<1x5x5x33xf32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: f32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: f32, %[[ARG4:[0-9a-zA-Z_]+]]: f32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[INIT]]
// CHECK: [[OUT:%.+]] = tensor.empty() // 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) inits([[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>) inits([[OUT]] : tensor<1x12x12x512xi32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: i32, %[[ARG4:[0-9a-zA-Z_]+]]: i32, %[[ARG5:[0-9a-zA-Z_]+]]: i32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: i32, %[[ARG4:[0-9a-zA-Z_]+]]: i32, %[[ARG5:[0-9a-zA-Z_]+]]: 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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[INIT]]
// CHECK: [[OUT:%.+]] = tensor.empty() // 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) inits([[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>) inits([[OUT]] : tensor<1x10x10x512xi32>) {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: i32, %[[ARG4:[0-9a-zA-Z_]+]]: i32, %[[ARG5:[0-9a-zA-Z_]+]]: i32): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: i32, %[[ARG4:[0-9a-zA-Z_]+]]: i32, %[[ARG5:[0-9a-zA-Z_]+]]: 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<1x10x10x512xi32> // CHECK: } -> tensor<1x10x10x512xi32>
%0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], quantization_info = #tosa.conv_quant<input_zp = -128, weight_zp = 42>, stride = [1, 1], dilation = [2, 2] } : (tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>) -> tensor<1x10x10x512xi32> %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], quantization_info = #tosa.conv_quant<input_zp = -128, weight_zp = 42>, stride = [1, 1], dilation = [2, 2] } : (tensor<1x14x14x4xi8>, tensor<3x3x4x128xi8>, tensor<512xi32>) -> tensor<1x10x10x512xi32>
return return
} }
// CHECK-LABEL: @depthwise_conv2d_dyn_w_h // CHECK-LABEL: @depthwise_conv2d_dyn_w_h
func.func @depthwise_conv2d_dyn_w_h(%arg0: tensor<2x?x?x3xf32>, %arg1: tensor<3x6x3x5xf32>, %arg2: tensor<15xf32>) { func.func @depthwise_conv2d_dyn_w_h(%arg0: tensor<2x?x?x3xf32>, %arg1: tensor<3x6x3x5xf32>, %arg2: tensor<15xf32>) {
// CHECK: arith.addi // CHECK: arith.addi
// CHECK: arith.subi // CHECK: arith.subi
// CHECK: arith.muli // CHECK: arith.muli
// CHECK: arith.divui // CHECK: arith.divui
// CHECK: %[[PADDED:.+]] = tensor.pad %arg0 low[0, 1, 3, 0] high[0, 2, 4, 0] { // CHECK: %[[PADDED:.+]] = tensor.pad %arg0 low[0, 1, 3, 0] high[0, 2, 4, 0] {
// CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: index, %[[ARG4:[0-9a-zA-Z_]+]]: index, %[[ARG5:[0-9a-zA-Z_]+]]: index, %[[ARG6:[0-9a-zA-Z_]+]]: index): // CHECK: ^bb0(%[[ARG3:[0-9a-zA-Z_]+]]: index, %[[ARG4:[0-9a-zA-Z_]+]]: index, %[[ARG5:[0-9a-zA-Z_]+]]: index, %[[ARG6:[0-9a-zA-Z_]+]]: index):
// CHECK: tensor.yield %cst : f32 // CHECK: tensor.yield %cst : f32
// CHECK: } : tensor<2x?x?x3xf32> to tensor<2x?x?x3xf32> // CHECK: } : tensor<2x?x?x3xf32> to tensor<2x?x?x3xf32>
// CHECK: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>} ins(%[[PADDED]], %arg1 : tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>) outs(%{{.*}} : tensor<2x?x?x3x5xf32>) -> tensor<2x?x?x3x5xf32> // CHECK: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwcm {dilations = dense<[2, 1]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>} ins(%[[PADDED]], %arg1 : tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>) inits(%{{.*}} : tensor<2x?x?x3x5xf32>) -> tensor<2x?x?x3x5xf32>
// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[CONV]] {{\[}}[0], [1], [2], [3, 4]] // CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[CONV]] {{\[}}[0], [1], [2], [3, 4]]
%0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 2, 3, 4], dilation = [2, 1], stride = [1, 2]} : (tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>, tensor<15xf32>) -> tensor<2x?x?x15xf32> %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 2, 3, 4], dilation = [2, 1], stride = [1, 2]} : (tensor<2x?x?x3xf32>, tensor<3x6x3x5xf32>, tensor<15xf32>) -> tensor<2x?x?x15xf32>
return return
} }

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

// RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg)" %s -o -| FileCheck %s // RUN: mlir-opt --split-input-file -pass-pipeline="func.func(tosa-to-linalg)" %s -o -| FileCheck %s
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d3)> // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 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: @broadcast_resize_nearest_fp // CHECK-LABEL: @broadcast_resize_nearest_fp
func.func @broadcast_resize_nearest_fp(%arg0 : tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> { func.func @broadcast_resize_nearest_fp(%arg0 : tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> {
// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0
// CHECK-SAME{literal}: [[0], [1, 2, 3]] // CHECK-SAME{literal}: [[0], [1, 2, 3]]
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xf32> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#map, #map1] // CHECK-SAME: indexing_maps = [#map, #map1]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xf32>) // CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xf32>)
// CHECK-SAME: outs(%[[EMPTY]] : tensor<3x15x13x7xf32>) // CHECK-SAME: inits(%[[EMPTY]] : tensor<3x15x13x7xf32>)
// CHECK-NEXT: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32): // CHECK-NEXT: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
// CHECK: linalg.yield %[[IN]] // CHECK: linalg.yield %[[IN]]
%resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> %resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32>
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
return %resize : tensor<3x15x13x7xf32> return %resize : tensor<3x15x13x7xf32>
} }
// ----- // -----
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d3)> // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 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: @broadcast_resize_bilinear_fp // CHECK-LABEL: @broadcast_resize_bilinear_fp
func.func @broadcast_resize_bilinear_fp(%arg0 : tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> { func.func @broadcast_resize_bilinear_fp(%arg0 : tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> {
// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0
// CHECK-SAME{literal}: [[0], [1, 2, 3]] // CHECK-SAME{literal}: [[0], [1, 2, 3]]
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xf32> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#map, #map1] // CHECK-SAME: indexing_maps = [#map, #map1]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xf32>) // CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xf32>)
// CHECK-SAME: outs(%[[EMPTY]] : tensor<3x15x13x7xf32>) // CHECK-SAME: inits(%[[EMPTY]] : tensor<3x15x13x7xf32>)
// CHECK-NEXT: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32): // CHECK-NEXT: ^bb0(%[[IN:.+]]: f32, %[[OUT:.+]]: f32):
// CHECK: linalg.yield %[[IN]] // CHECK: linalg.yield %[[IN]]
%resize = "tosa.resize"(%arg0) {mode = "BILINEAR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32> %resize = "tosa.resize"(%arg0) {mode = "BILINEAR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xf32>) -> tensor<3x15x13x7xf32>
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
return %resize : tensor<3x15x13x7xf32> return %resize : tensor<3x15x13x7xf32>
} }
// ----- // -----
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d3)> // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 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: @broadcast_resize_nearest_i8 // CHECK-LABEL: @broadcast_resize_nearest_i8
func.func @broadcast_resize_nearest_i8(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi8> { func.func @broadcast_resize_nearest_i8(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi8> {
// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0
// CHECK-SAME{literal}: [[0], [1, 2, 3]] // CHECK-SAME{literal}: [[0], [1, 2, 3]]
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi8> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi8>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#map, #map1] // CHECK-SAME: indexing_maps = [#map, #map1]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>) // CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>)
// CHECK-SAME: outs(%[[EMPTY]] : tensor<3x15x13x7xi8>) // CHECK-SAME: inits(%[[EMPTY]] : tensor<3x15x13x7xi8>)
// CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i8): // CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i8):
// CHECK: linalg.yield %[[IN]] // CHECK: linalg.yield %[[IN]]
%resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi8> %resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi8>
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
return %resize : tensor<3x15x13x7xi8> return %resize : tensor<3x15x13x7xi8>
} }
// ----- // -----
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d3)> // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 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: @broadcast_resize_nearest_i32 // CHECK-LABEL: @broadcast_resize_nearest_i32
func.func @broadcast_resize_nearest_i32(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> { func.func @broadcast_resize_nearest_i32(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> {
// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0
// CHECK-SAME{literal}: [[0], [1, 2, 3]] // CHECK-SAME{literal}: [[0], [1, 2, 3]]
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi32> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#map, #map1] // CHECK-SAME: indexing_maps = [#map, #map1]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>) // CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>)
// CHECK-SAME: outs(%[[EMPTY]] : tensor<3x15x13x7xi32>) // CHECK-SAME: inits(%[[EMPTY]] : tensor<3x15x13x7xi32>)
// CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i32): // CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i32):
// CHECK: %[[EXT:.+]] = arith.extsi %[[IN]] : i8 to i32 // CHECK: %[[EXT:.+]] = arith.extsi %[[IN]] : i8 to i32
// CHECK: linalg.yield %[[EXT]] // CHECK: linalg.yield %[[EXT]]
%resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> %resize = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32>
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
return %resize : tensor<3x15x13x7xi32> return %resize : tensor<3x15x13x7xi32>
} }
// ----- // -----
// CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, d3)> // CHECK: #map = affine_map<(d0, d1, d2, d3) -> (d0, 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: @broadcast_resize_bilinear_i32 // CHECK-LABEL: @broadcast_resize_bilinear_i32
func.func @broadcast_resize_bilinear_i32(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> { func.func @broadcast_resize_bilinear_i32(%arg0 : tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> {
// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0 // CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %arg0
// CHECK-SAME{literal}: [[0], [1, 2, 3]] // CHECK-SAME{literal}: [[0], [1, 2, 3]]
// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi32> // CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x15x13x7xi32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#map, #map1] // CHECK-SAME: indexing_maps = [#map, #map1]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>) // CHECK-SAME: ins(%[[COLLAPSE]] : tensor<3x7xi8>)
// CHECK-SAME: outs(%[[EMPTY]] : tensor<3x15x13x7xi32>) // CHECK-SAME: inits(%[[EMPTY]] : tensor<3x15x13x7xi32>)
// CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i32): // CHECK-NEXT: ^bb0(%[[IN:.+]]: i8, %[[OUT:.+]]: i32):
// CHECK: %[[EXT:.+]] = arith.extsi %[[IN]] : i8 to i32 // CHECK: %[[EXT:.+]] = arith.extsi %[[IN]] : i8 to i32
// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : i32 // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : i32
// CHECK: %[[MUL1:.+]] = arith.muli %[[EXT]], %[[C2]] : i32 // CHECK: %[[MUL1:.+]] = arith.muli %[[EXT]], %[[C2]] : i32
// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i32 // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : i32
// CHECK: %[[MUL2:.+]] = arith.muli %[[MUL1]], %[[C1]] : i32 // CHECK: %[[MUL2:.+]] = arith.muli %[[MUL1]], %[[C1]] : i32
// CHECK: linalg.yield %[[MUL2]] // CHECK: linalg.yield %[[MUL2]]
%resize = "tosa.resize"(%arg0) {mode = "BILINEAR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32> %resize = "tosa.resize"(%arg0) {mode = "BILINEAR", scale = [2, 2, 1, 1], offset = [0, 0], border = [0, 0]} : (tensor<3x1x1x7xi8>) -> tensor<3x15x13x7xi32>
▲ Show 20 Lines • Show All 375 Lines • Show Last 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
func.func @test_abs(%arg0: tensor<f32>) -> tensor<f32> { func.func @test_abs(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: [[INIT:%.+]] = tensor.empty() : 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>) inits([[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
func.func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> { func.func @test_abs(%arg0: tensor<2xf32>) -> tensor<2xf32> {
// CHECK: [[INIT:%.+]] = tensor.empty() : 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>) inits([[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
func.func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> { func.func @test_abs(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
// CHECK: [[INIT:%.+]] = tensor.empty() : 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>) inits([[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>
Show All 35 Lines
// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @test_broadcast // CHECK-LABEL: @test_broadcast
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<1xf32 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<1xf32
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: tensor<2xf32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: tensor<2xf32>
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:%.+]] = tensor.empty() : 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>) inits([[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
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<2xf32 // CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<2xf32
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: tensor<1xf32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: tensor<1xf32>
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:%.+]] = tensor.empty() : 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>) inits([[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]] // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
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:%.+]] = tensor.empty() : 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>) inits([[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>
return %0 : tensor<2x3xf32> return %0 : tensor<2x3xf32>
} }
▲ Show 20 Lines • Show All 498 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:%.+]] = tensor.empty() : tensor<2x3x1xi32> // 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>) inits([[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:.+]] = tensor.empty(%[[DIM]]) : tensor<?x4x1x3xi32> // 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>) inits([[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:.+]] = tensor.empty(%[[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>) inits([[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:%.+]] = tensor.empty() : tensor<4xf32> // 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]]{{.*}}inits([[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>) inits([[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:%.+]] = tensor.empty() : tensor<5xf32> // 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]]{{.*}}inits([[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>) inits([[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>
// CHECK: arith.constant 1.0 // CHECK: arith.constant 1.0
// CHECK: linalg.fill // CHECK: linalg.fill
Show All 22 Lines
// CHECK-LABEL: @reduce_float_dyn // CHECK-LABEL: @reduce_float_dyn
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?x5x4xf32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?x5x4xf32>
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:.+]] = tensor.empty(%[[DYN]]) : tensor<?x4xf32> // 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]]{{.*}}inits(%[[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>) inits(%[[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: tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]] : tensor<?x4xf32> into tensor<?x1x4xf32> // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]] : tensor<?x4xf32> 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) -> (d0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> ()> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> ()>
// CHECK-LABEL: @reduce_float_dyn_rank_1 // CHECK-LABEL: @reduce_float_dyn_rank_1
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?xf32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z_]*]]: tensor<?xf32>
func.func @reduce_float_dyn_rank_1(%arg0: tensor<?xf32>) -> () { func.func @reduce_float_dyn_rank_1(%arg0: tensor<?xf32>) -> () {
// CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<f32> // CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<f32>
// CHECK-DAG: %[[CST0:.+]] = arith.constant 0.0 // CHECK-DAG: %[[CST0:.+]] = arith.constant 0.0
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[CST0]]{{.*}}inits(%[[INIT]]
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction"]} ins(%[[ARG0]] : tensor<?xf32>) outs(%[[FILL]] : tensor<f32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], iterator_types = ["reduction"]} ins(%[[ARG0]] : tensor<?xf32>) inits(%[[FILL]] : tensor<f32>)
// 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]] {{\[}}] : tensor<f32> into tensor<1xf32> // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}] : tensor<f32> into tensor<1xf32>
%0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<?xf32>) -> tensor<1xf32> %0 = "tosa.reduce_sum"(%arg0) {axis = 0 : i64} : (tensor<?xf32>) -> tensor<1xf32>
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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[DYN]]) : tensor<5x?xf32> // 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]]{{.*}}inits(%[[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>) inits(%[[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: tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]] : tensor<5x?xf32> into tensor<5x?x1xf32> // CHECK: tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]] : tensor<5x?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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[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]]{{.*}}inits(%[[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>) inits(%[[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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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>) inits([[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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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>) inits([[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>
// CHECK: arith.constant 1 // CHECK: arith.constant 1
// CHECK: linalg.fill // CHECK: linalg.fill
Show All 22 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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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>) inits([[FILL]] : tensor<4xi1>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i1, %[[ARG2:[0-9a-zA-Z_]+]]: i1) // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i1, %[[ARG2:[0-9a-zA-Z_]+]]: 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>
// CHECK: arith.constant false // CHECK: arith.constant false
// CHECK: linalg.fill // CHECK: linalg.fill
Show All 12 Lines
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:%.+]] = tensor.empty() : tensor<11x1xf32> // 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]]{{.*}}inits([[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:%.+]] = tensor.empty() : tensor<5x2xf32> // 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]]{{.*}}inits([[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]] // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
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:.+]] = tensor.empty(%[[DYN]]) : tensor<11x?xf32> // 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]]{{.*}}inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[DYN]]) : tensor<?x3xf32> // 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]]{{.*}}inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:%.+]] = tensor.empty() // 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>) inits([[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:%.+]] = tensor.empty() // 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>) inits([[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]]
// CHECK-DAG: [[CMIN:%.+]] = arith.constant 0 // CHECK-DAG: [[CMIN:%.+]] = arith.constant 0
Show All 16 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x2xi8> // 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>) inits(%[[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:.+]] = tensor.empty(%[[BATCH]]) : tensor<?x2xui8> // 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>) inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[DIM1]], %[[DIM2]]) // 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>) inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:%.+]] = tensor.empty() // 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>) inits([[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}
// 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) -> (d0)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @rescale_per_channel // CHECK-LABEL: @rescale_per_channel
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:%.+]] = tensor.empty() // 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>) inits([[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}
// CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C252]] // CHECK-DAG: [[SCALED_ZEROED:%.+]] = arith.addi [[SCALED]], [[C252]]
Show All 36 Lines
// CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @reverse // CHECK-LABEL: @reverse
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty() // 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"]} inits(%[[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:.+]] = tensor.empty() // 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"]} inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[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"]} inits(%[[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
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3xi8> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x3xi8>
func.func @tile(%arg0 : tensor<2x3xi8>) -> () { func.func @tile(%arg0 : tensor<2x3xi8>) -> () {
// CHECK: [[INIT:%.+]] = tensor.empty() // 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>) inits([[INIT]] : tensor<2x2x1x3xi8>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8
// 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:%.+]] = tensor.empty() // 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>) inits([[INIT]] : tensor<1x2x2x3xi8>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8
// 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:%.+]] = tensor.empty() // 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>) inits([[INIT]] : tensor<5x2x7x3xi8>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i8
// 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[DYN]]) // 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>) inits(%[[INIT]] : tensor<2x?x1x3xi8>)
// CHECK: ^bb0(%[[ARG1:.+]]: i8, // CHECK: ^bb0(%[[ARG1:.+]]: i8,
// 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[DYN]]) // 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>) inits(%[[INIT]] : tensor<2x2x?x3xi8>)
// CHECK: ^bb0(%[[ARG1:.+]]: i8, // CHECK: ^bb0(%[[ARG1:.+]]: i8,
// 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: #[[$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:%.+]] = tensor.empty() // 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]]{{.*}}inits([[IDX_INIT]]
// CHECK: [[VAL_INIT:%.+]] = tensor.empty() // 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]]{{.*}}inits([[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>) inits([[IDX_FILL]], [[VAL_FILL]] : tensor<2xi32>, tensor<2xi32>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
// 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<3x2xi32>) -> (tensor<2xi32>) %0 = "tosa.argmax"(%arg0) { axis = 0 : i64} : (tensor<3x2xi32>) -> (tensor<2xi32>)
// CHECK: [[IDX_INIT:%.+]] = tensor.empty() // 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]]{{.*}}inits([[IDX_INIT]]
// CHECK: [[VAL_INIT:%.+]] = tensor.empty() // 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]]{{.*}}inits([[VAL_INIT]]
// CHECK: linalg.generic {indexing_maps = [#map, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x2xi32>) outs([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>) // CHECK: linalg.generic {indexing_maps = [#map, #map2, #map2], iterator_types = ["parallel", "reduction"]} ins(%[[ARG0]] : tensor<3x2xi32>) inits([[IDX_FILL]], [[VAL_FILL]] : tensor<3xi32>, tensor<3xi32>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
// 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]]
%1 = "tosa.argmax"(%arg0) { axis = 1 : i64} : (tensor<3x2xi32>) -> (tensor<3xi32>) %1 = "tosa.argmax"(%arg0) { axis = 1 : i64} : (tensor<3x2xi32>) -> (tensor<3xi32>)
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:.+]] = tensor.empty(%[[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]]{{.*}}inits(%[[IDX_INIT]]
// CHECK: %[[VAL_INIT:.+]] = tensor.empty(%[[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]]{{.*}}inits(%[[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>) inits(%[[IDX_FILL]], %[[VAL_FILL]] : tensor<?xi32>, tensor<?xi32>)
// CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32 // CHECK: ^bb0(%[[ARG1:[0-9a-zA-Z_]+]]: i32, %[[ARG2:[0-9a-zA-Z_]+]]: i32, %[[ARG3:[0-9a-zA-Z_]+]]: i32
// 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:.+]] = tensor.empty() // 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]]{{.*}}inits(%[[IDX_INIT]]
// CHECK: %[[VAL_INIT:.+]] = tensor.empty() // 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]]{{.*}}inits(%[[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>) inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]] // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
func.func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () { func.func @gather_float(%arg0: tensor<2x3x2xf32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty() // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) inits(%[[INIT]] : tensor<2x3x2xf32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32) // CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
// 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]] // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
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:.+]] = tensor.empty(%[[BATCH]]) // CHECK: %[[INIT:.+]] = tensor.empty(%[[BATCH]])
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<?x3xi32>) outs(%[[INIT]] : tensor<?x3x2xf32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<?x3xi32>) inits(%[[INIT]] : tensor<?x3x2xf32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32) // CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: f32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
// 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]] // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]] // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]
func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () { func.func @gather_int(%arg0: tensor<2x3x2xi32>, %arg1: tensor<2x3xi32>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty() // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) outs(%[[INIT]] : tensor<2x3x2xi32>) // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel"]} ins(%[[ARG1]] : tensor<2x3xi32>) inits(%[[INIT]] : tensor<2x3x2xi32>)
// CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32) // CHECK: ^bb0(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32)
// CHECK: %[[IDX0:.+]] = linalg.index 0 // CHECK: %[[IDX0:.+]] = linalg.index 0
// CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]] // CHECK: %[[CAST:.+]] = arith.index_cast %[[BBARG0]]
// 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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
func.func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () { func.func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty() // 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>) inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
func.func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () { func.func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty() // 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>) inits(%[[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]]
// CHECK: %[[IDX:.+]] = arith.shrui %[[INADD]], %[[C7]] // CHECK: %[[IDX:.+]] = arith.shrui %[[INADD]], %[[C7]]
Show All 18 Lines
// CHECK-LABEL: @table8_dyn // CHECK-LABEL: @table8_dyn
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
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:.+]] = tensor.empty(%[[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>) inits(%[[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
// CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]: // CHECK-SAME: (%[[ARG0:[0-9a-zA-Z_]*]]:
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]: // CHECK-SAME: %[[ARG1:[0-9a-zA-Z_]*]]:
func.func @table8_dyn_table(%arg0: tensor<6xi8>, %arg1: tensor<?xi8>) -> () { func.func @table8_dyn_table(%arg0: tensor<6xi8>, %arg1: tensor<?xi8>) -> () {
// CHECK: %[[INIT:.+]] = tensor.empty() // 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>) inits(%[[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
Show All 16 Lines

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

// RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | FileCheck %s // RUN: mlir-opt %s -eliminate-alloc-tensors -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs" -canonicalize -split-input-file | FileCheck %s
// CHECK: func @buffer_forwarding_conflict( // CHECK: func @buffer_forwarding_conflict(
// CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32> // CHECK-SAME: %[[FUNC_ARG:[0-9a-zA-Z]*]]: memref<?xf32>
// CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index // CHECK-SAME: %[[sz:[0-9a-zA-Z]*]]: index
func.func @buffer_forwarding_conflict( func.func @buffer_forwarding_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>) -> (tensor<?xf32>, tensor<?xf32>)
{ {
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
// CHECK: %[[EXTRACT_SLICE_ALLOC:.*]] = memref.alloc(%[[sz]]) // CHECK: %[[EXTRACT_SLICE_ALLOC:.*]] = memref.alloc(%[[sz]])
// CHECK: linalg.fill ins({{.*}} : f32) outs(%[[EXTRACT_SLICE_ALLOC]] : memref<?xf32>) // CHECK: linalg.fill ins({{.*}} : f32) inits(%[[EXTRACT_SLICE_ALLOC]] : memref<?xf32>)
// Alloc is needed for the **first** insert_slice (due to backward traversal during analysis). // Alloc is needed for the **first** insert_slice (due to backward traversal during analysis).
// CHECK: %[[DIM:.*]] = memref.dim %[[FUNC_ARG]] // CHECK: %[[DIM:.*]] = memref.dim %[[FUNC_ARG]]
// This allocs the whole dim to allow for a full clone of t. // This allocs the whole dim to allow for a full clone of t.
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]]) // CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM]])
// alloc_tensor itself does not alloc but forwards to the **second** // alloc_tensor itself does not alloc but forwards to the **second**
// insert_slice. AllocTensorOp replaces the alloc_tensor with an out-of-place // insert_slice. AllocTensorOp replaces the alloc_tensor with an out-of-place
// extract_slice. // extract_slice.
%a = bufferization.alloc_tensor(%sz) : tensor<?xf32> %a = bufferization.alloc_tensor(%sz) : tensor<?xf32>
%f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32> %f = linalg.fill ins(%f0 : f32) inits(%a : tensor<?xf32>) -> tensor<?xf32>
// CHECK: memref.copy %[[FUNC_ARG]], %[[ALLOC]] : memref<?xf32> to memref<?xf32> // CHECK: memref.copy %[[FUNC_ARG]], %[[ALLOC]] : memref<?xf32> to memref<?xf32>
// CHECK: %[[SV0_ALLOC:.*]] = memref.subview %[[ALLOC]][0] [%[[sz]]] [1] : memref<?xf32> to memref<?xf32, strided<[1]>> // CHECK: %[[SV0_ALLOC:.*]] = memref.subview %[[ALLOC]][0] [%[[sz]]] [1] : memref<?xf32> to memref<?xf32, strided<[1]>>
// CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[SV0_ALLOC]] : memref<?xf32> to memref<?xf32, strided<[1]>> // CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[SV0_ALLOC]] : memref<?xf32> to memref<?xf32, strided<[1]>>
%r0 = tensor.insert_slice %f into %t[0][%sz][1]: tensor<?xf32> into tensor<?xf32> %r0 = tensor.insert_slice %f into %t[0][%sz][1]: tensor<?xf32> into tensor<?xf32>
// CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1] // CHECK: %[[T_SUBVIEW:.*]] = memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
// CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[T_SUBVIEW]] // CHECK: memref.copy %[[EXTRACT_SLICE_ALLOC]], %[[T_SUBVIEW]]
Show All 15 Linesfunc.func @buffer_forwarding_no_conflict(
%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.
// AllocTensorOpElimination replaces the alloc_tensor with an inplace // AllocTensorOpElimination replaces the alloc_tensor with an inplace
// extract_slice. // 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) inits(%[[T_SUBVIEW]] : memref<?xf32
%f = linalg.fill ins(%f0 : f32) outs(%a : tensor<?xf32>) -> tensor<?xf32> %f = linalg.fill ins(%f0 : f32) inits(%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>
return %r1: tensor<?xf32> return %r1: tensor<?xf32>
} }
// ----- // -----
Show All 9 Linesfunc.func @insertion_point_inside_loop(%t : tensor<?xf32>, %sz : index) -> (tensor<?xf32>) {
%blank = bufferization.alloc_tensor() : tensor<5xf32> %blank = bufferization.alloc_tensor() : tensor<5xf32>
// CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} { // CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
%r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) { %r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
// CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[iv]]] [5] [1] // CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[iv]]] [5] [1]
%iv_i32 = arith.index_cast %iv : index to i32 %iv_i32 = arith.index_cast %iv : index to i32
%f = arith.sitofp %iv_i32 : i32 to f32 %f = arith.sitofp %iv_i32 : i32 to f32
// CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]] // CHECK: linalg.fill ins(%{{.*}}{{.*}}inits(%[[subview]]
%filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32> %filled = linalg.fill ins(%f : f32) inits(%blank : tensor<5xf32>) -> tensor<5xf32>
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
%inserted = tensor.insert_slice %filled into %bb[%iv][5][1] : tensor<5xf32> into tensor<?xf32> %inserted = tensor.insert_slice %filled into %bb[%iv][5][1] : tensor<5xf32> into tensor<?xf32>
scf.yield %inserted : tensor<?xf32> scf.yield %inserted : tensor<?xf32>
} }
return %r : tensor<?xf32> return %r : tensor<?xf32>
} }
Show All 12 Linesfunc.func @insertion_point_outside_loop(%t : tensor<?xf32>, %sz : index,
// CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[idx]]] [5] [1] // CHECK: %[[subview:.*]] = memref.subview %[[t]][%[[idx]]] [5] [1]
%blank = bufferization.alloc_tensor() : tensor<5xf32> %blank = bufferization.alloc_tensor() : tensor<5xf32>
// CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} { // CHECK: scf.for %[[iv:.*]] = %{{.*}} to %[[sz]] step %{{.*}} {
%r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) { %r = scf.for %iv = %c0 to %sz step %c5 iter_args(%bb = %t) -> (tensor<?xf32>) {
%iv_i32 = arith.index_cast %iv : index to i32 %iv_i32 = arith.index_cast %iv : index to i32
%f = arith.sitofp %iv_i32 : i32 to f32 %f = arith.sitofp %iv_i32 : i32 to f32
// CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[subview]] // CHECK: linalg.fill ins(%{{.*}}{{.*}}inits(%[[subview]]
%filled = linalg.fill ins(%f : f32) outs(%blank : tensor<5xf32>) -> tensor<5xf32> %filled = linalg.fill ins(%f : f32) inits(%blank : tensor<5xf32>) -> tensor<5xf32>
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
%inserted = tensor.insert_slice %filled into %bb[%idx][5][1] : tensor<5xf32> into tensor<?xf32> %inserted = tensor.insert_slice %filled into %bb[%idx][5][1] : tensor<5xf32> into tensor<?xf32>
scf.yield %inserted : tensor<?xf32> scf.yield %inserted : tensor<?xf32>
} }
return %r : tensor<?xf32> return %r : tensor<?xf32>
} }
// ----- // -----
// AllocTensorElimination does currently not apply to chains where the type is // AllocTensorElimination does currently not apply to chains where the type is
// changing. This test just ensures that we do not crash or generate IR that // changing. This test just ensures that we do not crash or generate IR that
// does not verify. // does not verify.
// CHECK-LABEL: func @shape_mismatch // CHECK-LABEL: func @shape_mismatch
func.func @shape_mismatch(%t: tensor<5x6x128xf32>) -> tensor<5x6x128xf32> { func.func @shape_mismatch(%t: tensor<5x6x128xf32>) -> tensor<5x6x128xf32> {
%cst = arith.constant 8.0 : f32 %cst = arith.constant 8.0 : f32
%0 = bufferization.alloc_tensor() : tensor<128xf32> %0 = bufferization.alloc_tensor() : tensor<128xf32>
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<128xf32>) -> tensor<128xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<128xf32>) -> tensor<128xf32>
%2 = tensor.expand_shape %1 [[0, 1, 2]] %2 = tensor.expand_shape %1 [[0, 1, 2]]
: tensor<128xf32> into tensor<1x1x128xf32> : tensor<128xf32> into tensor<1x1x128xf32>
%3 = tensor.insert_slice %2 into %t[2, 3, 0][1, 1, 128][1, 1, 1] %3 = tensor.insert_slice %2 into %t[2, 3, 0][1, 1, 128][1, 1, 1]
: tensor<1x1x128xf32> into tensor<5x6x128xf32> : tensor<1x1x128xf32> into tensor<5x6x128xf32>
return %3 : tensor<5x6x128xf32> return %3 : tensor<5x6x128xf32>
} }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-partial.mlir

Show First 20 Lines • Show All 131 Lines • ▼ Show 20 Lines
func.func @unknown_op_may_read(%v: vector<5xf32>) func.func @unknown_op_may_read(%v: vector<5xf32>)
-> (tensor<10xf32>, tensor<10xf32>) { -> (tensor<10xf32>, tensor<10xf32>) {
%idx = arith.constant 0 : index %idx = arith.constant 0 : index
%cst = arith.constant 5.0 : f32 %cst = arith.constant 5.0 : f32
// One alloc for the alloc_tensor, another one because the transfer_write // One alloc for the alloc_tensor, another one because the transfer_write
// bufferizes out-of-place. // bufferizes out-of-place.
// CHECK: %[[m1:.*]] = memref.alloc() {{.*}} : memref<10xf32> // CHECK: %[[m1:.*]] = memref.alloc() {{.*}} : memref<10xf32>
// CHECK: linalg.fill ins(%{{.*}}{{.*}}outs(%[[m1]] // CHECK: linalg.fill ins(%{{.*}}{{.*}}inits(%[[m1]]
// CHECK: %[[filled_tensor:.*]] = bufferization.to_tensor %[[m1]] // CHECK: %[[filled_tensor:.*]] = bufferization.to_tensor %[[m1]]
%t1 = bufferization.alloc_tensor() : tensor<10xf32> %t1 = bufferization.alloc_tensor() : tensor<10xf32>
%filled = linalg.fill ins(%cst : f32) outs(%t1 : tensor<10xf32>) -> tensor<10xf32> %filled = linalg.fill ins(%cst : f32) inits(%t1 : tensor<10xf32>) -> tensor<10xf32>
// The transfer_write is out-of-place because "dummy_op" may read. // The transfer_write is out-of-place because "dummy_op" may read.
// CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32> // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10xf32>
// CHECK: memref.copy %[[m1]], %[[alloc]] // CHECK: memref.copy %[[m1]], %[[alloc]]
// CHECK: vector.transfer_write %{{.*}}, %[[alloc]] // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
// CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]] // CHECK: %[[alloc_tensor:.*]] = bufferization.to_tensor %[[alloc]]
%1 = vector.transfer_write %v, %filled[%idx] : vector<5xf32>, tensor<10xf32> %1 = vector.transfer_write %v, %filled[%idx] : vector<5xf32>, tensor<10xf32>
▲ Show 20 Lines • Show All 71 Lines • Show Last 20 Lines

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-allow-return-allocs.mlir

Show First 20 Lines • Show All 45 Lines • ▼ Show 20 Linesfunc.func @return_slice(%t: tensor<?xf32>, %sz: index) -> (tensor<?xf32>) {
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// CHECK-LABEL: func @main( // CHECK-LABEL: func @main(
// CHECK-SAME: %[[t:.*]]: memref<?xf32 // CHECK-SAME: %[[t:.*]]: memref<?xf32
// CHECK: %[[call:.*]] = call @return_slice(%[[t]] // CHECK: %[[call:.*]] = call @return_slice(%[[t]]
// CHECK: %[[alloc:.*]] = memref.alloc // CHECK: %[[alloc:.*]] = memref.alloc
// CHECK: memref.copy %[[call]], %[[alloc]] // CHECK: memref.copy %[[call]], %[[alloc]]
// CHECK: linalg.fill ins({{.*}}) outs(%[[t]] // CHECK: linalg.fill ins({{.*}}) inits(%[[t]]
// CHECK: memref.load %[[alloc]] // CHECK: memref.load %[[alloc]]
// CHECK: memref.load %[[t]] // CHECK: memref.load %[[t]]
func.func @main(%t: tensor<?xf32>, %sz: index, %idx: index) -> (f32, f32) { func.func @main(%t: tensor<?xf32>, %sz: index, %idx: index) -> (f32, f32) {
%cst = arith.constant 1.0 : f32 %cst = arith.constant 1.0 : f32
%0 = call @return_slice(%t, %sz) : (tensor<?xf32>, index) -> (tensor<?xf32>) %0 = call @return_slice(%t, %sz) : (tensor<?xf32>, index) -> (tensor<?xf32>)
%filled = linalg.fill ins(%cst : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32> %filled = linalg.fill ins(%cst : f32) inits(%t : tensor<?xf32>) -> tensor<?xf32>
%r1 = tensor.extract %0[%idx] : tensor<?xf32> %r1 = tensor.extract %0[%idx] : tensor<?xf32>
%r2 = tensor.extract %filled[%idx] : tensor<?xf32> %r2 = tensor.extract %filled[%idx] : tensor<?xf32>
return %r1, %r2 : f32, f32 return %r1, %r2 : f32, f32
} }
// ----- // -----
func.func @return_arg(%A: tensor<?xf32>) -> tensor<?xf32> { func.func @return_arg(%A: tensor<?xf32>) -> tensor<?xf32> {
Show All 9 Lines

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-analysis.mlir

Show First 20 Lines • Show All 69 Lines • ▼ Show 20 Lines// CHECK-SAME: bufferization.access = "read"
%B : tensor<4x4xf32> {bufferization.writable = true}) %B : tensor<4x4xf32> {bufferization.writable = true})
// CHECK-SAME: bufferization.access = "read-write" // CHECK-SAME: bufferization.access = "read-write"
-> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) -> (tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>)
{ {
// matmul output operand interferes with input operand. // matmul output operand interferes with input operand.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
%C = linalg.matmul ins(%A, %B: tensor<4x4xf32>, tensor<4x4xf32>) %C = linalg.matmul ins(%A, %B: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%B: tensor<4x4xf32>) inits(%B: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// matmul output operand interferes with input operand. // matmul output operand interferes with input operand.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
%D = linalg.matmul ins(%B, %A: tensor<4x4xf32>, tensor<4x4xf32>) %D = linalg.matmul ins(%B, %A: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%B: tensor<4x4xf32>) inits(%B: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// matmul output operand does not interferes with input operand. // matmul output operand does not interferes with input operand.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%E = linalg.matmul ins(%A, %A: tensor<4x4xf32>, tensor<4x4xf32>) %E = linalg.matmul ins(%A, %A: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%B: tensor<4x4xf32>) inits(%B: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [-1, -1, 1] // CHECK-SAME: __equivalent_func_args__ = [-1, -1, 1]
return %C, %D, %E: tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32> return %C, %D, %E: tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>
} }
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
▲ Show 20 Lines • Show All 154 Lines • ▼ Show 20 Linesfunc.func @read_of_matching_insert_slice_source(
%cst2 = arith.constant 1.0 : f32 %cst2 = arith.constant 1.0 : f32
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
%0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32> %0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?xf32>) -> tensor<?xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<?xf32>) -> tensor<?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
%2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32> %2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
%3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32> %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
// CHECK: return // CHECK: return
Show All 15 Linesfunc.func @read_of_matching_insert_slice_source_interleaved(
%cst2 = arith.constant 1.0 : f32 %cst2 = arith.constant 1.0 : f32
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["false", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["false", "none", "none"]}
%0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32> %0 = tensor.extract_slice %A[%idx][%idx][1] : tensor<?xf32> to tensor<?xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?xf32>) -> tensor<?xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<?xf32>) -> tensor<?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
%2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32> %2 = tensor.insert_slice %1 into %A[%idx][%idx][1] : tensor<?xf32> into tensor<?xf32>
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
%4 = tensor.extract_slice %2[%idx3][%idx3][1] : tensor<?xf32> to tensor<?xf32> %4 = tensor.extract_slice %2[%idx3][%idx3][1] : tensor<?xf32> to tensor<?xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%5 = linalg.fill ins(%cst : f32) outs(%4 : tensor<?xf32>) -> tensor<?xf32> %5 = linalg.fill ins(%cst : f32) inits(%4 : tensor<?xf32>) -> tensor<?xf32>
%3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32> %3 = vector.transfer_read %1[%idx2], %cst2 : tensor<?xf32>, vector<5xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
%6 = tensor.insert_slice %5 into %2[%idx3][%idx3][1] : tensor<?xf32> into tensor<?xf32> %6 = tensor.insert_slice %5 into %2[%idx3][%idx3][1] : tensor<?xf32> into tensor<?xf32>
// CHECK: return // CHECK: return
Show All 15 Linesfunc.func @extract_slice_linalg_readonly_use(
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
%sA = tensor.extract_slice %A[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %sA = tensor.extract_slice %A[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
// matmul output operand is not inplaceable at the function boundary. // matmul output operand is not inplaceable at the function boundary.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
%D = linalg.matmul ins(%sA, %B: tensor<4x4xf32>, tensor<4x4xf32>) %D = linalg.matmul ins(%sA, %B: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%B: tensor<4x4xf32>) inits(%B: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// matmul output operand is inplaceable at the function boundary. // matmul output operand is inplaceable at the function boundary.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%E = linalg.matmul ins(%sA, %B: tensor<4x4xf32>, tensor<4x4xf32>) %E = linalg.matmul ins(%sA, %B: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%C: tensor<4x4xf32>) inits(%C: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [-1, 2] // CHECK-SAME: __equivalent_func_args__ = [-1, 2]
return %D, %E: tensor<4x4xf32>, tensor<4x4xf32> return %D, %E: tensor<4x4xf32>, tensor<4x4xf32>
} }
// ----- // -----
Show All 10 Linesfunc.func @extract_slice_to_linalg_write_use(
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
%sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
// Step 3. %sB has a read interference in %E, it does not bufferize inplace. // Step 3. %sB has a read interference in %E, it does not bufferize inplace.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]}
%D = linalg.matmul ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>) %D = linalg.matmul ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%sB: tensor<4x4xf32>) inits(%sB: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// Step 2. %sC forward propagates to an inplace write in %E. // Step 2. %sC forward propagates to an inplace write in %E.
// %sC backward propagates to %C which is inplaceable. // %sC backward propagates to %C which is inplaceable.
// As a consequence this is bufferized inplace. // As a consequence this is bufferized inplace.
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
%sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
// Step 1. %sC backprops to the tensor.extract_slice producer which is not // Step 1. %sC backprops to the tensor.extract_slice producer which is not
// considered an interference. This bufferizes inplace. // considered an interference. This bufferizes inplace.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%E = linalg.matmul ins(%A, %sB: tensor<4x4xf32>, tensor<4x4xf32>) %E = linalg.matmul ins(%A, %sB: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%sC: tensor<4x4xf32>) inits(%sC: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
return %D, %E: tensor<4x4xf32>, tensor<4x4xf32> return %D, %E: tensor<4x4xf32>, tensor<4x4xf32>
} }
// ----- // -----
// CHECK-LABEL: func @insert_slice_double_extract_slice // CHECK-LABEL: func @insert_slice_double_extract_slice
func.func @insert_slice_double_extract_slice( func.func @insert_slice_double_extract_slice(
%s1: index, %s1: index,
%s2: index, %s2: index,
%s3: index, %s3: index,
%s4: index, %s4: index,
%A: tensor<8x6xf32> {bufferization.writable = false}, %A: tensor<8x6xf32> {bufferization.writable = false},
%B: tensor<6x6xf32> {bufferization.writable = false}, %B: tensor<6x6xf32> {bufferization.writable = false},
%C: tensor<30x20xf32> {bufferization.writable = true}) %C: tensor<30x20xf32> {bufferization.writable = true})
-> tensor<30x20xf32> -> tensor<30x20xf32>
{ {
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none", "none", "none"]}
%15 = tensor.extract_slice %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<30x20xf32> to tensor<?x?xf32> %15 = tensor.extract_slice %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<30x20xf32> to tensor<?x?xf32>
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%18 = linalg.matmul ins(%A, %B : tensor<8x6xf32>, tensor<6x6xf32>) outs(%15 : tensor<?x?xf32>) -> tensor<?x?xf32> %18 = linalg.matmul ins(%A, %B : tensor<8x6xf32>, tensor<6x6xf32>) inits(%15 : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
%19 = tensor.extract_slice %18[0, 0] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %19 = tensor.extract_slice %18[0, 0] [%s1, %s2] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none", "none", "none"]}
%20 = tensor.insert_slice %19 into %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<?x?xf32> into tensor<30x20xf32> %20 = tensor.insert_slice %19 into %C[%s3, %s4] [%s1, %s2] [1, 1] : tensor<?x?xf32> into tensor<30x20xf32>
Show All 23 Linesfunc.func @extract_slice_to_linalg_write_use(
// CHECK-SAME: {__inplace_operands_attr__ = ["false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["false"]}
%sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %sB = tensor.extract_slice %B[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
// Step 3. %sB backprops to the tensor.extract_slice producer which is not // Step 3. %sB backprops to the tensor.extract_slice producer which is not
// considered an interference. This bufferizes inplace. // considered an interference. This bufferizes inplace.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%D = linalg.matmul ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>) %D = linalg.matmul ins(%B, %C: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%sB: tensor<4x4xf32>) inits(%sB: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
// Step 2. %sC forward propagates to an inplace write in %E. // Step 2. %sC forward propagates to an inplace write in %E.
// %sC backward propagates to %C which is inplaceable. // %sC backward propagates to %C which is inplaceable.
// As a consequence this is bufferized inplace. // As a consequence this is bufferized inplace.
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
%sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %sC = tensor.extract_slice %C[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
// Step 1. %sC backprops to the tensor.extract_slice producer which is not // Step 1. %sC backprops to the tensor.extract_slice producer which is not
// considered an interference. This bufferizes inplace. // considered an interference. This bufferizes inplace.
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%E = linalg.matmul ins(%A, %A: tensor<4x4xf32>, tensor<4x4xf32>) %E = linalg.matmul ins(%A, %A: tensor<4x4xf32>, tensor<4x4xf32>)
outs(%sC: tensor<4x4xf32>) inits(%sC: tensor<4x4xf32>)
-> tensor<4x4xf32> -> tensor<4x4xf32>
return %D, %E: tensor<4x4xf32>, tensor<4x4xf32> return %D, %E: tensor<4x4xf32>, tensor<4x4xf32>
} }
// ----- // -----
// CHECK-LABEL: func @nested_extract_slice_and_insert // CHECK-LABEL: func @nested_extract_slice_and_insert
Show All 22 Linesfunc.func @nested_extract_slice_and_insert(
// CHECK-NEXT: fill // CHECK-NEXT: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "false", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "false", "none", "none"]}
%sA = tensor.extract_slice %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %sA = tensor.extract_slice %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%ssA = tensor.extract_slice %sA[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %ssA = tensor.extract_slice %sA[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
%FA = linalg.fill ins(%f0 : f32) outs(%ssA : tensor<4x4xf32>) -> tensor<4x4xf32> %FA = linalg.fill ins(%f0 : f32) inits(%ssA : tensor<4x4xf32>) -> tensor<4x4xf32>
%rsA = tensor.insert_slice %FA into %sA[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<?x?xf32> %rsA = tensor.insert_slice %FA into %sA[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<?x?xf32>
%rA = tensor.insert_slice %rsA into %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %rA = tensor.insert_slice %rsA into %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
// 3-level matching tensor.extract_slice / tensor.insert_slice into // 3-level matching tensor.extract_slice / tensor.insert_slice into
// inplaceable %B. // inplaceable %B.
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none"]}
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
// CHECK-NEXT: fill // CHECK-NEXT: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
%sB = tensor.extract_slice %B[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %sB = tensor.extract_slice %B[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%ssB = tensor.extract_slice %sB[0, 0][4, %idx][1, 1] : tensor<?x?xf32> to tensor<4x?xf32> %ssB = tensor.extract_slice %sB[0, 0][4, %idx][1, 1] : tensor<?x?xf32> to tensor<4x?xf32>
%sssB = tensor.extract_slice %ssB[0, 0][4, 4][1, 1] : tensor<4x?xf32> to tensor<4x4xf32> %sssB = tensor.extract_slice %ssB[0, 0][4, 4][1, 1] : tensor<4x?xf32> to tensor<4x4xf32>
%FB = linalg.fill ins(%f0 : f32) outs(%sssB : tensor<4x4xf32>) -> tensor<4x4xf32> %FB = linalg.fill ins(%f0 : f32) inits(%sssB : tensor<4x4xf32>) -> tensor<4x4xf32>
%rssB = tensor.insert_slice %FB into %ssB[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<4x?xf32> %rssB = tensor.insert_slice %FB into %ssB[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<4x?xf32>
%rsB = tensor.insert_slice %rssB into %sB[0, 0][4, %idx][1, 1] : tensor<4x?xf32> into tensor<?x?xf32> %rsB = tensor.insert_slice %rssB into %sB[0, 0][4, %idx][1, 1] : tensor<4x?xf32> into tensor<?x?xf32>
%rB = tensor.insert_slice %rsB into %B[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %rB = tensor.insert_slice %rsB into %B[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
// 2-level matching tensor.extract_slice / tensor.insert_slice into // 2-level matching tensor.extract_slice / tensor.insert_slice into
// inplaceable %C with a twist. // inplaceable %C with a twist.
// Throw a wrench in the system: %rsC production sizes do not match %ssC. // Throw a wrench in the system: %rsC production sizes do not match %ssC.
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none", "none"]}
// The tensor.insert_slice that would be candidate for matching does not actually // The tensor.insert_slice that would be candidate for matching does not actually
// match. That tensor.insert_slice can still be bufferized inplace nonetheless // match. That tensor.insert_slice can still be bufferized inplace nonetheless
// but this tensor.extract_slice, which bufferizes to an inplace write, cannot. // but this tensor.extract_slice, which bufferizes to an inplace write, cannot.
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["false", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["false", "none"]}
// CHECK-NEXT: fill // CHECK-NEXT: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]}
// CHECK-NEXT: tensor.insert_slice // CHECK-NEXT: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none", "none"]}
%sC = tensor.extract_slice %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %sC = tensor.extract_slice %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%ssC = tensor.extract_slice %sC[0, 0][%sz1, 4][1, 1] : tensor<?x?xf32> to tensor<?x4xf32> %ssC = tensor.extract_slice %sC[0, 0][%sz1, 4][1, 1] : tensor<?x?xf32> to tensor<?x4xf32>
%FC = linalg.fill ins(%f0 : f32) outs(%ssC : tensor<?x4xf32>) -> tensor<?x4xf32> %FC = linalg.fill ins(%f0 : f32) inits(%ssC : tensor<?x4xf32>) -> tensor<?x4xf32>
%rsC = tensor.insert_slice %FC into %sC[0, 0][%sz2, 4][1, 1] : tensor<?x4xf32> into tensor<?x?xf32> %rsC = tensor.insert_slice %FC into %sC[0, 0][%sz2, 4][1, 1] : tensor<?x4xf32> into tensor<?x?xf32>
%rC = tensor.insert_slice %rsC into %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %rC = tensor.insert_slice %rsC into %C[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [-1, 1, 2] // CHECK-SAME: __equivalent_func_args__ = [-1, 1, 2]
return %rA, %rB, %rC: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32> return %rA, %rB, %rC: tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
} }
Show All 10 Linesfunc.func @dependence_through_call(%I : tensor<64xf32> {bufferization.writable = true}) {
%f1 = arith.constant 1.000000e+00 : f32 %f1 = arith.constant 1.000000e+00 : f32
%f2 = arith.constant 2.000000e+00 : f32 %f2 = arith.constant 2.000000e+00 : f32
// 2. %B already bufferizes inplace, %A would alias and have a different // 2. %B already bufferizes inplace, %A would alias and have a different
// value. The calls to `foo` are determined to read conservatively, so %A // value. The calls to `foo` are determined to read conservatively, so %A
// cannot bufferize inplace. // cannot bufferize inplace.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]}
%A = linalg.fill ins(%f1 : f32) outs(%I : tensor<64xf32>) -> tensor<64xf32> %A = linalg.fill ins(%f1 : f32) inits(%I : tensor<64xf32>) -> tensor<64xf32>
// 1. Bufferizes inplace: no alias to %A is yet possible. // 1. Bufferizes inplace: no alias to %A is yet possible.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%B = linalg.fill ins(%f2 : f32) outs(%I : tensor<64xf32>) -> tensor<64xf32> %B = linalg.fill ins(%f2 : f32) inits(%I : tensor<64xf32>) -> tensor<64xf32>
call @foo(%A) : (tensor<64xf32>) -> () call @foo(%A) : (tensor<64xf32>) -> ()
call @foo(%B) : (tensor<64xf32>) -> () call @foo(%B) : (tensor<64xf32>) -> ()
return return
} }
// ----- // -----
Show All 14 Linesfunc.func @read_dependence_through_scf_and_call(
%f1 = arith.constant 1.000000e+00 : f32 %f1 = arith.constant 1.000000e+00 : f32
%f2 = arith.constant 2.000000e+00 : f32 %f2 = arith.constant 2.000000e+00 : f32
// 5. %B bufferizes inplace, %A would alias and have a different value. // 5. %B bufferizes inplace, %A would alias and have a different value.
// The calls to `foo` are determined to read conservatively, so %A cannot // The calls to `foo` are determined to read conservatively, so %A cannot
// bufferize inplace. // bufferize inplace.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]}
%A = linalg.fill ins(%f1 : f32) outs(%I : tensor<64xf32>) -> tensor<64xf32> %A = linalg.fill ins(%f1 : f32) inits(%I : tensor<64xf32>) -> tensor<64xf32>
// 4. Bufferizes inplace: no alias to %A is yet possible. // 4. Bufferizes inplace: no alias to %A is yet possible.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%B = linalg.fill ins(%f2 : f32) outs(%I : tensor<64xf32>) -> tensor<64xf32> %B = linalg.fill ins(%f2 : f32) inits(%I : tensor<64xf32>) -> tensor<64xf32>
// 3. Does not read or write, bufferizes inplace. // 3. Does not read or write, bufferizes inplace.
// CHECK: scf.for // CHECK: scf.for
// CHECK-NEXT: scf.yield // CHECK-NEXT: scf.yield
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
// CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true", "true"]} // CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true", "true"]}
%r:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%0 = %A, %1 = %B) %r:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%0 = %A, %1 = %B)
-> (tensor<64xf32>, tensor<64xf32>) -> (tensor<64xf32>, tensor<64xf32>)
{ {
scf.yield %0, %1 : tensor<64xf32>, tensor<64xf32> scf.yield %0, %1 : tensor<64xf32>, tensor<64xf32>
} }
call @foo(%r#0) : (tensor<64xf32>) -> () call @foo(%r#0) : (tensor<64xf32>) -> ()
call @foo(%r#1) : (tensor<64xf32>) -> () call @foo(%r#1) : (tensor<64xf32>) -> ()
// 2. %B2 already bufferizes inplace, %A2 would alias and have a different // 2. %B2 already bufferizes inplace, %A2 would alias and have a different
// value. The calls to `foo` are determined to read conservatively, so %A2 // value. The calls to `foo` are determined to read conservatively, so %A2
// cannot bufferize inplace. // cannot bufferize inplace.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]}
%A2 = linalg.fill ins(%f1 : f32) outs(%I2 : tensor<64xf32>) -> tensor<64xf32> %A2 = linalg.fill ins(%f1 : f32) inits(%I2 : tensor<64xf32>) -> tensor<64xf32>
// 1. Bufferizes inplace: no alias to %A2 is yet possible. // 1. Bufferizes inplace: no alias to %A2 is yet possible.
// CHECK: fill // CHECK: fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%B2 = linalg.fill ins(%f2 : f32) outs(%I2 : tensor<64xf32>) -> tensor<64xf32> %B2 = linalg.fill ins(%f2 : f32) inits(%I2 : tensor<64xf32>) -> tensor<64xf32>
call @bar(%A2) : (tensor<64xf32>) -> () call @bar(%A2) : (tensor<64xf32>) -> ()
call @bar(%B2) : (tensor<64xf32>) -> () call @bar(%B2) : (tensor<64xf32>) -> ()
return return
} }
// ----- // -----
Show All 28 Linesfunc.func @matmul_on_tensors(
%cst_1 = arith.constant 1.000000e+00 : f32 %cst_1 = arith.constant 1.000000e+00 : f32
%7 = bufferization.alloc_tensor() : tensor<256x256xf32> %7 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]}
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
%8 = linalg.fill ins(%cst_0 : f32) outs(%7 : tensor<256x256xf32>) -> tensor<256x256xf32> %8 = linalg.fill ins(%cst_0 : f32) inits(%7 : tensor<256x256xf32>) -> tensor<256x256xf32>
%11 = linalg.fill ins(%cst_1 : f32) outs(%7 : tensor<256x256xf32>) -> tensor<256x256xf32> %11 = linalg.fill ins(%cst_1 : f32) inits(%7 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%sA = tensor.extract_slice %8[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32> %sA = tensor.extract_slice %8[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
%sB = tensor.extract_slice %11[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32> %sB = tensor.extract_slice %11[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
%r = linalg.matmul %r = linalg.matmul
ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>) ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [2] // CHECK-SAME: __equivalent_func_args__ = [2]
return %r : tensor<256x256xf32> return %r : tensor<256x256xf32>
} }
// ----- // -----
func.func @matmul_on_tensors( func.func @matmul_on_tensors(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32 %cst_0 = arith.constant 0.000000e+00 : f32
%cst_1 = arith.constant 1.000000e+00 : f32 %cst_1 = arith.constant 1.000000e+00 : f32
%7 = bufferization.alloc_tensor() : tensor<256x256xf32> %7 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "false"]}
// CHECK: vector.transfer_write // CHECK: vector.transfer_write
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"]
%8 = linalg.fill ins(%cst_0 : f32) outs(%7 : tensor<256x256xf32>) -> tensor<256x256xf32> %8 = linalg.fill ins(%cst_0 : f32) inits(%7 : tensor<256x256xf32>) -> tensor<256x256xf32>
%9 = vector.transfer_read %arg0[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32> %9 = vector.transfer_read %arg0[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32>
%10 = vector.transfer_write %9, %8[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32> %10 = vector.transfer_write %9, %8[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]}
// CHECK: vector.transfer_write // CHECK: vector.transfer_write
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"]
%11 = linalg.fill ins(%cst_1 : f32) outs(%7 : tensor<256x256xf32>) -> tensor<256x256xf32> %11 = linalg.fill ins(%cst_1 : f32) inits(%7 : tensor<256x256xf32>) -> tensor<256x256xf32>
%12 = vector.transfer_read %arg1[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32> %12 = vector.transfer_read %arg1[%c0, %c0], %cst_0 {in_bounds = [false, true]} : tensor<518x518xf32>, vector<256x256xf32>
%13 = vector.transfer_write %12, %11[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32> %13 = vector.transfer_write %12, %11[%c0, %c0] {in_bounds = [true, true]} : vector<256x256xf32>, tensor<256x256xf32>
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true"]}
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "true"]}
%sA = tensor.extract_slice %10[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32> %sA = tensor.extract_slice %10[0, 0][256, 16][1, 1]: tensor<256x256xf32> to tensor<256x16xf32>
%sB = tensor.extract_slice %13[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32> %sB = tensor.extract_slice %13[0, 0][16, 256][1, 1]: tensor<256x256xf32> to tensor<16x256xf32>
%r = linalg.matmul %r = linalg.matmul
ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>) ins(%sA, %sB : tensor<256x16xf32>, tensor<16x256xf32>)
outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [2] // CHECK-SAME: __equivalent_func_args__ = [2]
return %r : tensor<256x256xf32> return %r : tensor<256x256xf32>
} }
// ----- // -----
Show All 14 Lines
// CHECK-SAME: bufferization.access = "write" // CHECK-SAME: bufferization.access = "write"
-> tensor<62x90xf32> attributes {passthrough = [["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]} -> tensor<62x90xf32> attributes {passthrough = [["target-cpu", "skylake-avx512"], ["prefer-vector-width", "512"]]}
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]
%0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<62x90xf32>) -> tensor<62x90xf32> %0 = linalg.fill ins(%cst : f32) inits(%arg2 : tensor<62x90xf32>) -> tensor<62x90xf32>
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true"] // CHECK-SAME: {__inplace_operands_attr__ = ["true"]
%2 = tensor.extract_slice %0[0, 0] [32, 90] [1, 1] : tensor<62x90xf32> to tensor<32x90xf32> %2 = tensor.extract_slice %0[0, 0] [32, 90] [1, 1] : tensor<62x90xf32> to tensor<32x90xf32>
// CHECK: vector.transfer_write // CHECK: vector.transfer_write
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true", "none", "none"]
%7 = vector.transfer_write %v1, %2[%c0, %c0] {in_bounds = [true, true]} : vector<32x90xf32>, tensor<32x90xf32> %7 = vector.transfer_write %v1, %2[%c0, %c0] {in_bounds = [true, true]} : vector<32x90xf32>, tensor<32x90xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines
// CHECK-SAME: bufferization.access = "read" // CHECK-SAME: bufferization.access = "read"
%t2: tensor<?xf32> {bufferization.writable = true}) %t2: tensor<?xf32> {bufferization.writable = true})
// CHECK-SAME: bufferization.access = "write" // CHECK-SAME: bufferization.access = "write"
-> (tensor<?xf32>, tensor<?xf32>){ -> (tensor<?xf32>, tensor<?xf32>){
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false"]
%o:2 = linalg.generic #trait ins(%t1 : tensor<?xf32>) %o:2 = linalg.generic #trait ins(%t1 : tensor<?xf32>)
outs (%t2, %t2 : tensor<?xf32>, tensor<?xf32>) { inits (%t2, %t2 : tensor<?xf32>, tensor<?xf32>) {
^bb(%0: f32, %1: f32, %2 : f32) : ^bb(%0: f32, %1: f32, %2 : f32) :
linalg.yield %0, %0 : f32, f32 linalg.yield %0, %0 : f32, f32
} -> (tensor<?xf32>, tensor<?xf32>) } -> (tensor<?xf32>, tensor<?xf32>)
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [1, -1] // CHECK-SAME: __equivalent_func_args__ = [1, -1]
return %o#0, %o#1 : tensor<?xf32>, tensor<?xf32> return %o#0, %o#1 : tensor<?xf32>, tensor<?xf32>
} }
Show All 18 Lines// CHECK-SAME: bufferization.access = "read"
%t2: tensor<?xf32> {bufferization.writable = true}) %t2: tensor<?xf32> {bufferization.writable = true})
// CHECK-SAME: bufferization.access = "write" // CHECK-SAME: bufferization.access = "write"
-> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>){ -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>){
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false", "false"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "false", "false"]
%o:3 = linalg.generic #trait %o:3 = linalg.generic #trait
ins(%t1 : tensor<?xf32>) ins(%t1 : tensor<?xf32>)
outs (%t2, %t2, %t2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) { inits (%t2, %t2, %t2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {
^bb(%0: f32, %1: f32, %2 : f32, %3 : f32) : ^bb(%0: f32, %1: f32, %2 : f32, %3 : f32) :
linalg.yield %0, %0, %0 : f32, f32, f32 linalg.yield %0, %0, %0 : f32, f32, f32
} -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) } -> (tensor<?xf32>, tensor<?xf32>, tensor<?xf32>)
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [1, -1, -1] // CHECK-SAME: __equivalent_func_args__ = [1, -1, -1]
return %o#0, %o#1, %o#2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32> return %o#0, %o#1, %o#2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>
} }
▲ Show 20 Lines • Show All 379 Lines • Show Last 20 Lines

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir

Show First 20 Lines • Show All 216 Lines • ▼ Show 20 Linesfunc.func @unknown_op(%A : tensor<4xf32>) -> tensor<4xf32>
return %r: tensor<4xf32> return %r: tensor<4xf32>
} }
// ----- // -----
func.func @mini_test_case1() -> tensor<10x20xf32> { func.func @mini_test_case1() -> tensor<10x20xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%t = bufferization.alloc_tensor() : tensor<10x20xf32> %t = bufferization.alloc_tensor() : tensor<10x20xf32>
%r = linalg.fill ins(%f0 : f32) outs(%t : tensor<10x20xf32>) -> tensor<10x20xf32> %r = linalg.fill ins(%f0 : f32) inits(%t : tensor<10x20xf32>) -> tensor<10x20xf32>
// expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}} // expected-error @+1 {{operand #0 of ReturnLike op does not satisfy destination passing style}}
return %r : tensor<10x20xf32> return %r : tensor<10x20xf32>
} }
// ----- // -----
func.func @main() -> tensor<4xi32> { func.func @main() -> tensor<4xi32> {
%r = scf.execute_region -> tensor<4xi32> { %r = scf.execute_region -> tensor<4xi32> {
▲ Show 20 Lines • Show All 83 Lines • Show Last 20 Lines

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir

Show First 20 Lines • Show All 238 Lines • ▼ Show 20 Lines
// This function does not read, just write. We need an alloc, but no copy. // This function does not read, just write. We need an alloc, but no copy.
// CHECK-LABEL: func @does_not_read( // CHECK-LABEL: func @does_not_read(
// CHECK-NOT: alloc // CHECK-NOT: alloc
// CHECK-NOT: copy // CHECK-NOT: copy
func.func @does_not_read(%t: tensor<?xf32>) -> tensor<?xf32> { func.func @does_not_read(%t: tensor<?xf32>) -> tensor<?xf32> {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%r = linalg.fill ins(%f0 : f32) outs(%t : tensor<?xf32>) -> tensor<?xf32> %r = linalg.fill ins(%f0 : f32) inits(%t : tensor<?xf32>) -> tensor<?xf32>
return %r : tensor<?xf32> return %r : tensor<?xf32>
} }
// CHECK-LABEL: func @main( // CHECK-LABEL: func @main(
// CHECK-SAME: %[[t:.*]]: memref<?xf32 // CHECK-SAME: %[[t:.*]]: memref<?xf32
// CHECK: %[[alloc:.*]] = memref.alloc // CHECK: %[[alloc:.*]] = memref.alloc
// CHECK-NOT: copy // CHECK-NOT: copy
// CHECK: %[[casted:.*]] = memref.cast %[[alloc]] // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
▲ Show 20 Lines • Show All 155 Lines • ▼ Show 20 Lines
// CHECK: func @init_and_dot( // CHECK: func @init_and_dot(
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>> // CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>>
// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>> // CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: memref<64xf32, strided<[?], offset: ?>>
// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<f32, strided<[], offset: ?>> // CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: memref<f32, strided<[], offset: ?>>
func.func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> { func.func @init_and_dot(%a: tensor<64xf32>, %b: tensor<64xf32>, %c: tensor<f32>) -> tensor<f32> {
// CHECK-NEXT: %[[C0:.*]] = arith.constant 0{{.*}} : f32 // CHECK-NEXT: %[[C0:.*]] = arith.constant 0{{.*}} : f32
%v0 = arith.constant 0.0 : f32 %v0 = arith.constant 0.0 : f32
// CHECK-NEXT: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32, strided<[], offset: ?>>) // CHECK-NEXT: linalg.fill ins(%[[C0]] : f32) inits(%[[C]] : memref<f32, strided<[], offset: ?>>)
%d = linalg.fill ins(%v0 : f32) outs(%c : tensor<f32>) -> tensor<f32> %d = linalg.fill ins(%v0 : f32) inits(%c : tensor<f32>) -> tensor<f32>
// CHECK-NEXT: linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, strided<[?], offset: ?>>, memref<64xf32, strided<[?], offset: ?>>) outs(%[[C]] : memref<f32, strided<[], offset: ?>>) // CHECK-NEXT: linalg.dot ins(%[[A]], %[[B]] : memref<64xf32, strided<[?], offset: ?>>, memref<64xf32, strided<[?], offset: ?>>) inits(%[[C]] : memref<f32, strided<[], offset: ?>>)
%e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>) %e = linalg.dot ins(%a, %b : tensor<64xf32>,tensor<64xf32>)
outs(%d: tensor<f32>) -> tensor<f32> inits(%d: tensor<f32>) -> tensor<f32>
// CHECK-NEXT: return // CHECK-NEXT: return
return %e : tensor<f32> return %e : tensor<f32>
} }
// CHECK: func @main() // CHECK: func @main()
func.func @main() { func.func @main() {
// CHECK-DAG: %[[C0:.*]] = arith.constant 0{{.*}} : f32 // CHECK-DAG: %[[C0:.*]] = arith.constant 0{{.*}} : f32
// CHECK-DAG: %[[C1:.*]] = arith.constant 1{{.*}} : f32 // CHECK-DAG: %[[C1:.*]] = arith.constant 1{{.*}} : f32
// CHECK-DAG: %[[C2:.*]] = arith.constant 2{{.*}} : f32 // CHECK-DAG: %[[C2:.*]] = arith.constant 2{{.*}} : f32
%v0 = arith.constant 0.0 : f32 %v0 = arith.constant 0.0 : f32
%v1 = arith.constant 1.0 : f32 %v1 = arith.constant 1.0 : f32
%v2 = arith.constant 2.0 : f32 %v2 = arith.constant 2.0 : f32
// CHECK-NEXT: %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32> // CHECK-NEXT: %[[A:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
// CHECK-NEXT: %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32> // CHECK-NEXT: %[[B:.*]] = memref.alloc() {alignment = 128 : i64} : memref<64xf32>
// CHECK-NEXT: %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32> // CHECK-NEXT: %[[C:.*]] = memref.alloc() {alignment = 128 : i64} : memref<f32>
// CHECK-DAG: %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>> // CHECK-DAG: %[[cA:.*]] = memref.cast %[[A]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>
// CHECK-DAG: %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>> // CHECK-DAG: %[[cB:.*]] = memref.cast %[[B]] : memref<64xf32> to memref<64xf32, strided<[?], offset: ?>>
// CHECK-DAG: %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, strided<[], offset: ?>> // CHECK-DAG: %[[cC:.*]] = memref.cast %[[C]] : memref<f32> to memref<f32, strided<[], offset: ?>>
%A = bufferization.alloc_tensor() : tensor<64xf32> %A = bufferization.alloc_tensor() : tensor<64xf32>
%B = bufferization.alloc_tensor() : tensor<64xf32> %B = bufferization.alloc_tensor() : tensor<64xf32>
%C = bufferization.alloc_tensor() : tensor<f32> %C = bufferization.alloc_tensor() : tensor<f32>
// CHECK-DAG: linalg.fill ins(%[[C1]] : f32) outs(%[[A]] : memref<64xf32>) // CHECK-DAG: linalg.fill ins(%[[C1]] : f32) inits(%[[A]] : memref<64xf32>)
// CHECK-DAG: linalg.fill ins(%[[C2]] : f32) outs(%[[B]] : memref<64xf32>) // CHECK-DAG: linalg.fill ins(%[[C2]] : f32) inits(%[[B]] : memref<64xf32>)
// CHECK-DAG: linalg.fill ins(%[[C0]] : f32) outs(%[[C]] : memref<f32>) // CHECK-DAG: linalg.fill ins(%[[C0]] : f32) inits(%[[C]] : memref<f32>)
%AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32> %AA = linalg.fill ins(%v1 : f32) inits(%A : tensor<64xf32>) -> tensor<64xf32>
%BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32> %BB = linalg.fill ins(%v2 : f32) inits(%B : tensor<64xf32>) -> tensor<64xf32>
%CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32> %CC = linalg.fill ins(%v0 : f32) inits(%C : tensor<f32>) -> tensor<f32>
// CHECK-NEXT: call @init_and_dot(%[[cA]], %[[cB]], %[[cC]]) // CHECK-NEXT: call @init_and_dot(%[[cA]], %[[cB]], %[[cC]])
%res = call @init_and_dot(%AA, %BB, %CC) : %res = call @init_and_dot(%AA, %BB, %CC) :
(tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32> (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
// CHECK-NEXT: %[[dC:.*]] = memref.cast %[[cC]] : memref<f32, {{.*}}> to memref<*xf32> // CHECK-NEXT: %[[dC:.*]] = memref.cast %[[cC]] : memref<f32, {{.*}}> to memref<*xf32>
%res2 = tensor.cast %res: tensor<f32> to tensor<*xf32> %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
▲ Show 20 Lines • Show All 147 Lines • Show Last 20 Lines

mlir/test/Dialect/Bufferization/Transforms/tensor-copy-insertion.mlir

Show First 20 Lines • Show All 50 Lines • ▼ Show 20 Lines
// ----- // -----
// CHECK-LABEL: func @do_not_copy_when_overwritten // CHECK-LABEL: func @do_not_copy_when_overwritten
func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32) func.func @do_not_copy_when_overwritten(%t: tensor<5xf32>, %f: f32)
-> (tensor<5xf32>, tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>)
{ {
// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<5xf32> // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<5xf32>
// CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<5xf32>) // CHECK: linalg.generic {{.*}} inits(%[[alloc]] : tensor<5xf32>)
%r = linalg.generic { %r = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>], indexing_maps = [affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%t : tensor<5xf32>) { inits(%t : tensor<5xf32>) {
^bb0(%arg0 : f32) : ^bb0(%arg0 : f32) :
linalg.yield %f : f32 linalg.yield %f : f32
} -> tensor<5xf32> } -> tensor<5xf32>
return %t, %r : tensor<5xf32>, tensor<5xf32> return %t, %r : tensor<5xf32>, tensor<5xf32>
} }
// ----- // -----
// CHECK-LABEL: func @do_not_copy_when_result_not_read // CHECK-LABEL: func @do_not_copy_when_result_not_read
func.func @do_not_copy_when_result_not_read(%t: tensor<5xf32>, %f: f32) func.func @do_not_copy_when_result_not_read(%t: tensor<5xf32>, %f: f32)
-> (tensor<3xf32>) -> (tensor<3xf32>)
{ {
%0 = tensor.extract_slice %t[0][3][1] : tensor<5xf32> to tensor<3xf32> %0 = tensor.extract_slice %t[0][3][1] : tensor<5xf32> to tensor<3xf32>
// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<3xf32> // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() {bufferization.escape = [false], memory_space = 0 : ui64} : tensor<3xf32>
// CHECK: linalg.generic {{.*}} outs(%[[alloc]] : tensor<3xf32>) // CHECK: linalg.generic {{.*}} inits(%[[alloc]] : tensor<3xf32>)
%r = linalg.generic { %r = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>], indexing_maps = [affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%0 : tensor<3xf32>) { inits(%0 : tensor<3xf32>) {
^bb0(%arg0 : f32) : ^bb0(%arg0 : f32) :
linalg.yield %f : f32 linalg.yield %f : f32
} -> tensor<3xf32> } -> tensor<3xf32>
return %r : tensor<3xf32> return %r : tensor<3xf32>
} }

mlir/test/Dialect/GPU/transform-gpu-failing.mlir

Show First 20 Lines • Show All 145 Lines • ▼ Show 20 Lines
// ----- // -----
func.func @map_nested_foreach_to_threads_not_buffer(%x: tensor<32x32xf32>, %y: tensor<32x32xf32>, %z: tensor<32x32xf32>, %stream : !gpu.async.token) { func.func @map_nested_foreach_to_threads_not_buffer(%x: tensor<32x32xf32>, %y: tensor<32x32xf32>, %z: tensor<32x32xf32>, %stream : !gpu.async.token) {
%one = arith.constant 1 : index %one = arith.constant 1 : index
%name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one) %name = gpu.launch async[%stream] blocks(%arg3, %arg4, %arg5) in (%arg9 = %one, %arg10 = %one, %arg11 = %one)
threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one) threads(%arg6, %arg7, %arg8) in (%arg12 = %one, %arg13 = %one, %arg14 = %one)
{ {
%t = linalg.matmul ins(%x, %y: tensor<32x32xf32>, tensor<32x32xf32>) outs(%z : tensor<32x32xf32>) -> tensor<32x32xf32> %t = linalg.matmul ins(%x, %y: tensor<32x32xf32>, tensor<32x32xf32>) inits(%z : tensor<32x32xf32>) -> tensor<32x32xf32>
gpu.terminator gpu.terminator
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg0: !pdl.operation): ^bb1(%arg0: !pdl.operation):
%matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0 %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0
▲ Show 20 Lines • Show All 114 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/affine.mlir

// RUN: mlir-opt %s -convert-linalg-to-affine-loops | FileCheck %s // RUN: mlir-opt %s -convert-linalg-to-affine-loops | FileCheck %s
// Test that we can lower all the way to LLVM without crashing, don't check results here. // Test that we can lower all the way to LLVM without crashing, don't check results here.
// RUN: mlir-opt %s -convert-linalg-to-affine-loops -convert-linalg-to-llvm -o=/dev/null 2>&1 // RUN: mlir-opt %s -convert-linalg-to-affine-loops -convert-linalg-to-llvm -o=/dev/null 2>&1
func.func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) { func.func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%A = memref.view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32> %A = memref.view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32>
%B = memref.view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32> %B = memref.view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32>
%C = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32> %C = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32>
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C: memref<?x?xf32>) inits(%C: memref<?x?xf32>)
return return
} }
//----------------------------------------------------------------------------// //----------------------------------------------------------------------------//
// Named ops to loops. // Named ops to loops.
//----------------------------------------------------------------------------// //----------------------------------------------------------------------------//
func.func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) { func.func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) {
linalg.batch_matmul ins(%A, %B: memref<?x?x?xf32>, memref<?x?x?xf32>) linalg.batch_matmul ins(%A, %B: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs(%C : memref<?x?x?xf32>) inits(%C : memref<?x?x?xf32>)
return return
} }
// CHECK-LABEL: @named_batch_matmul // CHECK-LABEL: @named_batch_matmul
// CHECK-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK: %[[B:.*]] = memref.dim %[[mA]], %c0 : memref<?x?x?xf32> // CHECK: %[[B:.*]] = memref.dim %[[mA]], %c0 : memref<?x?x?xf32>
// CHECK: %[[M:.*]] = memref.dim %[[mA]], %c1 : memref<?x?x?xf32> // CHECK: %[[M:.*]] = memref.dim %[[mA]], %c1 : memref<?x?x?xf32>
Show All 12 Lines

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

//RUN: mlir-opt -test-linalg-transform-patterns=test-bubble-up-extract-slice-op-pattern -split-input-file %s | FileCheck %s //RUN: mlir-opt -test-linalg-transform-patterns=test-bubble-up-extract-slice-op-pattern -split-input-file %s | FileCheck %s
func.func @dynamic(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>, %arg2: index, %arg3: index, %arg4: index, %arg5:index) -> tensor<?x?xf32> { func.func @dynamic(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>, %arg2: index, %arg3: index, %arg4: index, %arg5:index) -> tensor<?x?xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>) } ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
outs(%arg0 : tensor<?x?xf32>) { inits(%arg0 : 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>
%1 = tensor.extract_slice %0 [%arg2, %arg3] [%arg4, %arg5] [1, 1] %1 = tensor.extract_slice %0 [%arg2, %arg3] [%arg4, %arg5] [1, 1]
: tensor<?x?xf32> to tensor<?x?xf32> : tensor<?x?xf32> to tensor<?x?xf32>
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
// CHECK: func @dynamic // CHECK: func @dynamic
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg3] [%arg5] [1] : tensor<?xf32> to tensor<?xf32> // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg3] [%arg5] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[%arg2, %arg3] [%arg4, %arg5] [1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x?xf32>, tensor<?xf32>) outs(%[[SLICE2]] : tensor<?x?xf32>) // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x?xf32>, tensor<?xf32>) inits(%[[SLICE2]] : tensor<?x?xf32>)
// CHECK: return %[[GENERIC]] : tensor<?x?xf32> // CHECK: return %[[GENERIC]] : tensor<?x?xf32>
//----- //-----
func.func @static(%arg0: tensor<16x8xf32>, %arg1: tensor<8xf32>) -> tensor<4x2xf32> { func.func @static(%arg0: tensor<16x8xf32>, %arg1: tensor<8xf32>) -> tensor<4x2xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<16x8xf32>, tensor<8xf32>) } ins(%arg0, %arg1 : tensor<16x8xf32>, tensor<8xf32>)
outs(%arg0 : tensor<16x8xf32>) { inits(%arg0 : tensor<16x8xf32>) {
^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<16x8xf32> } -> tensor<16x8xf32>
%1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1] %1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1]
: tensor<16x8xf32> to tensor<4x2xf32> : tensor<16x8xf32> to tensor<4x2xf32>
return %1 : tensor<4x2xf32> return %1 : tensor<4x2xf32>
} }
// CHECK: func @static // CHECK: func @static
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32> // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32>
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<8xf32> to tensor<2xf32> // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<8xf32> to tensor<2xf32>
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32> // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<16x8xf32> to tensor<4x2xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<4x2xf32>) // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) inits(%[[SLICE2]] : tensor<4x2xf32>)
// CHECK: return %[[GENERIC]] : tensor<4x2xf32> // CHECK: return %[[GENERIC]] : tensor<4x2xf32>
//----- //-----
func.func @mixed(%arg0: tensor<?x8xf32>, %arg1: tensor<8xf32>, %arg2: index, %arg3: index) -> tensor<?x2xf32> { func.func @mixed(%arg0: tensor<?x8xf32>, %arg1: tensor<8xf32>, %arg2: index, %arg3: index) -> tensor<?x2xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<?x8xf32>, tensor<8xf32>) } ins(%arg0, %arg1 : tensor<?x8xf32>, tensor<8xf32>)
outs(%arg0 : tensor<?x8xf32>) { inits(%arg0 : tensor<?x8xf32>) {
^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<?x8xf32> } -> tensor<?x8xf32>
%1 = tensor.extract_slice %0 [8, %arg2] [%arg3, 2] [1, 1] %1 = tensor.extract_slice %0 [8, %arg2] [%arg3, 2] [1, 1]
: tensor<?x8xf32> to tensor<?x2xf32> : tensor<?x8xf32> to tensor<?x2xf32>
return %1 : tensor<?x2xf32> return %1 : tensor<?x2xf32>
} }
// CHECK: func @mixed // CHECK: func @mixed
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32> // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32>
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg2] [2] [1] : tensor<8xf32> to tensor<2xf32> // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[%arg2] [2] [1] : tensor<8xf32> to tensor<2xf32>
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32> // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, %arg2] [%arg3, 2] [1, 1] : tensor<?x8xf32> to tensor<?x2xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<?x2xf32>) // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<?x2xf32>, tensor<2xf32>) inits(%[[SLICE2]] : tensor<?x2xf32>)
// CHECK: return %[[GENERIC]] : tensor<?x2xf32> // CHECK: return %[[GENERIC]] : tensor<?x2xf32>
//----- //-----
func.func @dynamic_to_static(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>) -> tensor<4x2xf32> { func.func @dynamic_to_static(%arg0: tensor<?x?xf32>, %arg1: tensor<?xf32>) -> tensor<4x2xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>) } ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
outs(%arg0 : tensor<?x?xf32>) { inits(%arg0 : 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>
%1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1] %1 = tensor.extract_slice %0 [8, 4] [4, 2] [1, 1]
: tensor<?x?xf32> to tensor<4x2xf32> : tensor<?x?xf32> to tensor<4x2xf32>
return %1 : tensor<4x2xf32> return %1 : tensor<4x2xf32>
} }
// CHECK: func @dynamic_to_static // CHECK: func @dynamic_to_static
// CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32> // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32>
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<?xf32> to tensor<2xf32> // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %arg1[4] [2] [1] : tensor<?xf32> to tensor<2xf32>
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32> // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %arg0[8, 4] [4, 2] [1, 1] : tensor<?x?xf32> to tensor<4x2xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]} // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map1, #map], iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) outs(%[[SLICE2]] : tensor<4x2xf32>) // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]] : tensor<4x2xf32>, tensor<2xf32>) inits(%[[SLICE2]] : tensor<4x2xf32>)
// CHECK: return %[[GENERIC]] : tensor<4x2xf32> // CHECK: return %[[GENERIC]] : tensor<4x2xf32>
//----- //-----
func.func @matmul_slice() -> tensor<2x2xf32> { func.func @matmul_slice() -> tensor<2x2xf32> {
%lhs = arith.constant dense<1.0> : tensor<4x4xf32> %lhs = arith.constant dense<1.0> : tensor<4x4xf32>
%rhs = arith.constant dense<1.0> : tensor<4x4xf32> %rhs = arith.constant dense<1.0> : tensor<4x4xf32>
%dst = arith.constant dense<[[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0]]> : tensor<4x4xf32> %dst = arith.constant dense<[[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0], [8.0, 9.0, 10.0, 11.0], [12.0, 13.0, 14.0, 15.0]]> : tensor<4x4xf32>
%0 = linalg.matmul ins(%lhs, %rhs : tensor<4x4xf32>, tensor<4x4xf32>) outs(%dst : tensor<4x4xf32>) -> tensor<4x4xf32> %0 = linalg.matmul ins(%lhs, %rhs : tensor<4x4xf32>, tensor<4x4xf32>) inits(%dst : tensor<4x4xf32>) -> tensor<4x4xf32>
%1 = tensor.extract_slice %0[1,1][2,2][1,1] : tensor<4x4xf32> to tensor<2x2xf32> %1 = tensor.extract_slice %0[1,1][2,2][1,1] : tensor<4x4xf32> to tensor<2x2xf32>
return %1 : tensor<2x2xf32> return %1 : tensor<2x2xf32>
} }
// CHECK: func @matmul_slice // CHECK: func @matmul_slice
// CHECK: %[[SLICE0:.+]] = arith.constant dense<1.000000e+00> : tensor<2x4xf32> // CHECK: %[[SLICE0:.+]] = arith.constant dense<1.000000e+00> : tensor<2x4xf32>
// CHECK: %[[SLICE1:.+]] = arith.constant dense<1.000000e+00> : tensor<4x2xf32> // CHECK: %[[SLICE1:.+]] = arith.constant dense<1.000000e+00> : tensor<4x2xf32>
// CHECK: %[[SLICE3:.+]] = tensor.extract_slice %[[CST:.+]][1, 1] [2, 2] [1, 1] : tensor<4x4xf32> to tensor<2x2xf32> // CHECK: %[[SLICE3:.+]] = tensor.extract_slice %[[CST:.+]][1, 1] [2, 2] [1, 1] : tensor<4x4xf32> to tensor<2x2xf32>
// CHECK: %[[MATMUL:.+]] = linalg.matmul ins(%[[SLICE0]], %[[SLICE1]] : tensor<2x4xf32>, tensor<4x2xf32>) outs(%[[SLICE3]] : tensor<2x2xf32>) -> tensor<2x2xf32> // CHECK: %[[MATMUL:.+]] = linalg.matmul ins(%[[SLICE0]], %[[SLICE1]] : tensor<2x4xf32>, tensor<4x2xf32>) inits(%[[SLICE3]] : tensor<2x2xf32>) -> tensor<2x2xf32>
// CHECK: return %[[MATMUL]] : tensor<2x2xf32> // CHECK: return %[[MATMUL]] : tensor<2x2xf32>
//----- //-----
func.func @conv_slice(%input: tensor<1x225x225x3xf32>, %filter: tensor<3x3x3x32xf32>) -> tensor<1x32x32x16xf32> { func.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 = tensor.empty() : 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) inits(%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> inits(%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:.+]] = tensor.empty() : 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) inits(%[[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>) inits(%[[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 = tensor.empty(%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) inits(%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:.+]] = tensor.empty // 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 10 Lines
// CHECK: #map = affine_map<(d0) -> (d0)> // CHECK: #map = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @basic( // CHECK-LABEL: func @basic(
// CHECK-SAME: %[[TENSOR:.*]]: tensor<4xf32>) -> tensor<4xf32> { // CHECK-SAME: %[[TENSOR:.*]]: tensor<4xf32>) -> tensor<4xf32> {
// CHECK-DAG: %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<4xf32> // CHECK-DAG: %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<4xf32>
// CHECK-DAG: %[[RESULT_MEMREF:.*]] = memref.alloc() {{.*}} : memref<4xf32> // CHECK-DAG: %[[RESULT_MEMREF:.*]] = memref.alloc() {{.*}} : memref<4xf32>
// CHECK: linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} // CHECK: linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]}
// CHECK-SAME: ins(%[[MEMREF]] : memref<4xf32>) // CHECK-SAME: ins(%[[MEMREF]] : memref<4xf32>)
// CHECK-SAME: outs(%[[RESULT_MEMREF]] : memref<4xf32>) { // CHECK-SAME: inits(%[[RESULT_MEMREF]] : memref<4xf32>) {
// CHECK: ^bb0(%[[RESULT1:.*]]: f32, %[[UNUSED:.*]]: f32): // CHECK: ^bb0(%[[RESULT1:.*]]: f32, %[[UNUSED:.*]]: f32):
// CHECK: %[[DIM1:.*]] = math.exp %[[RESULT1]] : f32 // CHECK: %[[DIM1:.*]] = math.exp %[[RESULT1]] : f32
// CHECK: linalg.yield %[[DIM1]] : f32 // CHECK: linalg.yield %[[DIM1]] : f32
// CHECK: } // CHECK: }
// CHECK: %[[RESULT:.*]] = bufferization.to_tensor %[[RESULT_MEMREF]] : memref<4xf32> // CHECK: %[[RESULT:.*]] = bufferization.to_tensor %[[RESULT_MEMREF]] : memref<4xf32>
// CHECK: return %[[RESULT]] : tensor<4xf32> // CHECK: return %[[RESULT]] : tensor<4xf32>
func.func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> { func.func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map0, #map0], indexing_maps = [#map0, #map0],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%arg0 : tensor<4xf32>) } ins(%arg0 : tensor<4xf32>)
outs(%arg0 : tensor<4xf32>) { inits(%arg0 : tensor<4xf32>) {
^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
} -> 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 tensor.empty op. // Same as above but with tensor.empty op.
// CHECK: #map = affine_map<(d0) -> (d0)> // CHECK: #map = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @empty_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: inits(%[[OUT_BUF]] : memref<?xf32>) {
func.func @empty_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> { func.func @empty_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> {
%init = tensor.empty(%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>) { inits(%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
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
#map0 = affine_map<(d0) -> (d0)> #map0 = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @multiple_results // CHECK-LABEL: func @multiple_results
// CHECK: %[[RESULT0:.*]] = memref.alloc() {{.*}} : memref<4xf32> // CHECK: %[[RESULT0:.*]] = memref.alloc() {{.*}} : memref<4xf32>
// CHECK: %[[RESULT1:.*]] = memref.alloc() {{.*}} : memref<4xf32> // CHECK: %[[RESULT1:.*]] = memref.alloc() {{.*}} : memref<4xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%{{.*}} : memref<4xf32>) // CHECK-SAME: ins(%{{.*}} : memref<4xf32>)
// CHECK-SAME: outs(%[[RESULT0]], %[[RESULT1]] : memref<4xf32>, memref<4xf32>) // CHECK-SAME: inits(%[[RESULT0]], %[[RESULT1]] : memref<4xf32>, memref<4xf32>)
// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): // CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32):
func.func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { func.func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
%0, %1 = linalg.generic { %0, %1 = linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%arg0 : tensor<4xf32>) } ins(%arg0 : tensor<4xf32>)
outs (%arg0, %arg0 : tensor<4xf32>, tensor<4xf32>) { inits (%arg0, %arg0 : tensor<4xf32>, tensor<4xf32>) {
^bb0(%gen_arg1: f32, %out1: f32, %out2: f32): ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32):
%tmp1 = math.exp %gen_arg1 : f32 %tmp1 = math.exp %gen_arg1 : f32
linalg.yield %tmp1, %tmp1 : f32, f32 linalg.yield %tmp1, %tmp1 : f32, f32
} -> (tensor<4xf32>, tensor<4xf32>) } -> (tensor<4xf32>, tensor<4xf32>)
return %0, %1 : tensor<4xf32>, tensor<4xf32> return %0, %1 : tensor<4xf32>, tensor<4xf32>
} }
// ----- // -----
Show All 9 Lines
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[DIM0:.*]] = tensor.dim %[[ARG]], %[[C0]] : tensor<?x?xf32> // CHECK-DAG: %[[DIM0:.*]] = tensor.dim %[[ARG]], %[[C0]] : tensor<?x?xf32>
// CHECK-DAG: %[[DIM1:.*]] = tensor.dim %[[ARG]], %[[C1]] : tensor<?x?xf32> // CHECK-DAG: %[[DIM1:.*]] = tensor.dim %[[ARG]], %[[C1]] : tensor<?x?xf32>
// CHECK-DAG: %[[RESULT0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {{.*}} : memref<?x?xf32> // CHECK-DAG: %[[RESULT0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {{.*}} : memref<?x?xf32>
// CHECK-DAG: %[[RESULT1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {{.*}} : memref<?x?xf32> // CHECK-DAG: %[[RESULT1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {{.*}} : memref<?x?xf32>
// CHECK-DAG: %[[MEMREF_ARG:.*]] = bufferization.to_memref %[[ARG]] : memref<?x?xf32> // CHECK-DAG: %[[MEMREF_ARG:.*]] = bufferization.to_memref %[[ARG]] : memref<?x?xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[MEMREF_ARG]] : memref<?x?xf32>) // CHECK-SAME: ins(%[[MEMREF_ARG]] : memref<?x?xf32>)
// CHECK-SAME: outs(%[[RESULT0]], %[[RESULT1]] : memref<?x?xf32>, memref<?x?xf32>) // CHECK-SAME: inits(%[[RESULT0]], %[[RESULT1]] : memref<?x?xf32>, memref<?x?xf32>)
func.func @dynamic_results(%arg0: tensor<?x?xf32>) func.func @dynamic_results(%arg0: tensor<?x?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>) { -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%0, %1 = linalg.generic { %0, %1 = linalg.generic {
indexing_maps = [#map_2d, #map_2d, #map_2d], indexing_maps = [#map_2d, #map_2d, #map_2d],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%arg0 : tensor<?x?xf32>) } ins(%arg0 : tensor<?x?xf32>)
outs (%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) { inits (%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%gen_arg1: f32, %out1: f32, %out2: f32): ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32):
%tmp1 = math.exp %gen_arg1 : f32 %tmp1 = math.exp %gen_arg1 : f32
linalg.yield %tmp1, %tmp1 : f32, f32 linalg.yield %tmp1, %tmp1 : f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>)
return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32> return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32>
} }
// ----- // -----
Show All 14 Lines
// CHECK-SAME: %[[ARG0_TENSOR:.*]]: tensor<2x3x4xvector<3x4xi4>>, // CHECK-SAME: %[[ARG0_TENSOR:.*]]: tensor<2x3x4xvector<3x4xi4>>,
// CHECK-SAME: %[[ARG1_TENSOR:.*]]: tensor<3x2xf32>) -> tensor<3x2xf32> { // CHECK-SAME: %[[ARG1_TENSOR:.*]]: tensor<3x2xf32>) -> tensor<3x2xf32> {
// CHECK-DAG: %[[INIT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<3x2xf32> // CHECK-DAG: %[[INIT_BUFFER:.*]] = memref.alloc() {{.*}} : memref<3x2xf32>
// CHECK-DAG: %[[ARG0_MEMREF:.*]] = bufferization.to_memref %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>> // CHECK-DAG: %[[ARG0_MEMREF:.*]] = bufferization.to_memref %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>>
// CHECK-DAG: %[[ARG1_MEMREF:.*]] = bufferization.to_memref %[[ARG1_TENSOR]] : memref<3x2xf32> // CHECK-DAG: %[[ARG1_MEMREF:.*]] = bufferization.to_memref %[[ARG1_TENSOR]] : memref<3x2xf32>
// CHECK: memref.copy %[[ARG1_MEMREF]], %[[INIT_BUFFER]] : memref<3x2xf32> to memref<3x2xf32> // CHECK: memref.copy %[[ARG1_MEMREF]], %[[INIT_BUFFER]] : memref<3x2xf32> to memref<3x2xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0_MEMREF]] : memref<2x3x4xvector<3x4xi4>>) // CHECK-SAME: ins(%[[ARG0_MEMREF]] : memref<2x3x4xvector<3x4xi4>>)
// CHECK-SAME: outs(%[[INIT_BUFFER]] : memref<3x2xf32>) { // CHECK-SAME: inits(%[[INIT_BUFFER]] : memref<3x2xf32>) {
func.func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>, func.func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>,
%arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) { %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>)
outs(%arg1 : tensor<3x2xf32>) { inits(%arg1 : tensor<3x2xf32>) {
^bb(%v0: vector<3x4xi4>, %v1: f32) : ^bb(%v0: vector<3x4xi4>, %v1: f32) :
linalg.yield %v1 : f32 linalg.yield %v1 : f32
} -> tensor<3x2xf32> } -> tensor<3x2xf32>
return %0 : tensor<3x2xf32> return %0 : tensor<3x2xf32>
} }
// ----- // -----
// CHECK-LABEL: func @bufferize_fill( // CHECK-LABEL: func @bufferize_fill(
// CHECK-SAME: %[[IN:.*]]: tensor<?xf32> // CHECK-SAME: %[[IN:.*]]: tensor<?xf32>
func.func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> { func.func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> {
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// CHECK: linalg.fill ins(%cst : f32) outs(%[[ALLOC]] : memref<?xf32>) // CHECK: linalg.fill ins(%cst : f32) inits(%[[ALLOC]] : memref<?xf32>)
// CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<?xf32> // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<?xf32>
// CHECK: return %[[TENSOR]] // CHECK: return %[[TENSOR]]
%0 = linalg.fill ins(%c0 : f32) outs(%arg0 : tensor<?xf32>) -> tensor<?xf32> %0 = linalg.fill ins(%c0 : f32) inits(%arg0 : tensor<?xf32>) -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @bufferize_dot // CHECK-LABEL: func @bufferize_dot
func.func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> { func.func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> {
%dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>) %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>)
outs(%out : tensor<f32>) -> tensor<f32> inits(%out : tensor<f32>) -> tensor<f32>
return %dot : tensor<f32> return %dot : tensor<f32>
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// TODO: The copy is not necessary. // TODO: The copy is not necessary.
// CHECK: memref.copy {{.*}}, %[[ALLOC]] // CHECK: memref.copy {{.*}}, %[[ALLOC]]
// CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>) // CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>)
// CHECK-SAME: outs(%[[ALLOC:.*]] : memref<f32>) // CHECK-SAME: inits(%[[ALLOC:.*]] : memref<f32>)
// CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<f32> // CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<f32>
// CHECK: return %[[OUT_TENSOR]] // CHECK: return %[[OUT_TENSOR]]
} }
// ----- // -----
// This is a regression test. The linalg-bufferize pass should ignore all func // This is a regression test. The linalg-bufferize pass should ignore all func
// dialect ops. // dialect ops.
// CHECK-LABEL: func private @csum(tensor<6xi64>) -> tensor<6xi64> // CHECK-LABEL: func private @csum(tensor<6xi64>) -> tensor<6xi64>
func.func private @csum(%arg0: tensor<6xi64>) -> tensor<6xi64> func.func private @csum(%arg0: tensor<6xi64>) -> tensor<6xi64>
// CHECK: func public @main(%[[arg0:.*]]: tensor<2x3xi1>) // CHECK: func public @main(%[[arg0:.*]]: tensor<2x3xi1>)
// CHECK: %[[collapse:.*]] = tensor.collapse_shape %[[arg0]] // CHECK: %[[collapse:.*]] = tensor.collapse_shape %[[arg0]]
// 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>) inits(%[[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 = tensor.empty() : 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>) inits(%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

// RUN: mlir-opt %s -split-input-file -canonicalize | FileCheck %s // RUN: mlir-opt %s -split-input-file -canonicalize | FileCheck %s
// Test case: Most basic case. Adding a vector to itself. // Test case: Most basic case. Adding a vector to itself.
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @basic // CHECK-LABEL: @basic
func.func @basic(%arg0: tensor<?xf32>) -> tensor<?xf32> { func.func @basic(%arg0: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: linalg.generic{{.*}}[#[[$MAP]], #[[$MAP]]] // CHECK: linalg.generic{{.*}}[#[[$MAP]], #[[$MAP]]]
// CHECK: attrs = {someattr} // CHECK: attrs = {someattr}
// CHECK: ^bb0(%[[BBARG:.*]]: f32, %{{.*}}: f32): // CHECK: ^bb0(%[[BBARG:.*]]: f32, %{{.*}}: f32):
// CHECK: arith.addf %[[BBARG]], %[[BBARG]] // CHECK: arith.addf %[[BBARG]], %[[BBARG]]
%0 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} %0 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]}
ins(%arg0, %arg0 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg0 : tensor<?xf32>, tensor<?xf32>)
outs(%arg0 : tensor<?xf32>) attrs = {someattr} { inits(%arg0 : tensor<?xf32>) attrs = {someattr} {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%1 = arith.addf %arg1, %arg2 : f32 %1 = arith.addf %arg1, %arg2 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Test case: Different indexing maps mean that args are not redundant, despite // Test case: Different indexing maps mean that args are not redundant, despite
// being the same Value. // being the same Value.
#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)>
// 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-LABEL: @distinct_affine_maps // CHECK-LABEL: @distinct_affine_maps
func.func @distinct_affine_maps(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @distinct_affine_maps(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: linalg.generic{{.*}}[#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]] // CHECK: linalg.generic{{.*}}[#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]]
%0 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %0 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg0 : tensor<?x?xf32>) { inits(%arg0 : tensor<?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):
%1 = arith.addf %arg1, %arg2 : f32 %1 = arith.addf %arg1, %arg2 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
// Test case: Check rewriting mechanics for mixed redundant and // Test case: Check rewriting mechanics for mixed redundant and
// non-redundant args. // non-redundant args.
#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)>
// 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-LABEL: @mixed_redundant_non_redundant // CHECK-LABEL: @mixed_redundant_non_redundant
func.func @mixed_redundant_non_redundant(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @mixed_redundant_non_redundant(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: linalg.generic{{.*}}[#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]] // CHECK: linalg.generic{{.*}}[#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]]
// CHECK: ^bb0(%[[BBARG0:.*]]: f32, %[[BBARG1:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32): // CHECK: ^bb0(%[[BBARG0:.*]]: f32, %[[BBARG1:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32):
// CHECK: "test.elementwise_mappable"(%[[BBARG0]], %[[BBARG1]], %[[BBARG0]]) // CHECK: "test.elementwise_mappable"(%[[BBARG0]], %[[BBARG1]], %[[BBARG0]])
%0 = linalg.generic {indexing_maps = [#map0, #map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %0 = linalg.generic {indexing_maps = [#map0, #map1, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg0 : tensor<?x?xf32>) { inits(%arg0 : tensor<?x?xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):
%1 = "test.elementwise_mappable"(%arg1, %arg2, %arg3) : (f32, f32, f32) -> f32 %1 = "test.elementwise_mappable"(%arg1, %arg2, %arg3) : (f32, f32, f32) -> f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
// Test case: Check rewriting mechanics for multiple different redundant args. // Test case: Check rewriting mechanics for multiple different redundant args.
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
// CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: @multiple_different_redundant_args // CHECK-LABEL: @multiple_different_redundant_args
func.func @multiple_different_redundant_args(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> { func.func @multiple_different_redundant_args(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: linalg.generic{{.*}}[#[[$MAP]], #[[$MAP]], #[[$MAP]]] // CHECK: linalg.generic{{.*}}[#[[$MAP]], #[[$MAP]], #[[$MAP]]]
// CHECK: ^bb0(%[[BBARG0:.*]]: f32, %[[BBARG1:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32): // CHECK: ^bb0(%[[BBARG0:.*]]: f32, %[[BBARG1:.*]]: f32, %{{[a-zA-Z0-9]+}}: f32):
// CHECK: "test.elementwise_mappable"(%[[BBARG0]], %[[BBARG1]], %[[BBARG0]], %[[BBARG1]]) // CHECK: "test.elementwise_mappable"(%[[BBARG0]], %[[BBARG1]], %[[BBARG0]], %[[BBARG1]])
%0 = linalg.generic {indexing_maps = [#map, #map, #map, #map, #map], iterator_types = ["parallel"]} %0 = linalg.generic {indexing_maps = [#map, #map, #map, #map, #map], iterator_types = ["parallel"]}
ins(%arg0, %arg1, %arg0, %arg1 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1, %arg0, %arg1 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>, tensor<?xf32>)
outs(%arg0 : tensor<?xf32>) { inits(%arg0 : tensor<?xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32, %arg5: f32, %arg6: f32): ^bb0(%arg2: f32, %arg3: f32, %arg4: f32, %arg5: f32, %arg6: f32):
%1 = "test.elementwise_mappable"(%arg2, %arg3, %arg4, %arg5) : (f32, f32, f32, f32) -> f32 %1 = "test.elementwise_mappable"(%arg2, %arg3, %arg4, %arg5) : (f32, f32, f32, f32) -> f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Drop dead result. // Drop dead result.
#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) -> (d1, d2, d0)> #map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
#map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
func.func @drop_dead_results(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) { func.func @drop_dead_results(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
%0:4 = linalg.generic { %0:4 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3, #map4], indexing_maps = [#map0, #map1, #map2, #map3, #map4],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<?x?x?xf32>) ins(%arg0 : tensor<?x?x?xf32>)
outs(%arg0, %arg0, %arg0, %arg0 inits(%arg0, %arg0, %arg0, %arg0
: tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) { : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
%1 = arith.addf %b0, %b0: f32 %1 = arith.addf %b0, %b0: f32
linalg.yield %1, %1, %1, %1 : f32, f32, f32, f32 linalg.yield %1, %1, %1, %1 : f32, f32, f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %0#0, %0#2 : tensor<?x?x?xf32>, tensor<?x?x?xf32> return %0#0, %0#2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
} }
// 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, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
// CHECK: func @drop_dead_results( // CHECK: func @drop_dead_results(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>)
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: outs(%[[ARG0]], %[[ARG0]] : // CHECK-SAME: inits(%[[ARG0]], %[[ARG0]] :
// CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1
// ----- // -----
// 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 = tensor.empty() : tensor<f32> %init0 = tensor.empty() : tensor<f32>
%init1 = tensor.empty() : 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>) { inits(%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:.+]] = tensor.empty() : tensor<f32> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<f32>
// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : 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: inits(%[[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 = tensor.empty() : 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) inits(%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>) inits(%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 = tensor.empty(%arg0) : tensor<?xi32> %init1 = tensor.empty(%arg0) : tensor<?xi32>
%init = tensor.empty() : 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) inits(%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>) { inits(%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:.+]] = tensor.empty(%[[ARG0]]) // CHECK: %[[INIT:.+]] = tensor.empty(%[[ARG0]])
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: outs(%[[INIT]] // CHECK-SAME: inits(%[[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)>
#map3 = affine_map<(d0, d1) -> (d1)> #map3 = affine_map<(d0, d1) -> (d1)>
func.func @multiple_redundant_args(%arg0 : tensor<?x?xi32>, %arg1 : tensor<?xi32>, func.func @multiple_redundant_args(%arg0 : tensor<?x?xi32>, %arg1 : tensor<?xi32>,
%arg2 : tensor<?xi32>, %arg3 : tensor<?x?xi32>, %arg4 : tensor<?xi32>) -> tensor<?xi32> { %arg2 : tensor<?xi32>, %arg3 : tensor<?x?xi32>, %arg4 : tensor<?xi32>) -> tensor<?xi32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map3, #map0, #map0, #map2, #map1, #map1, #map2], indexing_maps = [#map3, #map0, #map0, #map2, #map1, #map1, #map2],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg4, %arg0, %arg0, %arg1, %arg3, %arg3 ins(%arg4, %arg0, %arg0, %arg1, %arg3, %arg3
: tensor<?xi32>, tensor<?x?xi32>, tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xi32>, tensor<?x?xi32>) : tensor<?xi32>, tensor<?x?xi32>, tensor<?x?xi32>, tensor<?xi32>, tensor<?x?xi32>, tensor<?x?xi32>)
outs(%arg2 : tensor<?xi32>) { inits(%arg2 : tensor<?xi32>) {
^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32, %b4 : i32, %b5 : i32, %b6 : i32): ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32, %b4 : i32, %b5 : i32, %b6 : i32):
%1 = arith.addi %b0, %b1 : i32 %1 = arith.addi %b0, %b1 : i32
%2 = arith.addi %1, %b2 : i32 %2 = arith.addi %1, %b2 : i32
%3 = arith.addi %2, %b3 : i32 %3 = arith.addi %2, %b3 : i32
%4 = arith.addi %3, %b4 : i32 %4 = arith.addi %3, %b4 : i32
%5 = arith.addi %4, %b5 : i32 %5 = arith.addi %4, %b5 : i32
%6 = arith.addi %5, %b6 : i32 %6 = arith.addi %5, %b6 : i32
linalg.yield %6 : i32 linalg.yield %6 : i32
Show All 9 Lines
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?xi32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?xi32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?xi32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?xi32>
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xi32> // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xi32>
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: tensor<?xi32>) // CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: tensor<?xi32>)
// CHECK: %[[RETURN:.+]] = linalg.generic // CHECK: %[[RETURN:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "reduction"]
// CHECK-SAME: ins(%[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG3]] : // CHECK-SAME: ins(%[[ARG4]], %[[ARG0]], %[[ARG1]], %[[ARG3]] :
// CHECK-SAME: outs(%[[ARG2]] : // CHECK-SAME: inits(%[[ARG2]] :
// CHECK: ^{{.+}}(%[[B0:[a-zA-Z0-9]+]]: i32 // CHECK: ^{{.+}}(%[[B0:[a-zA-Z0-9]+]]: i32
// CHECK-SAME: %[[B1:[a-zA-Z0-9_]+]]: i32 // CHECK-SAME: %[[B1:[a-zA-Z0-9_]+]]: i32
// CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: i32 // CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: i32
// CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: i32 // CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: i32
// CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: i32) // CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: i32)
// CHECK: %[[T0:.+]] = arith.addi %[[B0]], %[[B1]] // CHECK: %[[T0:.+]] = arith.addi %[[B0]], %[[B1]]
// CHECK: %[[T1:.+]] = arith.addi %[[T0]], %[[B1]] // CHECK: %[[T1:.+]] = arith.addi %[[T0]], %[[B1]]
// CHECK: %[[T2:.+]] = arith.addi %[[T1]], %[[B2]] // CHECK: %[[T2:.+]] = arith.addi %[[T1]], %[[B2]]
Show All 9 Lines
#map = affine_map<(d0, d1) -> (d0, d1)> #map = affine_map<(d0, d1) -> (d0, d1)>
func.func @drop_redundant_results( func.func @drop_redundant_results(
%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) { %arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : tensor<?x?xf32>) ins(%arg0 : tensor<?x?xf32>)
outs(%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) { inits(%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):
%1 = arith.addf %b0, %b0 : f32 %1 = arith.addf %b0, %b0 : f32
linalg.yield %1, %1 : f32, f32 linalg.yield %1, %1 : f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>)
return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32> return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
} }
// CHECK: func @drop_redundant_results // CHECK: func @drop_redundant_results
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: outs(%[[ARG0]] : // CHECK-SAME: inits(%[[ARG0]] :
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
// ----- // -----
// Drop dead result with different tensors. // Drop dead result with different tensors.
#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) -> (d1, d2, d0)> #map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
#map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
func.func @drop_dead_results_with_different_tensors(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) { func.func @drop_dead_results_with_different_tensors(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32> %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
%0:4 = linalg.generic { %0:4 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3, #map4], indexing_maps = [#map0, #map1, #map2, #map3, #map4],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<?x?x?xf32>) ins(%arg0 : tensor<?x?x?xf32>)
outs(%arg0, %arg0, %init0, %init0 inits(%arg0, %arg0, %init0, %init0
: tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) { : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
linalg.yield %b0, %b0, %b3, %b4 : f32, f32, f32, f32 linalg.yield %b0, %b0, %b3, %b4 : f32, f32, f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32> return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
} }
// 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, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
// CHECK: func @drop_dead_results_with_different_tensors( // CHECK: func @drop_dead_results_with_different_tensors(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>)
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: outs(%[[ARG0]], %[[ARG0]] : // CHECK-SAME: inits(%[[ARG0]], %[[ARG0]] :
// CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1
// ----- // -----
// Drop dead result with unused cycles. // Drop dead result with unused cycles.
#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) -> (d1, d2, d0)> #map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
#map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map4 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
func.func @drop_dead_results_with_unused_cycles(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) { func.func @drop_dead_results_with_unused_cycles(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32> %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
%0:4 = linalg.generic { %0:4 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3, #map4], indexing_maps = [#map0, #map1, #map2, #map3, #map4],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<?x?x?xf32>) ins(%arg0 : tensor<?x?x?xf32>)
outs(%arg0, %arg0, %init0, %init0 inits(%arg0, %arg0, %init0, %init0
: tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) { : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32, %b4 : f32) :
%1 = arith.addf %b0, %b0: f32 %1 = arith.addf %b0, %b0: f32
%2 = arith.addf %b0, %b3: f32 %2 = arith.addf %b0, %b3: f32
%3 = arith.addf %b0, %b4: f32 %3 = arith.addf %b0, %b4: f32
linalg.yield %1, %1, %2, %3 : f32, f32, f32, f32 linalg.yield %1, %1, %2, %3 : f32, f32, f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32> return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>
} }
// 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, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
// CHECK: func @drop_dead_results_with_unused_cycles( // CHECK: func @drop_dead_results_with_unused_cycles(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>)
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: outs(%[[ARG0]], %[[ARG0]] : // CHECK-SAME: inits(%[[ARG0]], %[[ARG0]] :
// CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1
// ----- // -----
// Drop only the results not used by others. // Drop only the results not used by others.
#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) -> (d1, d2, d0)> #map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
func.func @drop_only_the_results_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) { func.func @drop_only_the_results_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32> %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
%0:3 = linalg.generic { %0:3 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<?x?x?xf32>) ins(%arg0 : tensor<?x?x?xf32>)
outs(%arg0, %init0, %init0 inits(%arg0, %init0, %init0
: tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) { : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) :
linalg.yield %b2, %b1, %b3 : f32, f32, f32 linalg.yield %b2, %b1, %b3 : f32, f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %0#0 : tensor<?x?x?xf32> return %0#0 : tensor<?x?x?xf32>
} }
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
// CHECK: func @drop_only_the_results_not_used_by_others( // CHECK: func @drop_only_the_results_not_used_by_others(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>)
// CHECK: %[[INIT:.+]] = tensor.empty // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]]
// CHECK-SAME: outs(%[[ARG0]], %[[INIT]] : // CHECK-SAME: inits(%[[ARG0]], %[[INIT]] :
// CHECK: return %[[GENERIC]]#0 // CHECK: return %[[GENERIC]]#0
// ----- // -----
// Drop only the cycles not used by others. // Drop only the cycles not used by others.
#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) -> (d1, d2, d0)> #map2 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)> #map3 = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
func.func @drop_only_the_cycles_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) { func.func @drop_only_the_cycles_not_used_by_others(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32> %d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32> %d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32> %init0 = tensor.empty(%d0, %d1, %d2) : tensor<?x?x?xf32>
%0:3 = linalg.generic { %0:3 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : tensor<?x?x?xf32>) ins(%arg0 : tensor<?x?x?xf32>)
outs(%arg0, %init0, %init0 inits(%arg0, %init0, %init0
: tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) { : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) :
%1 = arith.addf %b1, %b2: f32 %1 = arith.addf %b1, %b2: f32
%2 = arith.addf %b1, %b3 : f32 %2 = arith.addf %b1, %b3 : f32
linalg.yield %1, %b1, %2 : f32, f32, f32 linalg.yield %1, %b1, %2 : f32, f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>)
return %0#0 : tensor<?x?x?xf32> return %0#0 : tensor<?x?x?xf32>
} }
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
// CHECK: func @drop_only_the_cycles_not_used_by_others( // CHECK: func @drop_only_the_cycles_not_used_by_others(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>) // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>)
// CHECK: %[[INIT:.+]] = tensor.empty // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]]
// CHECK-SAME: outs(%[[ARG0]], %[[INIT]] : // CHECK-SAME: inits(%[[ARG0]], %[[INIT]] :
// CHECK: return %[[GENERIC]]#0 // CHECK: return %[[GENERIC]]#0

mlir/test/Dialect/Linalg/canonicalize.mlir

// RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s // RUN: mlir-opt %s -canonicalize -split-input-file | FileCheck %s
// CHECK-LABEL: func @memref_cast( // CHECK-LABEL: func @memref_cast(
func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> { func.func @memref_cast(%a: index, %b: index) -> memref<?x?xf32> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c8 = arith.constant 8 : index %c8 = arith.constant 8 : index
%c16 = arith.constant 16 : index %c16 = arith.constant 16 : index
%1 = memref.alloc (%b) : memref<?xi8> %1 = memref.alloc (%b) : memref<?xi8>
%2 = memref.view %1[%c0][] : memref<?xi8> to memref<16x16xf32> %2 = memref.view %1[%c0][] : memref<?xi8> to memref<16x16xf32>
%3 = memref.cast %2 : memref<16x16xf32> to memref<?x?xf32> %3 = memref.cast %2 : memref<16x16xf32> to memref<?x?xf32>
// CHECK: linalg.matmul ins({{.*}}memref<16x16xf32>, memref<16x16xf32>) outs({{.*}}memref<16x16xf32>) // CHECK: linalg.matmul ins({{.*}}memref<16x16xf32>, memref<16x16xf32>) inits({{.*}}memref<16x16xf32>)
linalg.matmul ins(%3, %3: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%3, %3: memref<?x?xf32>, memref<?x?xf32>)
outs(%3: memref<?x?xf32>) inits(%3: memref<?x?xf32>)
return %3: memref<?x?xf32> return %3: memref<?x?xf32>
} }
// ----- // -----
#accesses = [ #accesses = [
affine_map<(i) -> (i)> affine_map<(i) -> (i)>
] ]
#trait = { #trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
func.func @dce_zero_memref(%arg0 : memref<0xf32>, %arg1: tensor<0xf32>) -> tensor<0xf32> { func.func @dce_zero_memref(%arg0 : memref<0xf32>, %arg1: tensor<0xf32>) -> tensor<0xf32> {
// memref<0x32> is expected to be dce'ed // memref<0x32> is expected to be dce'ed
memref.copy %arg0, %arg0 : memref<0xf32> to memref<0xf32> memref.copy %arg0, %arg0 : memref<0xf32> to memref<0xf32>
// tensor<0xf32> cannot be dce'ed // tensor<0xf32> cannot be dce'ed
%1 = linalg.generic #trait outs(%arg1 : tensor<0xf32>) { %1 = linalg.generic #trait inits(%arg1 : tensor<0xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<0xf32> } -> tensor<0xf32>
return %1: tensor<0xf32> return %1: tensor<0xf32>
} }
// CHECK-LABEL: @dce_zero_memref // CHECK-LABEL: @dce_zero_memref
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<0xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<0xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<0xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<0xf32>
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
// CHECK-NEXT: return %[[ARG1]] // CHECK-NEXT: return %[[ARG1]]
// ----- // -----
// CHECK-LABEL: func @tensor.cast( // CHECK-LABEL: func @tensor.cast(
func.func @tensor.cast(%a : tensor<3x4xf32>, %b : tensor<4x?xf32>, %c : tensor<3x?xf32>) func.func @tensor.cast(%a : tensor<3x4xf32>, %b : tensor<4x?xf32>, %c : tensor<3x?xf32>)
-> tensor<3x?xf32> -> tensor<3x?xf32>
{ {
%ta = tensor.cast %a : tensor<3x4xf32> to tensor<?x?xf32> %ta = tensor.cast %a : tensor<3x4xf32> to tensor<?x?xf32>
%tb = tensor.cast %b : tensor<4x?xf32> to tensor<?x?xf32> %tb = tensor.cast %b : tensor<4x?xf32> to tensor<?x?xf32>
%tc = tensor.cast %c : tensor<3x?xf32> to tensor<?x?xf32> %tc = tensor.cast %c : tensor<3x?xf32> to tensor<?x?xf32>
// CHECK: linalg.matmul ins({{.*}}tensor<3x4xf32>, tensor<4x?xf32>) // CHECK: linalg.matmul ins({{.*}}tensor<3x4xf32>, tensor<4x?xf32>)
// CHECK-SAME: outs({{.*}}tensor<3x?xf32>) -> tensor<3x?xf32> // CHECK-SAME: inits({{.*}}tensor<3x?xf32>) -> tensor<3x?xf32>
%0 = linalg.matmul ins(%ta, %tb: tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%ta, %tb: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%tc: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%tc: tensor<?x?xf32>) -> tensor<?x?xf32>
%1 = tensor.cast %0 : tensor<?x?xf32> to tensor<3x?xf32> %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<3x?xf32>
return %1: tensor<3x?xf32> return %1: tensor<3x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @linalg_effects( // CHECK-LABEL: func @linalg_effects(
// CHECK-SAME: %[[A:[a-z0-9]*]]: tensor<?x?xf32> // CHECK-SAME: %[[A:[a-z0-9]*]]: tensor<?x?xf32>
// CHECK-SAME: %[[B:[a-z0-9]*]]: memref<?x?xf32> // CHECK-SAME: %[[B:[a-z0-9]*]]: memref<?x?xf32>
// CHECK-SAME: %[[C:[a-z0-9]*]]: tensor<?x?xf32> // CHECK-SAME: %[[C:[a-z0-9]*]]: tensor<?x?xf32>
func.func @linalg_effects(%a : tensor<?x?xf32>, %b : memref<?x?xf32>, %c : tensor<?x?xf32>) { func.func @linalg_effects(%a : tensor<?x?xf32>, %b : memref<?x?xf32>, %c : tensor<?x?xf32>) {
// CHECK-NOT: %{{.*}} = linalg.matmul // CHECK-NOT: %{{.*}} = linalg.matmul
%t = linalg.matmul ins(%a, %b : tensor<?x?xf32>, memref<?x?xf32>) %t = linalg.matmul ins(%a, %b : tensor<?x?xf32>, memref<?x?xf32>)
outs(%c : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%c : tensor<?x?xf32>) -> 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>) inits(%b : memref<?x?xf32>)
return return
} }
// ----- // -----
#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 = tensor.empty(%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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty(%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>) inits(%2 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3 : f32): ^bb0(%arg2: f32, %arg3 : f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %3 : tensor<?x?xf32> return %3 : tensor<?x?xf32>
} }
// CHECK-LABEL: func @keep_not_noop // CHECK-LABEL: func @keep_not_noop
// CHECK: %[[RESULT:.+]] = linalg.generic // CHECK: %[[RESULT:.+]] = linalg.generic
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
Show All 11 Linesfunc.func @keep_not_noop(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>)
%2 = tensor.empty(%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>) { inits(%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
Show All 15 Lines
// CHECK-NOT: linalg.generic // CHECK-NOT: linalg.generic
// CHECK-NOT: tensor.pad // CHECK-NOT: tensor.pad
// CHECK: return // CHECK: return
func.func @dead_linalg_tensor(%arg0 : tensor<7x7xi32>, %arg1 : tensor<7x7xf32>, func.func @dead_linalg_tensor(%arg0 : tensor<7x7xi32>, %arg1 : tensor<7x7xf32>,
%arg2: tensor<?x?xf32>, %high : index) { %arg2: tensor<?x?xf32>, %high : index) {
%c0_i32 = arith.constant 0 : i32 %c0_i32 = arith.constant 0 : i32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%0 = linalg.fill ins(%c0_i32 : i32) outs(%arg0 : tensor<7x7xi32>) -> tensor<7x7xi32> %0 = linalg.fill ins(%c0_i32 : i32) inits(%arg0 : tensor<7x7xi32>) -> tensor<7x7xi32>
%1 = linalg.matmul ins(%arg1, %arg1: tensor<7x7xf32>, tensor<7x7xf32>) %1 = linalg.matmul ins(%arg1, %arg1: tensor<7x7xf32>, tensor<7x7xf32>)
outs(%arg1: tensor<7x7xf32>) -> tensor<7x7xf32> inits(%arg1: tensor<7x7xf32>) -> tensor<7x7xf32>
%2 = linalg.generic #trait outs(%arg0 : tensor<7x7xi32>) { %2 = linalg.generic #trait inits(%arg0 : tensor<7x7xi32>) {
^bb(%3: i32) : ^bb(%3: i32) :
linalg.yield %3 : i32 linalg.yield %3 : i32
} -> tensor<7x7xi32> } -> tensor<7x7xi32>
%3 = tensor.pad %arg2 low[%c0, %c0] high[%high, %high] { %3 = tensor.pad %arg2 low[%c0, %c0] high[%high, %high] {
^bb0(%arg9: index, %arg10: index): ^bb0(%arg9: index, %arg10: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<?x?xf32> to tensor<2x4xf32> } : tensor<?x?xf32> to tensor<2x4xf32>
return return
} }
// ----- // -----
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 = tensor.empty(%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) inits(%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:.+]] = tensor.empty // CHECK: %[[INIT:.+]] = tensor.empty
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}inits(%[[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
func.func @self_copy(%arg0 : memref<2x3x?x4xf32>) { func.func @self_copy(%arg0 : memref<2x3x?x4xf32>) {
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
memref.copy %arg0, %arg0 : memref<2x3x?x4xf32> to memref<2x3x?x4xf32> memref.copy %arg0, %arg0 : memref<2x3x?x4xf32> to memref<2x3x?x4xf32>
// CHECK: return // CHECK: return
return return
} }
// ----- // -----
// CHECK-LABEL: func @remove_deadargs_generic_basic // CHECK-LABEL: func @remove_deadargs_generic_basic
// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> {
// 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: inits({{.*}} : 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 = tensor.empty(%0) : tensor<?xf32> %1 = tensor.empty(%0) : tensor<?xf32>
%2 = tensor.empty(%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>) inits (%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: inits({{.*}} : 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 = tensor.empty(%0, %1) : tensor<?x?xf32> %2 = tensor.empty(%0, %1) : tensor<?x?xf32>
%3 = tensor.empty(%1, %0) : tensor<?x?xf32> %3 = tensor.empty(%1, %0) : tensor<?x?xf32>
%4 = tensor.empty(%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>) inits (%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:.+]] = tensor.empty() : tensor<6x4xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<6x4xf32>
%empty = tensor.empty() : 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) inits(%[[INIT]] : tensor<6x4xf32>) -> tensor<6x4xf32>
%fill = linalg.fill ins(%zero : f32) outs(%empty : tensor<1x2x3x4xf32>) -> tensor<1x2x3x4xf32> %fill = linalg.fill ins(%zero : f32) inits(%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) inits(%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(%{{.+}}{{.*}}inits(%[[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>
} }
// ----- // -----
// 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>) { inits(%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:.+]] = tensor.empty() : tensor<412x276xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<412x276xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}inits(%[[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
%empty = tensor.empty() : tensor<400x273xf32> %empty = tensor.empty() : tensor<400x273xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<400x273xf32>) -> tensor<400x273xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
} }
// ----- // -----
// CHECK: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 9)> // CHECK: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 + 9)>
// CHECK: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 + 10)> // CHECK: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 + 10)>
// CHECK: #[[MAP2:.+]] = affine_map<()[s0] -> (s0 + 23)> // CHECK: #[[MAP2:.+]] = affine_map<()[s0] -> (s0 + 23)>
// CHECK: #[[MAP3:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 32)> // CHECK: #[[MAP3:.+]] = affine_map<()[s0, s1] -> (s0 + s1 + 32)>
// CHECK: func @fold_dynamic_pad_fill // CHECK: func @fold_dynamic_pad_fill
// CHECK-SAME: %[[SRC:.+]]: tensor<8x?x16x32xf32>, %[[LOW0:.+]]: index, %[[LOW3:.+]]: index, %[[HIGH2:.+]]: index, %[[HIGH3:.+]]: index // CHECK-SAME: %[[SRC:.+]]: tensor<8x?x16x32xf32>, %[[LOW0:.+]]: index, %[[LOW3:.+]]: index, %[[HIGH2:.+]]: index, %[[HIGH3:.+]]: index
// 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) inits(%[[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:.+]] = tensor.empty(%[[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]]{{.*}}inits(%[[INIT]]
// CHECK: return %[[FILL]] // CHECK: return %[[FILL]]
func.func @fold_dynamic_pad_fill(%empty: 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(%empty : tensor<8x?x16x32xf32>) -> tensor<8x?x16x32xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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
%empty = tensor.empty() : tensor<400x273xf32> %empty = tensor.empty() : tensor<400x273xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<400x273xf32>) -> tensor<400x273xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
} }
Show All 13 Linesfunc.func @static_input_without_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
%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 = tensor.empty(%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>) { inits(%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>)
%5 = tensor.cast %4 : tensor<?x?x?xf32> to tensor<2x3x4xf32> %5 = tensor.cast %4 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
return %5 : tensor<2x3x4xf32> return %5 : tensor<2x3x4xf32>
// CHECK: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32> // CHECK: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
// CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic // CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>) // CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
// CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>) // CHECK-SAME: inits({{.*}} : tensor<2x3x4xf32>)
} }
// ----- // -----
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @static_input_with_cast // CHECK-LABEL: func @static_input_with_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_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 = tensor.empty(%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>) { inits(%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>)
%6 = tensor.cast %5 : tensor<?x?x?xf32> to tensor<2x3x4xf32> %6 = tensor.cast %5 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
return %6: tensor<2x3x4xf32> return %6: tensor<2x3x4xf32>
// CHECK: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32> // CHECK: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
// CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic // CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>) // CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
// CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>) // CHECK-SAME: inits({{.*}} : tensor<2x3x4xf32>)
} }
// ----- // -----
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @static_output_with_cast // CHECK-LABEL: func @static_output_with_cast
// 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 = tensor.empty(%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>) { inits(%4 : tensor<2x3x4xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
%9 = arith.addf %arg3, %arg4 : f32 %9 = arith.addf %arg3, %arg4 : f32
linalg.yield %9 : f32 linalg.yield %9 : f32
} -> (tensor<2x3x4xf32>) } -> (tensor<2x3x4xf32>)
return %6: tensor<2x3x4xf32> return %6: tensor<2x3x4xf32>
// CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<2x3x4xf32> // CHECK: %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
// CHECK-NEXT: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32> // CHECK-NEXT: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
// CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic // CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%[[CAST_ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>) // CHECK-SAME: ins(%[[CAST_ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
// CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>) // CHECK-SAME: inits({{.*}} : tensor<2x3x4xf32>)
} }
// ----- // -----
// This test checks the folding of tensor.cast operation when the source value of cast // This test checks the folding of tensor.cast operation when the source value of cast
// has more static information than the destination value. // has more static information than the destination value.
#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @cast_source // CHECK-LABEL: func @cast_source
// 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 = tensor.empty(%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>) { inits(%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: inits({{.*}} : 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 = tensor.empty(%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>) { inits(%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
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
return %2 : tensor<?x?x?xf32> return %2 : tensor<?x?x?xf32>
// CHECK: %[[GENERIC_OP:.*]] = linalg.generic // CHECK: %[[GENERIC_OP:.*]] = linalg.generic
// CHECK-SAME: ins(%{{.*}}, %[[ARG1]] : tensor<1x?x?xf32>, tensor<1x?x?xf32>) // CHECK-SAME: ins(%{{.*}}, %[[ARG1]] : tensor<1x?x?xf32>, tensor<1x?x?xf32>)
// CHECK-SAME: outs(%{{.*}} : tensor<1x?x?xf32>) // CHECK-SAME: inits(%{{.*}} : tensor<1x?x?xf32>)
// CHECK: tensor.cast %[[GENERIC_OP]] : tensor<1x?x?xf32> to tensor<?x?x?xf32> // CHECK: tensor.cast %[[GENERIC_OP]] : tensor<1x?x?xf32> to tensor<?x?x?xf32>
} }
// ----- // -----
// 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:.+]] = tensor.empty() // CHECK: %[[INIT:.+]] = tensor.empty()
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]]{{.*}}inits(%[[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>
%empty = tensor.empty() : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
%empty = tensor.empty() : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
%empty = tensor.empty() : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
%empty = tensor.empty() : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
%empty = tensor.empty() : tensor<8x384x384xf32> %empty = tensor.empty() : tensor<8x384x384xf32>
%fill = linalg.fill ins(%f0 : f32) outs(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f0 : f32) inits(%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>
%empty = tensor.empty() : 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(%empty : tensor<8x384x384xf32>) -> tensor<8x384x384xf32> %fill = linalg.fill ins(%f1 : f32) inits(%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>
} }
// ----- // -----
func.func @fold_linalgop_with_cast_consumer(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, func.func @fold_linalgop_with_cast_consumer(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?x?xf32>) -> (tensor<4x8xf32>, tensor<?x?xf32>) { %arg2 : tensor<?x?xf32>) -> (tensor<4x8xf32>, tensor<?x?xf32>) {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x8xf32> %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x8xf32>
return %1, %0 : tensor<4x8xf32>, tensor<?x?xf32> return %1, %0 : tensor<4x8xf32>, tensor<?x?xf32>
} }
// CHECK: func @fold_linalgop_with_cast_consumer( // CHECK: func @fold_linalgop_with_cast_consumer(
// 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<?x?xf32>) // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32>)
// CHECK-DAG: %[[LHS_CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x?xf32> // CHECK-DAG: %[[LHS_CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?xf32> to tensor<4x?xf32>
// CHECK-DAG: %[[RHS_CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?x?xf32> to tensor<?x8xf32> // CHECK-DAG: %[[RHS_CAST:.+]] = tensor.cast %[[ARG1]] : tensor<?x?xf32> to tensor<?x8xf32>
// CHECK-DAG: %[[OUT_CAST:.+]] = tensor.cast %[[ARG2]] : tensor<?x?xf32> to tensor<4x8xf32> // CHECK-DAG: %[[OUT_CAST:.+]] = tensor.cast %[[ARG2]] : tensor<?x?xf32> to tensor<4x8xf32>
// CHECK: %[[MATMUL:.+]] = linalg.matmul // CHECK: %[[MATMUL:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS_CAST]], %[[RHS_CAST]] : // CHECK-SAME: ins(%[[LHS_CAST]], %[[RHS_CAST]] :
// CHECK-SAME: outs(%[[OUT_CAST]] : // CHECK-SAME: inits(%[[OUT_CAST]] :
// CHECK: %[[RESULT_CAST:.+]] = tensor.cast %[[MATMUL]] // CHECK: %[[RESULT_CAST:.+]] = tensor.cast %[[MATMUL]]
// CHECK: return %[[MATMUL]], %[[RESULT_CAST]] // CHECK: return %[[MATMUL]], %[[RESULT_CAST]]
// ----- // -----
func.func private @some_use(%0 : tensor<4x8xf32>) func.func private @some_use(%0 : tensor<4x8xf32>)
func.func @linalgop_with_cond_cast_consumer(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, func.func @linalgop_with_cond_cast_consumer(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?x?xf32>, %arg3 : i1) -> tensor<?x?xf32> { %arg2 : tensor<?x?xf32>, %arg3 : i1) -> tensor<?x?xf32> {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
scf.if %arg3 { scf.if %arg3 {
%1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x8xf32> %1 = tensor.cast %0 : tensor<?x?xf32> to tensor<4x8xf32>
func.call @some_use(%1) : (tensor<4x8xf32>) -> () func.call @some_use(%1) : (tensor<4x8xf32>) -> ()
} }
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// Check conditionally reachable cast is not folded into producer. // Check conditionally reachable cast is not folded into producer.
// CHECK-LABEL: func @linalgop_with_cond_cast_consumer // CHECK-LABEL: func @linalgop_with_cond_cast_consumer
// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>, %[[ARG3:.*]]: i1) // CHECK-SAME: (%[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x?xf32>, %[[ARG2:.*]]: tensor<?x?xf32>, %[[ARG3:.*]]: i1)
// CHECK: %[[RES:.*]] = linalg.matmul ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>) // CHECK: %[[RES:.*]] = linalg.matmul ins(%[[ARG0]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32> // CHECK-SAME: inits(%[[ARG2]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: scf.if %[[ARG3]] { // CHECK: scf.if %[[ARG3]] {
// CHECK: %[[CAST:.*]] = tensor.cast %[[RES]] : tensor<?x?xf32> to tensor<4x8xf32> // CHECK: %[[CAST:.*]] = tensor.cast %[[RES]] : tensor<?x?xf32> to tensor<4x8xf32>
// CHECK: func.call @some_use(%[[CAST]]) : (tensor<4x8xf32>) -> () // CHECK: func.call @some_use(%[[CAST]]) : (tensor<4x8xf32>) -> ()
// CHECK: } // CHECK: }
// CHECK: return %[[RES]] : tensor<?x?xf32> // CHECK: return %[[RES]] : tensor<?x?xf32>
// ----- // -----
func.func @fold_conv_op_with_cast_consumer(%arg0 : tensor<?x?x?x?xf32>, func.func @fold_conv_op_with_cast_consumer(%arg0 : tensor<?x?x?x?xf32>,
%arg1 : tensor<?x?x?x?xf32>, %arg2 : tensor<?x?x?x?xf32>) -> %arg1 : tensor<?x?x?x?xf32>, %arg2 : tensor<?x?x?x?xf32>) ->
(tensor<4x8x12x16xf32>, tensor<?x?x?x?xf32>) { (tensor<4x8x12x16xf32>, tensor<?x?x?x?xf32>) {
%0 = linalg.conv_2d_nchw_fchw ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) %0 = linalg.conv_2d_nchw_fchw ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
outs(%arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> inits(%arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
%1 = tensor.cast %0 : tensor<?x?x?x?xf32> to tensor<4x8x12x16xf32> %1 = tensor.cast %0 : tensor<?x?x?x?xf32> to tensor<4x8x12x16xf32>
return %1, %0 : tensor<4x8x12x16xf32>, tensor<?x?x?x?xf32> return %1, %0 : tensor<4x8x12x16xf32>, tensor<?x?x?x?xf32>
} }
// CHECK: func @fold_conv_op_with_cast_consumer( // CHECK: func @fold_conv_op_with_cast_consumer(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>) // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?x?x?xf32>)
// CHECK: %[[OUT_CAST:.+]] = tensor.cast %[[ARG2]] : tensor<?x?x?x?xf32> to tensor<4x8x12x16xf32> // CHECK: %[[OUT_CAST:.+]] = tensor.cast %[[ARG2]] : tensor<?x?x?x?xf32> to tensor<4x8x12x16xf32>
// CHECK: %[[CONV:.+]] = linalg.conv_2d_nchw_fchw // CHECK: %[[CONV:.+]] = linalg.conv_2d_nchw_fchw
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CHECK-SAME: outs(%[[OUT_CAST]] : // CHECK-SAME: inits(%[[OUT_CAST]] :
// CHECK: %[[RESULT_CAST:.+]] = tensor.cast %[[CONV]] // CHECK: %[[RESULT_CAST:.+]] = tensor.cast %[[CONV]]
// CHECK: return %[[CONV]], %[[RESULT_CAST]] // CHECK: return %[[CONV]], %[[RESULT_CAST]]
// ----- // -----
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>
%empty1 = tensor.empty(%d1, %d2, %d0) : tensor<?x?x?xf32> %empty1 = tensor.empty(%d1, %d2, %d0) : tensor<?x?x?xf32>
%empty2 = tensor.empty(%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(%empty1, %empty2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) { ins(%arg0 : tensor<?x?x?xf32>) inits(%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:.+]] = tensor.empty() : 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:.+]] = tensor.empty() : 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: inits(%[[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
// ----- // -----
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
func.func @identity_mixed(%arg0 : tensor<?xf32>, %arg1: memref<?xf32>) { func.func @identity_mixed(%arg0 : tensor<?xf32>, %arg1: memref<?xf32>) {
linalg.generic { linalg.generic {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%arg0 : tensor<?xf32>) } ins(%arg0 : tensor<?xf32>)
outs(%arg1 : memref<?xf32>) { inits(%arg1 : memref<?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} }
return return
} }
// There was a crash in EraseIdentityGenericOp for generic with mixed semantics. // There was a crash in EraseIdentityGenericOp for generic with mixed semantics.
// For now, check generic remained unchanged. // For now, check generic remained unchanged.
// CHECK-LABEL: func @identity_mixed // CHECK-LABEL: func @identity_mixed
// CHECK-SAME: (%[[ARG1:.*]]: tensor<?xf32>, %[[ARG2:.*]]: memref<?xf32>) // CHECK-SAME: (%[[ARG1:.*]]: tensor<?xf32>, %[[ARG2:.*]]: memref<?xf32>)
// CHECK: linalg.generic { // CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = [#map, #map], // CHECK-SAME: indexing_maps = [#map, #map],
// CHECK-SAME: iterator_types = ["parallel"] // CHECK-SAME: iterator_types = ["parallel"]
// CHECK-SAME: } ins(%[[ARG1]] : tensor<?xf32>) // CHECK-SAME: } ins(%[[ARG1]] : tensor<?xf32>)
// CHECK-SAME: outs(%[[ARG2]] : memref<?xf32>) { // CHECK-SAME: inits(%[[ARG2]] : memref<?xf32>) {

mlir/test/Dialect/Linalg/collapse-dim.mlir

// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=collapse-dimensions-control=2,3 -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=collapse-dimensions-control=2,3 -split-input-file | FileCheck %s
func.func @collapse_reduction( func.func @collapse_reduction(
%arg0: tensor<2x32x10x4096xf32>, %arg1: tensor<2x32xf32>) -> tensor<2x32xf32> { %arg0: tensor<2x32x10x4096xf32>, %arg1: tensor<2x32xf32>) -> tensor<2x32xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, 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<2x32x10x4096xf32>) outs(%arg1 : tensor<2x32xf32>) { ins(%arg0 : tensor<2x32x10x4096xf32>) inits(%arg1 : tensor<2x32xf32>) {
^bb0(%arg3: f32, %arg4: f32): ^bb0(%arg3: f32, %arg4: f32):
%1 = arith.addf %arg3, %arg4 : f32 %1 = arith.addf %arg3, %arg4 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<2x32xf32> } -> tensor<2x32xf32>
return %0 : tensor<2x32xf32> return %0 : tensor<2x32xf32>
} }
// 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: func @collapse_reduction // CHECK-LABEL: func @collapse_reduction
// CHECK: %[[T:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<2x32x10x4096xf32> into tensor<2x32x40960xf32> // CHECK: %[[T:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<2x32x10x4096xf32> into tensor<2x32x40960xf32>
// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], // CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
// CHECK-SAME: ins(%[[T]] : tensor<2x32x40960xf32>) outs(%{{.*}} : tensor<2x32xf32>) { // CHECK-SAME: ins(%[[T]] : tensor<2x32x40960xf32>) inits(%{{.*}} : tensor<2x32xf32>) {
// CHECK: } -> tensor<2x32xf32> // CHECK: } -> tensor<2x32xf32>
// ----- // -----
func.func @collapse_parallel( func.func @collapse_parallel(
%arg0: tensor<32x2x10x4096xf32>, %arg1: tensor<2x32x10x4096xf32>) -> tensor<2x32x10x4096xf32> { %arg0: tensor<32x2x10x4096xf32>, %arg1: tensor<2x32x10x4096xf32>) -> tensor<2x32x10x4096xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d1, d0, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d1, d0, d2, 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<32x2x10x4096xf32>) outs(%arg1 : tensor<2x32x10x4096xf32>) { ins(%arg0 : tensor<32x2x10x4096xf32>) inits(%arg1 : tensor<2x32x10x4096xf32>) {
^bb0(%arg3: f32, %arg4: f32): ^bb0(%arg3: f32, %arg4: f32):
%1 = arith.addf %arg3, %arg4 : f32 %1 = arith.addf %arg3, %arg4 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<2x32x10x4096xf32> } -> tensor<2x32x10x4096xf32>
return %0 : tensor<2x32x10x4096xf32> return %0 : tensor<2x32x10x4096xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @collapse_parallel // CHECK-LABEL: func @collapse_parallel
// CHECK-DAG: %[[S:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<32x2x10x4096xf32> into tensor<32x2x40960xf32> // CHECK-DAG: %[[S:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<32x2x10x4096xf32> into tensor<32x2x40960xf32>
// CHECK-DAG: %[[D:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<2x32x10x4096xf32> into tensor<2x32x40960xf32> // CHECK-DAG: %[[D:.*]] = tensor.collapse_shape %{{.*}} {{\[}}[0], [1], [2, 3]] : tensor<2x32x10x4096xf32> into tensor<2x32x40960xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]], // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP1]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[S]] : tensor<32x2x40960xf32>) outs(%[[D]] : tensor<2x32x40960xf32>) { // CHECK-SAME: ins(%[[S]] : tensor<32x2x40960xf32>) inits(%[[D]] : tensor<2x32x40960xf32>) {
// CHECK: } -> tensor<2x32x40960xf32> // CHECK: } -> tensor<2x32x40960xf32>
// CHECK: tensor.expand_shape %[[R]] {{\[}}[0], [1], [2, 3]] : tensor<2x32x40960xf32> into tensor<2x32x10x4096xf32> // CHECK: tensor.expand_shape %[[R]] {{\[}}[0], [1], [2, 3]] : tensor<2x32x40960xf32> into tensor<2x32x10x4096xf32>

mlir/test/Dialect/Linalg/conv-interface-invalid.mlir

Show All 12 Lines
// Check for number of operands being >= 2. // Check for number of operands being >= 2.
#map = affine_map<(d0) -> (d0)> #map = affine_map<(d0) -> (d0)>
func.func @test_conv_op_wrong_num_operands(%arg0 : tensor<?xf32>, func.func @test_conv_op_wrong_num_operands(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{expected op with 2 inputs and 1 output}} // expected-error @+1 {{expected op with 2 inputs and 1 output}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<?xf32>) { ins(%arg0 : tensor<?xf32>) inits(%arg1 : tensor<?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg3 : f32 linalg.yield %arg3 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
func.func @test_conv_op_wrong_input_indexing_map1(%arg0 : tensor<?xf32>, func.func @test_conv_op_wrong_input_indexing_map1(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{unexpected input index map for convolution}} // expected-error @+1 {{unexpected input index map for convolution}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d0 * 2)>, indexing_maps = [affine_map<(d0, d1) -> (d0 * 2)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
func.func @test_conv_op_wrong_input_indexing_map2(%arg0 : tensor<?x?xf32>, func.func @test_conv_op_wrong_input_indexing_map2(%arg0 : tensor<?x?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{unexpected input index map for convolution}} // expected-error @+1 {{unexpected input index map for convolution}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d0 + d1, d0)>, indexing_maps = [affine_map<(d0, d1) -> (d0 + d1, d0)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
func.func @test_conv_op_filter_index_map_not_projection(%arg0 : tensor<?xf32>, func.func @test_conv_op_filter_index_map_not_projection(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{expected output/filter indexing maps to be projected permutations}} // expected-error @+1 {{expected output/filter indexing maps to be projected permutations}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d1)>, indexing_maps = [affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d1 + d0)>, affine_map<(d0, d1) -> (d1 + d0)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
func.func @test_conv_op_output_index_map_not_projection(%arg0 : tensor<?xf32>, func.func @test_conv_op_output_index_map_not_projection(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{expected output/filter indexing maps to be projected permutations}} // expected-error @+1 {{expected output/filter indexing maps to be projected permutations}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d0)>, indexing_maps = [affine_map<(d0, d1) -> (d0)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0 + d1)>], affine_map<(d0, d1) -> (d0 + d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Convolution op illegal if a loop dimension is used to access // Convolution op illegal if a loop dimension is used to access
// output, filter and is convolved. // output, filter and is convolved.
func.func @test_conv_op_output_filter_convolved(%arg0 : tensor<?xf32>, func.func @test_conv_op_output_filter_convolved(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
// expected-error @+1 {{unexpected loop dimension for convolution op}} // expected-error @+1 {{unexpected loop dimension for convolution op}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?x?xf32>) { inits(%arg2 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
// Convolution op illegal if a loop dimension is used only in the output. // Convolution op illegal if a loop dimension is used only in the output.
func.func @test_conv_op_output_only_dim(%arg0 : tensor<?xf32>, func.func @test_conv_op_output_only_dim(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
// expected-error @+1 {{unexpected loop dimension for convolution op}} // expected-error @+1 {{unexpected loop dimension for convolution op}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0 + d1)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0 + d1)>,
affine_map<(d0, d1, d2) -> (d1)>, affine_map<(d0, d1, d2) -> (d1)>,
affine_map<(d0, d1, d2) -> (d0, d2)>], affine_map<(d0, d1, d2) -> (d0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]} iterator_types = ["parallel", "reduction", "parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?x?xf32>) { inits(%arg2 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
// Convolution op illegal if a loop dimension is used only in the filter. // Convolution op illegal if a loop dimension is used only in the filter.
func.func @test_conv_op_filter_only_dim(%arg0 : tensor<?xf32>, func.func @test_conv_op_filter_only_dim(%arg0 : tensor<?xf32>,
%arg1 : tensor<?x?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?x?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{unexpected loop dimension for convolution op}} // expected-error @+1 {{unexpected loop dimension for convolution op}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0 + d1)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0 + d1)>,
affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>,
affine_map<(d0, d1, d2) -> (d0)>], affine_map<(d0, d1, d2) -> (d0)>],
iterator_types = ["parallel", "reduction", "reduction"]} iterator_types = ["parallel", "reduction", "reduction"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Convolution op illegal if a loop dimension is used only in the input. // Convolution op illegal if a loop dimension is used only in the input.
func.func @test_conv_op_input_only_dim(%arg0 : tensor<?x?xf32>, func.func @test_conv_op_input_only_dim(%arg0 : tensor<?x?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{unexpected loop dimension for convolution op}} // expected-error @+1 {{unexpected loop dimension for convolution op}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
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)>, affine_map<(d0, d1, d2) -> (d1)>,
affine_map<(d0, d1, d2) -> (d0)>], affine_map<(d0, d1, d2) -> (d0)>],
iterator_types = ["parallel", "reduction", "reduction"]} iterator_types = ["parallel", "reduction", "reduction"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Convolution op illegal if a loop dimension accessing output is not parallel. // Convolution op illegal if a loop dimension accessing output is not parallel.
func.func @test_conv_op_non_output_access_loop_parallel(%arg0 : tensor<?xf32>, func.func @test_conv_op_non_output_access_loop_parallel(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{expected all iterators not used to access outputs to be reduction}} // expected-error @+1 {{expected all iterators not used to access outputs to be reduction}}
%0 = test.linalg_conv_op { %0 = test.linalg_conv_op {
indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0 + d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2 : tensor<?xf32>) { inits(%arg2 : tensor<?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
linalg.yield %arg5 : f32 linalg.yield %arg5 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }

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

// RUN: mlir-opt -pass-pipeline="func.func(convert-elementwise-to-linalg)" -split-input-file %s | FileCheck %s // RUN: mlir-opt -pass-pipeline="func.func(convert-elementwise-to-linalg)" -split-input-file %s | FileCheck %s
// In-depth checking of the linalg.generic op for a very trivial case. // In-depth checking of the linalg.generic op for a very trivial case.
// CHECK: #[[$MAP:.*]] = affine_map<() -> ()> // CHECK: #[[$MAP:.*]] = affine_map<() -> ()>
// CHECK-LABEL: func @addf_rank0 // CHECK-LABEL: func @addf_rank0
// 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 @addf_rank0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> { func.func @addf_rank0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
// CHECK: %{{.*}} = linalg.generic // CHECK: %{{.*}} = linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]] // CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]]
// CHECK-SAME: iterator_types = [] // CHECK-SAME: iterator_types = []
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]
// CHECK-SAME: outs(%[[ARG0]] // CHECK-SAME: inits(%[[ARG0]]
// CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32, %{{.*}}: f32): // CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32, %{{.*}}: f32):
// CHECK: %[[YIELD:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32 // CHECK: %[[YIELD:.*]] = arith.addf %[[LHS]], %[[RHS]] : f32
// CHECK: linalg.yield %[[YIELD]] : f32 // CHECK: linalg.yield %[[YIELD]] : f32
// CHECK: } -> tensor<f32> // CHECK: } -> tensor<f32>
%0 = arith.addf %arg0, %arg1 : tensor<f32> %0 = arith.addf %arg0, %arg1 : tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
// ----- // -----
// Check indexing maps and iterator types for the rank > 0 case. // Check indexing maps and iterator types for the rank > 0 case.
// CHECK-LABEL: func @addf_rank1 // CHECK-LABEL: func @addf_rank1
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<?xf32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<?xf32>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<?xf32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<?xf32>
func.func @addf_rank1(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> { func.func @addf_rank1(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: iterator_types = ["parallel"] // CHECK-SAME: iterator_types = ["parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]]
// CHECK-SAME: outs(%[[ARG0]] // CHECK-SAME: inits(%[[ARG0]]
%0 = arith.addf %arg0, %arg1 : tensor<?xf32> %0 = arith.addf %arg0, %arg1 : tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
// Check a unary op. // Check a unary op.
// CHECK-LABEL: func @exp // CHECK-LABEL: func @exp
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<f32>
func.func @exp(%arg0: tensor<f32>) -> tensor<f32> { func.func @exp(%arg0: tensor<f32>) -> tensor<f32> {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0]] // CHECK-SAME: ins(%[[ARG0]]
// CHECK-SAME: outs(%[[ARG0]] // CHECK-SAME: inits(%[[ARG0]]
// CHECK: ^bb0(%[[SCALAR:.*]]: f32, %{{.*}}: f32): // CHECK: ^bb0(%[[SCALAR:.*]]: f32, %{{.*}}: f32):
// CHECK: %[[YIELD:.*]] = math.exp %[[SCALAR]] : f32 // CHECK: %[[YIELD:.*]] = math.exp %[[SCALAR]] : f32
// CHECK: linalg.yield %[[YIELD]] : f32 // CHECK: linalg.yield %[[YIELD]] : f32
%0 = math.exp %arg0 : tensor<f32> %0 = math.exp %arg0 : tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
// ----- // -----
// Check a case with varying operand types. // Check a case with varying operand types.
// CHECK-LABEL: func @select // CHECK-LABEL: func @select
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<i1> // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: tensor<i1>
// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<i32> // CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: tensor<i32>
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]: tensor<i32> // CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]: tensor<i32>
func.func @select(%arg0: tensor<i1>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<i32> { func.func @select(%arg0: tensor<i1>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<i32> {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]] // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]]
// CHECK-SAME: outs(%[[ARG1]] // CHECK-SAME: inits(%[[ARG1]]
// CHECK: ^bb0(%[[PRED:.*]]: i1, %[[TRUE_VAL:.*]]: i32, %[[FALSE_VAL:.*]]: i32, %{{.*}}: i32): // CHECK: ^bb0(%[[PRED:.*]]: i1, %[[TRUE_VAL:.*]]: i32, %[[FALSE_VAL:.*]]: i32, %{{.*}}: i32):
// CHECK: arith.select %[[PRED]], %[[TRUE_VAL]], %[[FALSE_VAL]] : i32 // CHECK: arith.select %[[PRED]], %[[TRUE_VAL]], %[[FALSE_VAL]] : i32
%0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32> %0 = arith.select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
// ----- // -----
// 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:.*]] = tensor.empty() : 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: inits(%[[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:.*]] = tensor.empty(%[[D1]], %[[D2]], %[[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: inits(%[[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

Show All 9 Linesfunc.func @simple_op(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
%init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32> %init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32>
%init2 = tensor.empty(%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>) { inits(%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) :
%0 = arith.addf %b0, %b1 : f32 %0 = arith.addf %b0, %b1 : f32
%1 = arith.mulf %0, %b2 : f32 %1 = arith.mulf %0, %b2 : f32
linalg.yield %0, %1 : f32, f32 linalg.yield %0, %1 : f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>)
return %result#0, %result#1 : tensor<?x?xf32>, tensor<?x?xf32> return %result#0, %result#1 : tensor<?x?xf32>, tensor<?x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
Show All 9 Lines
// 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:.+]] = tensor.empty(%[[D1]], %[[D0]]) // CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[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: inits(%[[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
// CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32): // CHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32):
// CHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]] // CHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]]
// CHECK-NEXT: linalg.yield %[[S0]], %{{[a-zA-Z0-9_]+}}, %[[S0]] // CHECK-NEXT: linalg.yield %[[S0]], %{{[a-zA-Z0-9_]+}}, %[[S0]]
// CHECK: %[[GENERIC2:.+]]:2 = linalg.generic // CHECK: %[[GENERIC2:.+]]:2 = linalg.generic
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP3]], #[[MAP0]]] // CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP3]], #[[MAP0]]]
// CHECK-SAME: ["parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[GENERIC1]]#2 : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[GENERIC1]]#2 :
// CHECK-SAME: outs(%[[INIT1]], %[[INIT2]] : // CHECK-SAME: inits(%[[INIT1]], %[[INIT2]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B8:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B8:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B9:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B9:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B10:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B10:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f32): // CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f32):
// CHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B9]], %[[B8]] // CHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B9]], %[[B8]]
Show All 15 Lines
// 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:.+]] = tensor.empty(%[[D1]], %[[D0]]) // CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CANONICALIZECHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[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: inits(%[[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
// CANONICALIZECHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32): // CANONICALIZECHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32):
// CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]] // CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S0]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S0]]
// CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: [#[[MAP3]], #[[MAP2]], #[[MAP0]]] // CANONICALIZECHECK-SAME: [#[[MAP3]], #[[MAP2]], #[[MAP0]]]
// CANONICALIZECHECK-SAME: ["parallel", "parallel"] // CANONICALIZECHECK-SAME: ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG2]], %[[GENERIC1]] : // CANONICALIZECHECK-SAME: ins(%[[ARG2]], %[[GENERIC1]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT2]] : // CANONICALIZECHECK-SAME: inits(%[[INIT2]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32): // CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32):
// CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B4]], %[[B3]] // CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B4]], %[[B3]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S1]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S1]]
// CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]] // CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]]
Show All 9 Linesfunc.func @simple_op_permuted_outputs(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?xf32>)
%init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32> %init1 = tensor.empty(%d1, %d0) : tensor<?x?xf32>
%init2 = tensor.empty(%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>) { inits(%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) :
%0 = arith.addf %b0, %b1 : f32 %0 = arith.addf %b0, %b1 : f32
%1 = arith.mulf %0, %b2 : f32 %1 = arith.mulf %0, %b2 : f32
linalg.yield %0, %1, %0 : f32, f32, f32 linalg.yield %0, %1, %0 : f32, f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>)
return %result#0, %result#1, %result#2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32> return %result#0, %result#1, %result#2 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
Show All 9 Lines
// 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:.+]] = tensor.empty(%[[D1]], %[[D0]]) // CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[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: inits(%[[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
// CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B3:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32): // CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32):
// CHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]] // CHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]]
// CHECK-NEXT: linalg.yield %[[S0]], %{{[a-zA-Z0-9_]+}}, %[[S0]] // CHECK-NEXT: linalg.yield %[[S0]], %{{[a-zA-Z0-9_]+}}, %[[S0]]
// CHECK: %[[GENERIC2:.+]]:3 = linalg.generic // CHECK: %[[GENERIC2:.+]]:3 = linalg.generic
// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP3]], #[[MAP0]], #[[MAP0]]] // CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP3]], #[[MAP0]], #[[MAP0]]]
// CHECK-SAME: ["parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[GENERIC1]]#3 : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[GENERIC1]]#3 :
// CHECK-SAME: outs(%[[INIT1]], %[[INIT2]], %[[INIT2]] : // CHECK-SAME: inits(%[[INIT1]], %[[INIT2]], %[[INIT2]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B8:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B8:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B9:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B9:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B10:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B10:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[B12:[a-zA-Z0-9_]+]]: f32): // CHECK-SAME: %[[B12:[a-zA-Z0-9_]+]]: f32):
// CHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B10]], %[[B9]] // CHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B10]], %[[B9]]
Show All 13 Lines
// 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:.+]] = tensor.empty(%[[D1]], %[[D0]]) // CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[D1]], %[[D0]])
// CANONICALIZECHECK-DAG: %[[INIT2:.+]] = tensor.empty(%[[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: inits(%[[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
// CANONICALIZECHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32): // CANONICALIZECHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f32):
// CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]] // CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.addf %[[B0]], %[[B1]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S0]], %[[S0]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S0]], %[[S0]]
// CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: [#[[MAP3]], #[[MAP2]], #[[MAP0]]] // CANONICALIZECHECK-SAME: [#[[MAP3]], #[[MAP2]], #[[MAP0]]]
// CANONICALIZECHECK-SAME: ["parallel", "parallel"] // CANONICALIZECHECK-SAME: ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG2]], %[[GENERIC1]]#0 : // CANONICALIZECHECK-SAME: ins(%[[ARG2]], %[[GENERIC1]]#0 :
// CANONICALIZECHECK-SAME: outs(%[[INIT2]] : // CANONICALIZECHECK-SAME: inits(%[[INIT2]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32): // CANONICALIZECHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f32):
// CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B5]], %[[B4]] // CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.mulf %[[B5]], %[[B4]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S1]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S1]]
// 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 = tensor.empty() : 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>) { inits(%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:.+]] = tensor.empty() : tensor<20x10xf64> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<20x10xf64>
// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : 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: inits(%[[INIT0]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B0:.+]]: f32 // CHECK-SAME: %[[B0:.+]]: f32
// CHECK-SAME: %[[B1:.+]]: i32 // CHECK-SAME: %[[B1:.+]]: i32
// CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f64 // CHECK-SAME: %[[B2:[a-zA-Z0-9_]+]]: f64
// CHECK-SAME: %[[B3:.+]]: f64 // CHECK-SAME: %[[B3:.+]]: f64
// CHECK-NEXT: %[[S0:.+]] = arith.sitofp %[[B1]] : i32 to f64 // CHECK-NEXT: %[[S0:.+]] = arith.sitofp %[[B1]] : i32 to f64
// CHECK-NEXT: linalg.yield %{{.+}}, %[[S0]] // CHECK-NEXT: linalg.yield %{{.+}}, %[[S0]]
// CHECK: %[[GENERIC1:.+]]:2 = linalg.generic // CHECK: %[[GENERIC1:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP0]], #[[MAP2]], #[[MAP0]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP0]], #[[MAP2]], #[[MAP0]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC0]]#1 : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC0]]#1 :
// CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] : // CHECK-SAME: inits(%[[INIT0]], %[[INIT1]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B4:.+]]: f32 // CHECK-SAME: %[[B4:.+]]: f32
// CHECK-SAME: %[[B5:.+]]: i32 // CHECK-SAME: %[[B5:.+]]: i32
// CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f64 // CHECK-SAME: %[[B6:[a-zA-Z0-9_]+]]: f64
// CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f64 // CHECK-SAME: %[[B7:[a-zA-Z0-9_]+]]: f64
// CHECK-SAME: %[[B8:.+]]: f64 // CHECK-SAME: %[[B8:.+]]: f64
// CHECK-NEXT: %[[S1:.+]] = arith.extf %[[B4]] : f32 to f64 // CHECK-NEXT: %[[S1:.+]] = arith.extf %[[B4]] : f32 to f64
// CHECK-NEXT: linalg.yield %{{.+}}, %[[S1]] // CHECK-NEXT: linalg.yield %{{.+}}, %[[S1]]
// CHECK: %[[GENERIC2:.+]] = linalg.generic // CHECK: %[[GENERIC2:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP0]], #[[MAP0]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP0]], #[[MAP0]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC0]]#1, %[[GENERIC1]]#1 : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC0]]#1, %[[GENERIC1]]#1 :
// CHECK-SAME: outs(%[[INIT0]] : // CHECK-SAME: inits(%[[INIT0]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[B9:.+]]: f32 // CHECK-SAME: %[[B9:.+]]: f32
// CHECK-SAME: %[[B10:.+]]: i32 // CHECK-SAME: %[[B10:.+]]: i32
// CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f64 // CHECK-SAME: %[[B11:[a-zA-Z0-9_]+]]: f64
// CHECK-SAME: %[[B12:[a-zA-Z0-9_]+]]: f64 // CHECK-SAME: %[[B12:[a-zA-Z0-9_]+]]: f64
// CHECK-SAME: %[[B13:.+]]: f64 // CHECK-SAME: %[[B13:.+]]: f64
// CHECK-NEXT: %[[S2:.+]] = arith.addf %[[B11]], %[[B12]] : f64 // CHECK-NEXT: %[[S2:.+]] = arith.addf %[[B11]], %[[B12]] : f64
// CHECK-NEXT: linalg.yield %[[S2]] // CHECK-NEXT: linalg.yield %[[S2]]
// 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:.+]] = tensor.empty() : tensor<20x10xf64> // CANONICALIZECHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<20x10xf64>
// CANONICALIZECHECK-DAG: %[[INIT1:.+]] = tensor.empty() : 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: inits(%[[INIT1]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B0:.+]]: i32 // CANONICALIZECHECK-SAME: %[[B0:.+]]: i32
// CANONICALIZECHECK-SAME: %[[B1:.+]]: f64 // CANONICALIZECHECK-SAME: %[[B1:.+]]: f64
// CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.sitofp %[[B0]] : i32 to f64 // CANONICALIZECHECK-NEXT: %[[S0:.+]] = arith.sitofp %[[B0]] : i32 to f64
// CANONICALIZECHECK-NEXT: linalg.yield %[[S0]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S0]]
// CANONICALIZECHECK: %[[GENERIC1:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC1:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP1]]] // CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP1]]]
// CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"] // CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG0]] : // CANONICALIZECHECK-SAME: ins(%[[ARG0]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT1]] : // CANONICALIZECHECK-SAME: inits(%[[INIT1]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B2:.+]]: f32 // CANONICALIZECHECK-SAME: %[[B2:.+]]: f32
// CANONICALIZECHECK-SAME: %[[B3:.+]]: f64 // CANONICALIZECHECK-SAME: %[[B3:.+]]: f64
// CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.extf %[[B2]] : f32 to f64 // CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.extf %[[B2]] : f32 to f64
// CANONICALIZECHECK-NEXT: linalg.yield %[[S1]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S1]]
// CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC2:.+]] = linalg.generic
// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP1]], #[[MAP2]]] // CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP1]], #[[MAP2]]]
// CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"] // CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] : // CANONICALIZECHECK-SAME: ins(%[[GENERIC0]], %[[GENERIC1]] :
// CANONICALIZECHECK-SAME: outs(%[[INIT0]] : // CANONICALIZECHECK-SAME: inits(%[[INIT0]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f64 // CANONICALIZECHECK-SAME: %[[B4:[a-zA-Z0-9_]+]]: f64
// CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f64 // CANONICALIZECHECK-SAME: %[[B5:[a-zA-Z0-9_]+]]: f64
// CANONICALIZECHECK-SAME: %[[B6:.+]]: f64 // CANONICALIZECHECK-SAME: %[[B6:.+]]: f64
// CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.addf %[[B4]], %[[B5]] : f64 // CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.addf %[[B4]], %[[B5]] : f64
// CANONICALIZECHECK-NEXT: linalg.yield %[[S2]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S2]]
// CANONICALIZECHECK: return %[[GENERIC2]] // CANONICALIZECHECK: return %[[GENERIC2]]
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0)> #map0 = affine_map<(d0, d1) -> (d0)>
#map1 = affine_map<(d0, d1) -> (d1)> #map1 = affine_map<(d0, d1) -> (d1)>
#map2 = affine_map<(d0, d1) -> (d0, d1)> #map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1) -> (d1, d0)> #map3 = affine_map<(d0, d1) -> (d1, d0)>
func.func @destination_passing_style( func.func @destination_passing_style(
%arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>, %arg0 : tensor<?xf32>, %arg1 : tensor<?xf32>,
%arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>) %arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>) { -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map3], indexing_maps = [#map0, #map1, #map2, #map3],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>) ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)
outs(%arg2, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>) { inits(%arg2, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) : ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32) :
%1 = arith.addf %b0, %b2 : f32 %1 = arith.addf %b0, %b2 : f32
%2 = arith.mulf %b1, %b3 : f32 %2 = arith.mulf %b1, %b3 : f32
linalg.yield %1, %2 : f32, f32 linalg.yield %1, %2 : f32, f32
} -> (tensor<?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?xf32>, tensor<?x?xf32>)
return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32> return %0#0, %0#1 : tensor<?x?xf32>, tensor<?x?xf32>
} }
// 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-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, d0)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK: func.func @destination_passing_style( // CHECK: func.func @destination_passing_style(
// 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_]+]]: tensor<?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>) // CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>)
// CHECK: %[[GENERIC1:.+]]:3 = linalg.generic // CHECK: %[[GENERIC1:.+]]:3 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]], #[[MAP2]]]
// 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(%[[ARG2]], %[[ARG3]], %[[ARG2]] : // CHECK-SAME: inits(%[[ARG2]], %[[ARG3]], %[[ARG2]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: f32
// CHECK-NEXT: %[[S1:.+]] = arith.addf %[[ARG4]], %[[ARG6]] // CHECK-NEXT: %[[S1:.+]] = arith.addf %[[ARG4]], %[[ARG6]]
// CHECK-NEXT: linalg.yield %[[S1]], %{{.+}}, %[[S1]] // CHECK-NEXT: linalg.yield %[[S1]], %{{.+}}, %[[S1]]
// CHECK: %[[GENERIC2:.+]]:2 = linalg.generic // CHECK: %[[GENERIC2:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP2]], #[[MAP3]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC1]]#2 : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]], %[[GENERIC1]]#2 :
// CHECK-SAME: outs(%[[ARG2]], %[[ARG3]] : // CHECK-SAME: inits(%[[ARG2]], %[[ARG3]] :
// CHECK-NEXT: ^bb0( // CHECK-NEXT: ^bb0(
// CHECK-SAME: %[[ARG9:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG9:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG11:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG11:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG12:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG12:[a-zA-Z0-9_]+]]: f32
// CHECK-SAME: %[[ARG13:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG13:[a-zA-Z0-9_]+]]: f32
// CHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG10]], %[[ARG12]] // CHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG10]], %[[ARG12]]
// CHECK-NEXT: linalg.yield %[[ARG11]], %[[S2]] // CHECK-NEXT: linalg.yield %[[ARG11]], %[[S2]]
// CHECK: return %[[GENERIC1]]#0, %[[GENERIC2]]#1 // CHECK: return %[[GENERIC1]]#0, %[[GENERIC2]]#1
// 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)> // CANONICALIZECHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d1)>
// CANONICALIZECHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d1, d0)> // CANONICALIZECHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CANONICALIZECHECK: func.func @destination_passing_style( // CANONICALIZECHECK: func.func @destination_passing_style(
// CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?xf32> // CANONICALIZECHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?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<?x?xf32> // CANONICALIZECHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
// CANONICALIZECHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>) // CANONICALIZECHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: tensor<?x?xf32>)
// CANONICALIZECHECK: %[[GENERIC1:.+]] = linalg.generic // CANONICALIZECHECK: %[[GENERIC1:.+]] = 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(%[[ARG0]] : // CANONICALIZECHECK-SAME: ins(%[[ARG0]] :
// CANONICALIZECHECK-SAME: outs(%[[ARG2]] : // CANONICALIZECHECK-SAME: inits(%[[ARG2]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.addf %[[ARG4]], %[[ARG5]] // CANONICALIZECHECK-NEXT: %[[S1:.+]] = arith.addf %[[ARG4]], %[[ARG5]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[S1]] // CANONICALIZECHECK-NEXT: linalg.yield %[[S1]]
// CANONICALIZECHECK: %[[GENERIC2:.+]]:2 = linalg.generic // CANONICALIZECHECK: %[[GENERIC2:.+]]:2 = linalg.generic
// CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP1]], #[[MAP1]], #[[MAP3]]] // CANONICALIZECHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP1]], #[[MAP1]], #[[MAP3]]]
// CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"] // CANONICALIZECHECK-SAME: iterator_types = ["parallel", "parallel"]
// CANONICALIZECHECK-SAME: ins(%[[ARG1]], %[[GENERIC1]] : // CANONICALIZECHECK-SAME: ins(%[[ARG1]], %[[GENERIC1]] :
// CANONICALIZECHECK-SAME: outs(%[[ARG2]], %[[ARG3]] : // CANONICALIZECHECK-SAME: inits(%[[ARG2]], %[[ARG3]] :
// CANONICALIZECHECK-NEXT: ^bb0( // CANONICALIZECHECK-NEXT: ^bb0(
// CANONICALIZECHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG4:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG6:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: f32 // CANONICALIZECHECK-SAME: %[[ARG7:[a-zA-Z0-9_]+]]: f32
// CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG4]], %[[ARG6]] // CANONICALIZECHECK-NEXT: %[[S2:.+]] = arith.mulf %[[ARG4]], %[[ARG6]]
// CANONICALIZECHECK-NEXT: linalg.yield %[[ARG5]], %[[S2]] // CANONICALIZECHECK-NEXT: linalg.yield %[[ARG5]], %[[S2]]
// CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]]#1 // CANONICALIZECHECK: return %[[GENERIC1]], %[[GENERIC2]]#1

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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>
cf.br ^bb2(%7 : tensor<i32>) cf.br ^bb2(%7 : tensor<i32>)
^bb2(%8: tensor<i32>): ^bb2(%8: tensor<i32>):
return %8 : tensor<i32> return %8 : tensor<i32>
} }
Show All 11 Lines

mlir/test/Dialect/Linalg/detensorize_if.mlir

Show All 12 Linesfunc.func @main() -> (tensor<i32>) 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>
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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>)
^bb3(%10: tensor<i32>): // pred: ^bb1 ^bb3(%10: tensor<i32>): // pred: ^bb1
return %10 : tensor<i32> return %10 : tensor<i32>
Show All 32 Linesfunc.func @main() -> (tensor<i32>) 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>
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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>)
^bb3(%10: tensor<i32>): // pred: ^bb1 ^bb3(%10: tensor<i32>): // pred: ^bb1
cf.br ^bb4(%10 : tensor<i32>) cf.br ^bb4(%10 : tensor<i32>)
Show All 34 Linesfunc.func @main() -> (tensor<i32>) 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>
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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>)
^bb3(%10: tensor<i32>): // pred: ^bb1 ^bb3(%10: tensor<i32>): // pred: ^bb1
return %10 : tensor<i32> return %10 : tensor<i32>
Show All 16 Lines

mlir/test/Dialect/Linalg/detensorize_trivial.mlir

Show All 9 Lines
} }
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 = tensor.empty() : 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>) { inits(%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>
} }
Show All 18 Lines

mlir/test/Dialect/Linalg/detensorize_while.mlir

Show All 9 Lines
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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>)
^bb3(%7: tensor<i32>): // pred: ^bb1 ^bb3(%7: tensor<i32>): // pred: ^bb1
return %7 : tensor<i32> return %7 : tensor<i32>
Show All 33 Lines

mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir

Show All 22 Lines
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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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: tensor.empty() : 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>) inits(%{{.*}} : 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: tensor.empty() : 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>) inits(%{{.*}} : 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>
// DET-ALL: } // DET-ALL: }
// DET-CF-LABEL: func @main // DET-CF-LABEL: func @main
// DET-CF-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>) // DET-CF-SAME: (%{{.*}}: tensor<10xi32>, %{{.*}}: tensor<i32>)
// DET-CF: cf.br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>) // DET-CF: cf.br ^[[bb1:.*]](%{{.*}} : tensor<10xi32>)
// DET-CF: ^bb1(%{{.*}}: tensor<10xi32>) // DET-CF: ^bb1(%{{.*}}: tensor<10xi32>)
// DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) outs(%{{.*}} : tensor<i32>) { // DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<10xi32>) inits(%{{.*}} : tensor<i32>) {
// DET-CF: tensor.extract %{{.*}}[] : tensor<i32> // DET-CF: tensor.extract %{{.*}}[] : tensor<i32>
// DET-CF: cmpi slt, %{{.*}}, %{{.*}} : i32 // DET-CF: cmpi slt, %{{.*}}, %{{.*}} : i32
// DET-CF: cf.cond_br %{{.*}}, ^bb2(%{{.*}} : tensor<i32>), ^bb3(%{{.*}} : tensor<i32>) // DET-CF: cf.cond_br %{{.*}}, ^bb2(%{{.*}} : tensor<i32>), ^bb3(%{{.*}} : tensor<i32>)
// DET-CF: ^bb2(%{{.*}}: tensor<i32>) // DET-CF: ^bb2(%{{.*}}: tensor<i32>)
// DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) outs(%{{.*}} : tensor<10xi32>) { // DET-CF: %{{.*}} = linalg.generic {{{.*}}} ins(%{{.*}} : tensor<i32>) inits(%{{.*}} : tensor<10xi32>) {
// DET-CF: cf.br ^bb1(%{{.*}} : tensor<10xi32>) // DET-CF: cf.br ^bb1(%{{.*}} : tensor<10xi32>)
// DET-CF: ^bb3(%{{.*}}: tensor<i32>) // DET-CF: ^bb3(%{{.*}}: tensor<i32>)
// DET-CF: return %{{.*}} : tensor<i32> // DET-CF: return %{{.*}} : tensor<i32>
// DET-CF: } // DET-CF: }

mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir

Show All 14 Linesfunc.func @main() -> () attributes {} {
%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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>)
^bb3: // pred: ^bb1 ^bb3: // pred: ^bb1
return return
Show All 15 Lines

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

Show All 9 Lines#trait = {
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
indexing_maps = #accesses, indexing_maps = #accesses,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %arg1 : f32, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32> { func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %arg1 : f32, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0, %arg1 : tensor<?x1x?xf32>, f32) ins(%arg0, %arg1 : tensor<?x1x?xf32>, f32)
outs(%shape : tensor<?x1x?x1x?xf32>) { inits(%shape : tensor<?x1x?x1x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) : ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :
linalg.yield %arg3 : f32 linalg.yield %arg3 : f32
} -> tensor<?x1x?x1x?xf32> } -> tensor<?x1x?x1x?xf32>
return %0 : tensor<?x1x?x1x?xf32> return %0 : tensor<?x1x?x1x?xf32>
} }
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(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)>
Show All 17 Lines#trait = {
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_one_trip_loops_indexed func.func @drop_one_trip_loops_indexed
(%arg0 : tensor<?x1x?xi32>, %shape: tensor<?x1x?x1x?xi32>) -> tensor<?x1x?x1x?xi32> (%arg0 : tensor<?x1x?xi32>, %shape: tensor<?x1x?x1x?xi32>) -> tensor<?x1x?x1x?xi32>
{ {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<?x1x?xi32>) ins(%arg0 : tensor<?x1x?xi32>)
outs(%shape: tensor<?x1x?x1x?xi32>) { inits(%shape: tensor<?x1x?x1x?xi32>) {
^bb0(%arg6 : i32, %arg7 : i32) : ^bb0(%arg6 : i32, %arg7 : i32) :
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%idx2 = linalg.index 2 : index %idx2 = linalg.index 2 : index
%idx3 = linalg.index 3 : index %idx3 = linalg.index 3 : index
%idx4 = linalg.index 4 : index %idx4 = linalg.index 4 : index
%1 = arith.addi %idx0, %idx1 : index %1 = arith.addi %idx0, %idx1 : index
%2 = arith.subi %1, %idx2 : index %2 = arith.subi %1, %idx2 : index
Show All 29 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32> func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
{ {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x1xf32>) ins(%arg0 : tensor<1x1xf32>)
outs(%arg0 : tensor<1x1xf32>) { inits(%arg0 : tensor<1x1xf32>) {
^bb0(%arg1: f32, %arg2: f32) : ^bb0(%arg1: f32, %arg2: f32) :
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<1x1xf32> } -> tensor<1x1xf32>
return %0 : tensor<1x1xf32> return %0 : tensor<1x1xf32>
} }
// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()> // CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
// CHECK-LABEL: func @drop_all_loops // CHECK-LABEL: func @drop_all_loops
// CHECK: tensor.collapse_shape %{{.*}} [] // CHECK: tensor.collapse_shape %{{.*}} []
Show All 10 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops_indexed func.func @drop_all_loops_indexed
(%arg0 : tensor<1x1xi32>) -> tensor<1x1xi32>{ (%arg0 : tensor<1x1xi32>) -> tensor<1x1xi32>{
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x1xi32>) ins(%arg0 : tensor<1x1xi32>)
outs(%arg0 : tensor<1x1xi32>) { inits(%arg0 : tensor<1x1xi32>) {
^bb0(%arg3: i32, %arg4: i32) : ^bb0(%arg3: i32, %arg4: i32) :
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%1 = arith.addi %idx0, %idx1 : index %1 = arith.addi %idx0, %idx1 : index
%2 = arith.index_cast %1 : index to i32 %2 = arith.index_cast %1 : index to i32
%3 = arith.addi %2, %arg3 : i32 %3 = arith.addi %2, %arg3 : i32
linalg.yield %3 : i32 linalg.yield %3 : i32
} -> tensor<1x1xi32> } -> tensor<1x1xi32>
Show All 16 Lines#trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
iterator_types = ["parallel"], iterator_types = ["parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> { func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x5xf32>) ins(%arg0 : tensor<1x5xf32>)
outs(%shape : tensor<5xf32>) { inits(%shape : tensor<5xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<5xf32> } -> tensor<5xf32>
return %0 : tensor<5xf32> return %0 : tensor<5xf32>
} }
// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @leading_dim_1_canonicalization // CHECK-LABEL: func @leading_dim_1_canonicalization
Show All 17 Lines
} }
func.func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>, %shape : tensor<5x5xf32>) -> tensor<5x5xf32> func.func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>, %shape : tensor<5x5xf32>) -> tensor<5x5xf32>
{ {
%0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<5xf32> into tensor<1x5xf32> %0 = tensor.expand_shape %arg0 [[0, 1]] : tensor<5xf32> into tensor<1x5xf32>
%1 = tensor.expand_shape %arg1 [[0, 1]] : tensor<5xf32> into tensor<5x1xf32> %1 = tensor.expand_shape %arg1 [[0, 1]] : tensor<5xf32> into tensor<5x1xf32>
%2 = linalg.generic #trait %2 = linalg.generic #trait
ins(%0, %1 : tensor<1x5xf32>, tensor<5x1xf32>) ins(%0, %1 : tensor<1x5xf32>, tensor<5x1xf32>)
outs(%shape : tensor<5x5xf32>) { inits(%shape : tensor<5x5xf32>) {
^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<5x5xf32> } -> tensor<5x5xf32>
return %2 : tensor<5x5xf32> return %2 : tensor<5x5xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
Show All 17 Lines#trait = {
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @broadcast_scalar(%arg0 : tensor<1x1xf32>, %shape : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @broadcast_scalar(%arg0 : tensor<1x1xf32>, %shape : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x1xf32>) ins(%arg0 : tensor<1x1xf32>)
outs(%shape : tensor<?x?xf32>) { inits(%shape : 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>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: func @broadcast_scalar // CHECK-LABEL: func @broadcast_scalar
Show All 9 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) -> (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 = tensor.empty() : 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>) inits(%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_empty_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 = tensor.empty() : 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) inits(%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>)inits(%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_empty_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:.+]] = tensor.empty() : 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) inits(%[[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: inits(%[[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>
// ----- // -----
func.func @fold_slice( func.func @fold_slice(
%arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : tensor<1x?x?x?x?x1x1xf32>, %arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : tensor<1x?x?x?x?x1x1xf32>,
Show All 22 Lines
// ----- // -----
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 = tensor.empty(%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) inits(%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>) { inits(%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:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}inits(%[[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: inits(%[[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 = tensor.empty() : 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) inits(%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>) { inits(%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:.+]] = tensor.empty() : tensor<1xf32> // CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}inits(%[[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: inits(%[[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 = tensor.empty(%0) : 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) inits(%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>) { inits(%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:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32> // CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>
// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}inits(%[[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: inits(%[[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 @slice_unit_dims(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> { func.func @slice_unit_dims(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {
%0 = tensor.extract_slice %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32> %0 = tensor.extract_slice %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32>
return %0 : tensor<1x1xf32> return %0 : tensor<1x1xf32>
Show All 28 Lines#trait = {
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
indexing_maps = #accesses, indexing_maps = #accesses,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_one_trip_loops(%arg0 : memref<?x1x?xf32>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> { func.func @drop_one_trip_loops(%arg0 : memref<?x1x?xf32>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> {
linalg.generic #trait linalg.generic #trait
ins(%arg0, %arg1 : memref<?x1x?xf32>, f32) ins(%arg0, %arg1 : memref<?x1x?xf32>, f32)
outs(%shape : memref<?x1x?x1x?xf32>) { inits(%shape : memref<?x1x?x1x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) : ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :
linalg.yield %arg3 : f32 linalg.yield %arg3 : f32
} }
return %shape : memref<?x1x?x1x?xf32> return %shape : memref<?x1x?x1x?xf32>
} }
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(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)>
Show All 16 Lines#trait = {
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_one_trip_loops_indexed func.func @drop_one_trip_loops_indexed
(%arg0 : memref<?x1x?xi32>, %shape: memref<?x1x?x1x?xi32>) -> memref<?x1x?x1x?xi32> (%arg0 : memref<?x1x?xi32>, %shape: memref<?x1x?x1x?xi32>) -> memref<?x1x?x1x?xi32>
{ {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<?x1x?xi32>) ins(%arg0 : memref<?x1x?xi32>)
outs(%shape: memref<?x1x?x1x?xi32>) { inits(%shape: memref<?x1x?x1x?xi32>) {
^bb0(%arg6 : i32, %arg7 : i32) : ^bb0(%arg6 : i32, %arg7 : i32) :
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%idx2 = linalg.index 2 : index %idx2 = linalg.index 2 : index
%idx3 = linalg.index 3 : index %idx3 = linalg.index 3 : index
%idx4 = linalg.index 4 : index %idx4 = linalg.index 4 : index
%1 = arith.addi %idx0, %idx1 : index %1 = arith.addi %idx0, %idx1 : index
%2 = arith.subi %1, %idx2 : index %2 = arith.subi %1, %idx2 : index
Show All 29 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops(%arg0 : memref<1x1xf32>) -> memref<1x1xf32> func.func @drop_all_loops(%arg0 : memref<1x1xf32>) -> memref<1x1xf32>
{ {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x1xf32>) ins(%arg0 : memref<1x1xf32>)
outs(%arg0 : memref<1x1xf32>) { inits(%arg0 : memref<1x1xf32>) {
^bb0(%arg1: f32, %arg2: f32) : ^bb0(%arg1: f32, %arg2: f32) :
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} }
return %arg0 : memref<1x1xf32> return %arg0 : memref<1x1xf32>
} }
// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()> // CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>
// CHECK-LABEL: func @drop_all_loops // CHECK-LABEL: func @drop_all_loops
// CHECK: memref.collapse_shape %{{.*}} [] // CHECK: memref.collapse_shape %{{.*}} []
Show All 10 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops_indexed func.func @drop_all_loops_indexed
(%arg0 : memref<1x1xi32>) -> memref<1x1xi32>{ (%arg0 : memref<1x1xi32>) -> memref<1x1xi32>{
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x1xi32>) ins(%arg0 : memref<1x1xi32>)
outs(%arg0 : memref<1x1xi32>) { inits(%arg0 : memref<1x1xi32>) {
^bb0(%arg3: i32, %arg4: i32) : ^bb0(%arg3: i32, %arg4: i32) :
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%1 = arith.addi %idx0, %idx1 : index %1 = arith.addi %idx0, %idx1 : index
%2 = arith.index_cast %1 : index to i32 %2 = arith.index_cast %1 : index to i32
%3 = arith.addi %2, %arg3 : i32 %3 = arith.addi %2, %arg3 : i32
linalg.yield %3 : i32 linalg.yield %3 : i32
} }
Show All 16 Lines#trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
iterator_types = ["parallel"], iterator_types = ["parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @leading_dim_1_canonicalization(%arg0: memref<1x5xf32>, %shape: memref<5xf32>) -> memref<5xf32> { func.func @leading_dim_1_canonicalization(%arg0: memref<1x5xf32>, %shape: memref<5xf32>) -> memref<5xf32> {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x5xf32>) ins(%arg0 : memref<1x5xf32>)
outs(%shape : memref<5xf32>) { inits(%shape : memref<5xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} }
return %shape : memref<5xf32> return %shape : memref<5xf32>
} }
// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)> // CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @leading_dim_1_canonicalization // CHECK-LABEL: func @leading_dim_1_canonicalization
Show All 17 Lines
} }
func.func @broadcast_test(%arg0 : memref<5xf32>, %arg1 : memref<5xf32>, %shape : memref<5x5xf32>) -> memref<5x5xf32> func.func @broadcast_test(%arg0 : memref<5xf32>, %arg1 : memref<5xf32>, %shape : memref<5x5xf32>) -> memref<5x5xf32>
{ {
%0 = memref.expand_shape %arg0 [[0, 1]] : memref<5xf32> into memref<1x5xf32> %0 = memref.expand_shape %arg0 [[0, 1]] : memref<5xf32> into memref<1x5xf32>
%1 = memref.expand_shape %arg1 [[0, 1]] : memref<5xf32> into memref<5x1xf32> %1 = memref.expand_shape %arg1 [[0, 1]] : memref<5xf32> into memref<5x1xf32>
linalg.generic #trait linalg.generic #trait
ins(%0, %1 : memref<1x5xf32>, memref<5x1xf32>) ins(%0, %1 : memref<1x5xf32>, memref<5x1xf32>)
outs(%shape : memref<5x5xf32>) { inits(%shape : memref<5x5xf32>) {
^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
} }
return %shape : memref<5x5xf32> return %shape : memref<5x5xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>
Show All 17 Lines#trait = {
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @broadcast_scalar(%arg0 : memref<1x1xf32>, %shape : memref<?x?xf32>) -> memref<?x?xf32> func.func @broadcast_scalar(%arg0 : memref<1x1xf32>, %shape : memref<?x?xf32>) -> memref<?x?xf32>
{ {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x1xf32>) ins(%arg0 : memref<1x1xf32>)
outs(%shape : memref<?x?xf32>) { inits(%shape : memref<?x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} }
return %shape : memref<?x?xf32> return %shape : memref<?x?xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: func @broadcast_scalar // CHECK-LABEL: func @broadcast_scalar
Show All 9 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) -> (d2)> #map1 = affine_map<(d0, d1, d2) -> (d2)>
func.func @fold_unit_dim_memref_reshape_op(%arg0 : memref<5xf32>) -> memref<2x5xf32> func.func @fold_unit_dim_memref_reshape_op(%arg0 : memref<5xf32>) -> memref<2x5xf32>
{ {
%1 = memref.alloc() : memref<1x2x5xf32> %1 = memref.alloc() : memref<1x2x5xf32>
linalg.generic {i64, indexing_maps = [#map1, #map0], linalg.generic {i64, indexing_maps = [#map1, #map0],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0 : memref<5xf32>) outs(%1 : memref<1x2x5xf32>) { ins(%arg0 : memref<5xf32>) inits(%1 : memref<1x2x5xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} }
%3 = memref.collapse_shape %1 [[0, 1], [2]] %3 = memref.collapse_shape %1 [[0, 1], [2]]
: memref<1x2x5xf32> into memref<2x5xf32> : memref<1x2x5xf32> into memref<2x5xf32>
return %3 : memref<2x5xf32> return %3 : memref<2x5xf32>
} }
// CHECK-LABEL: func @fold_unit_dim_memref_reshape_op // CHECK-LABEL: func @fold_unit_dim_memref_reshape_op
// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1x2x5xf32> // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1x2x5xf32>
// CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ALLOC]] // CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ALLOC]]
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: outs(%[[OUT:.*]] : // CHECK-SAME: inits(%[[OUT:.*]] :
// CHECK: %[[RESULT:.*]] = memref.collapse_shape %[[ALLOC]] // CHECK: %[[RESULT:.*]] = memref.collapse_shape %[[ALLOC]]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
func.func @fold_unit_dim_for_init_memref(%input: memref<1x1000xf32>) -> memref<1xf32> { func.func @fold_unit_dim_for_init_memref(%input: memref<1x1000xf32>) -> memref<1xf32> {
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%init = memref.alloc() : memref<1xf32> %init = memref.alloc() : memref<1xf32>
linalg.generic { 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 : memref<1x1000xf32>)outs(%init : memref<1xf32>) { ins(%input : memref<1x1000xf32>)inits(%init : memref<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
} }
return %init : memref<1xf32> return %init : memref<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_memref // CHECK: func @fold_unit_dim_for_init_memref
// CHECK: %[[INIT:.+]] = memref.alloc() : memref<1xf32> // CHECK: %[[INIT:.+]] = memref.alloc() : memref<1xf32>
// CHECK: %[[INPUT_RESHAPE:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0, 1]] : memref<1x1000xf32> into memref<1000xf32> // CHECK: %[[INPUT_RESHAPE:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0, 1]] : memref<1x1000xf32> into memref<1000xf32>
// CHECK: %[[INIT_RESHAPE:.+]] = memref.collapse_shape %[[INIT]] [] : memref<1xf32> into memref<f32> // CHECK: %[[INIT_RESHAPE:.+]] = memref.collapse_shape %[[INIT]] [] : memref<1xf32> into memref<f32>
// CHECK: linalg.generic // CHECK: 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]] : memref<1000xf32>) // CHECK-SAME: ins(%[[INPUT_RESHAPE]] : memref<1000xf32>)
// CHECK-SAME: outs(%[[INIT_RESHAPE]] : memref<f32>) // CHECK-SAME: inits(%[[INIT_RESHAPE]] : memref<f32>)
// CHECK: return %[[INIT:.+]] : memref<1xf32> // CHECK: return %[[INIT:.+]] : memref<1xf32>
// ----- // -----
// Test that nothing changes and no assertions are fired for memrefs with affine // Test that nothing changes and no assertions are fired for memrefs with affine
// maps while still changing the other operations. // maps while still changing the other operations.
#accesses = [ #accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>, affine_map<(i, j, k, l, m) -> (i, k, m)>,
affine_map<(i, j, k, l, m) -> ()>, affine_map<(i, j, k, l, m) -> ()>,
affine_map<(i, j, k, l, m) -> (i, k, j, l, m)> affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>
] ]
#trait = { #trait = {
iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],
indexing_maps = #accesses, indexing_maps = #accesses,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @input_stays_same(%arg0 : memref<?x1x?xf32, strided<[?, 1, 1]>>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> { func.func @input_stays_same(%arg0 : memref<?x1x?xf32, strided<[?, 1, 1]>>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> {
linalg.generic #trait linalg.generic #trait
ins(%arg0, %arg1 : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32) ins(%arg0, %arg1 : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32)
outs(%shape : memref<?x1x?x1x?xf32>) { inits(%shape : memref<?x1x?x1x?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) : ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :
linalg.yield %arg3 : f32 linalg.yield %arg3 : f32
} }
return %shape : memref<?x1x?x1x?xf32> return %shape : memref<?x1x?x1x?xf32>
} }
// CHECK: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)> // CHECK: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
// CHECK: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()> // CHECK: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
// CHECK: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK: func @input_stays_same( // CHECK: func @input_stays_same(
// CHECK-SAME: %[[ARG0:.*]]: memref<?x1x?xf32, strided<[?, 1, 1]>>, // CHECK-SAME: %[[ARG0:.*]]: memref<?x1x?xf32, strided<[?, 1, 1]>>,
// CHECK-SAME: %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<?x1x?x1x?xf32>) // CHECK-SAME: %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<?x1x?x1x?xf32>)
// CHECK-SAME -> memref<?x1x?x1x?xf32> { // CHECK-SAME -> memref<?x1x?x1x?xf32> {
// CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ARG2]] {{\[}}[0, 1], [2, 3], [4]] // CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ARG2]] {{\[}}[0, 1], [2, 3], [4]]
// CHECK-SAME: : memref<?x1x?x1x?xf32> into memref<?x?x?xf32> // CHECK-SAME: : memref<?x1x?x1x?xf32> into memref<?x?x?xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]], // CHECK-SAME: {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32) // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32)
// CHECK-SAME: outs(%[[OUT]] : memref<?x?x?xf32>) { // CHECK-SAME: inits(%[[OUT]] : memref<?x?x?xf32>) {
// CHECK: ^bb0(%{{.*}}: f32, %[[ARG:.*]]: f32, %{{.*}}: f32): // CHECK: ^bb0(%{{.*}}: f32, %[[ARG:.*]]: f32, %{{.*}}: f32):
// CHECK: linalg.yield %[[ARG]] : f32 // CHECK: linalg.yield %[[ARG]] : f32
// CHECK: } // CHECK: }
// CHECK: return %[[ARG2]] : memref<?x1x?x1x?xf32> // CHECK: return %[[ARG2]] : memref<?x1x?x1x?xf32>
// ----- // -----
// Negative test for case with tensor encoding. // Negative test for case with tensor encoding.
#matvec = { #matvec = {
indexing_maps = [ indexing_maps = [
affine_map<(i,j) -> (i,j)>, // A affine_map<(i,j) -> (i,j)>, // A
affine_map<(i,j) -> (j)>, // b affine_map<(i,j) -> (j)>, // b
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 = tensor.empty() : 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>) { inits(%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: tensor.empty // 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 = tensor.empty() : 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 = tensor.empty() : 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) inits(%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>
} }
} }
return %res: tensor<4x2xf32> return %res: tensor<4x2xf32>
} }
// ----- // -----
#map0 = affine_map<(i, j) -> (i, j)> #map0 = affine_map<(i, j) -> (i, j)>
#access = [#map0, #map0] #access = [#map0, #map0]
#trait = { #trait = {
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops(%arg0 : memref<1x1xf32, 3>) -> memref<1x1xf32, 3> func.func @drop_all_loops(%arg0 : memref<1x1xf32, 3>) -> memref<1x1xf32, 3>
{ {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x1xf32, 3>) ins(%arg0 : memref<1x1xf32, 3>)
outs(%arg0 : memref<1x1xf32, 3>) { inits(%arg0 : memref<1x1xf32, 3>) {
^bb0(%arg1: f32, %arg2: f32) : ^bb0(%arg1: f32, %arg2: f32) :
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} }
return %arg0 : memref<1x1xf32, 3> return %arg0 : memref<1x1xf32, 3>
} }
// CHECK-LABEL: func @drop_all_loops // CHECK-LABEL: func @drop_all_loops
// CHECK: memref.collapse_shape // CHECK: memref.collapse_shape
// CHECK-SAME: [] : memref<1x1xf32, 3> into memref<f32, 3> // CHECK-SAME: [] : memref<1x1xf32, 3> into memref<f32, 3>
// CHECK: linalg.generic{{.*}}memref<f32, 3> // CHECK: linalg.generic{{.*}}memref<f32, 3>

mlir/test/Dialect/Linalg/fill-interface-invalid.mlir

Show All 12 Lines
#map0 = affine_map<(d0) -> ()> #map0 = affine_map<(d0) -> ()>
#map1 = affine_map<(d0) -> (d0)> #map1 = affine_map<(d0) -> (d0)>
func.func @test_fill_op_wrong_num_operands(%arg0 : f32, %arg1 : tensor<?xf32>) func.func @test_fill_op_wrong_num_operands(%arg0 : f32, %arg1 : tensor<?xf32>)
-> tensor<?xf32> { -> tensor<?xf32> {
// expected-error @+1 {{expected op with 1 input and 1 output}} // expected-error @+1 {{expected op with 1 input and 1 output}}
%0 = test.linalg_fill_op { %0 = test.linalg_fill_op {
indexing_maps = [#map0, #map0, #map1], indexing_maps = [#map0, #map0, #map1],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0, %arg0 : f32, f32) outs(%arg1 : tensor<?xf32>) { ins(%arg0, %arg0 : f32, f32) inits(%arg1 : tensor<?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32): ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }
// ----- // -----
#map1 = affine_map<(d0) -> (d0)> #map1 = affine_map<(d0) -> (d0)>
func.func @test_fill_op_non_scalar_input(%arg0 : tensor<?xf32>, func.func @test_fill_op_non_scalar_input(%arg0 : tensor<?xf32>,
%arg1 : tensor<?xf32>) -> tensor<?xf32> { %arg1 : tensor<?xf32>) -> tensor<?xf32> {
// expected-error @+1 {{expected op with scalar input}} // expected-error @+1 {{expected op with scalar input}}
%0 = test.linalg_fill_op { %0 = test.linalg_fill_op {
indexing_maps = [#map1, #map1], indexing_maps = [#map1, #map1],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<?xf32>) { ins(%arg0 : tensor<?xf32>) inits(%arg1 : tensor<?xf32>) {
^bb0(%arg2 : f32, %arg3 : f32): ^bb0(%arg2 : f32, %arg3 : f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }

mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir

Show All 9 Lines#trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32> func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32>
{ {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<?x1x?xf32>) ins(%arg0 : tensor<?x1x?xf32>)
outs(%shape : tensor<?x1x?x1x?xf32>) { inits(%shape : tensor<?x1x?x1x?xf32>) {
^bb0(%arg1 : f32, %arg2 : f32) : ^bb0(%arg1 : f32, %arg2 : f32) :
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<?x1x?x1x?xf32> } -> tensor<?x1x?x1x?xf32>
return %0 : tensor<?x1x?x1x?xf32> return %0 : tensor<?x1x?x1x?xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d1, 0, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d1, 0, d2)>
// CHECK-LABEL: func @drop_one_trip_loops // CHECK-LABEL: func @drop_one_trip_loops
Show All 10 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32> func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>
{ {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x1xf32>) ins(%arg0 : tensor<1x1xf32>)
outs(%arg0 : tensor<1x1xf32>) { inits(%arg0 : tensor<1x1xf32>) {
^bb0(%arg1: f32, %arg2: f32) : ^bb0(%arg1: f32, %arg2: f32) :
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<1x1xf32> } -> tensor<1x1xf32>
return %0 : tensor<1x1xf32> return %0 : tensor<1x1xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
// CHECK-LABEL: func @drop_all_loops // CHECK-LABEL: func @drop_all_loops
// CHECK: linalg.generic // CHECK: linalg.generic
Show All 9 Lines#trait = {
indexing_maps = #access, indexing_maps = #access,
library_call = "some_external_func" library_call = "some_external_func"
} }
func.func @drop_all_loops(%arg0 : memref<1x1xf32>, %arg1 : memref<1x1xf32>) func.func @drop_all_loops(%arg0 : memref<1x1xf32>, %arg1 : memref<1x1xf32>)
{ {
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<1x1xf32>) ins(%arg0 : memref<1x1xf32>)
outs(%arg1 : memref<1x1xf32>) { inits(%arg1 : memref<1x1xf32>) {
^bb0(%arg2: f32, %arg3 : f32) : ^bb0(%arg2: f32, %arg3 : f32) :
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} }
return return
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<() -> (0, 0)>
// CHECK-LABEL: func @drop_all_loops // CHECK-LABEL: func @drop_all_loops
// CHECK: linalg.generic // CHECK: linalg.generic
Show All 11 Lines#trait = {
indexing_maps = #accesses, indexing_maps = #accesses,
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> { func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arg0 : tensor<1x5xf32>) ins(%arg0 : tensor<1x5xf32>)
outs(%shape : tensor<5xf32>) { inits(%shape : tensor<5xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
linalg.yield %arg2 : f32 linalg.yield %arg2 : f32
} -> tensor<5xf32> } -> tensor<5xf32>
return %0 : tensor<5xf32> return %0 : tensor<5xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (0, d0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0) -> (0, d0)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
// CHECK-LABEL: func @leading_dim_1_canonicalization // CHECK-LABEL: func @leading_dim_1_canonicalization
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]] // CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
// CHECK-SAME: iterator_types = ["parallel"] // CHECK-SAME: iterator_types = ["parallel"]

mlir/test/Dialect/Linalg/forward-vector-transfers.mlir

Show All 23 Lines
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// CHECK: vector.transfer_read %[[ARG0]] // CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: in_bounds // CHECK-NOT: in_bounds
func.func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> { func.func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = arith.constant 0: index %c0 = arith.constant 0: index
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
%alloc = memref.alloc() : memref<32 x f32> %alloc = memref.alloc() : memref<32 x f32>
linalg.fill ins(%f0 : f32) outs(%alloc : memref<32 x f32>) linalg.fill ins(%f0 : f32) inits(%alloc : memref<32 x f32>)
%subview = memref.subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32> %subview = memref.subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
memref.copy %in, %subview : memref<? x f32> to memref<16 x f32> memref.copy %in, %subview : memref<? x f32> to memref<16 x f32>
%0 = vector.transfer_read %alloc[%c0], %f0 {in_bounds = [true]} : memref<32 x f32>, vector<32 x f32> %0 = vector.transfer_read %alloc[%c0], %f0 {in_bounds = [true]} : memref<32 x f32>, vector<32 x f32>
memref.dealloc %alloc : memref<32 x f32> memref.dealloc %alloc : memref<32 x f32>
return %0: vector<32 x f32> return %0: vector<32 x f32>
} }
// CHECK-LABEL: testViewRead // CHECK-LABEL: testViewRead
Show All 23 Lines
// CHECK: vector.transfer_read %[[ARG0]] // CHECK: vector.transfer_read %[[ARG0]]
// CHECK-NOT: in_bounds // CHECK-NOT: in_bounds
func.func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> { func.func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = arith.constant 0: index %c0 = arith.constant 0: index
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
%alloc = memref.alloc() : memref<128 x i8> %alloc = memref.alloc() : memref<128 x i8>
%view = memref.view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32> %view = memref.view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
%subview = memref.subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32> %subview = memref.subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
linalg.fill ins(%f0 : f32) outs(%view : memref<32 x f32>) linalg.fill ins(%f0 : f32) inits(%view : memref<32 x f32>)
memref.copy %in, %subview : memref<? x f32> to memref<16 x f32> memref.copy %in, %subview : memref<? x f32> to memref<16 x f32>
%0 = vector.transfer_read %view[%c0], %f0 {in_bounds = [true]} : memref<32 x f32>, vector<32 x f32> %0 = vector.transfer_read %view[%c0], %f0 {in_bounds = [true]} : memref<32 x f32>, vector<32 x f32>
memref.dealloc %alloc : memref<128 x i8> memref.dealloc %alloc : memref<128 x i8>
return %0: vector<32 x f32> return %0: vector<32 x f32>
} }
// CHECK-LABEL: testAllocWrite // CHECK-LABEL: testAllocWrite
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector // CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// CHECK: memref.copy // CHECK: memref.copy
// CHECK: vector.transfer_read %[[ALLOC]] // CHECK: vector.transfer_read %[[ALLOC]]
func.func @failAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> { func.func @failAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
%c0 = arith.constant 0: index %c0 = arith.constant 0: index
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
%f1 = arith.constant 1.0: f32 %f1 = arith.constant 1.0: f32
%alloc = memref.alloc() : memref<32 x f32> %alloc = memref.alloc() : memref<32 x f32>
linalg.fill ins(%f0 : f32) outs(%alloc : memref<32 x f32>) linalg.fill ins(%f0 : f32) inits(%alloc : memref<32 x f32>)
%subview = memref.subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32> %subview = memref.subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
memref.copy %in, %subview : memref<? x f32> to memref<16 x f32> memref.copy %in, %subview : memref<? x f32> to memref<16 x f32>
"some_interleaved_use"(%subview) : (memref<16 x f32>) -> () "some_interleaved_use"(%subview) : (memref<16 x f32>) -> ()
%0 = vector.transfer_read %alloc[%c0], %f1: memref<32 x f32>, vector<32 x f32> %0 = vector.transfer_read %alloc[%c0], %f1: memref<32 x f32>, vector<32 x f32>
memref.dealloc %alloc : memref<32 x f32> memref.dealloc %alloc : memref<32 x f32>
return %0: vector<32 x f32> return %0: vector<32 x f32>
} }
Show All 19 Lines

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

Show All 10 Linesfunc.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 = tensor.empty() : 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>) { inits(%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:.+]] = tensor.empty() // 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: inits(%[[INIT_RESHAPE]] :
// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}} // CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
// CONTROL: func @fuse_by_collapsing( // CONTROL: func @fuse_by_collapsing(
// CONTROL-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32> // CONTROL-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>
// CONTROL-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32> // CONTROL-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32>
// CONTROL-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32> // CONTROL-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32>
// CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] // CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]
Show All 11 Linesfunc.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 = tensor.empty() : 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>) { inits(%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
%t0 = arith.addi %iv0, %iv1 : index %t0 = arith.addi %iv0, %iv1 : index
%iv2 = linalg.index 2 : index %iv2 = linalg.index 2 : index
%t1 = arith.addi %t0, %iv2 : index %t1 = arith.addi %t0, %iv2 : index
%iv3 = linalg.index 3 : index %iv3 = linalg.index 3 : index
%t2 = arith.addi %t1, %iv3 : index %t2 = arith.addi %t1, %iv3 : index
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Linesfunc.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 = tensor.empty() : 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>) { inits(%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:.+]] = tensor.empty() // 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: inits(%[[INIT_RESHAPE]] :
// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}} // CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}}
// CHECK: return %[[RESULT_RESHAPE]] // CHECK: return %[[RESULT_RESHAPE]]
// ----- // -----
// Dynamic case. Only checks things not covered by `fuse_by_collapsing` test above. // Dynamic case. Only checks things not covered by `fuse_by_collapsing` test above.
#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)>
Show All 11 Linesfunc.func @fuse_by_collapsing_dynamic(%arg0 : tensor<?x?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 = tensor.empty(%d0, %d2, %d4, %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>) { inits(%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
%t0 = arith.addi %iv0, %iv1 : index %t0 = arith.addi %iv0, %iv1 : index
%iv2 = linalg.index 2 : index %iv2 = linalg.index 2 : index
%t1 = arith.addi %t0, %iv2 : index %t1 = arith.addi %t0, %iv2 : index
%iv3 = linalg.index 3 : index %iv3 = linalg.index 3 : index
%t2 = arith.addi %t1, %iv3 : index %t2 = arith.addi %t1, %iv3 : index
Show All 31 Lines
#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, d3)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3)>
func.func @fuse_reductions(%arg0 : tensor<2x?x5xf32>, %arg1 : tensor<2x5xf32>) -> tensor<2x5xf32> { func.func @fuse_reductions(%arg0 : tensor<2x?x5xf32>, %arg1 : tensor<2x5xf32>) -> tensor<2x5xf32> {
%0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x?x5xf32> into tensor<2x6x?x5xf32> %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] : tensor<2x?x5xf32> into tensor<2x6x?x5xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1], indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "reduction", "reduction", "parallel"]} iterator_types = ["parallel", "reduction", "reduction", "parallel"]}
ins(%0 : tensor<2x6x?x5xf32>) outs(%arg1 : tensor<2x5xf32>) { ins(%0 : tensor<2x6x?x5xf32>) inits(%arg1 : tensor<2x5xf32>) {
^bb0(%b0 : f32, %b1 : f32): ^bb0(%b0 : f32, %b1 : 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>
return %1 : tensor<2x5xf32> return %1 : tensor<2x5xf32>
} }
// 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: func @fuse_reductions( // CHECK: func @fuse_reductions(
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x?x5xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x?x5xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<2x5xf32>) -> tensor<2x5xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<2x5xf32>) -> tensor<2x5xf32>
// 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", "reduction", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"]
// CHECK-SAME: ins(%[[ARG0]] : tensor<2x?x5xf32>) // CHECK-SAME: ins(%[[ARG0]] : tensor<2x?x5xf32>)
// CHECK-SAME: outs(%[[ARG1]] : tensor<2x5xf32>) // CHECK-SAME: inits(%[[ARG1]] : tensor<2x5xf32>)
// ----- // -----
// 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 = tensor.empty(): 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>) inits(%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>
return %1 : tensor<2x3x4x5xf32> return %1 : tensor<2x3x4x5xf32>
} }
// CHECK: func @no_fuse_unpreserved_folding // CHECK: func @no_fuse_unpreserved_folding
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>
Show All 10 Lines
#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 = tensor.empty() : 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>) inits(%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>
return %1 : tensor<2x4x3x5xf32> return %1 : tensor<2x4x3x5xf32>
} }
// CHECK: func @no_fuse_unpreserved_folding_transpose // CHECK: func @no_fuse_unpreserved_folding_transpose
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>
Show All 10 Lines
#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 = tensor.empty() : 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>) inits(%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>
return %1 : tensor<2x5xf32> return %1 : tensor<2x5xf32>
} }
// CHECK: func @no_fuse_mismatched_iterator_types // CHECK: func @no_fuse_mismatched_iterator_types
// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>
Show All 12 Lines
#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 = tensor.empty() : 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>) inits(%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>
return %2 : tensor<2x3x4x5xf32> return %2 : tensor<2x3x4x5xf32>
} }
// 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-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: func @control_fusion( // CHECK: func @control_fusion(
// CHECK-SAME: %[[ARG0:.+]]: tensor<6xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<6xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<20xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<20xf32>
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// 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(%{{.+}}: tensor<6x20xf32>) // CHECK-SAME: inits(%{{.+}}: tensor<6x20xf32>)
// CHECK: %[[RESHAPE1:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]{{\]}} // CHECK: %[[RESHAPE1:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]{{\]}}
// CHECK: %[[RESHAPE2:.+]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1], [2], [3]{{\]}} // CHECK: %[[RESHAPE2:.+]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1], [2], [3]{{\]}}
// CHECK: return %[[RESHAPE2]] // CHECK: return %[[RESHAPE2]]
// 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:.+]] = tensor.empty() // 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: inits(%[[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 = tensor.empty() : 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>) inits(%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>
} }
// CHECK: func @zero_D_test // CHECK: func @zero_D_test
// CHECK-SAME: %[[ARG0:.+]]: tensor<f32> // CHECK-SAME: %[[ARG0:.+]]: tensor<f32>
// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] // CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: ins(%[[EXPAND]] : // CHECK-SAME: ins(%[[EXPAND]] :
// CHECK: return %[[GENERIC]] // CHECK: return %[[GENERIC]]
// ----- // -----
#map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d2, d3)> #map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
func.func @fuse_only_one_reassociation(%arg0 : tensor<?x?xf32>, %arg1 : tensor<4x?x?x8xf32>) -> tensor<4x?x?x8xf32> { func.func @fuse_only_one_reassociation(%arg0 : tensor<?x?xf32>, %arg1 : tensor<4x?x?x8xf32>) -> tensor<4x?x?x8xf32> {
%0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xf32> into tensor<?x4x?x8xf32> %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] : tensor<?x?xf32> into tensor<?x4x?x8xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map1], indexing_maps = [#map0, #map1, #map1],
iterator_types = ["parallel", "parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel", "parallel"]}
ins(%0, %arg1 : tensor<?x4x?x8xf32>, tensor<4x?x?x8xf32>) ins(%0, %arg1 : tensor<?x4x?x8xf32>, tensor<4x?x?x8xf32>)
outs(%arg1 : tensor<4x?x?x8xf32>) { inits(%arg1 : tensor<4x?x?x8xf32>) {
^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<4x?x?x8xf32> } -> tensor<4x?x?x8xf32>
return %1 : tensor<4x?x?x8xf32> return %1 : tensor<4x?x?x8xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK: func @fuse_only_one_reassociation( // CHECK: func @fuse_only_one_reassociation(
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>
// CHECK-SAME: %[[ARG1:.+]]: tensor<4x?x?x8xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<4x?x?x8xf32>
// CHECK-DAG: %[[EXPAND_ARG0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]{{\]}} // CHECK-DAG: %[[EXPAND_ARG0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]{{\]}}
// CHECK-DAG: %[[COLLAPSE_ARG0:.+]] = tensor.collapse_shape %[[EXPAND_ARG0]] {{\[}}[0], [1], [2, 3]{{\]}} // CHECK-DAG: %[[COLLAPSE_ARG0:.+]] = tensor.collapse_shape %[[EXPAND_ARG0]] {{\[}}[0], [1], [2, 3]{{\]}}
// CHECK-DAG: %[[COLLAPSE_ARG1_0:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}} // CHECK-DAG: %[[COLLAPSE_ARG1_0:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}
// CHECK-DAG: %[[COLLAPSE_ARG1_1:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}} // CHECK-DAG: %[[COLLAPSE_ARG1_1:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}
// CHECK: %[[GENERIC:.+]] = linalg.generic // CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP1]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP1]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[COLLAPSE_ARG0]], %[[COLLAPSE_ARG1_0]] : // CHECK-SAME: ins(%[[COLLAPSE_ARG0]], %[[COLLAPSE_ARG1_0]] :
// CHECK-SAME: outs(%[[COLLAPSE_ARG1_1]] : // CHECK-SAME: inits(%[[COLLAPSE_ARG1_1]] :
// CHECK: %[[EXPAND_GENERIC:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1], [2, 3]{{\]}} // CHECK: %[[EXPAND_GENERIC:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1], [2, 3]{{\]}}
// CHECK: return %[[EXPAND_GENERIC]] // CHECK: return %[[EXPAND_GENERIC]]
// ----- // -----
#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 = tensor.empty(%d1, %d0) : 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>) inits(%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
Show All 9 Lines
// 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:.+]] = tensor.empty // 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: inits(%[[COLLAPSE_INIT]] :
// CHECK-NEXT: ^bb{{[0-9]}} // CHECK-NEXT: ^bb{{[0-9]}}
// CHECK: %[[ID0:.+]] = linalg.index 0 // CHECK: %[[ID0:.+]] = linalg.index 0
// CHECK-DAG: %[[T0:.+]] = arith.remui %[[ID0]], %[[C4]] // CHECK-DAG: %[[T0:.+]] = arith.remui %[[ID0]], %[[C4]]
// CHECK-DAG: %[[T1:.+]] = arith.divui %[[ID0]], %[[C4]] // CHECK-DAG: %[[T1:.+]] = arith.divui %[[ID0]], %[[C4]]
// CHECK: %[[ID1:.+]] = linalg.index 1 // CHECK: %[[ID1:.+]] = linalg.index 1
// CHECK-DAG: %[[T2:.+]] = arith.remui %[[ID1]], %[[C8]] // CHECK-DAG: %[[T2:.+]] = arith.remui %[[ID1]], %[[C8]]
// CHECK-DAG: %[[T3:.+]] = arith.divui %[[ID1]], %[[C8]] // CHECK-DAG: %[[T3:.+]] = arith.divui %[[ID1]], %[[C8]]
// CHECK-DAG: %[[T4:.+]] = arith.addi %[[T1]], %[[T2]] // CHECK-DAG: %[[T4:.+]] = arith.addi %[[T1]], %[[T2]]
Show All 13 Lines
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 = tensor.empty() : 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>) inits(%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
%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
Show All 11 Lines

mlir/test/Dialect/Linalg/fusion-2-level.mlir

// RUN: mlir-opt %s -test-linalg-greedy-fusion | FileCheck %s // RUN: mlir-opt %s -test-linalg-greedy-fusion | FileCheck %s
func.func @f1(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, %C: memref<?x?xf32, strided<[?, 1], offset: ?>>, %D: memref<?x?xf32, strided<[?, 1], offset: ?>>, %E: memref<?x?xf32, strided<[?, 1], offset: ?>>) -> memref<?x?xf32, strided<[?, 1], offset: ?>> { func.func @f1(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, %C: memref<?x?xf32, strided<[?, 1], offset: ?>>, %D: memref<?x?xf32, strided<[?, 1], offset: ?>>, %E: memref<?x?xf32, strided<[?, 1], offset: ?>>) -> memref<?x?xf32, strided<[?, 1], offset: ?>> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c40 = arith.constant 40 : index %c40 = arith.constant 40 : index
%c30 = arith.constant 30 : index %c30 = arith.constant 30 : index
%c20 = arith.constant 20 : index %c20 = arith.constant 20 : index
%0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, 1], offset: ?>> %0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, 1], offset: ?>>
%1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, 1], offset: ?>> %1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, 1], offset: ?>>
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, 1], offset: ?>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, 1], offset: ?>>
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>) linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) inits(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>)
scf.for %arg5 = %c0 to %0 step %c20 { scf.for %arg5 = %c0 to %0 step %c20 {
scf.for %arg6 = %c0 to %2 step %c30 { scf.for %arg6 = %c0 to %2 step %c30 {
scf.for %arg7 = %c0 to %1 step %c40 { scf.for %arg7 = %c0 to %1 step %c40 {
%5 = memref.subview %C[%arg5, %arg7][%c20, %c40][%c1, %c1] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %5 = memref.subview %C[%arg5, %arg7][%c20, %c40][%c1, %c1] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %D[%arg7, %arg6][%c40, %c30][%c1, %c1]: memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %7 = memref.subview %D[%arg7, %arg6][%c40, %c30][%c1, %c1]: memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c20, %c40][%c1, %c1] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %8 = memref.subview %E[%arg5, %arg6][%c20, %c40][%c1, %c1] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%9 = memref.dim %5, %c0 : memref<?x?xf32, strided<[?, ?], offset: ?>> %9 = memref.dim %5, %c0 : memref<?x?xf32, strided<[?, ?], offset: ?>>
%10 = memref.dim %5, %c1 : memref<?x?xf32, strided<[?, ?], offset: ?>> %10 = memref.dim %5, %c1 : memref<?x?xf32, strided<[?, ?], offset: ?>>
%11 = memref.dim %7, %c1 : memref<?x?xf32, strided<[?, ?], offset: ?>> %11 = memref.dim %7, %c1 : memref<?x?xf32, strided<[?, ?], offset: ?>>
scf.for %arg8 = %c0 to %9 step %c2 { scf.for %arg8 = %c0 to %9 step %c2 {
scf.for %arg9 = %c0 to %11 step %c3 { scf.for %arg9 = %c0 to %11 step %c3 {
scf.for %arg10 = %c0 to %10 step %c4 { scf.for %arg10 = %c0 to %10 step %c4 {
%14 = memref.subview %5[%arg8, %arg10][%c2, %c4][%c1, %c1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %14 = memref.subview %5[%arg8, %arg10][%c2, %c4][%c1, %c1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%16 = memref.subview %7[%arg10, %arg9][%c4, %c3][%c1, %c1]: memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %16 = memref.subview %7[%arg10, %arg9][%c4, %c3][%c1, %c1]: memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%17 = memref.subview %8[%arg8, %arg9][%c2, %c3][%c1, %c1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> %17 = memref.subview %8[%arg8, %arg9][%c2, %c3][%c1, %c1] : memref<?x?xf32, strided<[?, ?], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%14, %16: memref<?x?xf32, strided<[?, ?], offset: ?>>, memref<?x?xf32, strided<[?, ?], offset: ?>>) linalg.matmul ins(%14, %16: memref<?x?xf32, strided<[?, ?], offset: ?>>, memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%17: memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%17: memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, 1], offset: ?>> return %E : memref<?x?xf32, strided<[?, 1], offset: ?>>
} }
Show All 10 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 = tensor.empty(%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>) { inits(%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 {
// CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}} // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}}
%4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>) { inits(%2 : tensor<?x?xf32>) {
// CHECK: ^{{[a-zA-Z0-9_]*}} // CHECK: ^{{[a-zA-Z0-9_]*}}
// CHECK-SAME: [[ARG0:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG0:%[a-zA-Z0-9_]*]]
// CHECK-SAME: [[ARG1:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG1:%[a-zA-Z0-9_]*]]
// CHECK-SAME: [[ARG2:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG2:%[a-zA-Z0-9_]*]]
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
// CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG0]], [[ARG1]] // CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG0]], [[ARG1]]
// CHECK-NOT: linalg.yield // CHECK-NOT: linalg.yield
// CHECK: arith.mulf [[T1]], [[ARG2]] // CHECK: arith.mulf [[T1]], [[ARG2]]
Show All 16 Lines
{ {
%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 = tensor.empty(%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>) { inits(%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 {
// CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP1]], [[$MAP0]]{{\]}} // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP1]], [[$MAP0]]{{\]}}
%4 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]} %4 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}
ins(%3, %arg2 : tensor<?x?xf32>, f32) ins(%3, %arg2 : tensor<?x?xf32>, f32)
outs(%2 : tensor<?x?xf32>) { inits(%2 : tensor<?x?xf32>) {
// CHECK: ^{{[a-zA-Z0-9_]*}} // CHECK: ^{{[a-zA-Z0-9_]*}}
// CHECK-SAME: [[ARG3:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG3:%[a-zA-Z0-9_]*]]
// CHECK-SAME: [[ARG4:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG4:%[a-zA-Z0-9_]*]]
// CHECK-SAME: [[ARG5:%[a-zA-Z0-9_]*]] // CHECK-SAME: [[ARG5:%[a-zA-Z0-9_]*]]
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
// CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG3]], [[ARG4]] // CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG3]], [[ARG4]]
// CHECK-NOT: linalg.yield // CHECK-NOT: linalg.yield
// CHECK: arith.mulf [[T1]], [[ARG5]] // CHECK: arith.mulf [[T1]], [[ARG5]]
Show All 16 Lines
{ {
%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 = tensor.empty(%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>) { inits(%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 {
// CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP0]], [[$MAP0]]{{\]}} // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP0]], [[$MAP0]]{{\]}}
%4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]} %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>) { inits(%2 : tensor<?x?xf32>) {
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
%5 = arith.mulf %arg5, %arg6 : f32 %5 = arith.mulf %arg5, %arg6 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %4 : tensor<?x?xf32> return %4 : tensor<?x?xf32>
} }
// ----- // -----
// 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)>
// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)>
#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)>
// 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 = tensor.empty(%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>) { inits(%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 {
// CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}} // CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}}
%4 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]} %4 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]}
ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%2 : tensor<?x?xf32>){ inits(%2 : tensor<?x?xf32>){
^bb0(%arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):
%5 = arith.mulf %arg5, %arg6 : f32 %5 = arith.mulf %arg5, %arg6 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %4 : tensor<?x?xf32> return %4 : tensor<?x?xf32>
} }
// ----- // -----
// 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)>
// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0)> // CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0)>
#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) -> (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 = tensor.empty(%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>) { inits(%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 = tensor.empty(%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>){ inits(%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 = tensor.empty() : 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>) { inits(%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 {
// CHECK: arith.addf // CHECK: arith.addf
// CHECK: arith.mulf // CHECK: arith.mulf
%2 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []} %2 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []}
ins(%1, %arg2 : tensor<f32>, tensor<f32>) ins(%1, %arg2 : tensor<f32>, tensor<f32>)
outs(%0 : tensor<f32>) { inits(%0 : tensor<f32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%3 = arith.mulf %arg3, %arg4 : f32 %3 = arith.mulf %arg3, %arg4 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<f32> } -> tensor<f32>
return %2 : tensor<f32> return %2 : tensor<f32>
} }
Show All 9 Linesfunc.func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32>
%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 = tensor.empty(%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>) { inits(%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
} -> tensor<5x?x?xf32> } -> tensor<5x?x?xf32>
return %3 : tensor<5x?x?xf32> return %3 : tensor<5x?x?xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @generic_op_constant_fusion // CHECK-LABEL: func @generic_op_constant_fusion
Show All 15 Linesfunc.func @generic_op_zero_dim_constant_fusion(%arg0 : tensor<5x?x?xf32>)
%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 = tensor.empty(%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>) { inits(%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
} -> tensor<5x?x?xf32> } -> tensor<5x?x?xf32>
return %3 : tensor<5x?x?xf32> return %3 : tensor<5x?x?xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK-LABEL: func @generic_op_zero_dim_constant_fusion // CHECK-LABEL: func @generic_op_zero_dim_constant_fusion
Show All 11 Linesfunc.func @producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>,
%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 = tensor.empty(%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>) { inits(%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
} -> tensor<?x?xi32> } -> tensor<?x?xi32>
%4 = linalg.generic { %4 = linalg.generic {
indexing_maps = [#map0, #map0], indexing_maps = [#map0, #map0],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%3 : tensor<?x?xi32>) ins(%3 : tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { inits(%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
%5 = arith.index_cast %idx0 : index to i32 %5 = arith.index_cast %idx0 : index to i32
%6 = arith.index_cast %idx1 : index to i32 %6 = arith.index_cast %idx1 : index to i32
%7 = arith.addi %arg2, %5 : i32 %7 = arith.addi %arg2, %5 : i32
%8 = arith.subi %7, %6 : i32 %8 = arith.subi %7, %6 : i32
linalg.yield %8 : i32 linalg.yield %8 : i32
Show All 25 Linesfunc.func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {
%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 = tensor.empty(%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>) { inits(%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
%4 = arith.index_cast %idx0 : index to i32 %4 = arith.index_cast %idx0 : index to i32
%5 = arith.index_cast %idx1 : index to i32 %5 = arith.index_cast %idx1 : index to i32
%6 = arith.addi %arg4, %4 : i32 %6 = arith.addi %arg4, %4 : i32
%7 = arith.subi %6, %5 : i32 %7 = arith.subi %6, %5 : i32
linalg.yield %7 : i32 linalg.yield %7 : i32
} -> tensor<?x?xi32> } -> tensor<?x?xi32>
%4 = linalg.generic { %4 = linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%3, %arg0 : tensor<?x?xi32>, tensor<?x?xi32>) ins(%3, %arg0 : tensor<?x?xi32>, tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { inits(%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
} -> tensor<?x?xi32> } -> tensor<?x?xi32>
return %4 : tensor<?x?xi32> return %4 : tensor<?x?xi32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: func @indexed_producer_consumer_fusion // CHECK-LABEL: func @indexed_producer_consumer_fusion
Show All 23 Linesfunc.func @indexed_producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>)
%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 = tensor.empty(%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>) { inits(%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
%4 = arith.index_cast %idx0 : index to i32 %4 = arith.index_cast %idx0 : index to i32
%5 = arith.index_cast %idx1 : index to i32 %5 = arith.index_cast %idx1 : index to i32
%6 = arith.addi %arg2, %4 : i32 %6 = arith.addi %arg2, %4 : i32
%7 = arith.subi %5, %6 : i32 %7 = arith.subi %5, %6 : i32
linalg.yield %7 : i32 linalg.yield %7 : i32
} -> tensor<?x?xi32> } -> tensor<?x?xi32>
%4= linalg.generic { %4= linalg.generic {
indexing_maps = [#map1, #map1], indexing_maps = [#map1, #map1],
iterator_types = ["parallel", "parallel"] } iterator_types = ["parallel", "parallel"] }
ins(%3 : tensor<?x?xi32>) ins(%3 : tensor<?x?xi32>)
outs(%2 : tensor<?x?xi32>) { inits(%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
%5 = arith.index_cast %idx0 : index to i32 %5 = arith.index_cast %idx0 : index to i32
%6 = arith.index_cast %idx1 : index to i32 %6 = arith.index_cast %idx1 : index to i32
%7 = arith.addi %arg2, %5 : i32 %7 = arith.addi %arg2, %5 : i32
%8 = arith.subi %7, %6 : i32 %8 = arith.subi %7, %6 : i32
linalg.yield %8 : i32 linalg.yield %8 : i32
Show All 30 Linesfunc.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 = tensor.empty(%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>) inits(%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 = tensor.empty(%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>) { inits(%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
} -> tensor<?x?xi32> } -> tensor<?x?xi32>
return %5 : tensor<?x?xi32> return %5 : tensor<?x?xi32>
} }
// 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)>
Show All 15 Linesfunc.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 = tensor.empty() : 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>) inits(%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 = tensor.empty() : 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>) inits(%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
} -> tensor<10xf32> } -> tensor<10xf32>
return %3 : tensor<10xf32> return %3 : tensor<10xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> ()> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> ()>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
Show All 13 Lines
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 = tensor.empty() : 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>) { inits (%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:.+]] = tensor.empty() : 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: inits(%[[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 = tensor.empty() : 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>) { inits(%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
} -> tensor<1x10xf32> } -> tensor<1x10xf32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [#map1, #map2], {indexing_maps = [#map1, #map2],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%0 : tensor<1x10xf32>) ins(%0 : tensor<1x10xf32>)
outs(%arg2 : tensor<1xf32>) { inits(%arg2 : tensor<1xf32>) {
^bb0(%arg3: f32, %arg4: f32): ^bb0(%arg3: f32, %arg4: f32):
%2 = arith.addf %arg3, %arg4 : f32 %2 = arith.addf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<1xf32> } -> tensor<1xf32>
return %1 : tensor<1xf32> return %1 : tensor<1xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (0, d0)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (0, d0)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (0)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (0)>
Show All 20 Linesfunc.func @sigmoid_dynamic_dim(%0: tensor<?x1xf32>) -> tensor<?x1xf32> {
%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 = tensor.empty(%extracted) : 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>) { inits(%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 = tensor.empty(%d0) : 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>) { inits(%init1 : tensor<?x1xf32>) {
^bb0(%a: f32, %b: f32, %c: f32): ^bb0(%a: f32, %b: f32, %c: f32):
%m = arith.mulf %a, %b : f32 %m = arith.mulf %a, %b : f32
linalg.yield %m : f32 linalg.yield %m : f32
} -> tensor<?x1xf32> } -> tensor<?x1xf32>
return %2 : tensor<?x1xf32> return %2 : tensor<?x1xf32>
} }
// ----- // -----
func.func private @compute1(%a: f64) -> f64 func.func private @compute1(%a: f64) -> f64
func.func private @compute2(%a: f64, %b: i32) -> i32 func.func private @compute2(%a: f64, %b: i32) -> i32
// CHECK-LABEL: func @generic_index_op2( // CHECK-LABEL: func @generic_index_op2(
func.func @generic_index_op2(%arg0: tensor<1x8xf64>, %arg1: tensor<1x8xi32>) -> tensor<1x8xi32> { func.func @generic_index_op2(%arg0: tensor<1x8xf64>, %arg1: tensor<1x8xi32>) -> tensor<1x8xi32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>], indexing_maps = [affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%arg0 : tensor<1x8xf64>) { inits(%arg0 : tensor<1x8xf64>) {
^bb0(%a: f64): ^bb0(%a: f64):
%r = func.call @compute1(%a) : (f64) -> f64 %r = func.call @compute1(%a) : (f64) -> f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<1x8xf64> } -> tensor<1x8xf64>
// CHECK-NEXT: %[[R:.*]]:2 = linalg.generic // CHECK-NEXT: %[[R:.*]]:2 = linalg.generic
// CHECK: bb0(%[[BBA:[0-9a-zA-Z_]*]]: f64, %[[BBB:[0-9a-zA-Z_]*]]: i32): // CHECK: bb0(%[[BBA:[0-9a-zA-Z_]*]]: f64, %[[BBB:[0-9a-zA-Z_]*]]: i32):
// CHECK-NEXT: %[[A:.*]] = func.call @compute1(%[[BBA]]) : (f64) -> f64 // CHECK-NEXT: %[[A:.*]] = func.call @compute1(%[[BBA]]) : (f64) -> f64
// CHECK-NEXT: %[[B:.*]] = func.call @compute2(%[[A]], %[[BBB]]) : (f64, i32) -> i32 // CHECK-NEXT: %[[B:.*]] = func.call @compute2(%[[A]], %[[BBB]]) : (f64, i32) -> i32
// CHECK-NEXT: linalg.yield %[[A]], %[[B]] : f64, i32 // CHECK-NEXT: linalg.yield %[[A]], %[[B]] : f64, i32
// CHECK-NEXT: } -> (tensor<1x8xf64>, tensor<1x8xi32>) // CHECK-NEXT: } -> (tensor<1x8xf64>, tensor<1x8xi32>)
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>, affine_map<(i, j) -> (i, j)>], indexing_maps = [affine_map<(i, j) -> (i, j)>, affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%0 : tensor<1x8xf64>) ins(%0 : tensor<1x8xf64>)
outs(%arg1 : tensor<1x8xi32>) { inits(%arg1 : tensor<1x8xi32>) {
^bb0(%a: f64, %b: i32): ^bb0(%a: f64, %b: i32):
%r = func.call @compute2(%a, %b) : (f64, i32) -> i32 %r = func.call @compute2(%a, %b) : (f64, i32) -> i32
linalg.yield %r : i32 linalg.yield %r : i32
} -> tensor<1x8xi32> } -> tensor<1x8xi32>
// CHECK-NEXT: return %[[R]]#1 : tensor<1x8xi32> // CHECK-NEXT: return %[[R]]#1 : tensor<1x8xi32>
return %1 : tensor<1x8xi32> return %1 : tensor<1x8xi32>
} }
// ----- // -----
// 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 = tensor.empty() : 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>) inits(%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
} -> tensor<3xf32> } -> tensor<3xf32>
return %result : tensor<3xf32> return %result : tensor<3xf32>
} }
// ----- // -----
#map = affine_map<(d0, d1) -> (d0, d1)> #map = affine_map<(d0, d1) -> (d0, d1)>
#trait = { #trait = {
indexing_maps = [#map, #map], indexing_maps = [#map, #map],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
func.func @break_outs_dependency(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> func.func @break_outs_dependency(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
{ {
%0 = linalg.generic #trait ins(%arg0 : tensor<?x?xf32>) outs(%arg0 : tensor<?x?xf32>) { %0 = linalg.generic #trait ins(%arg0 : tensor<?x?xf32>) inits(%arg0 : tensor<?x?xf32>) {
^bb0(%arg1 : f32, %arg2 : f32) : ^bb0(%arg1 : f32, %arg2 : f32) :
%1 = arith.addf %arg1, %arg1 : f32 %1 = arith.addf %arg1, %arg1 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%2 = linalg.generic #trait ins(%0 : tensor<?x?xf32>) outs(%0 : tensor<?x?xf32>) { %2 = linalg.generic #trait ins(%0 : tensor<?x?xf32>) inits(%0 : tensor<?x?xf32>) {
^bb0(%arg1 : f32, %arg2 : f32) : ^bb0(%arg1 : f32, %arg2 : f32) :
%3 = arith.mulf %arg1, %arg1 : f32 %3 = arith.mulf %arg1, %arg1 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> 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:.+]] = tensor.empty(%[[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: inits(%[[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:.+]] = tensor.empty(%[[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: inits(%[[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 = tensor.empty(%d0, %d1) : tensor<?x?xf32> %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>
%1 = tensor.empty(%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)
outs(%0, %1 : tensor<?x?xf32>, tensor<?x?xi32>) { inits(%0, %1 : tensor<?x?xf32>, tensor<?x?xi32>) {
^bb0(%arg1 : f32, %arg2 : f32, %arg3 : i32, %arg4 : f32, %arg5 : i32) : ^bb0(%arg1 : f32, %arg2 : f32, %arg3 : i32, %arg4 : f32, %arg5 : i32) :
%3 = arith.addf %arg1, %arg2 : f32 %3 = arith.addf %arg1, %arg2 : f32
linalg.yield %3, %arg3 : f32, i32 linalg.yield %3, %arg3 : f32, i32
} -> (tensor<?x?xf32>, tensor<?x?xi32>) } -> (tensor<?x?xf32>, tensor<?x?xi32>)
return %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xi32> return %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xi32>
} }
// CHECK-LABEL: func @fuse_scalar_constant // CHECK-LABEL: func @fuse_scalar_constant
// CHECK-DAG: %[[CST:.+]] = arith.constant 4.000000e+00 : f32 // CHECK-DAG: %[[CST:.+]] = arith.constant 4.000000e+00 : f32
// CHECK-DAG: %[[C42:.+]] = arith.constant 42 : i32 // CHECK-DAG: %[[C42:.+]] = arith.constant 42 : i32
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%{{.+}} : tensor<?x?xf32>) // CHECK-SAME: ins(%{{.+}} : tensor<?x?xf32>)
// CHECK: %[[YIELD:.+]] = arith.addf %{{.+}}, %[[CST]] : f32 // CHECK: %[[YIELD:.+]] = arith.addf %{{.+}}, %[[CST]] : f32
// CHECK: linalg.yield %[[YIELD]], %[[C42]] : f32, i32 // CHECK: linalg.yield %[[YIELD]], %[[C42]] : f32, i32
// ----- // -----
// CHECK-LABEL: @transpose_fold_2d_fp32 // CHECK-LABEL: @transpose_fold_2d_fp32
func.func @transpose_fold_2d_fp32(%init: tensor<3x2xf32>) -> tensor<3x2xf32> { func.func @transpose_fold_2d_fp32(%init: tensor<3x2xf32>) -> tensor<3x2xf32> {
%input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32> %input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32>
// CHECK: %[[CST:.+]] = arith.constant // CHECK: %[[CST:.+]] = arith.constant
// CHECK-SAME{LITERAL}: dense<[[0.000000e+00, 3.000000e+00], [1.000000e+00, 4.000000e+00], [2.000000e+00, 5.000000e+00]]> : tensor<3x2xf32> // CHECK-SAME{LITERAL}: dense<[[0.000000e+00, 3.000000e+00], [1.000000e+00, 4.000000e+00], [2.000000e+00, 5.000000e+00]]> : tensor<3x2xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<2x3xf32>) outs(%init : tensor<3x2xf32>) { } ins(%input : tensor<2x3xf32>) inits(%init : tensor<3x2xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<3x2xf32> } -> tensor<3x2xf32>
// CHECK: return %[[CST]] // CHECK: return %[[CST]]
return %1 : tensor<3x2xf32> return %1 : tensor<3x2xf32>
} }
// ----- // -----
// CHECK-LABEL: @transpose_fold_2d_fp64 // CHECK-LABEL: @transpose_fold_2d_fp64
func.func @transpose_fold_2d_fp64(%init: tensor<3x2xf64>) -> tensor<3x2xf64> { func.func @transpose_fold_2d_fp64(%init: tensor<3x2xf64>) -> tensor<3x2xf64> {
%input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf64> %input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf64>
// CHECK: %[[CST:.+]] = arith.constant // CHECK: %[[CST:.+]] = arith.constant
// CHECK-SAME{LITERAL}: dense<[[0.000000e+00, 3.000000e+00], [1.000000e+00, 4.000000e+00], [2.000000e+00, 5.000000e+00]]> : tensor<3x2xf64> // CHECK-SAME{LITERAL}: dense<[[0.000000e+00, 3.000000e+00], [1.000000e+00, 4.000000e+00], [2.000000e+00, 5.000000e+00]]> : tensor<3x2xf64>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<2x3xf64>) outs(%init : tensor<3x2xf64>) { } ins(%input : tensor<2x3xf64>) inits(%init : tensor<3x2xf64>) {
^bb0(%arg1: f64, %arg2: f64): ^bb0(%arg1: f64, %arg2: f64):
linalg.yield %arg1 : f64 linalg.yield %arg1 : f64
} -> tensor<3x2xf64> } -> tensor<3x2xf64>
// CHECK: return %[[CST]] // CHECK: return %[[CST]]
return %1 : tensor<3x2xf64> return %1 : tensor<3x2xf64>
} }
// ----- // -----
// CHECK-LABEL: @transpose_fold_4d_i32 // CHECK-LABEL: @transpose_fold_4d_i32
func.func @transpose_fold_4d_i32(%init: tensor<3x1x4x2xi32>) -> tensor<3x1x4x2xi32> { func.func @transpose_fold_4d_i32(%init: tensor<3x1x4x2xi32>) -> tensor<3x1x4x2xi32> {
%input = arith.constant dense<[[ %input = arith.constant dense<[[
[[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]], [[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]],
[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]] [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]
]]> : tensor<1x2x3x4xi32> ]]> : tensor<1x2x3x4xi32>
// CHECK: %[[CST:.+]] = arith.constant dense<[ // CHECK: %[[CST:.+]] = arith.constant dense<[
// CHECK-SAME{LITERAL}: [[[0, 12], [1, 13], [2, 14], [3, 15]]], // CHECK-SAME{LITERAL}: [[[0, 12], [1, 13], [2, 14], [3, 15]]],
// CHECK-SAME{LITERAL}: [[[4, 16], [5, 17], [6, 18], [7, 19]]], // CHECK-SAME{LITERAL}: [[[4, 16], [5, 17], [6, 18], [7, 19]]],
// CHECK-SAME{LITERAL}: [[[8, 20], [9, 21], [10, 22], [11, 23]]] // CHECK-SAME{LITERAL}: [[[8, 20], [9, 21], [10, 22], [11, 23]]]
// CHECK-SAME{LITERAL}: ]> // CHECK-SAME{LITERAL}: ]>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>], indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%input : tensor<1x2x3x4xi32>) outs(%init : tensor<3x1x4x2xi32>) { } ins(%input : tensor<1x2x3x4xi32>) inits(%init : tensor<3x1x4x2xi32>) {
^bb0(%arg1: i32, %arg2: i32): ^bb0(%arg1: i32, %arg2: i32):
linalg.yield %arg1 : i32 linalg.yield %arg1 : i32
} -> tensor<3x1x4x2xi32> } -> tensor<3x1x4x2xi32>
// CHECK: return %[[CST]] // CHECK: return %[[CST]]
return %1 : tensor<3x1x4x2xi32> return %1 : tensor<3x1x4x2xi32>
} }
// ----- // -----
// CHECK-LABEL: @transpose_fold_4d_i16 // CHECK-LABEL: @transpose_fold_4d_i16
func.func @transpose_fold_4d_i16(%init: tensor<3x1x4x2xi16>) -> tensor<3x1x4x2xi16> { func.func @transpose_fold_4d_i16(%init: tensor<3x1x4x2xi16>) -> tensor<3x1x4x2xi16> {
%input = arith.constant dense<[[ %input = arith.constant dense<[[
[[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]], [[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11]],
[[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]] [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]
]]> : tensor<1x2x3x4xi16> ]]> : tensor<1x2x3x4xi16>
// CHECK: %[[CST:.+]] = arith.constant dense<[ // CHECK: %[[CST:.+]] = arith.constant dense<[
// CHECK-SAME{LITERAL}: [[[0, 12], [1, 13], [2, 14], [3, 15]]], // CHECK-SAME{LITERAL}: [[[0, 12], [1, 13], [2, 14], [3, 15]]],
// CHECK-SAME{LITERAL}: [[[4, 16], [5, 17], [6, 18], [7, 19]]], // CHECK-SAME{LITERAL}: [[[4, 16], [5, 17], [6, 18], [7, 19]]],
// CHECK-SAME{LITERAL}: [[[8, 20], [9, 21], [10, 22], [11, 23]]] // CHECK-SAME{LITERAL}: [[[8, 20], [9, 21], [10, 22], [11, 23]]]
// CHECK-SAME{LITERAL}: ]> // CHECK-SAME{LITERAL}: ]>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>], indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d2, d0, d3, d1)>],
iterator_types = ["parallel", "parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%input : tensor<1x2x3x4xi16>) outs(%init : tensor<3x1x4x2xi16>) { } ins(%input : tensor<1x2x3x4xi16>) inits(%init : tensor<3x1x4x2xi16>) {
^bb0(%arg1: i16, %arg2: i16): ^bb0(%arg1: i16, %arg2: i16):
linalg.yield %arg1 : i16 linalg.yield %arg1 : i16
} -> tensor<3x1x4x2xi16> } -> tensor<3x1x4x2xi16>
// CHECK: return %[[CST]] // CHECK: return %[[CST]]
return %1 : tensor<3x1x4x2xi16> return %1 : tensor<3x1x4x2xi16>
} }
// ----- // -----
// CHECK-LABEL: @transpose_nofold_non_cst_input // CHECK-LABEL: @transpose_nofold_non_cst_input
func.func @transpose_nofold_non_cst_input(%input: tensor<2x3xf32>, %init: tensor<3x2xf32>) -> tensor<3x2xf32> { func.func @transpose_nofold_non_cst_input(%input: tensor<2x3xf32>, %init: tensor<3x2xf32>) -> tensor<3x2xf32> {
// CHECK: linalg.generic // CHECK: linalg.generic
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<2x3xf32>) outs(%init : tensor<3x2xf32>) { } ins(%input : tensor<2x3xf32>) inits(%init : tensor<3x2xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %arg1 : f32 linalg.yield %arg1 : f32
} -> tensor<3x2xf32> } -> tensor<3x2xf32>
return %1 : tensor<3x2xf32> return %1 : tensor<3x2xf32>
} }
// ----- // -----
// CHECK-LABEL: @transpose_nofold_yield_const // CHECK-LABEL: @transpose_nofold_yield_const
func.func @transpose_nofold_yield_const(%init: tensor<3x2xf32>) -> tensor<3x2xf32> { func.func @transpose_nofold_yield_const(%init: tensor<3x2xf32>) -> tensor<3x2xf32> {
%input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32> %input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32>
%cst = arith.constant 8.0 : f32 %cst = arith.constant 8.0 : f32
// CHECK: linalg.generic // CHECK: linalg.generic
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<2x3xf32>) outs(%init : tensor<3x2xf32>) { } ins(%input : tensor<2x3xf32>) inits(%init : tensor<3x2xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<3x2xf32> } -> tensor<3x2xf32>
return %1 : tensor<3x2xf32> return %1 : tensor<3x2xf32>
} }
// ----- // -----
// CHECK-LABEL: @transpose_nofold_multi_ops_in_region // CHECK-LABEL: @transpose_nofold_multi_ops_in_region
func.func @transpose_nofold_multi_ops_in_region(%init: tensor<3x2xf32>) -> tensor<3x2xf32> { func.func @transpose_nofold_multi_ops_in_region(%init: tensor<3x2xf32>) -> tensor<3x2xf32> {
%input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32> %input = arith.constant dense<[[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]> : tensor<2x3xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>], indexing_maps = [affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<2x3xf32>) outs(%init : tensor<3x2xf32>) { } ins(%input : tensor<2x3xf32>) inits(%init : tensor<3x2xf32>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%add = arith.addf %arg1, %arg1 : f32 %add = arith.addf %arg1, %arg1 : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<3x2xf32> } -> tensor<3x2xf32>
return %1 : tensor<3x2xf32> return %1 : tensor<3x2xf32>
} }
// ----- // -----
Show All 12 Lines
// 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 = tensor.empty(%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>) inits(%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 = tensor.empty(%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) inits(%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>) inits(%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
} -> tensor<?xf32> } -> tensor<?xf32>
return %8 : tensor<?xf32> return %8 : tensor<?xf32>
} }
// ----- // -----
func.func @fusion_different_axes(%arg0 : tensor<5000xi64>, %arg1 : tensor<5000xi32>) -> tensor<5000xi32> { func.func @fusion_different_axes(%arg0 : tensor<5000xi64>, %arg1 : tensor<5000xi32>) -> tensor<5000xi32> {
%c1_i32 = arith.constant 1 : i32 %c1_i32 = arith.constant 1 : i32
%0 = linalg.generic { %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>) { inits(%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 = tensor.empty() : 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>) inits(%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:.+]] = tensor.empty() : tensor<5000xi64> // CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<5000xi64>
// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : 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: inits(%[[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
// CHECK-DAG: %[[CAST2:.+]] = arith.index_cast %[[CAST1]] : i64 to index // CHECK-DAG: %[[CAST2:.+]] = arith.index_cast %[[CAST1]] : i64 to index
// CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG1]][%[[CAST2]]] // CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG1]][%[[CAST2]]]
// CHECK: linalg.yield %[[CAST1]], %[[EXTRACT]] // CHECK: linalg.yield %[[CAST1]], %[[EXTRACT]]
// CHECK: return %[[RESULT]]#1 // CHECK: return %[[RESULT]]#1
// ----- // -----
// CHECK-LABEL: func @fold_fill_generic_basic // CHECK-LABEL: func @fold_fill_generic_basic
// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> { // CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> {
// CHECK-NOT: linalg.fill // CHECK-NOT: linalg.fill
// 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: inits({{.*}} : 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 = tensor.empty(%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) inits(%1 : tensor<?xf32>) -> tensor<?xf32>
%3 = tensor.empty(%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>) inits (%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: inits({{.*}} : 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 = tensor.empty(%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) inits(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
%4 = tensor.empty(%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) inits(%4 : tensor<?x?xf32>) -> tensor<?x?xf32>
%6 = tensor.empty(%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>) inits (%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 = tensor.empty() : tensor<f32> %0 = tensor.empty() : tensor<f32>
%1 = tensor.empty() : 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>) inits(%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>) inits(%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:.+]] = tensor.empty // 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: inits(%[[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 15 Linesfunc.func @test_fusion_limit(
-> 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 = tensor.empty(%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>) { inits(%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
ins(%arg2, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg2, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%init : tensor<?x?xf32>) { inits(%init : tensor<?x?xf32>) {
^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32): ^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32):
%3 = arith.mulf %arg6, %arg7 : f32 %3 = arith.mulf %arg6, %arg7 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%4 = linalg.generic #binary2Dpointwise %4 = linalg.generic #binary2Dpointwise
ins(%arg4, %arg5 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg4, %arg5 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%init : tensor<?x?xf32>) { inits(%init : tensor<?x?xf32>) {
^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32): ^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32):
%5 = arith.mulf %arg6, %arg7 : f32 %5 = arith.mulf %arg6, %arg7 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%6 = linalg.generic #ternary2Dpointwise %6 = linalg.generic #ternary2Dpointwise
ins(%0, %2, %4 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>) ins(%0, %2, %4 : tensor<?x?xf32>, tensor<?x?xf32>, tensor<?x?xf32>)
outs(%init : tensor<?x?xf32>) { inits(%init : tensor<?x?xf32>) {
^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32, %arg9 : f32): ^bb0(%arg6 : f32, %arg7 : f32, %arg8 : f32, %arg9 : f32):
%7 = arith.addf %arg6, %arg7 : f32 %7 = arith.addf %arg6, %arg7 : f32
%8 = arith.addf %7, %arg8 : f32 %8 = arith.addf %7, %arg8 : f32
linalg.yield %8 : f32 linalg.yield %8 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %6 : tensor<?x?xf32> return %6 : tensor<?x?xf32>
} }
// CHECK-LABEL: func @test_fusion_limit // CHECK-LABEL: func @test_fusion_limit
Show All 22 Lines

mlir/test/Dialect/Linalg/fusion-indexed.mlir

// RUN: mlir-opt %s -test-linalg-greedy-fusion -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-greedy-fusion -split-input-file | FileCheck %s
#id_2d = affine_map<(d0, d1) -> (d0, d1)> #id_2d = affine_map<(d0, d1) -> (d0, d1)>
#pointwise_2d_trait = { #pointwise_2d_trait = {
indexing_maps = [#id_2d, #id_2d, #id_2d], indexing_maps = [#id_2d, #id_2d, #id_2d],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
func.func @fuse_indexed_consumer(%A: memref<?x?xf32>, func.func @fuse_indexed_consumer(%A: memref<?x?xf32>,
%B: memref<?x?xf32>, %B: memref<?x?xf32>,
%C: memref<?x?xf32>, %C: memref<?x?xf32>,
%D: memref<?x?xf32>) { %D: memref<?x?xf32>) {
linalg.generic #pointwise_2d_trait linalg.generic #pointwise_2d_trait
ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>) ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C : memref<?x?xf32>) { inits(%C : memref<?x?xf32>) {
^bb0(%e: f32, %arg5: f32, %arg6: f32): ^bb0(%e: f32, %arg5: f32, %arg6: f32):
%2 = arith.addf %e, %arg5 : f32 %2 = arith.addf %e, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} }
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c25 = arith.constant 25 : index %c25 = arith.constant 25 : index
%c10 = arith.constant 10 : index %c10 = arith.constant 10 : index
%0 = memref.dim %C, %c0 : memref<?x?xf32> %0 = memref.dim %C, %c0 : memref<?x?xf32>
%1 = memref.dim %C, %c1 : memref<?x?xf32> %1 = memref.dim %C, %c1 : memref<?x?xf32>
%2 = memref.dim %D, %c0 : memref<?x?xf32> %2 = memref.dim %D, %c0 : memref<?x?xf32>
%3 = memref.dim %D, %c1 : memref<?x?xf32> %3 = memref.dim %D, %c1 : memref<?x?xf32>
scf.for %arg2 = %c0 to %0 step %c10 { scf.for %arg2 = %c0 to %0 step %c10 {
scf.for %arg3 = %c0 to %1 step %c25 { scf.for %arg3 = %c0 to %1 step %c25 {
%4 = memref.subview %C[%arg2, %arg3][%c10, %c25][%c1, %c1] : %4 = memref.subview %C[%arg2, %arg3][%c10, %c25][%c1, %c1] :
memref<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%5 = memref.subview %D[%arg2, %arg3][%c10, %c25][%c1, %c1] : %5 = memref.subview %D[%arg2, %arg3][%c10, %c25][%c1, %c1] :
memref<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.generic { linalg.generic {
indexing_maps = [#id_2d, #id_2d], indexing_maps = [#id_2d, #id_2d],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%4 : memref<?x?xf32, strided<[?, ?], offset: ?>>) ins(%4 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%5 : memref<?x?xf32, strided<[?, ?], offset: ?>>) { inits(%5 : memref<?x?xf32, strided<[?, ?], offset: ?>>) {
^bb0(%arg4: f32, %arg5: f32): ^bb0(%arg4: f32, %arg5: f32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%6 = arith.addi %idx0, %arg2 : index %6 = arith.addi %idx0, %arg2 : index
%7 = arith.addi %idx1, %arg3 : index %7 = arith.addi %idx1, %arg3 : index
%8 = arith.index_cast %6 : index to i32 %8 = arith.index_cast %6 : index to i32
%9 = arith.sitofp %8 : i32 to f32 %9 = arith.sitofp %8 : i32 to f32
%10 = arith.index_cast %7 : index to i32 %10 = arith.index_cast %7 : index to i32
Show All 22 Linesfunc.func @fuse_indexed_producer(%A: memref<?x?xindex>,
%B: memref<?x?xindex>) { %B: memref<?x?xindex>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c25 = arith.constant 25 : index %c25 = arith.constant 25 : index
%c10 = arith.constant 10 : index %c10 = arith.constant 10 : index
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(i, j) -> (j, i)>], indexing_maps = [affine_map<(i, j) -> (j, i)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%A : memref<?x?xindex>) { inits(%A : memref<?x?xindex>) {
^bb0(%a: index): ^bb0(%a: index):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%0 = arith.addi %idx0, %idx1 : index %0 = arith.addi %idx0, %idx1 : index
linalg.yield %0 : index linalg.yield %0 : index
} }
%A_X = memref.dim %A, %c0 : memref<?x?xindex> %A_X = memref.dim %A, %c0 : memref<?x?xindex>
%A_Y = memref.dim %A, %c1 : memref<?x?xindex> %A_Y = memref.dim %A, %c1 : memref<?x?xindex>
scf.parallel (%arg2, %arg3) = (%c0, %c0) to (%A_X, %A_Y) step (%c10, %c25) { scf.parallel (%arg2, %arg3) = (%c0, %c0) to (%A_X, %A_Y) step (%c10, %c25) {
%A_view = memref.subview %A[%arg2, %arg3][%c10, %c25][%c1, %c1] : %A_view = memref.subview %A[%arg2, %arg3][%c10, %c25][%c1, %c1] :
memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>> memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>>
%B_view = memref.subview %B[%arg2, %arg3][%c10, %c25][%c1, %c1] : %B_view = memref.subview %B[%arg2, %arg3][%c10, %c25][%c1, %c1] :
memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>> memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>>
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>, indexing_maps = [affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i, j)>], affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%A_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) ins(%A_view : memref<?x?xindex, strided<[?, ?], offset: ?>>)
outs(%B_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) { inits(%B_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) {
^bb0(%a: index, %b: index): ^bb0(%a: index, %b: index):
linalg.yield %a : index linalg.yield %a : index
} }
} }
return return
} }
// CHECK: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)> // CHECK: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: func @fuse_indexed_producer // CHECK-LABEL: func @fuse_indexed_producer
Show All 12 Lines
func.func @fuse_indexed_producer_tiled_second_dim_only(%A: memref<?x?xindex>, func.func @fuse_indexed_producer_tiled_second_dim_only(%A: memref<?x?xindex>,
%B: memref<?x?xindex>) { %B: memref<?x?xindex>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c25 = arith.constant 25 : index %c25 = arith.constant 25 : index
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>], indexing_maps = [affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%A : memref<?x?xindex>) { inits(%A : memref<?x?xindex>) {
^bb0(%a: index): ^bb0(%a: index):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%0 = arith.addi %idx0, %idx1 : index %0 = arith.addi %idx0, %idx1 : index
linalg.yield %0 : index linalg.yield %0 : index
} }
%A_X = memref.dim %A, %c0 : memref<?x?xindex> %A_X = memref.dim %A, %c0 : memref<?x?xindex>
%A_Y = memref.dim %A, %c1 : memref<?x?xindex> %A_Y = memref.dim %A, %c1 : memref<?x?xindex>
scf.parallel (%arg3) = (%c0) to (%A_Y) step (%c25) { scf.parallel (%arg3) = (%c0) to (%A_Y) step (%c25) {
%A_view = memref.subview %A[%c0, %arg3][%A_X, %c25][%c1, %c1] : %A_view = memref.subview %A[%c0, %arg3][%A_X, %c25][%c1, %c1] :
memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>> memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>>
%B_view = memref.subview %B[%c0, %arg3][%A_X, %c25][%c1, %c1] : %B_view = memref.subview %B[%c0, %arg3][%A_X, %c25][%c1, %c1] :
memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>> memref<?x?xindex> to memref<?x?xindex, strided<[?, ?], offset: ?>>
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>, indexing_maps = [affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i, j)>], affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%A_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) ins(%A_view : memref<?x?xindex, strided<[?, ?], offset: ?>>)
outs(%B_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) { inits(%B_view : memref<?x?xindex, strided<[?, ?], offset: ?>>) {
^bb0(%a: index, %b: index): ^bb0(%a: index, %b: index):
linalg.yield %a : index linalg.yield %a : index
} }
} }
return return
} }
// CHECK: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)> // CHECK: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: func @fuse_indexed_producer_tiled_second_dim_only // CHECK-LABEL: func @fuse_indexed_producer_tiled_second_dim_only
Show All 9 Lines

mlir/test/Dialect/Linalg/fusion-push-reshape.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
// 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 // CHECK-LABEL: func @reshape
// CHECK-SAME: (%[[A:.*]]: tensor<?x16xf32>, %[[B:.*]]: tensor<16xf32>, %[[INIT:.*]]: tensor<?x112x16xf32>) // CHECK-SAME: (%[[A:.*]]: tensor<?x16xf32>, %[[B:.*]]: tensor<16xf32>, %[[INIT:.*]]: tensor<?x112x16xf32>)
// CHECK: %[[RI:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]] : tensor<?x112x16xf32> into tensor<?x16xf32> // CHECK: %[[RI:.*]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2]] : tensor<?x112x16xf32> into tensor<?x16xf32>
// CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP2]]], // CHECK: %[[R:.*]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP3]], #[[$MAP2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel"]} // CHECK-SAME: iterator_types = ["parallel", "parallel"]}
// CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x16xf32>, tensor<16xf32>) outs(%[[RI]] : tensor<?x16xf32>) // CHECK-SAME: ins(%[[A]], %[[B]] : tensor<?x16xf32>, tensor<16xf32>) inits(%[[RI]] : tensor<?x16xf32>)
// CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<?x16xf32> into tensor<?x112x16xf32> // CHECK: %[[RR:.*]] = tensor.expand_shape %[[R]] {{\[}}[0, 1], [2]] : tensor<?x16xf32> into tensor<?x112x16xf32>
// CHECK: return %[[RR]] : tensor<?x112x16xf32> // CHECK: return %[[RR]] : tensor<?x112x16xf32>
func.func @reshape(%A: tensor<?x16xf32>, %B: tensor<16xf32>, %init: tensor<?x112x16xf32>) -> tensor<?x112x16xf32> { func.func @reshape(%A: tensor<?x16xf32>, %B: tensor<16xf32>, %init: tensor<?x112x16xf32>) -> tensor<?x112x16xf32> {
%0 = tensor.expand_shape %A [[0, 1], [2]] %0 = tensor.expand_shape %A [[0, 1], [2]]
: tensor<?x16xf32> into tensor<?x112x16xf32> : tensor<?x16xf32> into tensor<?x112x16xf32>
%2 = linalg.generic {indexing_maps = [ %2 = linalg.generic {indexing_maps = [
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d2)>, affine_map<(d0, d1, d2) -> (d0, d1, 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, %B : tensor<?x112x16xf32>, tensor<16xf32>) ins(%0, %B : tensor<?x112x16xf32>, tensor<16xf32>)
outs(%init : tensor<?x112x16xf32>) { inits(%init : tensor<?x112x16xf32>) {
^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<?x112x16xf32> } -> tensor<?x112x16xf32>
return %2 : tensor<?x112x16xf32> return %2 : tensor<?x112x16xf32>
} }
// ----- // -----
// 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:.*]] = tensor.empty() : 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>) inits(%[[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 = tensor.empty() : 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>) { inits(%2 : tensor<112x112x16xf32>) {
^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):
%s = arith.subf %arg1, %arg2 : f32 %s = arith.subf %arg1, %arg2 : f32
%m = arith.mulf %s, %arg3 : f32 %m = arith.mulf %s, %arg3 : f32
linalg.yield %m : f32 linalg.yield %m : f32
} -> tensor<112x112x16xf32> } -> tensor<112x112x16xf32>
return %3 : tensor<112x112x16xf32> return %3 : tensor<112x112x16xf32>
} }
Show All 9 Linesfunc.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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%26 : tensor<2x3x5xf32>) {
^bb0(%arg6: i32, %arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg6: i32, %arg7: f32, %arg8: f32, %arg9: f32):
%29 = arith.sitofp %arg6 : i32 to f32 %29 = arith.sitofp %arg6 : i32 to f32
%30 = arith.addf %arg7, %cst_8 : f32 %30 = arith.addf %arg7, %cst_8 : f32
%31 = arith.divf %cst_7, %30 : f32 %31 = arith.divf %cst_7, %30 : f32
%32 = arith.divf %cst_6, %31 : f32 %32 = arith.divf %cst_6, %31 : f32
%33 = arith.mulf %29, %32 : f32 %33 = arith.mulf %29, %32 : f32
%34 = arith.addf %33, %arg8 : f32 %34 = arith.addf %33, %arg8 : f32
linalg.yield %34 : f32 linalg.yield %34 : f32
} -> tensor<2x3x5xf32> } -> tensor<2x3x5xf32>
return %28 : tensor<2x3x5xf32> return %28 : tensor<2x3x5xf32>
} }
// CHECK-LABEL: func @type_correctness // CHECK-LABEL: func @type_correctness
// CHECK: %[[OP:.+]] = linalg.generic // CHECK: %[[OP:.+]] = linalg.generic
// CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}}, %{{.+}} : tensor<6x5xi32>, tensor<5xf32>, tensor<5xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<6x5xf32>) // CHECK-SAME: inits(%{{.+}} : tensor<6x5xf32>)
// CHECK: tensor.expand_shape %[[OP]] // CHECK: tensor.expand_shape %[[OP]]
// CHECK-SAME: tensor<6x5xf32> into tensor<2x3x5xf32> // CHECK-SAME: tensor<6x5xf32> into tensor<2x3x5xf32>

mlir/test/Dialect/Linalg/fusion.mlir

Show All 10 Linesfunc.func @f1(%A: memref<?x?xf32, strided<[?, 1], offset: 0>>,
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, 1], offset: 0>> %0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, 1], offset: 0>>
%1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, 1], offset: 0>> %1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, 1], offset: 0>>
%2 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, 1], offset: 0>> %2 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, 1], offset: 0>>
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, 1], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, 1], offset: 0>>,
memref<?x?xf32, strided<[?, 1], offset: 0>>) memref<?x?xf32, strided<[?, 1], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, 1], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, 1], offset: 0>>)
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%5 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: 0>> to memref<?x?xf32, strided<[?, 1], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: 0>> to memref<?x?xf32, strided<[?, 1], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %C[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %C[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: 0>> to memref<?x?xf32, strided<[?, 1], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8: memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8: memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, 1], offset: 0>> return %E : memref<?x?xf32, strided<[?, 1], offset: 0>>
} }
// CHECK-LABEL: func @f1 // CHECK-LABEL: func @f1
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// CHECK: scf.for // CHECK: scf.for
Show All 12 Linesfunc.func @f2(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
) -> memref<?x?xf32, strided<[?, ?], offset: 0>> { ) -> memref<?x?xf32, strided<[?, ?], offset: 0>> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C: memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C: memref<?x?xf32, strided<[?, ?], offset: 0>>)
%0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f2 // CHECK-LABEL: func @f2
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// CHECK-DAG: %[[C_0:.*]] = memref.dim %[[C]], %c0{{[_0-9]*}} : memref<?x?xf32, strided<[?, ?]>> // CHECK-DAG: %[[C_0:.*]] = memref.dim %[[C]], %c0{{[_0-9]*}} : memref<?x?xf32, strided<[?, ?]>>
Show All 15 Linesfunc.func @f3(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
) -> memref<?x?xf32, strided<[?, ?], offset: 0>> { ) -> memref<?x?xf32, strided<[?, ?], offset: 0>> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
%0 = memref.dim %D, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %D, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%5 = memref.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f3 // CHECK-LABEL: func @f3
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
Show All 17 Linesfunc.func @f4(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
) -> memref<?x?xf32, strided<[?, ?], offset: 0>> { ) -> memref<?x?xf32, strided<[?, ?], offset: 0>> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>)
%0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f4 // CHECK-LABEL: func @f4
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
Show All 22 Linesfunc.func @f5(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %D, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %D, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
linalg.matmul ins(%C, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%C, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>)
scf.for %arg5 = %c0 to %1 step %c2 { scf.for %arg5 = %c0 to %1 step %c2 {
scf.for %arg6 = %c0 to %0 step %c3 { scf.for %arg6 = %c0 to %0 step %c3 {
scf.for %arg7 = %c0 to %2 step %c4 { scf.for %arg7 = %c0 to %2 step %c4 {
%5 = memref.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %D[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %B[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-DAG: #[[BOUND_2_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)> // CHECK-DAG: #[[BOUND_2_MAP:.+]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
// CHECK-DAG: #[[BOUND_2_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 2)> // CHECK-DAG: #[[BOUND_2_MAP_2:.+]] = affine_map<(d0)[s0, s1] -> (-d0 + s1, -d0 + s0, 2)>
Show All 17 Lines
// CHECK: scf.for %[[J:.*]] = %{{.*}} to %[[B_1]] step %{{.*}} { // CHECK: scf.for %[[J:.*]] = %{{.*}} to %[[B_1]] step %{{.*}} {
// CHECK: %[[E_IJ:.*]] = memref.subview %[[E]][%[[I]], %[[J]]] // CHECK: %[[E_IJ:.*]] = memref.subview %[[E]][%[[I]], %[[J]]]
// CHECK: scf.for %[[K:.*]] = %{{.*}} to %[[D_1]] step %{{.*}} { // CHECK: scf.for %[[K:.*]] = %{{.*}} to %[[D_1]] step %{{.*}} {
// CHECK: %[[D_IK:.*]] = memref.subview %[[D]][%[[I]], %[[K]]] [2, 4] // CHECK: %[[D_IK:.*]] = memref.subview %[[D]][%[[I]], %[[K]]] [2, 4]
// CHECK: %[[B_KJ:.*]] = memref.subview %[[B]][%[[K]], %[[J]]] // CHECK: %[[B_KJ:.*]] = memref.subview %[[B]][%[[K]], %[[J]]]
// CHECK: %[[BOUND_4_B1:.*]] = affine.min #[[BOUND_4_MAP]](%[[K]])[%[[B_1]]] // CHECK: %[[BOUND_4_B1:.*]] = affine.min #[[BOUND_4_MAP]](%[[K]])[%[[B_1]]]
// CHECK: %[[B_0K:.*]] = memref.subview %[[B]][0, %[[K]]] // CHECK: %[[B_0K:.*]] = memref.subview %[[B]][0, %[[K]]]
// CHECK: %[[D_IK_OUT:.+]] = memref.subview %[[D]][%[[I]], %[[K]]] [%[[BOUND_2_C0]], %[[BOUND_4_B1]]] // CHECK: %[[D_IK_OUT:.+]] = memref.subview %[[D]][%[[I]], %[[K]]] [%[[BOUND_2_C0]], %[[BOUND_4_B1]]]
// CHECK: linalg.matmul ins(%[[A_I0]], %[[B_00]]{{.*}} outs(%[[C_I0_OUT]] // CHECK: linalg.matmul ins(%[[A_I0]], %[[B_00]]{{.*}} inits(%[[C_I0_OUT]]
// CHECK: linalg.matmul ins(%[[C_I0]], %[[B_0K]]{{.*}} outs(%[[D_IK_OUT]] // CHECK: linalg.matmul ins(%[[C_I0]], %[[B_0K]]{{.*}} inits(%[[D_IK_OUT]]
// CHECK: linalg.matmul ins(%[[D_IK]], %[[B_KJ]]{{.*}} outs(%[[E_IJ]] // CHECK: linalg.matmul ins(%[[D_IK]], %[[B_KJ]]{{.*}} inits(%[[E_IJ]]
// ----- // -----
#map0 = affine_map<(d0) -> (d0 + 2)> #map0 = affine_map<(d0) -> (d0 + 2)>
#map1 = affine_map<(d0) -> (d0 + 4)> #map1 = affine_map<(d0) -> (d0 + 4)>
#map2 = affine_map<(d0) -> (d0 + 3)> #map2 = affine_map<(d0) -> (d0 + 3)>
func.func @f6(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>, func.func @f6(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%B: memref<?x?xf32, strided<[?, ?], offset: 0>>, %B: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%C: memref<?x?xf32, strided<[?, ?], offset: 0>>, %C: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%D: memref<?x?xf32, strided<[?, ?], offset: 0>>, %D: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%E: memref<?x?xf32, strided<[?, ?], offset: 0>> %E: memref<?x?xf32, strided<[?, ?], offset: 0>>
) -> memref<?x?xf32, strided<[?, ?], offset: 0>> { ) -> memref<?x?xf32, strided<[?, ?], offset: 0>> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%E : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%E : memref<?x?xf32, strided<[?, ?], offset: 0>>)
%1 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg5 = %c0 to %1 step %c2 { scf.for %arg5 = %c0 to %1 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %0 step %c4 { scf.for %arg7 = %c0 to %0 step %c4 {
%3 = affine.apply #map0(%arg5) %3 = affine.apply #map0(%arg5)
%4 = affine.apply #map1(%arg7) %4 = affine.apply #map1(%arg7)
%5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%6 = affine.apply #map2(%arg6) %6 = affine.apply #map2(%arg6)
%7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f6 // CHECK-LABEL: func @f6
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// Fuse the producer of E (WAW) then the producer of C (WAR). // Fuse the producer of E (WAW) then the producer of C (WAR).
Show All 19 Linesfunc.func @f7(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%2 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %C, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%3 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %3 = memref.dim %C, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%4 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %4 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%E : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%E : memref<?x?xf32, strided<[?, ?], offset: 0>>)
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%7 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] : %7 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%9 = memref.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] : %9 = memref.subview %C[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%10 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %10 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%7, %9 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%7, %9 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%10 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%10 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
scf.for %arg5 = %c0 to %3 step %c2 { scf.for %arg5 = %c0 to %3 step %c2 {
scf.for %arg6 = %c0 to %4 step %c3 { scf.for %arg6 = %c0 to %4 step %c3 {
scf.for %arg7 = %c0 to %2 step %c4 { scf.for %arg7 = %c0 to %2 step %c4 {
%7 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] : %7 = memref.subview %C[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%9 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : %9 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%10 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %10 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%7, %9 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%7, %9 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%10 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%10 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f7 // CHECK-LABEL: func @f7
// CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}}) // CHECK: (%[[A:.*]]:{{.*}}, %[[B:.*]]:{{.*}}, %[[C:.*]]:{{.*}}, %[[D:.*]]:{{.*}}, %[[E:.*]]:{{.*}})
// 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: %[[A_0:.*]] = memref.dim %[[A]], %[[C0:.*]] : memref<?x?xf32, strided<[?, ?]>> // CHECK: %[[A_0:.*]] = memref.dim %[[A]], %[[C0:.*]] : memref<?x?xf32, strided<[?, ?]>>
// CHECK: %[[A_1:.*]] = memref.dim %[[A]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>> // CHECK: %[[A_1:.*]] = memref.dim %[[A]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>>
// CHECK: %[[C_1:.*]] = memref.dim %[[C]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>> // CHECK: %[[C_1:.*]] = memref.dim %[[C]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>>
// CHECK: %[[C_0:.*]] = memref.dim %[[C]], %[[C0:.*]] : memref<?x?xf32, strided<[?, ?]>> // CHECK: %[[C_0:.*]] = memref.dim %[[C]], %[[C0:.*]] : memref<?x?xf32, strided<[?, ?]>>
// CHECK: %[[D_1:.*]] = memref.dim %[[D]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>> // CHECK: %[[D_1:.*]] = memref.dim %[[D]], %[[C1:.*]] : memref<?x?xf32, strided<[?, ?]>>
// CHECK: linalg.matmul ins(%[[A]], %[[C]]{{.*}} outs(%[[E]] // CHECK: linalg.matmul ins(%[[A]], %[[C]]{{.*}} inits(%[[E]]
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[A_0]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[A_0]] step %{{.*}} {
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[A_1]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[A_1]] step %{{.*}} {
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_0]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_0]] step %{{.*}} {
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[D_1]] step %{{.*}} {
// CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %[[C_1]] step %{{.*}} {
Show All 16 Linesfunc.func @f8(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %A, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %A, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %C : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%D : memref<?x?xf32, strided<[?, ?], offset: 0>>)
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
%2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %2 = memref.dim %D, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg5 = %c0 to %0 step %c2 { scf.for %arg5 = %c0 to %0 step %c2 {
scf.for %arg6 = %c0 to %2 step %c3 { scf.for %arg6 = %c0 to %2 step %c3 {
scf.for %arg7 = %c0 to %1 step %c4 { scf.for %arg7 = %c0 to %1 step %c4 {
%3 = affine.apply #map0(%arg5) %3 = affine.apply #map0(%arg5)
%4 = affine.apply #map1(%arg7) %4 = affine.apply #map1(%arg7)
%5 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] : %5 = memref.subview %A[%arg5, %arg7][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%6 = affine.apply #map2(%arg6) %6 = affine.apply #map2(%arg6)
%7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] : %7 = memref.subview %D[%arg7, %arg6][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] : %8 = memref.subview %E[%arg5, %arg6][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%5, %7 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%8 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return %E : memref<?x?xf32, strided<[?, ?], offset: 0>> return %E : memref<?x?xf32, strided<[?, ?], offset: 0>>
} }
// CHECK-LABEL: func @f8 // CHECK-LABEL: func @f8
// CHECK: (%[[A:.*]]: memref{{.*}}, %[[B:.*]]: memref{{.*}}, %[[C:.*]]: memref{{.*}}, %[[D:.*]]: memref{{.*}}, %[[E:.*]]: memref{{.*}}) // CHECK: (%[[A:.*]]: memref{{.*}}, %[[B:.*]]: memref{{.*}}, %[[C:.*]]: memref{{.*}}, %[[D:.*]]: memref{{.*}}, %[[E:.*]]: memref{{.*}})
// CHECK: linalg.matmul // CHECK: linalg.matmul
Show All 17 Linesfunc.func @pointwise(%A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
%D: memref<?x?xf32, strided<[?, ?], offset: 0>>) { %D: memref<?x?xf32, strided<[?, ?], offset: 0>>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
linalg.generic #pointwise_2d_trait linalg.generic #pointwise_2d_trait
ins(%A, %A: memref<?x?xf32, strided<[?, ?], offset: 0>>, ins(%A, %A: memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%B : memref<?x?xf32, strided<[?, ?], offset: 0>>) { inits(%B : memref<?x?xf32, strided<[?, ?], offset: 0>>) {
^bb0(%E: f32, %arg5: f32, %arg6: f32): ^bb0(%E: f32, %arg5: f32, %arg6: f32):
%2 = arith.addf %E, %arg5 : f32 %2 = arith.addf %E, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} }
%0 = memref.dim %B, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>> %0 = memref.dim %B, %c0 : memref<?x?xf32, strided<[?, ?], offset: 0>>
%1 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>> %1 = memref.dim %B, %c1 : memref<?x?xf32, strided<[?, ?], offset: 0>>
scf.for %arg4 = %c0 to %0 step %c2 { scf.for %arg4 = %c0 to %0 step %c2 {
scf.for %arg5 = %c0 to %1 step %c3 { scf.for %arg5 = %c0 to %1 step %c3 {
%4 = memref.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] : %4 = memref.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%5 = memref.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] : %5 = memref.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%6 = memref.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] : %6 = memref.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.generic #pointwise_2d_trait linalg.generic #pointwise_2d_trait
ins(%4, %5: memref<?x?xf32, strided<[?, ?], offset: ?>>, ins(%4, %5: memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>) { inits(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>) {
^bb0(%arg6: f32, %arg7: f32, %arg8: f32): ^bb0(%arg6: f32, %arg7: f32, %arg8: f32):
%7 = arith.mulf %arg6, %arg7 : f32 %7 = arith.mulf %arg6, %arg7 : f32
linalg.yield %7 : f32 linalg.yield %7 : f32
} }
} }
} }
return return
} }
Show All 20 Linesfunc.func @pointwise_no_view(%M: index, %N: index) {
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%A = memref.alloc (%M, %N): memref<?x?xf32> %A = memref.alloc (%M, %N): memref<?x?xf32>
%B = memref.alloc (%M, %N): memref<?x?xf32> %B = memref.alloc (%M, %N): memref<?x?xf32>
%C = memref.alloc (%M, %N): memref<?x?xf32> %C = memref.alloc (%M, %N): memref<?x?xf32>
%D = memref.alloc (%M, %N): memref<?x?xf32> %D = memref.alloc (%M, %N): memref<?x?xf32>
%E = memref.alloc (%M, %N): memref<?x?xf32> %E = memref.alloc (%M, %N): memref<?x?xf32>
linalg.generic #pointwise_2d_trait linalg.generic #pointwise_2d_trait
ins(%A, %A : memref<?x?xf32>, memref<?x?xf32>) ins(%A, %A : memref<?x?xf32>, memref<?x?xf32>)
outs(%B : memref<?x?xf32>) { inits(%B : memref<?x?xf32>) {
^bb0(%e: f32, %arg5: f32, %arg6: f32): ^bb0(%e: f32, %arg5: f32, %arg6: f32):
%2 = arith.addf %e, %arg5 : f32 %2 = arith.addf %e, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} }
%0 = memref.dim %B, %c0 : memref<?x?xf32> %0 = memref.dim %B, %c0 : memref<?x?xf32>
%1 = memref.dim %B, %c1 : memref<?x?xf32> %1 = memref.dim %B, %c1 : memref<?x?xf32>
scf.for %arg4 = %c0 to %0 step %c2 { scf.for %arg4 = %c0 to %0 step %c2 {
scf.for %arg5 = %c0 to %1 step %c3 { scf.for %arg5 = %c0 to %1 step %c3 {
%4 = memref.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] : %4 = memref.subview %B[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32> to memref<?x?xf32> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%5 = memref.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] : %5 = memref.subview %C[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32> to memref<?x?xf32> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%6 = memref.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] : %6 = memref.subview %D[%arg4, %arg5][%c2, %c3][%c1, %c1] :
memref<?x?xf32> to memref<?x?xf32> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.generic #pointwise_2d_trait linalg.generic #pointwise_2d_trait
ins(%4, %5: memref<?x?xf32, strided<[?, ?], offset: ?>>, ins(%4, %5: memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>) { inits(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>) {
^bb0(%arg6: f32, %arg7: f32, %arg8: f32): ^bb0(%arg6: f32, %arg7: f32, %arg8: f32):
%7 = arith.mulf %arg6, %arg7 : f32 %7 = arith.mulf %arg6, %arg7 : f32
linalg.yield %7 : f32 linalg.yield %7 : f32
} }
} }
} }
return return
} }
Show All 17 Linesfunc.func @fusion_of_three(%arg0: memref<100x10xf32>,
%arg2: memref<100x10xf32>) { %arg2: memref<100x10xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = memref.alloc() {temp = true} : memref<100x10xf32> %0 = memref.alloc() {temp = true} : memref<100x10xf32>
linalg.generic { linalg.generic {
indexing_maps = [#map0, #map1], indexing_maps = [#map0, #map1],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg1 : memref<100xf32>) ins(%arg1 : memref<100xf32>)
outs(%0 : memref<100x10xf32>) { inits(%0 : memref<100x10xf32>) {
^bb0(%arg3: f32, %arg4: f32): ^bb0(%arg3: f32, %arg4: f32):
linalg.yield %arg3 : f32 linalg.yield %arg3 : f32
} }
%1 = memref.alloc() {temp = true} : memref<100x10xf32> %1 = memref.alloc() {temp = true} : memref<100x10xf32>
linalg.generic { linalg.generic {
indexing_maps = [#map1, #map1, #map1], indexing_maps = [#map1, #map1, #map1],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %0: memref<100x10xf32>, memref<100x10xf32>) ins(%arg0, %0: memref<100x10xf32>, memref<100x10xf32>)
outs(%1 : memref<100x10xf32>) { inits(%1 : memref<100x10xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%2 = arith.subf %arg3, %arg4 : f32 %2 = arith.subf %arg3, %arg4 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} }
memref.dealloc %0 : memref<100x10xf32> memref.dealloc %0 : memref<100x10xf32>
%2 = memref.dim %1, %c0 : memref<100x10xf32> %2 = memref.dim %1, %c0 : memref<100x10xf32>
%3 = memref.dim %1, %c1 : memref<100x10xf32> %3 = memref.dim %1, %c1 : memref<100x10xf32>
%4 = memref.dim %arg2, %c0 : memref<100x10xf32> %4 = memref.dim %arg2, %c0 : memref<100x10xf32>
%5 = memref.dim %arg2, %c1 : memref<100x10xf32> %5 = memref.dim %arg2, %c1 : memref<100x10xf32>
scf.for %i = %c0 to %2 step %c1 { scf.for %i = %c0 to %2 step %c1 {
scf.for %j = %c0 to %3 step %c1 { scf.for %j = %c0 to %3 step %c1 {
%6 = memref.subview %1[%i, %j][%c1, %c1][%c1, %c1] : %6 = memref.subview %1[%i, %j][%c1, %c1][%c1, %c1] :
memref<100x10xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<100x10xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%7 = memref.subview %arg2[%i, %j][%c1, %c1][%c1, %c1] : %7 = memref.subview %arg2[%i, %j][%c1, %c1][%c1, %c1] :
memref<100x10xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<100x10xf32> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.generic { linalg.generic {
indexing_maps = [#map1, #map1], indexing_maps = [#map1, #map1],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>) ins(%6 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%7 : memref<?x?xf32, strided<[?, ?], offset: ?>>) { inits(%7 : memref<?x?xf32, strided<[?, ?], offset: ?>>) {
^bb0(%arg3: f32, %arg4: f32): ^bb0(%arg3: f32, %arg4: f32):
%8 = math.exp %arg3 : f32 %8 = math.exp %arg3 : f32
linalg.yield %8 : f32 linalg.yield %8 : f32
} }
} }
} }
memref.dealloc %1 : memref<100x10xf32> memref.dealloc %1 : memref<100x10xf32>
return return
Show All 21 Lines
#map4 = affine_map<(d0)[s0, s1] -> (s0 + 2, -d0 + s0 + s1)> #map4 = affine_map<(d0)[s0, s1] -> (s0 + 2, -d0 + s0 + s1)>
func.func @fill_and_conv(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) { func.func @fill_and_conv(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
linalg.fill ins(%cst : f32) outs(%arg0 : memref<?x?xf32>) linalg.fill ins(%cst : f32) inits(%arg0 : memref<?x?xf32>)
%2 = memref.dim %arg1, %c0 : memref<?x?xf32> %2 = memref.dim %arg1, %c0 : memref<?x?xf32>
%3 = memref.dim %arg1, %c1 : memref<?x?xf32> %3 = memref.dim %arg1, %c1 : memref<?x?xf32>
%4 = memref.dim %arg2, %c0 : memref<?x?xf32> %4 = memref.dim %arg2, %c0 : memref<?x?xf32>
%5 = memref.dim %arg2, %c1 : memref<?x?xf32> %5 = memref.dim %arg2, %c1 : memref<?x?xf32>
scf.for %arg3 = %c0 to %4 step %c2 { scf.for %arg3 = %c0 to %4 step %c2 {
scf.for %arg4 = %c0 to %5 step %c3 { scf.for %arg4 = %c0 to %5 step %c3 {
%6 = affine.min #map3(%arg3)[%2, %4] %6 = affine.min #map3(%arg3)[%2, %4]
%7 = affine.min #map4(%arg4)[%3, %5] %7 = affine.min #map4(%arg4)[%3, %5]
%8 = memref.subview %arg0[%arg3, %arg4] [%6, %7] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> %8 = memref.subview %arg0[%arg3, %arg4] [%6, %7] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
%9 = affine.min #map0(%arg3)[%4] %9 = affine.min #map0(%arg3)[%4]
%10 = affine.min #map1(%arg4)[%5] %10 = affine.min #map1(%arg4)[%5]
%11 = memref.subview %arg2[%arg3, %arg4] [%9, %10] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> %11 = memref.subview %arg2[%arg3, %arg4] [%9, %10] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
linalg.conv_2d ins(%8, %arg1 : memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32>) outs(%11 : memref<?x?xf32, strided<[?, 1], offset: ?>>) linalg.conv_2d ins(%8, %arg1 : memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32>) inits(%11 : memref<?x?xf32, strided<[?, 1], offset: ?>>)
} }
} }
return return
} }
// CHECK-LABEL: func @fill_and_conv // CHECK-LABEL: func @fill_and_conv
// CHECK: scf.for // CHECK: scf.for
// CHECK: scf.for // CHECK: scf.for
// CHECK: linalg.fill // CHECK: linalg.fill
Show All 10 Linesfunc.func @accept_different_alloc_ops(%dim: index, %s0 : index, %s1: index) {
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%A = memref.alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>> %A = memref.alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>>
%B = memref.alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>> %B = memref.alloca(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>>
%C = memref.alloc(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>> %C = memref.alloc(%dim, %dim)[%s0, %s1] : memref<?x?xf32, strided<[?, ?], offset: 0>>
linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>, linalg.matmul ins(%A, %B : memref<?x?xf32, strided<[?, ?], offset: 0>>,
memref<?x?xf32, strided<[?, ?], offset: 0>>) memref<?x?xf32, strided<[?, ?], offset: 0>>)
outs(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>) inits(%C : memref<?x?xf32, strided<[?, ?], offset: 0>>)
scf.for %i = %c0 to %dim step %c2 { scf.for %i = %c0 to %dim step %c2 {
scf.for %j = %c0 to %dim step %c3 { scf.for %j = %c0 to %dim step %c3 {
scf.for %k = %c0 to %dim step %c4 { scf.for %k = %c0 to %dim step %c4 {
%0 = memref.subview %A[%i, %k][%c2, %c4][%c1, %c1] : %0 = memref.subview %A[%i, %k][%c2, %c4][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%1 = memref.subview %B[%k, %j][%c4, %c3][%c1, %c1] : %1 = memref.subview %B[%k, %j][%c4, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
%2 = memref.subview %C[%i, %j][%c2, %c3][%c1, %c1] : %2 = memref.subview %C[%i, %j][%c2, %c3][%c1, %c1] :
memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, strided<[?, ?], offset: 0>> to
memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%0, %1 : memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%0, %1 : memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%2 : memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%2 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return return
} }
// CHECK-LABEL: func @accept_different_alloc_ops // CHECK-LABEL: func @accept_different_alloc_ops
// CHECK-COUNT-3: scf.for // CHECK-COUNT-3: scf.for
// CHECK-COUNT-2: linalg.matmul // CHECK-COUNT-2: linalg.matmul

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

// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s // RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
func.func @generalize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) { func.func @generalize_matmul_buffer(%A : memref<16x8xf32>, %B: memref<8x32xf32>, %C: memref<16x32xf32>) {
linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>) linalg.matmul ins(%A, %B: memref<16x8xf32>, memref<8x32xf32>)
outs(%C: memref<16x32xf32>) inits(%C: memref<16x32xf32>)
return return
} }
// CHECK: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> // CHECK: #[[A_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)> // CHECK: #[[B_MAP:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CHECK: #[[C_MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK: func @generalize_matmul_buffer // CHECK: func @generalize_matmul_buffer
// CHECK-SAME: %[[A:.+]]: memref<16x8xf32> // CHECK-SAME: %[[A:.+]]: memref<16x8xf32>
// CHECK-SAME: %[[B:.+]]: memref<8x32xf32> // CHECK-SAME: %[[B:.+]]: memref<8x32xf32>
// CHECK-SAME: %[[C:.+]]: memref<16x32xf32> // CHECK-SAME: %[[C:.+]]: memref<16x32xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]] // CHECK-SAME: indexing_maps = [#[[A_MAP]], #[[B_MAP]], #[[C_MAP]]]
// 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(%[[C]] // CHECK-SAME: inits(%[[C]]
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
// CHECK: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f32 // CHECK: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f32
// CHECK: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 // CHECK: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @generalize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_matmul_tensor(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xf32>, tensor<8x32xf32>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK: func @generalize_matmul_tensor // CHECK: func @generalize_matmul_tensor
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xf32>, tensor<8x32xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<16x32xf32>) // CHECK-SAME: inits(%{{.+}} : tensor<16x32xf32>)
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[A_ARG:.+]]: f32, %[[B_ARG:.+]]: f32, %[[C_ARG:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_ARG]], %[[B_ARG]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// CHECK-NEXT: -> tensor<16x32xf32> // CHECK-NEXT: -> tensor<16x32xf32>
// ----- // -----
func.func @generalize_matmul_tensor_complex(%A : tensor<16x8xcomplex<f32>>, func.func @generalize_matmul_tensor_complex(%A : tensor<16x8xcomplex<f32>>,
%B: tensor<8x32xcomplex<f32>>, %B: tensor<8x32xcomplex<f32>>,
%C: tensor<16x32xcomplex<f32>>) %C: tensor<16x32xcomplex<f32>>)
-> tensor<16x32xcomplex<f32>> { -> tensor<16x32xcomplex<f32>> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>)
outs(%C: tensor<16x32xcomplex<f32>>) -> tensor<16x32xcomplex<f32>> inits(%C: tensor<16x32xcomplex<f32>>) -> tensor<16x32xcomplex<f32>>
return %0: tensor<16x32xcomplex<f32>> return %0: tensor<16x32xcomplex<f32>>
} }
// CHECK: func @generalize_matmul_tensor_complex // CHECK: func @generalize_matmul_tensor_complex
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<16x8xcomplex<f32>>, tensor<8x32xcomplex<f32>>)
// CHECK-SAME: outs(%{{.+}} : tensor<16x32xcomplex<f32>>) // CHECK-SAME: inits(%{{.+}} : tensor<16x32xcomplex<f32>>)
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: complex<f32>, %[[B_ARG:.+]]: complex<f32>, %[[C_ARG:.+]]: complex<f32>) // CHECK: ^{{.*}}(%[[A_ARG:.+]]: complex<f32>, %[[B_ARG:.+]]: complex<f32>, %[[C_ARG:.+]]: complex<f32>)
// CHECK-NEXT: %[[MUL:.+]] = complex.mul %[[A_ARG]], %[[B_ARG]] : complex<f32> // CHECK-NEXT: %[[MUL:.+]] = complex.mul %[[A_ARG]], %[[B_ARG]] : complex<f32>
// CHECK-NEXT: %[[ADD:.+]] = complex.add %[[C_ARG]], %[[MUL]] : complex<f32> // CHECK-NEXT: %[[ADD:.+]] = complex.add %[[C_ARG]], %[[MUL]] : complex<f32>
// CHECK-NEXT: linalg.yield %[[ADD]] : complex<f32> // CHECK-NEXT: linalg.yield %[[ADD]] : complex<f32>
// CHECK-NEXT: -> tensor<16x32xcomplex<f32>> // CHECK-NEXT: -> tensor<16x32xcomplex<f32>>
// ----- // -----
func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x3x4x2x3xf32>) { func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x3x4x2x3xf32>) {
linalg.depthwise_conv_2d_nhwc_hwcm 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>) inits(%output : memref<2x3x4x2x3xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5, d2 + d6, d3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5, d2 + d6, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)>
// CHECK: func @depthwise_conv_2d_nhwc_hwcm // CHECK: func @depthwise_conv_2d_nhwc_hwcm
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<2x3x4x2x3xf32>) // CHECK-SAME: inits(%{{.+}} : memref<2x3x4x2x3xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x2x3x2x3xf32>) { func.func @depthwise_conv_2d_nhwc_hwcm(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x2x3x2x3xf32>) {
linalg.depthwise_conv_2d_nhwc_hwcm 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 : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>) ins(%input, %filter : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
outs(%output : memref<2x2x3x2x3xf32>) inits(%output : memref<2x2x3x2x3xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5 * 2, d2 + d6 * 2, d3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d5 * 2, d2 + d6 * 2, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d3, d4)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4)>
// CHECK: func @depthwise_conv_2d_nhwc_hwcm // CHECK: func @depthwise_conv_2d_nhwc_hwcm
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<2x2x3x2x3xf32>) // CHECK-SAME: inits(%{{.+}} : memref<2x2x3x2x3xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @depthwise_conv_2d_nhwc_hwc(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) { func.func @depthwise_conv_2d_nhwc_hwc(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: 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>) inits(%output: memref<1x56x56x96xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 2 + d5, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5, d3)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
// CHECK: func @depthwise_conv_2d_nhwc_hwc // CHECK: func @depthwise_conv_2d_nhwc_hwc
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x113x113x96xf32>, memref<3x3x96xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x56x56x96xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x56x56x96xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @conv_1d_nwc_wcf(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) { func.func @conv_1d_nwc_wcf(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {
linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>, linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>) ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%output: memref<?x?x?xf32>) inits (%output: memref<?x?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1 + d3, d4)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1 + d3, d4)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4, d2)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK: func @conv_1d_nwc_wcf // CHECK: func @conv_1d_nwc_wcf
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @conv_1d_ncw_fcw(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) { func.func @conv_1d_ncw_fcw(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {
linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>, linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>) ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%output: memref<?x?x?xf32>) inits (%output: memref<?x?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2 + d4)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2 + d4)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3, d4)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
// CHECK: func @conv_1d_ncw_fcw // CHECK: func @conv_1d_ncw_fcw
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @generalize_fill(%output: memref<?x?xf32>, %value : f32) { func.func @generalize_fill(%output: memref<?x?xf32>, %value : f32) {
linalg.fill ins(%value : f32) outs(%output : memref<?x?xf32>) linalg.fill ins(%value : f32) inits(%output : memref<?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: func @generalize_fill // CHECK: func @generalize_fill
// CHECK-SAME: (%[[ARG0:.+]]: memref<?x?xf32>, %[[VAL:.+]]: f32) // CHECK-SAME: (%[[ARG0:.+]]: memref<?x?xf32>, %[[VAL:.+]]: f32)
// CHECK: linalg.generic // CHECK: 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(%[[VAL]] : f32) // CHECK-SAME: ins(%[[VAL]] : f32)
// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32)
// CHECK-NEXT: linalg.yield %[[BBARG0]] : f32 // CHECK-NEXT: linalg.yield %[[BBARG0]] : f32
// ----- // -----
func.func @generalize_batch_matm_vec(%lhs : memref<?x?x?xi8>, %rhs: memref<?x?xi8>, %out: memref<?x?xf32>) { func.func @generalize_batch_matm_vec(%lhs : memref<?x?x?xi8>, %rhs: memref<?x?xi8>, %out: memref<?x?xf32>) {
linalg.batch_matvec ins(%lhs, %rhs: memref<?x?x?xi8>, memref<?x?xi8>) linalg.batch_matvec ins(%lhs, %rhs: memref<?x?x?xi8>, memref<?x?xi8>)
outs(%out: memref<?x?xf32>) inits(%out: memref<?x?xf32>)
return return
} }
// CHECK: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> // CHECK: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
// CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)> // CHECK: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
// CHECK: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> // CHECK: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK: @generalize_batch_matm_vec // CHECK: @generalize_batch_matm_vec
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]} // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xi8>, memref<?x?xi8>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xi8>, memref<?x?xi8>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?xf32>)
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: f32)
// CHECK: %[[BBARG0_F32:.+]] = arith.sitofp %[[BBARG0]] : i8 to f32 // CHECK: %[[BBARG0_F32:.+]] = arith.sitofp %[[BBARG0]] : i8 to f32
// CHECK: %[[BBARG1_F32:.+]] = arith.sitofp %[[BBARG1]] : i8 to f32 // CHECK: %[[BBARG1_F32:.+]] = arith.sitofp %[[BBARG1]] : i8 to f32
// CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0_F32]], %[[BBARG1_F32]] // CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0_F32]], %[[BBARG1_F32]]
// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] // CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]]
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32
// ----- // -----
func.func @batch_reduce_gemm(%lhs: memref<7x8x9xf32>, %rhs: memref<7x9x8xf32>, %out: memref<8x8xf32>) { func.func @batch_reduce_gemm(%lhs: memref<7x8x9xf32>, %rhs: memref<7x9x8xf32>, %out: memref<8x8xf32>) {
linalg.batch_reduce_matmul ins(%lhs, %rhs: memref<7x8x9xf32>, memref<7x9x8xf32>) linalg.batch_reduce_matmul ins(%lhs, %rhs: memref<7x8x9xf32>, memref<7x9x8xf32>)
outs(%out: memref<8x8xf32>) inits(%out: memref<8x8xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
// CHECK: @batch_reduce_gemm // CHECK: @batch_reduce_gemm
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]} // CHECK-SAME: iterator_types = ["reduction", "parallel", "parallel", "reduction"]}
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<7x8x9xf32>, memref<7x9x8xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<8x8xf32> // CHECK-SAME: inits(%{{.+}} : memref<8x8xf32>
// CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32) // CHECK: ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
// CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32 // CHECK: %[[MUL:.+]] = arith.mulf %[[BBARG0]], %[[BBARG1]] : f32
// CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32 // CHECK: %[[ADD:.+]] = arith.addf %[[BBARG2]], %[[MUL]] : f32
// CHECK: linalg.yield %[[ADD]] : f32 // CHECK: linalg.yield %[[ADD]] : f32

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

// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s // RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
// Verifies that different argument types is legal. // Verifies that different argument types is legal.
func.func @generalize_matmul_tensor_f16f64f32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_matmul_tensor_f16f64f32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-LABEL: @generalize_matmul_tensor_f16f64f32 // CHECK-LABEL: @generalize_matmul_tensor_f16f64f32
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)
// Verify floating point extension and truncation. // Verify floating point extension and truncation.
// CHECK-NEXT: %[[A_CAST:.+]] = arith.extf %[[A_ARG]] : f16 to f32 // CHECK-NEXT: %[[A_CAST:.+]] = arith.extf %[[A_ARG]] : f16 to f32
// CHECK-NEXT: %[[B_CAST:.+]] = arith.truncf %[[B_ARG]] : f64 to f32 // CHECK-NEXT: %[[B_CAST:.+]] = arith.truncf %[[B_ARG]] : f64 to f32
// CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32 // CHECK-NEXT: %[[MUL:.+]] = arith.mulf %[[A_CAST]], %[[B_CAST]] : f32
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[C_ARG]], %[[MUL]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// CHECK-NEXT: -> tensor<16x32xf32> // CHECK-NEXT: -> tensor<16x32xf32>
// ----- // -----
// Verifies that different argument types is legal. // Verifies that different argument types is legal.
func.func @generalize_matmul_tensor_i16i64i32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> { func.func @generalize_matmul_tensor_i16i64i32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32> inits(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32> return %0: tensor<16x32xi32>
} }
// CHECK-LABEL: @generalize_matmul_tensor_i16i64i32 // CHECK-LABEL: @generalize_matmul_tensor_i16i64i32
// CHECK: ^{{.*}}(%[[A_ARG:.+]]: i16, %[[B_ARG:.+]]: i64, %[[C_ARG:.+]]: i32) // CHECK: ^{{.*}}(%[[A_ARG:.+]]: i16, %[[B_ARG:.+]]: i64, %[[C_ARG:.+]]: i32)
// Verify signed integer extension and truncation. // Verify signed integer extension and truncation.
// CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i16 to i32 // CHECK-NEXT: %[[A_CAST:.+]] = arith.extsi %[[A_ARG]] : i16 to i32
// CHECK-NEXT: %[[B_CAST:.+]] = arith.trunci %[[B_ARG]] : i64 to i32 // CHECK-NEXT: %[[B_CAST:.+]] = arith.trunci %[[B_ARG]] : i64 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
// CHECK-NEXT: linalg.yield %[[ADD]] : i32 // CHECK-NEXT: linalg.yield %[[ADD]] : i32
// CHECK-NEXT: -> tensor<16x32xi32> // CHECK-NEXT: -> tensor<16x32xi32>
// ----- // -----
// Verifies that cast attributes control the cast operations used. // Verifies that cast attributes control the cast operations used.
func.func @generalize_matmul_tensor_i16i64i32_unsigned(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> { func.func @generalize_matmul_tensor_i16i64i32_unsigned(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>} %0 = linalg.matmul {cast = #linalg.type_fn<cast_unsigned>}
ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>) ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32> inits(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32> return %0: tensor<16x32xi32>
} }
// CHECK-LABEL: @generalize_matmul_tensor_i16i64i32_unsigned // CHECK-LABEL: @generalize_matmul_tensor_i16i64i32_unsigned
// CHECK: = arith.extui // CHECK: = arith.extui
// ----- // -----
func.func @generalize_matmul_tensor_i16i64f32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_matmul_tensor_i16i64f32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-LABEL: @generalize_matmul_tensor_i16i64f32 // CHECK-LABEL: @generalize_matmul_tensor_i16i64f32
// Verify signed integer to floating point cast. // Verify signed integer to floating point cast.
// CHECK: = arith.sitofp // CHECK: = arith.sitofp
// CHECK: = arith.sitofp // CHECK: = arith.sitofp
// ----- // -----
func.func @generalize_matmul_tensor_f16f64i32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> { func.func @generalize_matmul_tensor_f16f64i32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>) %0 = linalg.matmul ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32> inits(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32> return %0: tensor<16x32xi32>
} }
// CHECK-LABEL: @generalize_matmul_tensor_f16f64i32 // CHECK-LABEL: @generalize_matmul_tensor_f16f64i32
// Verify floating point to signed integer cast. // Verify floating point to signed integer cast.
// CHECK: = arith.fptosi // CHECK: = arith.fptosi
// CHECK: = arith.fptosi // CHECK: = arith.fptosi
// ----- // -----
func.func @generalize_matmul_unsigned_tensor_i16i64i32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> { func.func @generalize_matmul_unsigned_tensor_i16i64i32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>) %0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32> inits(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32> return %0: tensor<16x32xi32>
} }
// CHECK-LABEL: @generalize_matmul_unsigned_tensor_i16i64i32 // CHECK-LABEL: @generalize_matmul_unsigned_tensor_i16i64i32
// Verify unsigned integer extension and truncation. // Verify unsigned integer extension and truncation.
// CHECK: = arith.extui // CHECK: = arith.extui
// CHECK: = arith.trunci // CHECK: = arith.trunci
// ----- // -----
func.func @generalize_matmul_unsigned_tensor_i16i64f32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_matmul_unsigned_tensor_i16i64f32(%A : tensor<16x8xi16>, %B: tensor<8x32xi64>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>) %0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xi16>, tensor<8x32xi64>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-LABEL: @generalize_matmul_unsigned_tensor_i16i64f32 // CHECK-LABEL: @generalize_matmul_unsigned_tensor_i16i64f32
// Verify unsigned integer to floating point cast. // Verify unsigned integer to floating point cast.
// CHECK: = arith.uitofp // CHECK: = arith.uitofp
// CHECK: = arith.uitofp // CHECK: = arith.uitofp
// ----- // -----
func.func @generalize_matmul_unsigned_tensor_f16f64i32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> { func.func @generalize_matmul_unsigned_tensor_f16f64i32(%A : tensor<16x8xf16>, %B: tensor<8x32xf64>, %C: tensor<16x32xi32>) -> tensor<16x32xi32> {
%0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>) %0 = linalg.matmul_unsigned ins(%A, %B: tensor<16x8xf16>, tensor<8x32xf64>)
outs(%C: tensor<16x32xi32>) -> tensor<16x32xi32> inits(%C: tensor<16x32xi32>) -> tensor<16x32xi32>
return %0: tensor<16x32xi32> return %0: tensor<16x32xi32>
} }
// CHECK-LABEL: @generalize_matmul_unsigned_tensor_f16f64i32 // CHECK-LABEL: @generalize_matmul_unsigned_tensor_f16f64i32
// Verify floating point to unsigend integer cast. // Verify floating point to unsigend integer cast.
// CHECK: = arith.fptoui // CHECK: = arith.fptoui
// CHECK: = arith.fptoui // CHECK: = arith.fptoui
// ----- // -----
func.func @generalize_pooling_nhwc_max_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> { func.func @generalize_pooling_nhwc_max_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> {
%0 = linalg.pooling_nhwc_max {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_max {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) outs(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) inits(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
return %0: tensor<1x2x4x1xf32> return %0: tensor<1x2x4x1xf32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_max_f32 // CHECK-LABEL: @generalize_pooling_nhwc_max_f32
// CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32)
// CHECK-NEXT: %[[MAX:.+]] = arith.maxf %[[OUT_ARG]], %[[IN_ARG]] : f32 // CHECK-NEXT: %[[MAX:.+]] = arith.maxf %[[OUT_ARG]], %[[IN_ARG]] : f32
// CHECK-NEXT: linalg.yield %[[MAX]] : f32 // CHECK-NEXT: linalg.yield %[[MAX]] : f32
// CHECK-NEXT: -> tensor<1x2x4x1xf32> // CHECK-NEXT: -> tensor<1x2x4x1xf32>
// ----- // -----
func.func @generalize_pooling_nhwc_max_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> { func.func @generalize_pooling_nhwc_max_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
%0 = linalg.pooling_nhwc_max {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_max {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) inits(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
return %0: tensor<1x2x4x1xi32> return %0: tensor<1x2x4x1xi32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_max_i32 // CHECK-LABEL: @generalize_pooling_nhwc_max_i32
// Verify signed integer maximum. // Verify signed integer maximum.
// CHECK: = arith.maxsi // CHECK: = arith.maxsi
// ----- // -----
func.func @generalize_pooling_nhwc_max_unsigned_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> { func.func @generalize_pooling_nhwc_max_unsigned_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
%0 = linalg.pooling_nhwc_max_unsigned {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_max_unsigned {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) inits(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
return %0: tensor<1x2x4x1xi32> return %0: tensor<1x2x4x1xi32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_max_unsigned_i32 // CHECK-LABEL: @generalize_pooling_nhwc_max_unsigned_i32
// Verify unsigned integer minimum. // Verify unsigned integer minimum.
// CHECK: = arith.maxui // CHECK: = arith.maxui
// ----- // -----
func.func @generalize_pooling_nhwc_min_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> { func.func @generalize_pooling_nhwc_min_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> {
%0 = linalg.pooling_nhwc_min {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_min {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) outs(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) inits(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
return %0: tensor<1x2x4x1xf32> return %0: tensor<1x2x4x1xf32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_min_f32 // CHECK-LABEL: @generalize_pooling_nhwc_min_f32
// CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32)
// CHECK-NEXT: %[[MIN:.+]] = arith.minf %[[OUT_ARG]], %[[IN_ARG]] : f32 // CHECK-NEXT: %[[MIN:.+]] = arith.minf %[[OUT_ARG]], %[[IN_ARG]] : f32
// CHECK-NEXT: linalg.yield %[[MIN]] : f32 // CHECK-NEXT: linalg.yield %[[MIN]] : f32
// CHECK-NEXT: -> tensor<1x2x4x1xf32> // CHECK-NEXT: -> tensor<1x2x4x1xf32>
// ----- // -----
func.func @generalize_pooling_nhwc_min_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> { func.func @generalize_pooling_nhwc_min_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
%0 = linalg.pooling_nhwc_min {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_min {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) inits(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
return %0: tensor<1x2x4x1xi32> return %0: tensor<1x2x4x1xi32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_min_i32 // CHECK-LABEL: @generalize_pooling_nhwc_min_i32
// Verify signed integer minimum. // Verify signed integer minimum.
// CHECK: = arith.minsi // CHECK: = arith.minsi
// ----- // -----
func.func @generalize_pooling_nhwc_min_unsigned_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> { func.func @generalize_pooling_nhwc_min_unsigned_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
%0 = linalg.pooling_nhwc_min_unsigned {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_min_unsigned {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) inits(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
return %0: tensor<1x2x4x1xi32> return %0: tensor<1x2x4x1xi32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_min_unsigned_i32 // CHECK-LABEL: @generalize_pooling_nhwc_min_unsigned_i32
// Verify unsigned integer minimum. // Verify unsigned integer minimum.
// CHECK: = arith.minui // CHECK: = arith.minui
// ----- // -----
func.func @generalize_pooling_nhwc_sum_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> { func.func @generalize_pooling_nhwc_sum_f32(%input : tensor<1x4x16x1xf32>, %shape: tensor<2x2xf32>, %output: tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> {
%0 = linalg.pooling_nhwc_sum {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_sum {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) outs(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32> ins(%input, %shape : tensor<1x4x16x1xf32>, tensor<2x2xf32>) inits(%output : tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
return %0: tensor<1x2x4x1xf32> return %0: tensor<1x2x4x1xf32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_sum_f32 // CHECK-LABEL: @generalize_pooling_nhwc_sum_f32
// CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32) // CHECK: ^{{.*}}(%[[IN_ARG:.+]]: f32, %[[SHAPE_ARG:.+]]: f32, %[[OUT_ARG:.+]]: f32)
// CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[OUT_ARG]], %[[IN_ARG]] : f32 // CHECK-NEXT: %[[ADD:.+]] = arith.addf %[[OUT_ARG]], %[[IN_ARG]] : f32
// CHECK-NEXT: linalg.yield %[[ADD]] : f32 // CHECK-NEXT: linalg.yield %[[ADD]] : f32
// CHECK-NEXT: -> tensor<1x2x4x1xf32> // CHECK-NEXT: -> tensor<1x2x4x1xf32>
// ----- // -----
func.func @generalize_pooling_nhwc_sum_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> { func.func @generalize_pooling_nhwc_sum_i32(%input : tensor<1x4x16x1xi32>, %shape: tensor<2x2xi32>, %output: tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> {
%0 = linalg.pooling_nhwc_sum {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>} %0 = linalg.pooling_nhwc_sum {dilations = dense<[1, 2]> : tensor<2xi64>, strides = dense<[2, 4]> : tensor<2xi64>}
ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) outs(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32> ins(%input, %shape : tensor<1x4x16x1xi32>, tensor<2x2xi32>) inits(%output : tensor<1x2x4x1xi32>) -> tensor<1x2x4x1xi32>
return %0: tensor<1x2x4x1xi32> return %0: tensor<1x2x4x1xi32>
} }
// CHECK-LABEL: @generalize_pooling_nhwc_sum_i32 // CHECK-LABEL: @generalize_pooling_nhwc_sum_i32
// CHECK: ^{{.*}}(%[[IN_ARG:.+]]: i32, %[[SHAPE_ARG:.+]]: i32, %[[OUT_ARG:.+]]: i32) // CHECK: ^{{.*}}(%[[IN_ARG:.+]]: i32, %[[SHAPE_ARG:.+]]: i32, %[[OUT_ARG:.+]]: i32)
// CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[OUT_ARG]], %[[IN_ARG]] : i32 // CHECK-NEXT: %[[ADD:.+]] = arith.addi %[[OUT_ARG]], %[[IN_ARG]] : i32
// CHECK-NEXT: linalg.yield %[[ADD]] : i32 // CHECK-NEXT: linalg.yield %[[ADD]] : i32
// CHECK-NEXT: -> tensor<1x2x4x1xi32> // CHECK-NEXT: -> tensor<1x2x4x1xi32>
// ----- // -----
func.func @generalize_fill_0d(%value: f64, %O: tensor<f32>) -> tensor<f32> { func.func @generalize_fill_0d(%value: f64, %O: tensor<f32>) -> tensor<f32> {
%0 = linalg.fill ins(%value: f64) outs(%O : tensor<f32>) -> tensor<f32> %0 = linalg.fill ins(%value: f64) inits(%O : tensor<f32>) -> tensor<f32>
return %0: tensor<f32> return %0: tensor<f32>
} }
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<() -> ()> // CHECK-DAG: #[[$MAP0:.+]] = affine_map<() -> ()>
// CHECK-LABEL: @generalize_fill_0d // CHECK-LABEL: @generalize_fill_0d
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]] // CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]
// CHECK-SAME: iterator_types = [] // CHECK-SAME: iterator_types = []
// ----- // -----
func.func @generalize_fill_2d(%value: f64, %O: memref<16x32xf32>) { func.func @generalize_fill_2d(%value: f64, %O: memref<16x32xf32>) {
linalg.fill ins(%value: f64) outs(%O : memref<16x32xf32>) linalg.fill ins(%value: f64) inits(%O : memref<16x32xf32>)
return return
} }
// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1) -> ()> // CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1) -> ()>
// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)> // CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)>
// CHECK-LABEL: @generalize_fill // CHECK-LABEL: @generalize_fill
// CHECK: linalg.generic // CHECK: 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"]
// ----- // -----
func.func @generalize_index(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_index(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32> %0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) inits(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-LABEL: @generalize_index // CHECK-LABEL: @generalize_index
// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index // CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index
// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index // CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index
// CHECK-DAG: %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32 // CHECK-DAG: %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32
// CHECK-DAG: %[[IDX1_CAST:.+]] = arith.index_cast %[[IDX1]] : index to i32 // CHECK-DAG: %[[IDX1_CAST:.+]] = arith.index_cast %[[IDX1]] : index to i32
// ----- // -----
func.func @generalize_const(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @generalize_const(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32> %0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) inits(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-LABEL: @generalize_const // CHECK-LABEL: @generalize_const
// CHECK-DAG: %[[CST0:.+]] = arith.constant 1103515245 : i32 // CHECK-DAG: %[[CST0:.+]] = arith.constant 1103515245 : i32
// CHECK-DAG: %[[CST1:.+]] = arith.constant 12345 : i32 // CHECK-DAG: %[[CST1:.+]] = arith.constant 12345 : i32
// CHECK-DAG: %[[CST2:.+]] = arith.constant 2.3283063999999999E-10 : f64 // CHECK-DAG: %[[CST2:.+]] = arith.constant 2.3283063999999999E-10 : f64
// ----- // -----
// Verifies the default value of the fun attribute is an exp op. // Verifies the default value of the fun attribute is an exp op.
func.func @generalize_elemwise_exp(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_exp(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> %0 = linalg.elemwise_unary ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_exp // CHECK-LABEL: @generalize_elemwise_exp
// CHECK: = math.exp // CHECK: = math.exp
// ----- // -----
// Verifies the fun attribute controls the unary function used. // Verifies the fun attribute controls the unary function used.
func.func @generalize_elemwise_log(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_log(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary {fun = #linalg.unary_fn<log>} %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<log>}
ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_log // CHECK-LABEL: @generalize_elemwise_log
// CHECK: = math.log // CHECK: = math.log
// ----- // -----
// Verifies the fun attribute controls the unary function used. // Verifies the fun attribute controls the unary function used.
func.func @generalize_elemwise_abs(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_abs(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary {fun = #linalg.unary_fn<abs>} %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<abs>}
ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_abs // CHECK-LABEL: @generalize_elemwise_abs
// CHECK: = math.absf // CHECK: = math.absf
// ----- // -----
// Verifies the fun attribute controls the unary function used. // Verifies the fun attribute controls the unary function used.
func.func @generalize_elemwise_ceil(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_ceil(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary {fun = #linalg.unary_fn<ceil>} %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<ceil>}
ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_ceil // CHECK-LABEL: @generalize_elemwise_ceil
// CHECK: = math.ceil // CHECK: = math.ceil
// ----- // -----
// Verifies the fun attribute controls the unary function used. // Verifies the fun attribute controls the unary function used.
func.func @generalize_elemwise_floor(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_floor(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary {fun = #linalg.unary_fn<floor>} %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<floor>}
ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_floor // CHECK-LABEL: @generalize_elemwise_floor
// CHECK: = math.floor // CHECK: = math.floor
// ----- // -----
// Verifies the fun attribute controls the unary function used. // Verifies the fun attribute controls the unary function used.
func.func @generalize_elemwise_negf(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_negf(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_unary {fun = #linalg.unary_fn<negf>} %0 = linalg.elemwise_unary {fun = #linalg.unary_fn<negf>}
ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_negf // CHECK-LABEL: @generalize_elemwise_negf
// CHECK: = arith.negf // CHECK: = arith.negf
// ----- // -----
// Verifies the default value of the fun attribute is an add op. // Verifies the default value of the fun attribute is an add op.
func.func @generalize_elemwise_add(%lhs : tensor<4x8xf32>, %rhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_add(%lhs : tensor<4x8xf32>, %rhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_binary ins(%lhs, %rhs: tensor<4x8xf32>, tensor<4x8xf32>) %0 = linalg.elemwise_binary ins(%lhs, %rhs: tensor<4x8xf32>, tensor<4x8xf32>)
outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_add // CHECK-LABEL: @generalize_elemwise_add
// CHECK: = arith.addf // CHECK: = arith.addf
// ----- // -----
// Verifies the fun attribute controls the binary function used. // Verifies the fun attribute controls the binary function used.
func.func @generalize_elemwise_mul(%lhs : tensor<4x8xf32>, %rhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_mul(%lhs : tensor<4x8xf32>, %rhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_binary {fun = #linalg.binary_fn<mul>} %0 = linalg.elemwise_binary {fun = #linalg.binary_fn<mul>}
ins(%lhs, %rhs: tensor<4x8xf32>, tensor<4x8xf32>) ins(%lhs, %rhs: tensor<4x8xf32>, tensor<4x8xf32>)
outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_mul // CHECK-LABEL: @generalize_elemwise_mul
// CHECK: = arith.mulf // CHECK: = arith.mulf
// ----- // -----
// Verifies pointwise ops support rank zero input tensors // Verifies pointwise ops support rank zero input tensors
func.func @generalize_elemwise_rank_zero(%lhs : tensor<f32>, %rhs : tensor<f32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_elemwise_rank_zero(%lhs : tensor<f32>, %rhs : tensor<f32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.elemwise_binary {fun = #linalg.binary_fn<sub>} %0 = linalg.elemwise_binary {fun = #linalg.binary_fn<sub>}
ins(%lhs, %rhs: tensor<f32>, tensor<f32>) ins(%lhs, %rhs: tensor<f32>, tensor<f32>)
outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_elemwise_rank_zero // CHECK-LABEL: @generalize_elemwise_rank_zero
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: iterator_types = ["parallel", "parallel"] // CHECK-SAME: iterator_types = ["parallel", "parallel"]
// CHECK: = arith.subf // CHECK: = arith.subf
// ----- // -----
// Verifies the fun attribute controls the binary function used. // Verifies the fun attribute controls the binary function used.
func.func @generalize_copy(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> { func.func @generalize_copy(%lhs : tensor<4x8xf32>, %output : tensor<4x8xf32>) -> tensor<4x8xf32> {
%0 = linalg.copy ins(%lhs: tensor<4x8xf32>) outs(%output: tensor<4x8xf32>) -> tensor<4x8xf32> %0 = linalg.copy ins(%lhs: tensor<4x8xf32>) inits(%output: tensor<4x8xf32>) -> tensor<4x8xf32>
return %0: tensor<4x8xf32> return %0: tensor<4x8xf32>
} }
// CHECK-LABEL: @generalize_copy // CHECK-LABEL: @generalize_copy
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-NEXT: ^bb0(%[[I:[0-9a-zA-Z]*]]: f32 // CHECK-NEXT: ^bb0(%[[I:[0-9a-zA-Z]*]]: f32
// CHECK-NEXT: linalg.yield %[[I]] // CHECK-NEXT: linalg.yield %[[I]]

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:.*]] = tensor.empty() : 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) inits(%[[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:.*]] = tensor.empty(%[[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) inits(%[[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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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/inlining.mlir

Show All 17 Linesfunc.func @inline_into(%arg0: memref<?xf32>) {
call @inlined_fn(%arg0) : (memref<?xf32>) -> () call @inlined_fn(%arg0) : (memref<?xf32>) -> ()
return return
} }
func.func @inlined_fn(%arg0: memref<?xf32>) { func.func @inlined_fn(%arg0: memref<?xf32>) {
// CHECK: linalg.generic // CHECK: linalg.generic
linalg.generic #trait linalg.generic #trait
ins(%arg0 : memref<?xf32>) ins(%arg0 : memref<?xf32>)
outs(%arg0 : memref<?xf32>) { inits(%arg0 : memref<?xf32>) {
^bb(%0 : f32, %1 : f32) : ^bb(%0 : f32, %1 : f32) :
%2 = arith.addf %0, %0: f32 %2 = arith.addf %0, %0: f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} }
return return
} }

mlir/test/Dialect/Linalg/invalid.mlir

Show All 30 Lines
// ----- // -----
func.func @index_dim_lower_than_number_of_loops(%arg0: memref<f32>) { func.func @index_dim_lower_than_number_of_loops(%arg0: memref<f32>) {
// expected-error @+6 {{op expected dim (2) to be lower than the number of loops (0) of the enclosing LinalgOp}} // expected-error @+6 {{op expected dim (2) to be lower than the number of loops (0) of the enclosing LinalgOp}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
^bb(%0: f32): ^bb(%0: f32):
linalg.index 2 : index linalg.index 2 : index
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
} }
// ----- // -----
func.func @index_dim_negative(%arg0: memref<f32>) { func.func @index_dim_negative(%arg0: memref<f32>) {
// expected-error @+6 {{op attribute 'dim' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}} // expected-error @+6 {{op attribute 'dim' failed to satisfy constraint: 64-bit signless integer attribute whose minimum value is 0}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
^bb(%0: f32): ^bb(%0: f32):
linalg.index -1 : index linalg.index -1 : index
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
} }
// ----- // -----
func.func @generic_no_region(%arg0: memref<f32>) { func.func @generic_no_region(%arg0: memref<f32>) {
// expected-error @+4 {{expected '{' to begin a region}} // expected-error @+4 {{expected '{' to begin a region}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> (0)> ], indexing_maps = [ affine_map<() -> (0)> ],
iterator_types = [] iterator_types = []
} ins(%arg0 : memref<f32>) } ins(%arg0 : memref<f32>)
} }
// ----- // -----
func.func @generic_mismatched_num_returns(%arg0: memref<f32>) { func.func @generic_mismatched_num_returns(%arg0: memref<f32>) {
// expected-error @+6 {{op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (0)}} // expected-error @+6 {{op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (0)}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
^bb(%0: f32): ^bb(%0: f32):
linalg.yield linalg.yield
} }
} }
// ----- // -----
func.func @generic_wrong_dim_in_map(%arg0: memref<1xi32>) { func.func @generic_wrong_dim_in_map(%arg0: memref<1xi32>) {
// expected-error @+1 {{op expected indexing_map #0 to have 1 dim(s) to match the number of loops}} // expected-error @+1 {{op expected indexing_map #0 to have 1 dim(s) to match the number of loops}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> (0)> ], indexing_maps = [ affine_map<() -> (0)> ],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%arg0 : memref<1xi32>) { inits(%arg0 : memref<1xi32>) {
^bb(%i : i32): ^bb(%i : i32):
linalg.yield %i : i32 linalg.yield %i : i32
} }
} }
// ----- // -----
func.func @generic_wrong_iterator(%arg0: memref<1xi32>) { func.func @generic_wrong_iterator(%arg0: memref<1xi32>) {
// expected-error @+1 {{op unexpected iterator_type (random)}} // expected-error @+1 {{op unexpected iterator_type (random)}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<(i) -> (i)> ], indexing_maps = [ affine_map<(i) -> (i)> ],
iterator_types = ["random"]} iterator_types = ["random"]}
outs(%arg0 : memref<1xi32>) { inits(%arg0 : memref<1xi32>) {
^bb(%i : i32): ^bb(%i : i32):
linalg.yield %i : i32 linalg.yield %i : i32
} }
} }
// ----- // -----
func.func @generic_one_d_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) { func.func @generic_one_d_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
// expected-error @+1 {{expected operand rank (1) to match the result rank of indexing_map #0 (2)}} // expected-error @+1 {{expected operand rank (1) to match the result rank of indexing_map #0 (2)}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> (0, 0)> ], indexing_maps = [ affine_map<() -> (0, 0)> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) { inits(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
^bb(%f : f32): ^bb(%f : f32):
linalg.yield %f: f32 linalg.yield %f: f32
} }
} }
// ----- // -----
func.func @generic_scalar_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) { func.func @generic_scalar_view(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
// expected-error @+1 {{expected operand rank (0) to match the result rank of indexing_map #0 (1)}} // expected-error @+1 {{expected operand rank (0) to match the result rank of indexing_map #0 (1)}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> (0)>, affine_map<() -> (0, 0)> ], indexing_maps = [ affine_map<() -> (0)>, affine_map<() -> (0, 0)> ],
iterator_types = []} iterator_types = []}
ins(%cst : f32) ins(%cst : f32)
outs(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) { inits(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
^bb(%0 : f32, %1 : f32): ^bb(%0 : f32, %1 : f32):
linalg.yield %0: f32 linalg.yield %0: f32
} }
} }
// ----- // -----
func.func @generic_result_0_element_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) { func.func @generic_result_0_element_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
// expected-error @+7 {{'linalg.yield' op type of yield operand 1 ('i4') doesn't match the element type of the enclosing linalg.generic op ('f32')}} // expected-error @+7 {{'linalg.yield' op type of yield operand 1 ('i4') doesn't match the element type of the enclosing linalg.generic op ('f32')}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<(i) -> (i)> ], indexing_maps = [ affine_map<(i) -> (i)> ],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) { inits(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
^bb(%0: f32): ^bb(%0: f32):
%1 = arith.constant 1: i4 %1 = arith.constant 1: i4
linalg.yield %1: i4 linalg.yield %1: i4
} }
} }
// ----- // -----
func.func @generic_singular_maps(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>, %arg1: memref<?xf32, affine_map<(i)[off]->(off + i)>>) { func.func @generic_singular_maps(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>, %arg1: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
// expected-error @+1 {{expected the shape-to-loops map to be non-null}} // expected-error @+1 {{expected the shape-to-loops map to be non-null}}
linalg.generic { linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(i, j) -> (i + j)>, affine_map<(i, j) -> (i + j)>,
affine_map<(i, j) -> (i + j)> affine_map<(i, j) -> (i + j)>
], ],
iterator_types = ["parallel","parallel"]} iterator_types = ["parallel","parallel"]}
ins(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) ins(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>)
outs(%arg1 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) { inits(%arg1 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
^bb(%0: f32, %1: f32): ^bb(%0: f32, %1: f32):
linalg.yield %1: f32 linalg.yield %1: f32
} }
} }
//////////////////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////////
///////////////////////////// Region tests ///////////////////////////////////// ///////////////////////////// Region tests /////////////////////////////////////
//////////////////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////////
// ----- // -----
func.func @generic_empty_region(%arg0: memref<f32>) { func.func @generic_empty_region(%arg0: memref<f32>) {
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
// expected-error @+1 {{op expects region #0 to have 0 or 1 blocks}} // expected-error @+1 {{op expects region #0 to have 0 or 1 blocks}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()>, affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()>, affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
ins(%arg0 : memref<f32>) ins(%arg0 : memref<f32>)
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
^bb1: ^bb1:
linalg.yield %f0: f32 linalg.yield %f0: f32
^bb2: ^bb2:
linalg.yield %f0: f32 linalg.yield %f0: f32
} }
} }
// ----- // -----
func.func @generic_empty_region(%arg0: memref<f32>) { func.func @generic_empty_region(%arg0: memref<f32>) {
%f0 = arith.constant 0.0: f32 %f0 = arith.constant 0.0: f32
// expected-error @+1 {{op expects to have 1 region with 1 block}} // expected-error @+1 {{op expects to have 1 region with 1 block}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> , affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> , affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
ins(%arg0 : memref<f32>) ins(%arg0 : memref<f32>)
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
} }
} }
// ----- // -----
func.func @generic_mismatched_num_arguments(%arg0: memref<f32>) { func.func @generic_mismatched_num_arguments(%arg0: memref<f32>) {
// expected-error @+6 {{'linalg.yield' op expected number of yield values (2) to match the number of operands of the enclosing LinalgOp (1)}} // expected-error @+6 {{'linalg.yield' op expected number of yield values (2) to match the number of operands of the enclosing LinalgOp (1)}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()>, affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()>, affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0, %arg0 : memref<f32>, memref<f32>) { inits(%arg0, %arg0 : memref<f32>, memref<f32>) {
^bb(%f: f32): ^bb(%f: f32):
linalg.yield %f: f32 linalg.yield %f: f32
} }
} }
// ----- // -----
func.func @generic_shaped_operand_block_arg_type(%arg0: memref<f32>) { func.func @generic_shaped_operand_block_arg_type(%arg0: memref<f32>) {
// expected-error @+6 {{'linalg.yield' op type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}} // expected-error @+6 {{'linalg.yield' op type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : memref<f32>) { inits(%arg0 : memref<f32>) {
^bb(%i: i1): ^bb(%i: i1):
linalg.yield %i : i1 linalg.yield %i : i1
} }
} }
// ----- // -----
func.func @generic_scalar_operand_block_arg_type(%arg0: tensor<f32>) { func.func @generic_scalar_operand_block_arg_type(%arg0: tensor<f32>) {
// expected-error @+6 {{'linalg.yield' op type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}} // expected-error @+6 {{'linalg.yield' op type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<() -> ()> ], indexing_maps = [ affine_map<() -> ()> ],
iterator_types = []} iterator_types = []}
outs(%arg0 : tensor<f32>) { inits(%arg0 : tensor<f32>) {
^bb(%i: i1): ^bb(%i: i1):
linalg.yield %i : i1 linalg.yield %i : i1
} -> tensor<f32> } -> tensor<f32>
} }
// ----- // -----
func.func @generic_result_0_element_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) { func.func @generic_result_0_element_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
// expected-error @+7 {{type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}} // expected-error @+7 {{type of yield operand 1 ('i1') doesn't match the element type of the enclosing linalg.generic op ('f32')}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<(i) -> (i)> ], indexing_maps = [ affine_map<(i) -> (i)> ],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) { inits(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) {
^bb(%i: f32): ^bb(%i: f32):
%0 = arith.constant 0: i1 %0 = arith.constant 0: i1
linalg.yield %0: i1 linalg.yield %0: i1
} }
} }
// ----- // -----
func.func @generic_result_tensor_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>, func.func @generic_result_tensor_type(%arg0: memref<?xf32, affine_map<(i)[off]->(off + i)>>,
%arg1: tensor<?xf32>) { %arg1: tensor<?xf32>) {
// expected-error @+1 {{expected type of operand #1 ('tensor<?xf32>') to match type of corresponding result ('tensor<f32>')}} // expected-error @+1 {{expected type of operand #1 ('tensor<?xf32>') to match type of corresponding result ('tensor<f32>')}}
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ affine_map<(i) -> (i)> , affine_map<(i) -> (i)> ], indexing_maps = [ affine_map<(i) -> (i)> , affine_map<(i) -> (i)> ],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
ins(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>) ins(%arg0 : memref<?xf32, affine_map<(i)[off]->(off + i)>>)
outs(%arg1 : tensor<?xf32>) { inits(%arg1 : tensor<?xf32>) {
^bb(%i: f32, %j: f32): ^bb(%i: f32, %j: f32):
linalg.yield %i: f32 linalg.yield %i: f32
} -> tensor<f32> } -> tensor<f32>
} }
// ----- // -----
func.func @generic(%arg0: memref<?x?xf32>) { func.func @generic(%arg0: memref<?x?xf32>) {
// expected-error @+6 {{block with no terminator, has %0 = "arith.addf"(%arg1, %arg1) : (f32, f32) -> f32}} // expected-error @+6 {{block with no terminator, has %0 = "arith.addf"(%arg1, %arg1) : (f32, f32) -> f32}}
linalg.generic { linalg.generic {
indexing_maps = [ affine_map<(i, j) -> (i, j)> ], indexing_maps = [ affine_map<(i, j) -> (i, j)> ],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%arg0 : memref<?x?xf32>) { inits(%arg0 : memref<?x?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
%1 = arith.addf %0, %0: f32 %1 = arith.addf %0, %0: f32
} }
return return
} }
// ----- // -----
// This test is currently disabled: subject to verifier ordering issues. // This test is currently disabled: subject to verifier ordering issues.
// Instead, when the ranks are not greater than 2, an assertion will be triggered // Instead, when the ranks are not greater than 2, an assertion will be triggered
// in LinalgStructuredOps.td::ConvOp::iterator_types() for now because the // in LinalgStructuredOps.td::ConvOp::iterator_types() for now because the
// verifier inspects the iterator_types. This is slated to become an // verifier inspects the iterator_types. This is slated to become an
// autogenerated op in the future, alleviating the issue. // autogenerated op in the future, alleviating the issue.
// func @conv_rank_limit(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) { // func @conv_rank_limit(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
// // DISABLED_expected -error @+1 {{expects memref ranks to be greater than 2}} // // DISABLED_expected -error @+1 {{expects memref ranks to be greater than 2}}
// linalg.conv(%arg0, %arg1, %arg2) : memref<?xf32>, memref<?xf32>, memref<?xf32> // linalg.conv(%arg0, %arg1, %arg2) : memref<?xf32>, memref<?xf32>, memref<?xf32>
// } // }
// //
// // ----- // // -----
func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?xf32>, %c3: memref<?x?x?xf32>) { func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?xf32>, %c3: memref<?x?x?xf32>) {
// expected-error @+1 {{expected operand rank (2) to match the result rank of indexing_map #1 (3)}} // expected-error @+1 {{expected operand rank (2) to match the result rank of indexing_map #1 (3)}}
linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?xf32>) linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?xf32>)
outs(%c3 : memref<?x?x?xf32>) inits(%c3 : memref<?x?x?xf32>)
return return
} }
// ----- // -----
func.func @incorrect_region_arg_count(%m: memref<?x?xf32>) { func.func @incorrect_region_arg_count(%m: memref<?x?xf32>) {
// expected-error @+3 {{region expects 3 args, got 2}} // expected-error @+3 {{region expects 3 args, got 2}}
%res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>) %res = linalg.matmul ins(%m, %m : memref<?x?xf32>, memref<?x?xf32>)
-> (tensor<?x?xf32>, tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>)
return return
} }
// ----- // -----
func.func @matching_inits(%m: memref<?x?xf32>, %t: tensor<?x?xf32>) { func.func @matching_inits(%m: memref<?x?xf32>, %t: tensor<?x?xf32>) {
// expected-error @+1 {{expected type of operand #2 ('tensor<?x?xf32>') to match type of corresponding result ('tensor<?xf32>')}} // expected-error @+1 {{expected type of operand #2 ('tensor<?x?xf32>') to match type of corresponding result ('tensor<?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>) inits(%t : tensor<?x?xf32>)
-> tensor<?xf32> -> tensor<?xf32>
return 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 = tensor.empty(%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) inits(%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>
{ {
// expected-error @+1 {{expected the number of results (1) to be equal to the number of output tensors (0)}} // expected-error @+1 {{expected the number of results (1) to be equal to the number of output tensors (0)}}
%0 = linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?x?xf32>) -> tensor<?x?xf32> %0 = linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<?x?xf32>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
func.func @illegal_fill_tensor_with_memref_return func.func @illegal_fill_tensor_with_memref_return
(%arg0 : tensor<?x?xf32>, %arg1 : f32) -> memref<?x?xf32> (%arg0 : tensor<?x?xf32>, %arg1 : f32) -> memref<?x?xf32>
{ {
// expected-error @+1 {{result #0 must be ranked tensor of any type values, but got 'memref<?x?xf32>'}} // expected-error @+1 {{result #0 must be ranked tensor of any type values, but got 'memref<?x?xf32>'}}
%0 = linalg.fill ins(%arg1 : f32) outs(%arg0 : tensor<?x?xf32>) -> memref<?x?xf32> %0 = linalg.fill ins(%arg1 : f32) inits(%arg0 : tensor<?x?xf32>) -> memref<?x?xf32>
return %0 : memref<?x?xf32> return %0 : memref<?x?xf32>
} }
// ----- // -----
func.func @invalid_static_matmul(%arg0: memref<2x4xf32>, %arg1: memref<3x4xf32>, %arg2: memref<2x4xf32>) { func.func @invalid_static_matmul(%arg0: memref<2x4xf32>, %arg1: memref<3x4xf32>, %arg2: memref<2x4xf32>) {
// expected-error @+1 {{inferred input/output operand #1 has shape's dimension #0 to be 4, but found 3}} // expected-error @+1 {{inferred input/output operand #1 has shape's dimension #0 to be 4, but found 3}}
linalg.matmul ins(%arg0, %arg1 : memref<2x4xf32>, memref<3x4xf32>) linalg.matmul ins(%arg0, %arg1 : memref<2x4xf32>, memref<3x4xf32>)
outs(%arg2 :memref<2x4xf32>) inits(%arg2 :memref<2x4xf32>)
return return
} }
// ----- // -----
func.func @invalid_static_2d_conv(%input : memref<1x3x4x2xf32>, %filter: memref<3x2x2x1xf32>, %output: memref<1x2x3x1xf32>) { func.func @invalid_static_2d_conv(%input : memref<1x3x4x2xf32>, %filter: memref<3x2x2x1xf32>, %output: memref<1x2x3x1xf32>) {
// expected-error @+1 {{inferred input/output operand #0 has shape's dimension #1 to be greater than or equal to 4, but found 3}} // expected-error @+1 {{inferred input/output operand #0 has shape's dimension #1 to be greater than or equal to 4, but found 3}}
linalg.conv_2d_nhwc_hwcf linalg.conv_2d_nhwc_hwcf
{ dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %filter : memref<1x3x4x2xf32>, memref<3x2x2x1xf32>) ins(%input, %filter : memref<1x3x4x2xf32>, memref<3x2x2x1xf32>)
outs(%output : memref<1x2x3x1xf32>) inits(%output : memref<1x2x3x1xf32>)
return return
} }
// ----- // -----
#attrs = { #attrs = {
indexing_maps = [ indexing_maps = [
affine_map<(i) -> (3 - i)>, affine_map<(i) -> (3 - i)>,
affine_map<(i) -> (i)> affine_map<(i) -> (i)>
], ],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
func.func @invalid_reverse(%A: memref<5xf32>, %B: memref<5xf32>) { func.func @invalid_reverse(%A: memref<5xf32>, %B: memref<5xf32>) {
// expected-error @+1 {{unexpected result less than 0 at expression #0 in}} // expected-error @+1 {{unexpected result less than 0 at expression #0 in}}
linalg.generic #attrs ins(%A: memref<5xf32>) outs(%B: memref<5xf32>) { linalg.generic #attrs ins(%A: memref<5xf32>) inits(%B: memref<5xf32>) {
^bb0(%a: f32, %b: f32): ^bb0(%a: f32, %b: f32):
linalg.yield %a : f32 linalg.yield %a : f32
} }
return return
} }
// ----- // -----
func.func @map_binary_wrong_yield_operands( func.func @map_binary_wrong_yield_operands(
%lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>) %lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>)
-> tensor<64xf32> { -> tensor<64xf32> {
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
// expected-error @+1{{'linalg.yield' op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (2)}} // expected-error @+1{{'linalg.yield' op expected number of yield values (1) to match the number of operands of the enclosing LinalgOp (2)}}
linalg.yield %0, %0: f32, f32 linalg.yield %0, %0: f32, f32
} }
func.return %add : tensor<64xf32> func.return %add : tensor<64xf32>
} }
// ----- // -----
func.func @map_input_mapper_arity_mismatch( func.func @map_input_mapper_arity_mismatch(
%lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>) %lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>)
-> tensor<64xf32> { -> tensor<64xf32> {
// expected-error@+1{{'linalg.map' op expects number of operands to match the arity of mapper, but got: 2 and 3}} // expected-error@+1{{'linalg.map' op expects number of operands to match the arity of mapper, but got: 2 and 3}}
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32, %extra_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32, %extra_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %add : tensor<64xf32> func.return %add : tensor<64xf32>
} }
// ----- // -----
func.func @map_input_mapper_type_mismatch( func.func @map_input_mapper_type_mismatch(
%lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>) %lhs: tensor<64xf32>, %rhs: tensor<64xf32>, %init: tensor<64xf32>)
-> tensor<64xf32> { -> tensor<64xf32> {
// expected-error@+1{{'linalg.map' op expected element type of input 'f32' to match bbArg type 'f64'}} // expected-error@+1{{'linalg.map' op expected element type of input 'f32' to match bbArg type 'f64'}}
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%lhs_elem: f64, %rhs_elem: f64) { (%lhs_elem: f64, %rhs_elem: f64) {
%0 = arith.addf %lhs_elem, %rhs_elem: f64 %0 = arith.addf %lhs_elem, %rhs_elem: f64
linalg.yield %0: f64 linalg.yield %0: f64
} }
func.return %add : tensor<64xf32> func.return %add : tensor<64xf32>
} }
// ----- // -----
func.func @map_input_output_shape_mismatch( func.func @map_input_output_shape_mismatch(
%lhs: tensor<64x64xf32>, %rhs: tensor<64x64xf32>, %init: tensor<32xf32>) %lhs: tensor<64x64xf32>, %rhs: tensor<64x64xf32>, %init: tensor<32xf32>)
-> tensor<32xf32> { -> tensor<32xf32> {
// expected-error@+1{{'linalg.map' op expected shape of input (64, 64) to match shape of output (32)}} // expected-error@+1{{'linalg.map' op expected shape of input (64, 64) to match shape of output (32)}}
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs : tensor<64x64xf32>, tensor<64x64xf32>) ins(%lhs, %rhs : tensor<64x64xf32>, tensor<64x64xf32>)
outs(%init:tensor<32xf32>) inits(%init:tensor<32xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %add : tensor<32xf32> func.return %add : tensor<32xf32>
} }
// ----- // -----
func.func @reduce_input_vs_init_dimension_mismatch( func.func @reduce_input_vs_init_dimension_mismatch(
%input: tensor<16x32x64xf32>, %input: tensor<16x32x64xf32>,
%init: tensor<16x64xf32>) -> tensor<16x64xf32> { %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
// expected-error @+1 {{'linalg.reduce' op init dimensions [16, 64] doesn't match input dimensions after reduction [16, 32]}} // expected-error @+1 {{'linalg.reduce' op init dimensions [16, 64] doesn't match input dimensions after reduction [16, 32]}}
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [2] dimensions = [2]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16x64xf32> func.return %reduce : tensor<16x64xf32>
} }
// ----- // -----
func.func @reduce_dimensions_out_of_range(%input: tensor<16x32x64xf32>, func.func @reduce_dimensions_out_of_range(%input: tensor<16x32x64xf32>,
%init: tensor<16x64xf32>) -> tensor<16x64xf32> { %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
// expected-error @+1 {{'linalg.reduce' op dimensions for reduction should be in the range [0, 2].}} // expected-error @+1 {{'linalg.reduce' op dimensions for reduction should be in the range [0, 2].}}
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [3] dimensions = [3]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16x64xf32> func.return %reduce : tensor<16x64xf32>
} }
// ----- // -----
func.func @reduce_duplicate_dimensions(%input: tensor<16x32x64xf32>, func.func @reduce_duplicate_dimensions(%input: tensor<16x32x64xf32>,
%init: tensor<16xf32>) -> tensor<16xf32> { %init: tensor<16xf32>) -> tensor<16xf32> {
// expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}} // expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}}
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16xf32>) inits(%init:tensor<16xf32>)
dimensions = [1, 1] dimensions = [1, 1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16xf32> func.return %reduce : tensor<16xf32>
} }
// ----- // -----
func.func @reduce_non_increasing_dimensions(%input: tensor<16x32x64xf32>, func.func @reduce_non_increasing_dimensions(%input: tensor<16x32x64xf32>,
%init: tensor<16xf32>) -> tensor<16xf32> { %init: tensor<16xf32>) -> tensor<16xf32> {
// expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}} // expected-error @+1 {{'linalg.reduce' op attribute 'dimensions' failed to satisfy constraint: i64 dense array attribute should be in increasing order}}
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16xf32>) inits(%init:tensor<16xf32>)
dimensions = [2, 1] dimensions = [2, 1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16xf32> func.return %reduce : tensor<16xf32>
} }
// ----- // -----
func.func @reduce_reduced_input_init_rank_mismatch(%input: tensor<16x32x64xf32>, func.func @reduce_reduced_input_init_rank_mismatch(%input: tensor<16x32x64xf32>,
%init: tensor<16x64xf32>) -> tensor<16x64xf32> { %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
// expected-error @+1 {{'linalg.reduce' op number of dimensions after reduction 1 doesn't match the init rank 2}} // expected-error @+1 {{'linalg.reduce' op number of dimensions after reduction 1 doesn't match the init rank 2}}
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [1, 2] dimensions = [1, 2]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16x64xf32> func.return %reduce : tensor<16x64xf32>
} }
// ----- // -----
func.func @reduce_wrong_number_of_block_arguments( func.func @reduce_wrong_number_of_block_arguments(
%input1: tensor<16x32x64xf32>, %input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>, %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
%init2: tensor<16x64xf32>) -> (tensor<16x64xf32>, tensor<16x64xf32>) { %init2: tensor<16x64xf32>) -> (tensor<16x64xf32>, tensor<16x64xf32>) {
// expected-error @+1{{'linalg.reduce' op mismatching number of operands and block arguments}} // expected-error @+1{{'linalg.reduce' op mismatching number of operands and block arguments}}
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32>
} }
// ----- // -----
func.func @reduce_wrong_block_argument_input_type( func.func @reduce_wrong_block_argument_input_type(
%input1: tensor<16x32x64xf32>, %input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>, %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
%init2: tensor<16x64xf32>) -> (tensor<16x64xf32>, tensor<16x64xf32>) { %init2: tensor<16x64xf32>) -> (tensor<16x64xf32>, tensor<16x64xf32>) {
// expected-error @+1{{'linalg.reduce' op input element type 'f32' does not match corresponding block argument type 'f64'}} // expected-error @+1{{'linalg.reduce' op input element type 'f32' does not match corresponding block argument type 'f64'}}
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: f64, %out1: f32, %out2: f64) { (%in1: f32, %in2: f64, %out1: f32, %out2: f64) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
%1 = arith.addf %in2, %out2: f64 %1 = arith.addf %in2, %out2: f64
linalg.yield %0, %1: f32, f64 linalg.yield %0, %1: f32, f64
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf32>
} }
// ----- // -----
func.func @reduce_wrong_block_argument_output_type( func.func @reduce_wrong_block_argument_output_type(
%input1: tensor<16x32x64xf32>, %input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>, %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
%init2: tensor<16x64xf64>) -> (tensor<16x64xf32>, tensor<16x64xf32>) { %init2: tensor<16x64xf64>) -> (tensor<16x64xf32>, tensor<16x64xf32>) {
// expected-error @+1{{'linalg.reduce' op output element type 'f64' does not match corresponding block argument type 'f32'}} // expected-error @+1{{'linalg.reduce' op output element type 'f64' does not match corresponding block argument type 'f32'}}
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf64>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xf64>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: f32, %out1: f32, %out2: f32) { (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
linalg.yield %0, %out2: f32, f32 linalg.yield %0, %out2: f32, f32
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf64> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xf64>
} }
// ----- // -----
func.func @reduce_different_input_shapes(%input1: tensor<16x32x64xf32>, func.func @reduce_different_input_shapes(%input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<17x32x64xf32>, %init1: tensor<16x64xf32>, %input2: tensor<17x32x64xf32>,
%init2: tensor<17x64xf32>) -> (tensor<16x64xf32>, tensor<17x64xf32>) { %init2: tensor<17x64xf32>) -> (tensor<16x64xf32>, tensor<17x64xf32>) {
// expected-error @+1{{'linalg.reduce' op expects all inputs to have the same shapes. Shape at input-index 1 is not equal to the shape at input-index 0.}} // expected-error @+1{{'linalg.reduce' op expects all inputs to have the same shapes. Shape at input-index 1 is not equal to the shape at input-index 0.}}
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<17x32x64xf32>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<17x32x64xf32>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: f32, %out1: f32, %out2: f32) { (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
%1 = arith.addf %in2, %out2: f32 %1 = arith.addf %in2, %out2: f32
linalg.yield %0, %1: f32, f32 linalg.yield %0, %1: f32, f32
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32>
} }
// ----- // -----
func.func @reduce_different_output_shapes(%input1: tensor<16x32x64xf32>, func.func @reduce_different_output_shapes(%input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>, %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xf32>,
%init2: tensor<17x64xf32>) -> (tensor<16x64xf32>, tensor<17x64xf32>) { %init2: tensor<17x64xf32>) -> (tensor<16x64xf32>, tensor<17x64xf32>) {
// expected-error @+1{{'linalg.reduce' op expects all outputs to have the same shapes. Shape at output-index 1 is not equal to the shape at output-index 0.}} // expected-error @+1{{'linalg.reduce' op expects all outputs to have the same shapes. Shape at output-index 1 is not equal to the shape at output-index 0.}}
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xf32>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<17x64xf32>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: f32, %out1: f32, %out2: f32) { (%in1: f32, %in2: f32, %out1: f32, %out2: f32) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
%1 = arith.addf %in2, %out2: f32 %1 = arith.addf %in2, %out2: f32
linalg.yield %0, %1: f32, f32 linalg.yield %0, %1: f32, f32
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<17x64xf32>
} }
// ----- // -----
func.func @transpose_invalid_permutation(%input: tensor<16x32x64xf32>, func.func @transpose_invalid_permutation(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op permutation is not valid}} // expected-error @+1 {{'linalg.transpose' op permutation is not valid}}
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [1, 1, 2] permutation = [1, 1, 2]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }
// ----- // -----
func.func @transpose_permutated_dims_mismatch(%input: tensor<16x32x64xf32>, func.func @transpose_permutated_dims_mismatch(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op dim(result, 0) = 32 doesn't match dim(input, permutation[0]) = 16}} // expected-error @+1 {{'linalg.transpose' op dim(result, 0) = 32 doesn't match dim(input, permutation[0]) = 16}}
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [0, 1, 2] permutation = [0, 1, 2]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }
// ----- // -----
func.func @transpose_rank_permutation_size_mismatch( func.func @transpose_rank_permutation_size_mismatch(
%input: tensor<16x32x64xf32>, %input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op size of permutation 2 does not match the argument rank 3}} // expected-error @+1 {{'linalg.transpose' op size of permutation 2 does not match the argument rank 3}}
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [1, 0] permutation = [1, 0]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }
// ----- // -----
func.func @transpose_input_init_rank_mismatch(%input: tensor<16x32xf32>, func.func @transpose_input_init_rank_mismatch(%input: tensor<16x32xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// expected-error @+1 {{'linalg.transpose' op input rank 2 does not match init rank 3}} // expected-error @+1 {{'linalg.transpose' op input rank 2 does not match init rank 3}}
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32xf32>) ins(%input:tensor<16x32xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [1, 0, 2] permutation = [1, 0, 2]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }

mlir/test/Dialect/Linalg/library-calls.mlir

// RUN: mlir-opt %s -convert-linalg-to-std | FileCheck %s // RUN: mlir-opt %s -convert-linalg-to-std | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// CHECK: func private @linalg_fill_f32_viewsxsxf32(f32, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface} // CHECK: func private @linalg_fill_f32_viewsxsxf32(f32, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface}
// CHECK: func private @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32(memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface} // CHECK: func private @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32(memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface}
func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) { func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%x = memref.dim %A, %c0 : memref<?x?xf32> %x = memref.dim %A, %c0 : memref<?x?xf32>
%y = memref.dim %B, %c1 : memref<?x?xf32> %y = memref.dim %B, %c1 : memref<?x?xf32>
%C = memref.alloc(%x, %y) : memref<?x?xf32> %C = memref.alloc(%x, %y) : memref<?x?xf32>
// CHECK: call @linalg_fill_f32_viewsxsxf32({{.*}}) : (f32, memref<?x?xf32, {{.*}}>) // CHECK: call @linalg_fill_f32_viewsxsxf32({{.*}}) : (f32, memref<?x?xf32, {{.*}}>)
linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>) linalg.fill ins(%f0 : f32) inits(%C : memref<?x?xf32>)
// CHECK: call @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32({{.*}}) : (memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) -> () // CHECK: call @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32({{.*}}) : (memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) -> ()
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C: memref<?x?xf32>) inits(%C: memref<?x?xf32>)
return %C : memref<?x?xf32> return %C : memref<?x?xf32>
} }

mlir/test/Dialect/Linalg/loops.mlir

Show All 9 Lines
func.func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) { func.func @matmul(%arg0: memref<?xi8>, %M: index, %N: index, %K: index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%A = memref.view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32> %A = memref.view %arg0[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32>
%B = memref.view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32> %B = memref.view %arg0[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32>
%C = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32> %C = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32>
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C: memref<?x?xf32>) inits(%C: memref<?x?xf32>)
return return
} }
// CHECK-LABEL: func @matmul(%{{.*}}: memref<?xi8>, // CHECK-LABEL: func @matmul(%{{.*}}: memref<?xi8>,
// CHECK-SAME: [[M:arg[0-9]+]]: index // CHECK-SAME: [[M:arg[0-9]+]]: index
// CHECK-SAME: [[N:arg[0-9]+]]: index // CHECK-SAME: [[N:arg[0-9]+]]: index
// CHECK-SAME: [[K:arg[0-9]+]]: index // CHECK-SAME: [[K:arg[0-9]+]]: index
// CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32> // CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32>
// CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32> // CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32>
Show All 28 Lines
func.func @matvec(%arg0: memref<?xi8>, %M: index, %N: index) { func.func @matvec(%arg0: memref<?xi8>, %M: index, %N: index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%2 = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32> %2 = memref.view %arg0[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32>
%3 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32> %3 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32>
%4 = memref.view %arg0[%c0][%N] : memref<?xi8> to memref<?xf32> %4 = memref.view %arg0[%c0][%N] : memref<?xi8> to memref<?xf32>
linalg.matvec ins(%2, %3: memref<?x?xf32>, memref<?xf32>) linalg.matvec ins(%2, %3: memref<?x?xf32>, memref<?xf32>)
outs(%4 : memref<?xf32>) inits(%4 : memref<?xf32>)
return return
} }
// CHECK-LABEL: func @matvec(%{{.*}}: memref<?xi8>, // CHECK-LABEL: func @matvec(%{{.*}}: memref<?xi8>,
// CHECK-SAME: [[M:arg[0-9]+]]: index // CHECK-SAME: [[M:arg[0-9]+]]: index
// CHECK-SAME: [[K:arg[0-9]+]]: index // CHECK-SAME: [[K:arg[0-9]+]]: index
// CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32> // CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?x?xf32>
// CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?xf32> // CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?xf32>
// CHECK: %[[C:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?xf32> // CHECK: %[[C:.*]] = memref.view %{{.*}}[{{.*}}] : memref<?xi8> to memref<?xf32>
Show All 24 Lines
func.func @dot(%arg0: memref<?xi8>, %M: index) { func.func @dot(%arg0: memref<?xi8>, %M: index) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%1 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32> %1 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32>
%2 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32> %2 = memref.view %arg0[%c0][%M] : memref<?xi8> to memref<?xf32>
%3 = memref.view %arg0[%c0][] : memref<?xi8> to memref<f32> %3 = memref.view %arg0[%c0][] : memref<?xi8> to memref<f32>
linalg.dot ins(%1, %2 : memref<?xf32>, memref<?xf32>) linalg.dot ins(%1, %2 : memref<?xf32>, memref<?xf32>)
outs(%3 : memref<f32>) inits(%3 : memref<f32>)
return return
} }
// CHECK-LABEL: func @dot(%{{.*}}: memref<?xi8>, // CHECK-LABEL: func @dot(%{{.*}}: memref<?xi8>,
// CHECK-SAME: [[K:arg[0-9]+]]: index // CHECK-SAME: [[K:arg[0-9]+]]: index
// CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?xf32> // CHECK: %[[A:.*]] = memref.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?xf32>
// CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?xf32> // CHECK: %[[B:.*]] = memref.view %{{.*}}[{{.*}}][{{.*}}] : memref<?xi8> to memref<?xf32>
// CHECK: %[[C:.*]] = memref.view %{{.*}}[{{.*}}][] : memref<?xi8> to memref<f32> // CHECK: %[[C:.*]] = memref.view %{{.*}}[{{.*}}][] : memref<?xi8> to memref<f32>
// CHECK: scf.for {{.*}} to %[[K]] // CHECK: scf.for {{.*}} to %[[K]]
Show All 17 Lines
// CHECKPARALLEL-DAG: %[[res:.*]] = arith.addf %[[c]], %[[inc]] : f32 // CHECKPARALLEL-DAG: %[[res:.*]] = arith.addf %[[c]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %[[C]][] : memref<f32> // CHECKPARALLEL: store %[[res]], %[[C]][] : memref<f32>
func.func @dot_int(%arg0: memref<?xi32>, %arg1: memref<?xi32>, func.func @dot_int(%arg0: memref<?xi32>, %arg1: memref<?xi32>,
%arg3: memref<i32>) { %arg3: memref<i32>) {
// Verifies that we use the correct arith operations for integers. // Verifies that we use the correct arith operations for integers.
linalg.dot ins(%arg0, %arg1 : memref<?xi32>, memref<?xi32>) linalg.dot ins(%arg0, %arg1 : memref<?xi32>, memref<?xi32>)
outs(%arg3 : memref<i32>) inits(%arg3 : memref<i32>)
return return
} }
// CHECK-LABEL: func @dot_int( // CHECK-LABEL: func @dot_int(
// CHECK: %[[inc:.*]] = arith.muli {{.*}} : i32 // CHECK: %[[inc:.*]] = arith.muli {{.*}} : i32
// CHECK-NEXT: %[[res:.*]] = arith.addi {{.*}}, %[[inc]] : i32 // CHECK-NEXT: %[[res:.*]] = arith.addi {{.*}}, %[[inc]] : i32
// CHECK-NEXT: store %[[res]], {{.*}} : memref<i32> // CHECK-NEXT: store %[[res]], {{.*}} : memref<i32>
func.func @dot_bool(%arg0: memref<?xi1>, %arg1: memref<?xi1>, func.func @dot_bool(%arg0: memref<?xi1>, %arg1: memref<?xi1>,
%arg3: memref<i1>) { %arg3: memref<i1>) {
// Verifies that we use the correct (saturating) arith operations for booleans. // Verifies that we use the correct (saturating) arith operations for booleans.
linalg.dot ins(%arg0, %arg1 : memref<?xi1>, memref<?xi1>) linalg.dot ins(%arg0, %arg1 : memref<?xi1>, memref<?xi1>)
outs(%arg3 : memref<i1>) inits(%arg3 : memref<i1>)
return return
} }
// CHECK-LABEL: func @dot_bool( // CHECK-LABEL: func @dot_bool(
// CHECK: %[[inc:.*]] = arith.andi {{.*}} : i1 // CHECK: %[[inc:.*]] = arith.andi {{.*}} : i1
// CHECK-NEXT: %[[res:.*]] = arith.ori {{.*}}, %[[inc]] : i1 // CHECK-NEXT: %[[res:.*]] = arith.ori {{.*}}, %[[inc]] : i1
// CHECK-NEXT: store %[[res]], {{.*}} : memref<i1> // CHECK-NEXT: store %[[res]], {{.*}} : memref<i1>
func.func @dot_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>, %arg2: memref<f32>) { func.func @dot_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>, %arg2: memref<f32>) {
linalg.dot ins(%arg0, %arg1 : memref<?xf32, strided<[1], offset: ?>>, linalg.dot ins(%arg0, %arg1 : memref<?xf32, strided<[1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%arg2: memref<f32>) inits(%arg2: memref<f32>)
return return
} }
// CHECK-LABEL: func @dot_view( // CHECK-LABEL: func @dot_view(
// CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<f32>) { // CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<f32>) {
// CHECK: %[[K:.*]] = memref.dim %arg0, %c0 : memref<?xf32, strided<[1], offset: ?>> // CHECK: %[[K:.*]] = memref.dim %arg0, %c0 : memref<?xf32, strided<[1], offset: ?>>
// CHECK: scf.for {{.*}} to %[[K]] // CHECK: scf.for {{.*}} to %[[K]]
// CHECK-DAG: %[[a:.*]] = memref.load %arg0[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECK-DAG: %[[a:.*]] = memref.load %arg0[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
// CHECK-DAG: %[[b:.*]] = memref.load %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECK-DAG: %[[b:.*]] = memref.load %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
Show All 9 Lines
// CHECKPARALLEL-DAG: %[[a:.*]] = memref.load %arg0[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECKPARALLEL-DAG: %[[a:.*]] = memref.load %arg0[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
// CHECKPARALLEL-DAG: %[[b:.*]] = memref.load %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECKPARALLEL-DAG: %[[b:.*]] = memref.load %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
// CHECKPARALLEL-DAG: %[[inc:.*]] = arith.mulf %[[a]], %[[b]] : f32 // CHECKPARALLEL-DAG: %[[inc:.*]] = arith.mulf %[[a]], %[[b]] : f32
// CHECKPARALLEL-DAG: %[[c:.*]] = memref.load %{{.*}}[] : memref<f32> // CHECKPARALLEL-DAG: %[[c:.*]] = memref.load %{{.*}}[] : memref<f32>
// CHECKPARALLEL-DAG: %[[res:.*]] = arith.addf %[[c]], %[[inc]] : f32 // CHECKPARALLEL-DAG: %[[res:.*]] = arith.addf %[[c]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %{{.*}}[] : memref<f32> // CHECKPARALLEL: store %[[res]], %{{.*}}[] : memref<f32>
func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) { func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) {
linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?xf32, strided<[1], offset: ?>>) linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<?xf32, strided<[1], offset: ?>>)
return return
} }
// CHECK-LABEL: func @fill_view( // CHECK-LABEL: func @fill_view(
// CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) { // CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) {
// CHECK: scf.for {{.*}} to %{{.*}} // CHECK: scf.for {{.*}} to %{{.*}}
// CHECK: store %{{.*}}, %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECK: store %{{.*}}, %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
// CHECKPARALLEL-LABEL: func @fill_view( // CHECKPARALLEL-LABEL: func @fill_view(
// CHECKPARALLEL: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) { // CHECKPARALLEL: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) {
// CHECKPARALLEL: scf.parallel (%{{.*}}) = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) { // CHECKPARALLEL: scf.parallel (%{{.*}}) = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) {
// CHECKPARALLEL: store %{{.*}}, %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>> // CHECKPARALLEL: store %{{.*}}, %{{.*}}[%{{.*}}] : memref<?xf32, strided<[1], offset: ?>>
func.func @fill_view0(%arg0: memref<f32>, %arg1: f32) { func.func @fill_view0(%arg0: memref<f32>, %arg1: f32) {
linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<f32>) linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<f32>)
return return
} }
// CHECK-LABEL: func @fill_view0(%{{.*}}: memref<f32>, %{{.*}}: f32) { // CHECK-LABEL: func @fill_view0(%{{.*}}: memref<f32>, %{{.*}}: f32) {
// CHECK: store %{{.*}}, %{{.*}}[] : memref<f32> // CHECK: store %{{.*}}, %{{.*}}[] : memref<f32>
// CHECKPARALLEL-LABEL: func @fill_view0(%{{.*}}: memref<f32>, %{{.*}}: f32) { // CHECKPARALLEL-LABEL: func @fill_view0(%{{.*}}: memref<f32>, %{{.*}}: f32) {
// CHECKPARALLEL: store %{{.*}}, %{{.*}}[] : memref<f32> // CHECKPARALLEL: store %{{.*}}, %{{.*}}[] : memref<f32>
func.func @fill_view3(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: f32) { func.func @fill_view3(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: f32) {
linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
return return
} }
// CHECK-LABEL: func @fill_view3( // CHECK-LABEL: func @fill_view3(
// CHECK: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) { // CHECK: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) {
// CHECK: scf.for {{.*}} to %{{.*}} // CHECK: scf.for {{.*}} to %{{.*}}
// CHECK: scf.for {{.*}} to %{{.*}} // CHECK: scf.for {{.*}} to %{{.*}}
// CHECK: scf.for {{.*}} to %{{.*}} // CHECK: scf.for {{.*}} to %{{.*}}
// CHECK: store %{{.*}}, {{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> // CHECK: store %{{.*}}, {{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>
// CHECKPARALLEL-LABEL: func @fill_view3( // CHECKPARALLEL-LABEL: func @fill_view3(
// CHECKPARALLEL: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) { // CHECKPARALLEL: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) {
// CHECKPARALLEL: scf.parallel (%{{.*}}, %{{.*}}, %{{.*}}) = (%{{.*}}, %{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}}, %{{.*}}) { // CHECKPARALLEL: scf.parallel (%{{.*}}, %{{.*}}, %{{.*}}) = (%{{.*}}, %{{.*}}, %{{.*}}) to (%{{.*}}, %{{.*}}, %{{.*}}) step (%{{.*}}, %{{.*}}, %{{.*}}) {
// CHECKPARALLEL: store %{{.*}}, {{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> // CHECKPARALLEL: store %{{.*}}, {{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>
func.func @copy_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>) { func.func @copy_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>) {
linalg.generic { linalg.generic {
iterator_types = ["parallel"], iterator_types = ["parallel"],
indexing_maps = [ affine_map<(i) -> (i)>, affine_map<(i) -> (i)>] } indexing_maps = [ affine_map<(i) -> (i)>, affine_map<(i) -> (i)>] }
ins(%arg0: memref<?xf32, strided<[1], offset: ?>>) ins(%arg0: memref<?xf32, strided<[1], offset: ?>>)
outs(%arg1: memref<?xf32, strided<[1], offset: ?>>) { inits(%arg1: memref<?xf32, strided<[1], offset: ?>>) {
^bb0(%a: f32, %b: f32): ^bb0(%a: f32, %b: f32):
linalg.yield %a : f32 linalg.yield %a : f32
} }
return return
} }
// CHECK-LABEL: func @copy_view( // CHECK-LABEL: func @copy_view(
// CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>) { // CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>) {
// CHECK: scf.for {{.*}} to %{{.*}} // CHECK: scf.for {{.*}} to %{{.*}}
Show All 17 Lines#trait2 = {
iterator_types = ["parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel"],
indexing_maps = #accesses, indexing_maps = #accesses,
library_call = "some_external_function_name_2", library_call = "some_external_function_name_2",
doc = "B(i,j,k), C(i,k,j) = foo(A(i, j), B(i,j,k), C(i,k,j))" doc = "B(i,j,k), C(i,k,j) = foo(A(i, j), B(i,j,k), C(i,k,j))"
} }
func.func @generic_region(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>, %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg2: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { func.func @generic_region(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>, %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg2: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
linalg.generic #trait2 linalg.generic #trait2
ins(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>) ins(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%arg1, %arg2 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, inits(%arg1, %arg2 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>,
memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
^bb0(%a: f32, %b: f32, %c: f32): ^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b : f32 %d = arith.mulf %a, %b : f32
%e = arith.addf %c, %d : f32 %e = arith.addf %c, %d : f32
linalg.yield %d, %e : f32, f32 linalg.yield %d, %e : f32, f32
} }
return return
} }
Show All 28 Lines#trait4 = {
doc = "B(i,j,k), C(i,k,j) = foo(A(i, j) * B(i,j,k), i * j * k + C(i,k,j))" doc = "B(i,j,k), C(i,k,j) = foo(A(i, j) * B(i,j,k), i * j * k + C(i,k,j))"
} }
func.func @generic_index_region( func.func @generic_index_region(
%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>, %arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>,
%arg2: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { %arg2: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
linalg.generic #trait4 linalg.generic #trait4
ins(%arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>) ins(%arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%arg1, %arg2 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, inits(%arg1, %arg2 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>,
memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
^bb0(%a: f32, %b: f32, %c: f32): ^bb0(%a: f32, %b: f32, %c: f32):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%k = linalg.index 2 : index %k = linalg.index 2 : index
%result_1 = arith.mulf %a, %b : f32 %result_1 = arith.mulf %a, %b : f32
%ij = arith.addi %i, %j : index %ij = arith.addi %i, %j : index
▲ Show 20 Lines • Show All 51 Lines • ▼ Show 20 Lines#trait_broadcast = {
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
library_call = "some_broadcast_external_fn" library_call = "some_broadcast_external_fn"
} }
func.func @generic_op_zero_rank(%arg0: memref<f32>, %arg1: memref<3x4xf32>) func.func @generic_op_zero_rank(%arg0: memref<f32>, %arg1: memref<3x4xf32>)
{ {
linalg.generic #trait_broadcast linalg.generic #trait_broadcast
ins(%arg0 : memref<f32>) ins(%arg0 : memref<f32>)
outs(%arg1 : memref<3x4xf32>) { inits(%arg1 : memref<3x4xf32>) {
^bb(%a: f32, %b: f32) : ^bb(%a: f32, %b: f32) :
linalg.yield %a : f32 linalg.yield %a : f32
} }
return return
} }
// CHECK-LABEL: @generic_op_zero_rank // CHECK-LABEL: @generic_op_zero_rank
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<f32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<f32>
Show All 9 Lines
// CHECKPARALLEL: scf.parallel (%[[i:[a-zA-Z0-9_]*]], %[[j:[a-zA-Z0-9_]*]]) // CHECKPARALLEL: scf.parallel (%[[i:[a-zA-Z0-9_]*]], %[[j:[a-zA-Z0-9_]*]])
// CHECKPARALLEL: %[[a:.*]] = memref.load %[[ARG0]][] // CHECKPARALLEL: %[[a:.*]] = memref.load %[[ARG0]][]
// CHECKPARALLEL: store %[[a]], %[[ARG1]][%[[i]], %[[j]]] // CHECKPARALLEL: store %[[a]], %[[ARG1]][%[[i]], %[[j]]]
func.func @generic_op_scalar(%arg0: f32, %arg1: memref<3x4xf32>) func.func @generic_op_scalar(%arg0: f32, %arg1: memref<3x4xf32>)
{ {
linalg.generic #trait_broadcast linalg.generic #trait_broadcast
ins(%arg0 : f32) ins(%arg0 : f32)
outs(%arg1 : memref<3x4xf32>) { inits(%arg1 : memref<3x4xf32>) {
^bb(%a: f32, %b: f32) : ^bb(%a: f32, %b: f32) :
linalg.yield %a : f32 linalg.yield %a : f32
} }
return return
} }
// CHECK-LABEL: @generic_op_scalar // CHECK-LABEL: @generic_op_scalar
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: f32 // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: f32
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xf32>
// CHECK: scf.for %[[i:.*]] = {{.*}} // CHECK: scf.for %[[i:.*]] = {{.*}}
// CHECK: scf.for %[[j:.*]] = {{.*}} // CHECK: scf.for %[[j:.*]] = {{.*}}
// CHECK: store %[[ARG0]], %[[ARG1]][%[[i]], %[[j]]] // CHECK: store %[[ARG0]], %[[ARG1]][%[[i]], %[[j]]]
// CHECKPARALLEL-LABEL: @generic_op_scalar // CHECKPARALLEL-LABEL: @generic_op_scalar
// CHECKPARALLEL-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: f32 // CHECKPARALLEL-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: f32
// CHECKPARALLEL-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xf32> // CHECKPARALLEL-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<3x4xf32>
// CHECKPARALLEL: scf.parallel (%[[i:[a-zA-Z0-9_]*]], %[[j:[a-zA-Z0-9_]*]]) // CHECKPARALLEL: scf.parallel (%[[i:[a-zA-Z0-9_]*]], %[[j:[a-zA-Z0-9_]*]])
// CHECKPARALLEL: store %[[ARG0]], %[[ARG1]][%[[i]], %[[j]]] // CHECKPARALLEL: store %[[ARG0]], %[[ARG1]][%[[i]], %[[j]]]
func.func @generic_index_op_zero_rank(%arg0: memref<i32>, %arg1: memref<3x4xi32>) func.func @generic_index_op_zero_rank(%arg0: memref<i32>, %arg1: memref<3x4xi32>)
{ {
linalg.generic #trait_broadcast linalg.generic #trait_broadcast
ins(%arg0 : memref<i32>) ins(%arg0 : memref<i32>)
outs(%arg1 : memref<3x4xi32>) { inits(%arg1 : memref<3x4xi32>) {
^bb(%a: i32, %b: i32) : ^bb(%a: i32, %b: i32) :
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ij = arith.addi %i, %j : index %ij = arith.addi %i, %j : index
%ij_int = arith.index_cast %ij : index to i32 %ij_int = arith.index_cast %ij : index to i32
%result = arith.addi %a, %ij_int : i32 %result = arith.addi %a, %ij_int : i32
linalg.yield %result : i32 linalg.yield %result : i32
} }
Show All 33 Lines#trait_reduce_1D = {
iterator_types = ["reduction"], iterator_types = ["reduction"],
library_call = "some_reduce_external_fn" library_call = "some_reduce_external_fn"
} }
func.func @generic_op_1D_reduce(%arg0: memref<?xf32>, %arg1: memref<f32>) func.func @generic_op_1D_reduce(%arg0: memref<?xf32>, %arg1: memref<f32>)
{ {
linalg.generic #trait_reduce_1D linalg.generic #trait_reduce_1D
ins(%arg0 : memref<?xf32>) ins(%arg0 : memref<?xf32>)
outs(%arg1 : memref<f32>) { inits(%arg1 : memref<f32>) {
^bb(%a: f32, %b: f32) : ^bb(%a: f32, %b: f32) :
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
// CHECK-LABEL: @generic_op_1D_reduce // CHECK-LABEL: @generic_op_1D_reduce
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?xf32>
Show All 29 Lines
} }
func.func @generic_index_op_1D_reduce(%arg0: memref<?xf32>, func.func @generic_index_op_1D_reduce(%arg0: memref<?xf32>,
%arg1: memref<f32>, %arg1: memref<f32>,
%arg2: memref<f32>) %arg2: memref<f32>)
{ {
linalg.generic #trait_reduce_init_1D linalg.generic #trait_reduce_init_1D
ins(%arg0, %arg1 : memref<?xf32>, memref<f32>) ins(%arg0, %arg1 : memref<?xf32>, memref<f32>)
outs(%arg2 : memref<f32>) { inits(%arg2 : memref<f32>) {
^bb(%a: f32, %b: f32, %c: f32) : ^bb(%a: f32, %b: f32, %c: f32) :
%i = linalg.index 0 : index %i = linalg.index 0 : index
%0 = arith.constant 0 : index %0 = arith.constant 0 : index
%1 = arith.cmpi eq, %0, %i : index %1 = arith.cmpi eq, %0, %i : index
%2 = arith.select %1, %b, %c : f32 %2 = arith.select %1, %b, %c : f32
%3 = arith.addf %a, %2 : f32 %3 = arith.addf %a, %2 : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} }
Show All 27 Lines#trait_const_fill = {
args_in = 0, args_in = 0,
args_out = 1, args_out = 1,
indexing_maps = [affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>],
iterator_types = ["parallel"], iterator_types = ["parallel"],
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @generic_const_init(%arg0: memref<?xf32>) { func.func @generic_const_init(%arg0: memref<?xf32>) {
%cst = arith.constant 1.0 : f32 %cst = arith.constant 1.0 : f32
linalg.generic #trait_const_fill outs(%arg0 : memref<?xf32>) { linalg.generic #trait_const_fill inits(%arg0 : memref<?xf32>) {
^bb0(%arg1: f32): ^bb0(%arg1: f32):
linalg.yield %cst : f32 linalg.yield %cst : f32
} }
return return
} }
// CHECK-LABEL: @generic_const_init // CHECK-LABEL: @generic_const_init
// CHECK-SAME: %[[ARG0:.*]]: memref<?xf32> // CHECK-SAME: %[[ARG0:.*]]: memref<?xf32>
// CHECK: %[[CONST:.*]] = arith.constant 1.000000e+00 : f32 // CHECK: %[[CONST:.*]] = arith.constant 1.000000e+00 : f32
Show All 17 Lines#scalar_trait = {
iterator_types = [], iterator_types = [],
indexing_maps = #scalar_access, indexing_maps = #scalar_access,
library_call = "some_external_fn" library_call = "some_external_fn"
} }
func.func @scalar_code(%arg0: memref<f32>, %arg1 : memref<f32>, %arg2 : memref<f32>, %arg3 : i1) func.func @scalar_code(%arg0: memref<f32>, %arg1 : memref<f32>, %arg2 : memref<f32>, %arg3 : i1)
{ {
linalg.generic #scalar_trait linalg.generic #scalar_trait
ins(%arg0, %arg1 : memref<f32>, memref<f32>) ins(%arg0, %arg1 : memref<f32>, memref<f32>)
outs(%arg2 : memref<f32>) { inits(%arg2 : memref<f32>) {
^bb(%a : f32, %b : f32, %c : f32) : ^bb(%a : f32, %b : f32, %c : f32) :
%result = scf.if %arg3 -> (f32) { %result = scf.if %arg3 -> (f32) {
scf.yield %a : f32 scf.yield %a : f32
} else { } else {
scf.yield %b : f32 scf.yield %b : f32
} }
linalg.yield %result : f32 linalg.yield %result : f32
} }
Show All 25 Lines
// CHECKPARALLEL: scf.yield // CHECKPARALLEL: scf.yield
// CHECKPARALLEL: store %{{.*}}, %[[ARG2]][] // CHECKPARALLEL: store %{{.*}}, %[[ARG2]][]
//----------------------------------------------------------------------------// //----------------------------------------------------------------------------//
// Named ops to loops. // Named ops to loops.
//----------------------------------------------------------------------------// //----------------------------------------------------------------------------//
func.func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) { func.func @named_batch_matmul(%A: memref<?x?x?xf32>, %B: memref<?x?x?xf32>, %C: memref<?x?x?xf32>) {
linalg.batch_matmul ins(%A, %B : memref<?x?x?xf32>, memref<?x?x?xf32>) linalg.batch_matmul ins(%A, %B : memref<?x?x?xf32>, memref<?x?x?xf32>)
outs(%C : memref<?x?x?xf32>) inits(%C : memref<?x?x?xf32>)
return return
} }
// CHECK-LABEL: @named_batch_matmul // CHECK-LABEL: @named_batch_matmul
// CHECK-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mA:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mB:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[mC:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK: %[[B:.*]] = memref.dim %[[mA]], %c0 : memref<?x?x?xf32> // CHECK: %[[B:.*]] = memref.dim %[[mA]], %c0 : memref<?x?x?xf32>
// CHECK: %[[M:.*]] = memref.dim %[[mA]], %c1 : memref<?x?x?xf32> // CHECK: %[[M:.*]] = memref.dim %[[mA]], %c1 : memref<?x?x?xf32>
Show All 25 Lines
// CHECKPARALLEL: %[[vc:.*]] = memref.load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32> // CHECKPARALLEL: %[[vc:.*]] = memref.load %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
// CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[va]], %[[vb]] : f32 // CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[va]], %[[vb]] : f32
// CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32 // CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32> // CHECKPARALLEL: store %[[res]], %[[mC]][%[[b]], %[[m]], %[[n]]] : memref<?x?x?xf32>
func.func @conv1d_no_symbols(%in : memref<?xf32>, %filter : memref<?xf32>, %out : memref<?xf32>) -> () { func.func @conv1d_no_symbols(%in : memref<?xf32>, %filter : memref<?xf32>, %out : memref<?xf32>) -> () {
linalg.conv_1d ins(%in, %filter : memref<?xf32>, memref<?xf32>) linalg.conv_1d ins(%in, %filter : memref<?xf32>, memref<?xf32>)
outs(%out : memref<?xf32>) inits(%out : memref<?xf32>)
return return
} }
// CHECK-LABEL: @conv1d_no_symbols // CHECK-LABEL: @conv1d_no_symbols
// CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?xf32> // CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?xf32>
// CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?xf32> // CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?xf32>
// CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?xf32> // CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?xf32>
// CHECK: %[[c0:.*]] = arith.constant 0 : index // CHECK: %[[c0:.*]] = arith.constant 0 : index
Show All 26 Lines
// CHECKPARALLEL: %[[vc:.*]] = memref.load %[[arg2]][%[[b]]] : memref<?xf32> // CHECKPARALLEL: %[[vc:.*]] = memref.load %[[arg2]][%[[b]]] : memref<?xf32>
// CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[vb]], %[[va]] : f32 // CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[vb]], %[[va]] : f32
// CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32 // CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %[[arg2]][%[[b]]] : memref<?xf32> // CHECKPARALLEL: store %[[res]], %[[arg2]][%[[b]]] : memref<?xf32>
func.func @conv2d_no_symbols(%in : memref<?x?xf32>, %filter : memref<?x?xf32>, %out : memref<?x?xf32>) -> () { func.func @conv2d_no_symbols(%in : memref<?x?xf32>, %filter : memref<?x?xf32>, %out : memref<?x?xf32>) -> () {
linalg.conv_2d ins(%in, %filter : memref<?x?xf32>, memref<?x?xf32>) linalg.conv_2d ins(%in, %filter : memref<?x?xf32>, memref<?x?xf32>)
outs(%out: memref<?x?xf32>) inits(%out: memref<?x?xf32>)
return return
} }
// CHECK-LABEL: @conv2d_no_symbols // CHECK-LABEL: @conv2d_no_symbols
// CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK: %[[c0:.*]] = arith.constant 0 : index // CHECK: %[[c0:.*]] = arith.constant 0 : index
// CHECK: %[[c1:.*]] = arith.constant 1 : index // CHECK: %[[c1:.*]] = arith.constant 1 : index
Show All 36 Lines
// CHECKPARALLEL: %[[vc:.*]] = memref.load %[[arg2]][%[[arg3]], %[[arg4]]] : memref<?x?xf32> // CHECKPARALLEL: %[[vc:.*]] = memref.load %[[arg2]][%[[arg3]], %[[arg4]]] : memref<?x?xf32>
// CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[vb]], %[[va]] : f32 // CHECKPARALLEL: %[[inc:.*]] = arith.mulf %[[vb]], %[[va]] : f32
// CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32 // CHECKPARALLEL: %[[res:.*]] = arith.addf %[[vc]], %[[inc]] : f32
// CHECKPARALLEL: store %[[res]], %[[arg2]][%[[arg3]], %[[arg4]]] : memref<?x?xf32> // CHECKPARALLEL: store %[[res]], %[[arg2]][%[[arg3]], %[[arg4]]] : memref<?x?xf32>
func.func @conv3d_no_symbols(%in : memref<?x?x?xf32>, %filter : memref<?x?x?xf32>, %out : memref<?x?x?xf32>) -> () { func.func @conv3d_no_symbols(%in : memref<?x?x?xf32>, %filter : memref<?x?x?xf32>, %out : memref<?x?x?xf32>) -> () {
linalg.conv_3d ins(%in, %filter : memref<?x?x?xf32>, memref<?x?x?xf32>) linalg.conv_3d ins(%in, %filter : memref<?x?x?xf32>, memref<?x?x?xf32>)
outs(%out : memref<?x?x?xf32>) inits(%out : memref<?x?x?xf32>)
return return
} }
// CHECK-LABEL: @conv3d_no_symbols // CHECK-LABEL: @conv3d_no_symbols
// CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[arg0:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[arg1:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?xf32> // CHECK-SAME: %[[arg2:[a-zA-Z0-9]+]]: memref<?x?x?xf32>
// CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[c2:.*]] = arith.constant 2 : index
▲ Show 20 Lines • Show All 58 Lines • ▼ Show 20 Linesfunc.func @lower_to_loops_with_rank_reducing_subviews(
%0 = memref.subview %arg0[%arg2] [%arg3] [1] %0 = memref.subview %arg0[%arg2] [%arg3] [1]
: memref<?xi32> to memref<?xi32, strided<[1], offset: ?>> : memref<?xi32> to memref<?xi32, strided<[1], offset: ?>>
%1 = memref.subview %arg1[0, %arg4] [1, %arg3] [1, 1] %1 = memref.subview %arg1[0, %arg4] [1, %arg3] [1, 1]
: memref<?x?xi32> to memref<?xi32, strided<[1], offset: ?>> : memref<?x?xi32> to memref<?xi32, strided<[1], offset: ?>>
linalg.generic { linalg.generic {
iterator_types = ["parallel"], iterator_types = ["parallel"],
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>]} indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>]}
ins(%0: memref<?xi32, strided<[1], offset: ?>>) ins(%0: memref<?xi32, strided<[1], offset: ?>>)
outs(%1: memref<?xi32, strided<[1], offset: ?>>) { inits(%1: memref<?xi32, strided<[1], offset: ?>>) {
^bb0(%a: i32, %b: i32): ^bb0(%a: i32, %b: i32):
linalg.yield %a : i32 linalg.yield %a : i32
} }
return return
} }
// CHECK-LABEL: func @lower_to_loops_with_rank_reducing_subviews // CHECK-LABEL: func @lower_to_loops_with_rank_reducing_subviews
// CHECK: scf.for %[[IV:.+]] = %{{.+}} to %{{.+}} step %{{.+}} { // CHECK: scf.for %[[IV:.+]] = %{{.+}} to %{{.+}} step %{{.+}} {
// CHECK: %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]] // CHECK: %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]]
// CHECK: memref.store %[[VAL]], %{{.+}}[%[[IV]]] // CHECK: memref.store %[[VAL]], %{{.+}}[%[[IV]]]
// CHECK: } // CHECK: }
// CHECKPARALLEL-LABEL: func @lower_to_loops_with_rank_reducing_subviews // CHECKPARALLEL-LABEL: func @lower_to_loops_with_rank_reducing_subviews
// CHECKPARALLEL: scf.parallel (%[[IV:.+]]) = (%{{.+}}) to (%{{.+}}) step (%{{.+}}) { // CHECKPARALLEL: scf.parallel (%[[IV:.+]]) = (%{{.+}}) to (%{{.+}}) step (%{{.+}}) {
// CHECKPARALLEL: %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]] // CHECKPARALLEL: %[[VAL:.+]] = memref.load %{{.+}}[%[[IV]]]
// CHECKPARALLEL: memref.store %[[VAL]], %{{.+}}[%[[IV]]] // CHECKPARALLEL: memref.store %[[VAL]], %{{.+}}[%[[IV]]]
// CHECKPARALLEL: } // CHECKPARALLEL: }

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

Show First 20 Lines • Show All 51 Lines • ▼ Show 20 Lines
// 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:.+]] = tensor.empty() : 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(%{{.*}} : tensor<1x28x28x1xf32>) outs(%{{.*}} : tensor<1x32x32x1xf32>) // CHECK: ins(%{{.*}} : tensor<1x28x28x1xf32>) inits(%{{.*}} : tensor<1x32x32x1xf32>)
// CHECK: ^bb0(%[[VAL:.+]]: f32, %{{.*}}: f32) // CHECK: ^bb0(%[[VAL:.+]]: f32, %{{.*}}: f32)
// CHECK: linalg.yield %[[VAL]] : f32 // CHECK: linalg.yield %[[VAL]] : f32
// CHECK: return %[[R2c:.+]] // CHECK: return %[[R2c:.+]]

mlir/test/Dialect/Linalg/multisize-tiling-full.mlir

Show All 22 Lines
func.func @two_d(%arg0: tensor<10x34xf32>, func.func @two_d(%arg0: tensor<10x34xf32>,
%arg1: tensor<10x34xf32>) -> tensor<10x34xf32> { %arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>, indexing_maps = [affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i, j)>], affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%arg0: tensor<10x34xf32>) ins(%arg0: tensor<10x34xf32>)
outs(%arg1: tensor<10x34xf32>) { inits(%arg1: tensor<10x34xf32>) {
^bb0(%0: f32, %1: f32): ^bb0(%0: f32, %1: f32):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32 %call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32
linalg.yield %call_res : f32 linalg.yield %call_res : f32
} -> tensor<10x34xf32> } -> tensor<10x34xf32>
// 2D multi-size tiling should produce for quadrants with sizes // 2D multi-size tiling should produce for quadrants with sizes
Show All 10 Linesfunc.func @two_d(%arg0: tensor<10x34xf32>,
// CHECK: %[[OUTSLICE_1_IN:.+]] = tensor.extract_slice %[[SLICE_1_IN]][%[[I1]], 0] [2, 34] [1, 1] // CHECK: %[[OUTSLICE_1_IN:.+]] = tensor.extract_slice %[[SLICE_1_IN]][%[[I1]], 0] [2, 34] [1, 1]
// CHECK: %[[OUTSLICE_1:.+]] = tensor.extract_slice %[[ITERARG_1]][%[[I1]], 0] [2, 34] [1, 1] // CHECK: %[[OUTSLICE_1:.+]] = tensor.extract_slice %[[ITERARG_1]][%[[I1]], 0] [2, 34] [1, 1]
// CHECK: %[[SLICE_2_IN:.+]] = tensor.extract_slice %[[OUTSLICE_1_IN]][0, 0] [2, 16] [1, 1] // CHECK: %[[SLICE_2_IN:.+]] = tensor.extract_slice %[[OUTSLICE_1_IN]][0, 0] [2, 16] [1, 1]
// CHECK: %[[SLICE_2:.+]] = tensor.extract_slice %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1] // CHECK: %[[SLICE_2:.+]] = tensor.extract_slice %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1]
// CHECK: %[[LOOPRES:.+]] = scf.for %[[I2:.+]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ITERARG_2:.+]] = %[[SLICE_2]]) // CHECK: %[[LOOPRES:.+]] = scf.for %[[I2:.+]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ITERARG_2:.+]] = %[[SLICE_2]])
// CHECK: %[[INSLICE_2:.+]] = tensor.extract_slice %[[SLICE_2_IN]][0, %[[I2]]] [2, 8] [1, 1] // CHECK: %[[INSLICE_2:.+]] = tensor.extract_slice %[[SLICE_2_IN]][0, %[[I2]]] [2, 8] [1, 1]
// CHECK: %[[OUTSLICE_2:.+]] = tensor.extract_slice %[[ITERARG_2]][0, %[[I2]]] [2, 8] [1, 1] // CHECK: %[[OUTSLICE_2:.+]] = tensor.extract_slice %[[ITERARG_2]][0, %[[I2]]] [2, 8] [1, 1]
// CHECK: %[[RESSLICE_1:.+]] = linalg.generic {{.*}} ins(%[[INSLICE_2]] : tensor<2x8xf32>) outs(%[[OUTSLICE_2]] : tensor<2x8xf32>) // CHECK: %[[RESSLICE_1:.+]] = linalg.generic {{.*}} ins(%[[INSLICE_2]] : tensor<2x8xf32>) inits(%[[OUTSLICE_2]] : tensor<2x8xf32>)
// CHECK: %[[RESPARTIAL:.+]] = tensor.insert_slice %[[RESSLICE_1]] into %[[ITERARG_2]] // CHECK: %[[RESPARTIAL:.+]] = tensor.insert_slice %[[RESSLICE_1]] into %[[ITERARG_2]]
// CHECK: scf.yield %[[RESPARTIAL]] // CHECK: scf.yield %[[RESPARTIAL]]
// CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[LOOPRES]] into %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1] // CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[LOOPRES]] into %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1]
// CHECK: %[[OUTSLICE_3:.+]] = tensor.extract_slice %[[INSERTED]][0, 16] [2, 18] [1, 1] // CHECK: %[[OUTSLICE_3:.+]] = tensor.extract_slice %[[INSERTED]][0, 16] [2, 18] [1, 1]
// CHECK: scf.for %{{.*}} iter_args(%{{.*}} = %[[OUTSLICE_3]]) // CHECK: scf.for %{{.*}} iter_args(%{{.*}} = %[[OUTSLICE_3]])
// CHECK-COUNT-2: tensor.extract_slice // CHECK-COUNT-2: tensor.extract_slice
// CHECK: linalg.generic {{.*}} ins(%{{.*}} : tensor<2x9xf32>) // CHECK: linalg.generic {{.*}} ins(%{{.*}} : tensor<2x9xf32>)
Show All 29 Lines

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 = tensor.empty() : 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) inits(%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> inits(%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 = tensor.empty() : 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) inits(%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> inits(%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 = tensor.empty() : 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) inits(%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: inits(%{{.+}} : 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>)
outs(%fill : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32> inits(%fill : tensor<2x3x4x2x3xf32>) -> tensor<2x3x4x2x3xf32>
return %0 : tensor<2x3x4x2x3xf32> return %0 : tensor<2x3x4x2x3xf32>
} }
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_memref // CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_memref
func.func @depthwise_conv_2d_nhwc_hwcm_memref(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x3x4x2x3xf32>) { func.func @depthwise_conv_2d_nhwc_hwcm_memref(%input: memref<2x4x5x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<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(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x4x5x2xf32>, memref<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<2x3x4x2x3xf32>) // CHECK-SAME: inits(%{{.+}} : memref<2x3x4x2x3xf32>)
linalg.depthwise_conv_2d_nhwc_hwcm 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>) inits(%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 = tensor.empty() : 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: inits(%{{.+}} : 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> inits(%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 = tensor.empty() : 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: inits(%{{.+}} : 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> inits(%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: inits(%{{.+}} : 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>) inits(%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 = tensor.empty() : 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: inits(%{{.+}} : 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> inits(%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: inits(%{{.+}} : 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>) inits(%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 = tensor.empty() : 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) inits(%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: inits(%{{.+}} : 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>)
outs(%fill : tensor<2x6x7x2x3xf32>) -> tensor<2x6x7x2x3xf32> inits(%fill : tensor<2x6x7x2x3xf32>) -> tensor<2x6x7x2x3xf32>
return %0 : tensor<2x6x7x2x3xf32> return %0 : tensor<2x6x7x2x3xf32>
} }
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_memref_dilated // CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwcm_memref_dilated
func.func @depthwise_conv_2d_nhwc_hwcm_memref_dilated(%input: memref<2x8x9x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<2x6x7x2x3xf32>) { func.func @depthwise_conv_2d_nhwc_hwcm_memref_dilated(%input: memref<2x8x9x2xf32>, %filter: memref<2x2x2x3xf32>, %output: memref<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(%{{.+}}, %{{.+}} : memref<2x8x9x2xf32>, memref<2x2x2x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<2x8x9x2xf32>, memref<2x2x2x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<2x6x7x2x3xf32>) // CHECK-SAME: inits(%{{.+}} : memref<2x6x7x2x3xf32>)
linalg.depthwise_conv_2d_nhwc_hwcm 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 : memref<2x8x9x2xf32>, memref<2x2x2x3xf32>) ins(%input, %filter : memref<2x8x9x2xf32>, memref<2x2x2x3xf32>)
outs(%output : memref<2x6x7x2x3xf32>) inits(%output : memref<2x6x7x2x3xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: func @depthwise_conv_2d_input_nhwc_filter_default_attributes // CHECK-LABEL: func @depthwise_conv_2d_input_nhwc_filter_default_attributes
func.func @depthwise_conv_2d_input_nhwc_filter_default_attributes(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) { func.func @depthwise_conv_2d_input_nhwc_filter_default_attributes(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
// CHECK: linalg.depthwise_conv_2d_nhwc_hwc // CHECK: linalg.depthwise_conv_2d_nhwc_hwc
// CHECK-NOT: strides = // CHECK-NOT: strides =
// CHECK-NOT: dilations = // CHECK-NOT: dilations =
linalg.depthwise_conv_2d_nhwc_hwc linalg.depthwise_conv_2d_nhwc_hwc
ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>) ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
outs(%output: memref<1x56x56x96xf32>) inits(%output: memref<1x56x56x96xf32>)
return return
} }
// ----- // -----
func.func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_element_type(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) { func.func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_element_type(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
// expected-error @+1 {{op attribute 'strides' failed to satisfy constraint: 64-bit signless int elements attribute of shape [2]}} // expected-error @+1 {{op attribute 'strides' failed to satisfy constraint: 64-bit signless int elements attribute of shape [2]}}
linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2.0> : vector<2xf32>} linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2.0> : vector<2xf32>}
ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>) ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
outs(%output: memref<1x56x56x96xf32>) inits(%output: memref<1x56x56x96xf32>)
return return
} }
// ----- // -----
func.func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_size(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) { func.func @depthwise_conv_2d_input_nhwc_filter_wrong_stride_size(%input: memref<1x113x113x96xf32>, %filter: memref<3x3x96xf32>, %output: memref<1x56x56x96xf32>) {
// expected-error @+1 {{op attribute 'strides' failed to satisfy constraint: 64-bit signless int elements attribute of shape [2]}} // expected-error @+1 {{op attribute 'strides' failed to satisfy constraint: 64-bit signless int elements attribute of shape [2]}}
linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<3xi64> } linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<3xi64> }
ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>) ins(%input, %filter: memref<1x113x113x96xf32>, memref<3x3x96xf32>)
outs(%output: memref<1x56x56x96xf32>) inits(%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 = tensor.empty() : 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) inits(%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> inits(%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 = tensor.empty() : 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) inits(%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> inits(%fill : tensor<2x3x13x4x6xf32>) -> tensor<2x3x13x4x6xf32>
return %0 : tensor<2x3x13x4x6xf32> return %0 : tensor<2x3x13x4x6xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_1d_nwc_wcf // CHECK-LABEL: func @conv_1d_nwc_wcf
func.func @conv_1d_nwc_wcf(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { func.func @conv_1d_nwc_wcf(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_1d_nwc_wcf // CHECK: %{{.+}} = linalg.conv_1d_nwc_wcf
// CHECK-SAME: dilations = dense<1> : tensor<1xi64> // CHECK-SAME: dilations = dense<1> : tensor<1xi64>
// CHECK-SAME: strides = dense<1> : tensor<1xi64> // CHECK-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?xf32>, tensor<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%0 = linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>, %0 = linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>) ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
outs (%init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> inits (%init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32> return %0 : tensor<?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_1d_nwc_wcf // CHECK-LABEL: func @conv_1d_nwc_wcf
func.func @conv_1d_nwc_wcf(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) { func.func @conv_1d_nwc_wcf(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {
// CHECK: linalg.conv_1d_nwc_wcf // CHECK: linalg.conv_1d_nwc_wcf
// CHECK-SAME: dilations = dense<1> : tensor<1xi64> // CHECK-SAME: dilations = dense<1> : tensor<1xi64>
// CHECK-SAME: strides = dense<1> : tensor<1xi64> // CHECK-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?xf32>)
linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>, linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>) ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%output: memref<?x?x?xf32>) inits (%output: memref<?x?x?xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: func @conv_1d_ncw_fcw // CHECK-LABEL: func @conv_1d_ncw_fcw
func.func @conv_1d_ncw_fcw(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { func.func @conv_1d_ncw_fcw(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_1d_ncw_fcw // CHECK: %{{.+}} = linalg.conv_1d_ncw_fcw
// CHECK-SAME: dilations = dense<1> : tensor<1xi64> // CHECK-SAME: dilations = dense<1> : tensor<1xi64>
// CHECK-SAME: strides = dense<1> : tensor<1xi64> // CHECK-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?xf32>, tensor<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%0 = linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>, %0 = linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>) ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
outs (%init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> inits (%init: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32> return %0 : tensor<?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_1d_ncw_fcw // CHECK-LABEL: func @conv_1d_ncw_fcw
func.func @conv_1d_ncw_fcw(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) { func.func @conv_1d_ncw_fcw(%input: memref<?x?x?xf32>, %filter: memref<?x?x?xf32>, %output: memref<?x?x?xf32>) {
// CHECK: linalg.conv_1d_ncw_fcw // CHECK: linalg.conv_1d_ncw_fcw
// CHECK-SAME: dilations = dense<1> : tensor<1xi64> // CHECK-SAME: dilations = dense<1> : tensor<1xi64>
// CHECK-SAME: strides = dense<1> : tensor<1xi64> // CHECK-SAME: strides = dense<1> : tensor<1xi64>
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?xf32>)
linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>, linalg.conv_1d_ncw_fcw {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>) ins (%input, %filter: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%output: memref<?x?x?xf32>) inits (%output: memref<?x?x?xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_nhwc_hwcf // CHECK-LABEL: func @conv_2d_nhwc_hwcf
func.func @conv_2d_nhwc_hwcf(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> { func.func @conv_2d_nhwc_hwcf(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_2d_nhwc_hwcf // CHECK: %{{.+}} = linalg.conv_2d_nhwc_hwcf
// 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<?x?x?x?xf32>, tensor<?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
%0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : 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 (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> inits (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
return %0 : tensor<?x?x?x?xf32> return %0 : tensor<?x?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_ngchw_fgchw // CHECK-LABEL: func @conv_2d_ngchw_fgchw
func.func @conv_2d_ngchw_fgchw(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> { func.func @conv_2d_ngchw_fgchw(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_2d_ngchw_fgchw // CHECK: %{{.+}} = linalg.conv_2d_ngchw_fgchw
// 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<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
%0 = linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
outs (%init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> inits (%init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
return %0 : tensor<?x?x?x?x?xf32> return %0 : tensor<?x?x?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_nhwc_fhwc // CHECK-LABEL: func @conv_2d_nhwc_fhwc
func.func @conv_2d_nhwc_fhwc(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> { func.func @conv_2d_nhwc_fhwc(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_2d_nhwc_fhwc // CHECK: %{{.+}} = linalg.conv_2d_nhwc_fhwc
// 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<?x?x?x?xf32>, tensor<?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
%0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : 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 (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> inits (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
return %0 : tensor<?x?x?x?xf32> return %0 : tensor<?x?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_nhwc_fhwc_static // CHECK-LABEL: func @conv_2d_nhwc_fhwc_static
func.func @conv_2d_nhwc_fhwc_static(%input: tensor<?x128x128x32xf32>, %filter: tensor<64x3x3x32xf32>, %init: tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32> { func.func @conv_2d_nhwc_fhwc_static(%input: tensor<?x128x128x32xf32>, %filter: tensor<64x3x3x32xf32>, %init: tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32> {
// CHECK: %{{.+}} = linalg.conv_2d_nhwc_fhwc // CHECK: %{{.+}} = linalg.conv_2d_nhwc_fhwc
// 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<?x128x128x32xf32>, tensor<64x3x3x32xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x128x128x32xf32>, tensor<64x3x3x32xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32>
%0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x128x128x32xf32>, tensor<64x3x3x32xf32>) ins (%input, %filter: tensor<?x128x128x32xf32>, tensor<64x3x3x32xf32>)
outs (%init: tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32> inits (%init: tensor<?x126x126x64xf32>) -> tensor<?x126x126x64xf32>
return %0 : tensor<?x126x126x64xf32> return %0 : tensor<?x126x126x64xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_nhwc_hwcf // CHECK-LABEL: func @conv_2d_nhwc_hwcf
func.func @conv_2d_nhwc_hwcf(%input: memref<?x?x?x?xf32>, %filter: memref<?x?x?x?xf32>, %output: memref<?x?x?x?xf32>) { func.func @conv_2d_nhwc_hwcf(%input: memref<?x?x?x?xf32>, %filter: memref<?x?x?x?xf32>, %output: memref<?x?x?x?xf32>) {
// CHECK: linalg.conv_2d_nhwc_hwcf // CHECK: linalg.conv_2d_nhwc_hwcf
// 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(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?x?xf32>)
linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>) ins (%input, %filter: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
outs (%output: memref<?x?x?x?xf32>) inits (%output: memref<?x?x?x?xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: func @conv_2d_ngchw_fgchw // CHECK-LABEL: func @conv_2d_ngchw_fgchw
func.func @conv_2d_ngchw_fgchw(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) { func.func @conv_2d_ngchw_fgchw(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
// CHECK: linalg.conv_2d_ngchw_fgchw // CHECK: linalg.conv_2d_ngchw_fgchw
// 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(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?x?x?xf32>)
linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>, linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
outs (%output: memref<?x?x?x?x?xf32>) inits (%output: memref<?x?x?x?x?xf32>)
return return
} }
// ----- // -----
// CHECK-LABEL: func @conv_3d_ndhwc_dhwcf // CHECK-LABEL: func @conv_3d_ndhwc_dhwcf
func.func @conv_3d_ndhwc_dhwcf(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> { func.func @conv_3d_ndhwc_dhwcf(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> {
// CHECK: %{{.+}} = linalg.conv_3d_ndhwc_dhwcf // CHECK: %{{.+}} = linalg.conv_3d_ndhwc_dhwcf
// 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<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> // CHECK-SAME: inits(%{{.+}} : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
%0 = linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>, %0 = linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>,
strides = dense<1> : tensor<3xi64>} strides = dense<1> : tensor<3xi64>}
ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
outs (%init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> inits (%init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
return %0 : tensor<?x?x?x?x?xf32> return %0 : tensor<?x?x?x?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @conv_3d_ndhwc_dhwcf // CHECK-LABEL: func @conv_3d_ndhwc_dhwcf
func.func @conv_3d_ndhwc_dhwcf(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) { func.func @conv_3d_ndhwc_dhwcf(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
// CHECK: linalg.conv_3d_ndhwc_dhwcf // CHECK: linalg.conv_3d_ndhwc_dhwcf
// 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(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?x?x?x?xf32>)
linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>, linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>,
strides = dense<1> : tensor<3xi64>} strides = dense<1> : tensor<3xi64>}
ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
outs (%output: memref<?x?x?x?x?xf32>) inits (%output: memref<?x?x?x?x?xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32> return %res : tensor<1x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_sum // CHECK-LABEL: func @pooling_nhwc_sum
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xf32>)
func.func @pooling_nhwc_sum(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) { func.func @pooling_nhwc_sum(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>) ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>) inits(%output: memref<1x2x2x1xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
return %res : tensor<1x1x2x2xf32> return %res : tensor<1x1x2x2xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nchw_sum // CHECK-LABEL: func @pooling_nchw_sum
// 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(%{{.+}}, %{{.+}} : memref<1x1x4x4xf32>, memref<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x1x4x4xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x1x2x2xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x1x2x2xf32>)
func.func @pooling_nchw_sum(%input: memref<1x1x4x4xf32>, %fake: memref<3x3xf32>, %output: memref<1x1x2x2xf32>) { func.func @pooling_nchw_sum(%input: memref<1x1x4x4xf32>, %fake: memref<3x3xf32>, %output: memref<1x1x2x2xf32>) {
linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nchw_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x1x4x4xf32>, memref<3x3xf32>) ins(%input, %fake: memref<1x1x4x4xf32>, memref<3x3xf32>)
outs(%output: memref<1x1x2x2xf32>) inits(%output: memref<1x1x2x2xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x1x2x2xf32>) -> tensor<1x1x2x2xf32>
return %res : tensor<1x1x2x2xf32> return %res : tensor<1x1x2x2xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_max // CHECK-LABEL: func @pooling_nhwc_max
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xf32>)
func.func @pooling_nhwc_max(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) { func.func @pooling_nhwc_max(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>) ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>) inits(%output: memref<1x2x2x1xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xi8> %fake = tensor.empty() : tensor<3x3xi8>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
return %res : tensor<1x2x2x1xi8> return %res : tensor<1x2x2x1xi8>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_i8_max // CHECK-LABEL: func @pooling_nhwc_i8_max
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xi8>, memref<3x3xi8>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi8>, memref<3x3xi8>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xi8>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xi8>)
func.func @pooling_nhwc_i8_max(%input: memref<1x4x4x1xi8>, %fake: memref<3x3xi8>, %output: memref<1x2x2x1xi8>) { func.func @pooling_nhwc_i8_max(%input: memref<1x4x4x1xi8>, %fake: memref<3x3xi8>, %output: memref<1x2x2x1xi8>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xi8>, memref<3x3xi8>) ins(%input, %fake: memref<1x4x4x1xi8>, memref<3x3xi8>)
outs(%output: memref<1x2x2x1xi8>) inits(%output: memref<1x2x2x1xi8>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xi16> %fake = tensor.empty() : tensor<3x3xi16>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
return %res : tensor<1x2x2x1xi16> return %res : tensor<1x2x2x1xi16>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_i16_max // CHECK-LABEL: func @pooling_nhwc_i16_max
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xi16>, memref<3x3xi16>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi16>, memref<3x3xi16>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xi16>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xi16>)
func.func @pooling_nhwc_i16_max(%input: memref<1x4x4x1xi16>, %fake: memref<3x3xi16>, %output: memref<1x2x2x1xi16>) { func.func @pooling_nhwc_i16_max(%input: memref<1x4x4x1xi16>, %fake: memref<3x3xi16>, %output: memref<1x2x2x1xi16>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xi16>, memref<3x3xi16>) ins(%input, %fake: memref<1x4x4x1xi16>, memref<3x3xi16>)
outs(%output: memref<1x2x2x1xi16>) inits(%output: memref<1x2x2x1xi16>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xi32> %fake = tensor.empty() : tensor<3x3xi32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
return %res : tensor<1x2x2x1xi32> return %res : tensor<1x2x2x1xi32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_i32_max // CHECK-LABEL: func @pooling_nhwc_i32_max
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xi32>, memref<3x3xi32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi32>, memref<3x3xi32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xi32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xi32>)
func.func @pooling_nhwc_i32_max(%input: memref<1x4x4x1xi32>, %fake: memref<3x3xi32>, %output: memref<1x2x2x1xi32>) { func.func @pooling_nhwc_i32_max(%input: memref<1x4x4x1xi32>, %fake: memref<3x3xi32>, %output: memref<1x2x2x1xi32>) {
linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xi32>, memref<3x3xi32>) ins(%input, %fake: memref<1x4x4x1xi32>, memref<3x3xi32>)
outs(%output: memref<1x2x2x1xi32>) inits(%output: memref<1x2x2x1xi32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3xf32> %fake = tensor.empty() : tensor<3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
return %res : tensor<1x2x2x1xf32> return %res : tensor<1x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_nhwc_min // CHECK-LABEL: func @pooling_nhwc_min
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xf32>, memref<3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x1xf32>)
func.func @pooling_nhwc_min(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) { func.func @pooling_nhwc_min(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>) ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>)
outs(%output: memref<1x2x2x1xf32>) inits(%output: memref<1x2x2x1xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_ndhwc_sum // CHECK-LABEL: func @pooling_ndhwc_sum
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x2x1xf32>)
func.func @pooling_ndhwc_sum(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) { func.func @pooling_ndhwc_sum(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) {
linalg.pooling_ndhwc_sum {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} linalg.pooling_ndhwc_sum {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
outs(%output: memref<1x2x2x2x1xf32>) inits(%output: memref<1x2x2x2x1xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_ndhwc_max // CHECK-LABEL: func @pooling_ndhwc_max
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x2x1xf32>)
func.func @pooling_ndhwc_max(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) { func.func @pooling_ndhwc_max(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) {
linalg.pooling_ndhwc_max {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} linalg.pooling_ndhwc_max {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
outs(%output: memref<1x2x2x2x1xf32>) inits(%output: memref<1x2x2x2x1xf32>)
return return
} }
// ----- // -----
// 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: inits(%{{.+}} : 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 = tensor.empty() : tensor<3x3x3xf32> %fake = tensor.empty() : tensor<3x3x3xf32>
%init = tensor.empty() : 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) inits(%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> inits(%fill: tensor<1x2x2x2x1xf32>) -> tensor<1x2x2x2x1xf32>
return %res : tensor<1x2x2x2x1xf32> return %res : tensor<1x2x2x2x1xf32>
} }
// ----- // -----
// CHECK-LABEL: func @pooling_ndhwc_min // CHECK-LABEL: func @pooling_ndhwc_min
// 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(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<1x2x2x2x1xf32>) // CHECK-SAME: inits(%{{.+}} : memref<1x2x2x2x1xf32>)
func.func @pooling_ndhwc_min(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) { func.func @pooling_ndhwc_min(%input: memref<1x4x4x4x1xf32>, %fake: memref<3x3x3xf32>, %output: memref<1x2x2x2x1xf32>) {
linalg.pooling_ndhwc_min {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>} linalg.pooling_ndhwc_min {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}
ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>) ins(%input, %fake: memref<1x4x4x4x1xf32>, memref<3x3x3xf32>)
outs(%output: memref<1x2x2x2x1xf32>) inits(%output: memref<1x2x2x2x1xf32>)
return return
} }
// ----- // -----
#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 2, d2 * 2 + d5, d6)> #map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 2, d2 * 2 + d5, d6)>
#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)> #map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>
#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)> #map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
Show All 26 Linesfunc.func @conv_interface_wrong_num_operands(
return %0 : tensor<?x?x?x?xf32> return %0 : tensor<?x?x?x?xf32>
} }
// ----- // -----
func.func @batch_reduce_matmul(%arg0: tensor<8x128x256xf32>, %arg1: tensor<8x256x512xf32>, %arg2: tensor<128x512xf32>) -> tensor<128x512xf32> { func.func @batch_reduce_matmul(%arg0: tensor<8x128x256xf32>, %arg1: tensor<8x256x512xf32>, %arg2: tensor<128x512xf32>) -> tensor<128x512xf32> {
// CHECK: %{{.+}} = linalg.batch_reduce_matmul // CHECK: %{{.+}} = linalg.batch_reduce_matmul
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<8x128x256xf32>, tensor<8x256x512xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : tensor<8x128x256xf32>, tensor<8x256x512xf32>)
// CHECK-SAME: outs(%{{.+}} : tensor<128x512xf32>) -> tensor<128x512xf32> // CHECK-SAME: inits(%{{.+}} : tensor<128x512xf32>) -> tensor<128x512xf32>
%0 = linalg.batch_reduce_matmul ins(%arg0, %arg1 : tensor<8x128x256xf32>, tensor<8x256x512xf32>) outs(%arg2: tensor<128x512xf32>) -> tensor<128x512xf32> %0 = linalg.batch_reduce_matmul ins(%arg0, %arg1 : tensor<8x128x256xf32>, tensor<8x256x512xf32>) inits(%arg2: tensor<128x512xf32>) -> tensor<128x512xf32>
return %0: tensor<128x512xf32> return %0: tensor<128x512xf32>
} }
// ----- // -----
func.func @batch_reduce_matmul(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?xf32>) { func.func @batch_reduce_matmul(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?xf32>) {
// CHECK: linalg.batch_reduce_matmul // CHECK: linalg.batch_reduce_matmul
// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>) // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?xf32>, memref<?x?x?xf32>)
// CHECK-SAME: outs(%{{.+}} : memref<?x?xf32>) // CHECK-SAME: inits(%{{.+}} : memref<?x?xf32>)
linalg.batch_reduce_matmul ins(%arg0, %arg1 : memref<?x?x?xf32>, memref<?x?x?xf32>) outs(%arg2: memref<?x?xf32>) linalg.batch_reduce_matmul ins(%arg0, %arg1 : memref<?x?x?xf32>, memref<?x?x?xf32>) inits(%arg2: memref<?x?xf32>)
return return
} }

mlir/test/Dialect/Linalg/namedop_conversion.mlir

// RUN: mlir-opt %s -linalg-named-op-conversion -split-input-file | FileCheck %s // RUN: mlir-opt %s -linalg-named-op-conversion -split-input-file | FileCheck %s
// CHECK-LABEL: @depthwise_conv // CHECK-LABEL: @depthwise_conv
func.func @depthwise_conv(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x1xf32>, %arg2: tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32> { func.func @depthwise_conv(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x1xf32>, %arg2: tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32> {
// CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]] // CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
// CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]] // CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
// CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x?x?x?xf32>, tensor<?x?x?xf32>) outs(%[[INIT]] : tensor<?x?x?x?xf32>) // CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]] : tensor<?x?x?x?xf32>, tensor<?x?x?xf32>) inits(%[[INIT]] : tensor<?x?x?x?xf32>)
// CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]] // CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
%0 = linalg.depthwise_conv_2d_nhwc_hwcm {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x1xf32>) outs(%arg2 : tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32> %0 = linalg.depthwise_conv_2d_nhwc_hwcm {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x1xf32>) inits(%arg2 : tensor<?x?x?x?x1xf32>) -> tensor<?x?x?x?x1xf32>
return %0 : tensor<?x?x?x?x1xf32> return %0 : tensor<?x?x?x?x1xf32>
} }
// ----- // -----
// CHECK-LABEL: @depthwise_conv_q // CHECK-LABEL: @depthwise_conv_q
func.func @depthwise_conv_q(%arg0: tensor<?x?x?x?xi8>, %arg1: tensor<?x?x?x1xi8>, %arg2: tensor<?x?x?x?x1xi32>, %arg3 : i32, %arg4 : i32) -> tensor<?x?x?x?x1xi32> { func.func @depthwise_conv_q(%arg0: tensor<?x?x?x?xi8>, %arg1: tensor<?x?x?x1xi8>, %arg2: tensor<?x?x?x?x1xi32>, %arg3 : i32, %arg4 : i32) -> tensor<?x?x?x?x1xi32> {
// CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]] // CHECK-DAG: %[[KERNEL:.+]] = tensor.collapse_shape %arg1 {{\[\[}}0], [1], [2, 3]]
// CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]] // CHECK-DAG: %[[INIT:.+]] = tensor.collapse_shape %arg2 {{\[\[}}0], [1], [2], [3, 4]]
// CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc_q {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]], %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?xi8>, i32, i32) outs(%[[INIT]] : tensor<?x?x?x?xi32>) // CHECK-DAG: %[[CONV:.+]] = linalg.depthwise_conv_2d_nhwc_hwc_q {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %[[KERNEL]], %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?xi8>, i32, i32) inits(%[[INIT]] : tensor<?x?x?x?xi32>)
// CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]] // CHECK: %[[OUT:.+]] = tensor.expand_shape %[[CONV]] {{\[\[}}0], [1], [2], [3, 4]]
%0 = linalg.depthwise_conv_2d_nhwc_hwcm_q {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1, %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?x1xi8>, i32, i32) outs(%arg2 : tensor<?x?x?x?x1xi32>) -> tensor<?x?x?x?x1xi32> %0 = linalg.depthwise_conv_2d_nhwc_hwcm_q {_someattr, dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%arg0, %arg1, %arg3, %arg4 : tensor<?x?x?x?xi8>, tensor<?x?x?x1xi8>, i32, i32) inits(%arg2 : tensor<?x?x?x?x1xi32>) -> tensor<?x?x?x?x1xi32>
return %0 : tensor<?x?x?x?x1xi32> return %0 : tensor<?x?x?x?x1xi32>
} }

mlir/test/Dialect/Linalg/one-shot-bufferize-analysis-2fill-extract-matmul-all-perms.mlir

Show All 12 Linesfunc.func @fill_extract_matmul_1234(
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_1243( func.func @fill_extract_matmul_1243(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_1324( func.func @fill_extract_matmul_1324(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_1342( func.func @fill_extract_matmul_1342(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_1423( func.func @fill_extract_matmul_1423(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_1432( func.func @fill_extract_matmul_1432(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2134( func.func @fill_extract_matmul_2134(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2143( func.func @fill_extract_matmul_2143(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2314( func.func @fill_extract_matmul_2314(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2341( func.func @fill_extract_matmul_2341(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2413( func.func @fill_extract_matmul_2413(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_2431( func.func @fill_extract_matmul_2431(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "false"]} // CHECK: {__inplace_operands_attr__ = ["none", "false"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3124( func.func @fill_extract_matmul_3124(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3142( func.func @fill_extract_matmul_3142(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3214( func.func @fill_extract_matmul_3214(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) -> tensor<256x256xf32> %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3241( func.func @fill_extract_matmul_3241(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %2[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %4 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3412( func.func @fill_extract_matmul_3412(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_3421( func.func @fill_extract_matmul_3421(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4123( func.func @fill_extract_matmul_4123(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4132( func.func @fill_extract_matmul_4132(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4213( func.func @fill_extract_matmul_4213(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<256x256xf32>) -> tensor<256x256xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<256x256xf32>) -> tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %1[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%3, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4231( func.func @fill_extract_matmul_4231(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4312( func.func @fill_extract_matmul_4312(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }
// ----- // -----
// CHECK-LABEL: func @fill_extract_matmul_ // CHECK-LABEL: func @fill_extract_matmul_
func.func @fill_extract_matmul_4321( func.func @fill_extract_matmul_4321(
%arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg0: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false}, %arg1: tensor<518x518xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = false},
%arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true}) %arg2: tensor<256x256xf32> {bufferization.buffer_layout = affine_map<(d0, d1) -> (d0, d1)>, bufferization.writable = true})
-> tensor<256x256xf32> -> tensor<256x256xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst_0 = arith.constant 1.000000e+00 : f32 %cst_0 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<256x256xf32> %0 = bufferization.alloc_tensor() : tensor<256x256xf32>
// CHECK: {__inplace_operands_attr__ = ["false"]} // CHECK: {__inplace_operands_attr__ = ["false"]}
%4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32> %4 = tensor.extract_slice %0[0, 0] [16, 256] [1, 1] : tensor<256x256xf32> to tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["true"]} // CHECK: {__inplace_operands_attr__ = ["true"]}
%3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32> %3 = tensor.extract_slice %0[0, 0] [256, 16] [1, 1] : tensor<256x256xf32> to tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%2 = linalg.fill ins(%cst_0 : f32) outs(%4 : tensor<16x256xf32>) -> tensor<16x256xf32> %2 = linalg.fill ins(%cst_0 : f32) inits(%4 : tensor<16x256xf32>) -> tensor<16x256xf32>
// CHECK: {__inplace_operands_attr__ = ["none", "true"]} // CHECK: {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%3 : tensor<256x16xf32>) -> tensor<256x16xf32> %1 = linalg.fill ins(%cst : f32) inits(%3 : tensor<256x16xf32>) -> tensor<256x16xf32>
// CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]} // CHECK: {__inplace_operands_attr__ = ["true", "true", "true"]}
%5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) outs(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32> %5 = linalg.matmul ins(%1, %2 : tensor<256x16xf32>, tensor<16x256xf32>) inits(%arg2 : tensor<256x256xf32>) -> tensor<256x256xf32>
return %5 : tensor<256x256xf32> return %5 : tensor<256x256xf32>
} }

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
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// AllocTensorOp 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?
%0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32> %0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?xf32>) -> tensor<?xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<?xf32>) -> tensor<?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "false", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "false", "none"]
%2 = tensor.insert_slice %1 into %arg0[0] [%arg1] [1] : tensor<?xf32> into tensor<?xf32> %2 = tensor.insert_slice %1 into %arg0[0] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]
%3 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32> %3 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
Show All 10 Linesfunc.func @buffer_forwarding_no_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__ = ["true", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none"]
// Instead of allocating, share buffer with some inplace bufferization? // Instead of allocating, share buffer with some inplace bufferization?
%0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32> %0 = bufferization.alloc_tensor(%arg1) : tensor<?xf32>
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"] // CHECK-SAME: {__inplace_operands_attr__ = ["none", "true"]
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<?xf32>) -> tensor<?xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<?xf32>) -> tensor<?xf32>
// CHECK: tensor.insert_slice // CHECK: tensor.insert_slice
// CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"] // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true", "none"]
%2 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32> %2 = tensor.insert_slice %1 into %arg0[42] [%arg1] [1] : tensor<?xf32> into tensor<?xf32>
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [0, 0] // CHECK-SAME: __equivalent_func_args__ = [0, 0]
return %2, %2 : tensor<?xf32>, tensor<?xf32> return %2, %2 : tensor<?xf32>, tensor<?xf32>
} }

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

Show All 16 Linesfunc.func @fill_inplace(
%A : tensor<?xf32> {bufferization.writable = true}) %A : tensor<?xf32> {bufferization.writable = true})
-> tensor<?xf32> -> tensor<?xf32>
{ {
// CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
/// Inplaceable, no alloc /// Inplaceable, no alloc
// CHECK-NOT: alloc // CHECK-NOT: alloc
// CHECK: linalg.fill ins(%[[F0]] : f32) outs(%[[A]] : memref<?xf32, strided<[?], offset: ?>>) // CHECK: linalg.fill ins(%[[F0]] : f32) inits(%[[A]] : memref<?xf32, strided<[?], offset: ?>>)
%r = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32> %r = linalg.fill ins(%f0 : f32) inits(%A : tensor<?xf32>) -> tensor<?xf32>
// CHECK: return // CHECK: return
// CHECK-NOT: tensor // CHECK-NOT: tensor
return %r: tensor<?xf32> return %r: tensor<?xf32>
} }
// ----- // -----
/// No bufferization.writable flag, must allocate. /// No bufferization.writable flag, must allocate.
// CHECK-LABEL: func @not_inplace( // CHECK-LABEL: func @not_inplace(
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>) -> memref<?xf32> { // CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>) -> memref<?xf32> {
// CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?xf32>) -> memref<?xf32> // CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?xf32>) -> memref<?xf32>
func.func @not_inplace( func.func @not_inplace(
%A : tensor<?xf32> {bufferization.writable = false}) %A : tensor<?xf32> {bufferization.writable = false})
-> tensor<?xf32> -> tensor<?xf32>
{ {
// CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[F0:.*]] = arith.constant 0.000000e+00 : f32
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
// CHECK: %[[D0:.*]] = memref.dim %[[A]], {{.*}} : memref<?xf32, strided<[?], offset: ?>> // CHECK: %[[D0:.*]] = memref.dim %[[A]], {{.*}} : memref<?xf32, strided<[?], offset: ?>>
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[D0]]) {alignment = 128 : i64} : memref<?xf32> // CHECK: %[[ALLOC:.*]] = memref.alloc(%[[D0]]) {alignment = 128 : i64} : memref<?xf32>
// CHECK: linalg.fill ins(%[[F0]] : f32) outs(%[[ALLOC]] : memref<?xf32>) // CHECK: linalg.fill ins(%[[F0]] : f32) inits(%[[ALLOC]] : memref<?xf32>)
%r = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32> %r = linalg.fill ins(%f0 : f32) inits(%A : tensor<?xf32>) -> tensor<?xf32>
// CHECK-NOT: dealloc // CHECK-NOT: dealloc
// CHECK: return %[[ALLOC]] : memref<?xf32> // CHECK: return %[[ALLOC]] : memref<?xf32>
return %r: tensor<?xf32> return %r: tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @not_inplace // CHECK-LABEL: func @not_inplace
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?x?xf32, strided<[?, ?], offset: ?>>) { // CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?x?xf32, strided<[?, ?], offset: ?>>) {
// CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?x?xf32>) { // CHECK-NO-LAYOUT-MAP-LABEL: func @not_inplace(%{{.*}}: memref<?x?xf32>) {
func.func @not_inplace( func.func @not_inplace(
%A : tensor<?x?xf32> {bufferization.writable = true}) %A : tensor<?x?xf32> {bufferization.writable = true})
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
/// Cross-op multiple uses of %A, the first op which has interfering reads must alloc. /// Cross-op multiple uses of %A, the first op which has interfering reads must alloc.
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[ALLOC]] // CHECK: linalg.fill ins({{.*}}{{.*}}inits(%[[ALLOC]]
%f = linalg.fill ins(%f0 : f32) outs(%A : tensor<?x?xf32>) -> tensor<?x?xf32> %f = linalg.fill ins(%f0 : f32) inits(%A : tensor<?x?xf32>) -> tensor<?x?xf32>
/// The second op has no interfering reads and can reuse. /// The second op has no interfering reads and can reuse.
// CHECK-NOT: alloc // CHECK-NOT: alloc
// CHECK: linalg.matmul ins(%[[ALLOC]], %[[ALLOC]]{{.*}}) outs(%[[A]] // CHECK: linalg.matmul ins(%[[ALLOC]], %[[ALLOC]]{{.*}}) inits(%[[A]]
%r = linalg.matmul ins(%f, %f: tensor<?x?xf32>, tensor<?x?xf32>) %r = linalg.matmul ins(%f, %f: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%A: tensor<?x?xf32>) inits(%A: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: memref.dealloc %[[ALLOC]] // CHECK: memref.dealloc %[[ALLOC]]
// CHECK: return // CHECK: return
// CHECK-NOT: tensor // CHECK-NOT: tensor
return %r: tensor<?x?xf32> return %r: tensor<?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @not_inplace // CHECK-LABEL: func @not_inplace
func.func @not_inplace( func.func @not_inplace(
%A : tensor<?x?xf32> {bufferization.writable = true}) -> tensor<?x?xf32> { %A : tensor<?x?xf32> {bufferization.writable = true}) -> tensor<?x?xf32> {
/// Within op multiple uses of %A, must alloc. /// Within op multiple uses of %A, must alloc.
// CHECK: alloc // CHECK: alloc
%r = linalg.matmul ins(%A, %A: tensor<?x?xf32>, tensor<?x?xf32>) %r = linalg.matmul ins(%A, %A: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%A: tensor<?x?xf32>) inits(%A: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK-NOT: dealloc // CHECK-NOT: dealloc
return %r: tensor<?x?xf32> return %r: tensor<?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @vec_inplace // CHECK-LABEL: func @vec_inplace
func.func @vec_inplace( func.func @vec_inplace(
▲ Show 20 Lines • Show All 73 Lines • ▼ Show 20 Lines %2 = scf.for %arg5 = %c0 to %c192 step %c16 iter_args(%arg6 = %arg4) -> (tensor<128x192xf32>) {
// C was already replaced with a copy by preprocessing, so no copy is // C was already replaced with a copy by preprocessing, so no copy is
// needed here. // needed here.
// CHECK: %[[C_SLICE:.*]] = memref.subview %[[ALLOC]] // CHECK: %[[C_SLICE:.*]] = memref.subview %[[ALLOC]]
%4 = tensor.extract_slice %C[%arg3, %arg5] [8, 16] [1, 1] : %4 = tensor.extract_slice %C[%arg3, %arg5] [8, 16] [1, 1] :
tensor<128x192xf32> to tensor<8x16xf32> tensor<128x192xf32> to tensor<8x16xf32>
// linalg.fill is inplace. // linalg.fill is inplace.
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[C_SLICE]] // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[C_SLICE]]
%5 = linalg.fill ins(%cst : f32) outs(%4 : tensor<8x16xf32>) -> tensor<8x16xf32> %5 = linalg.fill ins(%cst : f32) inits(%4 : tensor<8x16xf32>) -> tensor<8x16xf32>
// CHECK: scf.for %[[K:.*]] = // CHECK: scf.for %[[K:.*]] =
%6 = scf.for %arg7 = %c0 to %c256 step %c32 iter_args(%arg8 = %5) -> (tensor<8x16xf32>) { %6 = scf.for %arg7 = %c0 to %c256 step %c32 iter_args(%arg8 = %5) -> (tensor<8x16xf32>) {
%8 = tensor.extract_slice %1[0, %arg7] [8, 32] [1, 1] : %8 = tensor.extract_slice %1[0, %arg7] [8, 32] [1, 1] :
tensor<8x256xf32> to tensor<8x32xf32> tensor<8x256xf32> to tensor<8x32xf32>
%9 = tensor.extract_slice %3[%arg7, 0] [32, 16] [1, 1] : %9 = tensor.extract_slice %3[%arg7, 0] [32, 16] [1, 1] :
tensor<256x16xf32> to tensor<32x16xf32> tensor<256x16xf32> to tensor<32x16xf32>
// linalg.matmul is inplace as well as the enclosing scf.for. // linalg.matmul is inplace as well as the enclosing scf.for.
// CHECK: linalg.matmul ins({{.*}} outs(%[[C_SLICE]] // CHECK: linalg.matmul ins({{.*}} inits(%[[C_SLICE]]
%10 = linalg.matmul ins(%8, %9 : tensor<8x32xf32>, tensor<32x16xf32>) %10 = linalg.matmul ins(%8, %9 : tensor<8x32xf32>, tensor<32x16xf32>)
outs(%arg8 : tensor<8x16xf32>) inits(%arg8 : tensor<8x16xf32>)
-> tensor<8x16xf32> -> tensor<8x16xf32>
scf.yield %10 : tensor<8x16xf32> scf.yield %10 : tensor<8x16xf32>
} }
// insert_slice is inplace but its source comes from an equivalent buffer // insert_slice is inplace but its source comes from an equivalent buffer
// that is not in place. So we must insert a copy of the small buffer into // that is not in place. So we must insert a copy of the small buffer into
// the bigger buffer. // the bigger buffer.
// CHECK: %[[T:.*]] = memref.subview %[[C]][%[[I]], %[[J]]] [8, 16] [1, 1] // CHECK: %[[T:.*]] = memref.subview %[[C]][%[[I]], %[[J]]] [8, 16] [1, 1]
Show All 20 Linesfunc.func @dominance_violation_bug_1(
%A : tensor<?x?xf32> {bufferization.writable = false}, %A : tensor<?x?xf32> {bufferization.writable = false},
%idx : index) %idx : index)
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%sA = tensor.extract_slice %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %sA = tensor.extract_slice %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%ssA = tensor.extract_slice %sA[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32> %ssA = tensor.extract_slice %sA[0, 0][4, 4][1, 1] : tensor<?x?xf32> to tensor<4x4xf32>
%FA = linalg.fill ins(%f0 : f32) outs(%ssA : tensor<4x4xf32>) -> tensor<4x4xf32> %FA = linalg.fill ins(%f0 : f32) inits(%ssA : tensor<4x4xf32>) -> tensor<4x4xf32>
%rsA = tensor.insert_slice %FA into %sA[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<?x?xf32> %rsA = tensor.insert_slice %FA into %sA[0, 0][4, 4][1, 1] : tensor<4x4xf32> into tensor<?x?xf32>
%rA = tensor.insert_slice %rsA into %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %rA = tensor.insert_slice %rsA into %A[0, 0][%idx, %idx][1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
return %rA : tensor<?x?xf32> return %rA : tensor<?x?xf32>
} }
// ----- // -----
func.func @gather_like( func.func @gather_like(
%arg0 : tensor<?x?xf32> {bufferization.writable = false}, %arg0 : tensor<?x?xf32> {bufferization.writable = false},
%arg1 : tensor<?xi32> {bufferization.writable = false}, %arg1 : tensor<?xi32> {bufferization.writable = false},
%arg2 : tensor<?x?xf32> {bufferization.writable = true}) %arg2 : tensor<?x?xf32> {bufferization.writable = true})
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
%0 = linalg.generic { %0 = 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(%arg1 : tensor<?xi32>) outs(%arg2 : tensor<?x?xf32>) { ins(%arg1 : tensor<?xi32>) inits(%arg2 : tensor<?x?xf32>) {
^bb0(%arg3: i32, %arg4 : f32): ^bb0(%arg3: i32, %arg4 : f32):
%iv1 = linalg.index 1 : index %iv1 = linalg.index 1 : index
%1 = arith.index_cast %arg3: i32 to index %1 = arith.index_cast %arg3: i32 to index
%2 = tensor.extract %arg0[%1, %iv1] : tensor<?x?xf32> %2 = tensor.extract %arg0[%1, %iv1] : tensor<?x?xf32>
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// CHECK-LABEL: func @gather_like( // CHECK-LABEL: func @gather_like(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32, // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32,
// CHECK-SAME: %[[ARG1:.+]]: memref<?xi32 // CHECK-SAME: %[[ARG1:.+]]: memref<?xi32
// CHECK-SAME: %[[ARG2:.+]]: memref<?x?xf32 // CHECK-SAME: %[[ARG2:.+]]: memref<?x?xf32
// CHECK-SAME: ) { // CHECK-SAME: ) {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[ARG1]] : // CHECK-SAME: ins(%[[ARG1]] :
// CHECK-SAME: outs(%[[ARG2]] : // CHECK-SAME: inits(%[[ARG2]] :
// CHECK: %[[YIELD:.+]] = memref.load %[[ARG0]] // CHECK: %[[YIELD:.+]] = memref.load %[[ARG0]]
// CHECK: linalg.yield %[[YIELD]] // CHECK: linalg.yield %[[YIELD]]
// ----- // -----
// CHECK-LABEL: func @linalg_op_bufferizes_inplace_with_input // CHECK-LABEL: func @linalg_op_bufferizes_inplace_with_input
// CHECK-SAME: %[[t1:.*]]: memref<?x?xf32, strided{{.*}}>, %[[t2:.*]]: memref<?xf32, strided{{.*}}>, %[[t3:.*]]: memref<?x?xf32, strided{{.*}}> // CHECK-SAME: %[[t1:.*]]: memref<?x?xf32, strided{{.*}}>, %[[t2:.*]]: memref<?xf32, strided{{.*}}>, %[[t3:.*]]: memref<?x?xf32, strided{{.*}}>
func.func @linalg_op_bufferizes_inplace_with_input( func.func @linalg_op_bufferizes_inplace_with_input(
%t1: tensor<?x?xf32> {bufferization.writable = true}, %t1: tensor<?x?xf32> {bufferization.writable = true},
%t2: tensor<?xf32> {bufferization.writable = true}, %t2: tensor<?xf32> {bufferization.writable = true},
%t3: tensor<?x?xf32> {bufferization.writable = true}, %t3: tensor<?x?xf32> {bufferization.writable = true},
%s1: index, %s2: index, %cst: f32) %s1: index, %s2: index, %cst: f32)
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
// CHECK: linalg.generic {{.*}} ins(%[[t1]], %[[t2]] : {{.*}}) outs(%[[t3]] : {{.*}}) // CHECK: linalg.generic {{.*}} ins(%[[t1]], %[[t2]] : {{.*}}) inits(%[[t3]] : {{.*}})
%r = linalg.generic { %r = linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1)-> (d0, d1)>], affine_map<(d0, d1)-> (d0, d1)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%t1, %t2 : tensor<?x?xf32>, tensor<?xf32>) ins(%t1, %t2 : tensor<?x?xf32>, tensor<?xf32>)
outs(%t3 : tensor<?x?xf32>) { inits(%t3 : tensor<?x?xf32>) {
^bb0(%arg0 : f32, %arg1 : f32, %arg2 : f32) : ^bb0(%arg0 : f32, %arg1 : f32, %arg2 : f32) :
%add = arith.addf %arg0, %arg1 : f32 %add = arith.addf %arg0, %arg1 : f32
linalg.yield %add : f32 linalg.yield %add : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %r : tensor<?x?xf32> return %r : tensor<?x?xf32>
} }
// ----- // -----
Show All 11 Lines
func.func @op_is_reading_but_following_ops_are_not( func.func @op_is_reading_but_following_ops_are_not(
%t0 : tensor<?xf32> {bufferization.writable = false}, %t0 : tensor<?xf32> {bufferization.writable = false},
%cst : f32) %cst : f32)
-> tensor<?xf32> -> tensor<?xf32>
{ {
// Make sure that a copy is inserted here. // Make sure that a copy is inserted here.
// CHECK: %[[ALLOC:.*]] = memref.alloc // CHECK: %[[ALLOC:.*]] = memref.alloc
// CHECK: memref.copy %[[t0]], %[[ALLOC]] // CHECK: memref.copy %[[t0]], %[[ALLOC]]
// CHECK: linalg.generic {{.*}} outs(%[[ALLOC]] : memref // CHECK: linalg.generic {{.*}} inits(%[[ALLOC]] : memref
%r0 =linalg.generic #trait outs (%t0 : tensor<?xf32>) { %r0 =linalg.generic #trait inits (%t0 : tensor<?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
%a = arith.addf %cst, %0 : f32 %a = arith.addf %cst, %0 : f32
linalg.yield %a : f32 linalg.yield %a : f32
} -> (tensor<?xf32>) } -> (tensor<?xf32>)
// CHECK: linalg.generic {{.*}} outs(%[[ALLOC]] : memref // CHECK: linalg.generic {{.*}} inits(%[[ALLOC]] : memref
%r1 = linalg.generic #trait outs (%r0 : tensor<?xf32>) { %r1 = linalg.generic #trait inits (%r0 : tensor<?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> (tensor<?xf32>) } -> (tensor<?xf32>)
// CHECK: return %[[ALLOC]] // CHECK: return %[[ALLOC]]
return %r1 : tensor<?xf32> return %r1 : tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @map_binary // CHECK-LABEL: func @map_binary
// CHECK-SAME: %[[LHS:[0-9a-zA-Z]*]]: memref<64xf32 // CHECK-SAME: %[[LHS:[0-9a-zA-Z]*]]: memref<64xf32
// CHECK-SAME: %[[RHS:[0-9a-zA-Z]*]]: memref<64xf32 // CHECK-SAME: %[[RHS:[0-9a-zA-Z]*]]: memref<64xf32
func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>, func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
%init: tensor<64xf32>) -> tensor<64xf32> { %init: tensor<64xf32>) -> tensor<64xf32> {
// CHECK: linalg.map // CHECK: linalg.map
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<64xf32 // CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<64xf32
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %add : tensor<64xf32> func.return %add : tensor<64xf32>
} }
// ----- // -----
// CHECK-LABEL: func @reduce // CHECK-LABEL: func @reduce
// CHECK-SAME: %[[INPUT:.*]]: memref<16x32x64xf32 // CHECK-SAME: %[[INPUT:.*]]: memref<16x32x64xf32
func.func @reduce(%input: tensor<16x32x64xf32>, func.func @reduce(%input: tensor<16x32x64xf32>,
%init: tensor<16x64xf32>) -> tensor<16x64xf32> { %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
// CHECK: linalg.reduce // CHECK: linalg.reduce
// CHECK-SAME: ins(%[[INPUT]] : memref<16x32x64xf32 // CHECK-SAME: ins(%[[INPUT]] : memref<16x32x64xf32
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16x64xf32> func.return %reduce : tensor<16x64xf32>
} }
// ----- // -----
// CHECK-LABEL: func @transpose // CHECK-LABEL: func @transpose
// CHECK-SAME: %[[ARG0:.*]]: memref<16x32x64xf32 // CHECK-SAME: %[[ARG0:.*]]: memref<16x32x64xf32
func.func @transpose(%input: tensor<16x32x64xf32>, func.func @transpose(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
// CHECK: linalg.transpose // CHECK: linalg.transpose
// CHECK-SAME: ins(%[[ARG0]] : memref<16x32x64xf32 // CHECK-SAME: ins(%[[ARG0]] : memref<16x32x64xf32
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [1, 2, 0] permutation = [1, 2, 0]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }
// ----- // -----
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// AllocTensorOp elimination would produce SSA violations for the example below. // AllocTensorOp elimination would produce SSA violations for the example below.
Show All 35 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 = tensor.empty(%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>) inits(%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>
%1 = tensor.pad %0 low [%arg1, %arg2] high [%arg3, %arg4] { %1 = tensor.pad %0 low [%arg1, %arg2] high [%arg3, %arg4] {
^bb0(%arg6: index, %arg7 : index): ^bb0(%arg6: index, %arg7 : index):
tensor.yield %arg5 : f32 tensor.yield %arg5 : f32
} : tensor<?x?xf32> to tensor<?x?xf32> } : tensor<?x?xf32> to tensor<?x?xf32>
Show All 11 Lines
// 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:.+]] = tensor.empty(%[[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]]{{.*}}inits(%[[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: inits(%[[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 = tensor.empty(%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>) inits(%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>
%1 = tensor.pad %0 low [3, %arg1] high [4, %arg2] { %1 = tensor.pad %0 low [3, %arg1] high [4, %arg2] {
^bb0(%arg4: index, %arg5 : index): ^bb0(%arg4: index, %arg5 : index):
tensor.yield %arg3 : f32 tensor.yield %arg3 : f32
} : tensor<42x?xf32> to tensor<49x?xf32> } : tensor<42x?xf32> to tensor<49x?xf32>
return %1 : tensor<49x?xf32> return %1 : tensor<49x?xf32>
} }
// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + s0 + s1)> // CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s2 + s0 + s1)>
// CHECK: func @mixed_pad_fusion // CHECK: func @mixed_pad_fusion
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x42xf32> // CHECK-SAME: %[[ARG0:.+]]: tensor<?x42xf32>
// 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:.+]] = tensor.empty(%[[TARGET_D1]]) // CHECK: %[[INIT:.+]] = tensor.empty(%[[TARGET_D1]])
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG3]]{{.*}}outs(%[[INIT]] // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG3]]{{.*}}inits(%[[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: inits(%[[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/parallel-loops.mlir

// RUN: mlir-opt %s -convert-linalg-to-parallel-loops -split-input-file | FileCheck %s // RUN: mlir-opt %s -convert-linalg-to-parallel-loops -split-input-file | FileCheck %s
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
func.func @linalg_generic_sum(%lhs: memref<2x2xf32>, func.func @linalg_generic_sum(%lhs: memref<2x2xf32>,
%rhs: memref<2x2xf32>, %rhs: memref<2x2xf32>,
%sum: memref<2x2xf32>) { %sum: memref<2x2xf32>) {
linalg.generic { linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%lhs, %rhs : memref<2x2xf32>, memref<2x2xf32>) ins(%lhs, %rhs : memref<2x2xf32>, memref<2x2xf32>)
outs(%sum : memref<2x2xf32>) { inits(%sum : memref<2x2xf32>) {
^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32): ^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32):
%0 = arith.addf %lhs_in, %rhs_in : f32 %0 = arith.addf %lhs_in, %rhs_in : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
// CHECK-LABEL: @linalg_generic_sum // CHECK-LABEL: @linalg_generic_sum
// CHECK: (%[[LHS:.*]]:{{.*}}, %[[RHS:.*]]:{{.*}}, %[[SUM:.*]]:{{.*}}) // CHECK: (%[[LHS:.*]]:{{.*}}, %[[RHS:.*]]:{{.*}}, %[[SUM:.*]]:{{.*}})
Show All 16 Lines
#trait = { #trait = {
iterator_types = ["parallel", "parallel", "reduction", "parallel"], iterator_types = ["parallel", "parallel", "reduction", "parallel"],
indexing_maps = #accesses indexing_maps = #accesses
} }
func.func @lower_outer_parallel(%A: memref<?x?x?x?xf32>, %B: memref<?x?x?xf32>) { func.func @lower_outer_parallel(%A: memref<?x?x?x?xf32>, %B: memref<?x?x?xf32>) {
linalg.generic #trait linalg.generic #trait
ins(%A : memref<?x?x?x?xf32>) ins(%A : memref<?x?x?x?xf32>)
outs(%B : memref<?x?x?xf32>) { inits(%B : memref<?x?x?xf32>) {
^bb0(%a: f32, %b: f32): ^bb0(%a: f32, %b: f32):
linalg.yield %a: f32 linalg.yield %a: f32
} }
return return
} }
// CHECK-LABEL: @lower_outer_parallel // CHECK-LABEL: @lower_outer_parallel
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 // CHECK-DAG: %[[C0:.*]] = arith.constant 0
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 // CHECK-DAG: %[[C1:.*]] = arith.constant 1
Show All 16 Lines
#trait = { #trait = {
iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"], iterator_types = ["parallel", "parallel", "reduction", "parallel", "parallel", "reduction"],
indexing_maps = #accesses indexing_maps = #accesses
} }
func.func @lower_mixed_parallel(%A: memref<?x?x?x?x?x?xf32>, %B: memref<?x?x?x?xf32>) { func.func @lower_mixed_parallel(%A: memref<?x?x?x?x?x?xf32>, %B: memref<?x?x?x?xf32>) {
linalg.generic #trait linalg.generic #trait
ins(%A : memref<?x?x?x?x?x?xf32>) ins(%A : memref<?x?x?x?x?x?xf32>)
outs(%B : memref<?x?x?x?xf32>) { inits(%B : memref<?x?x?x?xf32>) {
^bb0(%a: f32, %b: f32): ^bb0(%a: f32, %b: f32):
linalg.yield %a: f32 linalg.yield %a: f32
} }
return return
} }
// CHECK-LABEL: @lower_mixed_parallel // CHECK-LABEL: @lower_mixed_parallel
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 // CHECK-DAG: %[[C0:.*]] = arith.constant 0
// CHECK-DAG: %[[C1:.*]] = arith.constant 1 // CHECK-DAG: %[[C1:.*]] = arith.constant 1
Show All 12 Lines

mlir/test/Dialect/Linalg/promote.mlir

Show All 19 Lines scf.for %arg4 = %c0 to %6 step %c2 {
scf.for %arg5 = %c0 to %8 step %c3 { scf.for %arg5 = %c0 to %8 step %c3 {
scf.for %arg6 = %c0 to %7 step %c4 { scf.for %arg6 = %c0 to %7 step %c4 {
%11 = memref.subview %3[%arg4, %arg6][%c2, %c4][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> %11 = memref.subview %3[%arg4, %arg6][%c2, %c4][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
%14 = memref.subview %4[%arg6, %arg5][%c4, %c3][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> %14 = memref.subview %4[%arg6, %arg5][%c4, %c3][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
%17 = memref.subview %5[%arg4, %arg5][%c2, %c3][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> %17 = memref.subview %5[%arg4, %arg5][%c2, %c3][1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
linalg.matmul linalg.matmul
ins(%11, %14: memref<?x?xf32, strided<[?, 1], offset: ?>>, ins(%11, %14: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>) memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%17: memref<?x?xf32, strided<[?, 1], offset: ?>>) inits(%17: memref<?x?xf32, strided<[?, 1], offset: ?>>)
} }
} }
} }
return return
} }
// CHECK-LABEL: func @matmul_f32(%{{.*}}: memref<?xi8>, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) { // CHECK-LABEL: func @matmul_f32(%{{.*}}: memref<?xi8>, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
// CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
Show All 14 Lines
// CHECK: %[[tmpC:.*]] = memref.alloca() : memref<24xi8> // CHECK: %[[tmpC:.*]] = memref.alloca() : memref<24xi8>
// CHECK: %[[fullC:.*]] = memref.view %[[tmpC]][{{.*}}][{{.*}}] : memref<24xi8> to memref<?x?xf32> // CHECK: %[[fullC:.*]] = memref.view %[[tmpC]][{{.*}}][{{.*}}] : memref<24xi8> to memref<?x?xf32>
// CHECK: %[[partialC:.*]] = memref.subview %[[fullC]]{{.*}} : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: %[[partialC:.*]] = memref.subview %[[fullC]]{{.*}} : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vA]], %[[partialA]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vA]], %[[partialA]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vB]], %[[partialB]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vB]], %[[partialB]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vC]], %[[partialC]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vC]], %[[partialC]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// //
// CHECK: linalg.matmul ins(%[[partialA]], %[[partialB]]{{.*}} outs(%[[partialC]] // CHECK: linalg.matmul ins(%[[partialA]], %[[partialB]]{{.*}} inits(%[[partialC]]
// //
// CHECK: memref.copy %[[partialC]], %[[vC]] : // CHECK: memref.copy %[[partialC]], %[[vC]] :
// CHECK: memref<?x?xf32, strided<[?, 1], offset: ?>> to // CHECK: memref<?x?xf32, strided<[?, 1], offset: ?>> to
// CHECK: memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: memref<?x?xf32, strided<[?, 1], offset: ?>>
// //
// CHECK-NOT: memref.dealloc %[[tmpA]] : memref<32xi8> // CHECK-NOT: memref.dealloc %[[tmpA]] : memref<32xi8>
// CHECK-NOT: memref.dealloc %[[tmpB]] : memref<48xi8> // CHECK-NOT: memref.dealloc %[[tmpB]] : memref<48xi8>
// CHECK-NOT: memref.dealloc %[[tmpC]] : memref<24xi8> // CHECK-NOT: memref.dealloc %[[tmpC]] : memref<24xi8>
Show All 22 Lines scf.for %arg4 = %c0 to %6 step %c2 {
scf.for %arg5 = %c0 to %8 step %c3 { scf.for %arg5 = %c0 to %8 step %c3 {
scf.for %arg6 = %c0 to %7 step %c4 { scf.for %arg6 = %c0 to %7 step %c4 {
%11 = memref.subview %3[%arg4, %arg6][%c2, %c4][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>> %11 = memref.subview %3[%arg4, %arg6][%c2, %c4][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>>
%14 = memref.subview %4[%arg6, %arg5][%c4, %c3][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>> %14 = memref.subview %4[%arg6, %arg5][%c4, %c3][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>>
%17 = memref.subview %5[%arg4, %arg5][%c2, %c3][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>> %17 = memref.subview %5[%arg4, %arg5][%c2, %c3][1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>>
linalg.matmul linalg.matmul
ins(%11, %14: memref<?x?xf64, strided<[?, 1], offset: ?>>, ins(%11, %14: memref<?x?xf64, strided<[?, 1], offset: ?>>,
memref<?x?xf64, strided<[?, 1], offset: ?>>) memref<?x?xf64, strided<[?, 1], offset: ?>>)
outs(%17: memref<?x?xf64, strided<[?, 1], offset: ?>>) inits(%17: memref<?x?xf64, strided<[?, 1], offset: ?>>)
} }
} }
} }
return return
} }
// CHECK-LABEL: func @matmul_f64(%{{.*}}: memref<?xi8>, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) { // CHECK-LABEL: func @matmul_f64(%{{.*}}: memref<?xi8>, %{{.*}}: index, %{{.*}}: index, %{{.*}}: index) {
// CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} { // CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
Show All 14 Lines
// CHECK: %[[tmpC_f64:.*]] = memref.alloc() : memref<48xi8> // CHECK: %[[tmpC_f64:.*]] = memref.alloc() : memref<48xi8>
// CHECK: %[[fullC_f64:.*]] = memref.view %[[tmpC_f64]][{{.*}}][{{.*}}] : memref<48xi8> to memref<?x?xf64> // CHECK: %[[fullC_f64:.*]] = memref.view %[[tmpC_f64]][{{.*}}][{{.*}}] : memref<48xi8> to memref<?x?xf64>
// CHECK: %[[partialC_f64:.*]] = memref.subview %[[fullC_f64]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>> // CHECK: %[[partialC_f64:.*]] = memref.subview %[[fullC_f64]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf64> to memref<?x?xf64, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vA_f64]], %[[partialA_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vA_f64]], %[[partialA_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vB_f64]], %[[partialB_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vB_f64]], %[[partialB_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[vC_f64]], %[[partialC_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>> // CHECK: memref.copy %[[vC_f64]], %[[partialC_f64]] : memref<?x?xf64, strided<[?, 1], offset: ?>> to memref<?x?xf64, strided<[?, 1], offset: ?>>
// //
// CHECK: linalg.matmul ins(%[[partialA_f64]], %[[partialB_f64]]{{.*}} outs(%[[partialC_f64]] // CHECK: linalg.matmul ins(%[[partialA_f64]], %[[partialB_f64]]{{.*}} inits(%[[partialC_f64]]
// //
// CHECK: memref.copy %[[partialC_f64]], %[[vC_f64]] : // CHECK: memref.copy %[[partialC_f64]], %[[vC_f64]] :
// CHECK: memref<?x?xf64, strided<[?, 1], offset: ?>> to // CHECK: memref<?x?xf64, strided<[?, 1], offset: ?>> to
// CHECK: memref<?x?xf64, strided<[?, 1], offset: ?>> // CHECK: memref<?x?xf64, strided<[?, 1], offset: ?>>
// //
// CHECK: memref.dealloc %[[tmpA_f64]] : memref<64xi8> // CHECK: memref.dealloc %[[tmpA_f64]] : memref<64xi8>
// CHECK: memref.dealloc %[[tmpB_f64]] : memref<96xi8> // CHECK: memref.dealloc %[[tmpB_f64]] : memref<96xi8>
// CHECK: memref.dealloc %[[tmpC_f64]] : memref<48xi8> // CHECK: memref.dealloc %[[tmpC_f64]] : memref<48xi8>
Show All 27 Linesfunc.func @promote_rank_reducing_subviews(%arg0: memref<?x?x?x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg1: memref<128x3x3x64xf32, strided<[?, ?, ?, ?], offset: ?>>, %arg2: memref<?x?x?x128xf32>,
// CHECK: memref.alloc // CHECK: memref.alloc
// CHECK: %[[c_view:.+]] = memref.view // CHECK: %[[c_view:.+]] = memref.view
// CHECK: %[[c_pro_subview:.+]] = memref.subview %[[c_view]] // CHECK: %[[c_pro_subview:.+]] = memref.subview %[[c_view]]
// CHECK-COUNT-3: memref.copy // CHECK-COUNT-3: memref.copy
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[a_pro_subview]], %[[b_pro_subview]] // CHECK-SAME: ins(%[[a_pro_subview]], %[[b_pro_subview]]
// CHECK-SAME: outs(%[[c_pro_subview]] // CHECK-SAME: inits(%[[c_pro_subview]]
linalg.generic {indexing_maps = [#map6, #map7, #map8], iterator_types = ["parallel", "parallel", "reduction"]} ins(%13, %14 : memref<?x32xf32, strided<[?, ?], offset: ?>>, memref<128x32xf32, strided<[?, ?], offset: ?>>) outs(%9 : memref<?x128xf32, strided<[128, 1], offset: ?>>) { linalg.generic {indexing_maps = [#map6, #map7, #map8], iterator_types = ["parallel", "parallel", "reduction"]} ins(%13, %14 : memref<?x32xf32, strided<[?, ?], offset: ?>>, memref<128x32xf32, strided<[?, ?], offset: ?>>) inits(%9 : memref<?x128xf32, strided<[128, 1], offset: ?>>) {
^bb0(%arg9: f32, %arg10: f32, %arg11: f32): ^bb0(%arg9: f32, %arg10: f32, %arg11: f32):
%15 = arith.mulf %arg9, %arg10 : f32 %15 = arith.mulf %arg9, %arg10 : f32
%16 = arith.addf %arg11, %15 : f32 %16 = arith.addf %arg11, %15 : f32
linalg.yield %16 : f32 linalg.yield %16 : f32
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match interface{LinalgOp} in %arg1 %0 = transform.structured.match interface{LinalgOp} in %arg1
%1 = transform.structured.promote %0 %1 = transform.structured.promote %0
} }

mlir/test/Dialect/Linalg/promotion_options.mlir

// RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s
func.func @gemm(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>) func.func @gemm(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)
{ {
linalg.matmul ins(%a, %b: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%a, %b: memref<?x?xf32>, memref<?x?xf32>)
outs(%c: memref<?x?xf32>) inits(%c: memref<?x?xf32>)
return return
} }
// CHECK: func @gemm // CHECK: func @gemm
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK: scf.for // CHECK: scf.for
// CHECK: scf.for // CHECK: scf.for
// CHECK: scf.for // CHECK: scf.for
// CHECK: %[[T7:.+]] = memref.subview %[[ARG0]] // CHECK: %[[T7:.+]] = memref.subview %[[ARG0]]
// CHECK: %[[T12:.+]] = memref.subview %[[ARG1]] // CHECK: %[[T12:.+]] = memref.subview %[[ARG1]]
// CHECK: %[[T17:.+]] = memref.subview %[[ARG2]] // CHECK: %[[T17:.+]] = memref.subview %[[ARG2]]
// CHECK: %[[A0:.*]] = memref.alloc() : memref<1024xi8> // CHECK: %[[A0:.*]] = memref.alloc() : memref<1024xi8>
// CHECK: %[[V0:.*]] = memref.view %[[A0]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32> // CHECK: %[[V0:.*]] = memref.view %[[A0]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32>
// CHECK: %[[T19:.+]] = memref.subview %[[V0]] // CHECK: %[[T19:.+]] = memref.subview %[[V0]]
// CHECK: %[[A1:.*]] = memref.alloc() : memref<1024xi8> // CHECK: %[[A1:.*]] = memref.alloc() : memref<1024xi8>
// CHECK: %[[V1:.*]] = memref.view %[[A1]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32> // CHECK: %[[V1:.*]] = memref.view %[[A1]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32>
// CHECK: %[[T21:.+]] = memref.subview %[[V1]] // CHECK: %[[T21:.+]] = memref.subview %[[V1]]
// CHECK: memref.copy %[[T7]], %[[T19]] // CHECK: memref.copy %[[T7]], %[[T19]]
// CHECK: memref.copy %[[T17]], %[[T21]] // CHECK: memref.copy %[[T17]], %[[T21]]
// CHECK: linalg.matmul ins(%[[T19]], %[[T12]]{{.*}} outs(%[[T21]] // CHECK: linalg.matmul ins(%[[T19]], %[[T12]]{{.*}} inits(%[[T21]]
// CHECK: memref.copy %[[T21]], %[[T17]] // CHECK: memref.copy %[[T21]], %[[T17]]
// CHECK: memref.dealloc %[[A0]] // CHECK: memref.dealloc %[[A0]]
// CHECK: memref.dealloc %[[A1]] // CHECK: memref.dealloc %[[A1]]
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [16, 16, 16] %1, %loops:3 = transform.structured.tile %0 [16, 16, 16]
%2 = transform.structured.promote %1 { operands_to_promote = [0, 2], force_full_tiles = [false, false] } %2 = transform.structured.promote %1 { operands_to_promote = [0, 2], force_full_tiles = [false, false] }
} }

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 = tensor.empty(%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>) { inits(%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
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
// 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)>
Show All 16 Linesfunc.func @control_consumer_reshape_fusion(%arg0 : tensor<1x?x?xf32>, %arg1 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32> {
%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 = tensor.empty(%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>) { inits(%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>
%1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x?x?xf32>, tensor<1x?x?xf32>) %1 = linalg.batch_matmul ins(%arg0, %arg1 : tensor<1x?x?xf32>, tensor<1x?x?xf32>)
outs(%0 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32> inits(%0 : tensor<1x?x?xf32>) -> tensor<1x?x?xf32>
return %1 : tensor<1x?x?xf32> return %1 : tensor<1x?x?xf32>
} }
// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2) // CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)
// CHECK: func @control_consumer_reshape_fusion // CHECK: func @control_consumer_reshape_fusion
// CHECK: %[[FILL:.+]] = linalg.generic // CHECK: %[[FILL:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP]]] // CHECK-SAME: indexing_maps = [#[[MAP]]]
// CHECK-SAME: outs(%{{.+}} : tensor<1x?x?xf32>) // CHECK-SAME: inits(%{{.+}} : tensor<1x?x?xf32>)
// CHECK: linalg.batch_matmul // CHECK: linalg.batch_matmul
// CHECK-SAME: outs(%[[FILL]] : tensor<1x?x?xf32>) // CHECK-SAME: inits(%[[FILL]] : tensor<1x?x?xf32>)

mlir/test/Dialect/Linalg/reshape_fusion.mlir

// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-expansion -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-expansion -split-input-file | FileCheck %s
#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)> #map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
#map2 = affine_map<(d0, d1, d2) -> ()> #map2 = affine_map<(d0, d1, d2) -> ()>
func.func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xf32>, func.func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xf32>,
%arg1 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>,
%arg2 : f32) -> %arg2 : f32) ->
tensor<?x?x?xf32> tensor<?x?x?xf32>
{ {
%0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] : %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] :
tensor<?x?x4x?xf32> into tensor<?x?x?xf32> tensor<?x?x4x?xf32> into tensor<?x?x?xf32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map2, #map1], indexing_maps = [#map0, #map1, #map2, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%0, %arg1, %arg2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32) ins(%0, %arg1, %arg2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32)
outs(%arg1 : tensor<?x?x?xf32>) { inits(%arg1 : tensor<?x?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):
%1 = arith.mulf %arg3, %arg4 : f32 %1 = arith.mulf %arg3, %arg4 : f32
%2 = arith.addf %1, %arg5 : f32 %2 = arith.addf %1, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
return %1 : tensor<?x?x?xf32> return %1 : tensor<?x?x?xf32>
} }
// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)> // CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)> // CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>
// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> ()> // CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> ()>
// CHECK: func @generic_op_reshape_producer_fusion // CHECK: func @generic_op_reshape_producer_fusion
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: f32 // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: f32
// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] // CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]]
// CHECK-SAME: [0], [1], [2, 3] // CHECK-SAME: [0], [1], [2, 3]
// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG1]] // CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG1]]
// CHECK-SAME: [0], [1], [2, 3] // CHECK-SAME: [0], [1], [2, 3]
// CHECK: %[[T3:.+]] = linalg.generic // CHECK: %[[T3:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]], #[[MAP6]]] // CHECK-SAME: indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]], #[[MAP6]]]
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[ARG0]], %[[T1]], %[[ARG2]] : tensor<?x?x4x?xf32>, tensor<?x?x?x4xf32>, f32) // CHECK-SAME: ins(%[[ARG0]], %[[T1]], %[[ARG2]] : tensor<?x?x4x?xf32>, tensor<?x?x?x4xf32>, f32)
// CHECK-SAME: outs(%[[T2]] : tensor<?x?x?x4xf32>) // CHECK-SAME: inits(%[[T2]] : tensor<?x?x?x4xf32>)
// CHECK: %[[T4:.+]] = tensor.collapse_shape %[[T3]] // CHECK: %[[T4:.+]] = tensor.collapse_shape %[[T3]]
// CHECK-SAME: [0], [1], [2, 3] // CHECK-SAME: [0], [1], [2, 3]
// CHECK-SAME: tensor<?x?x?x4xf32> into tensor<?x?x?xf32> // CHECK-SAME: tensor<?x?x?x4xf32> into tensor<?x?x?xf32>
// CHECK: return %[[T4]] // CHECK: return %[[T4]]
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d0, d1)> #map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> ()> #map1 = affine_map<(d0, d1) -> ()>
func.func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>, func.func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,
%arg1 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : f32) -> %arg2 : f32) ->
tensor<?x4x?x5xf32> tensor<?x4x?x5xf32>
{ {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map0, #map0, #map1, #map0], indexing_maps = [#map0, #map0, #map1, #map0],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>, f32) ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>, f32)
outs(%arg0 : tensor<?x?xf32>) { inits(%arg0 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):
%1 = arith.mulf %arg3, %arg4 : f32 %1 = arith.mulf %arg3, %arg4 : f32
%2 = arith.addf %1, %arg5 : f32 %2 = arith.addf %1, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%1 = tensor.expand_shape %0 [[0], [1, 2, 3]] : %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] :
tensor<?x?xf32> into tensor<?x4x?x5xf32> tensor<?x?xf32> into tensor<?x4x?x5xf32>
return %1 : tensor<?x4x?x5xf32> return %1 : tensor<?x4x?x5xf32>
Show All 13 Lines
// CHECK-SAME: [0], [1, 2, 3] // CHECK-SAME: [0], [1, 2, 3]
// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] // CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK-SAME: [0], [1, 2, 3] // CHECK-SAME: [0], [1, 2, 3]
// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x?x5xf32> // CHECK-SAME: tensor<?x?xf32> into tensor<?x4x?x5xf32>
// CHECK: %[[T3:.+]] = linalg.generic // CHECK: %[[T3:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]], #[[MAP2]]]
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[T0]], %[[T1]], %[[ARG2]] : tensor<?x4x?x5xf32>, tensor<?x4x?x5xf32>, f32) // CHECK-SAME: ins(%[[T0]], %[[T1]], %[[ARG2]] : tensor<?x4x?x5xf32>, tensor<?x4x?x5xf32>, f32)
// CHECK-SAME: outs(%[[T2]] : tensor<?x4x?x5xf32>) // CHECK-SAME: inits(%[[T2]] : tensor<?x4x?x5xf32>)
// CHECK: return %[[T3]] : tensor<?x4x?x5xf32> // CHECK: return %[[T3]] : tensor<?x4x?x5xf32>
// ----- // -----
func.func @reshape_as_consumer_permutation func.func @reshape_as_consumer_permutation
(%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>) (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>)
-> tensor<?x2x?x3x4x?xf32> { -> tensor<?x2x?x3x4x?xf32> {
%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<?x?x?xf32>, tensor<?x?xf32>) ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>)
outs(%a : tensor<?x?x?xf32>) { inits(%a : tensor<?x?x?xf32>) {
^bb0(%arg0 : f32, %arg1: f32, %s: f32): ^bb0(%arg0 : f32, %arg1: f32, %s: f32):
%1 = arith.addf %arg0, %arg1 : f32 %1 = arith.addf %arg0, %arg1 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
%d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]] %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]]
: tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32> : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>
return %d : tensor<?x2x?x3x4x?xf32> return %d : tensor<?x2x?x3x4x?xf32>
} }
Show All 11 Lines
// CHECK-SAME: tensor<?x?xf32> into tensor<3x4x?x?xf32> // CHECK-SAME: tensor<?x?xf32> into tensor<3x4x?x?xf32>
// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] // CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK-SAME: [0, 1], [2], [3, 4, 5]] // CHECK-SAME: [0, 1], [2], [3, 4, 5]]
// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32> // CHECK-SAME: tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>
// CHECK: %[[T3:.+]] = linalg.generic // CHECK: %[[T3:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP8]], #[[MAP9]], #[[MAP10]]] // CHECK-SAME: indexing_maps = [#[[MAP8]], #[[MAP9]], #[[MAP10]]]
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<3x4x?x?x2x?xf32>, tensor<3x4x?x?xf32>) // CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<3x4x?x?x2x?xf32>, tensor<3x4x?x?xf32>)
// CHECK-SAME: outs(%[[T2]] : tensor<?x2x?x3x4x?xf32>) // CHECK-SAME: inits(%[[T2]] : tensor<?x2x?x3x4x?xf32>)
// CHECK: return %[[T3]] : tensor<?x2x?x3x4x?xf32> // CHECK: return %[[T3]] : tensor<?x2x?x3x4x?xf32>
// ----- // -----
#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 = tensor.empty() : 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>) { inits(%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>
} }
Show All 9 Lines
// 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]]]
// CHECK-SAME: ["parallel", "parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[T0]], %[[CST]] : // CHECK-SAME: ins(%[[T0]], %[[CST]] :
// CHECK-SAME: outs(%[[T1]] : tensor<8x33x4xf32>) // CHECK-SAME: inits(%[[T1]] : tensor<8x33x4xf32>)
// CHECK: return %[[T2]] : tensor<8x33x4xf32> // CHECK: return %[[T2]] : tensor<8x33x4xf32>
// ----- // -----
#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> #map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)> #map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>
func.func @indexed_consumer_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>, func.func @indexed_consumer_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>,
%arg1 : tensor<?x?x?xi32>) -> %arg1 : tensor<?x?x?xi32>) ->
tensor<?x?x?xi32> tensor<?x?x?xi32>
{ {
%0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]: %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]:
tensor<?x?x4x?xi32> into tensor<?x?x?xi32> tensor<?x?x4x?xi32> into tensor<?x?x?xi32>
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [#map0, #map1, #map1], indexing_maps = [#map0, #map1, #map1],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>) ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>)
outs(%0 : tensor<?x?x?xi32>) { inits(%0 : tensor<?x?x?xi32>) {
^bb0(%arg3: i32, %arg4: i32, %s: i32): ^bb0(%arg3: i32, %arg4: i32, %s: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%idx2 = linalg.index 2 : index %idx2 = linalg.index 2 : index
%1 = arith.muli %arg3, %arg4 : i32 %1 = arith.muli %arg3, %arg4 : i32
%2 = arith.index_cast %idx0 : index to i32 %2 = arith.index_cast %idx0 : index to i32
%3 = arith.addi %1, %2 : i32 %3 = arith.addi %1, %2 : i32
%4 = arith.index_cast %idx1 : index to i32 %4 = arith.index_cast %idx1 : index to i32
Show All 32 Lines
func.func @indexed_producer_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>, func.func @indexed_producer_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>,
%arg1 : tensor<?x?xi32>) -> %arg1 : tensor<?x?xi32>) ->
tensor<?x?x4x5xi32> tensor<?x?x4x5xi32>
{ {
%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<?x?xi32>, tensor<?x?xi32>) ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>)
outs(%arg0 : tensor<?x?xi32>) { inits(%arg0 : tensor<?x?xi32>) {
^bb0(%arg3: i32, %arg4: i32, %s: i32): ^bb0(%arg3: i32, %arg4: i32, %s: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%1 = arith.muli %arg3, %arg4 : i32 %1 = arith.muli %arg3, %arg4 : i32
%2 = arith.index_cast %idx0 : index to i32 %2 = arith.index_cast %idx0 : index to i32
%3 = arith.addi %1, %2 : i32 %3 = arith.addi %1, %2 : i32
%4 = arith.index_cast %idx1 : index to i32 %4 = arith.index_cast %idx1 : index to i32
%5 = arith.addi %3, %4 : i32 %5 = arith.addi %3, %4 : i32
Show All 32 Lines (%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>)
-> tensor<2x3x4x5x6x7xi32> { -> tensor<2x3x4x5x6x7xi32> {
%shape = tensor.empty() : 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>) { inits(%shape : tensor<6x4x210xi32>) {
^bb0(%arg3 : i32, %arg4: i32, %s: i32): ^bb0(%arg3 : i32, %arg4: i32, %s: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%idx2 = linalg.index 2 : index %idx2 = linalg.index 2 : index
%1 = arith.addi %arg3, %arg4 : i32 %1 = arith.addi %arg3, %arg4 : i32
%2 = arith.index_cast %idx0 : index to i32 %2 = arith.index_cast %idx0 : index to i32
%3 = arith.addi %1, %2 : i32 %3 = arith.addi %1, %2 : i32
%4 = arith.index_cast %idx1 : index to i32 %4 = arith.index_cast %idx1 : index to i32
Show All 24 Lines
// 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
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
// CHECK-SAME: ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>) // CHECK-SAME: ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>)
// CHECK-SAME: outs(%[[T3]] : tensor<2x3x4x5x6x7xi32>) // CHECK-SAME: inits(%[[T3]] : tensor<2x3x4x5x6x7xi32>)
// CHECK: ^{{.+}}( // CHECK: ^{{.+}}(
// CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: i32, %[[ARG9:[a-zA-Z0-9_]+]]: i32, // CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: i32, %[[ARG9:[a-zA-Z0-9_]+]]: i32,
// CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: i32) // CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: i32)
// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index // CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index
// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index // CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index
// CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index // CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index
// CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index // CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index
// CHECK-DAG: %[[IDX4:.+]] = linalg.index 4 : index // CHECK-DAG: %[[IDX4:.+]] = linalg.index 4 : index
Show All 16 Lines
{ {
%0 = tensor.collapse_shape %arg0 [[0, 1], [2]] %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]
: tensor<33x8x?xi32> into tensor<264x?xi32> : tensor<33x8x?xi32> into tensor<264x?xi32>
%1 = linalg.generic %1 = linalg.generic
{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, 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(%0 : tensor<264x?xi32>) ins(%0 : tensor<264x?xi32>)
outs(%shape : tensor<264x?x4xi32>) { inits(%shape : tensor<264x?x4xi32>) {
^bb0(%arg1: i32, %s: i32): ^bb0(%arg1: i32, %s: i32):
%idx0 = linalg.index 0 : index %idx0 = linalg.index 0 : index
%idx1 = linalg.index 1 : index %idx1 = linalg.index 1 : index
%idx2 = linalg.index 2 : index %idx2 = linalg.index 2 : index
%2 = arith.index_cast %idx0 : index to i32 %2 = arith.index_cast %idx0 : index to i32
%3 = arith.addi %arg1, %2 : i32 %3 = arith.addi %arg1, %2 : i32
%4 = arith.index_cast %idx1 : index to i32 %4 = arith.index_cast %idx1 : index to i32
%5 = arith.addi %3, %4 : i32 %5 = arith.addi %3, %4 : i32
Show All 39 Lines
func.func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>, func.func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>,
%arg1 : tensor<?x?xf32>) -> %arg1 : tensor<?x?xf32>) ->
tensor<?x?x4x5xf32> tensor<?x?x4x5xf32>
{ {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [#map0, #map0, #map1], indexing_maps = [#map0, #map0, #map1],
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(%arg0 : tensor<?x?xf32>) { inits(%arg0 : tensor<?x?xf32>) {
^bb0(%arg3: f32, %arg4: f32, %s: f32): ^bb0(%arg3: f32, %arg4: f32, %s: f32):
%1 = arith.mulf %arg3, %arg4 : f32 %1 = arith.mulf %arg3, %arg4 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%1 = tensor.expand_shape %0 [[0], [1, 2, 3]] : %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] :
tensor<?x?xf32> into tensor<?x?x4x5xf32> tensor<?x?xf32> into tensor<?x?x4x5xf32>
return %1 : tensor<?x?x4x5xf32> return %1 : tensor<?x?x4x5xf32>
} }
Show All 11 Lines
// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32> // CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32>
// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] // CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]]
// CHECK-SAME: [0], [1, 2, 3] // CHECK-SAME: [0], [1, 2, 3]
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32> // CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32>
// CHECK: %[[T3:.+]] = linalg.generic // CHECK: %[[T3:.+]] = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP4]], #[[MAP5]]] // CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP4]], #[[MAP5]]]
// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] // CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]
// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>) // CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>)
// CHECK-SAME: outs(%[[T2]] : tensor<?x?x4x5xf32>) // CHECK-SAME: inits(%[[T2]] : tensor<?x?x4x5xf32>)
// 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 = tensor.empty(%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>) inits(%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 = tensor.empty() : 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>) { inits(%1 : tensor<2xi64>) {
^bb0(%arg4: i64, %arg5: i64, %arg6: i64): ^bb0(%arg4: i64, %arg5: i64, %arg6: i64):
%3 = arith.addi %arg4, %arg5 : i64 %3 = arith.addi %arg4, %arg5 : i64
linalg.yield %3 : i64 linalg.yield %3 : i64
} -> tensor<2xi64> } -> tensor<2xi64>
return %2 : tensor<2xi64> return %2 : tensor<2xi64>
} }
// CHECK: func @no_fuse_mismatched_dynamism // CHECK: func @no_fuse_mismatched_dynamism
Show All 11 Lines (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>)
-> (tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32>) { -> (tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32>) {
%c:2 = linalg.generic { %c:2 = 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)>,
affine_map<(d0, d1, d2) -> (d2, d0, d1)>], affine_map<(d0, d1, d2) -> (d2, d0, d1)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>) ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>)
outs(%a, %a : tensor<?x?x?xf32>, tensor<?x?x?xf32>) { inits(%a, %a : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%arg0 : f32, %arg1: f32, %s: f32, %t : f32): ^bb0(%arg0 : f32, %arg1: f32, %s: f32, %t : f32):
%1 = arith.addf %arg0, %arg1 : f32 %1 = arith.addf %arg0, %arg1 : f32
linalg.yield %1, %1 : f32, f32 linalg.yield %1, %1 : f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
%d = tensor.expand_shape %c#0 [[0, 1], [2], [3, 4, 5]] %d = tensor.expand_shape %c#0 [[0, 1], [2], [3, 4, 5]]
: tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32> : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>
%e = tensor.expand_shape %c#1 [[0], [1, 2], [3, 4, 5]] %e = tensor.expand_shape %c#1 [[0], [1, 2], [3, 4, 5]]
: tensor<?x?x?xf32> into tensor<?x?x2x3x4x?xf32> : tensor<?x?x?xf32> into tensor<?x?x2x3x4x?xf32>
return %d, %e : tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32> return %d, %e : tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>
// 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, d2, d3, d4, d5) -> (d5, d0, d1, d2, d3, d4)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5, d0, d1, d2, d3, d4)>
// CHECK: func @reshape_as_consumer_permutation_with_multiple_results // CHECK: func @reshape_as_consumer_permutation_with_multiple_results
// 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?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>
// CHECK-DAG: %[[RESHAPE0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3, 4], [5]{{\]}} // CHECK-DAG: %[[RESHAPE0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3, 4], [5]{{\]}}
// CHECK-DAG: %[[RESHAPE1:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1, 2], [3]{{\]}} // CHECK-DAG: %[[RESHAPE1:.+]] = tensor.expand_shape %[[ARG1]] {{\[}}[0, 1, 2], [3]{{\]}}
// CHECK-DAG: %[[RESHAPE2:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3, 4, 5]{{\]}} // CHECK-DAG: %[[RESHAPE2:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2], [3, 4, 5]{{\]}}
// CHECK-DAG: %[[RESHAPE3:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3, 4, 5]{{\]}} // CHECK-DAG: %[[RESHAPE3:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3, 4, 5]{{\]}}
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]]
// CHECK-SAME: ins(%[[RESHAPE0]], %[[RESHAPE1]] : // CHECK-SAME: ins(%[[RESHAPE0]], %[[RESHAPE1]] :
// CHECK-SAME: outs(%[[RESHAPE2]], %[[RESHAPE3]] : // CHECK-SAME: inits(%[[RESHAPE2]], %[[RESHAPE3]] :
// CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1
// ----- // -----
#map0 = affine_map<(d0, d1) -> (d1)> #map0 = affine_map<(d0, d1) -> (d1)>
#map1 = affine_map<(d0, d1) -> (d0, d1)> #map1 = affine_map<(d0, d1) -> (d0, d1)>
module { module {
func.func @multi_result_op_expansion(%arg0: tensor<512xf32>, %arg1: tensor<512xf32>, func.func @multi_result_op_expansion(%arg0: tensor<512xf32>, %arg1: tensor<512xf32>,
%arg2: tensor<512xf32>, %arg3: tensor<200x512xf32>) -> tensor<25x8x1x512xf32> { %arg2: tensor<512xf32>, %arg3: tensor<200x512xf32>) -> tensor<25x8x1x512xf32> {
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [#map0, #map0, #map0, #map1], indexing_maps = [#map0, #map0, #map0, #map1],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1 : tensor<512xf32>, tensor<512xf32>) ins(%arg0, %arg1 : tensor<512xf32>, tensor<512xf32>)
outs(%arg2, %arg3 : tensor<512xf32>, tensor<200x512xf32>) { inits(%arg2, %arg3 : tensor<512xf32>, tensor<200x512xf32>) {
^bb0(%arg4: f32, %arg5: f32, %arg6: f32, %arg7: f32): ^bb0(%arg4: f32, %arg5: f32, %arg6: f32, %arg7: f32):
%2 = arith.addf %arg4, %arg5 : f32 %2 = arith.addf %arg4, %arg5 : f32
linalg.yield %2, %2 : f32, f32 linalg.yield %2, %2 : f32, f32
} -> (tensor<512xf32>, tensor<200x512xf32>) } -> (tensor<512xf32>, tensor<200x512xf32>)
%1 = tensor.expand_shape %0#1 [[0, 1, 2], [3]] : tensor<200x512xf32> into tensor<25x8x1x512xf32> %1 = tensor.expand_shape %0#1 [[0, 1, 2], [3]] : tensor<200x512xf32> into tensor<25x8x1x512xf32>
return %1 : tensor<25x8x1x512xf32> return %1 : tensor<25x8x1x512xf32>
} }
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (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: func.func @multi_result_op_expansion( // CHECK: func.func @multi_result_op_expansion(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<512xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<512xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<512xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<512xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<512xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<512xf32>
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: tensor<200x512xf32> // CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: tensor<200x512xf32>
// CHECK: %[[OUTS:.+]] = tensor.expand_shape %[[ARG3]] {{\[}}[0, 1, 2], [3]{{\]}} // CHECK: %[[OUTS:.+]] = tensor.expand_shape %[[ARG3]] {{\[}}[0, 1, 2], [3]{{\]}}
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic // CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]], #[[MAP1]]] // CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]], #[[MAP1]]]
// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : // CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :
// CHECK-SAME: outs(%[[ARG2]], %[[OUTS]] : // CHECK-SAME: inits(%[[ARG2]], %[[OUTS]] :
// CHECK: return %[[GENERIC]]#1 // CHECK: return %[[GENERIC]]#1

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

Show First 20 Lines • Show All 48 Lines • ▼ Show 20 Linesfunc.func @remove_dim_result_uses
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = linalg.generic %0 = linalg.generic
{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d2, d1)>,
affine_map<(d0, d1, d2) -> (d0 + d1, d1 - d0)>], affine_map<(d0, d1, d2) -> (d0 + d1, d1 - d0)>],
iterator_types = ["parallel", "parallel", "reduction"]} iterator_types = ["parallel", "parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) { inits(%arg2 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
%1 = arith.mulf %arg3, %arg4 : f32 %1 = arith.mulf %arg3, %arg4 : f32
%2 = arith.addf %1, %arg5 : f32 %2 = arith.addf %1, %arg5 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%3 = tensor.dim %0, %c0 : tensor<?x?xf32> %3 = tensor.dim %0, %c0 : tensor<?x?xf32>
%4 = tensor.dim %0, %c1 : tensor<?x?xf32> %4 = tensor.dim %0, %c1 : tensor<?x?xf32>
return %3, %4 : index, index return %3, %4 : index, index
Show All 21 Linesfunc.func @remove_dim_result_uses_outs
%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 = tensor.empty(%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>) inits(%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, %c1 : tensor<?x?xf32> %2 = tensor.dim %1, %c1 : tensor<?x?xf32>
return %2 : index return %2 : index
} }
// CHECK: func @remove_dim_result_uses_outs // CHECK: func @remove_dim_result_uses_outs
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index // CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK: return %[[ARG1]] // CHECK: return %[[ARG1]]
// ----- // -----
func.func @remove_dim_result_uses_sequence func.func @remove_dim_result_uses_sequence
(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, (%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?x?xf32>) -> (index, index, index, index) { %arg2 : tensor<?x?xf32>) -> (index, 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.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg2 : tensor<?x?xf32>) -> 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>
%3 = linalg.generic %3 = linalg.generic
{indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0)>, {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0)>,
affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>,
affine_map<(d0, d1, d2) -> (d0, d2)>], affine_map<(d0, d1, d2) -> (d0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"]} iterator_types = ["parallel", "reduction", "parallel"]}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%0 : tensor<?x?xf32>) { inits(%0 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
%4 = arith.mulf %arg3, %arg4 : f32 %4 = arith.mulf %arg3, %arg4 : f32
%5 = arith.addf %4, %arg5 : f32 %5 = arith.addf %4, %arg5 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
%6 = tensor.dim %3, %c0 : tensor<?x?xf32> %6 = tensor.dim %3, %c0 : tensor<?x?xf32>
%7 = tensor.dim %3, %c1 : tensor<?x?xf32> %7 = tensor.dim %3, %c1 : tensor<?x?xf32>
return %1, %2, %6, %7 : index, index, index, index return %1, %2, %6, %7 : index, index, index, index
Show All 17 Linesfunc.func @keep_result_dim_uses_sequence2
%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 = tensor.empty(%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>) inits(%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 @empty_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>) { inits(%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>
▲ Show 20 Lines • Show All 88 Lines • Show Last 20 Lines

mlir/test/Dialect/Linalg/roundtrip.mlir

Show All 26 Lines
// ----- // -----
func.func @ops(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @ops(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%arg1: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>,
%arg2: memref<?xf32, strided<[1], offset: ?>>, %arg2: memref<?xf32, strided<[1], offset: ?>>,
%arg3: memref<f32>) { %arg3: memref<f32>) {
linalg.matmul ins(%arg0, %arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>, linalg.matmul ins(%arg0, %arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>) memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>) inits(%arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>)
linalg.matvec ins(%arg0, %arg1: memref<?x?xf32, strided<[?, 1], offset: ?>>, linalg.matvec ins(%arg0, %arg1: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%arg2: memref<?xf32, strided<[1], offset: ?>>) inits(%arg2: memref<?xf32, strided<[1], offset: ?>>)
linalg.dot ins(%arg1, %arg2: memref<?xf32, strided<[1], offset: ?>>, linalg.dot ins(%arg1, %arg2: memref<?xf32, strided<[1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%arg3: memref<f32>) inits(%arg3: memref<f32>)
return return
} }
// CHECK-LABEL: func @ops(% // CHECK-LABEL: func @ops(%
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: ins(%{{.*}}, %{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>, // CHECK-SAME: ins(%{{.*}}, %{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>,
// CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK-SAME: outs(%{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK-SAME: inits(%{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK: linalg.matvec // CHECK: linalg.matvec
// CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, // CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>,
// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>) // CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>)
// CHECK-SAME: outs(%{{.*}}: memref<?xf32, strided<[1], offset: ?>>) // CHECK-SAME: inits(%{{.*}}: memref<?xf32, strided<[1], offset: ?>>)
// CHECK: linalg.dot // CHECK: linalg.dot
// CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, // CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>,
// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>) // CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>)
// CHECK-SAME: outs(%{{.*}}: memref<f32>) // CHECK-SAME: inits(%{{.*}}: memref<f32>)
// ----- // -----
func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) { func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) {
linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?xf32, strided<[1], offset: ?>>) linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<?xf32, strided<[1], offset: ?>>)
return return
} }
// CHECK-LABEL: func @fill_view( // CHECK-LABEL: func @fill_view(
// CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) { // CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) {
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : memref<?xf32, strided<[1], offset: ?>>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%{{.*}} : memref<?xf32, strided<[1], offset: ?>>)
// ----- // -----
func.func @memref_transpose(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { func.func @memref_transpose(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
%0 = memref.transpose %arg0 (i, j, k) -> (k, j, i) : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>> %0 = memref.transpose %arg0 (i, j, k) -> (k, j, i) : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>>
return return
} }
// CHECK-LABEL: func @memref_transpose // CHECK-LABEL: func @memref_transpose
// CHECK: memref.transpose %{{.*}} ([[i:.*]], [[j:.*]], [[k:.*]]) -> ([[k]], [[j]], [[i]]) : // CHECK: memref.transpose %{{.*}} ([[i:.*]], [[j:.*]], [[k:.*]]) -> ([[k]], [[j]], [[i]]) :
// CHECK-SAME: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>> // CHECK-SAME: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>>
// ----- // -----
func.func @fill_view3(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: f32) { func.func @fill_view3(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: f32) {
linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) linalg.fill ins(%arg1 : f32) inits(%arg0 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
return return
} }
// CHECK-LABEL: func @fill_view3( // CHECK-LABEL: func @fill_view3(
// CHECK: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) { // CHECK: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) {
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%{{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
// ----- // -----
#accesses_0 = [ #accesses_0 = [
affine_map<(i, j, k) -> (j, i)>, affine_map<(i, j, k) -> (j, i)>,
affine_map<(i, j, k) -> ()>, affine_map<(i, j, k) -> ()>,
affine_map<(i, j, k) -> (i, k, i + j)> affine_map<(i, j, k) -> (i, k, i + j)>
] ]
#trait_0 = { #trait_0 = {
indexing_maps = #accesses_0, indexing_maps = #accesses_0,
iterator_types = ["parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel"],
library_call = "some_external_function_name_1" library_call = "some_external_function_name_1"
} }
func.func @generic(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, func.func @generic(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>,
%arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
linalg.generic #trait_0 linalg.generic #trait_0
ins(%arg0, %cst : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32) ins(%arg0, %cst : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32)
outs(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) inits(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
attrs = {foo = 1} { attrs = {foo = 1} {
^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) : ^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) :
linalg.yield %1 : f32 linalg.yield %1 : f32
} }
return return
} }
// CHECK-LABEL: func @generic // CHECK-LABEL: func @generic
// CHECK: linalg.generic { // CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}, #{{[0-9a-z]*}}], // CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}, #{{[0-9a-z]*}}],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"], // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"],
// CHECK-SAME: library_call = "some_external_function_name_1"} // CHECK-SAME: library_call = "some_external_function_name_1"}
// CHECK-SAME: ins({{.*}}, {{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32) // CHECK-SAME: ins({{.*}}, {{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32)
// CHECK-SAME: outs({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) // CHECK-SAME: inits({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
// CHECK-SAME: {foo = 1 : i64} // CHECK-SAME: {foo = 1 : i64}
// ----- // -----
#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> #map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
func.func @generic_without_inputs(%arg0 : memref<?x?x?xf32>) { func.func @generic_without_inputs(%arg0 : memref<?x?x?xf32>) {
linalg.generic {indexing_maps = [#map0], linalg.generic {indexing_maps = [#map0],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
outs(%arg0 : memref<?x?x?xf32>) { inits(%arg0 : memref<?x?x?xf32>) {
^bb0(%arg3: f32): ^bb0(%arg3: f32):
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} }
return return
} }
// CHECK-LABEL: func @generic_without_inputs // CHECK-LABEL: func @generic_without_inputs
Show All 14 Lines#trait_1 = {
library_call = "some_external_function_name_1" library_call = "some_external_function_name_1"
} }
func.func @generic_with_tensor_input_and_output( func.func @generic_with_tensor_input_and_output(
%arg0: tensor<?x?xvector<3x4xi4>>, %arg1: tensor<?x?x?xf32>) %arg0: tensor<?x?xvector<3x4xi4>>, %arg1: tensor<?x?x?xf32>)
-> (tensor<?x?x?xf32>) { -> (tensor<?x?x?xf32>) {
%0 = linalg.generic #trait_1 %0 = linalg.generic #trait_1
ins(%arg0, %arg1 : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>)
outs(%arg1 : tensor<?x?x?xf32>) inits(%arg1 : tensor<?x?x?xf32>)
attrs = {foo = 1} { attrs = {foo = 1} {
^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) : ^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) :
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
linalg.yield %f0 : f32 linalg.yield %f0 : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32> return %0 : tensor<?x?x?xf32>
} }
// CHECK-LABEL: func @generic_with_tensor_input_and_output // CHECK-LABEL: func @generic_with_tensor_input_and_output
// CHECK: linalg.generic { // CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"], // CHECK-SAME: indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"],
// CHECK-SAME: library_call = "some_external_function_name_1"} // CHECK-SAME: library_call = "some_external_function_name_1"}
// CHECK-SAME: ins({{.*}} : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>) // CHECK-SAME: ins({{.*}} : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>)
// CHECK-SAME: outs({{.*}} : tensor<?x?x?xf32>) // CHECK-SAME: inits({{.*}} : tensor<?x?x?xf32>)
// CHECK-SAME: {foo = 1 : i64} // CHECK-SAME: {foo = 1 : i64}
// CHECK: -> tensor<?x?x?xf32> // CHECK: -> tensor<?x?x?xf32>
// 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 = tensor.empty() : 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) inits(%0 : tensor<i32>) -> tensor<i32>
%2 = tensor.empty() : 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) inits(%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>) { inits(%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
%6 = arith.select %5, %arg3, %arg5 : i32 %6 = arith.select %5, %arg3, %arg5 : i32
%7 = arith.cmpi eq, %arg3, %arg5 : i32 %7 = arith.cmpi eq, %arg3, %arg5 : i32
%8 = arith.cmpi slt, %arg4, %arg6 : i32 %8 = arith.cmpi slt, %arg4, %arg6 : i32
%9 = arith.select %8, %arg4, %arg6 : i32 %9 = arith.select %8, %arg4, %arg6 : i32
%10 = arith.select %5, %arg4, %arg6 : i32 %10 = arith.select %5, %arg4, %arg6 : i32
%11 = arith.select %7, %9, %10 : i32 %11 = arith.select %7, %9, %10 : i32
linalg.yield %6, %11 : i32, i32 linalg.yield %6, %11 : i32, i32
} -> (tensor<i32>, tensor<i32>) } -> (tensor<i32>, tensor<i32>)
return %4#0, %4#1 : tensor<i32>, tensor<i32> return %4#0, %4#1 : tensor<i32>, tensor<i32>
} }
// CHECK-LABEL: func @generic_with_multiple_tensor_outputs // CHECK-LABEL: func @generic_with_multiple_tensor_outputs
// CHECK: %{{.*}} = linalg.generic { // CHECK: %{{.*}} = linalg.generic {
// CHECK-SAME: ins({{.*}} : tensor<?xi32>, tensor<?xi32>) // CHECK-SAME: ins({{.*}} : tensor<?xi32>, tensor<?xi32>)
// CHECK-SAME: outs({{.*}} : tensor<i32>, tensor<i32>) // CHECK-SAME: inits({{.*}} : tensor<i32>, tensor<i32>)
// CHECK: } -> (tensor<i32>, tensor<i32>) // CHECK: } -> (tensor<i32>, tensor<i32>)
// ----- // -----
#broadcast_access = [ #broadcast_access = [
affine_map<(i, j) -> ()>, affine_map<(i, j) -> ()>,
affine_map<(i, j) -> (i, j)> affine_map<(i, j) -> (i, j)>
] ]
#trait_broadcast = { #trait_broadcast = {
indexing_maps = #broadcast_access, indexing_maps = #broadcast_access,
iterator_types = ["parallel", "parallel"], iterator_types = ["parallel", "parallel"],
library_call = "some_broadcast_external_fn" library_call = "some_broadcast_external_fn"
} }
func.func @generic_op_zero_rank(%arg0: tensor<f32>, %arg1 : tensor<3x4xf32>) -> (tensor<3x4xf32>) func.func @generic_op_zero_rank(%arg0: tensor<f32>, %arg1 : tensor<3x4xf32>) -> (tensor<3x4xf32>)
{ {
%0 = linalg.generic #trait_broadcast %0 = linalg.generic #trait_broadcast
ins(%arg0 : tensor<f32>) ins(%arg0 : tensor<f32>)
outs(%arg1 : tensor<3x4xf32>) { inits(%arg1 : tensor<3x4xf32>) {
^bb(%a: f32, %b: f32) : ^bb(%a: f32, %b: f32) :
linalg.yield %a : f32 linalg.yield %a : f32
} -> tensor<3x4xf32> } -> tensor<3x4xf32>
return %0 : tensor<3x4xf32> return %0 : tensor<3x4xf32>
} }
// ----- // -----
#accesses_3 = [ #accesses_3 = [
affine_map<(i, j, k) -> (j, i)>, affine_map<(i, j, k) -> (j, i)>,
affine_map<(i, j, k) -> (i, k, i + j)> affine_map<(i, j, k) -> (i, k, i + j)>
] ]
#trait_3 = { #trait_3 = {
indexing_maps = #accesses_3, indexing_maps = #accesses_3,
iterator_types = ["parallel", "parallel", "parallel"], iterator_types = ["parallel", "parallel", "parallel"],
library_call = "some_external_function_name_2" library_call = "some_external_function_name_2"
} }
func.func @generic_region(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, func.func @generic_region(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>,
%arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) { %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {
linalg.generic #trait_3 linalg.generic #trait_3
ins(%arg0 : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>) ins(%arg0 : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>)
outs(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) inits(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
attrs = {foo = 1} { attrs = {foo = 1} {
^bb(%a: vector<3x4xi4>, %b: f32) : ^bb(%a: vector<3x4xi4>, %b: f32) :
%0 = linalg.index 0 : index %0 = linalg.index 0 : index
%1 = linalg.index 1 : index %1 = linalg.index 1 : index
%2 = linalg.index 2 : index %2 = linalg.index 2 : index
linalg.yield %b : f32 linalg.yield %b : f32
} }
return return
} }
// CHECK-LABEL: func @generic_region // CHECK-LABEL: func @generic_region
// CHECK: linalg.generic { // CHECK: linalg.generic {
// CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}], // CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"], // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"],
// CHECK-SAME: library_call = "some_external_function_name_2" // CHECK-SAME: library_call = "some_external_function_name_2"
// CHECK-SAME: ins({{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>) // CHECK-SAME: ins({{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>)
// CHECK-SAME: outs({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) // CHECK-SAME: inits({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)
// CHECK-SAME: attrs = {foo = 1 : i64} { // CHECK-SAME: attrs = {foo = 1 : i64} {
// CHECK: ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32): // CHECK: ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32):
// CHECK: %{{.*}} = linalg.index 0 : index // CHECK: %{{.*}} = linalg.index 0 : index
// CHECK: %{{.*}} = linalg.index 1 : index // CHECK: %{{.*}} = linalg.index 1 : index
// CHECK: %{{.*}} = linalg.index 2 : index // CHECK: %{{.*}} = linalg.index 2 : index
// CHECK: linalg.yield %{{.*}} : f32 // CHECK: linalg.yield %{{.*}} : f32
// ----- // -----
func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>, func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,
%ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>) %ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>)
-> (tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>)
{ {
linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>) linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs(%c3: memref<?x?x?xf32>) inits(%c3: memref<?x?x?xf32>)
%res1 = linalg.batch_matmul %res1 = linalg.batch_matmul
ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>) ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>)
outs(%tc3: tensor<?x?x?xf32>) inits(%tc3: tensor<?x?x?xf32>)
-> tensor<?x?x?xf32> -> tensor<?x?x?xf32>
return %res1 : tensor<?x?x?xf32> return %res1 : tensor<?x?x?xf32>
} }
// 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
// ----- // -----
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 = tensor.empty(%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) inits(%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) inits(%{{.+}} : 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>,
%arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?x?xf32>, %arg3 : tensor<?x?xf32>) -> %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?x?xf32>, %arg3 : tensor<?x?xf32>) ->
(tensor<?x?x?xf32>, tensor<?x?xf32>) { (tensor<?x?x?xf32>, tensor<?x?xf32>) {
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction"]} iterator_types = ["parallel", "parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?xf32>)
outs(%arg2, %arg3 : tensor<?x?x?xf32>, tensor<?x?xf32>) { inits(%arg2, %arg3 : tensor<?x?x?xf32>, tensor<?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32):
%1 = arith.mulf %b0, %b1 : f32 %1 = arith.mulf %b0, %b1 : f32
%2 = arith.addf %1, %b3 : f32 %2 = arith.addf %1, %b3 : f32
linalg.yield %1, %2 : f32, f32 linalg.yield %1, %2 : f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?xf32>) } -> (tensor<?x?x?xf32>, tensor<?x?xf32>)
return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?xf32> return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?xf32>
} }
// CHECK-LABEL: func @mixed_parallel_reduced_results // CHECK-LABEL: func @mixed_parallel_reduced_results
// CHECK: linalg.generic // CHECK: linalg.generic
// ----- // -----
func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>, func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
%init: tensor<64xf32>) -> tensor<64xf32> { %init: tensor<64xf32>) -> tensor<64xf32> {
%add = linalg.map %add = linalg.map
ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>) ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %add : tensor<64xf32> func.return %add : tensor<64xf32>
} }
// CHECK-LABEL: func @map_binary // CHECK-LABEL: func @map_binary
// CHECK: linalg.map // CHECK: linalg.map
// ----- // -----
func.func @map_binary_memref(%lhs: memref<64xf32>, %rhs: memref<64xf32>, func.func @map_binary_memref(%lhs: memref<64xf32>, %rhs: memref<64xf32>,
%init: memref<64xf32>) { %init: memref<64xf32>) {
linalg.map linalg.map
ins(%lhs, %rhs: memref<64xf32>, memref<64xf32>) ins(%lhs, %rhs: memref<64xf32>, memref<64xf32>)
outs(%init:memref<64xf32>) inits(%init:memref<64xf32>)
(%lhs_elem: f32, %rhs_elem: f32) { (%lhs_elem: f32, %rhs_elem: f32) {
%0 = arith.addf %lhs_elem, %rhs_elem: f32 %0 = arith.addf %lhs_elem, %rhs_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return func.return
} }
// CHECK-LABEL: func @map_binary_memref // CHECK-LABEL: func @map_binary_memref
// CHECK: linalg.map // CHECK: linalg.map
// ----- // -----
func.func @map_unary(%input: tensor<64xf32>, %init: tensor<64xf32>) -> tensor<64xf32> { func.func @map_unary(%input: tensor<64xf32>, %init: tensor<64xf32>) -> tensor<64xf32> {
%abs = linalg.map %abs = linalg.map
ins(%input:tensor<64xf32>) ins(%input:tensor<64xf32>)
outs(%init:tensor<64xf32>) inits(%init:tensor<64xf32>)
(%input_elem: f32) { (%input_elem: f32) {
%0 = math.absf %input_elem: f32 %0 = math.absf %input_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %abs : tensor<64xf32> func.return %abs : tensor<64xf32>
} }
// CHECK-LABEL: func @map_unary // CHECK-LABEL: func @map_unary
// CHECK: linalg.map // CHECK: linalg.map
// ----- // -----
func.func @map_unary_memref(%input: memref<64xf32>, %init: memref<64xf32>) { func.func @map_unary_memref(%input: memref<64xf32>, %init: memref<64xf32>) {
linalg.map linalg.map
ins(%input:memref<64xf32>) ins(%input:memref<64xf32>)
outs(%init:memref<64xf32>) inits(%init:memref<64xf32>)
(%input_elem: f32) { (%input_elem: f32) {
%0 = math.absf %input_elem: f32 %0 = math.absf %input_elem: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return func.return
} }
// CHECK-LABEL: func @map_unary_memref // CHECK-LABEL: func @map_unary_memref
// CHECK: linalg.map // CHECK: linalg.map
// ----- // -----
func.func @reduce(%input: tensor<16x32x64xf32>, func.func @reduce(%input: tensor<16x32x64xf32>,
%init: tensor<16x64xf32>) -> tensor<16x64xf32> { %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
%reduce = linalg.reduce %reduce = linalg.reduce
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<16x64xf32>) inits(%init:tensor<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return %reduce : tensor<16x64xf32> func.return %reduce : tensor<16x64xf32>
} }
// CHECK-LABEL: func @reduce // CHECK-LABEL: func @reduce
// CHECK: linalg.reduce // CHECK: linalg.reduce
// ----- // -----
func.func @reduce_memref(%input: memref<16x32x64xf32>, func.func @reduce_memref(%input: memref<16x32x64xf32>,
%init: memref<16x64xf32>) { %init: memref<16x64xf32>) {
linalg.reduce linalg.reduce
ins(%input:memref<16x32x64xf32>) ins(%input:memref<16x32x64xf32>)
outs(%init:memref<16x64xf32>) inits(%init:memref<16x64xf32>)
dimensions = [1] dimensions = [1]
(%in: f32, %out: f32) { (%in: f32, %out: f32) {
%0 = arith.addf %in, %out: f32 %0 = arith.addf %in, %out: f32
linalg.yield %0: f32 linalg.yield %0: f32
} }
func.return func.return
} }
// CHECK-LABEL: func @reduce_memref // CHECK-LABEL: func @reduce_memref
// CHECK: linalg.reduce // CHECK: linalg.reduce
// ----- // -----
func.func @variadic_reduce(%input1: tensor<16x32x64xf32>, func.func @variadic_reduce(%input1: tensor<16x32x64xf32>,
%init1: tensor<16x64xf32>, %input2: tensor<16x32x64xi64>, %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xi64>,
%init2: tensor<16x64xi64>) -> (tensor<16x64xf32>, tensor<16x64xi64>) { %init2: tensor<16x64xi64>) -> (tensor<16x64xf32>, tensor<16x64xi64>) {
%reduce, %reduce2 = linalg.reduce %reduce, %reduce2 = linalg.reduce
ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xi64>) ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xi64>)
outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xi64>) inits(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xi64>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: i64, %out1: f32, %out2: i64) { (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
%1 = arith.addi %in2, %out2: i64 %1 = arith.addi %in2, %out2: i64
linalg.yield %0, %1: f32, i64 linalg.yield %0, %1: f32, i64
} }
func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xi64> func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xi64>
} }
// CHECK-LABEL: func @variadic_reduce // CHECK-LABEL: func @variadic_reduce
// CHECK: linalg.reduce // CHECK: linalg.reduce
// ----- // -----
func.func @variadic_reduce_memref(%input1: memref<16x32x64xf32>, func.func @variadic_reduce_memref(%input1: memref<16x32x64xf32>,
%init1: memref<16x64xf32>, %input2: memref<16x32x64xi64>, %init1: memref<16x64xf32>, %input2: memref<16x32x64xi64>,
%init2: memref<16x64xi64>) { %init2: memref<16x64xi64>) {
linalg.reduce linalg.reduce
ins(%input1, %input2 : memref<16x32x64xf32>, memref<16x32x64xi64>) ins(%input1, %input2 : memref<16x32x64xf32>, memref<16x32x64xi64>)
outs(%init1, %init2 : memref<16x64xf32>, memref<16x64xi64>) inits(%init1, %init2 : memref<16x64xf32>, memref<16x64xi64>)
dimensions = [1] dimensions = [1]
(%in1: f32, %in2: i64, %out1: f32, %out2: i64) { (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {
%0 = arith.addf %in1, %out1: f32 %0 = arith.addf %in1, %out1: f32
%1 = arith.addi %in2, %out2: i64 %1 = arith.addi %in2, %out2: i64
linalg.yield %0, %1: f32, i64 linalg.yield %0, %1: f32, i64
} }
func.return func.return
} }
// CHECK-LABEL: func @variadic_reduce_memref // CHECK-LABEL: func @variadic_reduce_memref
// CHECK: linalg.reduce // CHECK: linalg.reduce
// ----- // -----
func.func @transpose(%input: tensor<16x32x64xf32>, func.func @transpose(%input: tensor<16x32x64xf32>,
%init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> { %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
%transpose = linalg.transpose %transpose = linalg.transpose
ins(%input:tensor<16x32x64xf32>) ins(%input:tensor<16x32x64xf32>)
outs(%init:tensor<32x64x16xf32>) inits(%init:tensor<32x64x16xf32>)
permutation = [1, 2, 0] permutation = [1, 2, 0]
func.return %transpose : tensor<32x64x16xf32> func.return %transpose : tensor<32x64x16xf32>
} }
// CHECK-LABEL: func @transpose // CHECK-LABEL: func @transpose
// ----- // -----
func.func @transpose_memref(%input: memref<16x32x64xf32>, func.func @transpose_memref(%input: memref<16x32x64xf32>,
%init: memref<32x64x16xf32>) { %init: memref<32x64x16xf32>) {
linalg.transpose linalg.transpose
ins(%input:memref<16x32x64xf32>) ins(%input:memref<16x32x64xf32>)
outs(%init:memref<32x64x16xf32>) inits(%init:memref<32x64x16xf32>)
permutation = [1, 2, 0] permutation = [1, 2, 0]
func.return func.return
} }
// CHECK-LABEL: func @transpose_memref // CHECK-LABEL: func @transpose_memref

mlir/test/Dialect/Linalg/standard.mlir

// RUN: mlir-opt %s -convert-linalg-to-std | FileCheck %s // RUN: mlir-opt %s -convert-linalg-to-std | FileCheck %s
func.func @dot(%arg0: memref<?xf32, strided<[1], offset: ?>>, func.func @dot(%arg0: memref<?xf32, strided<[1], offset: ?>>,
%arg1: memref<?xf32, strided<[1], offset: ?>>, %arg1: memref<?xf32, strided<[1], offset: ?>>,
%arg2: memref<f32>) { %arg2: memref<f32>) {
linalg.dot ins(%arg0, %arg1: memref<?xf32, strided<[1], offset: ?>>, linalg.dot ins(%arg0, %arg1: memref<?xf32, strided<[1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%arg2: memref<f32>) inits(%arg2: memref<f32>)
return return
} }
// CHECK-LABEL: func @dot( // CHECK-LABEL: func @dot(
// CHECK-SAME: %[[arg0:[a-zA-z0-9]*]]: memref<?xf32, strided<[1], offset: ?>>, // CHECK-SAME: %[[arg0:[a-zA-z0-9]*]]: memref<?xf32, strided<[1], offset: ?>>,
// CHECK-SAME: %[[arg1:[a-zA-z0-9]*]]: memref<?xf32, strided<[1], offset: ?>>, // CHECK-SAME: %[[arg1:[a-zA-z0-9]*]]: memref<?xf32, strided<[1], offset: ?>>,
// CHECK-SAME: %[[arg2:[a-zA-z0-9]*]]: memref<f32>) { // CHECK-SAME: %[[arg2:[a-zA-z0-9]*]]: memref<f32>) {
// CHECK: %[[o0:.*]] = memref.cast %[[arg0]] : // CHECK: %[[o0:.*]] = memref.cast %[[arg0]] :
// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[?], offset: ?>> // CHECK-SAME: memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
Show All 22 Lines
!matrix_type_A = memref<?x?x!vector_type_A> !matrix_type_A = memref<?x?x!vector_type_A>
!matrix_type_B = memref<?x?x!vector_type_B> !matrix_type_B = memref<?x?x!vector_type_B>
!matrix_type_C = memref<?x?x!vector_type_C> !matrix_type_C = memref<?x?x!vector_type_C>
func.func @matmul_vec_impl(%A: !matrix_type_A, %B: !matrix_type_B, %C: !matrix_type_C) { func.func @matmul_vec_impl(%A: !matrix_type_A, %B: !matrix_type_B, %C: !matrix_type_C) {
linalg.generic #matmul_trait linalg.generic #matmul_trait
ins(%A, %B : !matrix_type_A, !matrix_type_B) ins(%A, %B : !matrix_type_A, !matrix_type_B)
outs(%C : !matrix_type_C) { inits(%C : !matrix_type_C) {
^bb0(%a: !vector_type_A, %b: !vector_type_B, %c: !vector_type_C): ^bb0(%a: !vector_type_A, %b: !vector_type_B, %c: !vector_type_C):
%d = vector.outerproduct %a, %b, %c: !vector_type_A, !vector_type_B %d = vector.outerproduct %a, %b, %c: !vector_type_A, !vector_type_B
linalg.yield %d: !vector_type_C linalg.yield %d: !vector_type_C
} }
return return
} }
// CHECK-LABEL: func @matmul_vec_impl( // CHECK-LABEL: func @matmul_vec_impl(
// CHECK: call @external_outerproduct_matmul(%{{.*}}) : // CHECK: call @external_outerproduct_matmul(%{{.*}}) :

mlir/test/Dialect/Linalg/swap-extract-slice-with-fill.mlir

//RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-swap-extract-slice-with-fill-pattern %s | FileCheck %s //RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-swap-extract-slice-with-fill-pattern %s | FileCheck %s
// CHECK-LABEL: func.func @swap_fill_insert_slice // CHECK-LABEL: func.func @swap_fill_insert_slice
// CHECK-SAME: (%[[INIT:.+]]: tensor<?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[SIZE1:.+]]: index) // CHECK-SAME: (%[[INIT:.+]]: tensor<?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[SIZE1:.+]]: index)
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[EXT:.+]] = tensor.extract_slice %[[INIT]][%[[OFFSET0]], 8, 4] [1, %[[SIZE1]], 6] [1, 3, 1] // CHECK: %[[EXT:.+]] = tensor.extract_slice %[[INIT]][%[[OFFSET0]], 8, 4] [1, %[[SIZE1]], 6] [1, 3, 1]
// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[EXT]] : tensor<?x6xf32>) -> tensor<?x6xf32> // CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) inits(%[[EXT]] : tensor<?x6xf32>) -> tensor<?x6xf32>
// CHECK: return %[[FILL]] // CHECK: return %[[FILL]]
func.func @swap_fill_insert_slice(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> tensor<?x6xf32> { func.func @swap_fill_insert_slice(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> tensor<?x6xf32> {
%f0 = arith.constant 0.000000e+00 : f32 %f0 = arith.constant 0.000000e+00 : f32
%0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> %0 = linalg.fill ins(%f0 : f32) inits(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%1 = tensor.extract_slice %0[%offset0, 8, 4] [1, %size1, 6] [1, 3, 1] %1 = tensor.extract_slice %0[%offset0, 8, 4] [1, %size1, 6] [1, 3, 1]
: tensor<?x?x?xf32> to tensor<?x6xf32> : tensor<?x?x?xf32> to tensor<?x6xf32>
return %1: tensor<?x6xf32> return %1: tensor<?x6xf32>
} }
// ----- // -----
// CHECK-LABEL: func.func @dont_swap_fill_insert_slice_multi_user // CHECK-LABEL: func.func @dont_swap_fill_insert_slice_multi_user
// CHECK: linalg.fill // CHECK: linalg.fill
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
func.func @dont_swap_fill_insert_slice_multi_user(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> (tensor<?x?x?xf32>, tensor<2x?x6xf32>) { func.func @dont_swap_fill_insert_slice_multi_user(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> (tensor<?x?x?xf32>, tensor<2x?x6xf32>) {
%f0 = arith.constant 0.000000e+00 : f32 %f0 = arith.constant 0.000000e+00 : f32
%0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> %0 = linalg.fill ins(%f0 : f32) inits(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
%1 = tensor.extract_slice %0[%offset0, 8, 4] [2, %size1, 6] [1, 3, 1] %1 = tensor.extract_slice %0[%offset0, 8, 4] [2, %size1, 6] [1, 3, 1]
: tensor<?x?x?xf32> to tensor<2x?x6xf32> : tensor<?x?x?xf32> to tensor<2x?x6xf32>
return %0, %1: tensor<?x?x?xf32>, tensor<2x?x6xf32> return %0, %1: tensor<?x?x?xf32>, tensor<2x?x6xf32>
} }

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

// RUN: mlir-opt %s -test-linalg-greedy-fusion -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-linalg-greedy-fusion -split-input-file | FileCheck %s
func.func @matmul_tensors(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @matmul_tensors(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
%t0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) %t0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2: tensor<?x?xf32>) inits(%arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%0 = tensor.dim %t0, %c0 : tensor<?x?xf32> %0 = tensor.dim %t0, %c0 : tensor<?x?xf32>
%1 = tensor.dim %t0, %c1 : tensor<?x?xf32> %1 = tensor.dim %t0, %c1 : tensor<?x?xf32>
%2 = tensor.dim %arg1, %c1 : tensor<?x?xf32> %2 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
%3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) { %3 = scf.for %arg3 = %c0 to %0 step %c2 iter_args(%arg4 = %arg2) -> (tensor<?x?xf32>) {
%4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) { %4 = scf.for %arg5 = %c0 to %2 step %c3 iter_args(%arg6 = %arg4) -> (tensor<?x?xf32>) {
%5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) { %5 = scf.for %arg7 = %c0 to %1 step %c4 iter_args(%arg8 = %arg6) -> (tensor<?x?xf32>) {
%6 = tensor.extract_slice %t0[%arg3, %arg7][%c2, 4][1, 1] : tensor<?x?xf32> to tensor<?x4xf32> %6 = tensor.extract_slice %t0[%arg3, %arg7][%c2, 4][1, 1] : tensor<?x?xf32> to tensor<?x4xf32>
%7 = tensor.extract_slice %arg1[%arg7, %arg5][4, %c3][1, 1] : tensor<?x?xf32> to tensor<4x?xf32> %7 = tensor.extract_slice %arg1[%arg7, %arg5][4, %c3][1, 1] : tensor<?x?xf32> to tensor<4x?xf32>
%8 = tensor.extract_slice %arg8[%arg3, %arg5][%c2, %c3][1, 1] : tensor<?x?xf32> to tensor<?x?xf32> %8 = tensor.extract_slice %arg8[%arg3, %arg5][%c2, %c3][1, 1] : tensor<?x?xf32> to tensor<?x?xf32>
%9 = linalg.matmul ins(%6, %7 : tensor<?x4xf32>, tensor<4x?xf32>) outs(%8 : tensor<?x?xf32>) -> tensor<?x?xf32> %9 = linalg.matmul ins(%6, %7 : tensor<?x4xf32>, tensor<4x?xf32>) inits(%8 : tensor<?x?xf32>) -> tensor<?x?xf32>
%10 = tensor.insert_slice %9 into %arg8[%arg3, %arg5] [%c2, %c3] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32> %10 = tensor.insert_slice %9 into %arg8[%arg3, %arg5] [%c2, %c3] [1, 1] : tensor<?x?xf32> into tensor<?x?xf32>
scf.yield %10 : tensor<?x?xf32> scf.yield %10 : tensor<?x?xf32>
} }
scf.yield %5 : tensor<?x?xf32> scf.yield %5 : tensor<?x?xf32>
} }
scf.yield %4 : tensor<?x?xf32> scf.yield %4 : tensor<?x?xf32>
} }
return %3 : tensor<?x?xf32> return %3 : tensor<?x?xf32>
Show All 14 Lines
// CHECK: scf.for %[[J:[0-9a-z]*]] // CHECK: scf.for %[[J:[0-9a-z]*]]
// CHECK-NEXT: scf.for %[[K:[0-9a-z]*]] {{.*}} iter_args(%[[RES:[0-9a-z]*]] // CHECK-NEXT: scf.for %[[K:[0-9a-z]*]] {{.*}} iter_args(%[[RES:[0-9a-z]*]]
// CHECK-DAG: %[[stB1:.*]] = tensor.extract_slice %[[B]][%[[K]], %[[J]]] [4, 3] [1, 1] : tensor<?x?xf32> to tensor<4x3xf32> // CHECK-DAG: %[[stB1:.*]] = tensor.extract_slice %[[B]][%[[K]], %[[J]]] [4, 3] [1, 1] : tensor<?x?xf32> to tensor<4x3xf32>
// CHECK-DAG: %[[stF:.*]] = tensor.extract_slice %[[RES]][%[[I]], %[[J]]] [2, 3] [1, 1] : tensor<?x?xf32> to tensor<2x3xf32> // CHECK-DAG: %[[stF:.*]] = tensor.extract_slice %[[RES]][%[[I]], %[[J]]] [2, 3] [1, 1] : tensor<?x?xf32> to tensor<2x3xf32>
// //
// slices of the producing matmul. // slices of the producing matmul.
// CHECK-DAG: %[[stB2:.*]] = tensor.extract_slice %[[B]][0, %[[K]]] [%[[dB0]], 4] [1, 1] : tensor<?x?xf32> to tensor<?x4xf32> // CHECK-DAG: %[[stB2:.*]] = tensor.extract_slice %[[B]][0, %[[K]]] [%[[dB0]], 4] [1, 1] : tensor<?x?xf32> to tensor<?x4xf32>
// CHECK-DAG: %[[stC:.*]] = tensor.extract_slice %[[C]][%[[I]], %[[K]]] [2, 4] [1, 1] : tensor<?x?xf32> to tensor<2x4xf32> // CHECK-DAG: %[[stC:.*]] = tensor.extract_slice %[[C]][%[[I]], %[[K]]] [2, 4] [1, 1] : tensor<?x?xf32> to tensor<2x4xf32>
// CHECK: %[[stD:.*]] = linalg.matmul ins(%[[stA]], %[[stB2]] : tensor<2x?xf32>, tensor<?x4xf32>) outs(%[[stC]] : tensor<2x4xf32>) -> tensor<2x4xf32> // CHECK: %[[stD:.*]] = linalg.matmul ins(%[[stA]], %[[stB2]] : tensor<2x?xf32>, tensor<?x4xf32>) inits(%[[stC]] : tensor<2x4xf32>) -> tensor<2x4xf32>
// CHECK-NEXT: %[[stG:.*]] = linalg.matmul ins(%[[stD]], %[[stB1]] : tensor<2x4xf32>, tensor<4x3xf32>) outs(%[[stF]] : tensor<2x3xf32>) -> tensor<2x3xf32> // CHECK-NEXT: %[[stG:.*]] = linalg.matmul ins(%[[stD]], %[[stB1]] : tensor<2x4xf32>, tensor<4x3xf32>) inits(%[[stF]] : tensor<2x3xf32>) -> tensor<2x3xf32>
// CHECK-NEXT: tensor.insert_slice %[[stG]] into %[[RES]][%[[I]], %[[J]]] // CHECK-NEXT: tensor.insert_slice %[[stG]] into %[[RES]][%[[I]], %[[J]]]
// ----- // -----
func.func @conv_tensors_static(%input: tensor<1x225x225x3xf32>, %filter: tensor<3x3x3x32xf32>, %elementwise: tensor<1x112x112x32xf32>) -> tensor<1x112x112x32xf32> { func.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 = tensor.empty() : 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) inits(%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> inits(%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> {
%for1 = scf.for %iv1 = %c0 to %c112 step %c16 iter_args(%arg1 = %arg0) -> tensor<1x112x112x32xf32> { %for1 = scf.for %iv1 = %c0 to %c112 step %c16 iter_args(%arg1 = %arg0) -> tensor<1x112x112x32xf32> {
%for2 = scf.for %iv2 = %c0 to %c32 step %c4 iter_args(%arg2 = %arg1) -> tensor<1x112x112x32xf32> { %for2 = scf.for %iv2 = %c0 to %c32 step %c4 iter_args(%arg2 = %arg1) -> tensor<1x112x112x32xf32> {
%0 = tensor.extract_slice %conv[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> %0 = tensor.extract_slice %conv[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
%1 = tensor.extract_slice %elementwise[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> %1 = tensor.extract_slice %elementwise[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
%2 = tensor.extract_slice %arg2[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> %2 = tensor.extract_slice %arg2[0, %iv0, %iv1, %iv2][1, 8, 16, 4][1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
%add = linalg.generic %add = linalg.generic
{ {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 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(%0, %1 : tensor<1x8x16x4xf32>, tensor<1x8x16x4xf32>) outs(%2 : tensor<1x8x16x4xf32>) { ins(%0, %1 : tensor<1x8x16x4xf32>, tensor<1x8x16x4xf32>) inits(%2 : tensor<1x8x16x4xf32>) {
^bb0(%arg3: f32, %arg4: f32, %arg5: f32): ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
%result = arith.addf %arg3, %arg4 : f32 %result = arith.addf %arg3, %arg4 : f32
linalg.yield %result : f32 linalg.yield %result : f32
} -> tensor<1x8x16x4xf32> } -> tensor<1x8x16x4xf32>
%insert = tensor.insert_slice %add into %arg2[0, %iv0, %iv1, %iv2] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x8x16x4xf32> into tensor<1x112x112x32xf32> %insert = tensor.insert_slice %add into %arg2[0, %iv0, %iv1, %iv2] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x8x16x4xf32> into tensor<1x112x112x32xf32>
scf.yield %insert : tensor<1x112x112x32xf32> scf.yield %insert : tensor<1x112x112x32xf32>
} }
scf.yield %for2 : tensor<1x112x112x32xf32> scf.yield %for2 : tensor<1x112x112x32xf32>
} }
scf.yield %for1 : tensor<1x112x112x32xf32> scf.yield %for1 : tensor<1x112x112x32xf32>
} }
return %for0 : tensor<1x112x112x32xf32> return %for0 : tensor<1x112x112x32xf32>
} }
// 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:.+]] = tensor.empty() : 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) inits(%[[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]])
// CHECK-NEXT: %[[ST_ELEM:.+]] = tensor.extract_slice %[[ELEM]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> // CHECK-NEXT: %[[ST_ELEM:.+]] = tensor.extract_slice %[[ELEM]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
// CHECK-NEXT: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> // CHECK-NEXT: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
// CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.extract_slice %[[FILTER]][0, 0, 0, %[[IV2]]] [3, 3, 3, 4] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x4xf32> // CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.extract_slice %[[FILTER]][0, 0, 0, %[[IV2]]] [3, 3, 3, 4] [1, 1, 1, 1] : tensor<3x3x3x32xf32> to tensor<3x3x3x4xf32>
// CHECK-NEXT: %[[ST_FILL:.+]] = tensor.extract_slice %[[FILL]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32> // CHECK-NEXT: %[[ST_FILL:.+]] = tensor.extract_slice %[[FILL]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] [1, 1, 1, 1] : tensor<1x112x112x32xf32> to tensor<1x8x16x4xf32>
// CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_nhwc_hwcf // CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_nhwc_hwcf
// CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor<1x17x33x3xf32>, tensor<3x3x3x4xf32>) // CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor<1x17x33x3xf32>, tensor<3x3x3x4xf32>)
// CHECK-SAME: outs(%[[ST_FILL]] : tensor<1x8x16x4xf32>) // CHECK-SAME: inits(%[[ST_FILL]] : tensor<1x8x16x4xf32>)
// CHECK-NEXT: %[[ADD:.+]] = linalg.generic // CHECK-NEXT: %[[ADD:.+]] = linalg.generic
// CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor<1x8x16x4xf32>, tensor<1x8x16x4xf32>) // CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor<1x8x16x4xf32>, tensor<1x8x16x4xf32>)
// CHECK-SAME: outs(%[[ST_ARG2]] : tensor<1x8x16x4xf32>) // CHECK-SAME: inits(%[[ST_ARG2]] : tensor<1x8x16x4xf32>)
// CHECK: tensor.insert_slice %[[ADD]] into %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4] // CHECK: tensor.insert_slice %[[ADD]] into %[[ARG2]][0, %[[IV0]], %[[IV1]], %[[IV2]]] [1, 8, 16, 4]
// ----- // -----
func.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> { func.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> {
%cst = arith.constant 0.0 : f32 %cst = arith.constant 0.0 : f32
%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
%c3 = arith.constant 3 : index %c3 = arith.constant 3 : index
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%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 = tensor.empty(%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) inits(%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> inits(%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> {
%for1 = scf.for %iv1 = %c0 to %oh step %c16 iter_args(%arg1 = %arg0) -> tensor<?x?x?x?xf32> { %for1 = scf.for %iv1 = %c0 to %oh step %c16 iter_args(%arg1 = %arg0) -> tensor<?x?x?x?xf32> {
%for2 = scf.for %iv2 = %c0 to %ow step %c4 iter_args(%arg2 = %arg1) -> tensor<?x?x?x?xf32> { %for2 = scf.for %iv2 = %c0 to %ow step %c4 iter_args(%arg2 = %arg1) -> tensor<?x?x?x?xf32> {
%for3 = scf.for %iv3 = %c0 to %oc step %c2 iter_args(%arg3 = %arg2) -> tensor<?x?x?x?xf32> { %for3 = scf.for %iv3 = %c0 to %oc step %c2 iter_args(%arg3 = %arg2) -> tensor<?x?x?x?xf32> {
%n_size = affine.min affine_map<(d0)[s0] -> (8, -d0 + s0)>(%iv0)[%n] %n_size = affine.min affine_map<(d0)[s0] -> (8, -d0 + s0)>(%iv0)[%n]
%oh_size = affine.min affine_map<(d0)[s0] -> (16, -d0 + s0)>(%iv1)[%oh] %oh_size = affine.min affine_map<(d0)[s0] -> (16, -d0 + s0)>(%iv1)[%oh]
%ow_size = affine.min affine_map<(d0)[s0] -> (4, -d0 + s0)>(%iv2)[%ow] %ow_size = affine.min affine_map<(d0)[s0] -> (4, -d0 + s0)>(%iv2)[%ow]
%oc_size = affine.min affine_map<(d0)[s0] -> (2, -d0 + s0)>(%iv2)[%oc] %oc_size = affine.min affine_map<(d0)[s0] -> (2, -d0 + s0)>(%iv2)[%oc]
%0 = tensor.extract_slice %conv[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32> %0 = tensor.extract_slice %conv[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>
%1 = tensor.extract_slice %elementwise[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32> %1 = tensor.extract_slice %elementwise[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>
%2 = tensor.extract_slice %arg3[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32> %2 = tensor.extract_slice %arg3[%iv0, %iv1, %iv2, %iv3][%n_size, %oh_size, %ow_size, %oc_size][1, 1, 1, 1] : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>
%add = linalg.generic %add = linalg.generic
{ {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 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(%0, %1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) outs(%2 : tensor<?x?x?x?xf32>) { ins(%0, %1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) inits(%2 : tensor<?x?x?x?xf32>) {
^bb0(%arg4: f32, %arg5: f32, %arg6: f32): ^bb0(%arg4: f32, %arg5: f32, %arg6: f32):
%result = arith.addf %arg4, %arg5 : f32 %result = arith.addf %arg4, %arg5 : f32
linalg.yield %result : f32 linalg.yield %result : f32
} -> tensor<?x?x?x?xf32> } -> tensor<?x?x?x?xf32>
%insert = tensor.insert_slice %add into %arg3[%iv0, %iv1, %iv2, %iv3] [%n_size, %oh_size, %ow_size, %oc_size] [1, 1, 1, 1] : tensor<?x?x?x?xf32> into tensor<?x?x?x?xf32> %insert = tensor.insert_slice %add into %arg3[%iv0, %iv1, %iv2, %iv3] [%n_size, %oh_size, %ow_size, %oc_size] [1, 1, 1, 1] : tensor<?x?x?x?xf32> into tensor<?x?x?x?xf32>
scf.yield %insert : tensor<?x?x?x?xf32> scf.yield %insert : tensor<?x?x?x?xf32>
} }
Show All 26 Lines
// 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:.+]] = tensor.empty(%[[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) inits(%[[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>
// CHECK-DAG: %[[FILL_H:.+]] = tensor.dim %[[FILL]], %[[C1]] : tensor<?x?x?x?xf32> // CHECK-DAG: %[[FILL_H:.+]] = tensor.dim %[[FILL]], %[[C1]] : tensor<?x?x?x?xf32>
Show All 22 Lines
// CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]]
// CHECK-NEXT: %[[SIZE_ELEM_OC_2:.+]] = affine.min #[[BOUND2_MAP_2]](%[[IV3]], %[[IV2]])[%[[FILTER_OC]], %[[ELEM_OC]]] // CHECK-NEXT: %[[SIZE_ELEM_OC_2:.+]] = affine.min #[[BOUND2_MAP_2]](%[[IV3]], %[[IV2]])[%[[FILTER_OC]], %[[ELEM_OC]]]
// CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.extract_slice %[[FILTER]][0, 0, 0, %[[IV3]]] // CHECK-NEXT: %[[ST_FILTER:.+]] = tensor.extract_slice %[[FILTER]][0, 0, 0, %[[IV3]]]
// CHECK-SAME: [%[[FILTER_H]], %[[FILTER_W]], %[[FILTER_IC]], %[[SIZE_ELEM_OC_2]]] // CHECK-SAME: [%[[FILTER_H]], %[[FILTER_W]], %[[FILTER_IC]], %[[SIZE_ELEM_OC_2]]]
// CHECK-NEXT: %[[ST_FILL:.+]] = tensor.extract_slice %[[FILL]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK-NEXT: %[[ST_FILL:.+]] = tensor.extract_slice %[[FILL]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
// CHECK-SAME: [%[[SIZE_INPUT_N]], %[[SIZE_ELEM_OH_2]], %[[SIZE_ELEM_OW_2]], %[[SIZE_ELEM_OC_2]]] // CHECK-SAME: [%[[SIZE_INPUT_N]], %[[SIZE_ELEM_OH_2]], %[[SIZE_ELEM_OW_2]], %[[SIZE_ELEM_OC_2]]]
// CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_nhwc_hwcf // CHECK-NEXT: %[[ST_CONV:.+]] = linalg.conv_2d_nhwc_hwcf
// CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) // CHECK-SAME: ins(%[[ST_INPUT]], %[[ST_FILTER]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
// CHECK-SAME: outs(%[[ST_FILL]] : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> // CHECK-SAME: inits(%[[ST_FILL]] : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
// CHECK-NEXT: %[[ST_ADD:.+]] = linalg.generic // CHECK-NEXT: %[[ST_ADD:.+]] = linalg.generic
// CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) // CHECK-SAME: ins(%[[ST_CONV]], %[[ST_ELEM]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
// CHECK-SAME: outs(%[[ST_ARG]] : tensor<?x?x?x?xf32>) // CHECK-SAME: inits(%[[ST_ARG]] : tensor<?x?x?x?xf32>)
// CHECK: tensor.insert_slice %[[ST_ADD]] into %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK: tensor.insert_slice %[[ST_ADD]] into %[[ARG]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
// CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]] // CHECK-SAME: [%[[SIZE_ELEM_N]], %[[SIZE_ELEM_OH]], %[[SIZE_ELEM_OW]], %[[SIZE_ELEM_OC]]]
// ----- // -----
#map = affine_map<(d0, d1) -> (d0, d1)> #map = affine_map<(d0, d1) -> (d0, d1)>
// CHECK: func @pad_generic_static // CHECK: func @pad_generic_static
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
Show All 25 Linesfunc.func @pad_generic_static(%small_input: tensor<58x1xf32>, %large_input: tensor<64x128xf32>) -> tensor<64x128xf32> {
%d0 = tensor.dim %large_input, %c0 : tensor<64x128xf32> %d0 = tensor.dim %large_input, %c0 : tensor<64x128xf32>
%d1 = tensor.dim %large_input, %c1 : tensor<64x128xf32> %d1 = tensor.dim %large_input, %c1 : tensor<64x128xf32>
%pad = tensor.pad %small_input low[4, 60] high[2, 67] { %pad = tensor.pad %small_input low[4, 60] high[2, 67] {
^bb0(%arg0: index, %arg1: index): ^bb0(%arg0: index, %arg1: index):
tensor.yield %zero : f32 tensor.yield %zero : f32
} : tensor<58x1xf32> to tensor<64x128xf32> } : tensor<58x1xf32> to tensor<64x128xf32>
%fill = linalg.fill ins(%zero : f32) outs(%large_input : tensor<64x128xf32>) -> tensor<64x128xf32> %fill = linalg.fill ins(%zero : f32) inits(%large_input : tensor<64x128xf32>) -> tensor<64x128xf32>
%for0 = scf.for %iv0 = %c0 to %d0 step %c16 iter_args(%arg0 = %fill) -> tensor<64x128xf32> { %for0 = scf.for %iv0 = %c0 to %d0 step %c16 iter_args(%arg0 = %fill) -> tensor<64x128xf32> {
%for1 = scf.for %iv1 = %c0 to %d1 step %c32 iter_args(%arg1 = %arg0) -> tensor<64x128xf32> { %for1 = scf.for %iv1 = %c0 to %d1 step %c32 iter_args(%arg1 = %arg0) -> tensor<64x128xf32> {
%0 = tensor.extract_slice %pad[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32> %0 = tensor.extract_slice %pad[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32>
%1 = tensor.extract_slice %large_input[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32> %1 = tensor.extract_slice %large_input[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32>
%2 = tensor.extract_slice %arg1[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32> %2 = tensor.extract_slice %arg1[%iv0, %iv1][16, 32][1, 1] : tensor<64x128xf32> to tensor<16x32xf32>
%add = linalg.generic %add = linalg.generic
{indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
ins(%0, %1 : tensor<16x32xf32>, tensor<16x32xf32>) outs(%2 : tensor<16x32xf32>) { ins(%0, %1 : tensor<16x32xf32>, tensor<16x32xf32>) inits(%2 : tensor<16x32xf32>) {
^bb0(%arg4: f32, %arg5: f32, %arg6: f32): ^bb0(%arg4: f32, %arg5: f32, %arg6: f32):
%result = arith.addf %arg4, %arg5 : f32 %result = arith.addf %arg4, %arg5 : f32
linalg.yield %result : f32 linalg.yield %result : f32
} -> tensor<16x32xf32> } -> tensor<16x32xf32>
%insert = tensor.insert_slice %add into %arg1[%iv0, %iv1] [16, 32] [1, 1] : tensor<16x32xf32> into tensor<64x128xf32> %insert = tensor.insert_slice %add into %arg1[%iv0, %iv1] [16, 32] [1, 1] : tensor<16x32xf32> into tensor<64x128xf32>
scf.yield %insert : tensor<64x128xf32> scf.yield %insert : tensor<64x128xf32>
} }
scf.yield %for1 : tensor<64x128xf32> scf.yield %for1 : tensor<64x128xf32>
} }
return %for0 : tensor<64x128xf32> return %for0 : tensor<64x128xf32>
} }

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

// RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize | FileCheck %s
// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 1)> // CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 1)>
// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 2)> // CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 - 1, s1 + 2)>
// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)> // CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 3)> // CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 3)>
func.func @conv(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>) { func.func @conv(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>) {
linalg.conv_2d ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>) linalg.conv_2d ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) inits(%arg2 : memref<?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1
%1, %loop:2 = transform.structured.tile %0 [2, 3] %1, %loop:2 = transform.structured.tile %0 [2, 3]
} }
Show All 15 Lines
// CHECK: %[[T4:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T2]], %[[T0]]] // CHECK: %[[T4:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T2]], %[[T0]]]
// CHECK: %[[T5:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T3]], %[[T1]]] // CHECK: %[[T5:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T3]], %[[T1]]]
// CHECK: %[[T6:.*]] = affine.min #[[MAP2]](%[[ARG3]])[%[[T2]] // CHECK: %[[T6:.*]] = affine.min #[[MAP2]](%[[ARG3]])[%[[T2]]
// CHECK: %[[T7:.*]] = affine.min #[[MAP3]](%[[ARG4]])[%[[T3]]] // CHECK: %[[T7:.*]] = affine.min #[[MAP3]](%[[ARG4]])[%[[T3]]]
// CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] [%[[T4]], %[[T5]]] // CHECK: %[[SV1:.*]] = memref.subview %[[ARG0]][%[[ARG3]], %[[ARG4]]] [%[[T4]], %[[T5]]]
// CHECK: %[[SV2:.*]] = memref.subview %[[ARG2]][%[[ARG3]], %[[ARG4]]] [%[[T6]], %[[T7]]] // CHECK: %[[SV2:.*]] = memref.subview %[[ARG2]][%[[ARG3]], %[[ARG4]]] [%[[T6]], %[[T7]]]
// CHECK: linalg.conv_2d // CHECK: linalg.conv_2d
// CHECK-SAME: ins(%[[SV1]], %[[ARG1]] // CHECK-SAME: ins(%[[SV1]], %[[ARG1]]
// CHECK-SAME: outs(%[[SV2]] // CHECK-SAME: inits(%[[SV2]]

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

// RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s -check-prefix=TILE-10n25 // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s -check-prefix=TILE-10n25
func.func @indexed_vector(%arg0: memref<50xindex>) { func.func @indexed_vector(%arg0: memref<50xindex>) {
linalg.generic {indexing_maps = [affine_map<(i) -> (i)>], linalg.generic {indexing_maps = [affine_map<(i) -> (i)>],
iterator_types = ["parallel"]} iterator_types = ["parallel"]}
outs(%arg0 : memref<50xindex>) { inits(%arg0 : memref<50xindex>) {
^bb0(%a: index): ^bb0(%a: index):
%i = linalg.index 0 : index %i = linalg.index 0 : index
linalg.yield %i : index linalg.yield %i : index
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
Show All 11 Lines
// TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[J]]) // TILE-10n25: %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[J]])
// TILE-10n25: linalg.yield %[[NEW_I]] : index // TILE-10n25: linalg.yield %[[NEW_I]] : index
// ----- // -----
func.func @indexed_matrix(%arg0: memref<50x50xindex>) { func.func @indexed_matrix(%arg0: memref<50x50xindex>) {
linalg.generic {indexing_maps = [affine_map<(i, j) -> (i, j)>], linalg.generic {indexing_maps = [affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
outs(%arg0 : memref<50x50xindex>) { inits(%arg0 : memref<50x50xindex>) {
^bb0(%a: index): ^bb0(%a: index):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%sum = arith.addi %i, %j : index %sum = arith.addi %i, %j : index
linalg.yield %sum : index linalg.yield %sum : index
} }
return return
} }
Show All 20 Lines

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

// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
// CHECK-LABEL: func @matmul_tensors( // CHECK-LABEL: func @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> // CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> { // CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
func.func @matmul_tensors( func.func @matmul_tensors(
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) { // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) {
// 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>) {
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) {
// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>) // CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32> // CHECK-SAME: inits(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32> // CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32>
// CHECK: scf.yield %[[TD]] : tensor<?x?xf32> // CHECK: scf.yield %[[TD]] : 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>
%0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2: tensor<?x?xf32>) inits(%arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: return %[[TD0]] : tensor<?x?xf32> // CHECK: return %[[TD0]] : tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
Show All 13 Linesfunc.func @generic_op_tensors(
%2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32> %2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%3 = tensor.empty(%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>) { inits(%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>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
Show All 9 Lines
// 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>)
// CHECK-SAME: outs(%[[STARG2]] : tensor<?x?x?xf32>) // CHECK-SAME: inits(%[[STARG2]] : tensor<?x?x?xf32>)
// CHECK: %[[TD:.+]] = tensor.insert_slice %[[STRETURN]] into %[[TC2]] // CHECK: %[[TD:.+]] = tensor.insert_slice %[[STRETURN]] into %[[TC2]]
// CHECK: scf.yield %[[TD]] // CHECK: scf.yield %[[TD]]
// CHECK: } // CHECK: }
// CHECK: scf.yield %[[TD2]] // CHECK: scf.yield %[[TD2]]
// CHECK: } // CHECK: }
// CHECK: scf.yield %[[TD1]] // CHECK: scf.yield %[[TD1]]
// CHECK: } // CHECK: }
// CHECK: return %[[TD0]] // CHECK: return %[[TD0]]
Show All 28 Linesfunc.func @fold_extract_slice(
// CHECK-SAME: 1, 1 // CHECK-SAME: 1, 1
// CHECK: {{.*}} = linalg.generic {{.*}} ins(%[[T0]] // CHECK: {{.*}} = linalg.generic {{.*}} ins(%[[T0]]
%2 = linalg.generic %2 = linalg.generic
{indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d0, d1)>], affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%1, %arg2 : tensor<?x42xf32>, tensor<?x42x?xf32>) ins(%1, %arg2 : tensor<?x42xf32>, tensor<?x42x?xf32>)
outs(%arg1 : tensor<?x42xf32>) { inits(%arg1 : tensor<?x42xf32>) {
^bb0(%arg3 : f32, %arg4: f32, %arg5: f32): ^bb0(%arg3 : f32, %arg4: f32, %arg5: f32):
%5 = arith.addf %arg3, %arg5 : f32 %5 = arith.addf %arg3, %arg5 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<?x42xf32> } -> tensor<?x42xf32>
return %2 : tensor<?x42xf32> return %2 : tensor<?x42xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [2, 3, 4] %1, %loops:3 = transform.structured.tile %0 [2, 3, 4]
} }

mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir

Show All 15 Lines// CHECK-SAME: %[[C:[0-9a-z]+]]: tensor<?x?xf32>
// 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: scf.foreach_thread ({{.*}}) in (%[[C10]], %[[C20]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) -> (tensor<?x?xf32>) { // CHECK: scf.foreach_thread ({{.*}}) in (%[[C10]], %[[C20]]) shared_outs(%[[C_BLK:.*]] = %[[C]]) -> (tensor<?x?xf32>) {
// CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[tC:.*]] = tensor.extract_slice %[[C_BLK]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32> // CHECK: %[[tC:.*]] = tensor.extract_slice %[[C_BLK]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
// CHECK: %[[RES:.*]] = linalg.matmul // CHECK: %[[RES:.*]] = linalg.matmul
// CHECK-SAME: ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>) // CHECK-SAME: ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[tC]] : tensor<?x?xf32>) -> tensor<?x?xf32> // CHECK-SAME: inits(%[[tC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: scf.foreach_thread.perform_concurrently { // CHECK: scf.foreach_thread.perform_concurrently {
// CHECK-NEXT: tensor.parallel_insert_slice %[[RES]] into %[[C_BLK]]{{.*}} : // CHECK-NEXT: tensor.parallel_insert_slice %[[RES]] into %[[C_BLK]]{{.*}} :
// CHECK-SAME: tensor<?x?xf32> into tensor<?x?xf32> // CHECK-SAME: tensor<?x?xf32> into tensor<?x?xf32>
// CHECK-NEXT: } // CHECK-NEXT: }
// CHECK-NEXT: } {thread_dim_mapping = [1, 0]} // CHECK-NEXT: } {thread_dim_mapping = [1, 0]}
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>) inits(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapped to dims [1, 0]) %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20] (mapped to dims [1, 0])
} }
Show All 24 Linesfunc.func @matmul_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> {
// CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]])
// CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] :
// CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] :
// CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] :
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: scf.foreach_thread.perform_concurrently // CHECK: scf.foreach_thread.perform_concurrently
// CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>) %0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>)
outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>) inits(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>)
return %0 : tensor<100x300xf32> return %0 : tensor<100x300xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 21] %1:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 21]
} }
Show All 24 Linesfunc.func @matmul_tile_size_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: %[[LB1:.+]] = affine.apply #[[$map6]](%[[IV1]]) // CHECK: %[[LB1:.+]] = affine.apply #[[$map6]](%[[IV1]])
// CHECK: tensor.extract_slice %[[A]] // CHECK: tensor.extract_slice %[[A]]
// CHECK: tensor.extract_slice %[[B]] // CHECK: tensor.extract_slice %[[B]]
// CHECK: tensor.extract_slice %[[C_BLK]] // CHECK: tensor.extract_slice %[[C_BLK]]
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: scf.foreach_thread.perform_concurrently // CHECK: scf.foreach_thread.perform_concurrently
// CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>) inits(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 20] %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 20]
} }
Show All 21 Linesfunc.func @matmul_tile_size_static(%A: tensor<100x200xf32>, %B: tensor<200x300xf32>, %C: tensor<100x300xf32>) -> tensor<100x300xf32> {
// CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]]) // CHECK: %[[LB1:.+]] = affine.apply #[[$map3]](%[[IV1]])
// CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] : // CHECK: %[[tA:.+]] = tensor.extract_slice %[[A]][%[[LB0]], 0] [10, 200] [1, 1] :
// CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] : // CHECK: %[[tB:.+]] = tensor.extract_slice %[[B]][0, %[[LB1]]] [200, %[[TS]]] [1, 1] :
// CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] : // CHECK: %[[tC:.+]] = tensor.extract_slice %[[C_BLK]][%[[LB0]], %[[LB1]]] [10, %[[TS]]] [1, 1] :
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: scf.foreach_thread.perform_concurrently // CHECK: scf.foreach_thread.perform_concurrently
// CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: tensor.parallel_insert_slice
%0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>) %0 = linalg.matmul ins(%A, %B : tensor<100x200xf32>, tensor<200x300xf32>)
outs(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>) inits(%C : tensor<100x300xf32>) -> (tensor<100x300xf32>)
return %0 : tensor<100x300xf32> return %0 : tensor<100x300xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 21] %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 21]
} }
// ----- // -----
module { module {
func.func @extract_source(%A: tensor<4xf32>, %B: tensor<16xf32>) -> tensor<4xf32> { func.func @extract_source(%A: tensor<4xf32>, %B: tensor<16xf32>) -> tensor<4xf32> {
%B1 = tensor.extract_slice %B[10] [4] [1] : tensor<16xf32> to tensor<4xf32> %B1 = tensor.extract_slice %B[10] [4] [1] : tensor<16xf32> to tensor<4xf32>
%result = linalg.generic {indexing_maps = [ %result = linalg.generic {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(%A : tensor<4xf32>) outs(%B1 : tensor<4xf32>) { ins(%A : tensor<4xf32>) inits(%B1 : tensor<4xf32>) {
^bb0(%arg3: f32, %arg4: f32): // no predecessors ^bb0(%arg3: f32, %arg4: f32): // no predecessors
%2 = arith.addf %arg3, %arg3 : f32 %2 = arith.addf %arg3, %arg3 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %result : tensor<4xf32> return %result : tensor<4xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
Show All 38 Linesfunc.func @matmul_tile_size_dynamic_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: tensor.extract_slice %[[A]] // CHECK: tensor.extract_slice %[[A]]
// CHECK: tensor.extract_slice %[[B]] // CHECK: tensor.extract_slice %[[B]]
// CHECK: tensor.extract_slice %[[C_BLK]] // CHECK: tensor.extract_slice %[[C_BLK]]
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: scf.foreach_thread.perform_concurrently // CHECK: scf.foreach_thread.perform_concurrently
// CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: tensor.parallel_insert_slice
%tile_size = "test.dummy"() : () -> (index) %tile_size = "test.dummy"() : () -> (index)
%0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>) inits(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%sz = transform.structured.match ops{["test.dummy"]} in %arg1 %sz = transform.structured.match ops{["test.dummy"]} in %arg1
%1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz, 20] %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz, 20]
Show All 32 Lines// CHECK-NEXT: tensor.parallel_insert_slice %[[RES1]]#1 into %[[OUT2]][%[[LB]]] [%[[TS]]] [1] :
%res1, %res2 = linalg.generic %res1, %res2 = 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)>,
affine_map<(d0) -> (d0)>], affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%IN1, %IN2 : tensor<100xf32>, tensor<100xf32>) } ins(%IN1, %IN2 : tensor<100xf32>, tensor<100xf32>)
outs(%OUT1, %OUT2 : tensor<100xf32>, tensor<100xf32>) inits(%OUT1, %OUT2 : tensor<100xf32>, tensor<100xf32>)
{ {
^bb0(%a1: f32, %a2: f32, %a3: f32, %a4: f32): ^bb0(%a1: f32, %a2: f32, %a3: f32, %a4: f32):
%1 = arith.addf %a1, %a3 : f32 %1 = arith.addf %a1, %a3 : f32
%2 = arith.addf %a2, %a4 : f32 %2 = arith.addf %a2, %a4 : f32
linalg.yield %1, %2 : f32,f32 linalg.yield %1, %2 : f32,f32
} -> (tensor<100xf32>, tensor<100xf32>) } -> (tensor<100xf32>, tensor<100xf32>)
return %res1, %res2 : tensor<100xf32>, tensor<100xf32> return %res1, %res2 : tensor<100xf32>, tensor<100xf32>
} }
Show All 36 Lines %res2, %res3 = linalg.generic {
indexing_maps = [affine_map<(d0,d1) -> (d1)>, indexing_maps = [affine_map<(d0,d1) -> (d1)>,
affine_map<(d0,d1) -> (d1,d0)>, affine_map<(d0,d1) -> (d1,d0)>,
affine_map<(d0,d1) -> (d0)>, affine_map<(d0,d1) -> (d0)>,
affine_map<(d0,d1) -> (d0,d1)>, affine_map<(d0,d1) -> (d0,d1)>,
affine_map<(d0,d1) -> (d0)> affine_map<(d0,d1) -> (d0)>
], ],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} ins(%IN1, %IN2, %IN3 : tensor<100xf32>, tensor<100x300xf32>, tensor<300xf32>) } ins(%IN1, %IN2, %IN3 : tensor<100xf32>, tensor<100x300xf32>, tensor<300xf32>)
outs(%OUT1, %OUT2: tensor<300x100xf32>, tensor<300xf32>) { inits(%OUT1, %OUT2: tensor<300x100xf32>, tensor<300xf32>) {
^bb0(%i1: f32, %i2: f32, %i3: f32, %o1: f32, %o2: f32): ^bb0(%i1: f32, %i2: f32, %i3: f32, %o1: f32, %o2: f32):
%1 = arith.addf %i1, %o1 : f32 %1 = arith.addf %i1, %o1 : f32
%2 = arith.addf %i2, %1 : f32 %2 = arith.addf %i2, %1 : f32
%3 = arith.addf %i3, %2 : f32 %3 = arith.addf %i3, %2 : f32
linalg.yield %3, %i3 : f32, f32 linalg.yield %3, %i3 : f32, f32
} -> (tensor<300x100xf32>, tensor<300xf32>) } -> (tensor<300x100xf32>, tensor<300xf32>)
return %res2, %res3 : tensor<300x100xf32>, tensor<300xf32> return %res2, %res3 : tensor<300x100xf32>, tensor<300xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%IN_MAT2: !pdl.operation): ^bb1(%IN_MAT2: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2 %0 = transform.structured.match ops{["linalg.generic"]} in %IN_MAT2
%foreach_thread, %tiled_generic = transform.structured.tile_to_foreach_thread_op %0 num_threads [4] %foreach_thread, %tiled_generic = transform.structured.tile_to_foreach_thread_op %0 num_threads [4]
} }

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

// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s // RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
// CHECK-LABEL: @conv_2d_nhwc_hwcf // CHECK-LABEL: @conv_2d_nhwc_hwcf
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x?xf32>, // CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x?xf32>,
// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?xf32> // CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?xf32>
// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32> // CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
func.func @conv_2d_nhwc_hwcf(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?x?x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> { func.func @conv_2d_nhwc_hwcf(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?x?x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
// 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: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]] // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
// CHECK: %[[SLICERES:.+]] = linalg.conv_1d_nwc_wcf // CHECK: %[[SLICERES:.+]] = linalg.conv_1d_nwc_wcf
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]] // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
%0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?x?x?xf32>) ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?x?x?xf32>)
outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> inits (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
return %0 : tensor<?x1x?x?xf32> return %0 : tensor<?x1x?x?xf32>
} }
// CHECK-LABEL: @conv_2d_nchw_fchw // CHECK-LABEL: @conv_2d_nchw_fchw
// CHECK-SAME: (%[[ARG0:[0-9a-z]+]]: tensor<?x?x1x?xf32>, // CHECK-SAME: (%[[ARG0:[0-9a-z]+]]: tensor<?x?x1x?xf32>,
// CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<?x?x1x?xf32>, // CHECK-SAME: %[[ARG1:[0-9a-z]+]]: tensor<?x?x1x?xf32>,
// CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<?x?x1x?xf32>) // CHECK-SAME: %[[ARG2:[0-9a-z]+]]: tensor<?x?x1x?xf32>)
func.func @conv_2d_nchw_fchw(%input: tensor<?x?x1x?xf32>, %filter: tensor<?x?x1x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> { func.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: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]] // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]] // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]] // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
// CHECK: %[[SLICERES:.+]] = linalg.conv_1d_ncw_fcw // CHECK: %[[SLICERES:.+]] = linalg.conv_1d_ncw_fcw
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]] // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
%0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>, %0 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%input, %filter: tensor<?x?x1x?xf32>, tensor<?x?x1x?xf32>) ins (%input, %filter: tensor<?x?x1x?xf32>, tensor<?x?x1x?xf32>)
outs (%init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> inits (%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:.+]] = tensor.empty // CHECK: %[[RES:.+]] = tensor.empty
%init = tensor.empty() : 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: inits(%[[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>}
ins(%input, %filter: tensor<1x1x113x96xf32>, tensor<1x3x96xf32>) ins(%input, %filter: tensor<1x1x113x96xf32>, tensor<1x3x96xf32>)
outs(%init: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32> inits(%init: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32>
// CHECK: %[[INSERTED]] // CHECK: %[[INSERTED]]
return %0: tensor<1x1x56x96xf32> return %0: tensor<1x1x56x96xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match interface{LinalgOp} in %arg1 %0 = transform.structured.match interface{LinalgOp} in %arg1
%1 = transform.structured.decompose %0 %1 = transform.structured.decompose %0
} }

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

// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s // RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
#map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)> #map0 = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
#map1 = affine_map<(d0)[s0] -> (d0 * s0)> #map1 = affine_map<(d0)[s0] -> (d0 * s0)>
#map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)> #map2 = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
module { module {
// CHECK-LABEL: func.func @fuse_tileable_op // CHECK-LABEL: func.func @fuse_tileable_op
// CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
// CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32> // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>
// CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32> // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>
func.func @fuse_tileable_op(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> { func.func @fuse_tileable_op(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = linalg.fill ins(%cst : f32) outs(%arg1 : tensor<?xf32>) -> tensor<?xf32> %0 = linalg.fill ins(%cst : f32) inits(%arg1 : tensor<?xf32>) -> tensor<?xf32>
%d0 = tensor.dim %arg1, %c0 : tensor<?xf32> %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>
%1 = affine.apply #map0()[%d0, %arg0] %1 = affine.apply #map0()[%d0, %arg0]
// CHECK: scf.foreach_thread {{.*}} { // CHECK: scf.foreach_thread {{.*}} {
%2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<?xf32>) { %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %arg2) -> (tensor<?xf32>) {
%3 = affine.apply #map1(%arg3)[%arg0] %3 = affine.apply #map1(%arg3)[%arg0]
%4 = affine.min #map2(%arg3)[%d0, %arg0] %4 = affine.min #map2(%arg3)[%d0, %arg0]
%5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}] // CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}]
// CHECK: %[[T1:.*]] = linalg.fill {{.*}} outs(%[[T0]] // CHECK: %[[T1:.*]] = linalg.fill {{.*}} inits(%[[T0]]
%6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %6 = tensor.extract_slice %0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]] // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]]
%7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32> %7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) inits(%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<?xf32> tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>
} }
} }
// CHECK: } // CHECK: }
func.return %2 : tensor<?xf32> func.return %2 : tensor<?xf32>
} }
Show All 30 Lines func.func @fuse_untileable_op(%arg0: index, %arg1: tensor<64xf32>, %arg2: tensor<64xf32>) -> tensor<64xf32> {
// 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:.*]] = tensor.empty // 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>) inits(%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>
} }
Show All 17 Lines
module { module {
// CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg // CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg
// CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index // CHECK-SAME: %[[CHUNK_SIZE:[0-9a-z]+]]: index
// CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32> // CHECK-SAME: %[[IN:[0-9a-z]+]]: tensor<?xf32>
// CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32> // CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<?xf32>
func.func @fuse_tileable_op_through_bbarg(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> { func.func @fuse_tileable_op_through_bbarg(%arg0: index, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<?xf32>) -> tensor<?xf32> %0 = linalg.fill ins(%cst : f32) inits(%arg2 : tensor<?xf32>) -> tensor<?xf32>
%d0 = tensor.dim %arg1, %c0 : tensor<?xf32> %d0 = tensor.dim %arg1, %c0 : tensor<?xf32>
%1 = affine.apply #map0()[%d0, %arg0] %1 = affine.apply #map0()[%d0, %arg0]
// CHECK: scf.foreach_thread {{.*}} shared_outs(%[[BBARGOUT:.*]] = %[[OUT]]) -> (tensor<?xf32>) { // CHECK: scf.foreach_thread {{.*}} shared_outs(%[[BBARGOUT:.*]] = %[[OUT]]) -> (tensor<?xf32>) {
%2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %0) -> (tensor<?xf32>) { %2 = scf.foreach_thread (%arg3) in (%1) shared_outs(%o = %0) -> (tensor<?xf32>) {
%3 = affine.apply #map1(%arg3)[%arg0] %3 = affine.apply #map1(%arg3)[%arg0]
%4 = affine.min #map2(%arg3)[%d0, %arg0] %4 = affine.min #map2(%arg3)[%d0, %arg0]
%5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[BBARGOUT]][%{{.*}}] [%{{.*}}] [{{.*}}] // CHECK: %[[T0:.*]] = tensor.extract_slice %[[BBARGOUT]][%{{.*}}] [%{{.*}}] [{{.*}}]
// CHECK: %[[T1:.*]] = linalg.fill {{.*}} outs(%[[T0]] // CHECK: %[[T1:.*]] = linalg.fill {{.*}} inits(%[[T0]]
%6 = tensor.extract_slice %arg1[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %6 = tensor.extract_slice %arg1[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T2:.*]] = linalg.elemwise_unary {{.*}} outs(%[[T1]] // CHECK: %[[T2:.*]] = linalg.elemwise_unary {{.*}} inits(%[[T1]]
%7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32> %7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) inits(%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<?xf32> tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>
} }
} }
// CHECK: } // CHECK: }
func.return %2 : tensor<?xf32> func.return %2 : tensor<?xf32>
} }
Show All 22 Linesmodule {
// CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32> // CHECK-SAME: %[[OUT_3:[0-9a-z]+]]: tensor<?xf32>
func.func @fuse_tileable_multi_output_op(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>) -> tensor<?xf32> { func.func @fuse_tileable_multi_output_op(%idx: index, %in: tensor<?xf32>, %out_1: tensor<?xf32>, %out_2: tensor<?xf32>, %out_3: tensor<?xf32>) -> tensor<?xf32> {
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} ins(%in : tensor<?xf32>) outs(%out_1, %out_3 : tensor<?xf32>, tensor<?xf32>) { } ins(%in : tensor<?xf32>) inits(%out_1, %out_3 : tensor<?xf32>, tensor<?xf32>) {
^bb0(%a: f32, %b: f32, %c: f32): ^bb0(%a: f32, %b: f32, %c: f32):
%d = arith.addf %a, %b : f32 %d = arith.addf %a, %b : f32
%e = arith.addf %d, %c : f32 %e = arith.addf %d, %c : f32
linalg.yield %d, %e : f32, f32 linalg.yield %d, %e : f32, f32
} -> (tensor<?xf32>, tensor<?xf32>) } -> (tensor<?xf32>, tensor<?xf32>)
%d0 = tensor.dim %out_1, %c0 : tensor<?xf32> %d0 = tensor.dim %out_1, %c0 : tensor<?xf32>
%1 = affine.apply #map0()[%d0, %idx] %1 = affine.apply #map0()[%d0, %idx]
// CHECK: scf.foreach_thread {{.*}} { // CHECK: scf.foreach_thread {{.*}} {
%2 = scf.foreach_thread (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) { %2 = scf.foreach_thread (%i) in (%1) shared_outs(%o = %out_2) -> (tensor<?xf32>) {
%3 = affine.apply #map1(%i)[%idx] %3 = affine.apply #map1(%i)[%idx]
%4 = affine.min #map2(%i)[%d0, %idx] %4 = affine.min #map2(%i)[%d0, %idx]
%5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %5 = tensor.extract_slice %o[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}] // CHECK: %[[T0:.*]] = tensor.extract_slice %[[IN]][%{{.*}}] [%{{.*}}] [{{.*}}]
// CHECK: %[[T1:.*]]:2 = linalg.generic {{.*}} ins(%[[T0]] // CHECK: %[[T1:.*]]:2 = linalg.generic {{.*}} ins(%[[T0]]
%6 = tensor.extract_slice %0#0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32> %6 = tensor.extract_slice %0#0[%3] [%4] [1] : tensor<?xf32> to tensor<?xf32>
// CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]]#0 // CHECK: %[[T2:.*]] = linalg.elemwise_unary ins(%[[T1]]#0
%7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) outs(%5 : tensor<?xf32>) -> tensor<?xf32> %7 = linalg.elemwise_unary ins(%6 : tensor<?xf32>) inits(%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<?xf32> tensor.parallel_insert_slice %7 into %o[%3] [%4] [1] : tensor<?xf32> into tensor<?xf32>
} }
} }
// CHECK: } // CHECK: }
func.return %2 : tensor<?xf32> func.return %2 : tensor<?xf32>
} }
Show All 9 Lines

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

// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s // RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
// CHECK-LABEL: func.func @fuse_unary // CHECK-LABEL: func.func @fuse_unary
func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: %[[RES:.*]] = scf.for // CHECK: %[[RES:.*]] = scf.for
// CHECK: scf.for // CHECK: scf.for
// CHECK: linalg.elemwise_unary // CHECK: linalg.elemwise_unary
// CHECK: linalg.elemwise_binary // CHECK: linalg.elemwise_binary
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
%0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>) %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
%1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) %1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1 %0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1
%1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]} %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
} }
// ----- // -----
// CHECK-LABEL: func.func @fuse_unary // CHECK-LABEL: func.func @fuse_unary
func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @fuse_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: %[[PARTIAL_RES:.*]] = scf.for // CHECK: %[[PARTIAL_RES:.*]] = scf.for
// CHECK: scf.for // CHECK: scf.for
// CHECK: linalg.elemwise_unary // CHECK: linalg.elemwise_unary
// CHECK: linalg.elemwise_binary // CHECK: linalg.elemwise_binary
// CHECK: %[[RES:.*]] = scf.for {{.*}}%[[PARTIAL_RES]] // CHECK: %[[RES:.*]] = scf.for {{.*}}%[[PARTIAL_RES]]
// CHECK: scf.for // CHECK: scf.for
// CHECK: linalg.elemwise_unary // CHECK: linalg.elemwise_unary
// CHECK: linalg.elemwise_binary // CHECK: linalg.elemwise_binary
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
%0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>) %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
%1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) %1 = linalg.elemwise_binary ins(%0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
return %1 : tensor<?x?xf32> return %1 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1 %0 = transform.structured.match ops{["linalg.elemwise_binary"]} in %arg1
%1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]} %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [32, 32], tile_interchange = [0, 1]}
%loop = transform.cast %loops#0 : !pdl.operation to !transform.op<"scf.for"> %loop = transform.cast %loops#0 : !pdl.operation to !transform.op<"scf.for">
Show All 10 Lines
// CHECK-DAG: %[[INIT:.+]] = tensor.empty() // CHECK-DAG: %[[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: %[[RES:.*]] = scf.for %[[IV0:.+]] = %{{.+}} to %{{.+}} step %[[C5]] iter_args(%[[FOR_ARG0:.+]] = %[[INIT]]) // CHECK: %[[RES:.*]] = 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 {{.+}} inits(%[[OUT_SLICE1]] : tensor<?x?xf32>)
// //
// Extra 4 constant is introduced, discard it. // Extra 4 constant is introduced, discard it.
// CHECK: arith.constant 4 : index // CHECK: arith.constant 4 : index
// CHECK: %[[C4:.+]] = arith.constant 4 : index // CHECK: %[[C4:.+]] = arith.constant 4 : index
// CHECK: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %[[C4]] iter_args(%[[FOR_ARG2:.+]] = %[[FILL]]) // CHECK: scf.for %[[IV2:.+]] = %{{.+}} to %{{.+}} step %[[C4]] iter_args(%[[FOR_ARG2:.+]] = %[[FILL]])
// CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[OUT_SLICE0]] // CHECK: %[[IN_SLICE:.+]] = tensor.extract_slice %[[OUT_SLICE0]]
// CHECK: %[[OUT_SLICE2:.+]] = tensor.extract_slice %[[FOR_ARG2]][0, 0] // CHECK: %[[OUT_SLICE2:.+]] = tensor.extract_slice %[[FOR_ARG2]][0, 0]
// CHECK: linalg.generic {{.+}} ins(%[[IN_SLICE]] : tensor<?x?x?xf32>) outs(%[[OUT_SLICE2]] : tensor<?x?xf32>) // CHECK: linalg.generic {{.+}} ins(%[[IN_SLICE]] : tensor<?x?x?xf32>) inits(%[[OUT_SLICE2]] : tensor<?x?xf32>)
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
%fill = linalg.fill ins(%five : f32) outs(%init : tensor<12x25xf32>) -> tensor<12x25xf32> %fill = linalg.fill ins(%five : f32) inits(%init : tensor<12x25xf32>) -> tensor<12x25xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>], indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2)>],
iterator_types = ["parallel", "reduction", "parallel"] iterator_types = ["parallel", "reduction", "parallel"]
} ins(%input : tensor<12x7x25xf32>) outs(%fill : tensor<12x25xf32>) { } ins(%input : tensor<12x7x25xf32>) inits(%fill : tensor<12x25xf32>) {
^bb0(%arg0: f32, %arg1: f32): ^bb0(%arg0: f32, %arg1: f32):
%2 = arith.addf %arg0, %arg1 : f32 %2 = arith.addf %arg0, %arg1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<12x25xf32> } -> tensor<12x25xf32>
func.return %0 : tensor<12x25xf32> func.return %0 : tensor<12x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [5, 0, 7], tile_interchange = [0, 2, 1]} %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [5, 0, 7], tile_interchange = [0, 2, 1]}
%2, %loops_2 = transform.structured.tile %1 [0, 4] %2, %loops_2 = transform.structured.tile %1 [0, 4]
} }

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

// RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s // RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
// CHECK-LABEL: func.func @generalize_unary // CHECK-LABEL: func.func @generalize_unary
func.func @generalize_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @generalize_unary(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK-NOT: linalg.elemwise_unary // CHECK-NOT: linalg.elemwise_unary
// CHECK: linalg.generic // CHECK: linalg.generic
%0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>) %0 = linalg.elemwise_unary ins(%arg0 : tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg1: tensor<?x?xf32>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.elemwise_unary"]} in %arg1 %0 = transform.structured.match ops{["linalg.elemwise_unary"]} in %arg1
%1 = transform.structured.generalize %0 %1 = transform.structured.generalize %0
} }

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

// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d1, d0)> // CHECK: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-LABEL: @interchange_generic // CHECK-LABEL: @interchange_generic
func.func @interchange_generic(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @interchange_generic(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]] // CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]]
%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(%arg1 : tensor<?x?xf32>) { } ins(%arg0 : tensor<?x?xf32>) inits(%arg1 : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
%1 = math.exp %arg2 : f32 %1 = math.exp %arg2 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
transform.structured.interchange %0 { iterator_interchange = [1, 0]} transform.structured.interchange %0 { iterator_interchange = [1, 0]}
} }
// ----- // -----
func.func @interchange_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @interchange_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) -> tensor<?x?xf32> {
// expected-note @below {{when applied to this op}} // expected-note @below {{when applied to this op}}
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
// expected-error @below {{transform applied to the wrong op kind}} // expected-error @below {{transform applied to the wrong op kind}}
transform.structured.interchange %0 { iterator_interchange = [1, 0]} transform.structured.interchange %0 { iterator_interchange = [1, 0]}
} }

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

Show First 20 Lines • Show All 42 Lines • ▼ Show 20 Lines
#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 = tensor.empty() : 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>) { inits(%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>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
Show All 13 Lines

mlir/test/Dialect/Linalg/transform-op-multitile-sizes.mlir

// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s // RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
// CHECK-DAG: #[[$MAP13:.+]] = affine_map<() -> (13)> // CHECK-DAG: #[[$MAP13:.+]] = affine_map<() -> (13)>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 } transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 }
} }
// CHECK-LABEL: @multitile_sizes_static // CHECK-LABEL: @multitile_sizes_static
func.func @multitile_sizes_static( func.func @multitile_sizes_static(
%arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>) %arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>)
-> tensor<13x42xf32> { -> tensor<13x42xf32> {
%0 = linalg.matmul ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>)
outs(%arg2: tensor<13x42xf32>) inits(%arg2: tensor<13x42xf32>)
-> tensor<13x42xf32> -> tensor<13x42xf32>
// The first application computes the total size. // The first application computes the total size.
// CHECK: %{{.*}} = affine.apply #[[$MAP13]]() // CHECK: %{{.*}} = affine.apply #[[$MAP13]]()
// CHECK: %[[SIZE:.+]] = affine.apply #[[$MAP13]]() // CHECK: %[[SIZE:.+]] = affine.apply #[[$MAP13]]()
// CHECK: %[[COND:.+]] = arith.cmpi eq, %[[SIZE]], %{{.*}} // CHECK: %[[COND:.+]] = arith.cmpi eq, %[[SIZE]], %{{.*}}
// CHECK: cf.assert %[[COND]], "could not compute dynamic multi-size tile shapes" // CHECK: cf.assert %[[COND]], "could not compute dynamic multi-size tile shapes"
return %0 : tensor<13x42xf32> return %0 : tensor<13x42xf32>
Show All 29 Linesfunc.func @multitile_sizes_dynamic(
// CHECK: affine.apply #[[$MAP_T]]() // CHECK: affine.apply #[[$MAP_T]]()
// CHECK: affine.apply #[[$MAP_D]]()[%[[DIM]]] // CHECK: affine.apply #[[$MAP_D]]()[%[[DIM]]]
// CHECK: affine.apply #[[$MAP_S]]()[%[[DIM]]] // CHECK: affine.apply #[[$MAP_S]]()[%[[DIM]]]
// CHECK: affine.apply #[[$MAP_V]]()[%[[DIM]]] // CHECK: affine.apply #[[$MAP_V]]()[%[[DIM]]]
// CHECK: affine.apply #[[$MAP_U]]()[%[[DIM]]] // CHECK: affine.apply #[[$MAP_U]]()[%[[DIM]]]
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>) %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> { -> tensor<?x?xf32> {
%0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2: tensor<?x?xf32>) inits(%arg2: tensor<?x?xf32>)
-> tensor<?x?xf32> -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }

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

Show All 19 Linesfunc.func @static_sizes_output_divisible(%arg0: tensor<24x12xf32>,
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[T3:.*]] = tensor.pad %[[T0]] nofold // CHECK: %[[T3:.*]] = tensor.pad %[[T0]] nofold
// CHECK: tensor.yield %[[CST]] // CHECK: tensor.yield %[[CST]]
// CHECK: %[[T4:.*]] = tensor.pad %[[T1]] nofold // CHECK: %[[T4:.*]] = tensor.pad %[[T1]] nofold
// CHECK: %[[T5:.*]] = linalg.matmul // CHECK: %[[T5:.*]] = linalg.matmul
// CHECK-SAME: ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>) // CHECK-SAME: ins(%[[T3]], %[[T4]] : tensor<4x7xf32>, tensor<7x5xf32>)
// CHECK-SAME: outs(%[[T2]] : tensor<4x5xf32>) // CHECK-SAME: inits(%[[T2]] : tensor<4x5xf32>)
%4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32> %4 = linalg.matmul ins(%1, %2 : tensor<4x?xf32>, tensor<?x5xf32>) inits(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
%5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32> %5 = tensor.insert_slice %4 into %arg2[%iv0, %iv1] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
func.return %5 : tensor<24x25xf32> func.return %5 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]} %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
} }
// ----- // -----
func.func @pad(%arg0: tensor<24x12xf32>, func.func @pad(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// expected-note @below {{when applied to this op}} // expected-note @below {{when applied to this op}}
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
// expected-error @below {{op expects a padding value of type 'f32', got 0 : i32}} // expected-error @below {{op expects a padding value of type 'f32', got 0 : i32}}
%1 = transform.structured.pad %0 {padding_values=[0: i32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]} %1 = transform.structured.pad %0 {padding_values=[0: i32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
} }
// ----- // -----
func.func @pad(%arg0: tensor<24x12xf32>, func.func @pad(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// expected-note @below {{when applied to this op}} // expected-note @below {{when applied to this op}}
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
// expected-error @below {{expects a padding that parses to 'f32', got "foo"}} // expected-error @below {{expects a padding that parses to 'f32', got "foo"}}
%1 = transform.structured.pad %0 {padding_values=["foo", 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]} %1 = transform.structured.pad %0 {padding_values=["foo", 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
} }
// ----- // -----
func.func @pad(%arg0: tensor<24x12xf32>, func.func @pad(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// This is attached to an error that is silenceable and is not reported by this transform // This is attached to an error that is silenceable and is not reported by this transform
// {{when applied to this op}} // {{when applied to this op}}
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(suppress) { transform.sequence failures(suppress) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
// This error is silenceable and is not reported by this transform // This error is silenceable and is not reported by this transform
// {{transform.structured.pad failed to apply}} // {{transform.structured.pad failed to apply}}
%1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]} %1 = transform.structured.pad %0 {padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32], padding_dimensions=[0, 1, 2], pack_paddings=[1, 1, 0]}
} }

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

// RUN: mlir-opt -test-transform-dialect-interpreter %s | FileCheck %s // RUN: mlir-opt -test-transform-dialect-interpreter %s | FileCheck %s
func.func @scalarize(%arg0: tensor<24x12xf32>, func.func @scalarize(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// The op is first tiled by 10 in the first dimension, which creates a // The op is first tiled by 10 in the first dimension, which creates a
// dynamic size, and then scalarized, which brings the dimension to static 1. // dynamic size, and then scalarized, which brings the dimension to static 1.
// CHECK: linalg.matmul ins(%{{.*}}, %{{.*}} : tensor<1x12 // CHECK: linalg.matmul ins(%{{.*}}, %{{.*}} : tensor<1x12
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1, %loops = transform.structured.tile %0 [10, 0, 0] %1, %loops = transform.structured.tile %0 [10, 0, 0]
%2 = transform.structured.scalarize %1 %2 = transform.structured.scalarize %1
} }

mlir/test/Dialect/Linalg/transform-op-split-reduction-by-scaling.mlir

// RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s // RUN: mlir-opt --test-transform-dialect-interpreter %s | FileCheck %s
// CHECK-LABEL: func.func @matmul_split // CHECK-LABEL: func.func @matmul_split
func.func @matmul_split(%A : tensor<?x256xf32>, %B: tensor<256x32xf32>, %C: tensor<?x32xf32>) -> tensor<?x32xf32> { func.func @matmul_split(%A : tensor<?x256xf32>, %B: tensor<256x32xf32>, %C: tensor<?x32xf32>) -> tensor<?x32xf32> {
// CHECK: bufferization.alloc_tensor({{.*}}) : tensor<?x32x64xf32> // CHECK: bufferization.alloc_tensor({{.*}}) : tensor<?x32x64xf32>
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]
// CHECK-SAME: ins(%{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}} : tensor<?x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>) // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}}, %{{[a-zA-Z0-9]*}} : tensor<?x256xf32>, tensor<256x32xf32>, tensor<64x4xi1>)
// CHECK-SAME: outs(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>) { // CHECK-SAME: inits(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>) {
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"] // CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
// CHECK-SAME: ins(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>) // CHECK-SAME: ins(%{{[a-zA-Z0-9]*}} : tensor<?x32x64xf32>)
// CHECK-SAME: outs(%{{[a-zA-Z0-9]*}} : tensor<?x32xf32>) { // CHECK-SAME: inits(%{{[a-zA-Z0-9]*}} : tensor<?x32xf32>) {
%0 = linalg.matmul ins(%A, %B: tensor<?x256xf32>, tensor<256x32xf32>) %0 = linalg.matmul ins(%A, %B: tensor<?x256xf32>, tensor<256x32xf32>)
outs(%C: tensor<?x32xf32>) -> tensor<?x32xf32> inits(%C: tensor<?x32xf32>) -> tensor<?x32xf32>
return %0: tensor<?x32xf32> return %0: tensor<?x32xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:4 = transform.structured.split_reduction %0 %1:4 = transform.structured.split_reduction %0
{ split_factor = 4, insert_split_dimension = 2, use_scaling_algorithm, use_alloc} { split_factor = 4, insert_split_dimension = 2, use_scaling_algorithm, use_alloc}
} }

mlir/test/Dialect/Linalg/transform-op-split-reduction.mlir

// RUN: mlir-opt --split-input-file --test-transform-dialect-interpreter %s | FileCheck %s // RUN: mlir-opt --split-input-file --test-transform-dialect-interpreter %s | FileCheck %s
func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>) %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// 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:.*]] = tensor.empty() : 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) inits(%[[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>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}
} }
// ----- // -----
func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> { func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> {
%red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> ()>, affine_map<(d0) -> ()>,
affine_map<(d0) -> ()>], affine_map<(d0) -> ()>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>) ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>)
outs(%out : tensor<f32>) { inits(%out : tensor<f32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
%40 = arith.subf %arg7, %arg8 : f32 %40 = arith.subf %arg7, %arg8 : f32
%41 = math.exp %40 : f32 %41 = math.exp %40 : f32
%42 = arith.mulf %41, %arg9 : f32 %42 = arith.mulf %41, %arg9 : f32
linalg.yield %42 : f32 linalg.yield %42 : f32
} -> tensor<f32> } -> tensor<f32>
return %red : tensor<f32> return %red : tensor<f32>
} }
// 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-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f32 // CHECK-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f32
// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<4x8xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<4x8xf32>
// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32> // CHECK-DAG: %[[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) inits(%[[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>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0}
} }
// ----- // -----
func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>) func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)
-> tensor<5x2xf32> -> tensor<5x2xf32>
{ {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2) -> (d1, d0)>, affine_map<(d0, d1, d2) -> (d1, d0)>,
affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d2, d1)>,
affine_map<(d0, d1, d2) -> (d2, d0)> affine_map<(d0, d1, d2) -> (d2, d0)>
], ],
iterator_types = ["parallel", "reduction", "parallel"] iterator_types = ["parallel", "reduction", "parallel"]
} ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) { } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) inits(%output : tensor<5x2xf32>) {
^bb0(%arg0: f32, %arg1: f32, %arg2: f32): ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
%3 = arith.addf %arg0, %arg1 : f32 %3 = arith.addf %arg0, %arg1 : f32
%4 = arith.maxf %3, %arg2 : f32 %4 = arith.maxf %3, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<5x2xf32> } -> tensor<5x2xf32>
return %0 : tensor<5x2xf32> return %0 : tensor<5x2xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d0)>
// 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-DAG: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32 // CHECK-DAG: %[[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-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32> // CHECK-DAG: %[[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) inits(%[[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>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2}
} }
// ----- // -----
func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> { func.func @matmul_split(%A : tensor<16x256xf32>, %B: tensor<256x32xf32>, %C: tensor<16x32xf32>) -> tensor<16x32xf32> {
%0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>) %0 = linalg.matmul ins(%A, %B: tensor<16x256xf32>, tensor<256x32xf32>)
outs(%C: tensor<16x32xf32>) -> tensor<16x32xf32> inits(%C: tensor<16x32xf32>) -> tensor<16x32xf32>
return %0: tensor<16x32xf32> return %0: tensor<16x32xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
// 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, d3)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// 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<16x64x4xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<16x256xf32> into tensor<16x64x4xf32>
// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<64x4x32xf32> // CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<256x32xf32> into tensor<64x4x32xf32>
// CHECK-DAG: %[[INI:.*]] = tensor.empty() : 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) inits(%[[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", "reduction", "parallel"]} // CHECK-SAME: , iterator_types = ["parallel", "parallel", "reduction", "parallel"]}
// CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) outs(%[[F]] : tensor<16x32x4xf32>) { // CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<16x64x4xf32>, tensor<64x4x32xf32>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}
} }
// ----- // -----
func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> { func.func @generic_split_1d(%arg0: tensor<32xf32>, %arg1: tensor<f32>, %out: tensor<f32>) -> tensor<f32> {
%red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, %red = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> ()>, affine_map<(d0) -> ()>,
affine_map<(d0) -> ()>], affine_map<(d0) -> ()>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>) ins(%arg0, %arg1 : tensor<32xf32>, tensor<f32>)
outs(%out : tensor<f32>) { inits(%out : tensor<f32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
%40 = arith.subf %arg7, %arg8 : f32 %40 = arith.subf %arg7, %arg8 : f32
%41 = math.exp %40 : f32 %41 = math.exp %40 : f32
%42 = arith.mulf %41, %arg9 : f32 %42 = arith.mulf %41, %arg9 : f32
linalg.yield %42 : f32 linalg.yield %42 : f32
} -> tensor<f32> } -> tensor<f32>
return %red : tensor<f32> return %red : tensor<f32>
} }
// 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) -> (d1)> // CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d1)>
// 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-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f32 // CHECK-DAG: %[[ID:.*]] = arith.constant 1.000000e+00 : f32
// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<8x4xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1]] : tensor<32xf32> into tensor<8x4xf32>
// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<4xf32> // CHECK-DAG: %[[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) inits(%[[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 = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) outs(%[[F]] : tensor<4xf32>) { // CHECK: iterator_types = ["reduction", "parallel"]} ins(%[[I1]], %{{.*}} : tensor<8x4xf32>, tensor<f32>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0, inner_parallel} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 0, inner_parallel}
} }
// ----- // -----
func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>) func.func @generic_split_3d(%input: tensor<32x2xf32>, %input_2: tensor<5x32xf32>, %output: tensor<5x2xf32>)
-> tensor<5x2xf32> -> tensor<5x2xf32>
{ {
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2) -> (d1, d0)>, affine_map<(d0, d1, d2) -> (d1, d0)>,
affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d2, d1)>,
affine_map<(d0, d1, d2) -> (d2, d0)> affine_map<(d0, d1, d2) -> (d2, d0)>
], ],
iterator_types = ["parallel", "reduction", "parallel"] iterator_types = ["parallel", "reduction", "parallel"]
} ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) outs(%output : tensor<5x2xf32>) { } ins(%input, %input_2 : tensor<32x2xf32>, tensor<5x32xf32>) inits(%output : tensor<5x2xf32>) {
^bb0(%arg0: f32, %arg1: f32, %arg2: f32): ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
%3 = arith.addf %arg0, %arg1 : f32 %3 = arith.addf %arg0, %arg1 : f32
%4 = arith.maxf %3, %arg2 : f32 %4 = arith.maxf %3, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<5x2xf32> } -> tensor<5x2xf32>
return %0 : tensor<5x2xf32> return %0 : tensor<5x2xf32>
} }
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)> // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d0)>
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)> // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d3, d1, d2)>
// 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-DAG: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32 // CHECK-DAG: %[[ID:.*]] = arith.constant -3.40282347E+38 : f32
// CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32> // CHECK-DAG: %[[I1:.*]] = tensor.expand_shape %{{.*}}[0, 1], [2]] : tensor<32x2xf32> into tensor<8x4x2xf32>
// CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32> // CHECK-DAG: %[[I2:.*]] = tensor.expand_shape %{{.*}}[0], [1, 2]] : tensor<5x32xf32> into tensor<5x8x4xf32>
// CHECK-DAG: %[[INI:.*]] = tensor.empty() : tensor<5x2x4xf32> // CHECK-DAG: %[[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) inits(%[[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<8x4x2xf32>, tensor<5x8x4xf32>) outs(%[[F]] : tensor<5x2x4xf32>) { // CHECK-SAME: ins(%[[I1]], %[[I2]] : tensor<8x4x2xf32>, tensor<5x8x4xf32>) inits(%[[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>) inits(%{{.*}} : 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>
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel} %1:4 = transform.structured.split_reduction %0 { split_factor = 4, insert_split_dimension = 2, inner_parallel}
} }

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

Show All 11 Lines
// CHECK-LABEL: @one_d_static // CHECK-LABEL: @one_d_static
// CHECK-SAME: %[[IN:.+]]: tensor<100xf32>, %[[OUT:.+]]: tensor<100xf32> // CHECK-SAME: %[[IN:.+]]: tensor<100xf32>, %[[OUT:.+]]: tensor<100xf32>
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> { func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
// CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN]][0] [42] [1] : tensor<100xf32> to tensor<42xf32> // CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN]][0] [42] [1] : tensor<100xf32> to tensor<42xf32>
// CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT]][0] [42] [1] : tensor<100xf32> to tensor<42xf32> // CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT]][0] [42] [1] : tensor<100xf32> to tensor<42xf32>
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
// CHECK: ins(%[[IN_SLICE_LOW]] // CHECK: ins(%[[IN_SLICE_LOW]]
// CHECK: outs(%[[OUT_SLICE_LOW]] // CHECK: inits(%[[OUT_SLICE_LOW]]
// CHECK: linalg.index 0 // CHECK: linalg.index 0
// CHECK: func.call @elem // CHECK: func.call @elem
// CHECK: %[[RES_PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [42] [1] // CHECK: %[[RES_PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [42] [1]
// //
// CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN]][42] [58] [1] : tensor<100xf32> to tensor<58xf32> // CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN]][42] [58] [1] : tensor<100xf32> to tensor<58xf32>
// CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[RES_PARTIAL]][42] [58] [1] : tensor<100xf32> to tensor<58xf32> // CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[RES_PARTIAL]][42] [58] [1] : tensor<100xf32> to tensor<58xf32>
// CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic // CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
// CHECK: ins(%[[IN_SLICE_HIGH]] // CHECK: ins(%[[IN_SLICE_HIGH]]
// CHECK: outs(%[[OUT_SLICE_HIGH]] // CHECK: inits(%[[OUT_SLICE_HIGH]]
// CHECK: %[[IDX:.+]] = linalg.index 0 // CHECK: %[[IDX:.+]] = linalg.index 0
// CHECK: affine.apply #[[$ADD_42_MAP]](%[[IDX]]) // CHECK: affine.apply #[[$ADD_42_MAP]](%[[IDX]])
// CHECK: func.call @elem // CHECK: func.call @elem
// CHECK: %[[RES:.+]] = tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[RES_PARTIAL]][42] [58] [1] // CHECK: %[[RES:.+]] = tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[RES_PARTIAL]][42] [58] [1]
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) { ins(%arg0: tensor<100xf32>) inits(%arg1: tensor<100xf32>) {
^bb0(%0: f32, %1: f32): ^bb0(%0: f32, %1: f32):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32 %call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32
linalg.yield %call_res : f32 linalg.yield %call_res : f32
} -> tensor<100xf32> } -> tensor<100xf32>
// CHECK: return %[[RES]] // CHECK: return %[[RES]]
return %0 : tensor<100xf32> return %0 : tensor<100xf32>
} }
// CHECK-LABEL: @one_d_static_overflow // CHECK-LABEL: @one_d_static_overflow
// CHECK-SAME: %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32> // CHECK-SAME: %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> { func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
// Folding is sufficiently powerful to detect the static overflow and avoid // Folding is sufficiently powerful to detect the static overflow and avoid
// the splitting altogether. // the splitting altogether.
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
// CHECK: ins(%[[IN]] // CHECK: ins(%[[IN]]
// CHECK: outs(%[[OUT]] // CHECK: inits(%[[OUT]]
// CHECK: linalg.index 0 // CHECK: linalg.index 0
// CHECK: func.call @elem // CHECK: func.call @elem
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<10xf32>) outs(%arg1: tensor<10xf32>) { ins(%arg0: tensor<10xf32>) inits(%arg1: tensor<10xf32>) {
^bb0(%0: f32, %1: f32): ^bb0(%0: f32, %1: f32):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32 %call_res = func.call @elem(%0, %i, %i) : (f32, index, index) -> f32
linalg.yield %call_res : f32 linalg.yield %call_res : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
Show All 15 Lines
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> { func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
// CHECK: %[[SPLIT:.+]] = call @get_size // CHECK: %[[SPLIT:.+]] = call @get_size
// CHECK: %[[SPLIT_LOW:.+]] = affine.min #[[$MAP_MIN_100]]()[%[[SPLIT]] // CHECK: %[[SPLIT_LOW:.+]] = affine.min #[[$MAP_MIN_100]]()[%[[SPLIT]]
// CHECK: %[[SPLIT_HIGH_1:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]] // CHECK: %[[SPLIT_HIGH_1:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
// CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32> // CHECK: %[[IN_SLICE_LOW:.+]] = tensor.extract_slice %[[IN:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
// CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32> // CHECK: %[[OUT_SLICE_LOW:.+]] = tensor.extract_slice %[[OUT:.+]][0] [%[[SPLIT_LOW]]] [1] : tensor<100xf32> to tensor<?xf32>
// CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic // CHECK: %[[RES_SLICE_LOW:.+]] = linalg.generic
// CHECK: ins(%[[IN_SLICE_LOW]] // CHECK: ins(%[[IN_SLICE_LOW]]
// CHECK: outs(%[[OUT_SLICE_LOW]] // CHECK: inits(%[[OUT_SLICE_LOW]]
// CHECK: %[[PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [%[[SPLIT_LOW]]] [1] // CHECK: %[[PARTIAL:.+]] = tensor.insert_slice %[[RES_SLICE_LOW]] into %[[OUT]][0] [%[[SPLIT_LOW]]] [1]
// //
// CHECK: %[[SPLIT_HIGH_2:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]] // CHECK: %[[SPLIT_HIGH_2:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
// CHECK: %[[SPLIT_HIGH_3:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]] // CHECK: %[[SPLIT_HIGH_3:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
// CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1] : tensor<100xf32> to tensor<?xf32> // CHECK: %[[IN_SLICE_HIGH:.+]] = tensor.extract_slice %[[IN:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_2]]] [1] : tensor<100xf32> to tensor<?xf32>
// CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[PARTIAL:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1] : tensor<100xf32> to tensor<?xf32> // CHECK: %[[OUT_SLICE_HIGH:.+]] = tensor.extract_slice %[[PARTIAL:.+]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_3]]] [1] : tensor<100xf32> to tensor<?xf32>
// CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic // CHECK: %[[RES_SLICE_HIGH:.+]] = linalg.generic
// CHECK: ins(%[[IN_SLICE_HIGH]] // CHECK: ins(%[[IN_SLICE_HIGH]]
// CHECK: outs(%[[OUT_SLICE_HIGH]] // CHECK: inits(%[[OUT_SLICE_HIGH]]
// CHECK: %[[SPLIT_HIGH_4:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]] // CHECK: %[[SPLIT_HIGH_4:.+]] = affine.apply #[[$MAP_S_MINUS_100]]()[%[[SPLIT_LOW]]]
// CHECK: tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[PARTIAL]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_4]]] [1] // CHECK: tensor.insert_slice %[[RES_SLICE_HIGH]] into %[[PARTIAL]][%[[SPLIT_LOW]]] [%[[SPLIT_HIGH_4]]] [1]
%0 = func.call @get_size() : () -> index %0 = func.call @get_size() : () -> index
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) { ins(%arg0: tensor<100xf32>) inits(%arg1: tensor<100xf32>) {
^bb0(%3: f32, %4: f32): ^bb0(%3: f32, %4: f32):
%5 = arith.addf %3, %4 : f32 %5 = arith.addf %3, %4 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<100xf32> } -> tensor<100xf32>
return %1 : tensor<100xf32> return %1 : tensor<100xf32>
} }
// ----- // -----
Show All 11 Lines
func.func @two_d(%arg0: tensor<10x34xf32>, func.func @two_d(%arg0: tensor<10x34xf32>,
%arg1: tensor<10x34xf32>) -> tensor<10x34xf32> { %arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
// Check the overall structure: split along the dimension 0, and then split // Check the overall structure: split along the dimension 0, and then split
// the second half only along the dimension 1. // the second half only along the dimension 1.
// CHECK: %[[IN_1:.+]] = tensor.extract_slice %[[IN:.+]][0, 0] // CHECK: %[[IN_1:.+]] = tensor.extract_slice %[[IN:.+]][0, 0]
// CHECK: %[[OUT_1:.+]] = tensor.extract_slice %[[OUT:.+]][0, 0] // CHECK: %[[OUT_1:.+]] = tensor.extract_slice %[[OUT:.+]][0, 0]
// CHECK: %[[RES_1:.+]] = linalg.generic // CHECK: %[[RES_1:.+]] = linalg.generic
// CHECK-SAME: ins(%[[IN_1]] : tensor<4x34xf32>) // CHECK-SAME: ins(%[[IN_1]] : tensor<4x34xf32>)
// CHECK-SAME: outs(%[[OUT_1]] : tensor<4x34xf32>) // CHECK-SAME: inits(%[[OUT_1]] : tensor<4x34xf32>)
// CHECK: %[[PARTIAL_1:.+]] = tensor.insert_slice %[[RES_1]] into %[[OUT]] // CHECK: %[[PARTIAL_1:.+]] = tensor.insert_slice %[[RES_1]] into %[[OUT]]
// //
// CHECK: %[[IN_2:.+]] = tensor.extract_slice %[[IN]] // CHECK: %[[IN_2:.+]] = tensor.extract_slice %[[IN]]
// CHECK: %[[OUT_2:.+]] = tensor.extract_slice %[[PARTIAL_1]] // CHECK: %[[OUT_2:.+]] = tensor.extract_slice %[[PARTIAL_1]]
// Note that `extract_slice` taking a slice from another `extract_slice` result // Note that `extract_slice` taking a slice from another `extract_slice` result
// is folded to use the operand of the first `extract_slice`. // is folded to use the operand of the first `extract_slice`.
// CHECK: %[[IN_21:.+]] = tensor.extract_slice %[[IN_2]] // CHECK: %[[IN_21:.+]] = tensor.extract_slice %[[IN_2]]
// CHECK: %[[OUT_21:.+]] = tensor.extract_slice %[[OUT_2]] // CHECK: %[[OUT_21:.+]] = tensor.extract_slice %[[OUT_2]]
// CHECK: %[[RES_21:.+]] = linalg.generic // CHECK: %[[RES_21:.+]] = linalg.generic
// CHECK-SAME: ins(%[[IN_21]] : tensor<6x16xf32>) // CHECK-SAME: ins(%[[IN_21]] : tensor<6x16xf32>)
// CHECK-SAME: outs(%[[OUT_21]] : tensor<6x16xf32>) // CHECK-SAME: inits(%[[OUT_21]] : tensor<6x16xf32>)
// CHECK: %[[PARTIAL_21:.+]] = tensor.insert_slice %[[RES_21]] into %[[OUT_2]] // CHECK: %[[PARTIAL_21:.+]] = tensor.insert_slice %[[RES_21]] into %[[OUT_2]]
// //
// CHECK: %[[IN_22:.+]] = tensor.extract_slice %[[IN_2]] // CHECK: %[[IN_22:.+]] = tensor.extract_slice %[[IN_2]]
// CHECK: %[[OUT_22:.+]] = tensor.extract_slice %[[PARTIAL_21]] // CHECK: %[[OUT_22:.+]] = tensor.extract_slice %[[PARTIAL_21]]
// CHECK: %[[RES_22:.+]] = linalg.generic // CHECK: %[[RES_22:.+]] = linalg.generic
// CHECK-SAME: ins(%[[IN_22]] : tensor<6x18xf32>) // CHECK-SAME: ins(%[[IN_22]] : tensor<6x18xf32>)
// CHECK-SAME: outs(%[[OUT_22]] : tensor<6x18xf32>) // CHECK-SAME: inits(%[[OUT_22]] : tensor<6x18xf32>)
// CHECK: %[[PARTIAL_22:.+]] = tensor.insert_slice %[[RES_22]] into %[[PARTIAL_21]] // CHECK: %[[PARTIAL_22:.+]] = tensor.insert_slice %[[RES_22]] into %[[PARTIAL_21]]
// CHECK: %[[PARTIAL_2:.+]] = tensor.insert_slice %[[PARTIAL_22]] into %[[PARTIAL_1]] // CHECK: %[[PARTIAL_2:.+]] = tensor.insert_slice %[[PARTIAL_22]] into %[[PARTIAL_1]]
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i, j) -> (i, j)>, indexing_maps = [affine_map<(i, j) -> (i, j)>,
affine_map<(i, j) -> (i, j)>], affine_map<(i, j) -> (i, j)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%arg0: tensor<10x34xf32>) ins(%arg0: tensor<10x34xf32>)
outs(%arg1: tensor<10x34xf32>) { inits(%arg1: tensor<10x34xf32>) {
^bb0(%0: f32, %1: f32): ^bb0(%0: f32, %1: f32):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32 %call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32
linalg.yield %call_res : f32 linalg.yield %call_res : f32
} -> tensor<10x34xf32> } -> tensor<10x34xf32>
return %0 : tensor<10x34xf32> return %0 : tensor<10x34xf32>
} }
Show All 20 Lines
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> { func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
// expected-note @below {{dynamic split point}} // expected-note @below {{dynamic split point}}
%0 = func.call @get_size() : () -> i64 %0 = func.call @get_size() : () -> i64
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) { ins(%arg0: tensor<100xf32>) inits(%arg1: tensor<100xf32>) {
^bb0(%3: f32, %4: f32): ^bb0(%3: f32, %4: f32):
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<100xf32> } -> tensor<100xf32>
return %1 : tensor<100xf32> return %1 : tensor<100xf32>
} }
// ----- // -----
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1 = transform.structured.match ops{["func.call"]} in %arg1 %1 = transform.structured.match ops{["func.call"]} in %arg1
// expected-error @below {{expected the dynamic split point handle to point to as many operations (0) as the target handle (1)}} // expected-error @below {{expected the dynamic split point handle to point to as many operations (0) as the target handle (1)}}
transform.structured.split %0 after %1 { dimension = 0 } transform.structured.split %0 after %1 { dimension = 0 }
} }
func.func private @get_size() -> i64 func.func private @get_size() -> i64
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> { func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
%1 = linalg.generic { %1 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) { ins(%arg0: tensor<100xf32>) inits(%arg1: tensor<100xf32>) {
^bb0(%3: f32, %4: f32): ^bb0(%3: f32, %4: f32):
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<100xf32> } -> tensor<100xf32>
return %1 : tensor<100xf32> return %1 : tensor<100xf32>
} }
// ----- // -----
Show All 19 Lines
} }
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> { func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
// expected-note @below {{target op}} // expected-note @below {{target op}}
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>], indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
iterator_types = ["parallel"] iterator_types = ["parallel"]
} }
ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) { ins(%arg0: tensor<100xf32>) inits(%arg1: tensor<100xf32>) {
^bb0(%0: f32, %1: f32): ^bb0(%0: f32, %1: f32):
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<100xf32> } -> tensor<100xf32>
return %0 : tensor<100xf32> return %0 : tensor<100xf32>
} }

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

Show All 15 Lines %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) {
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) {
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) {
// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32> // CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32> // CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32> // CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<4x4xf32>
// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<4x4xf32>, tensor<4x4xf32>) // CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<4x4xf32>, tensor<4x4xf32>)
// CHECK-SAME: outs(%[[sTC]] : tensor<4x4xf32>) -> tensor<4x4xf32> // CHECK-SAME: inits(%[[sTC]] : tensor<4x4xf32>) -> tensor<4x4xf32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<4x4xf32> into tensor<128x128xf32> // CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<4x4xf32> into tensor<128x128xf32>
// CHECK: scf.yield %[[TD]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD]] : tensor<128x128xf32>
// CHECK: scf.yield %[[TD2]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD2]] : tensor<128x128xf32>
// CHECK: scf.yield %[[TD1]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD1]] : tensor<128x128xf32>
%0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
// CHECK: return %[[TD0]] : tensor<128x128xf32> // CHECK: return %[[TD0]] : tensor<128x128xf32>
return %0 : tensor<128x128xf32> return %0 : tensor<128x128xf32>
} }
// ----- // -----
Show All 16 Lines %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<128x128xf32>) {
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<128x128xf32>) {
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) { // CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<128x128xf32>) {
// CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<?x4xf32> // CHECK: %[[sTA:.*]] = tensor.extract_slice %[[TA]][{{.*}}] : tensor<128x128xf32> to tensor<?x4xf32>
// CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x?xf32> // CHECK: %[[sTB:.*]] = tensor.extract_slice %[[TB]][{{.*}}] : tensor<128x128xf32> to tensor<4x?xf32>
// CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<?x?xf32> // CHECK: %[[sTC:.*]] = tensor.extract_slice %[[TC2]][{{.*}}] : tensor<128x128xf32> to tensor<?x?xf32>
// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x4xf32>, tensor<4x?xf32>) // CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x4xf32>, tensor<4x?xf32>)
// CHECK-SAME: outs(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32> // CHECK-SAME: inits(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
// CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<128x128xf32> // CHECK: %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<128x128xf32>
// CHECK: scf.yield %[[TD]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD]] : tensor<128x128xf32>
// CHECK: scf.yield %[[TD2]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD2]] : tensor<128x128xf32>
// CHECK: scf.yield %[[TD1]] : tensor<128x128xf32> // CHECK: scf.yield %[[TD1]] : tensor<128x128xf32>
%sz = func.call @get_dynamic_tile_size() : () -> index %sz = func.call @get_dynamic_tile_size() : () -> index
%0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
// CHECK: return %[[TD0]] : tensor<128x128xf32> // CHECK: return %[[TD0]] : tensor<128x128xf32>
return %0 : tensor<128x128xf32> return %0 : tensor<128x128xf32>
} }

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

// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file -verify-diagnostics | FileCheck %s
// CHECK-LABEL: @vectorize_matmul // CHECK-LABEL: @vectorize_matmul
// CHECK-SAME: %[[A:.*]]: tensor<24x12xf32> // CHECK-SAME: %[[A:.*]]: tensor<24x12xf32>
// CHECK-SAME: %[[B:.*]]: tensor<12x25xf32> // CHECK-SAME: %[[B:.*]]: tensor<12x25xf32>
// CHECK-SAME: %[[C:.*]]: tensor<24x25xf32> // CHECK-SAME: %[[C:.*]]: tensor<24x25xf32>
func.func @vectorize_matmul(%arg0: tensor<24x12xf32>, func.func @vectorize_matmul(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// CHECK: %[[vA:.+]] = vector.transfer_read %[[A]] // CHECK: %[[vA:.+]] = vector.transfer_read %[[A]]
// CHECK: %[[vB:.+]] = vector.transfer_read %[[B]] // CHECK: %[[vB:.+]] = vector.transfer_read %[[B]]
// CHECK: %[[vC:.+]] = vector.transfer_read %[[C]] // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
// CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]] // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
// CHECK: vector.transfer_write %[[vR]], %[[C]] // CHECK: vector.transfer_write %[[vR]], %[[C]]
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
Show All 28 Linesfunc.func @vectorize_keep_pad(
^bb0(%arg6: index, %arg7: index): ^bb0(%arg6: index, %arg7: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<?x5xf32> to tensor<7x5xf32> } : tensor<?x5xf32> to tensor<7x5xf32>
// CHECK: %[[vA:.+]] = vector.transfer_read %[[pA]] // CHECK: %[[vA:.+]] = vector.transfer_read %[[pA]]
// CHECK: %[[vB:.+]] = vector.transfer_read %[[pB]] // CHECK: %[[vB:.+]] = vector.transfer_read %[[pB]]
// CHECK: %[[vC:.+]] = vector.transfer_read %[[C]] // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
// CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]] // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
// CHECK: vector.transfer_write %[[vR]], %[[C]] // CHECK: vector.transfer_write %[[vR]], %[[C]]
%8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32> %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) inits(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
%9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32> %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
return %9 : tensor<24x25xf32> return %9 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
Show All 31 Linesfunc.func @vectorize_pad(
// CHECK: %[[vB:.+]] = vector.transfer_read %[[sB]] // CHECK: %[[vB:.+]] = vector.transfer_read %[[sB]]
%7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] { %7 = tensor.pad %2 nofold low[%c0, %c0] high[%6, %c0] {
^bb0(%arg6: index, %arg7: index): ^bb0(%arg6: index, %arg7: index):
tensor.yield %cst : f32 tensor.yield %cst : f32
} : tensor<?x5xf32> to tensor<7x5xf32> } : tensor<?x5xf32> to tensor<7x5xf32>
// CHECK: %[[vC:.+]] = vector.transfer_read %[[C]] // CHECK: %[[vC:.+]] = vector.transfer_read %[[C]]
// CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]] // CHECK: %[[vR:.+]] = vector.contract {{.*}} %[[vA]], %[[vB]], %[[vC]]
// CHECK: vector.transfer_write %[[vR]], %[[C]] // CHECK: vector.transfer_write %[[vR]], %[[C]]
%8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) outs(%3 : tensor<4x5xf32>) -> tensor<4x5xf32> %8 = linalg.matmul ins(%5, %7 : tensor<4x7xf32>, tensor<7x5xf32>) inits(%3 : tensor<4x5xf32>) -> tensor<4x5xf32>
%9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32> %9 = tensor.insert_slice %8 into %arg2[%arg3, %arg4] [4, 5] [1, 1] : tensor<4x5xf32> into tensor<24x25xf32>
return %9 : tensor<24x25xf32> return %9 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 {vectorize_padding} %2 = transform.structured.vectorize %1 {vectorize_padding}
} }
// ----- // -----
func.func @vectorize(%arg0: tensor<24x12xf32>, func.func @vectorize(%arg0: tensor<24x12xf32>,
%arg1: tensor<12x25xf32>, %arg1: tensor<12x25xf32>,
%arg2: tensor<24x25xf32>) -> tensor<24x25xf32> { %arg2: tensor<24x25xf32>) -> tensor<24x25xf32> {
// expected-note @below {{non-isolated target}} // expected-note @below {{non-isolated target}}
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) outs(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32> %0 = linalg.matmul ins(%arg0, %arg1 : tensor<24x12xf32>, tensor<12x25xf32>) inits(%arg2 : tensor<24x25xf32>) -> tensor<24x25xf32>
func.return %0 : tensor<24x25xf32> func.return %0 : tensor<24x25xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
// expected-error @below {{op requires isolated-from-above targets}} // expected-error @below {{op requires isolated-from-above targets}}
%2 = transform.structured.vectorize %0 %2 = transform.structured.vectorize %0
} }

mlir/test/Dialect/Linalg/transform-patterns.mlir

// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s
// ----- // -----
func.func @dot(%x: memref<?xf32, strided<[1], offset: ?>>, func.func @dot(%x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>, %y: memref<?xf32, strided<[1], offset: ?>>,
%v: memref<f32>) { %v: memref<f32>) {
linalg.dot ins(%x, %y: memref<?xf32, strided<[1], offset: ?>>, linalg.dot ins(%x, %y: memref<?xf32, strided<[1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%v: memref<f32>) inits(%v: memref<f32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.dot"]} in %arg1 %0 = transform.structured.match ops{["linalg.dot"]} in %arg1
%1, %loop = transform.structured.tile %0 [8000] %1, %loop = transform.structured.tile %0 [8000]
} }
// CHECK-LABEL: func @dot // CHECK-LABEL: func @dot
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c8000:.*]] = arith.constant 8000 : index // CHECK-DAG: %[[c8000:.*]] = arith.constant 8000 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] {
// CHECK: linalg.dot // CHECK: linalg.dot
// ----- // -----
func.func @matvec(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @matvec(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%x: memref<?xf32, strided<[1], offset: ?>>, %x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>) { %y: memref<?xf32, strided<[1], offset: ?>>) {
linalg.matvec linalg.matvec
ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>, ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%y: memref<?xf32, strided<[1], offset: ?>>) inits(%y: memref<?xf32, strided<[1], offset: ?>>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1
%1, %loops:2 = transform.structured.tile %0 [5, 6] %1, %loops:2 = transform.structured.tile %0 [5, 6]
} }
// CHECK-LABEL: func @matvec // CHECK-LABEL: func @matvec
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index // CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index // CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
// CHECK: scf.for {{.*}} step %[[c5]] // CHECK: scf.for {{.*}} step %[[c5]]
// CHECK: scf.for {{.*}} step %[[c6]] // CHECK: scf.for {{.*}} step %[[c6]]
// CHECK: linalg.matvec // CHECK: linalg.matvec
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>) // CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?xf32, strided<[1], offset: ?>>) // CHECK: inits({{.*}}: memref<?xf32, strided<[1], offset: ?>>)
// ----- // -----
func.func @matmul(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @matmul(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) { %C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>) memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) inits(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000]
%2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400]
Show All 24 Lines
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c2]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c3]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c4]] {
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK: inits({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>)
// ----- // -----
// Map corresponding to a 2D memory access where the stride along the last dim is known to be 1. // Map corresponding to a 2D memory access where the stride along the last dim is known to be 1.
// CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)> // CHECK-DAG: #[[$kn:.*]] = affine_map<(d0, d1, d2) -> (d2, d1)>
// CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)> // CHECK-DAG: #[[$nm:.*]] = affine_map<(d0, d1, d2) -> (d1, d0)>
// CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)> // CHECK-DAG: #[[$km:.*]] = affine_map<(d0, d1, d2) -> (d2, d0)>
Show All 10 Lines#generic_matmul_trait = {
iterator_types = ["parallel", "parallel", "reduction"] iterator_types = ["parallel", "parallel", "reduction"]
} }
func.func @permute_generic(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @permute_generic(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) { %C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.generic #generic_matmul_trait linalg.generic #generic_matmul_trait
ins(%A, %B : memref<?x?xf32, strided<[?, 1], offset: ?>>, ins(%A, %B : memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>) memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>) { inits(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>) {
^bb(%a: f32, %b: f32, %c: f32): ^bb(%a: f32, %b: f32, %c: f32):
%d = arith.mulf %a, %b: f32 %d = arith.mulf %a, %b: f32
%e = arith.addf %c, %d: f32 %e = arith.addf %c, %d: f32
linalg.yield %e: f32 linalg.yield %e: f32
} }
return return
} }
Show All 14 Lines
// ----- // -----
func.func @matvec_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @matvec_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%x: memref<?xf32, strided<[1], offset: ?>>, %x: memref<?xf32, strided<[1], offset: ?>>,
%y: memref<?xf32, strided<[1], offset: ?>>) { %y: memref<?xf32, strided<[1], offset: ?>>) {
linalg.matvec ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>, linalg.matvec ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?xf32, strided<[1], offset: ?>>) memref<?xf32, strided<[1], offset: ?>>)
outs(%y: memref<?xf32, strided<[1], offset: ?>>) inits(%y: memref<?xf32, strided<[1], offset: ?>>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1
%1, %loops:2 = transform.structured.tile %0 [5, 6] {interchange = [1, 0]} %1, %loops:2 = transform.structured.tile %0 [5, 6] {interchange = [1, 0]}
} }
// CHECK-LABEL: func @matvec_perm // CHECK-LABEL: func @matvec_perm
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index // CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index
// CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index // CHECK-DAG: %[[c6:.*]] = arith.constant 6 : index
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c6]] // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c6]]
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c5]] // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c5]]
// CHECK: linalg.matvec // CHECK: linalg.matvec
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>) // CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?xf32, strided<[1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?xf32, strided<[1], offset: ?>>) // CHECK: inits({{.*}}: memref<?xf32, strided<[1], offset: ?>>)
// ----- // -----
func.func @matmul_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>, func.func @matmul_perm(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%B: memref<?x?xf32, strided<[?, 1], offset: ?>>, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
%C: memref<?x?xf32, strided<[?, 1], offset: ?>>) { %C: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>, linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
memref<?x?xf32, strided<[?, 1], offset: ?>>) memref<?x?xf32, strided<[?, 1], offset: ?>>)
outs(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>) inits(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] {interchange=[1, 2, 0]} %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] {interchange=[1, 2, 0]}
%2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] {interchange=[1, 0, 2]} %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] {interchange=[1, 0, 2]}
Show All 17 Lines
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c300]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c200]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c400]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c20]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c30]] {
// CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] { // CHECK: scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c40]] {
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK: ins({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>, memref<?x?xf32, strided<[?, 1], offset: ?>>)
// CHECK: outs({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>) // CHECK: inits({{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>)

mlir/test/Dialect/Linalg/transform-promotion.mlir

Show All 16 Lines scf.for %arg4 = %c0 to %2 step %c3000 {
%3 = memref.subview %arg0[%arg3, %arg5][%c2000, %c4000][%c1, %c1] : %3 = memref.subview %arg0[%arg3, %arg5][%c2000, %c4000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%4 = memref.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] : %4 = memref.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%5 = memref.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] : %5 = memref.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul ins(%3, %4: memref<?x?xf32, strided<[?, ?], offset: ?>>, linalg.matmul ins(%3, %4: memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%5: memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%5: memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return return
} }
// CHECK-LABEL: func @promote_subview_matmul // CHECK-LABEL: func @promote_subview_matmul
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index // CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
Show All 17 Lines
// CHECK: %[[v2:.*]] = memref.view %[[a2]]{{.*}} : memref<24000000xi8> to memref<?x?xf32> // CHECK: %[[v2:.*]] = memref.view %[[a2]]{{.*}} : memref<24000000xi8> to memref<?x?xf32>
// CHECK: %[[l2:.*]] = memref.subview %[[v2]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] // CHECK: %[[l2:.*]] = memref.subview %[[v2]][0, 0] [%{{.*}}, %{{.*}}] [1, 1]
// CHECK-SAME: memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK-SAME: memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK: memref.copy %[[s1]], %[[l1]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: memref.copy %[[s1]], %[[l1]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK: memref.copy %[[s2]], %[[l2]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: memref.copy %[[s2]], %[[l2]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: ins(%[[v0]], %[[v1]] : memref<?x?xf32>, memref<?x?xf32>) // CHECK-SAME: ins(%[[v0]], %[[v1]] : memref<?x?xf32>, memref<?x?xf32>)
// CHECK-SAME: outs(%[[v2]] : memref<?x?xf32>) // CHECK-SAME: inits(%[[v2]] : memref<?x?xf32>)
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = transform.structured.promote %0 { operands_to_promote = [0, 1, 2], use_full_tiles_by_default } %1 = transform.structured.promote %0 { operands_to_promote = [0, 1, 2], use_full_tiles_by_default }
} }
// ----- // -----
Show All 16 Lines scf.for %arg4 = %c0 to %2 step %c3000 {
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%4 = memref.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] : %4 = memref.subview %arg1[%arg5, %arg4][%c4000, %c3000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
%5 = memref.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] : %5 = memref.subview %arg2[%arg3, %arg4][%c2000, %c3000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.matmul {__internal_linalg_transform__ = "_promote_first_view_"} linalg.matmul {__internal_linalg_transform__ = "_promote_first_view_"}
ins(%3, %4: memref<?x?xf32, strided<[?, ?], offset: ?>>, ins(%3, %4: memref<?x?xf32, strided<[?, ?], offset: ?>>,
memref<?x?xf32, strided<[?, ?], offset: ?>>) memref<?x?xf32, strided<[?, ?], offset: ?>>)
outs(%5: memref<?x?xf32, strided<[?, ?], offset: ?>>) inits(%5: memref<?x?xf32, strided<[?, ?], offset: ?>>)
} }
} }
} }
return return
} }
// CHECK-LABEL: func @promote_first_subview_matmul // CHECK-LABEL: func @promote_first_subview_matmul
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index // CHECK-DAG: %[[c2000:.*]] = arith.constant 2000 : index
Show All 10 Lines
// CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK-NOT: memref.alloc // CHECK-NOT: memref.alloc
// CHECK-NOT: memref.view // CHECK-NOT: memref.view
// CHECK-NOT: memref.subview // CHECK-NOT: memref.subview
// CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: ins(%[[v0]], %[[s1]] : memref<?x?xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>) // CHECK-SAME: ins(%[[v0]], %[[s1]] : memref<?x?xf32>, memref<?x?xf32, strided<[?, ?], offset: ?>>)
// CHECK-SAME: outs(%[[s2]] : memref<?x?xf32, strided<[?, ?], offset: ?>>) // CHECK-SAME: inits(%[[s2]] : memref<?x?xf32, strided<[?, ?], offset: ?>>)
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
sequence %arg0 : !pdl.operation failures(propagate) { sequence %arg0 : !pdl.operation failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = transform.structured.promote %0 { operands_to_promote = [0], use_full_tiles_by_default } %1 = transform.structured.promote %0 { operands_to_promote = [0], use_full_tiles_by_default }
} }
} }
// ----- // -----
func.func @aligned_promote_fill(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>) { func.func @aligned_promote_fill(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>) {
%c2000 = arith.constant 2000 : index %c2000 = arith.constant 2000 : index
%c4000 = arith.constant 4000 : index %c4000 = arith.constant 4000 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cf = arith.constant 1.0 : f32 %cf = arith.constant 1.0 : f32
%3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] : %3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] :
memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>> memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, ?], offset: ?>>
linalg.fill linalg.fill
ins(%cf : f32) outs(%3 : memref<?x?xf32, strided<[?, ?], offset: ?>>) ins(%cf : f32) inits(%3 : memref<?x?xf32, strided<[?, ?], offset: ?>>)
return return
} }
// CHECK-LABEL: func @aligned_promote_fill // CHECK-LABEL: func @aligned_promote_fill
// CHECK: %[[cf:.*]] = arith.constant 1.{{.*}} : f32 // CHECK: %[[cf:.*]] = arith.constant 1.{{.*}} : f32
// CHECK: %[[s0:.*]] = memref.subview {{.*}}: memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: %[[s0:.*]] = memref.subview {{.*}}: memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK: %[[a0:.*]] = memref.alloc() {alignment = 32 : i64} : memref<32000000xi8> // CHECK: %[[a0:.*]] = memref.alloc() {alignment = 32 : i64} : memref<32000000xi8>
// CHECK: %[[v0:.*]] = memref.view %[[a0]]{{.*}} : memref<32000000xi8> to memref<?x?xf32> // CHECK: %[[v0:.*]] = memref.view %[[a0]]{{.*}} : memref<32000000xi8> to memref<?x?xf32>
// CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>> // CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xf32> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// CHECK: linalg.fill ins({{.*}} : f32) outs(%[[v0]] : memref<?x?xf32>) // CHECK: linalg.fill ins({{.*}} : f32) inits(%[[v0]] : memref<?x?xf32>)
// CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}> // CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xf32, strided{{.*}}> to memref<?x?xf32, strided{{.*}}>
// CHECK: linalg.fill ins(%[[cf]] : f32) outs(%[[v0]] : memref<?x?xf32>) // CHECK: linalg.fill ins(%[[cf]] : f32) inits(%[[v0]] : memref<?x?xf32>)
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
sequence %arg0 : !pdl.operation failures(propagate) { sequence %arg0 : !pdl.operation failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.fill"]} in %arg1 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1
%1 = transform.structured.promote %0 { operands_to_promote = [1], use_full_tile_buffers = [false, true], alignment = 32} %1 = transform.structured.promote %0 { operands_to_promote = [1], use_full_tile_buffers = [false, true], alignment = 32}
} }
} }
// ----- // -----
func.func @aligned_promote_fill_complex(%arg0: memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>>) { func.func @aligned_promote_fill_complex(%arg0: memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>>) {
%c2000 = arith.constant 2000 : index %c2000 = arith.constant 2000 : index
%c4000 = arith.constant 4000 : index %c4000 = arith.constant 4000 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cf = arith.constant 1.0 : f32 %cf = arith.constant 1.0 : f32
%cc = complex.create %cf, %cf : complex<f32> %cc = complex.create %cf, %cf : complex<f32>
%3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] : %3 = memref.subview %arg0[%c0, %c0][%c2000, %c4000][%c1, %c1] :
memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>> to memref<?x?xcomplex<f32>, strided<[?, ?], offset: ?>> memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>> to memref<?x?xcomplex<f32>, strided<[?, ?], offset: ?>>
linalg.fill ins(%cc : complex<f32>) linalg.fill ins(%cc : complex<f32>)
outs(%3 : memref<?x?xcomplex<f32>, strided<[?, ?], offset: ?>>) inits(%3 : memref<?x?xcomplex<f32>, strided<[?, ?], offset: ?>>)
return return
} }
// CHECK-LABEL: func @aligned_promote_fill_complex // CHECK-LABEL: func @aligned_promote_fill_complex
// CHECK: %[[cc:.*]] = complex.create {{.*}} : complex<f32> // CHECK: %[[cc:.*]] = complex.create {{.*}} : complex<f32>
// CHECK: %[[s0:.*]] = memref.subview {{.*}}: memref<?x?xcomplex<f32>, strided{{.*}}> to memref<?x?xcomplex<f32>, strided{{.*}}> // CHECK: %[[s0:.*]] = memref.subview {{.*}}: memref<?x?xcomplex<f32>, strided{{.*}}> to memref<?x?xcomplex<f32>, strided{{.*}}>
// CHECK: %[[a0:.*]] = memref.alloc() {alignment = 32 : i64} : memref<64000000xi8> // CHECK: %[[a0:.*]] = memref.alloc() {alignment = 32 : i64} : memref<64000000xi8>
// CHECK: %[[v0:.*]] = memref.view %[[a0]]{{.*}} : memref<64000000xi8> to memref<?x?xcomplex<f32>> // CHECK: %[[v0:.*]] = memref.view %[[a0]]{{.*}} : memref<64000000xi8> to memref<?x?xcomplex<f32>>
// CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xcomplex<f32>> to memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>> // CHECK: %[[l0:.*]] = memref.subview %[[v0]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : memref<?x?xcomplex<f32>> to memref<?x?xcomplex<f32>, strided<[?, 1], offset: ?>>
// CHECK: linalg.fill ins({{.*}} : complex<f32>) outs(%[[v0]] : memref<?x?xcomplex<f32>>) // CHECK: linalg.fill ins({{.*}} : complex<f32>) inits(%[[v0]] : memref<?x?xcomplex<f32>>)
// CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xcomplex<f32>, strided{{.*}}> to memref<?x?xcomplex<f32>, strided{{.*}}> // CHECK: memref.copy %[[s0]], %[[l0]] : memref<?x?xcomplex<f32>, strided{{.*}}> to memref<?x?xcomplex<f32>, strided{{.*}}>
// CHECK: linalg.fill ins(%[[cc]] : complex<f32>) outs(%[[v0]] : memref<?x?xcomplex<f32>>) // CHECK: linalg.fill ins(%[[cc]] : complex<f32>) inits(%[[v0]] : memref<?x?xcomplex<f32>>)
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
sequence %arg0 : !pdl.operation failures(propagate) { sequence %arg0 : !pdl.operation failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.fill"]} in %arg1 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1
%1 = transform.structured.promote %0 { operands_to_promote = [1], use_full_tile_buffers = [false, true], alignment = 32} %1 = transform.structured.promote %0 { operands_to_promote = [1], use_full_tile_buffers = [false, true], alignment = 32}
} }
} }

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

Show All 11 Linesmodule {
func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>, func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>,
%D: tensor<?x?xf32>, %sz0: index, %sz1: index) %D: tensor<?x?xf32>, %sz0: index, %sz1: index)
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%5 = linalg.fill %5 = linalg.fill
{__producer__} {__producer__}
ins(%cst : f32) ins(%cst : f32)
outs(%D : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%D : tensor<?x?xf32>) -> tensor<?x?xf32>
%6 = linalg.matmul %6 = linalg.matmul
{__producer__} {__producer__}
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%5 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%5 : tensor<?x?xf32>) -> tensor<?x?xf32>
%7 = linalg.generic %7 = linalg.generic
{__root__, {__root__,
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)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>) ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>)
outs(%D : tensor<?x?xf32>) { inits(%D : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
%16 = arith.maxf %arg3, %cst : f32 %16 = arith.maxf %arg3, %cst : f32
%17 = arith.cmpf ogt, %arg2, %cst : f32 %17 = arith.cmpf ogt, %arg2, %cst : f32
%18 = arith.select %17, %cst, %16 : f32 %18 = arith.select %17, %cst, %16 : f32
linalg.yield %18 : f32 linalg.yield %18 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %7 : tensor<?x?xf32> return %7 : tensor<?x?xf32>
} }
Show All 27 Linesmodule {
func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>, func.func @foo(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?xf32>,
%D: tensor<?x?xf32>, %sz0: index, %sz1: index) %D: tensor<?x?xf32>, %sz0: index, %sz1: index)
-> tensor<?x?xf32> -> tensor<?x?xf32>
{ {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%5 = linalg.fill %5 = linalg.fill
{__producer__} {__producer__}
ins(%cst : f32) ins(%cst : f32)
outs(%D : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%D : tensor<?x?xf32>) -> tensor<?x?xf32>
%6 = linalg.matmul %6 = linalg.matmul
{__producer__} {__producer__}
ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>) ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%5 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%5 : tensor<?x?xf32>) -> tensor<?x?xf32>
%7 = linalg.generic %7 = linalg.generic
{__root__, {__root__,
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)>,
affine_map<(d0, d1) -> (d0, d1)>], affine_map<(d0, d1) -> (d0, d1)>],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>) ins(%C, %6 : tensor<?xf32>, tensor<?x?xf32>)
outs(%D : tensor<?x?xf32>) { inits(%D : tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):
%16 = arith.maxf %arg3, %cst : f32 %16 = arith.maxf %arg3, %cst : f32
%17 = arith.cmpf ogt, %arg2, %cst : f32 %17 = arith.cmpf ogt, %arg2, %cst : f32
%18 = arith.select %17, %cst, %16 : f32 %18 = arith.select %17, %cst, %16 : f32
linalg.yield %18 : f32 linalg.yield %18 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %7 : tensor<?x?xf32> return %7 : tensor<?x?xf32>
} }
Show All 15 Lines

mlir/test/Dialect/Linalg/vectorization.mlir

// RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter -split-input-file | FileCheck %s
// ----- // -----
// CHECK-LABEL: contraction_dot // CHECK-LABEL: contraction_dot
func.func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) { func.func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) {
// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584xf32> // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584xf32>
// CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [0] : vector<1584xf32> to f32 // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [0] : vector<1584xf32> to f32
linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>) linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>)
outs(%C: memref<f32>) inits(%C: memref<f32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.dot"]} in %arg1 %0 = transform.structured.match ops{["linalg.dot"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns }
} }
// ----- // -----
// CHECK-LABEL: contraction_matvec // CHECK-LABEL: contraction_matvec
func.func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) { func.func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) {
// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584xf32> // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584xf32>
// CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [1] : vector<1584x1584xf32> to vector<1584xf32> // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [1] : vector<1584x1584xf32> to vector<1584xf32>
linalg.matvec ins(%A, %B: memref<1584x1584xf32>, memref<1584xf32>) linalg.matvec ins(%A, %B: memref<1584x1584xf32>, memref<1584xf32>)
outs(%C: memref<1584xf32>) inits(%C: memref<1584xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matvec"]} in %arg1 %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns }
} }
// ----- // -----
// CHECK-LABEL: contraction_matmul // CHECK-LABEL: contraction_matmul
func.func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) { func.func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) {
// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584xf32> // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584xf32>
// CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [2] : vector<1584x1584x1584xf32> to vector<1584x1584xf32> // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [2] : vector<1584x1584x1584xf32> to vector<1584x1584xf32>
linalg.matmul ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>) linalg.matmul ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>)
outs(%C: memref<1584x1584xf32>) inits(%C: memref<1584x1584xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns }
} }
// ----- // -----
// CHECK-LABEL: contraction_batch_matmul // CHECK-LABEL: contraction_batch_matmul
func.func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) { func.func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) {
// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584x1584xf32> // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584x1584xf32>
// CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [3] : vector<1584x1584x1584x1584xf32> to vector<1584x1584x1584xf32> // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [3] : vector<1584x1584x1584x1584xf32> to vector<1584x1584x1584xf32>
linalg.batch_matmul linalg.batch_matmul
ins(%A, %B: memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>) ins(%A, %B: memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>)
outs(%C: memref<1584x1584x1584xf32>) inits(%C: memref<1584x1584x1584xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.batch_matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.batch_matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns }
Show All 18 Linesfunc.func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
// CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32> // CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32>
// CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32> // CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32>
// CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<8x32xf32>, vector<8x32xf32> // CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<8x32xf32>, vector<8x32xf32>
// CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> // CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32>
// CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xf32> to vector<8x32xf32> // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xf32> to vector<8x32xf32>
// CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<8x32xf32> // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<8x32xf32>
linalg.generic #matmul_trait linalg.generic #matmul_trait
ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>)
outs(%C : memref<8x32xf32>) { inits(%C : memref<8x32xf32>) {
^bb(%a: f32, %b: f32, %c: f32) : ^bb(%a: f32, %b: f32, %c: f32) :
%d = arith.mulf %a, %b: f32 %d = arith.mulf %a, %b: f32
%e = arith.addf %c, %d: f32 %e = arith.addf %c, %d: f32
linalg.yield %e : f32 linalg.yield %e : f32
} }
return return
} }
Show All 23 Linesfunc.func @generic_output_transpose(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
// CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32> // CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32>
// CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32> // CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32>
// CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<32x8xf32>, vector<8x32xf32> // CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<32x8xf32>, vector<8x32xf32>
// CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> // CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32>
// CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xf32> to vector<8x32xf32> // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xf32> to vector<8x32xf32>
// CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<32x8xf32> // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<32x8xf32>
linalg.generic #matmul_transpose_out_trait linalg.generic #matmul_transpose_out_trait
ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>)
outs(%C : memref<32x8xf32>) { inits(%C : memref<32x8xf32>) {
^bb(%a: f32, %b: f32, %c: f32) : ^bb(%a: f32, %b: f32, %c: f32) :
%d = arith.mulf %a, %b: f32 %d = arith.mulf %a, %b: f32
%e = arith.addf %c, %d: f32 %e = arith.addf %c, %d: f32
linalg.yield %e : f32 linalg.yield %e : f32
} }
return return
} }
Show All 12 Lines
// 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 = tensor.empty() : 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>) { inits(%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>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
Show All 21 Linesfunc.func @vectorization_test_integer(%A: memref<8x16xi32>, %B: memref<16x32xi32>,
// CHECK: vector.transfer_read %{{.*}} : memref<8x16xi32>, vector<8x32x16xi32> // CHECK: vector.transfer_read %{{.*}} : memref<8x16xi32>, vector<8x32x16xi32>
// CHECK: vector.transfer_read %{{.*}} : memref<16x32xi32>, vector<8x32x16xi32> // CHECK: vector.transfer_read %{{.*}} : memref<16x32xi32>, vector<8x32x16xi32>
// CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<8x32xi32>, vector<8x32xi32> // CHECK: %[[ACC:.*]] = vector.transfer_read %{{.*}} : memref<8x32xi32>, vector<8x32xi32>
// CHECK: %[[MUL:.*]] = arith.muli %{{.*}}, %{{.*}} : vector<8x32x16xi32> // CHECK: %[[MUL:.*]] = arith.muli %{{.*}}, %{{.*}} : vector<8x32x16xi32>
// CHECK: vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xi32> to vector<8x32xi32> // CHECK: vector.multi_reduction <add>, %[[MUL]], %[[ACC]] [2] : vector<8x32x16xi32> to vector<8x32xi32>
// CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xi32>, memref<8x32xi32> // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xi32>, memref<8x32xi32>
linalg.generic #matmul_trait linalg.generic #matmul_trait
ins(%A, %B : memref<8x16xi32>, memref<16x32xi32>) ins(%A, %B : memref<8x16xi32>, memref<16x32xi32>)
outs(%C : memref<8x32xi32>) { inits(%C : memref<8x32xi32>) {
^bb(%a: i32, %b: i32, %c: i32) : ^bb(%a: i32, %b: i32, %c: i32) :
%d = arith.muli %a, %b: i32 %d = arith.muli %a, %b: i32
%e = arith.addi %c, %d: i32 %e = arith.addi %c, %d: i32
linalg.yield %e : i32 linalg.yield %e : i32
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns, disable_transfer_permutation_map_lowering_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns, disable_transfer_permutation_map_lowering_patterns }
} }
// ----- // -----
// CHECK-LABEL: func @vectorization_test_2 // CHECK-LABEL: func @vectorization_test_2
func.func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>, func.func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>,
%C: memref<8x32xf32>) { %C: memref<8x32xf32>) {
// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32>
// CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [2] : vector<8x32x16xf32> to vector<8x32xf32> // CHECK: vector.multi_reduction <add>, %{{.*}}, {{.*}} [2] : vector<8x32x16xf32> to vector<8x32xf32>
linalg.matmul linalg.matmul
ins(%A, %B: memref<8x16xf32>, memref<16x32xf32>) ins(%A, %B: memref<8x16xf32>, memref<16x32xf32>)
outs(%C: memref<8x32xf32>) inits(%C: memref<8x32xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns } %2 = transform.structured.vectorize %1 { disable_multi_reduction_to_contract_patterns }
} }
// ----- // -----
// CHECK-LABEL: func @test_vectorize_scalar_input // CHECK-LABEL: func @test_vectorize_scalar_input
func.func @test_vectorize_scalar_input(%A : memref<8x16xf32>, %arg0 : f32) { func.func @test_vectorize_scalar_input(%A : memref<8x16xf32>, %arg0 : f32) {
// CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32> // CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32>
// CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32> // CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>], indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : f32) ins(%arg0 : f32)
outs(%A: memref<8x16xf32>) { inits(%A: memref<8x16xf32>) {
^bb(%0: f32, %1: f32) : ^bb(%0: f32, %1: f32) :
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
} }
// ----- // -----
// CHECK-LABEL: func @test_do_not_vectorize_unsupported_element_types // CHECK-LABEL: func @test_do_not_vectorize_unsupported_element_types
func.func @test_do_not_vectorize_unsupported_element_types(%A : memref<8x16xcomplex<f32>>, %arg0 : complex<f32>) { func.func @test_do_not_vectorize_unsupported_element_types(%A : memref<8x16xcomplex<f32>>, %arg0 : complex<f32>) {
// CHECK-NOT: vector.broadcast // CHECK-NOT: vector.broadcast
// CHECK-NOT: vector.transfer_write // CHECK-NOT: vector.transfer_write
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>], indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0 : complex<f32>) ins(%arg0 : complex<f32>)
outs(%A: memref<8x16xcomplex<f32>>) { inits(%A: memref<8x16xcomplex<f32>>) {
^bb(%0: complex<f32>, %1: complex<f32>) : ^bb(%0: complex<f32>, %1: complex<f32>) :
linalg.yield %0 : complex<f32> linalg.yield %0 : complex<f32>
} }
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
} }
// ----- // -----
// CHECK-LABEL: func @test_vectorize_fill // CHECK-LABEL: func @test_vectorize_fill
func.func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) { func.func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) {
// CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32> // CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32>
// CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32> // CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32>
linalg.fill ins(%arg0 : f32) outs(%A : memref<8x16xf32>) linalg.fill ins(%arg0 : f32) inits(%A : memref<8x16xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.fill"]} in %arg1 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
} }
// ----- // -----
// CHECK-LABEL: func @test_vectorize_fill // CHECK-LABEL: func @test_vectorize_fill
func.func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) { func.func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) {
// CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[val:.*]]: f32) // CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[val:.*]]: f32)
// CHECK: %[[VEC:.*]] = vector.broadcast %[[val]] : f32 to vector<f32> // CHECK: %[[VEC:.*]] = vector.broadcast %[[val]] : f32 to vector<f32>
// CHECK: vector.transfer_write %[[VEC]], %[[M]][] : vector<f32>, memref<f32> // CHECK: vector.transfer_write %[[VEC]], %[[M]][] : vector<f32>, memref<f32>
linalg.fill ins(%arg0 : f32) outs(%A : memref<f32>) linalg.fill ins(%arg0 : f32) inits(%A : memref<f32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.fill"]} in %arg1 %0 = transform.structured.match ops{["linalg.fill"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines// CHECK-LABEL: func @test_vectorize_trailing_index
// CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>) // CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>)
func.func @test_vectorize_trailing_index(%arg0: memref<1x2x4x8xindex>) { func.func @test_vectorize_trailing_index(%arg0: memref<1x2x4x8xindex>) {
// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex> // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
linalg.generic { linalg.generic {
indexing_maps = [ indexing_maps = [
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"]}
outs(%arg0: memref<1x2x4x8xindex>) { inits(%arg0: memref<1x2x4x8xindex>) {
^bb0(%arg1: index): ^bb0(%arg1: index):
// CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<8xindex> to vector<1x2x4x8xindex> // CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<8xindex> to vector<1x2x4x8xindex>
// CHECK: vector.transfer_write %[[BCST]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex> // CHECK: vector.transfer_write %[[BCST]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex>
%0 = linalg.index 3 : index %0 = linalg.index 3 : index
linalg.yield %0 : index linalg.yield %0 : index
} }
return return
} }
Show All 11 Lines// CHECK-LABEL: func @test_vectorize_inner_index
// CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>) // CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>)
func.func @test_vectorize_inner_index(%arg0: memref<1x2x4x8xindex>) { func.func @test_vectorize_inner_index(%arg0: memref<1x2x4x8xindex>) {
// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1]> : vector<2xindex> // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1]> : vector<2xindex>
// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
linalg.generic { linalg.generic {
indexing_maps = [ indexing_maps = [
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"]}
outs(%arg0: memref<1x2x4x8xindex>) { inits(%arg0: memref<1x2x4x8xindex>) {
^bb0(%arg1: index): ^bb0(%arg1: index):
// CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<2xindex> to vector<1x8x4x2xindex> // CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<2xindex> to vector<1x8x4x2xindex>
// CHECK: %[[TRAN:.*]] = vector.transpose %[[BCST]], [0, 3, 2, 1] : vector<1x8x4x2xindex> to vector<1x2x4x8xindex> // CHECK: %[[TRAN:.*]] = vector.transpose %[[BCST]], [0, 3, 2, 1] : vector<1x8x4x2xindex> to vector<1x2x4x8xindex>
// CHECK: vector.transfer_write %[[TRAN]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex> // CHECK: vector.transfer_write %[[TRAN]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex>
%0 = linalg.index 1 : index %0 = linalg.index 1 : index
linalg.yield %0 : index linalg.yield %0 : index
} }
return return
Show All 31 Lines 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)>,
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(%arg1, %arg2: memref<4x256xf32>, memref<256xf32>) ins(%arg1, %arg2: memref<4x256xf32>, memref<256xf32>)
outs( inits(
%arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 : %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 :
memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>,
memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>,
memref<4x256xf32>, memref<4x256xf32>) { memref<4x256xf32>, memref<4x256xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, ^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32,
// CHECK: %[[V2:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> // CHECK: %[[V2:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32>
// CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG2]][%[[C0]]], {{.*}} : memref<256xf32>, vector<4x256xf32> // CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG2]][%[[C0]]], {{.*}} : memref<256xf32>, vector<4x256xf32>
// CHECK: %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> // CHECK: %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32>
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines 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)>,
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(%arg1, %arg2: tensor<4x256xf32>, tensor<256xf32>) ins(%arg1, %arg2: tensor<4x256xf32>, tensor<256xf32>)
outs( inits(
%arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 : %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 :
tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>,
tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>,
tensor<4x256xf32>, tensor<4x256xf32>) { tensor<4x256xf32>, tensor<4x256xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, ^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32,
%arg9 : f32, %arg10 : f32, %arg11 : f32, %arg12 : f32, %arg13 : f32, %arg9 : f32, %arg10 : f32, %arg11 : f32, %arg12 : f32, %arg13 : f32,
%arg14 : f32): %arg14 : f32):
// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<4x256xf32> // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<4x256xf32>
▲ Show 20 Lines • Show All 74 Lines • ▼ Show 20 Linesfunc.func @generic_vectorize_broadcast_transpose(
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d3)>, indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d3)>,
affine_map<(d0, d1, d2, d3) -> (d0)>, affine_map<(d0, d1, d2, d3) -> (d0)>,
affine_map<(d0, d1, d2, d3) -> (d2)>, affine_map<(d0, d1, d2, d3) -> (d2)>,
affine_map<(d0, d1, d2, d3) -> (d2, d0)>, affine_map<(d0, d1, d2, d3) -> (d2, d0)>,
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(%B, %A, %A, %B: memref<4x4xf32>, memref<4xf32>, memref<4xf32>, memref<4x4xf32>) ins(%B, %A, %A, %B: memref<4x4xf32>, memref<4xf32>, memref<4xf32>, memref<4x4xf32>)
outs(%C : memref<4x4x4x4xf32>) { inits(%C : memref<4x4x4x4xf32>) {
^bb0(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32): ^bb0(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32):
%s = arith.subf %arg0, %arg1 : f32 %s = arith.subf %arg0, %arg1 : f32
%a = arith.addf %arg2, %s : f32 %a = arith.addf %arg2, %s : f32
%b = arith.addf %arg3, %a : f32 %b = arith.addf %arg3, %a : f32
linalg.yield %b : f32 linalg.yield %b : f32
} }
return return
} }
Show All 27 Lines
// CHECK: vector.transfer_read {{.*}}{in_bounds = [true, true, true, true], permutation_map = #[[MAP2]]} : memref<16x14x7x8xf32>, vector<7x14x8x16xf32> // CHECK: vector.transfer_read {{.*}}{in_bounds = [true, true, true, true], permutation_map = #[[MAP2]]} : memref<16x14x7x8xf32>, vector<7x14x8x16xf32>
// CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32> // CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32>
// CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32> // CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32>
// CHECK: vector.transfer_write {{.*}} : vector<7x14x8x16xf32>, memref<7x14x8x16xf32> // CHECK: vector.transfer_write {{.*}} : vector<7x14x8x16xf32>, memref<7x14x8x16xf32>
func.func @vectorization_transpose(%A: memref<14x7xf32>, %B: memref<16x14xf32>, func.func @vectorization_transpose(%A: memref<14x7xf32>, %B: memref<16x14xf32>,
%C: memref<16x14x7x8xf32>, %D: memref<7x14x8x16xf32>) { %C: memref<16x14x7x8xf32>, %D: memref<7x14x8x16xf32>) {
linalg.generic #matmul_trait linalg.generic #matmul_trait
ins(%A, %B, %C : memref<14x7xf32>, memref<16x14xf32>, memref<16x14x7x8xf32>) ins(%A, %B, %C : memref<14x7xf32>, memref<16x14xf32>, memref<16x14x7x8xf32>)
outs(%D : memref<7x14x8x16xf32>) { inits(%D : memref<7x14x8x16xf32>) {
^bb(%a: f32, %b: f32, %c: f32, %d: f32) : ^bb(%a: f32, %b: f32, %c: f32, %d: f32) :
%e = arith.addf %a, %b: f32 %e = arith.addf %a, %b: f32
%f = arith.addf %e, %c: f32 %f = arith.addf %e, %c: f32
linalg.yield %f : f32 linalg.yield %f : f32
} }
return return
} }
Show All 18 Linesfunc.func @matmul_tensors(
// CHECK-DAG: %[[V2:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], {{.*}} : tensor<8x12xf32>, vector<8x12xf32> // CHECK-DAG: %[[V2:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], {{.*}} : tensor<8x12xf32>, vector<8x12xf32>
// //
// linalg matmul lowers gets expanded to a 3D reduction, canonicalization later // linalg matmul lowers gets expanded to a 3D reduction, canonicalization later
// convert it to a 2D contract. // convert it to a 2D contract.
// CHECK: %[[MUL:.*]] = arith.mulf %[[V0]], %[[V1]] : vector<8x12x4xf32> // CHECK: %[[MUL:.*]] = arith.mulf %[[V0]], %[[V1]] : vector<8x12x4xf32>
// CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[V2]] [2] : vector<8x12x4xf32> to vector<8x12xf32> // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]], %[[V2]] [2] : vector<8x12x4xf32> to vector<8x12xf32>
// CHECK: %[[W:.*]] = vector.transfer_write %[[R]], %[[ARG2]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x12xf32>, tensor<8x12xf32> // CHECK: %[[W:.*]] = vector.transfer_write %[[R]], %[[ARG2]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x12xf32>, tensor<8x12xf32>
%0 = linalg.matmul ins(%arg0, %arg1: tensor<8x4xf32>, tensor<4x12xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<8x4xf32>, tensor<4x12xf32>)
outs(%arg2: tensor<8x12xf32>) inits(%arg2: tensor<8x12xf32>)
-> tensor<8x12xf32> -> tensor<8x12xf32>
// CHECK: return %[[W]] : tensor<8x12xf32> // CHECK: return %[[W]] : tensor<8x12xf32>
return %0 : tensor<8x12xf32> return %0 : tensor<8x12xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
▲ Show 20 Lines • Show All 70 Lines • ▼ Show 20 Lines
// 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:.*]] = tensor.empty(%[[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) inits(%[[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):
tensor.yield %pad_value : f32 tensor.yield %pad_value : f32
▲ Show 20 Lines • Show All 222 Lines • ▼ Show 20 Linesfunc.func @sum_exp(%input: tensor<4x16x8xf32>, %output: tensor<4x16xf32>)
// CHECK: vector.transfer_write {{.*}} : vector<4x16xf32>, tensor<4x16xf32> // CHECK: vector.transfer_write {{.*}} : vector<4x16xf32>, tensor<4x16xf32>
// CHECK: return {{.*}} : tensor<4x16xf32> // CHECK: return {{.*}} : tensor<4x16xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ 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)> affine_map<(d0, d1, d2) -> (d0, d1)>
], ],
iterator_types = ["parallel", "parallel", "reduction"] iterator_types = ["parallel", "parallel", "reduction"]
} ins(%input : tensor<4x16x8xf32>) outs(%output : tensor<4x16xf32>) { } ins(%input : tensor<4x16x8xf32>) inits(%output : tensor<4x16xf32>) {
^bb0(%arg0: f32, %arg1: f32): ^bb0(%arg0: f32, %arg1: f32):
%1 = math.exp %arg0 : f32 %1 = math.exp %arg0 : f32
%2 = arith.addf %1, %arg1 : f32 %2 = arith.addf %1, %arg1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<4x16xf32> } -> tensor<4x16xf32>
return %0 : tensor<4x16xf32> return %0 : tensor<4x16xf32>
} }
Show All 26 Linesfunc.func @sum_exp_2(%input: tensor<3x2xf32>, %input_2: tensor<5x4xf32>, %output: tensor<5x2xf32>)
// CHECK: return {{.*}} : tensor<5x2xf32> // CHECK: return {{.*}} : tensor<5x2xf32>
%0 = linalg.generic { %0 = linalg.generic {
indexing_maps = [ indexing_maps = [
affine_map<(d0, d1, d2, d3) -> (d1, d0)>, affine_map<(d0, d1, d2, d3) -> (d1, d0)>,
affine_map<(d0, d1, d2, d3) -> (d3, d2)>, affine_map<(d0, d1, d2, d3) -> (d3, d2)>,
affine_map<(d0, d1, d2, d3) -> (d3, d0)> affine_map<(d0, d1, d2, d3) -> (d3, d0)>
], ],
iterator_types = ["parallel", "reduction", "reduction", "parallel"] iterator_types = ["parallel", "reduction", "reduction", "parallel"]
} ins(%input, %input_2 : tensor<3x2xf32>, tensor<5x4xf32>) outs(%output : tensor<5x2xf32>) { } ins(%input, %input_2 : tensor<3x2xf32>, tensor<5x4xf32>) inits(%output : tensor<5x2xf32>) {
^bb0(%arg0: f32, %arg1: f32, %arg2: f32): ^bb0(%arg0: f32, %arg1: f32, %arg2: f32):
%1 = math.exp %arg0 : f32 %1 = math.exp %arg0 : f32
%2 = math.exp %arg1 : f32 %2 = math.exp %arg1 : f32
%3 = arith.addf %1, %2 : f32 %3 = arith.addf %1, %2 : f32
%4 = arith.addf %3, %arg2 : f32 %4 = arith.addf %3, %arg2 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<5x2xf32> } -> tensor<5x2xf32>
return %0 : tensor<5x2xf32> return %0 : tensor<5x2xf32>
Show All 12 Lines
// 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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xf32> } -> tensor<4xf32>
return %red : tensor<4xf32> return %red : tensor<4xf32>
} }
Show All 10 Lines
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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xf32> } -> tensor<4xf32>
return %red : tensor<4xf32> return %red : tensor<4xf32>
} }
Show All 9 Lines
// 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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xf32> } -> tensor<4xf32>
return %red : tensor<4xf32> return %red : tensor<4xf32>
} }
Show All 9 Lines
// 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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xi1> } -> tensor<4xi1>
return %red : tensor<4xi1> return %red : tensor<4xi1>
} }
Show All 9 Lines
// 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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xi1> } -> tensor<4xi1>
return %red : tensor<4xi1> return %red : tensor<4xi1>
} }
Show All 9 Lines
// 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: tensor.empty() : 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 = tensor.empty() : 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) inits(%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>) inits(%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
} -> tensor<4xi1> } -> tensor<4xi1>
return %red : tensor<4xi1> return %red : tensor<4xi1>
} }
Show All 11 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 = tensor.empty() : 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) inits(%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>) { inits(%fill : tensor<4x4xf32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
%40 = arith.subf %arg7, %arg8 : f32 %40 = arith.subf %arg7, %arg8 : f32
linalg.yield %40 : f32 linalg.yield %40 : f32
} -> tensor<4x4xf32> } -> tensor<4x4xf32>
return %red : tensor<4x4xf32> return %red : tensor<4x4xf32>
} }
Show All 13 Linesfunc.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 = tensor.empty() : 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) inits(%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>) { inits(%fill : tensor<4xf32>) {
^bb0(%arg7: f32, %arg8: f32, %arg9: f32): ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
%40 = arith.subf %arg7, %arg8 : f32 %40 = arith.subf %arg7, %arg8 : f32
%41 = math.exp %40 : f32 %41 = math.exp %40 : f32
%42 = arith.addf %41, %arg9 : f32 %42 = arith.addf %41, %arg9 : f32
linalg.yield %42 : f32 linalg.yield %42 : f32
} -> tensor<4xf32> } -> tensor<4xf32>
return %red : tensor<4xf32> return %red : tensor<4xf32>
} }
Show All 18 Linesfunc.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:.*]] = tensor.empty() : tensor<f32> // CHECK: %[[init:.*]] = tensor.empty() : tensor<f32>
%0 = tensor.empty() : 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) inits(%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>
// CHECK: %[[res:.*]] = vector.transfer_write %[[red_v1]], %[[init]][] // CHECK: %[[res:.*]] = vector.transfer_write %[[red_v1]], %[[init]][]
// CHECK-SAME: : vector<f32>, tensor<f32> // CHECK-SAME: : vector<f32>, tensor<f32>
%2 = linalg.generic { %2 = linalg.generic {
indexing_maps = [affine_map<(d0) -> (d0)>, indexing_maps = [affine_map<(d0) -> (d0)>,
affine_map<(d0) -> ()>], affine_map<(d0) -> ()>],
iterator_types = ["reduction"]} iterator_types = ["reduction"]}
ins(%arg0 : tensor<32xf32>) ins(%arg0 : tensor<32xf32>)
outs(%1 : tensor<f32>) { inits(%1 : tensor<f32>) {
^bb0(%a: f32, %b: f32): ^bb0(%a: f32, %b: f32):
%3 = arith.addf %a, %b : f32 %3 = arith.addf %a, %b : f32
linalg.yield %3 : f32 linalg.yield %3 : f32
} -> tensor<f32> } -> tensor<f32>
return %2 : tensor<f32> return %2 : tensor<f32>
} }
Show All 10 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 = tensor.empty() : 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) inits(%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>) { inits(%fill : tensor<6x6x3x3xf32>) {
^bb0(%arg7: f32, %arg9: f32): ^bb0(%arg7: f32, %arg9: f32):
linalg.yield %arg7 : f32 linalg.yield %arg7 : f32
} -> tensor<6x6x3x3xf32> } -> tensor<6x6x3x3xf32>
return %result : tensor<6x6x3x3xf32> return %result : tensor<6x6x3x3xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation): ^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.generic"]} in %arg1 %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
%1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation %1 = get_closest_isolated_parent %0 : (!pdl.operation) -> !pdl.operation
%2 = transform.structured.vectorize %1 %2 = transform.structured.vectorize %1
} }
// ----- // -----
// Check vectorization can handle cases where outputs are a mix of reduced and non-reduced values. // Check vectorization can handle cases where outputs are a mix of reduced and non-reduced values.
func.func @mixed_parallel_reduced_results(%arg0 : tensor<2x4x8xf32>, func.func @mixed_parallel_reduced_results(%arg0 : tensor<2x4x8xf32>,
%arg1 : tensor<2x4xf32>, %arg2 : tensor<2x4x8xf32>, %arg3 : tensor<2x4xf32>) -> %arg1 : tensor<2x4xf32>, %arg2 : tensor<2x4x8xf32>, %arg3 : tensor<2x4xf32>) ->
(tensor<2x4x8xf32>, tensor<2x4xf32>) { (tensor<2x4x8xf32>, tensor<2x4xf32>) {
%0:2 = linalg.generic { %0:2 = linalg.generic {
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>,
affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>], affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],
iterator_types = ["parallel", "parallel", "reduction"]} iterator_types = ["parallel", "parallel", "reduction"]}
ins(%arg0, %arg1 : tensor<2x4x8xf32>, tensor<2x4xf32>) ins(%arg0, %arg1 : tensor<2x4x8xf32>, tensor<2x4xf32>)
outs(%arg2, %arg3 : tensor<2x4x8xf32>, tensor<2x4xf32>) { inits(%arg2, %arg3 : tensor<2x4x8xf32>, tensor<2x4xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32): ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32):
%1 = arith.mulf %b0, %b1 : f32 %1 = arith.mulf %b0, %b1 : f32
%2 = arith.addf %1, %b3 : f32 %2 = arith.addf %1, %b3 : f32
linalg.yield %1, %2 : f32, f32 linalg.yield %1, %2 : f32, f32
} -> (tensor<2x4x8xf32>, tensor<2x4xf32>) } -> (tensor<2x4x8xf32>, tensor<2x4xf32>)
return %0#0, %0#1 : tensor<2x4x8xf32>, tensor<2x4xf32> return %0#0, %0#1 : tensor<2x4x8xf32>, tensor<2x4xf32>
} }
// CHECK-LABEL: func @mixed_parallel_reduced_results( // CHECK-LABEL: func @mixed_parallel_reduced_results(
Show All 20 Lines
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, 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 @vectorize_1d_tensor_extract(%arg0: tensor<3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x7x2xf32>, %arg3: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> { func.func @vectorize_1d_tensor_extract(%arg0: tensor<3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x7x2xf32>, %arg3: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {
%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(%arg1, %arg2 : tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%arg3 : tensor<4x7x3x2xf32>) { } ins(%arg1, %arg2 : tensor<4x3xi32>, tensor<4x7x2xf32>) inits(%arg3 : tensor<4x7x3x2xf32>) {
^bb0(%arg4: i32, %arg5: f32, %arg6: f32): ^bb0(%arg4: i32, %arg5: f32, %arg6: f32):
%3 = arith.index_cast %arg4 : i32 to index %3 = arith.index_cast %arg4 : i32 to index
%7 = tensor.extract %arg0[%3] : tensor<3xf32> %7 = tensor.extract %arg0[%3] : tensor<3xf32>
linalg.yield %7 : f32 linalg.yield %7 : f32
} -> tensor<4x7x3x2xf32> } -> tensor<4x7x3x2xf32>
return %2 : tensor<4x7x3x2xf32> return %2 : tensor<4x7x3x2xf32>
} }
// CHECK-LABEL: func.func @vectorize_1d_tensor_extract // CHECK-LABEL: func.func @vectorize_1d_tensor_extract
Show All 20 Lines
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)> #map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> #map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, 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 @not_vectorize_nd_tensor_extract(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>, %arg3: tensor<4x7x2xf32>, %arg4: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> { func.func @not_vectorize_nd_tensor_extract(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>, %arg3: tensor<4x7x2xf32>, %arg4: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {
%2 = linalg.generic { %2 = linalg.generic {
indexing_maps = [#map0, #map0, #map1, #map2], indexing_maps = [#map0, #map0, #map1, #map2],
iterator_types = ["parallel", "parallel", "parallel", "parallel"] iterator_types = ["parallel", "parallel", "parallel", "parallel"]
} ins(%arg1, %arg2, %arg3 : tensor<4x3xi32>, tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%arg4 : tensor<4x7x3x2xf32>) { } ins(%arg1, %arg2, %arg3 : tensor<4x3xi32>, tensor<4x3xi32>, tensor<4x7x2xf32>) inits(%arg4 : tensor<4x7x3x2xf32>) {
^bb0(%arg5: i32, %arg6: i32, %arg7: f32, %arg8: f32): ^bb0(%arg5: i32, %arg6: i32, %arg7: f32, %arg8: f32):
%3 = arith.index_cast %arg5 : i32 to index %3 = arith.index_cast %arg5 : i32 to index
%4 = arith.index_cast %arg6 : i32 to index %4 = arith.index_cast %arg6 : i32 to index
%7 = tensor.extract %arg0[%3, %4] : tensor<3x3xf32> %7 = tensor.extract %arg0[%3, %4] : tensor<3x3xf32>
linalg.yield %7 : f32 linalg.yield %7 : f32
} -> tensor<4x7x3x2xf32> } -> tensor<4x7x3x2xf32>
return %2 : tensor<4x7x3x2xf32> return %2 : tensor<4x7x3x2xf32>
} }
Show All 9 Lines

mlir/test/Dialect/Linalg/vectorize-convolution.mlir

// RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-linalg-to-vector-patterns %s | FileCheck %s // RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-linalg-to-vector-patterns %s | FileCheck %s
func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<1x3x8xf32>, %output: memref<4x2x8xf32>) { func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<1x3x8xf32>, %output: memref<4x2x8xf32>) {
linalg.conv_1d_nwc_wcf linalg.conv_1d_nwc_wcf
{dilations = dense<1> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>} {dilations = dense<1> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>}
ins(%input, %filter : memref<4x6x3xf32>, memref<1x3x8xf32>) ins(%input, %filter : memref<4x6x3xf32>, memref<1x3x8xf32>)
outs(%output : memref<4x2x8xf32>) inits(%output : memref<4x2x8xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_nwc_4x2x8_memref // CHECK: func @conv1d_nwc_4x2x8_memref
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines
// CHECK: vector.transfer_write %[[RES_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]] // CHECK: vector.transfer_write %[[RES_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]]
// ----- // -----
func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<2x3x8xf32>, %output: memref<4x2x8xf32>) { func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<2x3x8xf32>, %output: memref<4x2x8xf32>) {
linalg.conv_1d_nwc_wcf linalg.conv_1d_nwc_wcf
{dilations = dense<2> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>} {dilations = dense<2> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>}
ins(%input, %filter : memref<4x6x3xf32>, memref<2x3x8xf32>) ins(%input, %filter : memref<4x6x3xf32>, memref<2x3x8xf32>)
outs(%output : memref<4x2x8xf32>) inits(%output : memref<4x2x8xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_nwc_4x2x8_memref // CHECK: func @conv1d_nwc_4x2x8_memref
▲ Show 20 Lines • Show All 60 Lines • ▼ Show 20 Lines
// CHECK: vector.transfer_write %[[RES_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]] // CHECK: vector.transfer_write %[[RES_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]]
// ----- // -----
func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<2x3x8xf32>, %output: memref<4x2x8xf32>) { func.func @conv1d_nwc_4x2x8_memref(%input: memref<4x6x3xf32>, %filter: memref<2x3x8xf32>, %output: memref<4x2x8xf32>) {
linalg.conv_1d_nwc_wcf linalg.conv_1d_nwc_wcf
{dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} {dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
ins(%input, %filter : memref<4x6x3xf32>, memref<2x3x8xf32>) ins(%input, %filter : memref<4x6x3xf32>, memref<2x3x8xf32>)
outs(%output : memref<4x2x8xf32>) inits(%output : memref<4x2x8xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_nwc_4x2x8_memref // CHECK: func @conv1d_nwc_4x2x8_memref
Show All 32 Lines
// CHECK: vector.transfer_write %[[CONTRACT_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]] // CHECK: vector.transfer_write %[[CONTRACT_1]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]]
// ----- // -----
func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x1xf32>, %output: memref<4x8x2xf32>) { func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x1xf32>, %output: memref<4x8x2xf32>) {
linalg.conv_1d_ncw_fcw linalg.conv_1d_ncw_fcw
{dilations = dense<1> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>} {dilations = dense<1> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>}
ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x1xf32>) ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x1xf32>)
outs(%output : memref<4x8x2xf32>) inits(%output : memref<4x8x2xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_ncw_4x8x2_memref // CHECK: func @conv1d_ncw_4x8x2_memref
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Lines
// CHECK: vector.transfer_write %[[RES_2]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]] // CHECK: vector.transfer_write %[[RES_2]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]]
// ----- // -----
func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x2xf32>, %output: memref<4x8x2xf32>) { func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x2xf32>, %output: memref<4x8x2xf32>) {
linalg.conv_1d_ncw_fcw linalg.conv_1d_ncw_fcw
{dilations = dense<2> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>} {dilations = dense<2> : tensor<1xi64>, strides = dense<3> : tensor<1xi64>}
ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x2xf32>) ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x2xf32>)
outs(%output : memref<4x8x2xf32>) inits(%output : memref<4x8x2xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_ncw_4x8x2_memref // CHECK: func @conv1d_ncw_4x8x2_memref
▲ Show 20 Lines • Show All 68 Lines • ▼ Show 20 Lines
// CHECK: vector.transfer_write %[[RES_2]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]] // CHECK: vector.transfer_write %[[RES_2]], %[[OUTPUT]][%[[C0]], %[[C0]], %[[C0]]]
// ----- // -----
func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x2xf32>, %output: memref<4x8x2xf32>) { func.func @conv1d_ncw_4x8x2_memref(%input: memref<4x3x6xf32>, %filter: memref<8x3x2xf32>, %output: memref<4x8x2xf32>) {
linalg.conv_1d_ncw_fcw linalg.conv_1d_ncw_fcw
{dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} {dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x2xf32>) ins(%input, %filter : memref<4x3x6xf32>, memref<8x3x2xf32>)
outs(%output : memref<4x8x2xf32>) inits(%output : memref<4x8x2xf32>)
return return
} }
// CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)> // CHECK: #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
// CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)> // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>
// CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> // CHECK: #[[OUTPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
// CHECK: func @conv1d_ncw_4x8x2_memref // CHECK: func @conv1d_ncw_4x8x2_memref
▲ Show 20 Lines • Show All 41 Lines • ▼ Show 20 Lines
// ----- // -----
func.func @depthwise_conv1d_nwc_wc_3x5x4_memref(%input: memref<3x5x4xf32>, %filter: memref<2x4xf32>, %output: memref<3x2x4xf32>) { func.func @depthwise_conv1d_nwc_wc_3x5x4_memref(%input: memref<3x5x4xf32>, %filter: memref<2x4xf32>, %output: memref<3x2x4xf32>) {
linalg.depthwise_conv_1d_nwc_wc linalg.depthwise_conv_1d_nwc_wc
{dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} {dilations = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}
ins(%input, %filter : memref<3x5x4xf32>, memref<2x4xf32>) ins(%input, %filter : memref<3x5x4xf32>, memref<2x4xf32>)
outs(%output : memref<3x2x4xf32>) inits(%output : memref<3x2x4xf32>)
return return
} }
// CHECK: func @depthwise_conv1d_nwc_wc_3x5x4_memref // CHECK: func @depthwise_conv1d_nwc_wc_3x5x4_memref
// CHECK-SAME: (%[[INPUT:[0-9a-z]+]]: memref<3x5x4xf32>, %[[FILTER:[0-9a-z]+]]: memref<2x4xf32>, %[[OUTPUT:[0-9a-z]+]]: memref<3x2x4xf32>) // CHECK-SAME: (%[[INPUT:[0-9a-z]+]]: memref<3x5x4xf32>, %[[FILTER:[0-9a-z]+]]: memref<2x4xf32>, %[[OUTPUT:[0-9a-z]+]]: memref<3x2x4xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
Show All 24 Lines
// ----- // -----
func.func @conv_1d_nwc_wcf_mixed_type_memref(%input: memref<1x2x3xf16>, %filter: memref<1x3x2xf16>, %output: memref<1x2x2xf32>) { func.func @conv_1d_nwc_wcf_mixed_type_memref(%input: memref<1x2x3xf16>, %filter: memref<1x3x2xf16>, %output: memref<1x2x2xf32>) {
linalg.conv_1d_nwc_wcf linalg.conv_1d_nwc_wcf
{dilations = dense<1> : vector<1xi64>, strides = dense<1> : vector<1xi64>} {dilations = dense<1> : vector<1xi64>, strides = dense<1> : vector<1xi64>}
ins(%input, %filter : memref<1x2x3xf16>, memref<1x3x2xf16>) ins(%input, %filter : memref<1x2x3xf16>, memref<1x3x2xf16>)
outs(%output : memref<1x2x2xf32>) inits(%output : memref<1x2x2xf32>)
return return
} }
// CHECK: func @conv_1d_nwc_wcf_mixed_type_memref // CHECK: func @conv_1d_nwc_wcf_mixed_type_memref
// CHECK-SAME: (%[[INPUT:[0-9a-z]+]]: memref<1x2x3xf16>, %[[FILTER:[0-9a-z]+]]: memref<1x3x2xf16>, %[[OUTPUT:[0-9a-z]+]]: memref<1x2x2xf32>) // CHECK-SAME: (%[[INPUT:[0-9a-z]+]]: memref<1x2x3xf16>, %[[FILTER:[0-9a-z]+]]: memref<1x3x2xf16>, %[[OUTPUT:[0-9a-z]+]]: memref<1x2x2xf32>)
// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
Show All 9 Lines

mlir/test/Dialect/SCF/foreach-thread-canonicalization.mlir

// RUN: mlir-opt %s -scf-for-loop-canonicalization | FileCheck %s // RUN: mlir-opt %s -scf-for-loop-canonicalization | FileCheck %s
func.func @reduce() { func.func @reduce() {
// CHECK: %[[C64:.*]] = arith.constant 64 : index // CHECK: %[[C64:.*]] = arith.constant 64 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%cst_0 = arith.constant -0.000000e+00 : f32 %cst_0 = arith.constant -0.000000e+00 : f32
%0 = memref.alloc() : memref<128x384xf32> %0 = memref.alloc() : memref<128x384xf32>
linalg.fill ins(%cst_0 : f32) outs(%0 : memref<128x384xf32>) linalg.fill ins(%cst_0 : f32) inits(%0 : memref<128x384xf32>)
%2 = memref.alloc() : memref<128xf32> %2 = memref.alloc() : memref<128xf32>
linalg.fill ins(%cst_0 : f32) outs(%2 : memref<128xf32>) linalg.fill ins(%cst_0 : f32) inits(%2 : memref<128xf32>)
scf.foreach_thread (%arg0) in (%c2) { scf.foreach_thread (%arg0) in (%c2) {
%7 = affine.min affine_map<(d0) -> (d0 * -64 + 128, 64)>(%arg0) %7 = affine.min affine_map<(d0) -> (d0 * -64 + 128, 64)>(%arg0)
%8 = affine.max affine_map<(d0) -> (0, d0)>(%7) %8 = affine.max affine_map<(d0) -> (0, d0)>(%7)
%9 = affine.apply affine_map<(d0) -> (d0 * 64)>(%arg0) %9 = affine.apply affine_map<(d0) -> (d0 * 64)>(%arg0)
%10 = affine.min affine_map<(d0, d1) -> (d1 * -64 + 128, d0)>(%8, %arg0) %10 = affine.min affine_map<(d0, d1) -> (d1 * -64 + 128, d0)>(%8, %arg0)
// CHECK: memref.subview %{{.*}}[%{{.*}}, 0] [%[[C64]], 384] [1, 1] : memref<128x384xf32> to memref<?x384xf32, {{.*}}> // CHECK: memref.subview %{{.*}}[%{{.*}}, 0] [%[[C64]], 384] [1, 1] : memref<128x384xf32> to memref<?x384xf32, {{.*}}>
// CHECK: memref.subview %{{.*}}[%{{.*}}] [%[[C64]]] [1] : memref<128xf32> to memref<?xf32, {{.*}}> // CHECK: memref.subview %{{.*}}[%{{.*}}] [%[[C64]]] [1] : memref<128xf32> to memref<?xf32, {{.*}}>
%11 = memref.subview %0[%9, 0] [%10, 384] [1, 1] : %11 = memref.subview %0[%9, 0] [%10, 384] [1, 1] :
memref<128x384xf32> to memref<?x384xf32, affine_map<(d0, d1)[s0] -> (d0 * 384 + s0 + d1)>> memref<128x384xf32> to memref<?x384xf32, affine_map<(d0, d1)[s0] -> (d0 * 384 + s0 + d1)>>
%12 = memref.subview %2[%9] [%10] [1] : %12 = memref.subview %2[%9] [%10] [1] :
memref<128xf32> to memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>> memref<128xf32> to memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
// CHECK: linalg.generic {{.*}} ins(%{{.*}} : memref<?x384xf32, {{.*}}>) outs(%{{.*}} : memref<?xf32, {{.*}}>) // CHECK: linalg.generic {{.*}} ins(%{{.*}} : memref<?x384xf32, {{.*}}>) inits(%{{.*}} : memref<?xf32, {{.*}}>)
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)>], affine_map<(d0, d1) -> (d0)>],
iterator_types = ["parallel", "reduction"]} iterator_types = ["parallel", "reduction"]}
ins(%11 : memref<?x384xf32, affine_map<(d0, d1)[s0] -> (d0 * 384 + s0 + d1)>>) ins(%11 : memref<?x384xf32, affine_map<(d0, d1)[s0] -> (d0 * 384 + s0 + d1)>>)
outs(%12 : memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>>) { inits(%12 : memref<?xf32, affine_map<(d0)[s0] -> (d0 + s0)>>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%14 = arith.addf %arg1, %arg2 : f32 %14 = arith.addf %arg1, %arg2 : f32
linalg.yield %14 : f32 linalg.yield %14 : f32
} }
} }
return return
} }

mlir/test/Dialect/SCF/invalid.mlir

Show First 20 Lines • Show All 522 Lines • ▼ Show 20 Lines
// ----- // -----
func.func @wrong_num_results(%in: tensor<100xf32>, %out: tensor<100xf32>) { func.func @wrong_num_results(%in: tensor<100xf32>, %out: tensor<100xf32>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%num_threads = arith.constant 100 : index %num_threads = arith.constant 100 : index
// expected-error @+1 {{1 operands present, but expected 2}} // expected-error @+1 {{1 operands present, but expected 2}}
%result:2 = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>, tensor<100xf32>) { %result:2 = scf.foreach_thread (%thread_idx) in (%num_threads) shared_inits(%o = %out) -> (tensor<100xf32>, tensor<100xf32>) {
%1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
tensor<1xf32> into tensor<100xf32> tensor<1xf32> into tensor<100xf32>
} }
} }
return return
} }
// ----- // -----
func.func @invalid_insert_dest(%in: tensor<100xf32>, %out: tensor<100xf32>) { func.func @invalid_insert_dest(%in: tensor<100xf32>, %out: tensor<100xf32>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%num_threads = arith.constant 100 : index %num_threads = arith.constant 100 : index
%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>) { %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_inits(%o = %out) -> (tensor<100xf32>) {
%1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
// expected-error @+1 {{may only insert into an output block argument}} // expected-error @+1 {{may only insert into an output block argument}}
tensor.parallel_insert_slice %1 into %out[%thread_idx][1][1] : tensor.parallel_insert_slice %1 into %out[%thread_idx][1][1] :
tensor<1xf32> into tensor<100xf32> tensor<1xf32> into tensor<100xf32>
} }
} }
return return
} }
// ----- // -----
func.func @wrong_terminator_op(%in: tensor<100xf32>, %out: tensor<100xf32>) { func.func @wrong_terminator_op(%in: tensor<100xf32>, %out: tensor<100xf32>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%num_threads = arith.constant 100 : index %num_threads = arith.constant 100 : index
%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> (tensor<100xf32>) { %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_inits(%o = %out) -> (tensor<100xf32>) {
%1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
// expected-error @+1 {{expected only tensor.parallel_insert_slice ops}} // expected-error @+1 {{expected only tensor.parallel_insert_slice ops}}
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
tensor<1xf32> into tensor<100xf32> tensor<1xf32> into tensor<100xf32>
%0 = arith.constant 1: index %0 = arith.constant 1: index
} }
} }
▲ Show 20 Lines • Show All 56 Lines • Show Last 20 Lines

mlir/test/Dialect/SCF/loop-pipelining.mlir

Show First 20 Lines • Show All 552 Lines • ▼ Show 20 Lines
// CHECK: memref.copy %[[ASV0]], %[[PAV0]] : // CHECK: memref.copy %[[ASV0]], %[[PAV0]] :
// CHECK: memref.copy %[[BSV0]], %[[PBV0]] : // CHECK: memref.copy %[[BSV0]], %[[PBV0]] :
// Kernel: // Kernel:
// CHECK: %[[R:.+]]:2 = scf.for %[[IV:.+]] = %[[C0]] to %[[C3]] step %[[C1]] // CHECK: %[[R:.+]]:2 = scf.for %[[IV:.+]] = %[[C0]] to %[[C3]] step %[[C1]]
// CHECK-SAME: iter_args(%[[IA:.+]] = %[[PAV0]], %[[IB:.+]] = %[[PBV0:.+]]) // CHECK-SAME: iter_args(%[[IA:.+]] = %[[PAV0]], %[[IB:.+]] = %[[PBV0:.+]])
// CHECK: %[[CV:.+]] = memref.subview %[[ARG2]] // CHECK: %[[CV:.+]] = memref.subview %[[ARG2]]
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[IA]], %[[IB]], %{{.*}} : {{.*}}) outs(%[[CV]] : // CHECK-SAME: ins(%[[IA]], %[[IB]], %{{.*}} : {{.*}}) inits(%[[CV]] :
// CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]] // CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]]
// CHECK: %[[ASV:.+]] = memref.subview %[[ARG0]][%[[NEXT]]] [8] [1] : // CHECK: %[[ASV:.+]] = memref.subview %[[ARG0]][%[[NEXT]]] [8] [1] :
// CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]] : // CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]] :
// CHECK: %[[BSV:.+]] = memref.subview %[[ARG1]][%[[NEXT]]] [8] [1] : // CHECK: %[[BSV:.+]] = memref.subview %[[ARG1]][%[[NEXT]]] [8] [1] :
// CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]] : // CHECK: %[[NEXT:.+]] = arith.addi %[[IV]], %[[C1]] :
// CHECK: %[[BUFIDX:.+]] = affine.apply // CHECK: %[[BUFIDX:.+]] = affine.apply
// CHECK: %[[APROSV:.+]] = memref.subview %[[APRO]][%[[BUFIDX]], 0] [1, 8] [1, 1] : // CHECK: %[[APROSV:.+]] = memref.subview %[[APRO]][%[[BUFIDX]], 0] [1, 8] [1, 1] :
// CHECK: %[[BPROSV:.+]] = memref.subview %[[BPRO]][%[[BUFIDX]], 0] [1, 8] [1, 1] : // CHECK: %[[BPROSV:.+]] = memref.subview %[[BPRO]][%[[BUFIDX]], 0] [1, 8] [1, 1] :
// CHECK: memref.copy %[[ASV]], %[[APROSV]] : // CHECK: memref.copy %[[ASV]], %[[APROSV]] :
// CHECK: memref.copy %[[BSV]], %[[BPROSV]] : // CHECK: memref.copy %[[BSV]], %[[BPROSV]] :
// CHECK: scf.yield %[[APROSV]], %[[BPROSV]] : // CHECK: scf.yield %[[APROSV]], %[[BPROSV]] :
// CHECK: } // CHECK: }
// CHECK: %[[CV:.+]] = memref.subview %[[ARG2]][%[[C3]]] [8] [1] : // CHECK: %[[CV:.+]] = memref.subview %[[ARG2]][%[[C3]]] [8] [1] :
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: ins(%[[R]]#0, %[[R]]#1, %{{.*}} : {{.*}}) outs(%[[CV]] : // CHECK-SAME: ins(%[[R]]#0, %[[R]]#1, %{{.*}} : {{.*}}) inits(%[[CV]] :
#map = affine_map<(d0)[s0]->(d0 + s0)> #map = affine_map<(d0)[s0]->(d0 + s0)>
#map1 = affine_map<(d0)->(d0)> #map1 = affine_map<(d0)->(d0)>
#map2 = affine_map<(d0)->()> #map2 = affine_map<(d0)->()>
#linalg_attrs = { #linalg_attrs = {
indexing_maps = [ indexing_maps = [
#map1, #map1,
Show All 18 Lines scf.for %i0 = %c0 to %c4 step %c1 {
%buf_idx = affine.apply affine_map<(d0)->(d0 mod 2)> (%i0)[] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 5 } %buf_idx = affine.apply affine_map<(d0)->(d0 mod 2)> (%i0)[] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 5 }
%a_buf_view = memref.subview %a_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 6 } : memref<2x8xf32> to memref<8xf32, #map> %a_buf_view = memref.subview %a_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 6 } : memref<2x8xf32> to memref<8xf32, #map>
%b_buf_view = memref.subview %b_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 7 } : memref<2x8xf32> to memref<8xf32, #map> %b_buf_view = memref.subview %b_buf[%buf_idx,0][1,8][1,1] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 7 } : memref<2x8xf32> to memref<8xf32, #map>
memref.copy %A_view , %a_buf_view {__test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 8} : memref<8xf32, #map> to memref<8xf32, #map> memref.copy %A_view , %a_buf_view {__test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 8} : memref<8xf32, #map> to memref<8xf32, #map>
memref.copy %B_view , %b_buf_view {__test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 9} : memref<8xf32, #map> to memref<8xf32, #map> memref.copy %B_view , %b_buf_view {__test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 9} : memref<8xf32, #map> to memref<8xf32, #map>
%C_view = memref.subview %result[%i0][8][1] { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : memref<?xf32> to memref<8xf32, #map> %C_view = memref.subview %result[%i0][8][1] { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 0 } : memref<?xf32> to memref<8xf32, #map>
%scalar = arith.addf %cf, %cf {__test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1} : f32 %scalar = arith.addf %cf, %cf {__test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1} : f32
linalg.generic #linalg_attrs ins(%a_buf_view, %b_buf_view, %scalar : memref<8xf32, #map>, memref<8xf32, #map>, f32) linalg.generic #linalg_attrs ins(%a_buf_view, %b_buf_view, %scalar : memref<8xf32, #map>, memref<8xf32, #map>, f32)
outs(%C_view: memref<8xf32, #map>) { inits(%C_view: memref<8xf32, #map>) {
^bb0(%a: f32, %b: f32, %s: f32, %c: f32): ^bb0(%a: f32, %b: f32, %s: f32, %c: f32):
%add = arith.addf %a, %b : f32 %add = arith.addf %a, %b : f32
%accum = arith.addf %add, %c : f32 %accum = arith.addf %add, %c : f32
%accum1 = arith.addf %scalar, %accum : f32 %accum1 = arith.addf %scalar, %accum : f32
%accum2 = arith.addf %s, %accum1 : f32 %accum2 = arith.addf %s, %accum1 : f32
linalg.yield %accum2 : f32 linalg.yield %accum2 : f32
} }
scf.yield scf.yield
} { __test_pipelining_loop__ } } { __test_pipelining_loop__ }
return return
} }

mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir

Show First 20 Lines • Show All 147 Lines • ▼ Show 20 Lines %r, %v3 = scf.for %i = %c0 to %s step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) {
%v2 = vector.transfer_read %e[%s], %cst : tensor<?xf32>, vector<5xf32> %v2 = vector.transfer_read %e[%s], %cst : tensor<?xf32>, vector<5xf32>
scf.yield %t2, %v2 : tensor<?xf32>, vector<5xf32> scf.yield %t2, %v2 : tensor<?xf32>, vector<5xf32>
} }
// CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true", "none"]} // CHECK: } {__inplace_operands_attr__ = ["none", "none", "none", "true", "none"]}
// Use %t3 in some way without reading it, so that it does not get DCE'd. // Use %t3 in some way without reading it, so that it does not get DCE'd.
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: __inplace_operands_attr__ = ["true"] // CHECK-SAME: __inplace_operands_attr__ = ["true"]
%o = linalg.generic #trait outs (%t3 : tensor<?xf32>) { %o = linalg.generic #trait inits (%t3 : tensor<?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> (tensor<?xf32>) } -> (tensor<?xf32>)
return %o, %v3 : tensor<?xf32>, vector<5xf32> return %o, %v3 : tensor<?xf32>, vector<5xf32>
} }
// ----- // -----
Show All 21 Linesfunc.func @non_reading_scf_for(%t1: tensor<?xf32> {bufferization.writable = true},
%t3 = vector.transfer_write %v, %t1[%s] : vector<5xf32>, tensor<?xf32> %t3 = vector.transfer_write %v, %t1[%s] : vector<5xf32>, tensor<?xf32>
// This loop does not read from %t1. It only writes to it. // This loop does not read from %t1. It only writes to it.
// CHECK: scf.for // CHECK: scf.for
%r, %v3 = scf.for %i = %c0 to %c10 step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) { %r, %v3 = scf.for %i = %c0 to %c10 step %c1 iter_args(%t2 = %t1, %v0 = %v) -> (tensor<?xf32>, vector<5xf32>) {
// Write to %t1 via %t2. (Overwrite %t3.) // Write to %t1 via %t2. (Overwrite %t3.)
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: __inplace_operands_attr__ = ["true"] // CHECK-SAME: __inplace_operands_attr__ = ["true"]
%o2 = linalg.generic #trait outs (%t2 : tensor<?xf32>) { %o2 = linalg.generic #trait inits (%t2 : tensor<?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> (tensor<?xf32>) } -> (tensor<?xf32>)
// Read overwritten value. This is not a read of %t1. // Read overwritten value. This is not a read of %t1.
%v2 = vector.transfer_read %o2[%s], %cst : tensor<?xf32>, vector<5xf32> %v2 = vector.transfer_read %o2[%s], %cst : tensor<?xf32>, vector<5xf32>
scf.yield %o2, %v2 : tensor<?xf32>, vector<5xf32> scf.yield %o2, %v2 : tensor<?xf32>, vector<5xf32>
} }
// Use %t3 in some way without reading it, so that it does not get DCE'd. // Use %t3 in some way without reading it, so that it does not get DCE'd.
// CHECK: linalg.generic // CHECK: linalg.generic
// CHECK-SAME: __inplace_operands_attr__ = ["true"] // CHECK-SAME: __inplace_operands_attr__ = ["true"]
%o = linalg.generic #trait outs (%t3 : tensor<?xf32>) { %o = linalg.generic #trait inits (%t3 : tensor<?xf32>) {
^bb(%0: f32) : ^bb(%0: f32) :
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> (tensor<?xf32>) } -> (tensor<?xf32>)
// CHECK: return // CHECK: return
// CHECK-SAME: __equivalent_func_args__ = [0, -1] // CHECK-SAME: __equivalent_func_args__ = [0, -1]
return %o, %v3 : tensor<?xf32>, vector<5xf32> return %o, %v3 : tensor<?xf32>, vector<5xf32>
} }
▲ Show 20 Lines • Show All 399 Lines • ▼ Show 20 Linesfunc.func @same_enclosing_repetitive_region(%2: tensor<320xf32>,
%cst = arith.constant -0.000000e+00 : f32 %cst = arith.constant -0.000000e+00 : f32
%c320 = arith.constant 320 : index %c320 = arith.constant 320 : index
%4 = scf.foreach_thread (%arg0) in (%c320) shared_outs(%arg1 = %2) -> (tensor<320xf32>) { %4 = scf.foreach_thread (%arg0) in (%c320) shared_outs(%arg1 = %2) -> (tensor<320xf32>) {
// CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]} // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
%5 = tensor.extract_slice %3[%arg0, 0] [1, 10240] [1, 1] : tensor<320x10240xf32> to tensor<1x10240xf32> %5 = tensor.extract_slice %3[%arg0, 0] [1, 10240] [1, 1] : tensor<320x10240xf32> to tensor<1x10240xf32>
// CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]} // CHECK: tensor.extract_slice {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
%6 = tensor.extract_slice %arg1[%arg0] [1] [1] : tensor<320xf32> to tensor<1xf32> %6 = tensor.extract_slice %arg1[%arg0] [1] [1] : tensor<320xf32> to tensor<1xf32>
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
%7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<1xf32>) -> tensor<1xf32> %7 = linalg.fill ins(%cst : f32) inits(%6 : tensor<1xf32>) -> tensor<1xf32>
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
%8 = linalg.fill ins(%cst : f32) outs(%7 : tensor<1xf32>) -> tensor<1xf32> %8 = linalg.fill ins(%cst : f32) inits(%7 : tensor<1xf32>) -> tensor<1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
// CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]} // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]}
tensor.parallel_insert_slice %8 into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32> tensor.parallel_insert_slice %8 into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32>
} }
} {thread_dim_mapping = []} } {thread_dim_mapping = []}
return %4 : tensor<320xf32> return %4 : tensor<320xf32>
} }
// ----- // -----
// CHECK-LABEL: different_repetitive_region_via_alias // CHECK-LABEL: different_repetitive_region_via_alias
func.func @different_repetitive_region_via_alias(%arg0: tensor<4xf32>, func.func @different_repetitive_region_via_alias(%arg0: tensor<4xf32>,
%arg1: tensor<4xf32>, %arg1: tensor<4xf32>,
%arg2: index, %arg2: index,
%arg3: index, %arg3: index,
%arg4: index) %arg4: index)
-> (tensor<4xf32>) -> (tensor<4xf32>)
{ {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst2 = arith.constant 1.000000e+00 : f32 %cst2 = arith.constant 1.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<4xf32> %0 = bufferization.alloc_tensor() : tensor<4xf32>
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<4xf32>) -> tensor<4xf32>
%2 = scf.for %arg5 = %arg2 to %arg3 step %arg4 iter_args(%arg6 = %arg1) -> (tensor<4xf32>) { %2 = scf.for %arg5 = %arg2 to %arg3 step %arg4 iter_args(%arg6 = %arg1) -> (tensor<4xf32>) {
// CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]} // CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
%4 = tensor.extract %1[%arg4] : tensor<4xf32> %4 = tensor.extract %1[%arg4] : tensor<4xf32>
vector.print %4 : f32 vector.print %4 : f32
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
%5 = linalg.fill ins(%cst2 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32> %5 = linalg.fill ins(%cst2 : f32) inits(%0 : tensor<4xf32>) -> tensor<4xf32>
scf.yield %5 : tensor<4xf32> scf.yield %5 : tensor<4xf32>
} }
return %2 : tensor<4xf32> return %2 : tensor<4xf32>
} }
// ----- // -----
// CHECK-LABEL: no_raw_conflict_after_repetitive_use // CHECK-LABEL: no_raw_conflict_after_repetitive_use
func.func @no_raw_conflict_after_repetitive_use(%arg0: tensor<4xf32>, func.func @no_raw_conflict_after_repetitive_use(%arg0: tensor<4xf32>,
%arg1: tensor<4xf32>, %arg1: tensor<4xf32>,
%arg2: index, %arg2: index,
%arg3: index, %arg3: index,
%arg4: index) %arg4: index)
-> (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>)
{ {
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%cst2 = arith.constant 1.000000e+00 : f32 %cst2 = arith.constant 1.000000e+00 : f32
%cst3 = arith.constant 2.000000e+00 : f32 %cst3 = arith.constant 2.000000e+00 : f32
%0 = bufferization.alloc_tensor() : tensor<4xf32> %0 = bufferization.alloc_tensor() : tensor<4xf32>
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
%1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32> %1 = linalg.fill ins(%cst : f32) inits(%0 : tensor<4xf32>) -> tensor<4xf32>
%2 = scf.for %arg5 = %arg2 to %arg3 step %arg4 iter_args(%arg6 = %arg1) -> (tensor<4xf32>) { %2 = scf.for %arg5 = %arg2 to %arg3 step %arg4 iter_args(%arg6 = %arg1) -> (tensor<4xf32>) {
// CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]} // CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
%4 = tensor.extract %1[%arg4] : tensor<4xf32> %4 = tensor.extract %1[%arg4] : tensor<4xf32>
vector.print %4 : f32 vector.print %4 : f32
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}
%5 = linalg.fill ins(%cst2 : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %5 = linalg.fill ins(%cst2 : f32) inits(%1 : tensor<4xf32>) -> tensor<4xf32>
scf.yield %5 : tensor<4xf32> scf.yield %5 : tensor<4xf32>
} }
// The following is *not* a RaW conflict. // The following is *not* a RaW conflict.
// CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]} // CHECK: tensor.extract {{.*}} {__inplace_operands_attr__ = ["true", "none"]}
%6 = tensor.extract %1[%arg4] : tensor<4xf32> %6 = tensor.extract %1[%arg4] : tensor<4xf32>
vector.print %6 : f32 vector.print %6 : f32
// CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]} // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
%7 = linalg.fill ins(%cst3 : f32) outs(%1 : tensor<4xf32>) -> tensor<4xf32> %7 = linalg.fill ins(%cst3 : f32) inits(%1 : tensor<4xf32>) -> tensor<4xf32>
return %2, %7 : tensor<4xf32>, tensor<4xf32> return %2, %7 : tensor<4xf32>, tensor<4xf32>
} }

mlir/test/Dialect/SCF/one-shot-bufferize-tensor-copy-insertion.mlir

Show First 20 Lines • Show All 112 Lines • ▼ Show 20 Lines
// CHECK-FUNC-LABEL: func @scf_foreach_thread_out_of_place( // CHECK-FUNC-LABEL: func @scf_foreach_thread_out_of_place(
func.func @scf_foreach_thread_out_of_place(%in: tensor<100xf32>, func.func @scf_foreach_thread_out_of_place(%in: tensor<100xf32>,
%out: tensor<100xf32>) { %out: tensor<100xf32>) {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%num_threads = arith.constant 100 : index %num_threads = arith.constant 100 : index
// CHECK-FUNC-NOT: alloc_tensor // CHECK-FUNC-NOT: alloc_tensor
// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {bufferization.escape = [false]} : tensor<100xf32> // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[arg1]]) {bufferization.escape = [false]} : tensor<100xf32>
// CHECK: scf.foreach_thread {{.*}} shared_outs(%[[o:.*]] = %[[alloc]]) // CHECK: scf.foreach_thread {{.*}} shared_inits(%[[o:.*]] = %[[alloc]])
%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> { %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_inits(%o = %out) -> tensor<100xf32> {
// CHECK: tensor.extract_slice // CHECK: tensor.extract_slice
// CHECK: scf.foreach_thread.perform_concurrently // CHECK: scf.foreach_thread.perform_concurrently
// CHECK: tensor.parallel_insert_slice %{{.*}} into %[[o]] // CHECK: tensor.parallel_insert_slice %{{.*}} into %[[o]]
%1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
tensor<1xf32> into tensor<100xf32> tensor<1xf32> into tensor<100xf32>
} }
// CHECK: } {thread_dim_mapping = [5]} // CHECK: } {thread_dim_mapping = [5]}
} {thread_dim_mapping = [5]} } {thread_dim_mapping = [5]}
return return
} }

mlir/test/Dialect/SCF/one-shot-bufferize.mlir

Show First 20 Lines • Show All 48 Lines • ▼ Show 20 Linesfunc.func @scf_for_is_reading(%A : tensor<?xf32>, %B : tensor<?xf32>,
// This is a regression test to make sure that an alloc + copy is emitted. // This is a regression test to make sure that an alloc + copy is emitted.
// CHECK: %[[alloc:.*]] = memref.alloc // CHECK: %[[alloc:.*]] = memref.alloc
// CHECK: memref.copy %[[A]], %[[alloc]] // CHECK: memref.copy %[[A]], %[[alloc]]
// CHECK: %[[clone:.*]] = bufferization.clone %[[alloc]] // CHECK: %[[clone:.*]] = bufferization.clone %[[alloc]]
// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[clone]]) // CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[clone]])
%0 = scf.for %iv = %lb to %ub step %c1 iter_args(%1 = %A) -> tensor<?xf32> { %0 = scf.for %iv = %lb to %ub step %c1 iter_args(%1 = %A) -> tensor<?xf32> {
%r = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32> %r = linalg.fill ins(%cst : f32) inits(%1 : tensor<?xf32>) -> tensor<?xf32>
scf.yield %B : tensor<?xf32> scf.yield %B : tensor<?xf32>
} }
%1 = tensor.extract %0[%c1] : tensor<?xf32> %1 = tensor.extract %0[%c1] : tensor<?xf32>
%2 = tensor.extract %A[%c1] : tensor<?xf32> %2 = tensor.extract %A[%c1] : tensor<?xf32>
return %1, %2 : f32, f32 return %1, %2 : f32, f32
} }
// ----- // -----
▲ Show 20 Lines • Show All 475 Lines • ▼ Show 20 Linesfunc.func @parallel_insert_slice_no_conflict(
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
// CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]]) // CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]])
%2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) { %2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) {
// CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1] // CHECK: %[[subview:.*]] = memref.subview %[[arg2]][5] [%[[idx]]] [1]
%6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32> %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
// CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview]] : memref<?xf32 // CHECK: linalg.fill ins(%{{.*}}) inits(%[[subview]] : memref<?xf32
%8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32> %8 = linalg.fill ins(%cst : f32) inits(%6 : tensor<?xf32>) -> tensor<?xf32>
// Self-copy will DCE away later. // Self-copy will DCE away later.
// CHECK: memref.copy %[[subview]], %[[subview]] // CHECK: memref.copy %[[subview]], %[[subview]]
// Empty terminator is elided from pretty-printing. // Empty terminator is elided from pretty-printing.
// CHECK-NOT: scf.foreach_thread.perform_concurrently // CHECK-NOT: scf.foreach_thread.perform_concurrently
// CHECK-NOT: parallel_insert_slice // CHECK-NOT: parallel_insert_slice
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %8 into %o[5] [%idx] [%c1] : tensor.parallel_insert_slice %8 into %o[5] [%idx] [%c1] :
Show All 10 Lines
// ----- // -----
// CHECK-LABEL: func.func @parallel_insert_slice_with_conflict( // CHECK-LABEL: func.func @parallel_insert_slice_with_conflict(
// CHECK-SAME: %[[idx:.*]]: index, %[[idx2:.*]]: index, // CHECK-SAME: %[[idx:.*]]: index, %[[idx2:.*]]: index,
// CHECK-SAME: %[[arg1:.*]]: memref<?xf32, strided{{.*}}>, // CHECK-SAME: %[[arg1:.*]]: memref<?xf32, strided{{.*}}>,
// CHECK-SAME: %[[arg2:.*]]: memref<?xf32, strided{{.*}}> // CHECK-SAME: %[[arg2:.*]]: memref<?xf32, strided{{.*}}>
func.func @parallel_insert_slice_with_conflict( func.func @parallel_insert_slice_with_conflict(
%idx: index, %idx: index,
%idx2: index, %idx2: index,
%arg1: tensor<?xf32> {bufferization.writable = true}, %arg1: tensor<?xf32> {bufferization.writable = true},
%arg2: tensor<?xf32> {bufferization.writable = true}) -> (f32, f32) %arg2: tensor<?xf32> {bufferization.writable = true}) -> (f32, f32)
{ {
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
// The parallel_insert_slice_op bufferizes out-of-place due to a RAW conflict // The parallel_insert_slice_op bufferizes out-of-place due to a RAW conflict
// on %arg2, so we need an allocation. // on %arg2, so we need an allocation.
// CHECK: %[[alloc1:.*]] = memref.alloc // CHECK: %[[alloc1:.*]] = memref.alloc
// CHECK: memref.copy %[[arg2]], %[[alloc1]] // CHECK: memref.copy %[[arg2]], %[[alloc1]]
// CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]]) // CHECK: scf.foreach_thread (%[[tidx:.*]]) in (%[[idx2]])
%2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) { %2 = scf.foreach_thread (%arg3) in (%idx2) shared_outs(%o = %arg2) -> (tensor<?xf32>) {
// CHECK: %[[subview1:.*]] = memref.subview %[[alloc1]][5] [%[[idx]]] [1] // CHECK: %[[subview1:.*]] = memref.subview %[[alloc1]][5] [%[[idx]]] [1]
%6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32> %6 = tensor.extract_slice %o[5] [%idx] [%c1] : tensor<?xf32> to tensor<?xf32>
// CHECK: linalg.fill ins(%{{.*}}) outs(%[[subview1]] : memref<?xf32 // CHECK: linalg.fill ins(%{{.*}}) inits(%[[subview1]] : memref<?xf32
%8 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32> %8 = linalg.fill ins(%cst : f32) inits(%6 : tensor<?xf32>) -> tensor<?xf32>
// Now the copy of the actual insert_slice. (It will fold away.) // Now the copy of the actual insert_slice. (It will fold away.)
// CHECK: memref.copy %[[subview1]], %[[subview1]] // CHECK: memref.copy %[[subview1]], %[[subview1]]
// Empty terminator is elided from pretty-printing. // Empty terminator is elided from pretty-printing.
// CHECK-NOT: scf.foreach_thread.perform_concurrently // CHECK-NOT: scf.foreach_thread.perform_concurrently
// CHECK-NOT: parallel_insert_slice // CHECK-NOT: parallel_insert_slice
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
Show All 25 Linesfunc.func @matmul(%arg0: tensor<8x8xf32>, %arg1: tensor<8x8xf32>, %arg2: tensor<8x8xf32> {bufferization.writable = true}) -> tensor<8x8xf32> {
// CHECK: scf.foreach_thread {{.*}} // CHECK: scf.foreach_thread {{.*}}
%0 = scf.foreach_thread (%arg3, %arg4) in (%c2, %c4) shared_outs(%o = %arg2) -> (tensor<8x8xf32>) { %0 = scf.foreach_thread (%arg3, %arg4) in (%c2, %c4) shared_outs(%o = %arg2) -> (tensor<8x8xf32>) {
%1 = affine.apply #map0(%arg3) %1 = affine.apply #map0(%arg3)
%3 = tensor.extract_slice %arg0[%1, 0] [4, 8] [1, 1] : tensor<8x8xf32> to tensor<4x8xf32> %3 = tensor.extract_slice %arg0[%1, 0] [4, 8] [1, 1] : tensor<8x8xf32> to tensor<4x8xf32>
%4 = affine.apply #map1(%arg4) %4 = affine.apply #map1(%arg4)
%6 = tensor.extract_slice %arg1[0, %4] [8, 4] [1, 1] : tensor<8x8xf32> to tensor<8x4xf32> %6 = tensor.extract_slice %arg1[0, %4] [8, 4] [1, 1] : tensor<8x8xf32> to tensor<8x4xf32>
%7 = tensor.extract_slice %o[%1, %4] [4, 4] [1, 1] : tensor<8x8xf32> to tensor<4x4xf32> %7 = tensor.extract_slice %o[%1, %4] [4, 4] [1, 1] : tensor<8x8xf32> to tensor<4x4xf32>
// CHECK: linalg.matmul ins({{.*}}memref<4x8xf32, strided<[?, ?], offset: ?>>, memref<8x4xf32, strided<[?, ?], offset: ?>>) outs({{.*}} : memref<4x4xf32, strided<[?, ?], offset: ?>>) // CHECK: linalg.matmul ins({{.*}}memref<4x8xf32, strided<[?, ?], offset: ?>>, memref<8x4xf32, strided<[?, ?], offset: ?>>) inits({{.*}} : memref<4x4xf32, strided<[?, ?], offset: ?>>)
%8 = linalg.matmul ins(%3, %6 : tensor<4x8xf32>, tensor<8x4xf32>) outs(%7 : tensor<4x4xf32>) -> tensor<4x4xf32> %8 = linalg.matmul ins(%3, %6 : tensor<4x8xf32>, tensor<8x4xf32>) inits(%7 : tensor<4x4xf32>) -> tensor<4x4xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %8 into %o[%1, %4] [4, 4] [1, 1] : tensor<4x4xf32> into tensor<8x8xf32> tensor.parallel_insert_slice %8 into %o[%1, %4] [4, 4] [1, 1] : tensor<4x4xf32> into tensor<8x8xf32>
} }
} }
return %0 : tensor<8x8xf32> return %0 : tensor<8x8xf32>
} }
// ----- // -----
▲ Show 20 Lines • Show All 218 Lines • ▼ Show 20 Linesfunc.func @scf_while_buffer_type_mismatch(%sz: index, %sz2: index) -> f32 {
} }
%x = tensor.extract %r[%c1] : tensor<?xf32> %x = tensor.extract %r[%c1] : tensor<?xf32>
return %x : f32 return %x : f32
} }
// ----- // -----
// CHECK-LABEL: func @non_tensor_for_arg // CHECK-LABEL: func @non_tensor_for_arg
func.func @non_tensor_for_arg(%A : tensor<?xf32> {bufferization.writable = true}) func.func @non_tensor_for_arg(%A : tensor<?xf32> {bufferization.writable = true})
-> tensor<?xf32> { -> tensor<?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.0 : f32 %c2 = arith.constant 2.0 : f32
%c10 = arith.constant 10 : index %c10 = arith.constant 10 : index
%r1:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%idx = %c1, %t = %A) -> (index, tensor<?xf32>) { %r1:2 = scf.for %i = %c0 to %c10 step %c1 iter_args(%idx = %c1, %t = %A) -> (index, tensor<?xf32>) {
%t2 = tensor.insert %c2 into %t[%idx] : tensor<?xf32> %t2 = tensor.insert %c2 into %t[%idx] : tensor<?xf32>
scf.yield %idx, %t2 : index, tensor<?xf32> scf.yield %idx, %t2 : index, tensor<?xf32>
} }
return %r1#1 : tensor<?xf32> return %r1#1 : tensor<?xf32>
} }
// ----- // -----
// This is a regression test. Just check that the IR bufferizes. // This is a regression test. Just check that the IR bufferizes.
// CHECK-LABEL: func @buffer_type_of_collapse_shape // CHECK-LABEL: func @buffer_type_of_collapse_shape
func.func @buffer_type_of_collapse_shape(%arg0: tensor<f64>) { func.func @buffer_type_of_collapse_shape(%arg0: tensor<f64>) {
%true = arith.constant true %true = arith.constant true
%0 = scf.while (%arg1 = %arg0) : (tensor<f64>) -> (tensor<f64>) { %0 = scf.while (%arg1 = %arg0) : (tensor<f64>) -> (tensor<f64>) {
scf.condition(%true) %arg1 : tensor<f64> scf.condition(%true) %arg1 : tensor<f64>
} do { } do {
^bb0(%_: tensor<f64>): ^bb0(%_: tensor<f64>):
%3 = bufferization.alloc_tensor() : tensor<1xf64> %3 = bufferization.alloc_tensor() : tensor<1xf64>
%16 = tensor.collapse_shape %3 [] : tensor<1xf64> into tensor<f64> %16 = tensor.collapse_shape %3 [] : tensor<1xf64> into tensor<f64>
scf.yield %16 : tensor<f64> scf.yield %16 : tensor<f64>
} }
return return
} }
// ----- // -----
// This is a regression test. Just check that the IR bufferizes. // This is a regression test. Just check that the IR bufferizes.
// CHECK-LABEL: func @non_block_argument_yield // CHECK-LABEL: func @non_block_argument_yield
func.func @non_block_argument_yield() { func.func @non_block_argument_yield() {
%true = arith.constant true %true = arith.constant true
%0 = bufferization.alloc_tensor() : tensor<i32> %0 = bufferization.alloc_tensor() : tensor<i32>
%1 = scf.while (%arg0 = %0) : (tensor<i32>) -> (tensor<i32>) { %1 = scf.while (%arg0 = %0) : (tensor<i32>) -> (tensor<i32>) {
scf.condition(%true) %arg0 : tensor<i32> scf.condition(%true) %arg0 : tensor<i32>
} do { } do {
^bb0(%arg0: tensor<i32>): ^bb0(%arg0: tensor<i32>):
%ret = scf.while (%arg1 = %0) : (tensor<i32>) -> (tensor<i32>) { %ret = scf.while (%arg1 = %0) : (tensor<i32>) -> (tensor<i32>) {
scf.condition(%true) %arg1 : tensor<i32> scf.condition(%true) %arg1 : tensor<i32>
} do { } do {
^bb0(%arg7: tensor<i32>): ^bb0(%arg7: tensor<i32>):
scf.yield %0 : tensor<i32> scf.yield %0 : tensor<i32>
} }
scf.yield %ret : tensor<i32> scf.yield %ret : tensor<i32>
} }
return return
} }

mlir/test/Dialect/SCF/ops.mlir

Show First 20 Lines • Show All 317 Lines • ▼ Show 20 Linesfunc.func @simple_example(%in: tensor<100xf32>, %out: tensor<100xf32>) {
// CHECK: scf.foreach_thread // CHECK: scf.foreach_thread
// CHECK-NEXT: tensor.extract_slice // CHECK-NEXT: tensor.extract_slice
// CHECK-NEXT: scf.foreach_thread.perform_concurrently // CHECK-NEXT: scf.foreach_thread.perform_concurrently
// CHECK-NEXT: tensor.parallel_insert_slice // CHECK-NEXT: tensor.parallel_insert_slice
// CHECK-NEXT: } // CHECK-NEXT: }
// CHECK-NEXT: } // CHECK-NEXT: }
// CHECK-NEXT: return // CHECK-NEXT: return
%result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_outs(%o = %out) -> tensor<100xf32> { %result = scf.foreach_thread (%thread_idx) in (%num_threads) shared_inits(%o = %out) -> tensor<100xf32> {
%1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32> %1 = tensor.extract_slice %in[%thread_idx][1][1] : tensor<100xf32> to tensor<1xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] : tensor.parallel_insert_slice %1 into %o[%thread_idx][1][1] :
tensor<1xf32> into tensor<100xf32> tensor<1xf32> into tensor<100xf32>
} }
} }
return return
} }
▲ Show 20 Lines • Show All 47 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/codegen.mlir

Show First 20 Lines • Show All 280 Lines • ▼ Show 20 Lines
// CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex> // CHECK: memref.store %[[A]], %[[T0]][%[[C0]]] : memref<2xindex>
// CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex> // CHECK: memref.store %[[C10]], %[[T0]][%[[C1]]] : memref<2xindex>
// CHECK: %[[T2:.*]] = memref.alloc() : memref<1xindex> // CHECK: %[[T2:.*]] = memref.alloc() : memref<1xindex>
// CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<1xindex> to memref<?xindex> // CHECK: %[[T3:.*]] = memref.cast %[[T2]] : memref<1xindex> to memref<?xindex>
// CHECK: %[[T4:.*]] = memref.alloc() : memref<1xindex> // CHECK: %[[T4:.*]] = memref.alloc() : memref<1xindex>
// CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<1xindex> to memref<?xindex> // CHECK: %[[T5:.*]] = memref.cast %[[T4]] : memref<1xindex> to memref<?xindex>
// CHECK: %[[T6:.*]] = memref.alloc() : memref<1xf64> // CHECK: %[[T6:.*]] = memref.alloc() : memref<1xf64>
// CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<1xf64> to memref<?xf64> // CHECK: %[[T7:.*]] = memref.cast %[[T6]] : memref<1xf64> to memref<?xf64>
// CHECK: linalg.fill ins(%[[C0]] : index) outs(%[[T1]] : memref<3xindex>) // CHECK: linalg.fill ins(%[[C0]] : index) inits(%[[T1]] : memref<3xindex>)
// CHECK: return %[[T0]], %[[T1]], %[[T3]], %[[T5]], %[[T7]] // CHECK: return %[[T0]], %[[T1]], %[[T3]], %[[T5]], %[[T7]]
func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> { func.func @sparse_alloc_csc(%arg0: index) -> tensor<10x?xf64, #CSC> {
%0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC> %0 = bufferization.alloc_tensor(%arg0) : tensor<10x?xf64, #CSC>
%1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC> %1 = sparse_tensor.load %0 : tensor<10x?xf64, #CSC>
return %1 : tensor<10x?xf64, #CSC> return %1 : tensor<10x?xf64, #CSC>
} }
// CHECK-LABEL: func @sparse_alloc_3d() -> // CHECK-LABEL: func @sparse_alloc_3d() ->
Show All 20 Linesfunc.func @sparse_alloc_3d() -> tensor<10x20x30xf64, #Dense3D> {
return %1 : tensor<10x20x30xf64, #Dense3D> return %1 : tensor<10x20x30xf64, #Dense3D>
} }
// CHECK-LABEL: func.func @sparse_expansion1() // CHECK-LABEL: func.func @sparse_expansion1()
// CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64> // CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
// CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1> // CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1>
// CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex> // CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex>
// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex> // CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<8xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<8xf64>)
// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<8xi1>) // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<8xi1>)
// CHECK: return %[[D]] : memref<?xindex> // CHECK: return %[[D]] : memref<?xindex>
func.func @sparse_expansion1() -> memref<?xindex> { func.func @sparse_expansion1() -> memref<?xindex> {
%0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR> %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
// CHECK-LABEL: func.func @sparse_expansion2() // CHECK-LABEL: func.func @sparse_expansion2()
// CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64> // CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64>
// CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1> // CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1>
// CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex> // CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex>
// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex> // CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<4xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<4xf64>)
// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<4xi1>) // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<4xi1>)
// CHECK: return %[[D]] : memref<?xindex> // CHECK: return %[[D]] : memref<?xindex>
func.func @sparse_expansion2() -> memref<?xindex> { func.func @sparse_expansion2() -> memref<?xindex> {
%0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC> %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
// CHECK-LABEL: func.func @sparse_expansion3( // CHECK-LABEL: func.func @sparse_expansion3(
// CHECK-SAME: %[[D0:.*]]: index, // CHECK-SAME: %[[D0:.*]]: index,
// CHECK-SAME: %{{.*}}: index) -> memref<?xindex> { // CHECK-SAME: %{{.*}}: index) -> memref<?xindex> {
// CHECK: %[[C1:.*]] = arith.constant 1 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex> // CHECK: %[[S0:.*]] = memref.alloc() : memref<2xindex>
// CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex> // CHECK: memref.store %[[D0]], %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
// CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex> // CHECK: %[[D1:.*]] = memref.load %[[S0]]{{\[}}%[[C1]]] : memref<2xindex>
// CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64> // CHECK: %[[V:.*]] = memref.alloc(%[[D1]]) : memref<?xf64>
// CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1> // CHECK: %[[B:.*]] = memref.alloc(%[[D1]]) : memref<?xi1>
// CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex> // CHECK: %[[D:.*]] = memref.alloc(%[[D1]]) : memref<?xindex>
// CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[V]] : memref<?xf64>) // CHECK: linalg.fill ins(%{{.*}} : f64) inits(%[[V]] : memref<?xf64>)
// CHECK: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) // CHECK: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<?xi1>)
// CHECK: return %[[D]] : memref<?xindex> // CHECK: return %[[D]] : memref<?xindex>
func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> { func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
%0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC> %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
▲ Show 20 Lines • Show All 176 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/conversion.mlir

Show First 20 Lines • Show All 315 Lines • ▼ Show 20 Lines
} }
// CHECK-LABEL: func @sparse_expansion1() // CHECK-LABEL: func @sparse_expansion1()
// CHECK: %[[N:.*]] = call @newSparseTensor // CHECK: %[[N:.*]] = call @newSparseTensor
// CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64> // CHECK: %[[A:.*]] = memref.alloc() : memref<8xf64>
// CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1> // CHECK: %[[B:.*]] = memref.alloc() : memref<8xi1>
// CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex> // CHECK: %[[C:.*]] = memref.alloc() : memref<8xindex>
// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex> // CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<8xindex> to memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<8xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<8xf64>)
// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<8xi1>) // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<8xi1>)
// CHECK: return %[[D]] : memref<?xindex> // CHECK: return %[[D]] : memref<?xindex>
func.func @sparse_expansion1() -> memref<?xindex> { func.func @sparse_expansion1() -> memref<?xindex> {
%0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR> %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSR>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<4x8xf64, #CSR> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
// CHECK-LABEL: func @sparse_expansion2() // CHECK-LABEL: func @sparse_expansion2()
// CHECK: %[[N:.*]] = call @newSparseTensor // CHECK: %[[N:.*]] = call @newSparseTensor
// CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64> // CHECK: %[[A:.*]] = memref.alloc() : memref<4xf64>
// CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1> // CHECK: %[[B:.*]] = memref.alloc() : memref<4xi1>
// CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex> // CHECK: %[[C:.*]] = memref.alloc() : memref<4xindex>
// CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex> // CHECK: %[[D:.*]] = memref.cast %[[C]] : memref<4xindex> to memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<4xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<4xf64>)
// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<4xi1>) // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<4xi1>)
// CHECK: return %[[D]] : memref<?xindex> // CHECK: return %[[D]] : memref<?xindex>
func.func @sparse_expansion2() -> memref<?xindex> { func.func @sparse_expansion2() -> memref<?xindex> {
%0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC> %0 = bufferization.alloc_tensor() : tensor<4x8xf64, #CSC>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
// CHECK-LABEL: func @sparse_expansion3( // CHECK-LABEL: func @sparse_expansion3(
// CHECK: %[[C1:.*]] = arith.constant 1 : index // CHECK: %[[C1:.*]] = arith.constant 1 : index
// CHECK: %[[N:.*]] = call @newSparseTensor // CHECK: %[[N:.*]] = call @newSparseTensor
// CHECK: %[[S:.*]] = call @sparseDimSize(%[[N]], %c1) : (!llvm.ptr<i8>, index) -> index // CHECK: %[[S:.*]] = call @sparseDimSize(%[[N]], %c1) : (!llvm.ptr<i8>, index) -> index
// CHECK: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64> // CHECK: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
// CHECK: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1> // CHECK: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
// CHECK: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex> // CHECK: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<?xf64>)
// CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) // CHECK-DAG: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<?xi1>)
// CHECK: return %[[C]] : memref<?xindex> // CHECK: return %[[C]] : memref<?xindex>
func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> { func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
%0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC> %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>
%values, %filled, %added, %count = sparse_tensor.expand %0 %values, %filled, %added, %count = sparse_tensor.expand %0
: tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex> : tensor<?x?xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
return %added : memref<?xindex> return %added : memref<?xindex>
} }
▲ Show 20 Lines • Show All 66 Lines • Show Last 20 Lines

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

Show All 28 Lines
// CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex> // CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
// CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
// CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<13xi32> // CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<13xi32>
// CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) outs(%[[M]] : memref<13xi32>) // CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) inits(%[[M]] : memref<13xi32>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<13xi32> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<13xi32>
// CHECK: scf.yield // CHECK: scf.yield
Show All 21 Lines
// CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex> // CHECK-DAG: memref.store %[[I0]], %[[PermS]][%[[I0]]] : memref<1xindex>
// CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[zeroI32:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ElemTpActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %[[zeroI32]], %[[zeroI32]], %[[ElemTpActionToIter]], %[[ElemTpActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<1xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<1xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<i32>
// CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?xi32> // CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?xi32>
// CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) outs(%[[M]] : memref<?xi32>) // CHECK-DAG: linalg.fill ins(%[[zeroI32]] : i32) inits(%[[M]] : memref<?xi32>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextI32(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<i32>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<1xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<i32>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<?xi32> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]]] : memref<?xi32>
// CHECK: scf.yield // CHECK: scf.yield
Show All 26 Lines
// CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex> // CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex>
// CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64>
// CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<2x4xf64> // CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<2x4xf64>
// CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: linalg.fill ins(%[[E0]] : f64) outs(%[[M]] : memref<2x4xf64>) // CHECK-DAG: linalg.fill ins(%[[E0]] : f64) inits(%[[M]] : memref<2x4xf64>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
// CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex> // CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<2x4xf64> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<2x4xf64>
Show All 27 Lines
// CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex> // CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex>
// CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64>
// CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?x4xf64> // CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]]) : memref<?x4xf64>
// CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: linalg.fill ins(%[[E0]] : f64) outs(%[[M]] : memref<?x4xf64>) // CHECK-DAG: linalg.fill ins(%[[E0]] : f64) inits(%[[M]] : memref<?x4xf64>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
// CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex> // CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<?x4xf64> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<?x4xf64>
Show All 27 Lines
// CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex> // CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex>
// CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64>
// CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI1]]) : memref<2x?xf64> // CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI1]]) : memref<2x?xf64>
// CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: linalg.fill ins(%[[E0]] : f64) outs(%[[M]] : memref<2x?xf64>) // CHECK-DAG: linalg.fill ins(%[[E0]] : f64) inits(%[[M]] : memref<2x?xf64>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
// CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex> // CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<2x?xf64> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<2x?xf64>
Show All 27 Lines
// CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex> // CHECK-DAG: memref.store %[[I1]], %[[PermS]][%[[I1]]] : memref<2xindex>
// CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<2xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<2xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64>
// CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]], %[[SizeI1]]) : memref<?x?xf64> // CHECK-DAG: %[[M:.*]] = memref.alloc(%[[SizeI0]], %[[SizeI1]]) : memref<?x?xf64>
// CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: linalg.fill ins(%[[E0]] : f64) outs(%[[M]] : memref<?x?xf64>) // CHECK-DAG: linalg.fill ins(%[[E0]] : f64) inits(%[[M]] : memref<?x?xf64>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<2xindex>
// CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex> // CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<2xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
// CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<?x?xf64> // CHECK: memref.store %[[ElemVal]], %[[M]][%[[Iv0]], %[[Iv1]]] : memref<?x?xf64>
Show All 31 Lines
// CHECK-DAG: memref.store %[[I2]], %[[PermS]][%[[I2]]] : memref<3xindex> // CHECK-DAG: memref.store %[[I2]], %[[PermS]][%[[I2]]] : memref<3xindex>
// CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32 // CHECK-DAG: %[[ActionToIter:.*]] = arith.constant 6 : i32
// CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK-DAG: %[[Iter:.*]] = call @newSparseTensor(%[[AttrsD]], %[[SizesD]], %[[PermD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ActionToIter]], %[[Arg]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<3xindex> // CHECK-DAG: %[[IndS:.*]] = memref.alloca() : memref<3xindex>
// CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<3xindex> to memref<?xindex> // CHECK-DAG: %[[IndD:.*]] = memref.cast %[[IndS]] : memref<3xindex> to memref<?xindex>
// CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64> // CHECK-DAG: %[[ElemBuffer:.*]] = memref.alloca() : memref<f64>
// CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<2x3x4xf64> // CHECK-DAG: %[[M:.*]] = memref.alloc() : memref<2x3x4xf64>
// CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: linalg.fill ins(%[[E0]] : f64) outs(%[[M]] : memref<2x3x4xf64>) // CHECK-DAG: linalg.fill ins(%[[E0]] : f64) inits(%[[M]] : memref<2x3x4xf64>)
// CHECK: scf.while : () -> () { // CHECK: scf.while : () -> () {
// CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1 // CHECK: %[[Cond:.*]] = func.call @getNextF64(%[[Iter]], %[[IndD]], %[[ElemBuffer]]) : (!llvm.ptr<i8>, memref<?xindex>, memref<f64>) -> i1
// CHECK: scf.condition(%[[Cond]]) // CHECK: scf.condition(%[[Cond]])
// CHECK: } do { // CHECK: } do {
// CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<3xindex> // CHECK: %[[Iv0:.*]] = memref.load %[[IndS]][%[[I0]]] : memref<3xindex>
// CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<3xindex> // CHECK: %[[Iv1:.*]] = memref.load %[[IndS]][%[[I1]]] : memref<3xindex>
// CHECK: %[[Iv2:.*]] = memref.load %[[IndS]][%[[I2]]] : memref<3xindex> // CHECK: %[[Iv2:.*]] = memref.load %[[IndS]][%[[I2]]] : memref<3xindex>
// CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64> // CHECK: %[[ElemVal:.*]] = memref.load %[[ElemBuffer]][] : memref<f64>
Show All 9 Lines

mlir/test/Dialect/SparseTensor/dense.mlir

Show First 20 Lines • Show All 53 Lines • ▼ Show 20 Lines
// CHECK: return %[[VAL_15]] : tensor<32x16xf32> // CHECK: return %[[VAL_15]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @dense1(%arga: tensor<32x16xf32, #DenseMatrix>, func.func @dense1(%arga: tensor<32x16xf32, #DenseMatrix>,
%argx: tensor<32x16xf32>) %argx: tensor<32x16xf32>)
-> tensor<32x16xf32> { -> tensor<32x16xf32> {
%c = arith.constant 1.0 : f32 %c = arith.constant 1.0 : f32
%0 = linalg.generic #trait_2d %0 = linalg.generic #trait_2d
ins(%arga: tensor<32x16xf32, #DenseMatrix>) ins(%arga: tensor<32x16xf32, #DenseMatrix>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%1 = arith.addf %a, %c : f32 %1 = arith.addf %a, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// //
Show All 23 Lines
// CHECK: return %[[VAL_15]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: return %[[VAL_15]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: } // CHECK: }
func.func @dense2(%arga: tensor<32x16xf32>, func.func @dense2(%arga: tensor<32x16xf32>,
%argx: tensor<32x16xf32, #DenseMatrix>) %argx: tensor<32x16xf32, #DenseMatrix>)
-> tensor<32x16xf32, #DenseMatrix> { -> tensor<32x16xf32, #DenseMatrix> {
%c = arith.constant 1.0 : f32 %c = arith.constant 1.0 : f32
%0 = linalg.generic #trait_2d %0 = linalg.generic #trait_2d
ins(%arga: tensor<32x16xf32>) ins(%arga: tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32, #DenseMatrix>) { inits(%argx: tensor<32x16xf32, #DenseMatrix>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%1 = arith.addf %a, %c : f32 %1 = arith.addf %a, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32, #DenseMatrix> } -> tensor<32x16xf32, #DenseMatrix>
return %0 : tensor<32x16xf32, #DenseMatrix> return %0 : tensor<32x16xf32, #DenseMatrix>
} }
Show All 29 Lines
// CHECK: %[[VAL_20:.*]] = sparse_tensor.load %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_20:.*]] = sparse_tensor.load %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: return %[[VAL_20]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: return %[[VAL_20]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: } // CHECK: }
func.func @dense3(%arga: tensor<32x16x8xf32>, func.func @dense3(%arga: tensor<32x16x8xf32>,
%argx: tensor<32x16xf32, #DenseMatrix>) %argx: tensor<32x16xf32, #DenseMatrix>)
-> tensor<32x16xf32, #DenseMatrix> { -> tensor<32x16xf32, #DenseMatrix> {
%0 = linalg.generic #trait_3d %0 = linalg.generic #trait_3d
ins(%arga: tensor<32x16x8xf32>) ins(%arga: tensor<32x16x8xf32>)
outs(%argx: tensor<32x16xf32, #DenseMatrix>) { inits(%argx: tensor<32x16xf32, #DenseMatrix>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%1 = arith.addf %x, %a : f32 %1 = arith.addf %x, %a : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32, #DenseMatrix> } -> tensor<32x16xf32, #DenseMatrix>
return %0 : tensor<32x16xf32, #DenseMatrix> return %0 : tensor<32x16xf32, #DenseMatrix>
} }

mlir/test/Dialect/SparseTensor/one_shot_bufferize_tensor_copy_insertion.mlir

Show First 20 Lines • Show All 53 Lines • ▼ Show 20 Lines
// CHECK-LABEL: func @update_notinplace( // CHECK-LABEL: func @update_notinplace(
// CHECK-SAME: %[[argb:.*]]: tensor<10xf32> // CHECK-SAME: %[[argb:.*]]: tensor<10xf32>
// CHECK-FUNC-LABEL: func @update_notinplace( // CHECK-FUNC-LABEL: func @update_notinplace(
// CHECK-FUNC-SAME: %[[argb:.*]]: tensor<10xf32> // CHECK-FUNC-SAME: %[[argb:.*]]: tensor<10xf32>
func.func @update_notinplace(%argb: tensor<10xf32>, %arga: tensor<10xf32, #SV>) func.func @update_notinplace(%argb: tensor<10xf32>, %arga: tensor<10xf32, #SV>)
-> (tensor<10xf32>, tensor<10xf32>) -> (tensor<10xf32>, tensor<10xf32>)
{ {
// CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) {bufferization.escape = [false]} : tensor<10xf32> // CHECK: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) {bufferization.escape = [false]} : tensor<10xf32>
// CHECK: linalg.generic {{.*}} outs(%[[alloc]] // CHECK: linalg.generic {{.*}} inits(%[[alloc]]
// CHECK-FUNC: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) {bufferization.escape = [true]} : tensor<10xf32> // CHECK-FUNC: %[[alloc:.*]] = bufferization.alloc_tensor() copy(%[[argb]]) {bufferization.escape = [true]} : tensor<10xf32>
// CHECK-FUNC: linalg.generic {{.*}} outs(%[[alloc]] // CHECK-FUNC: linalg.generic {{.*}} inits(%[[alloc]]
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<10xf32, #SV>) ins(%arga: tensor<10xf32, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: f32, %x : f32): ^bb(%a: f32, %x : f32):
%up = arith.addf %a, %x : f32 %up = arith.addf %a, %x : f32
linalg.yield %up : f32 linalg.yield %up : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0, %argb : tensor<10xf32>, tensor<10xf32> return %0, %argb : tensor<10xf32>, tensor<10xf32>
} }

mlir/test/Dialect/SparseTensor/one_trip.mlir

Show All 19 Lines
// CHECK: %[[VAL_4:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<?xf32> // CHECK: %[[VAL_4:.*]] = memref.load %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<?xf32>
// CHECK: %[[VAL_5:.*]] = arith.mulf %[[VAL_4]], %[[VAL_2]] : f32 // CHECK: %[[VAL_5:.*]] = arith.mulf %[[VAL_4]], %[[VAL_2]] : f32
// CHECK: memref.store %[[VAL_5]], %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<?xf32> // CHECK: memref.store %[[VAL_5]], %[[VAL_3]]{{\[}}%[[VAL_1]]] : memref<?xf32>
// CHECK: %[[VAL_6:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<1x1xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> // CHECK: %[[VAL_6:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<1x1xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>
// CHECK: return %[[VAL_6]] : tensor<1x1xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> // CHECK: return %[[VAL_6]] : tensor<1x1xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>
func.func @sparse_scale(%argx: tensor<1x1xf32, #Dense>) -> tensor<1x1xf32, #Dense> { func.func @sparse_scale(%argx: tensor<1x1xf32, #Dense>) -> tensor<1x1xf32, #Dense> {
%c = arith.constant 2.0 : f32 %c = arith.constant 2.0 : f32
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
outs(%argx: tensor<1x1xf32, #Dense>) { inits(%argx: tensor<1x1xf32, #Dense>) {
^bb(%x: f32): ^bb(%x: f32):
%1 = arith.mulf %x, %c : f32 %1 = arith.mulf %x, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<1x1xf32, #Dense> } -> tensor<1x1xf32, #Dense>
return %0 : tensor<1x1xf32, #Dense> return %0 : tensor<1x1xf32, #Dense>
} }

mlir/test/Dialect/SparseTensor/sorted_coo.mlir

Show First 20 Lines • Show All 50 Lines • ▼ Show 20 Lines
// CHECK: memref.store %[[VAL_10]], %[[VAL_5]]{{\[}}%[[VAL_8]]] : memref<?xf32> // CHECK: memref.store %[[VAL_10]], %[[VAL_5]]{{\[}}%[[VAL_8]]] : memref<?xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_11:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> // CHECK: %[[VAL_11:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
// CHECK: return %[[VAL_11]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> // CHECK: return %[[VAL_11]] : tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>
// CHECK: } // CHECK: }
func.func @sparse_scale(%argx: tensor<?x?xf32, #SortedCOO>) -> tensor<?x?xf32, #SortedCOO> { func.func @sparse_scale(%argx: tensor<?x?xf32, #SortedCOO>) -> tensor<?x?xf32, #SortedCOO> {
%c = arith.constant 2.0 : f32 %c = arith.constant 2.0 : f32
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
outs(%argx: tensor<?x?xf32, #SortedCOO>) { inits(%argx: tensor<?x?xf32, #SortedCOO>) {
^bb(%x: f32): ^bb(%x: f32):
%1 = arith.mulf %x, %c : f32 %1 = arith.mulf %x, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?x?xf32, #SortedCOO> } -> tensor<?x?xf32, #SortedCOO>
return %0 : tensor<?x?xf32, #SortedCOO> return %0 : tensor<?x?xf32, #SortedCOO>
} }
// CHECK-LABEL: func.func @matvec( // CHECK-LABEL: func.func @matvec(
Show All 23 Lines
// CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf64> // CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf64>
// CHECK: return %[[VAL_21]] : tensor<32xf64> // CHECK: return %[[VAL_21]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @matvec(%arga: tensor<32x64xf64, #SortedCOO>, func.func @matvec(%arga: tensor<32x64xf64, #SortedCOO>,
%argb: tensor<64xf64>, %argb: tensor<64xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%arga, %argb : tensor<32x64xf64, #SortedCOO>, tensor<64xf64>) ins(%arga, %argb : tensor<32x64xf64, #SortedCOO>, tensor<64xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%A: f64, %b: f64, %x: f64): ^bb(%A: f64, %b: f64, %x: f64):
%0 = arith.mulf %A, %b : f64 %0 = arith.mulf %A, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func.func @mateltmul( // CHECK-LABEL: func.func @mateltmul(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x64xf64>) -> tensor<32x64xf64> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x64xf64>) -> tensor<32x64xf64> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed-nu", "singleton" ] }>> to memref<?xf64>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x64xf64> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x64xf64>
// CHECK-DAG: linalg.fill ins(%[[VAL_3]] : f64) outs(%[[VAL_14]] : memref<32x64xf64>) // CHECK-DAG: linalg.fill ins(%[[VAL_3]] : f64) inits(%[[VAL_14]] : memref<32x64xf64>)
// CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_15]], %[[VAL_21:.*]] = %[[VAL_17]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_15]], %[[VAL_21:.*]] = %[[VAL_17]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_16]] : index // CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_16]] : index
// CHECK: %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_18]] : index // CHECK: %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_18]] : index
// CHECK: %[[VAL_24:.*]] = arith.andi %[[VAL_22]], %[[VAL_23]] : i1 // CHECK: %[[VAL_24:.*]] = arith.andi %[[VAL_22]], %[[VAL_23]] : i1
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_66:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x64xf64> // CHECK: %[[VAL_66:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x64xf64>
// CHECK: return %[[VAL_66]] : tensor<32x64xf64> // CHECK: return %[[VAL_66]] : tensor<32x64xf64>
// CHECK: } // CHECK: }
func.func @mateltmul(%argx: tensor<32x64xf64, #SortedCOO>, func.func @mateltmul(%argx: tensor<32x64xf64, #SortedCOO>,
%argy: tensor<32x64xf64, #SortedCOO>, %argy: tensor<32x64xf64, #SortedCOO>,
%argz: tensor<32x64xf64>) -> tensor<32x64xf64> { %argz: tensor<32x64xf64>) -> tensor<32x64xf64> {
%0 = linalg.generic #trait_mul %0 = linalg.generic #trait_mul
ins(%argx, %argy : tensor<32x64xf64, #SortedCOO>, tensor<32x64xf64, #SortedCOO>) ins(%argx, %argy : tensor<32x64xf64, #SortedCOO>, tensor<32x64xf64, #SortedCOO>)
outs(%argz: tensor<32x64xf64>) { inits(%argz: tensor<32x64xf64>) {
^bb(%x: f64, %y: f64, %z: f64): ^bb(%x: f64, %y: f64, %z: f64):
%1 = arith.mulf %x, %y : f64 %1 = arith.mulf %x, %y : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32x64xf64> } -> tensor<32x64xf64>
return %0 : tensor<32x64xf64> return %0 : tensor<32x64xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_1d.mlir

Show All 16 Lines
// 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, // CHECK-SAME: %[[VAL_1:.*]]: f32,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_8]] : memref<32xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_8]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xf32> // CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xf32>
// CHECK: %[[VAL_11:.*]] = arith.addf %[[VAL_10]], %[[VAL_1]] : f32 // CHECK: %[[VAL_11:.*]] = arith.addf %[[VAL_10]], %[[VAL_1]] : f32
// CHECK: memref.store %[[VAL_11]], %[[VAL_8]]{{\[}}%[[VAL_9]]] : memref<32xf32> // CHECK: memref.store %[[VAL_11]], %[[VAL_8]]{{\[}}%[[VAL_9]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_12:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32> // CHECK: %[[VAL_12:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32>
// CHECK: return %[[VAL_12]] : tensor<32xf32> // CHECK: return %[[VAL_12]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_d(%arga: tensor<32xf32, #DV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_d(%arga: tensor<32xf32, #DV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #DV>) ins(%arga: tensor<32xf32, #DV>)
outs(%argx: tensor<32xf32>) { inits(%argx: 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>
} }
// 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:.*]] = tensor.empty() : 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) inits(%[[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 = tensor.empty() : 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>) { inits(%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>
} }
// CHECK-LABEL: func @mul_d( // CHECK-LABEL: func @mul_d(
// 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, // CHECK-SAME: %[[VAL_1:.*]]: f32,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_8]] : memref<32xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_8]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xf32> // CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xf32>
// CHECK: %[[VAL_11:.*]] = arith.mulf %[[VAL_10]], %[[VAL_1]] : f32 // CHECK: %[[VAL_11:.*]] = arith.mulf %[[VAL_10]], %[[VAL_1]] : f32
// CHECK: memref.store %[[VAL_11]], %[[VAL_8]]{{\[}}%[[VAL_9]]] : memref<32xf32> // CHECK: memref.store %[[VAL_11]], %[[VAL_8]]{{\[}}%[[VAL_9]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_12:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32> // CHECK: %[[VAL_12:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32>
// CHECK: return %[[VAL_12]] : tensor<32xf32> // CHECK: return %[[VAL_12]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_d(%arga: tensor<32xf32, #DV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_d(%arga: tensor<32xf32, #DV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #DV>) ins(%arga: tensor<32xf32, #DV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %a, %argb : f32 %0 = arith.mulf %a, %argb : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @add_s( // CHECK-LABEL: func @add_s(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: f32, // CHECK-SAME: %[[VAL_1:.*]]: f32,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant true // CHECK-DAG: %[[VAL_5:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// 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-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_11]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_11]] : memref<32xf32>)
// CHECK: %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_14:.*]]:2 = scf.while (%[[VAL_15:.*]] = %[[VAL_12]], %[[VAL_16:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index // CHECK: %[[VAL_17:.*]] = arith.cmpi ult, %[[VAL_15]], %[[VAL_13]] : index
// CHECK: scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index // CHECK: scf.condition(%[[VAL_17]]) %[[VAL_15]], %[[VAL_16]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index): // CHECK: ^bb0(%[[VAL_18:.*]]: index, %[[VAL_19:.*]]: index):
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi eq, %[[VAL_20]], %[[VAL_19]] : index
// CHECK: scf.if %[[VAL_21]] { // CHECK: scf.if %[[VAL_21]] {
Show All 16 Lines
// CHECK: memref.store %[[VAL_1]], %[[VAL_11]]{{\[}}%[[VAL_28]]] : memref<32xf32> // CHECK: memref.store %[[VAL_1]], %[[VAL_11]]{{\[}}%[[VAL_28]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf32> // CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf32>
// CHECK: return %[[VAL_30]] : tensor<32xf32> // CHECK: return %[[VAL_30]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_s(%arga: tensor<32xf32, #SV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_s(%arga: tensor<32xf32, #SV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #SV>) ins(%arga: tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: 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>
} }
// CHECK-LABEL: func @repeated_add_s( // CHECK-LABEL: func @repeated_add_s(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_4:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_4:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]]
// CHECK-DAG: %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_8]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_8]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_3]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_3]] {
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xindex> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xindex>
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
// CHECK: %[[VAL_15:.*]] = arith.addf %[[VAL_13]], %[[VAL_14]] : f32 // CHECK: %[[VAL_15:.*]] = arith.addf %[[VAL_13]], %[[VAL_14]] : f32
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xf32>
// CHECK: %[[VAL_18:.*]] = arith.addf %[[VAL_16]], %[[VAL_17]] : f32 // CHECK: %[[VAL_18:.*]] = arith.addf %[[VAL_16]], %[[VAL_17]] : f32
// CHECK: %[[VAL_19:.*]] = arith.addf %[[VAL_15]], %[[VAL_18]] : f32 // CHECK: %[[VAL_19:.*]] = arith.addf %[[VAL_15]], %[[VAL_18]] : f32
// CHECK: memref.store %[[VAL_19]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xf32> // CHECK: memref.store %[[VAL_19]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32> // CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf32>
// CHECK: return %[[VAL_20]] : tensor<32xf32> // CHECK: return %[[VAL_20]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @repeated_add_s(%arga: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @repeated_add_s(%arga: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #SV>) ins(%arga: tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %a, %a : f32 // same tensor %0 = arith.addf %a, %a : f32 // same tensor
%1 = arith.addf %a, %a : f32 // should yield %1 = arith.addf %a, %a : f32 // should yield
%2 = arith.addf %0, %1 : f32 // one guard %2 = arith.addf %0, %1 : f32 // one guard
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @mul_s( // CHECK-LABEL: func @mul_s(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: f32, // CHECK-SAME: %[[VAL_1:.*]]: f32,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_9]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_9]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_4]] {
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xf32> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_12]]] : memref<?xf32>
// CHECK: %[[VAL_15:.*]] = arith.mulf %[[VAL_14]], %[[VAL_1]] : f32 // CHECK: %[[VAL_15:.*]] = arith.mulf %[[VAL_14]], %[[VAL_1]] : f32
// CHECK: memref.store %[[VAL_15]], %[[VAL_9]]{{\[}}%[[VAL_13]]] : memref<32xf32> // CHECK: memref.store %[[VAL_15]], %[[VAL_9]]{{\[}}%[[VAL_13]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32> // CHECK: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32>
// CHECK: return %[[VAL_16]] : tensor<32xf32> // CHECK: return %[[VAL_16]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_s(%arga: tensor<32xf32, #SV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_s(%arga: tensor<32xf32, #SV>, %argb: f32, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf32, #SV>) ins(%arga: tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %a, %argb : f32 %0 = arith.mulf %a, %argb : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
#trait2 = { #trait2 = {
Show All 11 Lines
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32> // CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_9]] : memref<32xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_9]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_10]]] : memref<?xf32> // CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_10]]] : memref<?xf32>
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_10]]] : memref<32xf32> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_10]]] : memref<32xf32>
// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_11]], %[[VAL_12]] : f32 // CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_11]], %[[VAL_12]] : f32
// CHECK: memref.store %[[VAL_13]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<32xf32> // CHECK: memref.store %[[VAL_13]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32>
// CHECK: return %[[VAL_14]] : tensor<32xf32> // CHECK: return %[[VAL_14]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_dd(%arga: tensor<32xf32, #DV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_dd(%arga: tensor<32xf32, #DV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #DV>, tensor<32xf32>) ins(%arga, %argb: tensor<32xf32, #DV>, tensor<32xf32>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @mul_dd( // CHECK-LABEL: func @mul_dd(
// 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:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32> // CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_9]] : memref<32xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_9]] : memref<32xf32>)
// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_10]]] : memref<?xf32> // CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_10]]] : memref<?xf32>
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_10]]] : memref<32xf32> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_10]]] : memref<32xf32>
// CHECK: %[[VAL_13:.*]] = arith.mulf %[[VAL_11]], %[[VAL_12]] : f32 // CHECK: %[[VAL_13:.*]] = arith.mulf %[[VAL_11]], %[[VAL_12]] : f32
// CHECK: memref.store %[[VAL_13]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<32xf32> // CHECK: memref.store %[[VAL_13]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf32>
// CHECK: return %[[VAL_14]] : tensor<32xf32> // CHECK: return %[[VAL_14]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_dd(%arga: tensor<32xf32, #DV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_dd(%arga: tensor<32xf32, #DV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #DV>, tensor<32xf32>) ins(%arga, %argb: tensor<32xf32, #DV>, tensor<32xf32>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @add_ds( // CHECK-LABEL: func @add_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant true // CHECK-DAG: %[[VAL_5:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32> // CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_12]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index // CHECK: %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index
// CHECK: scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index // CHECK: scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index): // CHECK: ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
Show All 21 Lines
// CHECK: memref.store %[[VAL_33]], %[[VAL_12]]{{\[}}%[[VAL_31]]] : memref<32xf32> // CHECK: memref.store %[[VAL_33]], %[[VAL_12]]{{\[}}%[[VAL_31]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32> // CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32>
// CHECK: return %[[VAL_34]] : tensor<32xf32> // CHECK: return %[[VAL_34]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_ds(%arga: tensor<32xf32>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_ds(%arga: tensor<32xf32>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32>, tensor<32xf32, #SV>) ins(%arga, %argb: tensor<32xf32>, tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @mul_ds( // CHECK-LABEL: func @mul_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32> // CHECK-DAG: %[[VAL_5:.*]] = bufferization.to_memref %[[VAL_0]] : memref<32xf32>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_10]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<32xf32> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<32xf32>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_13]]] : memref<?xf32>
// CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32 // CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32
// CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32> // CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32>
// CHECK: return %[[VAL_18]] : tensor<32xf32> // CHECK: return %[[VAL_18]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_ds(%arga: tensor<32xf32>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_ds(%arga: tensor<32xf32>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32>, tensor<32xf32, #SV>) ins(%arga, %argb: tensor<32xf32>, tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @add_sd( // CHECK-LABEL: func @add_sd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant true // CHECK-DAG: %[[VAL_5:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_12]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_15:.*]]:2 = scf.while (%[[VAL_16:.*]] = %[[VAL_13]], %[[VAL_17:.*]] = %[[VAL_4]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index // CHECK: %[[VAL_18:.*]] = arith.cmpi ult, %[[VAL_16]], %[[VAL_14]] : index
// CHECK: scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index // CHECK: scf.condition(%[[VAL_18]]) %[[VAL_16]], %[[VAL_17]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index): // CHECK: ^bb0(%[[VAL_19:.*]]: index, %[[VAL_20:.*]]: index):
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
Show All 21 Lines
// CHECK: memref.store %[[VAL_33]], %[[VAL_12]]{{\[}}%[[VAL_31]]] : memref<32xf32> // CHECK: memref.store %[[VAL_33]], %[[VAL_12]]{{\[}}%[[VAL_31]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32> // CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32>
// CHECK: return %[[VAL_34]] : tensor<32xf32> // CHECK: return %[[VAL_34]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_sd(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_sd(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32>) ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @mul_sd( // CHECK-LABEL: func @mul_sd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32> // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_10]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_11]] to %[[VAL_12]] step %[[VAL_4]] {
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xf32> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xf32>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<32xf32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<32xf32>
// CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32 // CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32
// CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_14]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32> // CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32xf32>
// CHECK: return %[[VAL_18]] : tensor<32xf32> // CHECK: return %[[VAL_18]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_sd(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_sd(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32>) ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @add_ss( // CHECK-LABEL: func @add_ss(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_12]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index // CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index
// CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1 // CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
// CHECK: memref.store %[[VAL_52]], %[[VAL_12]]{{\[}}%[[VAL_51]]] : memref<32xf32> // CHECK: memref.store %[[VAL_52]], %[[VAL_12]]{{\[}}%[[VAL_51]]] : memref<32xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_53:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32> // CHECK: %[[VAL_53:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32>
// CHECK: return %[[VAL_53]] : tensor<32xf32> // CHECK: return %[[VAL_53]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @add_ss(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @add_ss(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32, #SV>) ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @mul_ss( // CHECK-LABEL: func @mul_ss(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<32xf32>) -> tensor<32xf32> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<32xf32>) -> tensor<32xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_12]] : memref<32xf32>)
// CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_13]], %[[VAL_19:.*]] = %[[VAL_15]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index // CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_14]] : index
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_16]] : index
// CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1 // CHECK: %[[VAL_22:.*]] = arith.andi %[[VAL_20]], %[[VAL_21]] : i1
Show All 23 Lines
// CHECK: scf.yield %[[VAL_37]], %[[VAL_40]] : index, index // CHECK: scf.yield %[[VAL_37]], %[[VAL_40]] : index, index
// CHECK: } // CHECK: }
// CHECK: %[[VAL_41:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32> // CHECK: %[[VAL_41:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xf32>
// CHECK: return %[[VAL_41]] : tensor<32xf32> // CHECK: return %[[VAL_41]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_ss(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> { func.func @mul_ss(%arga: tensor<32xf32, #SV>, %argb: tensor<32xf32, #SV>, %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32, #SV>) ins(%arga, %argb: tensor<32xf32, #SV>, tensor<32xf32, #SV>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
// CHECK-LABEL: func @two_way_inv( // CHECK-LABEL: func @two_way_inv(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: f32, // CHECK-SAME: %[[VAL_2:.*2]]: f32,
// CHECK-SAME: %[[VAL_3:.*3]]: tensor<16xf32>) -> tensor<16xf32> { // CHECK-SAME: %[[VAL_3:.*3]]: tensor<16xf32>) -> tensor<16xf32> {
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]] // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_13]] : memref<16xf32>)
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index
// CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index // CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index
// CHECK: %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1 // CHECK: %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_60:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<16xf32> // CHECK: %[[VAL_60:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<16xf32>
// CHECK: return %[[VAL_60]] : tensor<16xf32> // CHECK: return %[[VAL_60]] : tensor<16xf32>
// CHECK: } // CHECK: }
func.func @two_way_inv(%arga: tensor<16xf32, #SV>, %argb: tensor<16xf32, #SV>, %argc: f32, %argx: tensor<16xf32>) -> tensor<16xf32> { func.func @two_way_inv(%arga: tensor<16xf32, #SV>, %argb: tensor<16xf32, #SV>, %argc: f32, %argx: tensor<16xf32>) -> tensor<16xf32> {
// Kernel "x(i) = a(i) * c + b(i) * c". // Kernel "x(i) = a(i) * c + b(i) * c".
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>) ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>)
outs(%argx: tensor<16xf32>) { inits(%argx: tensor<16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %argc : f32 %0 = arith.mulf %a, %argc : f32
%1 = arith.mulf %b, %argc : f32 %1 = arith.mulf %b, %argc : f32
%2 = arith.addf %0, %1 : f32 %2 = arith.addf %0, %1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<16xf32> } -> tensor<16xf32>
return %0 : tensor<16xf32> return %0 : tensor<16xf32>
} }
// CHECK-LABEL: func @two_way_inv_alt( // CHECK-LABEL: func @two_way_inv_alt(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: f32, // CHECK-SAME: %[[VAL_2:.*2]]: f32,
// CHECK-SAME: %[[VAL_3:.*3]]: tensor<16xf32>) -> tensor<16xf32> { // CHECK-SAME: %[[VAL_3:.*3]]: tensor<16xf32>) -> tensor<16xf32> {
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]] // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<16xf32>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_13]] : memref<16xf32>)
// CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_14]], %[[VAL_20:.*]] = %[[VAL_16]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_15]] : index
// CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index // CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_17]] : index
// CHECK: %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1 // CHECK: %[[VAL_23:.*]] = arith.andi %[[VAL_21]], %[[VAL_22]] : i1
▲ Show 20 Lines • Show All 52 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_59:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<16xf32> // CHECK: %[[VAL_59:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<16xf32>
// CHECK: return %[[VAL_59]] : tensor<16xf32> // CHECK: return %[[VAL_59]] : tensor<16xf32>
// CHECK: } // CHECK: }
func.func @two_way_inv_alt(%arga: tensor<16xf32, #SV>, func.func @two_way_inv_alt(%arga: tensor<16xf32, #SV>,
%argb: tensor<16xf32, #SV>, %argc: f32, %argx: tensor<16xf32>) -> tensor<16xf32> { %argb: tensor<16xf32, #SV>, %argc: f32, %argx: tensor<16xf32>) -> tensor<16xf32> {
// Same kernel, but now expressed as "x(i) = (a(i) + b(i)) * c". // Same kernel, but now expressed as "x(i) = (a(i) + b(i)) * c".
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>) ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>)
outs(%argx: tensor<16xf32>) { inits(%argx: tensor<16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
%1 = arith.mulf %0, %argc : f32 %1 = arith.mulf %0, %argc : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<16xf32> } -> tensor<16xf32>
return %0 : tensor<16xf32> return %0 : tensor<16xf32>
} }
Show All 24 Lines
// CHECK: } // CHECK: }
// CHECK: memref.store %[[VAL_11]], %[[VAL_6]][] : memref<f32> // CHECK: memref.store %[[VAL_11]], %[[VAL_6]][] : memref<f32>
// CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<f32> // CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<f32>
// CHECK: return %[[VAL_17]] : tensor<f32> // CHECK: return %[[VAL_17]] : tensor<f32>
// CHECK: } // CHECK: }
func.func @sum_reduction(%arga: tensor<?xf32, #SV>, %argx: tensor<f32>) -> tensor<f32> { func.func @sum_reduction(%arga: tensor<?xf32, #SV>, %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_sum_reduction %0 = linalg.generic #trait_sum_reduction
ins(%arga: tensor<?xf32, #SV>) ins(%arga: tensor<?xf32, #SV>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %x, %a : f32 %0 = arith.addf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
#trait_sum_reduction2 = { #trait_sum_reduction2 = {
▲ Show 20 Lines • Show All 87 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
func.func @sum_reduction_ss(%arga: tensor<16xf32, #SV>, func.func @sum_reduction_ss(%arga: tensor<16xf32, #SV>,
%argb: tensor<16xf32, #SV>, %argb: tensor<16xf32, #SV>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
// Just for testing. This case would be better expressed // Just for testing. This case would be better expressed
// as two separate reductions kernels. // as two separate reductions kernels.
%0 = linalg.generic #trait_sum_reduction2 %0 = linalg.generic #trait_sum_reduction2
ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>) ins(%arga, %argb: tensor<16xf32, #SV>, tensor<16xf32, #SV>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
%1 = arith.addf %x, %0 : f32 %1 = arith.addf %x, %0 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
▲ Show 20 Lines • Show All 96 Lines • ▼ Show 20 Lines
func.func @sum_reduction_inv(%arga: tensor<16xf32, #SV>, func.func @sum_reduction_inv(%arga: tensor<16xf32, #SV>,
%argb: tensor<f32>, %argb: tensor<f32>,
%argc: tensor<16xf32, #SV>, %argc: tensor<16xf32, #SV>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
// Just for testing. This case would be better expressed // Just for testing. This case would be better expressed
// as two separate reductions kernels. // as two separate reductions kernels.
%0 = linalg.generic #trait_sum_reduction_inv %0 = linalg.generic #trait_sum_reduction_inv
ins(%arga, %argb, %argc : tensor<16xf32, #SV>, tensor<f32>, tensor<16xf32, #SV>) ins(%arga, %argb, %argc : tensor<16xf32, #SV>, tensor<f32>, tensor<16xf32, #SV>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32, %c: f32, %x: f32): ^bb(%a: f32, %b: f32, %c: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.addf %0, %c : f32 %1 = arith.addf %0, %c : f32
%2 = arith.addf %x, %1 : f32 %2 = arith.addf %x, %1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
Show All 24 Lines
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?xf64> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<?xf64>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.pointers %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.pointers %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.indices %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.indices %[[VAL_3]] {dimension = 0 : index} : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.values %[[VAL_3]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64> // CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.values %[[VAL_3]] : tensor<?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> to memref<?xf64>
// CHECK-DAG: %[[VAL_16:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?xf64> // CHECK-DAG: %[[VAL_16:.*]] = tensor.dim %[[VAL_0]], %[[VAL_5]] : tensor<?xf64>
// CHECK-DAG: %[[VAL_18:.*]] = bufferization.to_memref %[[VAL_4]] // CHECK-DAG: %[[VAL_18:.*]] = bufferization.to_memref %[[VAL_4]]
// CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[VAL_18]] : memref<?xf64>) // CHECK-DAG: linalg.fill ins(%{{.*}} : f64) inits(%[[VAL_18]] : memref<?xf64>)
// CHECK-DAG: %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK-DAG: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_21:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK-DAG: %[[VAL_21:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK-DAG: %[[VAL_22:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK-DAG: %[[VAL_22:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_23:.*]]:3 = scf.while (%[[VAL_24:.*]] = %[[VAL_19]], %[[VAL_25:.*]] = %[[VAL_21]], %[[VAL_26:.*]] = %[[VAL_5]]) : (index, index, index) -> (index, index, index) { // CHECK: %[[VAL_23:.*]]:3 = scf.while (%[[VAL_24:.*]] = %[[VAL_19]], %[[VAL_25:.*]] = %[[VAL_21]], %[[VAL_26:.*]] = %[[VAL_5]]) : (index, index, index) -> (index, index, index) {
// CHECK: %[[VAL_27:.*]] = arith.cmpi ult, %[[VAL_24]], %[[VAL_20]] : index // CHECK: %[[VAL_27:.*]] = arith.cmpi ult, %[[VAL_24]], %[[VAL_20]] : index
// CHECK: %[[VAL_28:.*]] = arith.cmpi ult, %[[VAL_25]], %[[VAL_22]] : index // CHECK: %[[VAL_28:.*]] = arith.cmpi ult, %[[VAL_25]], %[[VAL_22]] : index
// CHECK: %[[VAL_29:.*]] = arith.andi %[[VAL_27]], %[[VAL_28]] : i1 // CHECK: %[[VAL_29:.*]] = arith.andi %[[VAL_27]], %[[VAL_28]] : i1
▲ Show 20 Lines • Show All 121 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
func.func @four_tensors_op(%arga: tensor<?xf64>, func.func @four_tensors_op(%arga: tensor<?xf64>,
%argb: tensor<?xf64, #SV>, %argb: tensor<?xf64, #SV>,
%argc: tensor<?xf64>, %argc: tensor<?xf64>,
%argd: tensor<?xf64, #SV>, %argd: tensor<?xf64, #SV>,
%argx: tensor<?xf64>) -> tensor<?xf64> { %argx: tensor<?xf64>) -> tensor<?xf64> {
%r = linalg.generic #trait_four_tensors %r = linalg.generic #trait_four_tensors
ins(%arga, %argb, %argc, %argd: tensor<?xf64>, tensor<?xf64, #SV>, tensor<?xf64>, tensor<?xf64, #SV>) ins(%arga, %argb, %argc, %argd: tensor<?xf64>, tensor<?xf64, #SV>, tensor<?xf64>, tensor<?xf64, #SV>)
outs(%argx: tensor<?xf64>) { inits(%argx: tensor<?xf64>) {
^bb(%a: f64, %b: f64, %c: f64, %d: f64, %x: f64): ^bb(%a: f64, %b: f64, %c: f64, %d: f64, %x: f64):
%0 = arith.addf %a, %b : f64 %0 = arith.addf %a, %b : f64
%1 = arith.addf %c, %d : f64 %1 = arith.addf %c, %d : f64
%2 = arith.addf %0, %1 : f64 %2 = arith.addf %0, %1 : f64
linalg.yield %2 : f64 linalg.yield %2 : f64
} -> tensor<?xf64> } -> tensor<?xf64>
return %r : tensor<?xf64> return %r : tensor<?xf64>
} }
▲ Show 20 Lines • Show All 296 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_252:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<f64> // CHECK: %[[VAL_252:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<f64>
// CHECK: return %[[VAL_252]] : tensor<f64> // CHECK: return %[[VAL_252]] : tensor<f64>
// CHECK: } // CHECK: }
func.func @red3s(%arga: tensor<?xf64, #SV>, func.func @red3s(%arga: tensor<?xf64, #SV>,
%argb: tensor<?xf64, #SV>, %argb: tensor<?xf64, #SV>,
%argc: tensor<?xf64, #SV>, %argx: tensor<f64>) ->tensor<f64>{ %argc: tensor<?xf64, #SV>, %argx: tensor<f64>) ->tensor<f64>{
%0 = linalg.generic #trait_red3s %0 = linalg.generic #trait_red3s
ins(%arga, %argb, %argc: tensor<?xf64, #SV>, tensor<?xf64, #SV>, tensor<?xf64, #SV>) ins(%arga, %argb, %argc: tensor<?xf64, #SV>, tensor<?xf64, #SV>, tensor<?xf64, #SV>)
outs(%argx: tensor<f64>) { inits(%argx: tensor<f64>) {
^bb(%a: f64,%b: f64,%c: f64,%x: f64): ^bb(%a: f64,%b: f64,%c: f64,%x: f64):
%0 = arith.addf %x, %a : f64 %0 = arith.addf %x, %a : f64
%1 = arith.addf %0, %b : f64 %1 = arith.addf %0, %b : f64
%2 = arith.addf %1, %c : f64 %2 = arith.addf %1, %c : f64
linalg.yield %2 : f64 linalg.yield %2 : f64
} -> tensor<f64> } -> tensor<f64>
return %0 : tensor<f64> return %0 : tensor<f64>
} }

mlir/test/Dialect/SparseTensor/sparse_2d.mlir

Show All 21 Lines
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_10]] : memref<32x16xf32>)
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_11]], %[[VAL_4]] : index // CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_11]], %[[VAL_4]] : index
// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index // CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32>
// CHECK: %[[VAL_17:.*]] = arith.addf %[[VAL_15]], %[[VAL_16]] : f32 // CHECK: %[[VAL_17:.*]] = arith.addf %[[VAL_15]], %[[VAL_16]] : f32
// CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32x16xf32> // CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32x16xf32>
// CHECK: return %[[VAL_18]] : tensor<32x16xf32> // CHECK: return %[[VAL_18]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_dd(%arga: tensor<32x16xf32, #Tdd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_dd(%arga: tensor<32x16xf32, #Tdd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tdd>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tdd>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_dd( // CHECK-LABEL: func @mul_dd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_10]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_10]] : memref<32x16xf32>)
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_11]], %[[VAL_4]] : index // CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_11]], %[[VAL_4]] : index
// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index // CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32>
// CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32 // CHECK: %[[VAL_17:.*]] = arith.mulf %[[VAL_15]], %[[VAL_16]] : f32
// CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_10]]{{\[}}%[[VAL_11]], %[[VAL_12]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32x16xf32> // CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<32x16xf32>
// CHECK: return %[[VAL_18]] : tensor<32x16xf32> // CHECK: return %[[VAL_18]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_dd(%arga: tensor<32x16xf32, #Tdd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_dd(%arga: tensor<32x16xf32, #Tdd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tdd>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tdd>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @add_ds( // CHECK-LABEL: func @add_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant true // CHECK-DAG: %[[VAL_6:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_13]] : memref<32x16xf32>)
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_7]] {
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<?xindex>
// CHECK: %[[VAL_16:.*]] = arith.addi %[[VAL_14]], %[[VAL_7]] : index // CHECK: %[[VAL_16:.*]] = arith.addi %[[VAL_14]], %[[VAL_7]] : index
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_15]], %[[VAL_20:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_15]], %[[VAL_20:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index
// CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index // CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index
// CHECK: } do { // CHECK: } do {
Show All 24 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_37:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32> // CHECK: %[[VAL_37:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32>
// CHECK: return %[[VAL_37]] : tensor<32x16xf32> // CHECK: return %[[VAL_37]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_ds(%arga: tensor<32x16xf32, #Tds>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_ds(%arga: tensor<32x16xf32, #Tds>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tds>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tds>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_ds( // CHECK-LABEL: func @mul_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_11]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_11]] : memref<32x16xf32>)
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_5]] : index // CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_5]] : index
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_14]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_14]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_5]] {
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xf32> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xf32>
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<32x16xf32> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<32x16xf32>
// CHECK: %[[VAL_20:.*]] = arith.mulf %[[VAL_18]], %[[VAL_19]] : f32 // CHECK: %[[VAL_20:.*]] = arith.mulf %[[VAL_18]], %[[VAL_19]] : f32
// CHECK: memref.store %[[VAL_20]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_20]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16xf32> // CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16xf32>
// CHECK: return %[[VAL_21]] : tensor<32x16xf32> // CHECK: return %[[VAL_21]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_ds(%arga: tensor<32x16xf32, #Tds>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_ds(%arga: tensor<32x16xf32, #Tds>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tds>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tds>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @add_sd( // CHECK-LABEL: func @add_sd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant true // CHECK-DAG: %[[VAL_5:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_13]] : memref<32x16xf32>)
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_16:.*]]:2 = scf.while (%[[VAL_17:.*]] = %[[VAL_14]], %[[VAL_18:.*]] = %[[VAL_6]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_16:.*]]:2 = scf.while (%[[VAL_17:.*]] = %[[VAL_14]], %[[VAL_18:.*]] = %[[VAL_6]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_19:.*]] = arith.cmpi ult, %[[VAL_17]], %[[VAL_15]] : index // CHECK: %[[VAL_19:.*]] = arith.cmpi ult, %[[VAL_17]], %[[VAL_15]] : index
// CHECK: scf.condition(%[[VAL_19]]) %[[VAL_17]], %[[VAL_18]] : index, index // CHECK: scf.condition(%[[VAL_19]]) %[[VAL_17]], %[[VAL_18]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_20:.*]]: index, %[[VAL_21:.*]]: index): // CHECK: ^bb0(%[[VAL_20:.*]]: index, %[[VAL_21:.*]]: index):
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<?xindex>
Show All 29 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_40:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32> // CHECK: %[[VAL_40:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32>
// CHECK: return %[[VAL_40]] : tensor<32x16xf32> // CHECK: return %[[VAL_40]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_sd(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_sd(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_sd( // CHECK-LABEL: func @mul_sd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_11]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_11]] : memref<32x16xf32>)
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_13]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_13]] step %[[VAL_5]] {
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_14]], %[[VAL_3]] : index // CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_14]], %[[VAL_3]] : index
// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index // CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_15]], %[[VAL_16]]] : memref<32x16xf32> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_15]], %[[VAL_16]]] : memref<32x16xf32>
// CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32 // CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32
// CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_15]], %[[VAL_16]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_15]], %[[VAL_16]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16xf32> // CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16xf32>
// CHECK: return %[[VAL_22]] : tensor<32x16xf32> // CHECK: return %[[VAL_22]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_sd(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_sd(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @add_ss( // CHECK-LABEL: func @add_ss(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant true // CHECK-DAG: %[[VAL_5:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_15]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_15]] : memref<32x16xf32>)
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_16]], %[[VAL_20:.*]] = %[[VAL_6]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_16]], %[[VAL_20:.*]] = %[[VAL_6]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index
// CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index // CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index): // CHECK: ^bb0(%[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index):
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex> // CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex>
▲ Show 20 Lines • Show All 53 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_58:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16xf32> // CHECK: %[[VAL_58:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16xf32>
// CHECK: return %[[VAL_58]] : tensor<32x16xf32> // CHECK: return %[[VAL_58]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_ss( // CHECK-LABEL: func @mul_ss(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_12]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_12]] : memref<32x16xf32>)
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_13]] to %[[VAL_14]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_13]] to %[[VAL_14]] step %[[VAL_4]] {
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_15]], %[[VAL_4]] : index // CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_15]], %[[VAL_4]] : index
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_20:.*]] = %[[VAL_17]] to %[[VAL_19]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_20:.*]] = %[[VAL_17]] to %[[VAL_19]] step %[[VAL_4]] {
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_20]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_20]]] : memref<?xindex>
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<?xf32> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<?xf32>
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]], %[[VAL_21]]] : memref<32x16xf32> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]], %[[VAL_21]]] : memref<32x16xf32>
// CHECK: %[[VAL_24:.*]] = arith.mulf %[[VAL_22]], %[[VAL_23]] : f32 // CHECK: %[[VAL_24:.*]] = arith.mulf %[[VAL_22]], %[[VAL_23]] : f32
// CHECK: memref.store %[[VAL_24]], %[[VAL_12]]{{\[}}%[[VAL_16]], %[[VAL_21]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_24]], %[[VAL_12]]{{\[}}%[[VAL_16]], %[[VAL_21]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_25:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16xf32> // CHECK: %[[VAL_25:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16xf32>
// CHECK: return %[[VAL_25]] : tensor<32x16xf32> // CHECK: return %[[VAL_25]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32>) ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @add_ss_ss( // CHECK-LABEL: func @add_ss_ss(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_16]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_16]] : memref<32x16xf32>)
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]]:2 = scf.while (%[[VAL_22:.*]] = %[[VAL_17]], %[[VAL_23:.*]] = %[[VAL_19]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_21:.*]]:2 = scf.while (%[[VAL_22:.*]] = %[[VAL_17]], %[[VAL_23:.*]] = %[[VAL_19]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_18]] : index // CHECK: %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_18]] : index
// CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_20]] : index // CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_20]] : index
// CHECK: %[[VAL_26:.*]] = arith.andi %[[VAL_24]], %[[VAL_25]] : i1 // CHECK: %[[VAL_26:.*]] = arith.andi %[[VAL_24]], %[[VAL_25]] : i1
▲ Show 20 Lines • Show All 122 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_116:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16xf32> // CHECK: %[[VAL_116:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16xf32>
// CHECK: return %[[VAL_116]] : tensor<32x16xf32> // CHECK: return %[[VAL_116]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_ss_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32, #Tss>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_ss_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32, #Tss>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32, #Tss>) ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32, #Tss>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_ss_ss( // CHECK-LABEL: func @mul_ss_ss(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_16]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_16]] : memref<32x16xf32>)
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]]:2 = scf.while (%[[VAL_22:.*]] = %[[VAL_17]], %[[VAL_23:.*]] = %[[VAL_19]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_21:.*]]:2 = scf.while (%[[VAL_22:.*]] = %[[VAL_17]], %[[VAL_23:.*]] = %[[VAL_19]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_18]] : index // CHECK: %[[VAL_24:.*]] = arith.cmpi ult, %[[VAL_22]], %[[VAL_18]] : index
// CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_20]] : index // CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_20]] : index
// CHECK: %[[VAL_26:.*]] = arith.andi %[[VAL_24]], %[[VAL_25]] : i1 // CHECK: %[[VAL_26:.*]] = arith.andi %[[VAL_24]], %[[VAL_25]] : i1
▲ Show 20 Lines • Show All 54 Lines • ▼ Show 20 Lines
// CHECK: scf.yield %[[VAL_68]], %[[VAL_71]] : index, index // CHECK: scf.yield %[[VAL_68]], %[[VAL_71]] : index, index
// CHECK: } // CHECK: }
// CHECK: %[[VAL_72:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16xf32> // CHECK: %[[VAL_72:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16xf32>
// CHECK: return %[[VAL_72]] : tensor<32x16xf32> // CHECK: return %[[VAL_72]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_ss_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32, #Tss>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_ss_ss(%arga: tensor<32x16xf32, #Tss>, %argb: tensor<32x16xf32, #Tss>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32, #Tss>) ins(%arga, %argb: tensor<32x16xf32, #Tss>, tensor<32x16xf32, #Tss>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @add_sd_ds( // CHECK-LABEL: func @add_sd_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant true // CHECK-DAG: %[[VAL_6:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_15]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_15]] : memref<32x16xf32>)
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_16]], %[[VAL_20:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_18:.*]]:2 = scf.while (%[[VAL_19:.*]] = %[[VAL_16]], %[[VAL_20:.*]] = %[[VAL_5]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index // CHECK: %[[VAL_21:.*]] = arith.cmpi ult, %[[VAL_19]], %[[VAL_17]] : index
// CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index // CHECK: scf.condition(%[[VAL_21]]) %[[VAL_19]], %[[VAL_20]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index): // CHECK: ^bb0(%[[VAL_22:.*]]: index, %[[VAL_23:.*]]: index):
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex> // CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex>
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_70:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16xf32> // CHECK: %[[VAL_70:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16xf32>
// CHECK: return %[[VAL_70]] : tensor<32x16xf32> // CHECK: return %[[VAL_70]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @add_sd_ds(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32, #Tds>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @add_sd_ds(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32, #Tds>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32, #Tds>) ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32, #Tds>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
// CHECK-LABEL: func @mul_sd_ds( // CHECK-LABEL: func @mul_sd_ds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16xf32>) -> tensor<32x16xf32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 1 : index} : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<32x16xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16xf32>
// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[VAL_13]] : memref<32x16xf32>) // CHECK: linalg.fill ins(%{{.*}} : f32) inits(%[[VAL_13]] : memref<32x16xf32>)
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_14]] to %[[VAL_15]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_14]] to %[[VAL_15]] step %[[VAL_5]] {
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_5]] : index // CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_5]] : index
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_5]] {
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex>
// CHECK: %[[VAL_23:.*]] = arith.muli %[[VAL_16]], %[[VAL_3]] : index // CHECK: %[[VAL_23:.*]] = arith.muli %[[VAL_16]], %[[VAL_3]] : index
// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_22]] : index // CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_22]] : index
// CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_24]]] : memref<?xf32> // CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_24]]] : memref<?xf32>
// CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_21]]] : memref<?xf32> // CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_21]]] : memref<?xf32>
// CHECK: %[[VAL_27:.*]] = arith.mulf %[[VAL_25]], %[[VAL_26]] : f32 // CHECK: %[[VAL_27:.*]] = arith.mulf %[[VAL_25]], %[[VAL_26]] : f32
// CHECK: memref.store %[[VAL_27]], %[[VAL_13]]{{\[}}%[[VAL_17]], %[[VAL_22]]] : memref<32x16xf32> // CHECK: memref.store %[[VAL_27]], %[[VAL_13]]{{\[}}%[[VAL_17]], %[[VAL_22]]] : memref<32x16xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_28:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32> // CHECK: %[[VAL_28:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf32>
// CHECK: return %[[VAL_28]] : tensor<32x16xf32> // CHECK: return %[[VAL_28]] : tensor<32x16xf32>
// CHECK: } // CHECK: }
func.func @mul_sd_ds(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32, #Tds>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> { func.func @mul_sd_ds(%arga: tensor<32x16xf32, #Tsd>, %argb: tensor<32x16xf32, #Tds>, %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32, #Tds>) ins(%arga, %argb: tensor<32x16xf32, #Tsd>, tensor<32x16xf32, #Tds>)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }
#trait_matvec = { #trait_matvec = {
Show All 34 Lines
// CHECK: memref.store %[[VAL_17]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<16xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<16xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<16xf32> // CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<16xf32>
// CHECK: return %[[VAL_26]] : tensor<16xf32> // CHECK: return %[[VAL_26]] : tensor<16xf32>
// CHECK: } // CHECK: }
func.func @matvec(%argA: tensor<16x32xf32, #Tds>, %argb: tensor<32xf32>, %argx: tensor<16xf32>) -> tensor<16xf32> { func.func @matvec(%argA: tensor<16x32xf32, #Tds>, %argb: tensor<32xf32>, %argx: tensor<16xf32>) -> tensor<16xf32> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%argA, %argb: tensor<16x32xf32, #Tds>, tensor<32xf32>) ins(%argA, %argb: tensor<16x32xf32, #Tds>, tensor<32xf32>)
outs(%argx: tensor<16xf32>) { inits(%argx: tensor<16xf32>) {
^bb(%A: f32, %b: f32, %x: f32): ^bb(%A: f32, %b: f32, %x: f32):
%0 = arith.mulf %A, %b : f32 %0 = arith.mulf %A, %b : f32
%1 = arith.addf %0, %x : f32 %1 = arith.addf %0, %x : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<16xf32> } -> tensor<16xf32>
return %0 : tensor<16xf32> return %0 : tensor<16xf32>
} }
Show All 29 Lines
// CHECK: } // CHECK: }
// CHECK: memref.store %[[VAL_10]], %[[VAL_8]][] : memref<f32> // CHECK: memref.store %[[VAL_10]], %[[VAL_8]][] : memref<f32>
// CHECK: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<f32> // CHECK: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<f32>
// CHECK: return %[[VAL_23]] : tensor<f32> // CHECK: return %[[VAL_23]] : tensor<f32>
// CHECK: } // CHECK: }
func.func @sum_reduction(%arga: tensor<10x20xf32, #Tds>, %argx: tensor<f32>) -> tensor<f32> { func.func @sum_reduction(%arga: tensor<10x20xf32, #Tds>, %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_sum_reduction %0 = linalg.generic #trait_sum_reduction
ins(%arga: tensor<10x20xf32, #Tds>) ins(%arga: tensor<10x20xf32, #Tds>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %x, %a : f32 %0 = arith.addf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
#trait_scale = { #trait_scale = {
Show All 11 Lines
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 2.000000e+00 : f64 // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 2.000000e+00 : f64
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf64> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> to memref<?xf64>
// CHECK-DAG: %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf64, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf64, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?xf64> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?xf64>
// CHECK: linalg.fill ins(%{{.*}} : f64) outs(%[[VAL_11]] : memref<?x?xf64>) // CHECK: linalg.fill ins(%{{.*}} : f64) inits(%[[VAL_11]] : memref<?x?xf64>)
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] {
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_12]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_12]]] : memref<?xindex>
// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_4]] : index // CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_12]], %[[VAL_4]] : index
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_14]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_4]] {
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xf64> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xf64>
// CHECK: %[[VAL_19:.*]] = arith.mulf %[[VAL_18]], %[[VAL_2]] : f64 // CHECK: %[[VAL_19:.*]] = arith.mulf %[[VAL_18]], %[[VAL_2]] : f64
// CHECK: memref.store %[[VAL_19]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<?x?xf64> // CHECK: memref.store %[[VAL_19]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_17]]] : memref<?x?xf64>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<?x?xf64> // CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<?x?xf64>
// CHECK: return %[[VAL_20]] : tensor<?x?xf64> // CHECK: return %[[VAL_20]] : tensor<?x?xf64>
// CHECK: } // CHECK: }
func.func @scale(%arga: tensor<?x?xf64, #Tds>, %argx: tensor<?x?xf64>) -> tensor<?x?xf64> { func.func @scale(%arga: tensor<?x?xf64, #Tds>, %argx: tensor<?x?xf64>) -> tensor<?x?xf64> {
%0 = arith.constant 2.0 : f64 %0 = arith.constant 2.0 : f64
%1 = linalg.generic #trait_scale %1 = linalg.generic #trait_scale
ins(%arga: tensor<?x?xf64, #Tds>) ins(%arga: tensor<?x?xf64, #Tds>)
outs(%argx: tensor<?x?xf64>) { inits(%argx: tensor<?x?xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%2 = arith.mulf %a, %0 : f64 %2 = arith.mulf %a, %0 : f64
linalg.yield %2 : f64 linalg.yield %2 : f64
} -> tensor<?x?xf64> } -> tensor<?x?xf64>
return %1 : tensor<?x?xf64> return %1 : tensor<?x?xf64>
} }
#trait_sampled_dense_dense = { #trait_sampled_dense_dense = {
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines
// CHECK: return %[[VAL_38]] : tensor<?x?xf32> // CHECK: return %[[VAL_38]] : tensor<?x?xf32>
// CHECK: } // CHECK: }
func.func @sampled_dense_dense(%args: tensor<?x?xf32, #Tss>, func.func @sampled_dense_dense(%args: tensor<?x?xf32, #Tss>,
%arga: tensor<?x?xf32>, %arga: tensor<?x?xf32>,
%argb: tensor<?x?xf32>, %argb: tensor<?x?xf32>,
%argx: tensor<?x?xf32>) -> tensor<?x?xf32> { %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.generic #trait_sampled_dense_dense %0 = linalg.generic #trait_sampled_dense_dense
ins(%args, %arga, %argb: tensor<?x?xf32, #Tss>, tensor<?x?xf32>, tensor<?x?xf32>) ins(%args, %arga, %argb: tensor<?x?xf32, #Tss>, tensor<?x?xf32>, tensor<?x?xf32>)
outs(%argx: tensor<?x?xf32>) { inits(%argx: tensor<?x?xf32>) {
^bb(%s: f32, %a: f32, %b: f32, %x: f32): ^bb(%s: f32, %a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.mulf %s, %0 : f32 %1 = arith.mulf %s, %0 : f32
%2 = arith.addf %x, %1 : f32 %2 = arith.addf %x, %1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
▲ Show 20 Lines • Show All 207 Lines • ▼ Show 20 Linesfunc.func @sum_kernel_with_inv(%arga: tensor<?x?xf32, #Tss>,
%arge: tensor<f32>, %arge: tensor<f32>,
%argx: tensor<?xf32>) -> tensor<?xf32> { %argx: tensor<?xf32>) -> tensor<?xf32> {
%0 = linalg.generic #trait_sum_kernel_with_inv %0 = linalg.generic #trait_sum_kernel_with_inv
ins(%arga, %argb, %argc, %argd, %arge : tensor<?x?xf32, #Tss>, ins(%arga, %argb, %argc, %argd, %arge : tensor<?x?xf32, #Tss>,
tensor<?x?xf32, #Tds>, tensor<?x?xf32, #Tds>,
tensor<?x?xf32, #Tds>, tensor<?x?xf32, #Tds>,
tensor<?xf32>, tensor<?xf32>,
tensor<f32>) tensor<f32>)
outs(%argx: tensor<?xf32>) { inits(%argx: tensor<?xf32>) {
^bb(%a: f32, %b: f32, %c: f32, %d: f32, %e: f32, %x: f32): ^bb(%a: f32, %b: f32, %c: f32, %d: f32, %e: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.mulf %0, %d : f32 %1 = arith.mulf %0, %d : f32
%2 = arith.mulf %1, %e : f32 %2 = arith.mulf %1, %e : f32
%3 = arith.addf %2, %c : f32 %3 = arith.addf %2, %c : f32
%4 = arith.addf %x, %3 : f32 %4 = arith.addf %x, %3 : f32
linalg.yield %4 : f32 linalg.yield %4 : f32
} -> tensor<?xf32> } -> tensor<?xf32>
return %0 : tensor<?xf32> return %0 : tensor<?xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_3d.mlir

Show All 29 Lines
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_11]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_11]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_6]] to %[[VAL_3]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_6]] to %[[VAL_3]] step %[[VAL_7]] {
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] {
// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_12]], %[[VAL_4]] : index // CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_12]], %[[VAL_4]] : index
// CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_13]] : index // CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_13]] : index
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] {
// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_5]] : index // CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_5]] : index
// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index // CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_21:.*]] = arith.addf %[[VAL_19]], %[[VAL_20]] : f32 // CHECK: %[[VAL_21:.*]] = arith.addf %[[VAL_19]], %[[VAL_20]] : f32
// CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16x8xf32> // CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_22]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_22]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_ddd(%arga: tensor<32x16x8xf32, #Tddd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_ddd(%arga: tensor<32x16x8xf32, #Tddd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tddd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tddd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_ddd( // CHECK-LABEL: func @mul_ddd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_11]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_11]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_6]] to %[[VAL_3]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_6]] to %[[VAL_3]] step %[[VAL_7]] {
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] {
// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_12]], %[[VAL_4]] : index // CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_12]], %[[VAL_4]] : index
// CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_13]] : index // CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_13]] : index
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] {
// CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_5]] : index // CHECK: %[[VAL_17:.*]] = arith.muli %[[VAL_15]], %[[VAL_5]] : index
// CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index // CHECK: %[[VAL_18:.*]] = arith.addi %[[VAL_17]], %[[VAL_16]] : index
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf32>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32 // CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32
// CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_21]], %[[VAL_11]]{{\[}}%[[VAL_12]], %[[VAL_13]], %[[VAL_16]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16x8xf32> // CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_22]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_22]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_ddd(%arga: tensor<32x16x8xf32, #Tddd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_ddd(%arga: tensor<32x16x8xf32, #Tddd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tddd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tddd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_dds( // CHECK-LABEL: func @add_dds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_8:.*]] = arith.constant true // CHECK-DAG: %[[VAL_8:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_15]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_15]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_9]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_9]] {
// CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_7]] to %[[VAL_5]] step %[[VAL_9]] { // CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_7]] to %[[VAL_5]] step %[[VAL_9]] {
// CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_16]], %[[VAL_5]] : index // CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_16]], %[[VAL_5]] : index
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_17]] : index // CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_17]] : index
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_9]] : index // CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_19]], %[[VAL_9]] : index
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex>
// CHECK: %[[VAL_23:.*]]:2 = scf.while (%[[VAL_24:.*]] = %[[VAL_20]], %[[VAL_25:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_23:.*]]:2 = scf.while (%[[VAL_24:.*]] = %[[VAL_20]], %[[VAL_25:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) {
Show All 28 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_42:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16x8xf32> // CHECK: %[[VAL_42:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_42]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_42]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_dds(%arga: tensor<32x16x8xf32, #Tdds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_dds(%arga: tensor<32x16x8xf32, #Tdds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdds>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdds>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_dds( // CHECK-LABEL: func @mul_dds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_13]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_13]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_6]] to %[[VAL_4]] step %[[VAL_7]] {
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_6]] to %[[VAL_5]] step %[[VAL_7]] {
// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_14]], %[[VAL_5]] : index // CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_14]], %[[VAL_5]] : index
// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_16]], %[[VAL_15]] : index // CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_16]], %[[VAL_15]] : index
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_7]] : index // CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_17]], %[[VAL_7]] : index
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_7]] { // CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_7]] {
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex>
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xf32> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xf32>
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_14]], %[[VAL_15]], %[[VAL_22]]] : memref<32x16x8xf32> // CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_14]], %[[VAL_15]], %[[VAL_22]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_25:.*]] = arith.mulf %[[VAL_23]], %[[VAL_24]] : f32 // CHECK: %[[VAL_25:.*]] = arith.mulf %[[VAL_23]], %[[VAL_24]] : f32
// CHECK: memref.store %[[VAL_25]], %[[VAL_13]]{{\[}}%[[VAL_14]], %[[VAL_15]], %[[VAL_22]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_25]], %[[VAL_13]]{{\[}}%[[VAL_14]], %[[VAL_15]], %[[VAL_22]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16x8xf32> // CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_26]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_26]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_dds(%arga: tensor<32x16x8xf32, #Tdds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_dds(%arga: tensor<32x16x8xf32, #Tdds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdds>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdds>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_dsd( // CHECK-LABEL: func @add_dsd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant true // CHECK-DAG: %[[VAL_6:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_14]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_14]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_7]] to %[[VAL_3]] step %[[VAL_8]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_7]] to %[[VAL_3]] step %[[VAL_8]] {
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_8]] : index // CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_8]] : index
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_16]], %[[VAL_21:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_16]], %[[VAL_21:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_18]] : index // CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_18]] : index
// CHECK: scf.condition(%[[VAL_22]]) %[[VAL_20]], %[[VAL_21]] : index, index // CHECK: scf.condition(%[[VAL_22]]) %[[VAL_20]], %[[VAL_21]] : index, index
// CHECK: } do { // CHECK: } do {
Show All 32 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_43:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32> // CHECK: %[[VAL_43:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_43]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_43]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_dsd(%arga: tensor<32x16x8xf32, #Tdsd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_dsd(%arga: tensor<32x16x8xf32, #Tdsd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdsd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdsd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_dsd( // CHECK-LABEL: func @mul_dsd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// 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:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_12]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_12]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_13:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_13]], %[[VAL_6]] : index // CHECK: %[[VAL_15:.*]] = arith.addi %[[VAL_13]], %[[VAL_6]] : index
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_14]] to %[[VAL_16]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_14]] to %[[VAL_16]] step %[[VAL_6]] {
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_17]], %[[VAL_4]] : index // CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_17]], %[[VAL_4]] : index
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : index // CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : index
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xf32> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xf32>
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_13]], %[[VAL_18]], %[[VAL_19]]] : memref<32x16x8xf32> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_13]], %[[VAL_18]], %[[VAL_19]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_24:.*]] = arith.mulf %[[VAL_22]], %[[VAL_23]] : f32 // CHECK: %[[VAL_24:.*]] = arith.mulf %[[VAL_22]], %[[VAL_23]] : f32
// CHECK: memref.store %[[VAL_24]], %[[VAL_12]]{{\[}}%[[VAL_13]], %[[VAL_18]], %[[VAL_19]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_24]], %[[VAL_12]]{{\[}}%[[VAL_13]], %[[VAL_18]], %[[VAL_19]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_25:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16x8xf32> // CHECK: %[[VAL_25:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_25]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_25]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_dsd(%arga: tensor<32x16x8xf32, #Tdsd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_dsd(%arga: tensor<32x16x8xf32, #Tdsd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdsd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdsd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_dss( // CHECK-LABEL: func @add_dss(
Show All 9 Lines
// CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_17]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_17]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_8]] to %[[VAL_4]] step %[[VAL_9]] { // CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_8]] to %[[VAL_4]] step %[[VAL_9]] {
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref<?xindex>
// CHECK: %[[VAL_20:.*]] = arith.addi %[[VAL_18]], %[[VAL_9]] : index // CHECK: %[[VAL_20:.*]] = arith.addi %[[VAL_18]], %[[VAL_9]] : index
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_20]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_20]]] : memref<?xindex>
// CHECK: %[[VAL_22:.*]]:2 = scf.while (%[[VAL_23:.*]] = %[[VAL_19]], %[[VAL_24:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_22:.*]]:2 = scf.while (%[[VAL_23:.*]] = %[[VAL_19]], %[[VAL_24:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_21]] : index // CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_21]] : index
// CHECK: scf.condition(%[[VAL_25]]) %[[VAL_23]], %[[VAL_24]] : index, index // CHECK: scf.condition(%[[VAL_25]]) %[[VAL_23]], %[[VAL_24]] : index, index
// CHECK: } do { // CHECK: } do {
▲ Show 20 Lines • Show All 56 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_62:.*]] = bufferization.to_tensor %[[VAL_17]] : memref<32x16x8xf32> // CHECK: %[[VAL_62:.*]] = bufferization.to_tensor %[[VAL_17]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_62]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_62]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_dss(%arga: tensor<32x16x8xf32, #Tdss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_dss(%arga: tensor<32x16x8xf32, #Tdss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdss>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdss>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_dss( // CHECK-LABEL: func @mul_dss(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_14]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_14]] : memref<32x16x8xf32>)
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_6]] : index // CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_6]] : index
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_16]] to %[[VAL_18]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_16]] to %[[VAL_18]] step %[[VAL_6]] {
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index // CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_19]], %[[VAL_6]] : index
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_24:.*]] = %[[VAL_21]] to %[[VAL_23]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_24:.*]] = %[[VAL_21]] to %[[VAL_23]] step %[[VAL_6]] {
// CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_24]]] : memref<?xindex> // CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_24]]] : memref<?xindex>
// CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_24]]] : memref<?xf32> // CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_24]]] : memref<?xf32>
// CHECK: %[[VAL_27:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_15]], %[[VAL_20]], %[[VAL_25]]] : memref<32x16x8xf32> // CHECK: %[[VAL_27:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_15]], %[[VAL_20]], %[[VAL_25]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_28:.*]] = arith.mulf %[[VAL_26]], %[[VAL_27]] : f32 // CHECK: %[[VAL_28:.*]] = arith.mulf %[[VAL_26]], %[[VAL_27]] : f32
// CHECK: memref.store %[[VAL_28]], %[[VAL_14]]{{\[}}%[[VAL_15]], %[[VAL_20]], %[[VAL_25]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_28]], %[[VAL_14]]{{\[}}%[[VAL_15]], %[[VAL_20]], %[[VAL_25]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_29:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32> // CHECK: %[[VAL_29:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_29]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_29]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_dss(%arga: tensor<32x16x8xf32, #Tdss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_dss(%arga: tensor<32x16x8xf32, #Tdss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tdss>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tdss>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_sdd( // CHECK-LABEL: func @add_sdd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant true // CHECK-DAG: %[[VAL_6:.*]] = arith.constant true
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_14]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_14]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_8]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_8]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_15]], %[[VAL_19:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_17:.*]]:2 = scf.while (%[[VAL_18:.*]] = %[[VAL_15]], %[[VAL_19:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_16]] : index // CHECK: %[[VAL_20:.*]] = arith.cmpi ult, %[[VAL_18]], %[[VAL_16]] : index
// CHECK: scf.condition(%[[VAL_20]]) %[[VAL_18]], %[[VAL_19]] : index, index // CHECK: scf.condition(%[[VAL_20]]) %[[VAL_18]], %[[VAL_19]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_21:.*]]: index, %[[VAL_22:.*]]: index): // CHECK: ^bb0(%[[VAL_21:.*]]: index, %[[VAL_22:.*]]: index):
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<?xindex>
Show All 37 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_46:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32> // CHECK: %[[VAL_46:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_46]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_46]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_sdd(%arga: tensor<32x16x8xf32, #Tsdd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_sdd(%arga: tensor<32x16x8xf32, #Tsdd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsdd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsdd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_sdd( // CHECK-LABEL: func @mul_sdd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 16 : index
// 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:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_12]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_12]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_13]] to %[[VAL_14]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_13]] to %[[VAL_14]] step %[[VAL_6]] {
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_15]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_15]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
// CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_15]], %[[VAL_3]] : index // CHECK: %[[VAL_18:.*]] = arith.muli %[[VAL_15]], %[[VAL_3]] : index
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_17]] : index // CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_17]] : index
// CHECK: scf.for %[[VAL_20:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_20:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_21:.*]] = arith.muli %[[VAL_19]], %[[VAL_4]] : index // CHECK: %[[VAL_21:.*]] = arith.muli %[[VAL_19]], %[[VAL_4]] : index
// CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_21]], %[[VAL_20]] : index // CHECK: %[[VAL_22:.*]] = arith.addi %[[VAL_21]], %[[VAL_20]] : index
// CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xf32> // CHECK: %[[VAL_23:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<?xf32>
// CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]], %[[VAL_17]], %[[VAL_20]]] : memref<32x16x8xf32> // CHECK: %[[VAL_24:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]], %[[VAL_17]], %[[VAL_20]]] : memref<32x16x8xf32>
// CHECK: %[[VAL_25:.*]] = arith.mulf %[[VAL_23]], %[[VAL_24]] : f32 // CHECK: %[[VAL_25:.*]] = arith.mulf %[[VAL_23]], %[[VAL_24]] : f32
// CHECK: memref.store %[[VAL_25]], %[[VAL_12]]{{\[}}%[[VAL_16]], %[[VAL_17]], %[[VAL_20]]] : memref<32x16x8xf32> // CHECK: memref.store %[[VAL_25]], %[[VAL_12]]{{\[}}%[[VAL_16]], %[[VAL_17]], %[[VAL_20]]] : memref<32x16x8xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16x8xf32> // CHECK: %[[VAL_26:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_26]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_26]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_sdd(%arga: tensor<32x16x8xf32, #Tsdd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_sdd(%arga: tensor<32x16x8xf32, #Tsdd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsdd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsdd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_sds( // CHECK-LABEL: func @add_sds(
Show All 9 Lines
// CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_9:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_17]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_17]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_8]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_8]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_9]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_9]]] : memref<?xindex>
// CHECK: %[[VAL_20:.*]]:2 = scf.while (%[[VAL_21:.*]] = %[[VAL_18]], %[[VAL_22:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_20:.*]]:2 = scf.while (%[[VAL_21:.*]] = %[[VAL_18]], %[[VAL_22:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_19]] : index // CHECK: %[[VAL_23:.*]] = arith.cmpi ult, %[[VAL_21]], %[[VAL_19]] : index
// CHECK: scf.condition(%[[VAL_23]]) %[[VAL_21]], %[[VAL_22]] : index, index // CHECK: scf.condition(%[[VAL_23]]) %[[VAL_21]], %[[VAL_22]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_24:.*]]: index, %[[VAL_25:.*]]: index): // CHECK: ^bb0(%[[VAL_24:.*]]: index, %[[VAL_25:.*]]: index):
// CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_24]]] : memref<?xindex> // CHECK: %[[VAL_26:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_24]]] : memref<?xindex>
▲ Show 20 Lines • Show All 61 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_65:.*]] = bufferization.to_tensor %[[VAL_17]] : memref<32x16x8xf32> // CHECK: %[[VAL_65:.*]] = bufferization.to_tensor %[[VAL_17]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_65]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_65]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_sds(%arga: tensor<32x16x8xf32, #Tsds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_sds(%arga: tensor<32x16x8xf32, #Tsds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsds>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsds>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_sds( // CHECK-LABEL: func @mul_sds(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "dense", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_12:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_14]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_14]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_6]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_15]] to %[[VAL_16]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_17:.*]] = %[[VAL_15]] to %[[VAL_16]] step %[[VAL_6]] {
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_19:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_17]], %[[VAL_4]] : index // CHECK: %[[VAL_20:.*]] = arith.muli %[[VAL_17]], %[[VAL_4]] : index
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : index // CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_20]], %[[VAL_19]] : index
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_21]]] : memref<?xindex>
Show All 9 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32> // CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_14]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_30]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_30]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_sds(%arga: tensor<32x16x8xf32, #Tsds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_sds(%arga: tensor<32x16x8xf32, #Tsds>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsds>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsds>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_ssd( // CHECK-LABEL: func @add_ssd(
Show All 9 Lines
// CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_14:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_16:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_16]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_16]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_7]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_8]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_8]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_17]], %[[VAL_21:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_19:.*]]:2 = scf.while (%[[VAL_20:.*]] = %[[VAL_17]], %[[VAL_21:.*]] = %[[VAL_7]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_18]] : index // CHECK: %[[VAL_22:.*]] = arith.cmpi ult, %[[VAL_20]], %[[VAL_18]] : index
// CHECK: scf.condition(%[[VAL_22]]) %[[VAL_20]], %[[VAL_21]] : index, index // CHECK: scf.condition(%[[VAL_22]]) %[[VAL_20]], %[[VAL_21]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index): // CHECK: ^bb0(%[[VAL_23:.*]]: index, %[[VAL_24:.*]]: index):
// CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_23]]] : memref<?xindex> // CHECK: %[[VAL_25:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_23]]] : memref<?xindex>
▲ Show 20 Lines • Show All 65 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_66:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16x8xf32> // CHECK: %[[VAL_66:.*]] = bufferization.to_tensor %[[VAL_16]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_66]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_66]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_ssd(%arga: tensor<32x16x8xf32, #Tssd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_ssd(%arga: tensor<32x16x8xf32, #Tssd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tssd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tssd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_ssd( // CHECK-LABEL: func @mul_ssd(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 8 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 8 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_13]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_13]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_14]] to %[[VAL_15]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_14]] to %[[VAL_15]] step %[[VAL_5]] {
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xindex> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_16]]] : memref<?xindex>
// CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_16]], %[[VAL_5]] : index // CHECK: %[[VAL_19:.*]] = arith.addi %[[VAL_16]], %[[VAL_5]] : index
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_18]] to %[[VAL_20]] step %[[VAL_5]] {
Show All 9 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_29:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16x8xf32> // CHECK: %[[VAL_29:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_29]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_29]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_ssd(%arga: tensor<32x16x8xf32, #Tssd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_ssd(%arga: tensor<32x16x8xf32, #Tssd>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tssd>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tssd>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @add_sss( // CHECK-LABEL: func @add_sss(
Show All 11 Lines
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_16:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_17:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_19:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_19:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_19]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_19]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_8]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_8]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_9]]] : memref<?xindex> // CHECK: %[[VAL_21:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_9]]] : memref<?xindex>
// CHECK: %[[VAL_22:.*]]:2 = scf.while (%[[VAL_23:.*]] = %[[VAL_20]], %[[VAL_24:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) { // CHECK: %[[VAL_22:.*]]:2 = scf.while (%[[VAL_23:.*]] = %[[VAL_20]], %[[VAL_24:.*]] = %[[VAL_8]]) : (index, index) -> (index, index) {
// CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_21]] : index // CHECK: %[[VAL_25:.*]] = arith.cmpi ult, %[[VAL_23]], %[[VAL_21]] : index
// CHECK: scf.condition(%[[VAL_25]]) %[[VAL_23]], %[[VAL_24]] : index, index // CHECK: scf.condition(%[[VAL_25]]) %[[VAL_23]], %[[VAL_24]] : index, index
// CHECK: } do { // CHECK: } do {
// CHECK: ^bb0(%[[VAL_26:.*]]: index, %[[VAL_27:.*]]: index): // CHECK: ^bb0(%[[VAL_26:.*]]: index, %[[VAL_27:.*]]: index):
// CHECK: %[[VAL_28:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_26]]] : memref<?xindex> // CHECK: %[[VAL_28:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_26]]] : memref<?xindex>
▲ Show 20 Lines • Show All 89 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_85:.*]] = bufferization.to_tensor %[[VAL_19]] : memref<32x16x8xf32> // CHECK: %[[VAL_85:.*]] = bufferization.to_tensor %[[VAL_19]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_85]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_85]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @add_sss(%arga: tensor<32x16x8xf32, #Tsss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @add_sss(%arga: tensor<32x16x8xf32, #Tsss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsss>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsss>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.addf %a, %b : f32 %0 = arith.addf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
// CHECK-LABEL: func @mul_sss( // CHECK-LABEL: func @mul_sss(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<32x16x8xf32>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 1 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 2 : index} : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_12:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x16x8xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed", "compressed" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_1]] : memref<32x16x8xf32>
// CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32> // CHECK-DAG: %[[VAL_15:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32x16x8xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_15]] : memref<32x16x8xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_15]] : memref<32x16x8xf32>)
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_16]] to %[[VAL_17]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_16]] to %[[VAL_17]] step %[[VAL_5]] {
// CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex> // CHECK: %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex>
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
// CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_18]], %[[VAL_5]] : index // CHECK: %[[VAL_21:.*]] = arith.addi %[[VAL_18]], %[[VAL_5]] : index
// CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_21]]] : memref<?xindex> // CHECK: %[[VAL_22:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_21]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_23:.*]] = %[[VAL_20]] to %[[VAL_22]] step %[[VAL_5]] { // CHECK: scf.for %[[VAL_23:.*]] = %[[VAL_20]] to %[[VAL_22]] step %[[VAL_5]] {
Show All 11 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16x8xf32> // CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_15]] : memref<32x16x8xf32>
// CHECK: return %[[VAL_33]] : tensor<32x16x8xf32> // CHECK: return %[[VAL_33]] : tensor<32x16x8xf32>
// CHECK: } // CHECK: }
func.func @mul_sss(%arga: tensor<32x16x8xf32, #Tsss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> { func.func @mul_sss(%arga: tensor<32x16x8xf32, #Tsss>, %argb: tensor<32x16x8xf32>, %argx: tensor<32x16x8xf32>) -> tensor<32x16x8xf32> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16x8xf32, #Tsss>, tensor<32x16x8xf32>) ins(%arga, %argb: tensor<32x16x8xf32, #Tsss>, tensor<32x16x8xf32>)
outs(%argx: tensor<32x16x8xf32>) { inits(%argx: tensor<32x16x8xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<32x16x8xf32> } -> tensor<32x16x8xf32>
return %0 : tensor<32x16x8xf32> return %0 : tensor<32x16x8xf32>
} }
#trait_kernel_3d = { #trait_kernel_3d = {
▲ Show 20 Lines • Show All 49 Lines • ▼ Show 20 Lines
// CHECK: return %[[VAL_34]] : tensor<?x?xf32> // CHECK: return %[[VAL_34]] : tensor<?x?xf32>
// CHECK: } // CHECK: }
func.func @kernel_3d(%arga: tensor<?x?xf32>, func.func @kernel_3d(%arga: tensor<?x?xf32>,
%argb: tensor<?x?x?xf32, #Tdds>, %argb: tensor<?x?x?xf32, #Tdds>,
%argc: tensor<?x?xf32>, %argc: tensor<?x?xf32>,
%argd: tensor<?x?xf32>) -> tensor<?x?xf32> { %argd: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.generic #trait_kernel_3d %0 = linalg.generic #trait_kernel_3d
ins(%argb, %argc, %argd: tensor<?x?x?xf32, #Tdds>, tensor<?x?xf32>, tensor<?x?xf32>) ins(%argb, %argc, %argd: tensor<?x?x?xf32, #Tdds>, tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arga: tensor<?x?xf32>) { inits(%arga: tensor<?x?xf32>) {
^bb(%b: f32, %c: f32, %d: f32, %a: f32): ^bb(%b: f32, %c: f32, %d: f32, %a: f32):
%0 = arith.mulf %b, %c : f32 %0 = arith.mulf %b, %c : f32
%1 = arith.mulf %0, %d : f32 %1 = arith.mulf %0, %d : f32
%2 = arith.addf %1, %a : f32 %2 = arith.addf %1, %a : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
Show All 39 Lines
// CHECK: } // CHECK: }
// CHECK: memref.store %[[VAL_14]], %[[VAL_10]][] : memref<f32> // CHECK: memref.store %[[VAL_14]], %[[VAL_10]][] : memref<f32>
// CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<f32> // CHECK: %[[VAL_34:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<f32>
// CHECK: return %[[VAL_34]] : tensor<f32> // CHECK: return %[[VAL_34]] : tensor<f32>
// CHECK: } // CHECK: }
func.func @sum_reduction(%arga: tensor<10x20x30xf32, #Tsss>, %argx: tensor<f32>) -> tensor<f32> { func.func @sum_reduction(%arga: tensor<10x20x30xf32, #Tsss>, %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_sum_reduction %0 = linalg.generic #trait_sum_reduction
ins(%arga: tensor<10x20x30xf32, #Tsss>) ins(%arga: tensor<10x20x30xf32, #Tsss>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %x, %a : f32 %0 = arith.addf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
#trait_sum_reduction_inv = { #trait_sum_reduction_inv = {
Show All 37 Lines
// CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<f32> // CHECK: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<f32>
// CHECK: return %[[VAL_30]] : tensor<f32> // CHECK: return %[[VAL_30]] : tensor<f32>
// CHECK: } // CHECK: }
func.func @sum_reduction_inv(%arga: tensor<?x?x?xf32>, func.func @sum_reduction_inv(%arga: tensor<?x?x?xf32>,
%argb: tensor<?xf32, #Td>, %argb: tensor<?xf32, #Td>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_sum_reduction_inv %0 = linalg.generic #trait_sum_reduction_inv
ins(%arga, %argb: tensor<?x?x?xf32>, tensor<?xf32, #Td>) ins(%arga, %argb: tensor<?x?x?xf32>, tensor<?xf32, #Td>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.addf %x, %0 : f32 %1 = arith.addf %x, %0 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
Show All 18 Lines
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 20 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 20 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 30 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 30 : index
// CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_7:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_8:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<20xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<20xf32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<30xf32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<30xf32>
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]] : memref<10x20x30xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_3]] : memref<10x20x30xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_13]] : memref<10x20x30xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_13]] : memref<10x20x30xf32>)
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_8]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_7]] to %[[VAL_4]] step %[[VAL_8]] {
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_14]]] : memref<?xf32> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_14]]] : memref<?xf32>
// CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_7]] to %[[VAL_5]] step %[[VAL_8]] { // CHECK: scf.for %[[VAL_16:.*]] = %[[VAL_7]] to %[[VAL_5]] step %[[VAL_8]] {
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]]] : memref<20xf32> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_16]]] : memref<20xf32>
// CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_7]] to %[[VAL_6]] step %[[VAL_8]] { // CHECK: scf.for %[[VAL_18:.*]] = %[[VAL_7]] to %[[VAL_6]] step %[[VAL_8]] {
// CHECK: %[[VAL_19:.*]] = arith.mulf %[[VAL_15]], %[[VAL_17]] : f32 // CHECK: %[[VAL_19:.*]] = arith.mulf %[[VAL_15]], %[[VAL_17]] : f32
// CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<30xf32> // CHECK: %[[VAL_20:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<30xf32>
// CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32 // CHECK: %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32
// CHECK: memref.store %[[VAL_21]], %[[VAL_13]]{{\[}}%[[VAL_14]], %[[VAL_16]], %[[VAL_18]]] : memref<10x20x30xf32> // CHECK: memref.store %[[VAL_21]], %[[VAL_13]]{{\[}}%[[VAL_14]], %[[VAL_16]], %[[VAL_18]]] : memref<10x20x30xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<10x20x30xf32> // CHECK: %[[VAL_22:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<10x20x30xf32>
// CHECK: return %[[VAL_22]] : tensor<10x20x30xf32> // CHECK: return %[[VAL_22]] : tensor<10x20x30xf32>
// CHECK: } // CHECK: }
func.func @invariants(%arga: tensor<10xf32, #Td>, func.func @invariants(%arga: tensor<10xf32, #Td>,
%argb: tensor<20xf32>, %argb: tensor<20xf32>,
%argc: tensor<30xf32>, %argc: tensor<30xf32>,
%argx: tensor<10x20x30xf32>) -> tensor<10x20x30xf32> { %argx: tensor<10x20x30xf32>) -> tensor<10x20x30xf32> {
%0 = linalg.generic #trait_invariants %0 = linalg.generic #trait_invariants
ins(%arga, %argb, %argc : tensor<10xf32, #Td>, tensor<20xf32>, tensor<30xf32>) ins(%arga, %argb, %argc : tensor<10xf32, #Td>, tensor<20xf32>, tensor<30xf32>)
outs(%argx: tensor<10x20x30xf32>) { inits(%argx: tensor<10x20x30xf32>) {
^bb(%a: f32, %b: f32, %c: f32, %x: f32): ^bb(%a: f32, %b: f32, %c: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.mulf %0, %c : f32 %1 = arith.mulf %0, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<10x20x30xf32> } -> tensor<10x20x30xf32>
return %0 : tensor<10x20x30xf32> return %0 : tensor<10x20x30xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_affine.mlir

Show All 39 Lines
// CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf32> // CHECK: %[[VAL_21:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf32>
// CHECK: return %[[VAL_21]] : tensor<32xf32> // CHECK: return %[[VAL_21]] : tensor<32xf32>
// CHECK: } // CHECK: }
func.func @mul_inv_dense1d(%arga: tensor<32xf32, #SpVec>, func.func @mul_inv_dense1d(%arga: tensor<32xf32, #SpVec>,
%argb: tensor<4xf32>, %argb: tensor<4xf32>,
%argx: tensor<32xf32>) -> tensor<32xf32> { %argx: tensor<32xf32>) -> tensor<32xf32> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga, %argb: tensor<32xf32, #SpVec>, tensor<4xf32>) ins(%arga, %argb: tensor<32xf32, #SpVec>, tensor<4xf32>)
outs(%argx: tensor<32xf32>) { inits(%argx: tensor<32xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.addf %x, %0 : f32 %1 = arith.addf %x, %0 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32xf32> } -> tensor<32xf32>
return %0 : tensor<32xf32> return %0 : tensor<32xf32>
} }
Show All 15 Lines
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 2 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32xi32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<34xi32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<34xi32>
// CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xi32> // CHECK-DAG: %[[VAL_11:.*]] = bufferization.to_memref %[[VAL_2]] : memref<32xi32>
// CHECK: linalg.fill ins(%[[ZERO]] : i32) outs(%[[VAL_11]] : memref<32xi32>) // CHECK: linalg.fill ins(%[[ZERO]] : i32) inits(%[[VAL_11]] : memref<32xi32>)
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_13]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_12]] to %[[VAL_13]] step %[[VAL_4]] {
// CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xindex> // CHECK: %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xindex>
// CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<?xi32> // CHECK: %[[VAL_16:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_14]]] : memref<?xi32>
// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_5]] : index // CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_15]], %[[VAL_5]] : index
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<34xi32> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_17]]] : memref<34xi32>
// CHECK: %[[VAL_19:.*]] = arith.andi %[[VAL_16]], %[[VAL_18]] : i32 // CHECK: %[[VAL_19:.*]] = arith.andi %[[VAL_16]], %[[VAL_18]] : i32
// CHECK: memref.store %[[VAL_19]], %[[VAL_11]]{{\[}}%[[VAL_15]]] : memref<32xi32> // CHECK: memref.store %[[VAL_19]], %[[VAL_11]]{{\[}}%[[VAL_15]]] : memref<32xi32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi32> // CHECK: %[[VAL_20:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi32>
// CHECK: return %[[VAL_20]] : tensor<32xi32> // CHECK: return %[[VAL_20]] : tensor<32xi32>
// CHECK: } // CHECK: }
func.func @and_affine_dense1d(%arga: tensor<32xi32, #SpVec>, func.func @and_affine_dense1d(%arga: tensor<32xi32, #SpVec>,
%argb: tensor<34xi32>, %argb: tensor<34xi32>,
%argx: tensor<32xi32>) -> tensor<32xi32> { %argx: tensor<32xi32>) -> tensor<32xi32> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi32, #SpVec>, tensor<34xi32>) ins(%arga, %argb: tensor<32xi32, #SpVec>, tensor<34xi32>)
outs(%argx: tensor<32xi32>) { inits(%argx: tensor<32xi32>) {
^bb(%a: i32, %b: i32, %x: i32): ^bb(%a: i32, %b: i32, %x: i32):
%0 = arith.andi %a, %b : i32 %0 = arith.andi %a, %b : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<32xi32> } -> tensor<32xi32>
return %0 : tensor<32xi32> return %0 : tensor<32xi32>
} }
#trait3 = { #trait3 = {
Show All 39 Lines
// CHECK: %[[VAL_27:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf64> // CHECK: %[[VAL_27:.*]] = bufferization.to_tensor %[[VAL_13]] : memref<32x16xf64>
// CHECK: return %[[VAL_27]] : tensor<32x16xf64> // CHECK: return %[[VAL_27]] : tensor<32x16xf64>
// CHECK: } // CHECK: }
func.func @mul_affine_dense2d(%arga: tensor<32x16xf64, #CSR>, func.func @mul_affine_dense2d(%arga: tensor<32x16xf64, #CSR>,
%argb: tensor<34x19xf64>, %argb: tensor<34x19xf64>,
%argx: tensor<32x16xf64>) -> tensor<32x16xf64> { %argx: tensor<32x16xf64>) -> tensor<32x16xf64> {
%0 = linalg.generic #trait3 %0 = linalg.generic #trait3
ins(%arga, %argb: tensor<32x16xf64, #CSR>, tensor<34x19xf64>) ins(%arga, %argb: tensor<32x16xf64, #CSR>, tensor<34x19xf64>)
outs(%argx: tensor<32x16xf64>) { inits(%argx: tensor<32x16xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.mulf %a, %b : f64 %0 = arith.mulf %a, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32x16xf64> } -> tensor<32x16xf64>
return %0 : tensor<32x16xf64> return %0 : tensor<32x16xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_concat.mlir

Show All 15 Lines
// CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8 // CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8
// CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index // CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index
// CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<5x4xf64> // CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<5x4xf64>
// CHECK: linalg.fill ins(%[[TMP_cst]] : f64) outs(%[[TMP_0]] : memref<5x4xf64>) // CHECK: linalg.fill ins(%[[TMP_cst]] : f64) inits(%[[TMP_0]] : memref<5x4xf64>)
// CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] {
// CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c4]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c4]] step %[[TMP_c1]] {
// CHECK: %[[TMP_12:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<2x4xf64> // CHECK: %[[TMP_12:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<2x4xf64>
// CHECK: %[[TMP_13:.*]] = arith.cmpf une, %[[TMP_12]], %[[TMP_cst]] : f64 // CHECK: %[[TMP_13:.*]] = arith.cmpf une, %[[TMP_12]], %[[TMP_cst]] : f64
// CHECK: scf.if %[[TMP_13]] { // CHECK: scf.if %[[TMP_13]] {
// CHECK: memref.store %[[TMP_12]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<5x4xf64> // CHECK: memref.store %[[TMP_12]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<5x4xf64>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
▲ Show 20 Lines • Show All 212 Lines • ▼ Show 20 Lines
// CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8 // CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8
// CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index // CHECK-DAG: %[[TMP_c4:.*]] = arith.constant 4 : index
// CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<4x5xf64> // CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<4x5xf64>
// CHECK: linalg.fill ins(%[[TMP_cst]] : f64) outs(%[[TMP_0]] : memref<4x5xf64>) // CHECK: linalg.fill ins(%[[TMP_cst]] : f64) inits(%[[TMP_0]] : memref<4x5xf64>)
// CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c4]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c4]] step %[[TMP_c1]] {
// CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] {
// CHECK: %[[TMP_12:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<4x2xf64> // CHECK: %[[TMP_12:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<4x2xf64>
// CHECK: %[[TMP_13:.*]] = arith.cmpf une, %[[TMP_12]], %[[TMP_cst]] : f64 // CHECK: %[[TMP_13:.*]] = arith.cmpf une, %[[TMP_12]], %[[TMP_cst]] : f64
// CHECK: scf.if %[[TMP_13]] { // CHECK: scf.if %[[TMP_13]] {
// CHECK: memref.store %[[TMP_12]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<4x5xf64> // CHECK: memref.store %[[TMP_12]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<4x5xf64>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
▲ Show 20 Lines • Show All 44 Lines • ▼ Show 20 Lines
// CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32 // CHECK-DAG: %[[TMP_c0_i32:.*]] = arith.constant 0 : i32
// CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8 // CHECK-DAG: %[[TMP_c8_i8:.*]] = arith.constant 8 : i8
// CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64 // CHECK-DAG: %[[TMP_cst:.*]] = arith.constant 0.000000e+00 : f64
// CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index // CHECK-DAG: %[[TMP_c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index // CHECK-DAG: %[[TMP_c3:.*]] = arith.constant 3 : index
// CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index // CHECK-DAG: %[[TMP_c1:.*]] = arith.constant 1 : index
// CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<3x5xf64> // CHECK: %[[TMP_0:.*]] = memref.alloc() : memref<3x5xf64>
// CHECK: %[[TMP_1:.*]] = memref.cast %[[TMP_0]] : memref<3x5xf64> to memref<?x?xf64> // CHECK: %[[TMP_1:.*]] = memref.cast %[[TMP_0]] : memref<3x5xf64> to memref<?x?xf64>
// CHECK: linalg.fill ins(%[[TMP_cst]] : f64) outs(%[[TMP_0]] : memref<3x5xf64>) // CHECK: linalg.fill ins(%[[TMP_cst]] : f64) inits(%[[TMP_0]] : memref<3x5xf64>)
// CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c3]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg2:.*]] = %[[TMP_c0]] to %[[TMP_c3]] step %[[TMP_c1]] {
// CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] { // CHECK: scf.for %[[TMP_arg3:.*]] = %[[TMP_c0]] to %[[TMP_c2]] step %[[TMP_c1]] {
// CHECK: %[[TMP_13:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<3x2xf64> // CHECK: %[[TMP_13:.*]] = tensor.extract %[[TMP_arg0]][%[[TMP_arg2]], %[[TMP_arg3]]] : tensor<3x2xf64>
// CHECK: %[[TMP_14:.*]] = arith.cmpf une, %[[TMP_13]], %[[TMP_cst]] : f64 // CHECK: %[[TMP_14:.*]] = arith.cmpf une, %[[TMP_13]], %[[TMP_cst]] : f64
// CHECK: scf.if %[[TMP_14]] { // CHECK: scf.if %[[TMP_14]] {
// CHECK: memref.store %[[TMP_13]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<3x5xf64> // CHECK: memref.store %[[TMP_13]], %[[TMP_0]][%[[TMP_arg2]], %[[TMP_arg3]]] : memref<3x5xf64>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
Show All 38 Lines

mlir/test/Dialect/SparseTensor/sparse_expand.mlir

Show First 20 Lines • Show All 47 Lines • ▼ Show 20 Lines
// CHECK-CONVERT-LABEL: func @kernel( // CHECK-CONVERT-LABEL: func @kernel(
// CHECK-CONVERT: %[[C:.*]] = arith.constant 0 : index // CHECK-CONVERT: %[[C:.*]] = arith.constant 0 : index
// CHECK-CONVERT: %{{.*}} = call @sparseDimSize // CHECK-CONVERT: %{{.*}} = call @sparseDimSize
// CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor // CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor
// CHECK-CONVERT: %[[S:.*]] = call @sparseDimSize(%[[N]], %[[C]]) // CHECK-CONVERT: %[[S:.*]] = call @sparseDimSize(%[[N]], %[[C]])
// CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64> // CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
// CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1> // CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
// CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex> // CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<?xf64>)
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<?xi1>)
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: call @expInsertF64 // CHECK-CONVERT: call @expInsertF64
// CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64> // CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64>
// CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1> // CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1>
// CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex> // CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex>
// CHECK-CONVERT: call @endInsert // CHECK-CONVERT: call @endInsert
// //
func.func @kernel(%arga: tensor<?x?xf64, #DCSC>) -> tensor<?xf64, #SV> { func.func @kernel(%arga: tensor<?x?xf64, #DCSC>) -> tensor<?xf64, #SV> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%n = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSC> %n = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSC>
%v = bufferization.alloc_tensor(%n) : tensor<?xf64, #SV> %v = bufferization.alloc_tensor(%n) : tensor<?xf64, #SV>
%0 = linalg.generic #rowsum %0 = linalg.generic #rowsum
ins(%arga: tensor<?x?xf64, #DCSC>) ins(%arga: tensor<?x?xf64, #DCSC>)
outs(%v: tensor<?xf64, #SV>) { inits(%v: tensor<?xf64, #SV>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = arith.addf %x, %a : f64 %1 = arith.addf %x, %a : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SV> } -> tensor<?xf64, #SV>
return %0 : tensor<?xf64, #SV> return %0 : tensor<?xf64, #SV>
} }
// //
Show All 16 Lines
// CHECK-CONVERT-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-CONVERT-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-CONVERT-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-CONVERT-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-CONVERT-DAG: %[[C4:.*]] = arith.constant 4 : index // CHECK-CONVERT-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-CONVERT-DAG: %[[C8:.*]] = arith.constant 8 : index // CHECK-CONVERT-DAG: %[[C8:.*]] = arith.constant 8 : index
// CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor // CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor
// CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[C4]]) : memref<?xf64> // CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[C4]]) : memref<?xf64>
// CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[C4]]) : memref<?xi1> // CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[C4]]) : memref<?xi1>
// CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[C4]]) : memref<?xindex> // CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[C4]]) : memref<?xindex>
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<?xf64>)
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<?xi1>)
// CHECK-CONVERT: scf.for %{{.*}} = %[[C0]] to %[[C8]] step %[[C1]] { // CHECK-CONVERT: scf.for %{{.*}} = %[[C0]] to %[[C8]] step %[[C1]] {
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: call @expInsertF64 // CHECK-CONVERT: call @expInsertF64
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64> // CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64>
// CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1> // CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1>
// CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex> // CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex>
// CHECK-CONVERT: call @endInsert // CHECK-CONVERT: call @endInsert
// //
func.func @matmul1(%A: tensor<8x2xf64, #CSR>, func.func @matmul1(%A: tensor<8x2xf64, #CSR>,
%B: tensor<2x4xf64, #CSR>) -> tensor<8x4xf64, #CSR> { %B: tensor<2x4xf64, #CSR>) -> tensor<8x4xf64, #CSR> {
%C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSR> %C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSR>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<8x2xf64, #CSR>, tensor<2x4xf64, #CSR>) ins(%A, %B: tensor<8x2xf64, #CSR>, tensor<2x4xf64, #CSR>)
outs(%C: tensor<8x4xf64, #CSR>) -> tensor<8x4xf64, #CSR> inits(%C: tensor<8x4xf64, #CSR>) -> tensor<8x4xf64, #CSR>
return %D: tensor<8x4xf64, #CSR> return %D: tensor<8x4xf64, #CSR>
} }
// //
// CHECK-SPARSE-LABEL: func @matmul2( // CHECK-SPARSE-LABEL: func @matmul2(
// CHECK-SPARSE-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-SPARSE-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-SPARSE-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-SPARSE-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-SPARSE-DAG: %[[C4:.*]] = arith.constant 4 : index // CHECK-SPARSE-DAG: %[[C4:.*]] = arith.constant 4 : index
Show All 12 Lines
// CHECK-CONVERT-DAG: %[[C0:.*]] = arith.constant 0 : index // CHECK-CONVERT-DAG: %[[C0:.*]] = arith.constant 0 : index
// CHECK-CONVERT-DAG: %[[C1:.*]] = arith.constant 1 : index // CHECK-CONVERT-DAG: %[[C1:.*]] = arith.constant 1 : index
// CHECK-CONVERT-DAG: %[[C4:.*]] = arith.constant 4 : index // CHECK-CONVERT-DAG: %[[C4:.*]] = arith.constant 4 : index
// CHECK-CONVERT-DAG: %[[C8:.*]] = arith.constant 8 : index // CHECK-CONVERT-DAG: %[[C8:.*]] = arith.constant 8 : index
// CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor // CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor
// CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[C8]]) : memref<?xf64> // CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[C8]]) : memref<?xf64>
// CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[C8]]) : memref<?xi1> // CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[C8]]) : memref<?xi1>
// CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[C8]]) : memref<?xindex> // CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[C8]]) : memref<?xindex>
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) inits(%[[A]] : memref<?xf64>)
// CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>) // CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) inits(%[[B]] : memref<?xi1>)
// CHECK-CONVERT: scf.for %{{.*}} = %[[C0]] to %[[C4]] step %[[C1]] { // CHECK-CONVERT: scf.for %{{.*}} = %[[C0]] to %[[C4]] step %[[C1]] {
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: scf.for {{.*}} { // CHECK-CONVERT: scf.for {{.*}} {
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: call @expInsertF64 // CHECK-CONVERT: call @expInsertF64
// CHECK-CONVERT: } // CHECK-CONVERT: }
// CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64> // CHECK-CONVERT: memref.dealloc %[[A]] : memref<?xf64>
// CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1> // CHECK-CONVERT: memref.dealloc %[[B]] : memref<?xi1>
// CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex> // CHECK-CONVERT: memref.dealloc %[[C]] : memref<?xindex>
// CHECK-CONVERT: call @endInsert // CHECK-CONVERT: call @endInsert
// //
func.func @matmul2(%A: tensor<8x2xf64, #CSC>, func.func @matmul2(%A: tensor<8x2xf64, #CSC>,
%B: tensor<2x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> { %B: tensor<2x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> {
%C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSC> %C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSC>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<8x2xf64, #CSC>, tensor<2x4xf64, #CSC>) ins(%A, %B: tensor<8x2xf64, #CSC>, tensor<2x4xf64, #CSC>)
outs(%C: tensor<8x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> inits(%C: tensor<8x4xf64, #CSC>) -> tensor<8x4xf64, #CSC>
return %D: tensor<8x4xf64, #CSC> return %D: tensor<8x4xf64, #CSC>
} }

mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir

Show All 29 Lines
// CHECK: %[[VAL_18:.*]] = llvm.mlir.null : !llvm.ptr<i8> // CHECK: %[[VAL_18:.*]] = llvm.mlir.null : !llvm.ptr<i8>
// CHECK: %[[VAL_19:.*]] = call @newSparseTensor(%[[VAL_13]], %[[VAL_15]], %[[VAL_17]], %[[VAL_4]], %[[VAL_4]], %[[VAL_3]], %[[VAL_4]], %[[VAL_18]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8> // CHECK: %[[VAL_19:.*]] = call @newSparseTensor(%[[VAL_13]], %[[VAL_15]], %[[VAL_17]], %[[VAL_4]], %[[VAL_4]], %[[VAL_3]], %[[VAL_4]], %[[VAL_18]]) : (memref<?xi8>, memref<?xindex>, memref<?xindex>, i32, i32, i32, i32, !llvm.ptr<i8>) -> !llvm.ptr<i8>
// CHECK: %[[VAL_20:.*]] = memref.alloc() : memref<300xf64> // CHECK: %[[VAL_20:.*]] = memref.alloc() : memref<300xf64>
// CHECK: %[[VAL_21:.*]] = memref.cast %[[VAL_20]] : memref<300xf64> to memref<?xf64> // CHECK: %[[VAL_21:.*]] = memref.cast %[[VAL_20]] : memref<300xf64> to memref<?xf64>
// CHECK: %[[VAL_22:.*]] = memref.alloc() : memref<300xi1> // CHECK: %[[VAL_22:.*]] = memref.alloc() : memref<300xi1>
// CHECK: %[[VAL_23:.*]] = memref.cast %[[VAL_22]] : memref<300xi1> to memref<?xi1> // CHECK: %[[VAL_23:.*]] = memref.cast %[[VAL_22]] : memref<300xi1> to memref<?xi1>
// CHECK: %[[VAL_24:.*]] = memref.alloc() : memref<300xindex> // CHECK: %[[VAL_24:.*]] = memref.alloc() : memref<300xindex>
// CHECK: %[[VAL_25:.*]] = memref.cast %[[VAL_24]] : memref<300xindex> to memref<?xindex> // CHECK: %[[VAL_25:.*]] = memref.cast %[[VAL_24]] : memref<300xindex> to memref<?xindex>
// CHECK: linalg.fill ins(%[[VAL_2]] : f64) outs(%[[VAL_20]] : memref<300xf64>) // CHECK: linalg.fill ins(%[[VAL_2]] : f64) inits(%[[VAL_20]] : memref<300xf64>)
// CHECK: linalg.fill ins(%[[VAL_7]] : i1) outs(%[[VAL_22]] : memref<300xi1>) // CHECK: linalg.fill ins(%[[VAL_7]] : i1) inits(%[[VAL_22]] : memref<300xi1>)
// CHECK: %[[VAL_26:.*]] = call @sparsePointers0(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_26:.*]] = call @sparsePointers0(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_27:.*]] = call @sparseIndices0(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_27:.*]] = call @sparseIndices0(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_28:.*]] = call @sparsePointers0(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_28:.*]] = call @sparsePointers0(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_29:.*]] = call @sparseIndices0(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_29:.*]] = call @sparseIndices0(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_30:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64> // CHECK: %[[VAL_30:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
// CHECK: %[[VAL_31:.*]] = call @sparsePointers0(%[[VAL_1]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_31:.*]] = call @sparsePointers0(%[[VAL_1]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_32:.*]] = call @sparseIndices0(%[[VAL_1]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_32:.*]] = call @sparseIndices0(%[[VAL_1]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
// CHECK: %[[VAL_33:.*]] = call @sparsePointers0(%[[VAL_1]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex> // CHECK: %[[VAL_33:.*]] = call @sparsePointers0(%[[VAL_1]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
▲ Show 20 Lines • Show All 67 Lines • ▼ Show 20 Lines
// CHECK: call @endInsert(%[[VAL_19]]) : (!llvm.ptr<i8>) -> () // CHECK: call @endInsert(%[[VAL_19]]) : (!llvm.ptr<i8>) -> ()
// CHECK: return %[[VAL_19]] : !llvm.ptr<i8> // CHECK: return %[[VAL_19]] : !llvm.ptr<i8>
// CHECK: } // CHECK: }
func.func @fill_zero_after_alloc(%arg0: tensor<100x200xf64, #DCSR>, func.func @fill_zero_after_alloc(%arg0: tensor<100x200xf64, #DCSR>,
%arg1: tensor<200x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR> { %arg1: tensor<200x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR> {
%0 = bufferization.alloc_tensor() : tensor<100x300xf64, #DCSR> %0 = bufferization.alloc_tensor() : tensor<100x300xf64, #DCSR>
%cst = arith.constant 0.000000e+00 : f64 %cst = arith.constant 0.000000e+00 : f64
%1 = linalg.fill ins(%cst : f64) %1 = linalg.fill ins(%cst : f64)
outs(%0 : tensor<100x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR> inits(%0 : tensor<100x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR>
%2 = linalg.matmul ins(%arg0, %arg1 : tensor<100x200xf64, #DCSR>, tensor<200x300xf64, #DCSR>) %2 = linalg.matmul ins(%arg0, %arg1 : tensor<100x200xf64, #DCSR>, tensor<200x300xf64, #DCSR>)
outs(%1 : tensor<100x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR> inits(%1 : tensor<100x300xf64, #DCSR>) -> tensor<100x300xf64, #DCSR>
return %2 : tensor<100x300xf64, #DCSR> return %2 : tensor<100x300xf64, #DCSR>
} }

mlir/test/Dialect/SparseTensor/sparse_fp_ops.mlir

Show First 20 Lines • Show All 48 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64>
// CHECK: return %[[VAL_14]] : tensor<32xf64> // CHECK: return %[[VAL_14]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @abs(%arga: tensor<32xf64, #SV>, func.func @abs(%arga: tensor<32xf64, #SV>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = math.absf %a : f64 %0 = math.absf %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @ceil( // CHECK-LABEL: func @ceil(
Show All 15 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64>
// CHECK: return %[[VAL_14]] : tensor<32xf64> // CHECK: return %[[VAL_14]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @ceil(%arga: tensor<32xf64, #SV>, func.func @ceil(%arga: tensor<32xf64, #SV>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = math.ceil %a : f64 %0 = math.ceil %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @floor( // CHECK-LABEL: func @floor(
Show All 15 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64>
// CHECK: return %[[VAL_14]] : tensor<32xf64> // CHECK: return %[[VAL_14]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @floor(%arga: tensor<32xf64, #SV>, func.func @floor(%arga: tensor<32xf64, #SV>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = math.floor %a : f64 %0 = math.floor %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @neg( // CHECK-LABEL: func @neg(
Show All 15 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64> // CHECK: %[[VAL_14:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<32xf64>
// CHECK: return %[[VAL_14]] : tensor<32xf64> // CHECK: return %[[VAL_14]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @neg(%arga: tensor<32xf64, #SV>, func.func @neg(%arga: tensor<32xf64, #SV>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.negf %a : f64 %0 = arith.negf %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @add( // CHECK-LABEL: func @add(
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf64> // CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf64>
// CHECK: return %[[VAL_33]] : tensor<32xf64> // CHECK: return %[[VAL_33]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @add(%arga: tensor<32xf64, #SV>, func.func @add(%arga: tensor<32xf64, #SV>,
%argb: tensor<32xf64>, %argb: tensor<32xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>) ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.addf %a, %b : f64 %0 = arith.addf %a, %b : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @sub( // CHECK-LABEL: func @sub(
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_35:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf64> // CHECK: %[[VAL_35:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xf64>
// CHECK: return %[[VAL_35]] : tensor<32xf64> // CHECK: return %[[VAL_35]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @sub(%arga: tensor<32xf64, #SV>, func.func @sub(%arga: tensor<32xf64, #SV>,
%argb: tensor<32xf64>, %argb: tensor<32xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>) ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.subf %a, %b : f64 %0 = arith.subf %a, %b : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @mul( // CHECK-LABEL: func @mul(
Show All 19 Lines
// CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf64> // CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xf64>
// CHECK: return %[[VAL_17]] : tensor<32xf64> // CHECK: return %[[VAL_17]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @mul(%arga: tensor<32xf64, #SV>, func.func @mul(%arga: tensor<32xf64, #SV>,
%argb: tensor<32xf64>, %argb: tensor<32xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>) ins(%arga, %argb: tensor<32xf64, #SV>, tensor<32xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.mulf %a, %b : f64 %0 = arith.mulf %a, %b : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @divbyc( // CHECK-LABEL: func @divbyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xf64>
// CHECK: return %[[VAL_15]] : tensor<32xf64> // CHECK: return %[[VAL_15]] : tensor<32xf64>
// CHECK: } // CHECK: }
func.func @divbyc(%arga: tensor<32xf64, #SV>, func.func @divbyc(%arga: tensor<32xf64, #SV>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%c = arith.constant 2.0 : f64 %c = arith.constant 2.0 : f64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.divf %a, %c : f64 %0 = arith.divf %a, %c : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func.func @zero_preserving_math( // CHECK-LABEL: func.func @zero_preserving_math(
Show All 22 Lines
// CHECK: %[[VAL_20:.*]] = sparse_tensor.load %[[VAL_3]] hasInserts : tensor<32xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> // CHECK: %[[VAL_20:.*]] = sparse_tensor.load %[[VAL_3]] hasInserts : tensor<32xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
// CHECK: return %[[VAL_20]] : tensor<32xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> // CHECK: return %[[VAL_20]] : tensor<32xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
// CHECK: } // CHECK: }
func.func @zero_preserving_math(%arga: tensor<32xf64, #SV>) -> tensor<32xf64, #SV> { func.func @zero_preserving_math(%arga: tensor<32xf64, #SV>) -> tensor<32xf64, #SV> {
%c32 = arith.constant 32 : index %c32 = arith.constant 32 : index
%xinp = bufferization.alloc_tensor() : tensor<32xf64, #SV> %xinp = bufferization.alloc_tensor() : tensor<32xf64, #SV>
%0 = linalg.generic #trait1 %0 = linalg.generic #trait1
ins(%arga: tensor<32xf64, #SV>) ins(%arga: tensor<32xf64, #SV>)
outs(%xinp: tensor<32xf64, #SV>) { inits(%xinp: tensor<32xf64, #SV>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = math.absf %a : f64 %0 = math.absf %a : f64
%1 = math.ceil %0 : f64 %1 = math.ceil %0 : f64
%2 = math.floor %1 : f64 %2 = math.floor %1 : f64
%3 = math.sqrt %2 : f64 %3 = math.sqrt %2 : f64
%4 = math.expm1 %3 : f64 %4 = math.expm1 %3 : f64
%5 = math.log1p %4 : f64 %5 = math.log1p %4 : f64
%6 = math.sin %5 : f64 %6 = math.sin %5 : f64
Show All 23 Lines
// CHECK: %[[VAL_14:.*]] = sparse_tensor.load %[[VAL_4]] hasInserts : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> // CHECK: %[[VAL_14:.*]] = sparse_tensor.load %[[VAL_4]] hasInserts : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
// CHECK: return %[[VAL_14]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>> // CHECK: return %[[VAL_14]] : tensor<32xcomplex<f64>, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>
// CHECK: } // CHECK: }
func.func @complex_divbyc(%arg0: tensor<32xcomplex<f64>, #SV>) -> tensor<32xcomplex<f64>, #SV> { func.func @complex_divbyc(%arg0: tensor<32xcomplex<f64>, #SV>) -> tensor<32xcomplex<f64>, #SV> {
%c = complex.constant [0.0, 1.0] : complex<f64> %c = complex.constant [0.0, 1.0] : complex<f64>
%init = bufferization.alloc_tensor() : tensor<32xcomplex<f64>, #SV> %init = bufferization.alloc_tensor() : tensor<32xcomplex<f64>, #SV>
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arg0: tensor<32xcomplex<f64>, #SV>) ins(%arg0: tensor<32xcomplex<f64>, #SV>)
outs(%init: tensor<32xcomplex<f64>, #SV>) { inits(%init: tensor<32xcomplex<f64>, #SV>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%0 = complex.div %a, %c : complex<f64> %0 = complex.div %a, %c : complex<f64>
linalg.yield %0 : complex<f64> linalg.yield %0 : complex<f64>
} -> tensor<32xcomplex<f64>, #SV> } -> tensor<32xcomplex<f64>, #SV>
return %0 : tensor<32xcomplex<f64>, #SV> return %0 : tensor<32xcomplex<f64>, #SV>
} }

mlir/test/Dialect/SparseTensor/sparse_index.mlir

Show First 20 Lines • Show All 49 Lines • ▼ Show 20 Linesfunc.func @dense_index(%arga: tensor<?x?xi64, #DenseMatrix>)
-> tensor<?x?xi64, #DenseMatrix> { -> tensor<?x?xi64, #DenseMatrix> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 0 : index %c1 = arith.constant 0 : index
%0 = tensor.dim %arga, %c0 : tensor<?x?xi64, #DenseMatrix> %0 = tensor.dim %arga, %c0 : tensor<?x?xi64, #DenseMatrix>
%1 = tensor.dim %arga, %c1 : tensor<?x?xi64, #DenseMatrix> %1 = tensor.dim %arga, %c1 : tensor<?x?xi64, #DenseMatrix>
%init = bufferization.alloc_tensor(%0, %1) : tensor<?x?xi64, #DenseMatrix> %init = bufferization.alloc_tensor(%0, %1) : tensor<?x?xi64, #DenseMatrix>
%r = linalg.generic #trait %r = linalg.generic #trait
ins(%arga: tensor<?x?xi64, #DenseMatrix>) ins(%arga: tensor<?x?xi64, #DenseMatrix>)
outs(%init: tensor<?x?xi64, #DenseMatrix>) { inits(%init: tensor<?x?xi64, #DenseMatrix>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.muli %ii, %a : i64 %m1 = arith.muli %ii, %a : i64
%m2 = arith.muli %jj, %m1 : i64 %m2 = arith.muli %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
Show All 38 Linesfunc.func @sparse_index(%arga: tensor<?x?xi64, #SparseMatrix>)
-> tensor<?x?xi64, #SparseMatrix> { -> tensor<?x?xi64, #SparseMatrix> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 0 : index %c1 = arith.constant 0 : index
%0 = tensor.dim %arga, %c0 : tensor<?x?xi64, #SparseMatrix> %0 = tensor.dim %arga, %c0 : tensor<?x?xi64, #SparseMatrix>
%1 = tensor.dim %arga, %c1 : tensor<?x?xi64, #SparseMatrix> %1 = tensor.dim %arga, %c1 : tensor<?x?xi64, #SparseMatrix>
%init = bufferization.alloc_tensor(%0, %1) : tensor<?x?xi64, #SparseMatrix> %init = bufferization.alloc_tensor(%0, %1) : tensor<?x?xi64, #SparseMatrix>
%r = linalg.generic #trait %r = linalg.generic #trait
ins(%arga: tensor<?x?xi64, #SparseMatrix>) ins(%arga: tensor<?x?xi64, #SparseMatrix>)
outs(%init: tensor<?x?xi64, #SparseMatrix>) { inits(%init: tensor<?x?xi64, #SparseMatrix>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.muli %ii, %a : i64 %m1 = arith.muli %ii, %a : i64
%m2 = arith.muli %jj, %m1 : i64 %m2 = arith.muli %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
} -> tensor<?x?xi64, #SparseMatrix> } -> tensor<?x?xi64, #SparseMatrix>
return %r : tensor<?x?xi64, #SparseMatrix> return %r : tensor<?x?xi64, #SparseMatrix>
} }

mlir/test/Dialect/SparseTensor/sparse_int_ops.mlir

Show First 20 Lines • Show All 68 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64> // CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
// CHECK: return %[[VAL_33]] : tensor<32xi64> // CHECK: return %[[VAL_33]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @add(%arga: tensor<32xi64, #SV>, func.func @add(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.addi %a, %b : i64 %0 = arith.addi %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @sub( // CHECK-LABEL: func @sub(
▲ Show 20 Lines • Show All 46 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_36:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xi64> // CHECK: %[[VAL_36:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<32xi64>
// CHECK: return %[[VAL_36]] : tensor<32xi64> // CHECK: return %[[VAL_36]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @sub(%arga: tensor<32xi64, #SV>, func.func @sub(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.subi %a, %b : i64 %0 = arith.subi %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @mul( // CHECK-LABEL: func @mul(
Show All 19 Lines
// CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64> // CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64>
// CHECK: return %[[VAL_17]] : tensor<32xi64> // CHECK: return %[[VAL_17]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @mul(%arga: tensor<32xi64, #SV>, func.func @mul(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.muli %a, %b : i64 %0 = arith.muli %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @divsbyc( // CHECK-LABEL: func @divsbyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
// CHECK: return %[[VAL_15]] : tensor<32xi64> // CHECK: return %[[VAL_15]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @divsbyc(%arga: tensor<32xi64, #SV>, func.func @divsbyc(%arga: tensor<32xi64, #SV>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%c = arith.constant 2 : i64 %c = arith.constant 2 : i64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xi64, #SV>) ins(%arga: tensor<32xi64, #SV>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%0 = arith.divsi %a, %c : i64 %0 = arith.divsi %a, %c : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @divubyc( // CHECK-LABEL: func @divubyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
// CHECK: return %[[VAL_15]] : tensor<32xi64> // CHECK: return %[[VAL_15]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @divubyc(%arga: tensor<32xi64, #SV>, func.func @divubyc(%arga: tensor<32xi64, #SV>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%c = arith.constant 2 : i64 %c = arith.constant 2 : i64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xi64, #SV>) ins(%arga: tensor<32xi64, #SV>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%0 = arith.divui %a, %c : i64 %0 = arith.divui %a, %c : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @and( // CHECK-LABEL: func @and(
Show All 19 Lines
// CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64> // CHECK: %[[VAL_17:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<32xi64>
// CHECK: return %[[VAL_17]] : tensor<32xi64> // CHECK: return %[[VAL_17]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @and(%arga: tensor<32xi64, #SV>, func.func @and(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.andi %a, %b : i64 %0 = arith.andi %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @or( // CHECK-LABEL: func @or(
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64> // CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
// CHECK: return %[[VAL_33]] : tensor<32xi64> // CHECK: return %[[VAL_33]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @or(%arga: tensor<32xi64, #SV>, func.func @or(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.ori %a, %b : i64 %0 = arith.ori %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @xor( // CHECK-LABEL: func @xor(
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64> // CHECK: %[[VAL_33:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<32xi64>
// CHECK: return %[[VAL_33]] : tensor<32xi64> // CHECK: return %[[VAL_33]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @xor(%arga: tensor<32xi64, #SV>, func.func @xor(%arga: tensor<32xi64, #SV>,
%argb: tensor<32xi64>, %argb: tensor<32xi64>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%0 = linalg.generic #trait2 %0 = linalg.generic #trait2
ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>) ins(%arga, %argb: tensor<32xi64, #SV>, tensor<32xi64>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %b: i64, %x: i64): ^bb(%a: i64, %b: i64, %x: i64):
%0 = arith.xori %a, %b : i64 %0 = arith.xori %a, %b : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @ashrbyc( // CHECK-LABEL: func @ashrbyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
// CHECK: return %[[VAL_15]] : tensor<32xi64> // CHECK: return %[[VAL_15]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @ashrbyc(%arga: tensor<32xi64, #SV>, func.func @ashrbyc(%arga: tensor<32xi64, #SV>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%c = arith.constant 2 : i64 %c = arith.constant 2 : i64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xi64, #SV>) ins(%arga: tensor<32xi64, #SV>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%0 = arith.shrsi %a, %c : i64 %0 = arith.shrsi %a, %c : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @lsrbyc( // CHECK-LABEL: func @lsrbyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
// CHECK: return %[[VAL_15]] : tensor<32xi64> // CHECK: return %[[VAL_15]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @lsrbyc(%arga: tensor<32xi64, #SV>, func.func @lsrbyc(%arga: tensor<32xi64, #SV>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%c = arith.constant 2 : i64 %c = arith.constant 2 : i64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xi64, #SV>) ins(%arga: tensor<32xi64, #SV>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%0 = arith.shrui %a, %c : i64 %0 = arith.shrui %a, %c : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }
// CHECK-LABEL: func @lslbyc( // CHECK-LABEL: func @lslbyc(
Show All 17 Lines
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<32xi64>
// CHECK: return %[[VAL_15]] : tensor<32xi64> // CHECK: return %[[VAL_15]] : tensor<32xi64>
// CHECK: } // CHECK: }
func.func @lslbyc(%arga: tensor<32xi64, #SV>, func.func @lslbyc(%arga: tensor<32xi64, #SV>,
%argx: tensor<32xi64>) -> tensor<32xi64> { %argx: tensor<32xi64>) -> tensor<32xi64> {
%c = arith.constant 2 : i64 %c = arith.constant 2 : i64
%0 = linalg.generic #traitc %0 = linalg.generic #traitc
ins(%arga: tensor<32xi64, #SV>) ins(%arga: tensor<32xi64, #SV>)
outs(%argx: tensor<32xi64>) { inits(%argx: tensor<32xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%0 = arith.shli %a, %c : i64 %0 = arith.shli %a, %c : i64
linalg.yield %0 : i64 linalg.yield %0 : i64
} -> tensor<32xi64> } -> tensor<32xi64>
return %0 : tensor<32xi64> return %0 : tensor<32xi64>
} }

mlir/test/Dialect/SparseTensor/sparse_kernels.mlir

Show First 20 Lines • Show All 41 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_28:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<10x30xf32> // CHECK: %[[VAL_28:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<10x30xf32>
// CHECK: return %[[VAL_28]] : tensor<10x30xf32> // CHECK: return %[[VAL_28]] : tensor<10x30xf32>
// CHECK: } // CHECK: }
func.func @matmul1(%a: tensor<10x20xf32, #DCSR>, func.func @matmul1(%a: tensor<10x20xf32, #DCSR>,
%b: tensor<20x30xf32>, %b: tensor<20x30xf32>,
%c: tensor<10x30xf32>) -> tensor<10x30xf32> { %c: tensor<10x30xf32>) -> tensor<10x30xf32> {
%0 = linalg.matmul %0 = linalg.matmul
ins(%a, %b: tensor<10x20xf32, #DCSR>, tensor<20x30xf32>) ins(%a, %b: tensor<10x20xf32, #DCSR>, tensor<20x30xf32>)
outs(%c: tensor<10x30xf32>) -> tensor<10x30xf32> inits(%c: tensor<10x30xf32>) -> tensor<10x30xf32>
return %0 : tensor<10x30xf32> return %0 : tensor<10x30xf32>
} }
// //
// Computes C = A x B with all matrices sparse (SpMSpM) in DCSR. // Computes C = A x B with all matrices sparse (SpMSpM) in DCSR.
// //
// CHECK-LABEL: func.func @matmul2( // CHECK-LABEL: func.func @matmul2(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<4x8xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<4x8xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
▲ Show 20 Lines • Show All 79 Lines • ▼ Show 20 Lines
// CHECK: return %[[VAL_74]] : tensor<4x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> // CHECK: return %[[VAL_74]] : tensor<4x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
// CHECK: } // CHECK: }
func.func @matmul2(%A: tensor<4x8xf64, #DCSR>, func.func @matmul2(%A: tensor<4x8xf64, #DCSR>,
%B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%c4 = arith.constant 4 : index %c4 = arith.constant 4 : index
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>) ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> inits(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
return %D: tensor<4x4xf64, #DCSR> return %D: tensor<4x4xf64, #DCSR>
} }
// CHECK-LABEL: func.func @conv2d( // CHECK-LABEL: func.func @conv2d(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8xi32>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8xi32>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3xi32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x3xi32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<6x6xi32>) -> tensor<6x6xi32> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<6x6xi32>) -> tensor<6x6xi32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 6 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 6 : index
Show All 32 Lines
// CHECK: %[[VAL_31:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<6x6xi32> // CHECK: %[[VAL_31:.*]] = bufferization.to_tensor %[[VAL_12]] : memref<6x6xi32>
// CHECK: return %[[VAL_31]] : tensor<6x6xi32> // CHECK: return %[[VAL_31]] : tensor<6x6xi32>
// CHECK: } // CHECK: }
func.func @conv2d(%input: tensor<8x8xi32>, func.func @conv2d(%input: tensor<8x8xi32>,
%filter: tensor<3x3xi32, #DCSR>, %filter: tensor<3x3xi32, #DCSR>,
%output: tensor<6x6xi32>) -> tensor<6x6xi32> { %output: tensor<6x6xi32>) -> tensor<6x6xi32> {
%0 = linalg.conv_2d %0 = linalg.conv_2d
ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>) ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>)
outs (%output: tensor<6x6xi32>) -> tensor<6x6xi32> inits (%output: tensor<6x6xi32>) -> tensor<6x6xi32>
return %0 : tensor<6x6xi32> return %0 : tensor<6x6xi32>
} }
// CHECK-LABEL: func.func @quantized_matmul( // CHECK-LABEL: func.func @quantized_matmul(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<5x3xi8>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<5x3xi8>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<3x6xi8, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*]]: tensor<3x6xi8, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*]]: tensor<5x6xi64>) -> tensor<5x6xi64> { // CHECK-SAME: %[[VAL_2:.*]]: tensor<5x6xi64>) -> tensor<5x6xi64> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 5 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 5 : index
Show All 34 Lines
// CHECK: } // CHECK: }
func.func @quantized_matmul(%input1: tensor<5x3xi8>, func.func @quantized_matmul(%input1: tensor<5x3xi8>,
%input2: tensor<3x6xi8, #DCSR>, %input2: tensor<3x6xi8, #DCSR>,
%output: tensor<5x6xi64>) -> tensor<5x6xi64> { %output: tensor<5x6xi64>) -> tensor<5x6xi64> {
%c0 = arith.constant 0 : i32 %c0 = arith.constant 0 : i32
%c2 = arith.constant 2 : i32 %c2 = arith.constant 2 : i32
%0 = linalg.quantized_matmul %0 = linalg.quantized_matmul
ins(%input1, %input2, %c2, %c0 : tensor<5x3xi8>, tensor<3x6xi8, #DCSR>, i32, i32) ins(%input1, %input2, %c2, %c0 : tensor<5x3xi8>, tensor<3x6xi8, #DCSR>, i32, i32)
outs(%output : tensor<5x6xi64>) -> tensor<5x6xi64> inits(%output : tensor<5x6xi64>) -> tensor<5x6xi64>
return %0: tensor<5x6xi64> return %0: tensor<5x6xi64>
} }
// CHECK-LABEL: func.func @sparse_dot( // CHECK-LABEL: func.func @sparse_dot(
// CHECK-SAME: %[[VAL_0:.*0]]: tensor<1024xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_0:.*0]]: tensor<1024xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_1:.*1]]: tensor<1024xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>, // CHECK-SAME: %[[VAL_1:.*1]]: tensor<1024xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>,
// CHECK-SAME: %[[VAL_2:.*2]]: tensor<f32>) -> tensor<f32> { // CHECK-SAME: %[[VAL_2:.*2]]: tensor<f32>) -> tensor<f32> {
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
▲ Show 20 Lines • Show All 45 Lines • ▼ Show 20 Lines
// CHECK: %[[VAL_47:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<f32> // CHECK: %[[VAL_47:.*]] = bufferization.to_tensor %[[VAL_11]] : memref<f32>
// CHECK: return %[[VAL_47]] : tensor<f32> // CHECK: return %[[VAL_47]] : tensor<f32>
// CHECK: } // CHECK: }
func.func @sparse_dot(%a: tensor<1024xf32, #SparseVector>, func.func @sparse_dot(%a: tensor<1024xf32, #SparseVector>,
%b: tensor<1024xf32, #SparseVector>, %b: tensor<1024xf32, #SparseVector>,
%x: tensor<f32>) -> tensor<f32> { %x: tensor<f32>) -> tensor<f32> {
%dot = linalg.dot ins(%a, %b: tensor<1024xf32, #SparseVector>, %dot = linalg.dot ins(%a, %b: tensor<1024xf32, #SparseVector>,
tensor<1024xf32, #SparseVector>) tensor<1024xf32, #SparseVector>)
outs(%x: tensor<f32>) -> tensor<f32> inits(%x: tensor<f32>) -> tensor<f32>
return %dot : tensor<f32> return %dot : tensor<f32>
} }

mlir/test/Dialect/SparseTensor/sparse_lower.mlir

Show First 20 Lines • Show All 109 Lines • ▼ Show 20 Lines
// CHECK-LIR: return %[[VAL_2]] : memref<32xf64> // CHECK-LIR: return %[[VAL_2]] : memref<32xf64>
// CHECK-LIR: } // CHECK-LIR: }
func.func @matvec(%arga: tensor<32x64xf64, #CSR>, func.func @matvec(%arga: tensor<32x64xf64, #CSR>,
%argb: tensor<64xf64>, %argb: tensor<64xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>) ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%A: f64, %b: f64, %x: f64): ^bb(%A: f64, %b: f64, %x: f64):
%0 = arith.mulf %A, %b : f64 %0 = arith.mulf %A, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_lower_col.mlir

Show First 20 Lines • Show All 109 Lines • ▼ Show 20 Lines
// CHECK-LIR: return %[[VAL_2]] : memref<32xf64> // CHECK-LIR: return %[[VAL_2]] : memref<32xf64>
// CHECK-LIR: } // CHECK-LIR: }
func.func @matvec(%arga: tensor<32x64xf64, #CSC>, func.func @matvec(%arga: tensor<32x64xf64, #CSC>,
%argb: tensor<64xf64>, %argb: tensor<64xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%arga, %argb : tensor<32x64xf64, #CSC>, tensor<64xf64>) ins(%arga, %argb : tensor<32x64xf64, #CSC>, tensor<64xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%A: f64, %b: f64, %x: f64): ^bb(%A: f64, %b: f64, %x: f64):
%0 = arith.mulf %A, %b : f64 %0 = arith.mulf %A, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_lower_inplace.mlir

Show First 20 Lines • Show All 109 Lines • ▼ Show 20 Lines
// CHECK-LIR: return %[[VAL_2]] : memref<32xf64> // CHECK-LIR: return %[[VAL_2]] : memref<32xf64>
// CHECK-LIR: } // CHECK-LIR: }
func.func @matvec(%arga: tensor<32x64xf64, #CSR>, func.func @matvec(%arga: tensor<32x64xf64, #CSR>,
%argb: tensor<64xf64>, %argb: tensor<64xf64>,
%argx: tensor<32xf64>) -> tensor<32xf64> { %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>) ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%A: f64, %b: f64, %x: f64): ^bb(%A: f64, %b: f64, %x: f64):
%0 = arith.mulf %A, %b : f64 %0 = arith.mulf %A, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_nd.mlir

Show All 34 Lines
// CHECK-DAG: %[[VAL_12:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_12:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_0]] : memref<10x20x30x40x50x60x70x80xf32> // CHECK-DAG: %[[VAL_13:.*]] = bufferization.to_memref %[[VAL_0]] : memref<10x20x30x40x50x60x70x80xf32>
// CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 3 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 3 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 3 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 3 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_16:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 4 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_16:.*]] = sparse_tensor.pointers %[[VAL_1]] {dimension = 4 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_17:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 4 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex> // CHECK-DAG: %[[VAL_17:.*]] = sparse_tensor.indices %[[VAL_1]] {dimension = 4 : index} : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xindex>
// CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xf32> // CHECK-DAG: %[[VAL_18:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<80x70x60x50x40x30x20x10xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "dense", "dense", "compressed", "compressed", "dense", "dense", "dense" ] }>> to memref<?xf32>
// CHECK-DAG: %[[VAL_20:.*]] = bufferization.to_memref %[[VAL_2]] : memref<10x20x30x40x50x60x70x80xf32> // CHECK-DAG: %[[VAL_20:.*]] = bufferization.to_memref %[[VAL_2]] : memref<10x20x30x40x50x60x70x80xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_20]] : memref<10x20x30x40x50x60x70x80xf32> // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_20]] : memref<10x20x30x40x50x60x70x80xf32>
// CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_11]] to %[[VAL_10]] step %[[VAL_12]] { // CHECK: scf.for %[[VAL_21:.*]] = %[[VAL_11]] to %[[VAL_10]] step %[[VAL_12]] {
// CHECK: scf.for %[[VAL_22:.*]] = %[[VAL_11]] to %[[VAL_9]] step %[[VAL_12]] { // CHECK: scf.for %[[VAL_22:.*]] = %[[VAL_11]] to %[[VAL_9]] step %[[VAL_12]] {
// CHECK: %[[VAL_23:.*]] = arith.muli %[[VAL_21]], %[[VAL_9]] : index // CHECK: %[[VAL_23:.*]] = arith.muli %[[VAL_21]], %[[VAL_9]] : index
// CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_22]] : index // CHECK: %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_22]] : index
// CHECK: scf.for %[[VAL_25:.*]] = %[[VAL_11]] to %[[VAL_8]] step %[[VAL_12]] { // CHECK: scf.for %[[VAL_25:.*]] = %[[VAL_11]] to %[[VAL_8]] step %[[VAL_12]] {
// CHECK: %[[VAL_26:.*]] = arith.muli %[[VAL_24]], %[[VAL_8]] : index // CHECK: %[[VAL_26:.*]] = arith.muli %[[VAL_24]], %[[VAL_8]] : index
// CHECK: %[[VAL_27:.*]] = arith.addi %[[VAL_26]], %[[VAL_25]] : index // CHECK: %[[VAL_27:.*]] = arith.addi %[[VAL_26]], %[[VAL_25]] : index
// CHECK: %[[VAL_28:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_27]]] : memref<?xindex> // CHECK: %[[VAL_28:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_27]]] : memref<?xindex>
Show All 32 Lines
// CHECK: } // CHECK: }
func.func @mul(%arga: tensor<10x20x30x40x50x60x70x80xf32>, func.func @mul(%arga: tensor<10x20x30x40x50x60x70x80xf32>,
%argb: tensor<80x70x60x50x40x30x20x10xf32, #SparseTensor>, %argb: tensor<80x70x60x50x40x30x20x10xf32, #SparseTensor>,
%argx: tensor<10x20x30x40x50x60x70x80xf32>) %argx: tensor<10x20x30x40x50x60x70x80xf32>)
-> tensor<10x20x30x40x50x60x70x80xf32> { -> tensor<10x20x30x40x50x60x70x80xf32> {
%0 = linalg.generic #trait_mul %0 = linalg.generic #trait_mul
ins(%arga, %argb: tensor<10x20x30x40x50x60x70x80xf32>, ins(%arga, %argb: tensor<10x20x30x40x50x60x70x80xf32>,
tensor<80x70x60x50x40x30x20x10xf32, #SparseTensor>) tensor<80x70x60x50x40x30x20x10xf32, #SparseTensor>)
outs(%argx: tensor<10x20x30x40x50x60x70x80xf32>) { inits(%argx: tensor<10x20x30x40x50x60x70x80xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<10x20x30x40x50x60x70x80xf32> } -> tensor<10x20x30x40x50x60x70x80xf32>
return %0 : tensor<10x20x30x40x50x60x70x80xf32> return %0 : tensor<10x20x30x40x50x60x70x80xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_out.mlir

Show First 20 Lines • Show All 42 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_18:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: return %[[VAL_18]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: return %[[VAL_18]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: } // CHECK: }
func.func @sparse_simply_dynamic1(%argx: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> { func.func @sparse_simply_dynamic1(%argx: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> {
%c = arith.constant 2.0 : f32 %c = arith.constant 2.0 : f32
%0 = linalg.generic #trait_scale_inpl %0 = linalg.generic #trait_scale_inpl
outs(%argx: tensor<32x16xf32, #DCSR>) { inits(%argx: tensor<32x16xf32, #DCSR>) {
^bb(%x: f32): ^bb(%x: f32):
%1 = arith.mulf %x, %c : f32 %1 = arith.mulf %x, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32, #DCSR> } -> tensor<32x16xf32, #DCSR>
return %0 : tensor<32x16xf32, #DCSR> return %0 : tensor<32x16xf32, #DCSR>
} }
// CHECK-LABEL: func.func @sparse_simply_dynamic2( // CHECK-LABEL: func.func @sparse_simply_dynamic2(
Show All 16 Lines
// CHECK: memref.store %[[VAL_15]], %[[VAL_5]]{{\[}}%[[VAL_12]]] : memref<?xf32> // CHECK: memref.store %[[VAL_15]], %[[VAL_5]]{{\[}}%[[VAL_12]]] : memref<?xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_16:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: %[[VAL_16:.*]] = sparse_tensor.load %[[VAL_0]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: return %[[VAL_16]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>> // CHECK: return %[[VAL_16]] : tensor<32x16xf32, #sparse_tensor.encoding<{{.*}}>>
// CHECK: } // CHECK: }
func.func @sparse_simply_dynamic2(%argx: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> { func.func @sparse_simply_dynamic2(%argx: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> {
%0 = linalg.generic #trait_scale_inpl %0 = linalg.generic #trait_scale_inpl
outs(%argx: tensor<32x16xf32, #DCSR>) { inits(%argx: tensor<32x16xf32, #DCSR>) {
^bb(%x: f32): ^bb(%x: f32):
%1 = arith.addf %x, %x : f32 %1 = arith.addf %x, %x : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32, #DCSR> } -> tensor<32x16xf32, #DCSR>
return %0 : tensor<32x16xf32, #DCSR> return %0 : tensor<32x16xf32, #DCSR>
} }
#trait_scale = { #trait_scale = {
Show All 29 Lines
// CHECK: %[[VAL_17:.*]] = sparse_tensor.load %[[VAL_5]] hasInserts : tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> // CHECK: %[[VAL_17:.*]] = sparse_tensor.load %[[VAL_5]] hasInserts : tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
// CHECK: return %[[VAL_17]] : tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> // CHECK: return %[[VAL_17]] : tensor<10x20xf32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
// CHECK: } // CHECK: }
func.func @sparse_truly_dynamic(%arga: tensor<10x20xf32, #CSR>) -> tensor<10x20xf32, #DCSR> { func.func @sparse_truly_dynamic(%arga: tensor<10x20xf32, #CSR>) -> tensor<10x20xf32, #DCSR> {
%s = arith.constant 2.0 : f32 %s = arith.constant 2.0 : f32
%xm = bufferization.alloc_tensor() : tensor<10x20xf32, #DCSR> %xm = bufferization.alloc_tensor() : tensor<10x20xf32, #DCSR>
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
ins(%arga: tensor<10x20xf32, #CSR>) ins(%arga: tensor<10x20xf32, #CSR>)
outs(%xm: tensor<10x20xf32, #DCSR>) { inits(%xm: tensor<10x20xf32, #DCSR>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%1 = arith.mulf %a, %s : f32 %1 = arith.mulf %a, %s : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<10x20xf32, #DCSR> } -> tensor<10x20xf32, #DCSR>
return %0 : tensor<10x20xf32, #DCSR> return %0 : tensor<10x20xf32, #DCSR>
} }
#trait_sumred = { #trait_sumred = {
▲ Show 20 Lines • Show All 135 Lines • ▼ Show 20 Linesfunc.func @sumred(%arga: tensor<?x?x?xi32, #SparseTensor>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?x?xi32, #SparseTensor> %d0 = tensor.dim %arga, %c0 : tensor<?x?x?xi32, #SparseTensor>
%d1 = tensor.dim %arga, %c1 : tensor<?x?x?xi32, #SparseTensor> %d1 = tensor.dim %arga, %c1 : tensor<?x?x?xi32, #SparseTensor>
%xinit = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xi32, #DCSR> %xinit = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xi32, #DCSR>
%0 = linalg.generic #trait_sumred %0 = linalg.generic #trait_sumred
ins(%arga, %argb: tensor<?x?x?xi32, #SparseTensor>, ins(%arga, %argb: tensor<?x?x?xi32, #SparseTensor>,
tensor<?x?x?xi32, #SparseTensor>) tensor<?x?x?xi32, #SparseTensor>)
outs(%xinit: tensor<?x?xi32, #DCSR>) { inits(%xinit: tensor<?x?xi32, #DCSR>) {
^bb(%a: i32, %b: i32, %x: i32): ^bb(%a: i32, %b: i32, %x: i32):
%0 = arith.muli %a, %b : i32 %0 = arith.muli %a, %b : i32
%1 = arith.addi %x, %0 : i32 %1 = arith.addi %x, %0 : i32
linalg.yield %1 : i32 linalg.yield %1 : i32
} -> tensor<?x?xi32, #DCSR> } -> tensor<?x?xi32, #DCSR>
return %0 : tensor<?x?xi32, #DCSR> return %0 : tensor<?x?xi32, #DCSR>
} }
▲ Show 20 Lines • Show All 97 Lines • ▼ Show 20 Linesfunc.func @matmat(%arga: tensor<?x?xf32, #DCSR>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf32, #DCSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf32, #DCSR>
%d1 = tensor.dim %argb, %c1 : tensor<?x?xf32, #DCSR> %d1 = tensor.dim %argb, %c1 : tensor<?x?xf32, #DCSR>
%cinit = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf32, #DCSR> %cinit = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf32, #DCSR>
%0 = linalg.generic #trait_matmat %0 = linalg.generic #trait_matmat
ins(%arga, %argb: tensor<?x?xf32, #DCSR>, ins(%arga, %argb: tensor<?x?xf32, #DCSR>,
tensor<?x?xf32, #DCSR>) tensor<?x?xf32, #DCSR>)
outs(%cinit: tensor<?x?xf32, #DCSR>) { inits(%cinit: tensor<?x?xf32, #DCSR>) {
^bb(%a: f32, %b: f32, %c: f32): ^bb(%a: f32, %b: f32, %c: f32):
%1 = arith.mulf %a, %b : f32 %1 = arith.mulf %a, %b : f32
%2 = arith.addf %c, %1 : f32 %2 = arith.addf %c, %1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32, #DCSR> } -> tensor<?x?xf32, #DCSR>
return %0 : tensor<?x?xf32, #DCSR> return %0 : tensor<?x?xf32, #DCSR>
} }

mlir/test/Dialect/SparseTensor/sparse_outbuf.mlir

Show All 14 Lines
// CHECK-SAME: %[[VAL_1:.*]]: tensor<10xf32>) -> tensor<10xf32> { // CHECK-SAME: %[[VAL_1:.*]]: tensor<10xf32>) -> tensor<10xf32> {
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex> // CHECK: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex>
// CHECK: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex> // CHECK: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex>
// CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xi32, #{{.*}}> to memref<?xi32> // CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xi32, #{{.*}}> to memref<?xi32>
// CHECK: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<10xf32> // CHECK: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_1]] : memref<10xf32>
// CHECK: linalg.fill ins(%[[VAL_3]] : f32) outs(%[[VAL_8]] : memref<10xf32>) // CHECK: linalg.fill ins(%[[VAL_3]] : f32) inits(%[[VAL_8]] : memref<10xf32>)
// CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_2]]] : memref<?xindex> // CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_2]]] : memref<?xindex>
// CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex> // CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] {
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi32> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi32>
// CHECK: %[[VAL_14:.*]] = arith.sitofp %[[VAL_13]] : i32 to f32 // CHECK: %[[VAL_14:.*]] = arith.sitofp %[[VAL_13]] : i32 to f32
// CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<10xf32> // CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<10xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<10xf32> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<10xf32>
// CHECK: return %[[VAL_15]] : tensor<10xf32> // CHECK: return %[[VAL_15]] : tensor<10xf32>
// CHECK: } // CHECK: }
func.func @allout_inplace(%arga: tensor<10xi32, #SV>, func.func @allout_inplace(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xf32>) -> tensor<10xf32> { %argb: tensor<10xf32>) -> tensor<10xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: i32, %x : f32): ^bb(%a: i32, %x : f32):
%cst = arith.sitofp %a : i32 to f32 %cst = arith.sitofp %a : i32 to f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
// CHECK-LABEL: func.func @allout_materialize( // CHECK-LABEL: func.func @allout_materialize(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<10xi32, #{{.*}}>) -> tensor<10xf32> { // CHECK-SAME: %[[VAL_0:.*]]: tensor<10xi32, #{{.*}}>) -> tensor<10xf32> {
// CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_1:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 1 : index
// CHECK: %[[VAL_4:.*]] = bufferization.alloc_tensor() : tensor<10xf32> // CHECK: %[[VAL_4:.*]] = bufferization.alloc_tensor() : tensor<10xf32>
// CHECK: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex> // CHECK: %[[VAL_5:.*]] = sparse_tensor.pointers %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex>
// CHECK: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex> // CHECK: %[[VAL_6:.*]] = sparse_tensor.indices %[[VAL_0]] {dimension = 0 : index} : tensor<10xi32, #{{.*}}> to memref<?xindex>
// CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xi32, #{{.*}}> to memref<?xi32> // CHECK: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10xi32, #{{.*}}> to memref<?xi32>
// CHECK: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_4]] : memref<10xf32> // CHECK: %[[VAL_8:.*]] = bufferization.to_memref %[[VAL_4]] : memref<10xf32>
// CHECK: linalg.fill ins(%[[VAL_2]] : f32) outs(%[[VAL_8]] : memref<10xf32>) // CHECK: linalg.fill ins(%[[VAL_2]] : f32) inits(%[[VAL_8]] : memref<10xf32>)
// CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_1]]] : memref<?xindex> // CHECK: %[[VAL_9:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_1]]] : memref<?xindex>
// CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex> // CHECK: %[[VAL_10:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_3]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_3]] {
// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex> // CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
// CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi32> // CHECK: %[[VAL_13:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xi32>
// CHECK: %[[VAL_14:.*]] = arith.sitofp %[[VAL_13]] : i32 to f32 // CHECK: %[[VAL_14:.*]] = arith.sitofp %[[VAL_13]] : i32 to f32
// CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<10xf32> // CHECK: memref.store %[[VAL_14]], %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<10xf32>
// CHECK: } // CHECK: }
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<10xf32> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<10xf32>
// CHECK: return %[[VAL_15]] : tensor<10xf32> // CHECK: return %[[VAL_15]] : tensor<10xf32>
// CHECK: } // CHECK: }
func.func @allout_materialize(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> { func.func @allout_materialize(%arga: tensor<10xi32, #SV>) -> tensor<10xf32> {
%m = bufferization.alloc_tensor() : tensor<10xf32> %m = bufferization.alloc_tensor() : tensor<10xf32>
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%m: tensor<10xf32>) { inits(%m: tensor<10xf32>) {
^bb(%a: i32, %x : f32): ^bb(%a: i32, %x : f32):
%cst = arith.sitofp %a : i32 to f32 %cst = arith.sitofp %a : i32 to f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
// CHECK-LABEL: func.func @update_inplace( // CHECK-LABEL: func.func @update_inplace(
Show All 16 Lines
// CHECK: } // CHECK: }
// CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<10xf32> // CHECK: %[[VAL_15:.*]] = bufferization.to_tensor %[[VAL_7]] : memref<10xf32>
// CHECK: return %[[VAL_15]] : tensor<10xf32> // CHECK: return %[[VAL_15]] : tensor<10xf32>
// CHECK: } // CHECK: }
func.func @update_inplace(%arga: tensor<10xf32, #SV>, func.func @update_inplace(%arga: tensor<10xf32, #SV>,
%argb: tensor<10xf32>) -> tensor<10xf32> { %argb: tensor<10xf32>) -> tensor<10xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<10xf32, #SV>) ins(%arga: tensor<10xf32, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: f32, %x : f32): ^bb(%a: f32, %x : f32):
%up = arith.addf %a, %x : f32 %up = arith.addf %a, %x : f32
linalg.yield %up : f32 linalg.yield %up : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

Show First 20 Lines • Show All 56 Lines • ▼ Show 20 Lines
// CHECK-PAR4: scf.parallel // CHECK-PAR4: scf.parallel
// CHECK-PAR4: return // CHECK-PAR4: return
// //
func.func @scale_dd(%scale: f32, func.func @scale_dd(%scale: f32,
%arga: tensor<?x?xf32, #DenseMatrix>, %arga: tensor<?x?xf32, #DenseMatrix>,
%argx: tensor<?x?xf32>) -> tensor<?x?xf32> { %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.generic #trait_dd %0 = linalg.generic #trait_dd
ins(%arga: tensor<?x?xf32, #DenseMatrix>) ins(%arga: tensor<?x?xf32, #DenseMatrix>)
outs(%argx: tensor<?x?xf32>) { inits(%argx: tensor<?x?xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %a, %scale : f32 %0 = arith.mulf %a, %scale : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
#trait_ss = { #trait_ss = {
Show All 31 Lines
// CHECK-PAR4: scf.parallel // CHECK-PAR4: scf.parallel
// CHECK-PAR4: return // CHECK-PAR4: return
// //
func.func @scale_ss(%scale: f32, func.func @scale_ss(%scale: f32,
%arga: tensor<?x?xf32, #SparseMatrix>, %arga: tensor<?x?xf32, #SparseMatrix>,
%argx: tensor<?x?xf32>) -> tensor<?x?xf32> { %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.generic #trait_ss %0 = linalg.generic #trait_ss
ins(%arga: tensor<?x?xf32, #SparseMatrix>) ins(%arga: tensor<?x?xf32, #SparseMatrix>)
outs(%argx: tensor<?x?xf32>) { inits(%argx: tensor<?x?xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %a, %scale : f32 %0 = arith.mulf %a, %scale : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
#trait_matvec = { #trait_matvec = {
Show All 32 Lines
// CHECK-PAR4: scf.for // CHECK-PAR4: scf.for
// CHECK-PAR4: return // CHECK-PAR4: return
// //
func.func @matvec(%arga: tensor<16x32xf32, #CSR>, func.func @matvec(%arga: tensor<16x32xf32, #CSR>,
%argb: tensor<32xf32>, %argb: tensor<32xf32>,
%argx: tensor<16xf32>) -> tensor<16xf32> { %argx: tensor<16xf32>) -> tensor<16xf32> {
%0 = linalg.generic #trait_matvec %0 = linalg.generic #trait_matvec
ins(%arga, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>) ins(%arga, %argb : tensor<16x32xf32, #CSR>, tensor<32xf32>)
outs(%argx: tensor<16xf32>) { inits(%argx: tensor<16xf32>) {
^bb(%A: f32, %b: f32, %x: f32): ^bb(%A: f32, %b: f32, %x: f32):
%0 = arith.mulf %A, %b : f32 %0 = arith.mulf %A, %b : f32
%1 = arith.addf %0, %x : f32 %1 = arith.addf %0, %x : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<16xf32> } -> tensor<16xf32>
return %0 : tensor<16xf32> return %0 : tensor<16xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_perm.mlir

Show All 19 Lines
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 20 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 20 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 30 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 30 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 10 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 10 : index
// CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_5:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_6:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10x20x30xf32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<10x20x30xf32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<20x30x10xf32> // CHECK-DAG: %[[VAL_9:.*]] = bufferization.to_memref %[[VAL_1]] : memref<20x30x10xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_9]] : memref<20x30x10xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_9]] : memref<20x30x10xf32>)
// CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_10:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
// CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_10]], %[[VAL_4]] : index // CHECK: %[[VAL_12:.*]] = arith.muli %[[VAL_10]], %[[VAL_4]] : index
// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_11]] : index // CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_12]], %[[VAL_11]] : index
// CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_5]] to %[[VAL_2]] step %[[VAL_6]] { // CHECK: scf.for %[[VAL_14:.*]] = %[[VAL_5]] to %[[VAL_2]] step %[[VAL_6]] {
// CHECK: %[[VAL_15:.*]] = arith.muli %[[VAL_13]], %[[VAL_2]] : index // CHECK: %[[VAL_15:.*]] = arith.muli %[[VAL_13]], %[[VAL_2]] : index
// CHECK: %[[VAL_16:.*]] = arith.addi %[[VAL_15]], %[[VAL_14]] : index // CHECK: %[[VAL_16:.*]] = arith.addi %[[VAL_15]], %[[VAL_14]] : index
// CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xf32> // CHECK: %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xf32>
// CHECK: memref.store %[[VAL_17]], %[[VAL_9]]{{\[}}%[[VAL_14]], %[[VAL_10]], %[[VAL_11]]] : memref<20x30x10xf32> // CHECK: memref.store %[[VAL_17]], %[[VAL_9]]{{\[}}%[[VAL_14]], %[[VAL_10]], %[[VAL_11]]] : memref<20x30x10xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<20x30x10xf32> // CHECK: %[[VAL_18:.*]] = bufferization.to_tensor %[[VAL_9]] : memref<20x30x10xf32>
// CHECK: return %[[VAL_18]] : tensor<20x30x10xf32> // CHECK: return %[[VAL_18]] : tensor<20x30x10xf32>
// CHECK: } // CHECK: }
func.func @sparse_static_dims(%arga: tensor<10x20x30xf32, #X>, func.func @sparse_static_dims(%arga: tensor<10x20x30xf32, #X>,
%argx: tensor<20x30x10xf32>) -> tensor<20x30x10xf32> { %argx: tensor<20x30x10xf32>) -> tensor<20x30x10xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<10x20x30xf32, #X>) ins(%arga: tensor<10x20x30xf32, #X>)
outs(%argx: tensor<20x30x10xf32>) { inits(%argx: tensor<20x30x10xf32>) {
^bb(%a : f32, %x: f32): ^bb(%a : f32, %x: f32):
linalg.yield %a : f32 linalg.yield %a : f32
} -> tensor<20x30x10xf32> } -> tensor<20x30x10xf32>
return %0 : tensor<20x30x10xf32> return %0 : tensor<20x30x10xf32>
} }
// CHECK-LABEL: func @sparse_dynamic_dims( // CHECK-LABEL: func @sparse_dynamic_dims(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>,
// CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { // CHECK-SAME: %[[VAL_1:.*]]: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32 // CHECK-DAG: %[[ZERO:.*]] = arith.constant 0.000000e+00 : f32
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 2 : index // CHECK-DAG: %[[VAL_2:.*]] = arith.constant 2 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index // CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index // CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_6:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_3]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> // CHECK-DAG: %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_4]] : tensor<?x?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
// CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?x?xf32> // CHECK-DAG: %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<?x?x?xf32>
// CHECK: linalg.fill ins(%[[ZERO]] : f32) outs(%[[VAL_10]] : memref<?x?x?xf32>) // CHECK: linalg.fill ins(%[[ZERO]] : f32) inits(%[[VAL_10]] : memref<?x?x?xf32>)
// CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_3]] to %[[VAL_6]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_11:.*]] = %[[VAL_3]] to %[[VAL_6]] step %[[VAL_4]] {
// CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_7]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_12:.*]] = %[[VAL_3]] to %[[VAL_7]] step %[[VAL_4]] {
// CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_7]], %[[VAL_11]] : index // CHECK: %[[VAL_13:.*]] = arith.muli %[[VAL_7]], %[[VAL_11]] : index
// CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index // CHECK: %[[VAL_14:.*]] = arith.addi %[[VAL_13]], %[[VAL_12]] : index
// CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] { // CHECK: scf.for %[[VAL_15:.*]] = %[[VAL_3]] to %[[VAL_8]] step %[[VAL_4]] {
// CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_8]], %[[VAL_14]] : index // CHECK: %[[VAL_16:.*]] = arith.muli %[[VAL_8]], %[[VAL_14]] : index
// CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_16]], %[[VAL_15]] : index // CHECK: %[[VAL_17:.*]] = arith.addi %[[VAL_16]], %[[VAL_15]] : index
// CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_17]]] : memref<?xf32> // CHECK: %[[VAL_18:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_17]]] : memref<?xf32>
// CHECK: memref.store %[[VAL_18]], %[[VAL_10]]{{\[}}%[[VAL_15]], %[[VAL_11]], %[[VAL_12]]] : memref<?x?x?xf32> // CHECK: memref.store %[[VAL_18]], %[[VAL_10]]{{\[}}%[[VAL_15]], %[[VAL_11]], %[[VAL_12]]] : memref<?x?x?xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: %[[VAL_19:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<?x?x?xf32> // CHECK: %[[VAL_19:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<?x?x?xf32>
// CHECK: return %[[VAL_19]] : tensor<?x?x?xf32> // CHECK: return %[[VAL_19]] : tensor<?x?x?xf32>
// CHECK: } // CHECK: }
func.func @sparse_dynamic_dims(%arga: tensor<?x?x?xf32, #X>, func.func @sparse_dynamic_dims(%arga: tensor<?x?x?xf32, #X>,
%argx: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> { %argx: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<?x?x?xf32, #X>) ins(%arga: tensor<?x?x?xf32, #X>)
outs(%argx: tensor<?x?x?xf32>) { inits(%argx: tensor<?x?x?xf32>) {
^bb(%a : f32, %x: f32): ^bb(%a : f32, %x: f32):
linalg.yield %a : f32 linalg.yield %a : f32
} -> tensor<?x?x?xf32> } -> tensor<?x?x?xf32>
return %0 : tensor<?x?x?xf32> return %0 : tensor<?x?x?xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_perm_lower.mlir

Show First 20 Lines • Show All 77 Lines • ▼ Show 20 Lines
// CHECK-MIR: memref.store %[[VAL_12]], %[[VAL_10]][] : memref<f32> // CHECK-MIR: memref.store %[[VAL_12]], %[[VAL_10]][] : memref<f32>
// CHECK-MIR: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<f32> // CHECK-MIR: %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<f32>
// CHECK-MIR: return %[[VAL_30]] : tensor<f32> // CHECK-MIR: return %[[VAL_30]] : tensor<f32>
// CHECK-MIR: } // CHECK-MIR: }
func.func @sparse_dynamic_dims(%arga: tensor<?x?x?xf32, #X>, func.func @sparse_dynamic_dims(%arga: tensor<?x?x?xf32, #X>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga: tensor<?x?x?xf32, #X>) ins(%arga: tensor<?x?x?xf32, #X>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a : f32, %x: f32): ^bb(%a : f32, %x: f32):
%0 = arith.addf %x, %a : f32 %0 = arith.addf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }

mlir/test/Dialect/SparseTensor/sparse_scalars.mlir

Show First 20 Lines • Show All 62 Lines • ▼ Show 20 Linesfunc.func @mul(%arga: tensor<32x16xf32, #SparseMatrix>,
%argp: tensor<f32>, %argp: tensor<f32>,
%argq: f32, %argq: f32,
%argr: f32, %argr: f32,
%argx: tensor<32x16xf32>) -> tensor<32x16xf32> { %argx: tensor<32x16xf32>) -> tensor<32x16xf32> {
%s = arith.addf %argq, %argr : f32 %s = arith.addf %argq, %argr : f32
%c = arith.constant 2.2 : f32 %c = arith.constant 2.2 : f32
%0 = linalg.generic #trait %0 = linalg.generic #trait
ins(%arga, %argp, %argq: tensor<32x16xf32, #SparseMatrix>, tensor<f32>, f32) ins(%arga, %argp, %argq: tensor<32x16xf32, #SparseMatrix>, tensor<f32>, f32)
outs(%argx: tensor<32x16xf32>) { inits(%argx: tensor<32x16xf32>) {
^bb(%a: f32, %p: f32, %q: f32, %x: f32): ^bb(%a: f32, %p: f32, %q: f32, %x: f32):
%0 = arith.mulf %a, %p : f32 // scalar tensor argument %0 = arith.mulf %a, %p : f32 // scalar tensor argument
%1 = arith.mulf %0, %q : f32 // scalar argument %1 = arith.mulf %0, %q : f32 // scalar argument
%2 = arith.mulf %1, %argr : f32 // scalar argument from outside block %2 = arith.mulf %1, %argr : f32 // scalar argument from outside block
%3 = arith.mulf %2, %s : f32 // scalar value from outside block %3 = arith.mulf %2, %s : f32 // scalar value from outside block
%4 = arith.mulf %3, %c : f32 // direct constant from outside block %4 = arith.mulf %3, %c : f32 // direct constant from outside block
%5 = arith.addf %4, %x : f32 %5 = arith.addf %4, %x : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<32x16xf32> } -> tensor<32x16xf32>
return %0 : tensor<32x16xf32> return %0 : tensor<32x16xf32>
} }

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

Show All 26 Lines
// 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 = tensor.empty() : 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>) { inits(%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 = tensor.empty() : 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>) { inits(%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>
} }
// CHECK-LABEL: func.func @sampled_dd_unfused( // CHECK-LABEL: func.func @sampled_dd_unfused(
// CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>, // CHECK-SAME: %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
Show All 40 Lines
// CHECK: } // CHECK: }
func.func @sampled_dd_unfused(%args: tensor<8x8xf64, #SM>, func.func @sampled_dd_unfused(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64> {
// Perform dense-dense matrix matrix multiplication. // Perform dense-dense matrix matrix multiplication.
%1 = arith.constant dense<0.0> : tensor<8x8xf64> %1 = arith.constant dense<0.0> : tensor<8x8xf64>
%2 = linalg.generic #trait_matmul %2 = linalg.generic #trait_matmul
ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>) ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%a: f64, %b: f64, %x: f64): ^bb0(%a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.addf %x, %p : f64 %q = arith.addf %x, %p : f64
linalg.yield %q : f64 linalg.yield %q : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
// Sample the result with elements-wise multiplication with sparse matrix. // Sample the result with elements-wise multiplication with sparse matrix.
%3 = linalg.generic #trait_scale %3 = linalg.generic #trait_scale
ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>) ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%t: f64, %s: f64, %x: f64): ^bb0(%t: f64, %s: f64, %x: f64):
%r = arith.mulf %t, %s : f64 %r = arith.mulf %t, %s : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
return %3 : tensor<8x8xf64> return %3 : tensor<8x8xf64>
} }
// CHECK-LABEL: func.func @sparse_sampled_dd_unfused( // CHECK-LABEL: func.func @sparse_sampled_dd_unfused(
▲ Show 20 Lines • Show All 55 Lines • ▼ Show 20 Lines
// CHECK: } // CHECK: }
func.func @sparse_sampled_dd_unfused(%args: tensor<8x8xf64, #SM>, func.func @sparse_sampled_dd_unfused(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> {
// Perform dense-dense matrix matrix multiplication. // Perform dense-dense matrix matrix multiplication.
%1 = arith.constant dense<0.0> : tensor<8x8xf64> %1 = arith.constant dense<0.0> : tensor<8x8xf64>
%2 = linalg.generic #trait_matmul %2 = linalg.generic #trait_matmul
ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>) ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%a: f64, %b: f64, %x: f64): ^bb0(%a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.addf %x, %p : f64 %q = arith.addf %x, %p : f64
linalg.yield %q : f64 linalg.yield %q : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
// Sample the result with elements-wise multiplication with sparse matrix. // Sample the result with elements-wise multiplication with sparse matrix.
%3 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM> %3 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM>
%4 = linalg.generic #trait_scale %4 = linalg.generic #trait_scale
ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>) ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>)
outs(%3 : tensor<8x8xf64, #SM>) { inits(%3 : tensor<8x8xf64, #SM>) {
^bb0(%t: f64, %s: f64, %x: f64): ^bb0(%t: f64, %s: f64, %x: f64):
%r = arith.mulf %t, %s : f64 %r = arith.mulf %t, %s : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64, #SM> } -> tensor<8x8xf64, #SM>
return %4 : tensor<8x8xf64, #SM> return %4 : tensor<8x8xf64, #SM>
} }

mlir/test/Dialect/SparseTensor/sparse_storage.mlir

Show All 34 Lines
// CHECK: %[[VAL0:.*]] = memref.load %{{.*}}[%[[I]]] : memref<?xf64> // CHECK: %[[VAL0:.*]] = memref.load %{{.*}}[%[[I]]] : memref<?xf64>
// CHECK: %[[VAL1:.*]] = memref.load %{{.*}}[%[[INDC]]] : memref<32xf64> // CHECK: %[[VAL1:.*]] = memref.load %{{.*}}[%[[INDC]]] : memref<32xf64>
// CHECK: %[[MUL:.*]] = arith.mulf %[[VAL0]], %[[VAL1]] : f64 // CHECK: %[[MUL:.*]] = arith.mulf %[[VAL0]], %[[VAL1]] : f64
// CHECK: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> // CHECK: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64>
// CHECK: } // CHECK: }
func.func @mul64(%arga: tensor<32xf64, #SparseVector64>, %argb: tensor<32xf64>, %argx: tensor<32xf64>) -> tensor<32xf64> { func.func @mul64(%arga: tensor<32xf64, #SparseVector64>, %argb: tensor<32xf64>, %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_mul %0 = linalg.generic #trait_mul
ins(%arga, %argb: tensor<32xf64, #SparseVector64>, tensor<32xf64>) ins(%arga, %argb: tensor<32xf64, #SparseVector64>, tensor<32xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.mulf %a, %b : f64 %0 = arith.mulf %a, %b : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }
// CHECK-LABEL: func @mul32( // CHECK-LABEL: func @mul32(
Show All 12 Lines
// CHECK: %[[VAL0:.*]] = memref.load %{{.*}}[%[[I]]] : memref<?xf64> // CHECK: %[[VAL0:.*]] = memref.load %{{.*}}[%[[I]]] : memref<?xf64>
// CHECK: %[[VAL1:.*]] = memref.load %{{.*}}[%[[INDC]]] : memref<32xf64> // CHECK: %[[VAL1:.*]] = memref.load %{{.*}}[%[[INDC]]] : memref<32xf64>
// CHECK: %[[MUL:.*]] = arith.mulf %[[VAL0]], %[[VAL1]] : f64 // CHECK: %[[MUL:.*]] = arith.mulf %[[VAL0]], %[[VAL1]] : f64
// CHECK: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64> // CHECK: store %[[MUL]], %{{.*}}[%[[INDC]]] : memref<32xf64>
// CHECK: } // CHECK: }
func.func @mul32(%arga: tensor<32xf64, #SparseVector32>, %argb: tensor<32xf64>, %argx: tensor<32xf64>) -> tensor<32xf64> { func.func @mul32(%arga: tensor<32xf64, #SparseVector32>, %argb: tensor<32xf64>, %argx: tensor<32xf64>) -> tensor<32xf64> {
%0 = linalg.generic #trait_mul %0 = linalg.generic #trait_mul
ins(%arga, %argb: tensor<32xf64, #SparseVector32>, tensor<32xf64>) ins(%arga, %argb: tensor<32xf64, #SparseVector32>, tensor<32xf64>)
outs(%argx: tensor<32xf64>) { inits(%argx: tensor<32xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.mulf %a, %b : f64 %0 = arith.mulf %a, %b : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<32xf64> } -> tensor<32xf64>
return %0 : tensor<32xf64> return %0 : tensor<32xf64>
} }

mlir/test/Dialect/SparseTensor/sparse_transpose.mlir

Show First 20 Lines • Show All 42 Lines • ▼ Show 20 Lines
// CHECK: bufferization.dealloc_tensor %[[VAL_4]] : tensor<3x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>> // CHECK: bufferization.dealloc_tensor %[[VAL_4]] : tensor<3x4xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)> }>>
// CHECK: return %[[VAL_20]] : tensor<4x3xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>> // CHECK: return %[[VAL_20]] : tensor<4x3xf64, #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>>
// CHECK: } // CHECK: }
func.func @sparse_transpose_auto(%arga: tensor<3x4xf64, #DCSR>) func.func @sparse_transpose_auto(%arga: tensor<3x4xf64, #DCSR>)
-> tensor<4x3xf64, #DCSR> { -> tensor<4x3xf64, #DCSR> {
%i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR> %i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR>
%0 = linalg.generic #transpose_trait %0 = linalg.generic #transpose_trait
ins(%arga: tensor<3x4xf64, #DCSR>) ins(%arga: tensor<3x4xf64, #DCSR>)
outs(%i: tensor<4x3xf64, #DCSR>) { inits(%i: tensor<4x3xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
linalg.yield %a : f64 linalg.yield %a : f64
} -> tensor<4x3xf64, #DCSR> } -> tensor<4x3xf64, #DCSR>
return %0 : tensor<4x3xf64, #DCSR> return %0 : tensor<4x3xf64, #DCSR>
} }

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

Show All 22 Lines
// CHECK: store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32> // CHECK: store %[[m]], %{{.*}}[%[[i]]] : memref<1024xf32>
// CHECK: } // CHECK: }
// CHECK: return // CHECK: return
// //
func.func @scale_d(%arga: tensor<1024xf32, #DenseVector>, %b: f32, %argx: tensor<1024xf32>) -> tensor<1024xf32> { func.func @scale_d(%arga: tensor<1024xf32, #DenseVector>, %b: f32, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
%0 = linalg.generic #trait_scale_d %0 = linalg.generic #trait_scale_d
ins(%arga: tensor<1024xf32, #DenseVector>) ins(%arga: tensor<1024xf32, #DenseVector>)
outs(%argx: tensor<1024xf32>) { inits(%argx: tensor<1024xf32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<1024xf32> } -> tensor<1024xf32>
return %0 : tensor<1024xf32> return %0 : tensor<1024xf32>
} }
// ----- // -----
Show All 33 Lines
// CHECK: %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : f32 // CHECK: %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : f32
// CHECK: store %[[m]], %{{.*}}[%[[ci]]] : memref<1024xf32> // CHECK: store %[[m]], %{{.*}}[%[[ci]]] : memref<1024xf32>
// CHECK: } // CHECK: }
// CHECK: return // CHECK: return
// //
func.func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> { func.func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
%0 = linalg.generic #trait_mul_s %0 = linalg.generic #trait_mul_s
ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>) ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>)
outs(%argx: tensor<1024xf32>) { inits(%argx: tensor<1024xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<1024xf32> } -> tensor<1024xf32>
return %0 : tensor<1024xf32> return %0 : tensor<1024xf32>
} }
// ----- // -----
Show All 22 Lines
// CHECK: %[[a:.*]] = arith.addf %[[red_in]], %[[m]] : f32 // CHECK: %[[a:.*]] = arith.addf %[[red_in]], %[[m]] : f32
// CHECK: scf.yield %[[a]] : f32 // CHECK: scf.yield %[[a]] : f32
// CHECK: } // CHECK: }
// CHECK: return // CHECK: return
// //
func.func @reduction_d(%arga: tensor<1024xf32, #DenseVector>, %argb: tensor<1024xf32>, %argx: tensor<f32>) -> tensor<f32> { func.func @reduction_d(%arga: tensor<1024xf32, #DenseVector>, %argb: tensor<1024xf32>, %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_reduction_d %0 = linalg.generic #trait_reduction_d
ins(%arga, %argb: tensor<1024xf32, #DenseVector>, tensor<1024xf32>) ins(%arga, %argb: tensor<1024xf32, #DenseVector>, tensor<1024xf32>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.addf %x, %0 : f32 %1 = arith.addf %x, %0 : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
Show All 38 Lines
// CHECK: store %[[m]], %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32> // CHECK: store %[[m]], %{{.*}}[%[[i]], %[[cj]]] : memref<512x1024xf32>
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: return // CHECK: return
// //
func.func @mul_ds(%arga: tensor<512x1024xf32, #SparseMatrix>, %argb: tensor<512x1024xf32>, %argx: tensor<512x1024xf32>) -> tensor<512x1024xf32> { func.func @mul_ds(%arga: tensor<512x1024xf32, #SparseMatrix>, %argb: tensor<512x1024xf32>, %argx: tensor<512x1024xf32>) -> tensor<512x1024xf32> {
%0 = linalg.generic #trait_mul_ds %0 = linalg.generic #trait_mul_ds
ins(%arga, %argb: tensor<512x1024xf32, #SparseMatrix>, tensor<512x1024xf32>) ins(%arga, %argb: tensor<512x1024xf32, #SparseMatrix>, tensor<512x1024xf32>)
outs(%argx: tensor<512x1024xf32>) { inits(%argx: tensor<512x1024xf32>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<512x1024xf32> } -> tensor<512x1024xf32>
return %0 : tensor<512x1024xf32> return %0 : tensor<512x1024xf32>
} }
// ----- // -----
Show All 27 Lines
// CHECK: } // CHECK: }
// CHECK: } // CHECK: }
// CHECK: return // CHECK: return
// //
func.func @add_dense(%arga: tensor<32x64xf64, #SparseMatrix>, func.func @add_dense(%arga: tensor<32x64xf64, #SparseMatrix>,
%argx: tensor<33x64xf64>) -> tensor<33x64xf64> { %argx: tensor<33x64xf64>) -> tensor<33x64xf64> {
%0 = linalg.generic #trait_affine %0 = linalg.generic #trait_affine
ins(%arga: tensor<32x64xf64, #SparseMatrix>) ins(%arga: tensor<32x64xf64, #SparseMatrix>)
outs(%argx: tensor<33x64xf64>) { inits(%argx: tensor<33x64xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.addf %x, %a : f64 %0 = arith.addf %x, %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<33x64xf64> } -> tensor<33x64xf64>
return %0 : tensor<33x64xf64> return %0 : tensor<33x64xf64>
} }

mlir/test/Dialect/Tensor/canonicalize.mlir

Show First 20 Lines • Show All 1,449 Lines • ▼ Show 20 Linesfunc.func @canonicalize_parallel_insert_slice_indices(
%arg0 : tensor<1x5xf32>, %arg1: tensor<?x?xf32>, %arg0 : tensor<1x5xf32>, %arg1: tensor<?x?xf32>,
%num_threads : index) -> tensor<?x?xf32> %num_threads : index) -> tensor<?x?xf32>
{ {
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
// CHECK-NOT: tensor.cast // CHECK-NOT: tensor.cast
// CHECK: scf.foreach_thread (%[[tidx:[0-9a-z]*]]) in (%[[num_threads]]) shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor<?x?xf32>) { // CHECK: scf.foreach_thread (%[[tidx:[0-9a-z]*]]) in (%[[num_threads]]) shared_inits(%[[o:.*]] = %[[arg1]]) -> (tensor<?x?xf32>) {
// CHECK-NEXT: scf.foreach_thread.perform_concurrently { // CHECK-NEXT: scf.foreach_thread.perform_concurrently {
// CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][%[[tidx]], 0] [1, 5] [1, 1] // CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][%[[tidx]], 0] [1, 5] [1, 1]
%2 = scf.foreach_thread (%tidx) in (%num_threads) shared_outs(%o = %arg1) -> (tensor<?x?xf32>) { %2 = scf.foreach_thread (%tidx) in (%num_threads) shared_inits(%o = %arg1) -> (tensor<?x?xf32>) {
%3 = tensor.cast %arg0 : tensor<1x5xf32> to tensor<?x5xf32> %3 = tensor.cast %arg0 : tensor<1x5xf32> to tensor<?x5xf32>
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %3 into %o[%tidx, %c0] [%c1, 5] [%c1, %c1] : tensor<?x5xf32> into tensor<?x?xf32> tensor.parallel_insert_slice %3 into %o[%tidx, %c0] [%c1, 5] [%c1, %c1] : tensor<?x5xf32> into tensor<?x?xf32>
} }
} }
return %2 : tensor<?x?xf32> return %2 : tensor<?x?xf32>
} }
// ----- // -----
// CHECK-LABEL: func.func @dont_fold_parallel_insert_slice( // CHECK-LABEL: func.func @dont_fold_parallel_insert_slice(
// CHECK-SAME: %[[arg0:[0-9a-z]*]]: tensor<1x5xf32>, // CHECK-SAME: %[[arg0:[0-9a-z]*]]: tensor<1x5xf32>,
// CHECK-SAME: %[[arg1:[0-9a-z]*]]: tensor<1x5xf32>) // CHECK-SAME: %[[arg1:[0-9a-z]*]]: tensor<1x5xf32>)
func.func @dont_fold_parallel_insert_slice( func.func @dont_fold_parallel_insert_slice(
%arg0 : tensor<1x5xf32>, %arg1: tensor<1x5xf32>) -> tensor<1x5xf32> %arg0 : tensor<1x5xf32>, %arg1: tensor<1x5xf32>) -> tensor<1x5xf32>
{ {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
// CHECK: scf.foreach_thread () in () shared_outs(%[[o:.*]] = %[[arg1]]) -> (tensor<1x5xf32>) { // CHECK: scf.foreach_thread () in () shared_inits(%[[o:.*]] = %[[arg1]]) -> (tensor<1x5xf32>) {
// CHECK-NEXT: scf.foreach_thread.perform_concurrently { // CHECK-NEXT: scf.foreach_thread.perform_concurrently {
// CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][0, 0] [1, 5] [1, 1] : tensor<1x5xf32> into tensor<1x5xf32> // CHECK-NEXT: tensor.parallel_insert_slice %[[arg0]] into %[[o]][0, 0] [1, 5] [1, 1] : tensor<1x5xf32> into tensor<1x5xf32>
%2 = scf.foreach_thread () in () shared_outs(%o = %arg1) -> (tensor<1x5xf32>) { %2 = scf.foreach_thread () in () shared_inits(%o = %arg1) -> (tensor<1x5xf32>) {
scf.foreach_thread.perform_concurrently { scf.foreach_thread.perform_concurrently {
tensor.parallel_insert_slice %arg0 into %o[%c0, %c0] [1, 5] [%c1, %c1] : tensor<1x5xf32> into tensor<1x5xf32> tensor.parallel_insert_slice %arg0 into %o[%c0, %c0] [1, 5] [%c1, %c1] : tensor<1x5xf32> into tensor<1x5xf32>
} }
} }
return %2 : tensor<1x5xf32> return %2 : tensor<1x5xf32>
} }
// ----- // -----
▲ Show 20 Lines • Show All 139 Lines • Show Last 20 Lines

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

Show All 23 Lines
// 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:.+]] = tensor.empty() : tensor<20x11xf32> // 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_inits(%[[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 20 Lines • Show All 91 Lines • ▼ Show 20 Lines
// 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:.+]] = tensor.empty(%[[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_inits(%[[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]]) :
// FOREACH: %[[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] : // FOREACH: %[[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] :
// FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} : // FOREACH: %[[sliceFlat:.+]] = tensor.collapse_shape %[[slice]] {{\[}}[0, 1, 2], [3, 4]{{\]}} :
// FOREACH: perform_concurrently // FOREACH: perform_concurrently
//FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[tid1]], %[[tid2]]] [1, 1] [1, 1] : //FOREACH-NEXT: tensor.parallel_insert_slice %[[sliceFlat]] into %[[dest]][%[[tid1]], %[[tid2]]] [1, 1] [1, 1] :
▲ Show 20 Lines • Show All 93 Lines • Show Last 20 Lines

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

Show First 20 Lines • Show All 64 Lines • ▼ Show 20 Linesfunc.func @insert_slice_fun(
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
// CHECK-NOT: alloc // CHECK-NOT: alloc
// CHECK: %[[SV_A:.*]] = memref.subview %[[A]] // CHECK: %[[SV_A:.*]] = memref.subview %[[A]]
// CHECK: memref.copy %[[t]], %[[SV_A]] // CHECK: memref.copy %[[t]], %[[SV_A]]
%r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32> %r0 = tensor.insert_slice %t into %A[0][4][1] : tensor<4xf32> into tensor<?xf32>
/// Overwrite A inplace. /// Overwrite A inplace.
// CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]] // CHECK: linalg.fill ins({{.*}}{{.*}}inits(%[[A]]
%r1 = linalg.fill ins(%f0 : f32) outs(%r0 : tensor<?xf32>) -> tensor<?xf32> %r1 = linalg.fill ins(%f0 : f32) inits(%r0 : tensor<?xf32>) -> tensor<?xf32>
// CHECK: return // CHECK: return
// CHECK-NOT: tensor // CHECK-NOT: tensor
return %r1: tensor<?xf32> return %r1: tensor<?xf32>
} }
// ----- // -----
// CHECK-LABEL: func @insert_slice_fun // CHECK-LABEL: func @insert_slice_fun
// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>> // CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref<?xf32, strided<[?], offset: ?>>
// CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>> // CHECK-SAME: %[[t:[a-zA-Z0-9]*]]: memref<4xf32, strided<[?], offset: ?>>
func.func @insert_slice_fun( func.func @insert_slice_fun(
%A : tensor<?xf32> {bufferization.writable = true}, %A : tensor<?xf32> {bufferization.writable = true},
%t : tensor<4xf32> {bufferization.writable = false}) %t : tensor<4xf32> {bufferization.writable = false})
-> tensor<?xf32> -> tensor<?xf32>
{ {
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
// CHECK: linalg.fill ins({{.*}}{{.*}}outs(%[[A]] // CHECK: linalg.fill ins({{.*}}{{.*}}inits(%[[A]]
%r0 = linalg.fill ins(%f0 : f32) outs(%A : tensor<?xf32>) -> tensor<?xf32> %r0 = linalg.fill ins(%f0 : f32) inits(%A : tensor<?xf32>) -> tensor<?xf32>
// CHECK-NOT: alloc // CHECK-NOT: alloc
// CHECK: %[[SV_A:.*]] = memref.subview %[[A]] // CHECK: %[[SV_A:.*]] = memref.subview %[[A]]
/// Overwrite A inplace by copying into the subview. /// Overwrite A inplace by copying into the subview.
// CHECK: memref.copy %[[t]], %[[SV_A]] // CHECK: memref.copy %[[t]], %[[SV_A]]
%r1 = tensor.insert_slice %t into %r0[0][4][1] : tensor<4xf32> into tensor<?xf32> %r1 = tensor.insert_slice %t into %r0[0][4][1] : tensor<4xf32> into tensor<?xf32>
// CHECK: return // CHECK: return
▲ Show 20 Lines • Show All 140 Lines • ▼ Show 20 Lines
// ----- // -----
// CHECK-LABEL: func @insert_equivalent_tensor // CHECK-LABEL: func @insert_equivalent_tensor
func.func @insert_equivalent_tensor(%t: tensor<10xf32>) -> tensor<10xf32> { func.func @insert_equivalent_tensor(%t: tensor<10xf32>) -> tensor<10xf32> {
// CHECK-NOT: memref.alloc // CHECK-NOT: memref.alloc
%cst = arith.constant 4.200000e+01 : f32 %cst = arith.constant 4.200000e+01 : f32
// CHECK: linalg.fill // CHECK: linalg.fill
%0 = linalg.fill ins(%cst : f32) outs(%t : tensor<10xf32>) -> tensor<10xf32> %0 = linalg.fill ins(%cst : f32) inits(%t : tensor<10xf32>) -> tensor<10xf32>
// CHECK-NOT: memref.copy // CHECK-NOT: memref.copy
%1 = tensor.insert_slice %0 into %t[0][10][1] : tensor<10xf32> into tensor<10xf32> %1 = tensor.insert_slice %0 into %t[0][10][1] : tensor<10xf32> into tensor<10xf32>
return %1 : tensor<10xf32> return %1 : tensor<10xf32>
} }

mlir/test/Dialect/Transform/selective-targeting.mlir

// RUN: mlir-opt %s -test-transform-dialect-interpreter --split-input-file | FileCheck %s // RUN: mlir-opt %s -test-transform-dialect-interpreter --split-input-file | FileCheck %s
// CHECK-LABEL: func.func @matmul_tensors_1( // CHECK-LABEL: func.func @matmul_tensors_1(
func.func @matmul_tensors_1( func.func @matmul_tensors_1(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
%arg2: tensor<128x128xf32>) %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// This operation is marked for tiling only. // This operation is marked for tiling only.
// CHECK-COUNT-3: scf.for // CHECK-COUNT-3: scf.for
// CHECK-COUNT-3: tensor.extract_slice // CHECK-COUNT-3: tensor.extract_slice
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: -> tensor<4x4xf32> // CHECK-SAME: -> tensor<4x4xf32>
%0 = linalg.matmul { test.attrA } %0 = linalg.matmul { test.attrA }
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
func.return %0 : tensor<128x128xf32> func.return %0 : tensor<128x128xf32>
} }
func.func @matmul_tensors_2( func.func @matmul_tensors_2(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
%arg2: tensor<128x128xf32>) %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// This operation is marked f // This operation is marked f
// This operation is marked for tiling and vectorization. // This operation is marked for tiling and vectorization.
// CHECK-COUNT-3: scf.for // CHECK-COUNT-3: scf.for
// CHECK-COUNT-3: vector.transfer_read // CHECK-COUNT-3: vector.transfer_read
// CHECK: vector.contract // CHECK: vector.contract
// CHECK-NOT: linalg.matmul // CHECK-NOT: linalg.matmul
// CHECK: vector.transfer_write // CHECK: vector.transfer_write
%0 = linalg.matmul { test.attrA, test.attrC } %0 = linalg.matmul { test.attrA, test.attrC }
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
func.return %0 : tensor<128x128xf32> func.return %0 : tensor<128x128xf32>
} }
func.func @matmul_tensors_3( func.func @matmul_tensors_3(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
%arg2: tensor<128x128xf32>) %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// This operation is marked for vectorization only. // This operation is marked for vectorization only.
// CHECK-NOT: scf.for // CHECK-NOT: scf.for
// CHECK-COUNT-3: vector.transfer_read // CHECK-COUNT-3: vector.transfer_read
// CHECK: vector.contract // CHECK: vector.contract
// CHECK-SAME: into vector<128x128xf32> // CHECK-SAME: into vector<128x128xf32>
// CHECK: vector.transfer_write // CHECK: vector.transfer_write
%0 = linalg.matmul { test.attrC } %0 = linalg.matmul { test.attrC }
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
func.return %0 : tensor<128x128xf32> func.return %0 : tensor<128x128xf32>
} }
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
// Match matmul operations inside @matmul_tensors with test.attrA set. // Match matmul operations inside @matmul_tensors with test.attrA set.
pdl.pattern @pdl_target_attrA : benefit(1) { pdl.pattern @pdl_target_attrA : benefit(1) {
Show All 30 Lines
// CHECK-LABEL: @vectorize_one // CHECK-LABEL: @vectorize_one
func.func @vectorize_one( func.func @vectorize_one(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
%arg2: tensor<128x128xf32>) %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// CHECK: vector.contract // CHECK: vector.contract
%0 = linalg.matmul {test.attrA} %0 = linalg.matmul {test.attrA}
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
func.return %0 : tensor<128x128xf32> func.return %0 : tensor<128x128xf32>
} }
func.func @vectorize_none( func.func @vectorize_none(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>,
%arg2: tensor<128x128xf32>) %arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// CHECK: linalg.matmul // CHECK: linalg.matmul
%0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) %0 = linalg.matmul ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
func.return %0 : tensor<128x128xf32> func.return %0 : tensor<128x128xf32>
} }
transform.with_pdl_patterns { transform.with_pdl_patterns {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
pdl.pattern @pdl_target : benefit(1) { pdl.pattern @pdl_target : benefit(1) {
%args = operands %args = operands
Show All 17 Lines
// CHECK-LABEL: @vectorize_all // CHECK-LABEL: @vectorize_all
func.func @vectorize_all( func.func @vectorize_all(
%arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>, %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>,
%arg3: tensor<128x128xf32>) %arg3: tensor<128x128xf32>)
-> tensor<128x128xf32> { -> tensor<128x128xf32> {
// CHECK: vector.contract // CHECK: vector.contract
%0 = linalg.matmul {test.attrA} %0 = linalg.matmul {test.attrA}
ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>) ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg2: tensor<128x128xf32>) inits(%arg2: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
// CHECK: vector.contract // CHECK: vector.contract
%1 = linalg.matmul ins(%arg0, %0: tensor<128x128xf32>, tensor<128x128xf32>) %1 = linalg.matmul ins(%arg0, %0: tensor<128x128xf32>, tensor<128x128xf32>)
outs(%arg3: tensor<128x128xf32>) inits(%arg3: tensor<128x128xf32>)
-> tensor<128x128xf32> -> tensor<128x128xf32>
return %1 : tensor<128x128xf32> return %1 : tensor<128x128xf32>
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg0: !pdl.operation): ^bb0(%arg0: !pdl.operation):
transform.structured.vectorize %arg0 transform.structured.vectorize %arg0
} }

mlir/test/Dialect/Vector/vector-transfer-full-partial-split.mlir

Show First 20 Lines • Show All 66 Lines • ▼ Show 20 Linesfunc.func @split_vector_transfer_read_2d(%A: memref<?x8xf32>, %i: index, %j: index) -> vector<4x8xf32> {
// LINALG: %[[cmp1:.*]] = arith.cmpi sle, %[[idx1]], %[[c8]] : index // LINALG: %[[cmp1:.*]] = arith.cmpi sle, %[[idx1]], %[[c8]] : index
// are both conds true // are both conds true
// LINALG: %[[cond:.*]] = arith.andi %[[cmp0]], %[[cmp1]] : i1 // LINALG: %[[cond:.*]] = arith.andi %[[cmp0]], %[[cmp1]] : i1
// LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32>, index, index) { // LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32>, index, index) {
// inBounds, just yield %A // inBounds, just yield %A
// LINALG: scf.yield %[[A]], %[[i]], %[[j]] : memref<?x8xf32>, index, index // LINALG: scf.yield %[[A]], %[[i]], %[[j]] : memref<?x8xf32>, index, index
// LINALG: } else { // LINALG: } else {
// slow path, fill tmp alloc and yield a memref_casted version of it // slow path, fill tmp alloc and yield a memref_casted version of it
// LINALG: linalg.fill ins(%cst : f32) outs(%[[alloc]] : memref<4x8xf32>) // LINALG: linalg.fill ins(%cst : f32) inits(%[[alloc]] : memref<4x8xf32>)
// LINALG: %[[d0:.*]] = memref.dim %[[A]], %[[c0]] : memref<?x8xf32> // LINALG: %[[d0:.*]] = memref.dim %[[A]], %[[c0]] : memref<?x8xf32>
// LINALG: %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[d0]], %[[i]], %[[c4]]) // LINALG: %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[d0]], %[[i]], %[[c4]])
// LINALG: %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]]) // LINALG: %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]])
// LINALG: %[[sv:.*]] = memref.subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [1, 1] // LINALG: %[[sv:.*]] = memref.subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [1, 1]
// LINALG-SAME: memref<?x8xf32> to memref<?x?xf32, strided<[8, 1], offset: ?>> // LINALG-SAME: memref<?x8xf32> to memref<?x?xf32, strided<[8, 1], offset: ?>>
// LINALG: %[[alloc_view:.*]] = memref.subview %[[alloc]][0, 0] [%[[sv0]], %[[sv1]]] [1, 1] // LINALG: %[[alloc_view:.*]] = memref.subview %[[alloc]][0, 0] [%[[sv0]], %[[sv1]]] [1, 1]
// LINALG: memref.copy %[[sv]], %[[alloc_view]] : memref<?x?xf32, strided<[8, 1], offset: ?>> to memref<?x?xf32, strided{{.*}}> // LINALG: memref.copy %[[sv]], %[[alloc_view]] : memref<?x?xf32, strided<[8, 1], offset: ?>> to memref<?x?xf32, strided{{.*}}>
// LINALG: %[[yielded:.*]] = memref.cast %[[alloc]] : // LINALG: %[[yielded:.*]] = memref.cast %[[alloc]] :
▲ Show 20 Lines • Show All 76 Lines • ▼ Show 20 Linesfunc.func @split_vector_transfer_read_strided_2d(
// LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32, strided<[?, 1], offset: ?>>, index, index) { // LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32, strided<[?, 1], offset: ?>>, index, index) {
// inBounds but not cast-compatible: yield a memref_casted form of %A // inBounds but not cast-compatible: yield a memref_casted form of %A
// LINALG: %[[casted:.*]] = memref.cast %arg0 : // LINALG: %[[casted:.*]] = memref.cast %arg0 :
// LINALG-SAME: memref<7x8xf32, strided<[?, 1], offset: ?>> to memref<?x8xf32, strided<[?, 1], offset: ?>> // LINALG-SAME: memref<7x8xf32, strided<[?, 1], offset: ?>> to memref<?x8xf32, strided<[?, 1], offset: ?>>
// LINALG: scf.yield %[[casted]], %[[i]], %[[j]] : // LINALG: scf.yield %[[casted]], %[[i]], %[[j]] :
// LINALG-SAME: memref<?x8xf32, strided<[?, 1], offset: ?>>, index, index // LINALG-SAME: memref<?x8xf32, strided<[?, 1], offset: ?>>, index, index
// LINALG: } else { // LINALG: } else {
// slow path, fill tmp alloc and yield a memref_casted version of it // slow path, fill tmp alloc and yield a memref_casted version of it
// LINALG: linalg.fill ins(%cst : f32) outs(%[[alloc]] : memref<4x8xf32>) // LINALG: linalg.fill ins(%cst : f32) inits(%[[alloc]] : memref<4x8xf32>)
// LINALG: %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[c7]], %[[i]], %[[c4]]) // LINALG: %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[c7]], %[[i]], %[[c4]])
// LINALG: %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]]) // LINALG: %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]])
// LINALG: %[[sv:.*]] = memref.subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [1, 1] // LINALG: %[[sv:.*]] = memref.subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [1, 1]
// LINALG-SAME: memref<7x8xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>> // LINALG-SAME: memref<7x8xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided<[?, 1], offset: ?>>
// LINALG: %[[alloc_view:.*]] = memref.subview %[[alloc]][0, 0] [%[[sv0]], %[[sv1]]] [1, 1] // LINALG: %[[alloc_view:.*]] = memref.subview %[[alloc]][0, 0] [%[[sv0]], %[[sv1]]] [1, 1]
// LINALG: memref.copy %[[sv]], %[[alloc_view]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided{{.*}}> // LINALG: memref.copy %[[sv]], %[[alloc_view]] : memref<?x?xf32, strided<[?, 1], offset: ?>> to memref<?x?xf32, strided{{.*}}>
// LINALG: %[[yielded:.*]] = memref.cast %[[alloc]] : // LINALG: %[[yielded:.*]] = memref.cast %[[alloc]] :
// LINALG-SAME: memref<4x8xf32> to memref<?x8xf32, strided<[?, 1], offset: ?>> // LINALG-SAME: memref<4x8xf32> to memref<?x8xf32, strided<[?, 1], offset: ?>>
▲ Show 20 Lines • Show All 251 Lines • Show Last 20 Lines

mlir/test/IR/slice.mlir

// RUN: mlir-opt -slice-analysis-test %s | FileCheck %s // RUN: mlir-opt -slice-analysis-test %s | FileCheck %s
func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) { func.func @slicing_linalg_op(%arg0 : index, %arg1 : index, %arg2 : index) {
%a = memref.alloc(%arg0, %arg2) : memref<?x?xf32> %a = memref.alloc(%arg0, %arg2) : memref<?x?xf32>
%b = memref.alloc(%arg2, %arg1) : memref<?x?xf32> %b = memref.alloc(%arg2, %arg1) : memref<?x?xf32>
%c = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %c = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
%d = memref.alloc(%arg0, %arg1) : memref<?x?xf32> %d = memref.alloc(%arg0, %arg1) : memref<?x?xf32>
linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
outs(%c : memref<?x?xf32>) inits(%c : memref<?x?xf32>)
linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%a, %b : memref<?x?xf32>, memref<?x?xf32>)
outs(%d : memref<?x?xf32>) inits(%d : memref<?x?xf32>)
memref.dealloc %c : memref<?x?xf32> memref.dealloc %c : memref<?x?xf32>
memref.dealloc %b : memref<?x?xf32> memref.dealloc %b : memref<?x?xf32>
memref.dealloc %a : memref<?x?xf32> memref.dealloc %a : memref<?x?xf32>
memref.dealloc %d : memref<?x?xf32> memref.dealloc %d : memref<?x?xf32>
return return
} }
// CHECK-LABEL: func @slicing_linalg_op__backward_slice__0 // CHECK-LABEL: func @slicing_linalg_op__backward_slice__0
Show All 16 Lines

mlir/test/Integration/Dialect/Async/CPU/microbench-linalg-async-parallel-for.mlir

Show All 37 Lines
func.func @linalg_generic(%lhs: memref<?x?xf32>, func.func @linalg_generic(%lhs: memref<?x?xf32>,
%rhs: memref<?x?xf32>, %rhs: memref<?x?xf32>,
%sum: memref<?x?xf32>) { %sum: memref<?x?xf32>) {
linalg.generic { linalg.generic {
indexing_maps = [#map0, #map0, #map0], indexing_maps = [#map0, #map0, #map0],
iterator_types = ["parallel", "parallel"] iterator_types = ["parallel", "parallel"]
} }
ins(%lhs, %rhs : memref<?x?xf32>, memref<?x?xf32>) ins(%lhs, %rhs : memref<?x?xf32>, memref<?x?xf32>)
outs(%sum : memref<?x?xf32>) inits(%sum : memref<?x?xf32>)
{ {
^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32): ^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32):
%0 = arith.addf %lhs_in, %rhs_in : f32 %0 = arith.addf %lhs_in, %rhs_in : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
func.func @entry() { func.func @entry() {
%f1 = arith.constant 1.0 : f32 %f1 = arith.constant 1.0 : f32
%f4 = arith.constant 4.0 : f32 %f4 = arith.constant 4.0 : f32
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cN = arith.constant 50 : index %cN = arith.constant 50 : index
// //
// Sanity check for the function under test. // Sanity check for the function under test.
// //
%LHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %LHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
%RHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %RHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
%DST10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %DST10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
linalg.fill ins(%f1 : f32) outs(%LHS10 : memref<1x10xf32>) linalg.fill ins(%f1 : f32) inits(%LHS10 : memref<1x10xf32>)
linalg.fill ins(%f1 : f32) outs(%RHS10 : memref<1x10xf32>) linalg.fill ins(%f1 : f32) inits(%RHS10 : memref<1x10xf32>)
%LHS = memref.cast %LHS10 : memref<1x10xf32> to memref<?x?xf32> %LHS = memref.cast %LHS10 : memref<1x10xf32> to memref<?x?xf32>
%RHS = memref.cast %RHS10 : memref<1x10xf32> to memref<?x?xf32> %RHS = memref.cast %RHS10 : memref<1x10xf32> to memref<?x?xf32>
%DST = memref.cast %DST10 : memref<1x10xf32> to memref<?x?xf32> %DST = memref.cast %DST10 : memref<1x10xf32> to memref<?x?xf32>
call @linalg_generic(%LHS, %RHS, %DST) call @linalg_generic(%LHS, %RHS, %DST)
: (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> () : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
▲ Show 20 Lines • Show All 54 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Async/CPU/microbench-scf-async-parallel-for.mlir

Show First 20 Lines • Show All 84 Lines • ▼ Show 20 Linesfunc.func @entry() {
// //
// Sanity check for the function under test. // Sanity check for the function under test.
// //
%LHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %LHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
%RHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %RHS10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
%DST10 = memref.alloc() {alignment = 64} : memref<1x10xf32> %DST10 = memref.alloc() {alignment = 64} : memref<1x10xf32>
linalg.fill ins(%f1 : f32) outs(%LHS10 : memref<1x10xf32>) linalg.fill ins(%f1 : f32) inits(%LHS10 : memref<1x10xf32>)
linalg.fill ins(%f1 : f32) outs(%RHS10 : memref<1x10xf32>) linalg.fill ins(%f1 : f32) inits(%RHS10 : memref<1x10xf32>)
%LHS = memref.cast %LHS10 : memref<1x10xf32> to memref<?x?xf32> %LHS = memref.cast %LHS10 : memref<1x10xf32> to memref<?x?xf32>
%RHS = memref.cast %RHS10 : memref<1x10xf32> to memref<?x?xf32> %RHS = memref.cast %RHS10 : memref<1x10xf32> to memref<?x?xf32>
%DST = memref.cast %DST10 : memref<1x10xf32> to memref<?x?xf32> %DST = memref.cast %DST10 : memref<1x10xf32> to memref<?x?xf32>
call @scf_parallel(%LHS, %RHS, %DST) call @scf_parallel(%LHS, %RHS, %DST)
: (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> () : (memref<?x?xf32>, memref<?x?xf32>, memref<?x?xf32>) -> ()
▲ Show 20 Lines • Show All 54 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/matmul-vs-matvec.mlir

// RUN: mlir-opt %s -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-opt %s -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 -O3 -e main -entry-point-result=void \ // RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) { func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%x = memref.dim %A, %c0 : memref<?x?xf32> %x = memref.dim %A, %c0 : memref<?x?xf32>
%y = memref.dim %B, %c1 : memref<?x?xf32> %y = memref.dim %B, %c1 : memref<?x?xf32>
%C = memref.alloc(%x, %y) : memref<?x?xf32> %C = memref.alloc(%x, %y) : memref<?x?xf32>
linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>) linalg.fill ins(%f0 : f32) inits(%C : memref<?x?xf32>)
linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
outs(%C: memref<?x?xf32>) inits(%C: memref<?x?xf32>)
return %C : memref<?x?xf32> return %C : memref<?x?xf32>
} }
func.func @matvec(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) { func.func @matvec(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%f0 = arith.constant 0.0 : f32 %f0 = arith.constant 0.0 : f32
%m = memref.dim %A, %c0 : memref<?x?xf32> %m = memref.dim %A, %c0 : memref<?x?xf32>
%x = memref.dim %A, %c1 : memref<?x?xf32> %x = memref.dim %A, %c1 : memref<?x?xf32>
%n = memref.dim %B, %c1 : memref<?x?xf32> %n = memref.dim %B, %c1 : memref<?x?xf32>
%C = memref.alloc(%m, %n) : memref<?x?xf32> %C = memref.alloc(%m, %n) : memref<?x?xf32>
linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>) linalg.fill ins(%f0 : f32) inits(%C : memref<?x?xf32>)
scf.for %i = %c0 to %n step %c1 { scf.for %i = %c0 to %n step %c1 {
%b = memref.subview %B[0, %i][%x, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>> %b = memref.subview %B[0, %i][%x, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>>
%c = memref.subview %C[0, %i][%m, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>> %c = memref.subview %C[0, %i][%m, 1][1, 1] : memref<?x?xf32> to memref<?xf32, strided<[?], offset: ?>>
linalg.matvec ins(%A, %b: memref<?x?xf32>, memref<?xf32, strided<[?], offset: ?>>) linalg.matvec ins(%A, %b: memref<?x?xf32>, memref<?xf32, strided<[?], offset: ?>>)
outs(%c: memref<?xf32, strided<[?], offset: ?>>) inits(%c: memref<?xf32, strided<[?], offset: ?>>)
} }
return %C : memref<?x?xf32> return %C : memref<?x?xf32>
} }
func.func @main() { func.func @main() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%m = arith.constant 5 : index %m = arith.constant 5 : index
%x = arith.constant 3 : index %x = arith.constant 3 : index
%n = arith.constant 2 : index %n = arith.constant 2 : index
%val1 = arith.constant 13.0 : f32 %val1 = arith.constant 13.0 : f32
%val2 = arith.constant 17.0 : f32 %val2 = arith.constant 17.0 : f32
%A = memref.alloc(%m, %x) : memref<?x?xf32> %A = memref.alloc(%m, %x) : memref<?x?xf32>
%B = memref.alloc(%x, %n) : memref<?x?xf32> %B = memref.alloc(%x, %n) : memref<?x?xf32>
linalg.fill ins(%val1 : f32) outs(%A : memref<?x?xf32>) linalg.fill ins(%val1 : f32) inits(%A : memref<?x?xf32>)
linalg.fill ins(%val2 : f32) outs(%B : memref<?x?xf32>) linalg.fill ins(%val2 : f32) inits(%B : memref<?x?xf32>)
memref.store %val1, %B[%c0, %c0] : memref<?x?xf32> memref.store %val1, %B[%c0, %c0] : memref<?x?xf32>
%C1 = call @matmul(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32> %C1 = call @matmul(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
%C2 = call @matvec(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32> %C2 = call @matvec(%A, %B) : (memref<?x?xf32>, memref<?x?xf32>) -> memref<?x?xf32>
scf.for %i = %c0 to %m step %c1 { scf.for %i = %c0 to %m step %c1 {
scf.for %j = %c0 to %n step %c1 { scf.for %j = %c0 to %n step %c1 {
%e1 = memref.load %C1[%i, %j] : memref<?x?xf32> %e1 = memref.load %C1[%i, %j] : memref<?x?xf32>
%e2 = memref.load %C2[%i, %j] : memref<?x?xf32> %e2 = memref.load %C2[%i, %j] : memref<?x?xf32>
%c = arith.cmpf oeq, %e1, %e2 : f32 %c = arith.cmpf oeq, %e1, %e2 : f32
Show All 18 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-1d-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns a 1-D buffer of size %s1 filled with the value %f // Creates and returns a 1-D buffer of size %s1 filled with the value %f
func.func @alloc_1d_filled_f32(%s1 : index, %f : f32) -> memref<?xf32> { func.func @alloc_1d_filled_f32(%s1 : index, %f : f32) -> memref<?xf32> {
%buf = memref.alloc(%s1) : memref<?xf32> %buf = memref.alloc(%s1) : memref<?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?xf32>)
return %buf : memref<?xf32> return %buf : memref<?xf32>
} }
func.func @conv_1d(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) { func.func @conv_1d(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %arg2: memref<?xf32>) {
linalg.conv_1d ins (%arg0, %arg1: memref<?xf32>, memref<?xf32>) linalg.conv_1d ins (%arg0, %arg1: memref<?xf32>, memref<?xf32>)
outs (%arg2: memref<?xf32>) inits (%arg2: memref<?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_1d"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_1d"]} in %arg1
%1, %loop = transform.structured.tile %0 [4] %1, %loop = transform.structured.tile %0 [4]
} }
Show All 26 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-1d-nwc-wcf-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f // Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f
func.func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> { func.func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> {
%buf = memref.alloc(%s1, %s2, %s3) : memref<?x?x?xf32> %buf = memref.alloc(%s1, %s2, %s3) : memref<?x?x?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?x?x?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?x?x?xf32>)
return %buf : memref<?x?x?xf32> return %buf : memref<?x?x?xf32>
} }
func.func @conv_1d_nwc_wcf(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) { func.func @conv_1d_nwc_wcf(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) {
linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>, linalg.conv_1d_nwc_wcf {dilations = dense<1> : tensor<1xi64>,
strides = dense<1> : tensor<1xi64>} strides = dense<1> : tensor<1xi64>}
ins (%arg0, %arg1: memref<?x?x?xf32>, memref<?x?x?xf32>) ins (%arg0, %arg1: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%arg2: memref<?x?x?xf32>) inits (%arg2: memref<?x?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_1d_nwc_wcf"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_1d_nwc_wcf"]} in %arg1
%1, %loops:2 = transform.structured.tile %0 [2, 4] %1, %loops:2 = transform.structured.tile %0 [2, 4]
} }
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns a 2-D buffer of size (%s1, %s2) filled with the value %f // Creates and returns a 2-D buffer of size (%s1, %s2) filled with the value %f
func.func @alloc_2d_filled_f32(%s1 : index, %s2 : index, %f : f32) -> memref<?x?xf32> { func.func @alloc_2d_filled_f32(%s1 : index, %s2 : index, %f : f32) -> memref<?x?xf32> {
%buf = memref.alloc(%s1, %s2) : memref<?x?xf32> %buf = memref.alloc(%s1, %s2) : memref<?x?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?x?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?x?xf32>)
return %buf : memref<?x?xf32> return %buf : memref<?x?xf32>
} }
func.func @conv_2d(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) { func.func @conv_2d(%arg0: memref<?x?xf32>, %arg1: memref<?x?xf32>, %arg2: memref<?x?xf32>) {
linalg.conv_2d ins (%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>) linalg.conv_2d ins (%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>)
outs (%arg2: memref<?x?xf32>) inits (%arg2: memref<?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1
%1, %loops:2 = transform.structured.tile %0 [2, 2] %1, %loops:2 = transform.structured.tile %0 [2, 2]
} }
Show All 31 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-2d-nhwc-hwcf-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f // Creates and returns 4-D buffer of size (%s1, %s2, %s3, %s4) filled with the value %f
func.func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> memref<?x?x?x?xf32> { func.func @alloc_4d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %f : f32) -> memref<?x?x?x?xf32> {
%buf = memref.alloc(%s1, %s2, %s3, %s4) : memref<?x?x?x?xf32> %buf = memref.alloc(%s1, %s2, %s3, %s4) : memref<?x?x?x?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?x?x?x?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?x?x?x?xf32>)
return %buf : memref<?x?x?x?xf32> return %buf : memref<?x?x?x?xf32>
} }
func.func @conv_2d_nhwc_hwcf(%arg0: memref<?x?x?x?xf32>, %arg1: memref<?x?x?x?xf32>, %arg2: memref<?x?x?x?xf32>) { func.func @conv_2d_nhwc_hwcf(%arg0: memref<?x?x?x?xf32>, %arg1: memref<?x?x?x?xf32>, %arg2: memref<?x?x?x?xf32>) {
linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
strides = dense<1> : tensor<2xi64>} strides = dense<1> : tensor<2xi64>}
ins (%arg0, %arg1: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>) ins (%arg0, %arg1: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
outs (%arg2: memref<?x?x?x?xf32>) inits (%arg2: memref<?x?x?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1
%1, %loops:4 = transform.structured.tile %0 [2, 3, 3, 2] %1, %loops:4 = transform.structured.tile %0 [2, 3, 3, 2]
} }
▲ Show 20 Lines • Show All 91 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-3d-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f // Creates and returns 3-D buffer of size (%s1, %s2, %s3) filled with the value %f
func.func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> { func.func @alloc_3d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %f : f32) -> memref<?x?x?xf32> {
%buf = memref.alloc(%s1, %s2, %s3) : memref<?x?x?xf32> %buf = memref.alloc(%s1, %s2, %s3) : memref<?x?x?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?x?x?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?x?x?xf32>)
return %buf : memref<?x?x?xf32> return %buf : memref<?x?x?xf32>
} }
func.func @conv_3d(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) { func.func @conv_3d(%arg0: memref<?x?x?xf32>, %arg1: memref<?x?x?xf32>, %arg2: memref<?x?x?xf32>) {
linalg.conv_3d ins (%arg0, %arg1: memref<?x?x?xf32>, memref<?x?x?xf32>) linalg.conv_3d ins (%arg0, %arg1: memref<?x?x?xf32>, memref<?x?x?xf32>)
outs (%arg2: memref<?x?x?xf32>) inits (%arg2: memref<?x?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_3d"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_3d"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [2, 2, 2] %1, %loops:3 = transform.structured.tile %0 [2, 2, 2]
} }
▲ Show 20 Lines • Show All 48 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-conv-3d-ndhwc-dhwcf-call.mlir

// RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: mlir-opt %s -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
// RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \ // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-transform-dialect-erase-schedule -convert-linalg-to-loops -convert-scf-to-cf \
// RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --convert-memref-to-llvm -convert-func-to-llvm -reconcile-unrealized-casts | \ // RUN: -convert-linalg-to-llvm -lower-affine -convert-scf-to-cf --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_runner_utils%shlibext \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext \
// RUN: | FileCheck %s // RUN: | FileCheck %s
func.func private @printMemrefF32(memref<*xf32>) func.func private @printMemrefF32(memref<*xf32>)
// Creates and returns 5-D buffer of size (%s1, %s2, %s3, %s4, %s5) filled with the value %f // Creates and returns 5-D buffer of size (%s1, %s2, %s3, %s4, %s5) filled with the value %f
func.func @alloc_5d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %s5 : index, %f : f32) -> memref<?x?x?x?x?xf32> { func.func @alloc_5d_filled_f32(%s1 : index, %s2 : index, %s3 : index, %s4 : index, %s5 : index, %f : f32) -> memref<?x?x?x?x?xf32> {
%buf = memref.alloc(%s1, %s2, %s3, %s4, %s5) : memref<?x?x?x?x?xf32> %buf = memref.alloc(%s1, %s2, %s3, %s4, %s5) : memref<?x?x?x?x?xf32>
linalg.fill ins(%f : f32) outs(%buf : memref<?x?x?x?x?xf32>) linalg.fill ins(%f : f32) inits(%buf : memref<?x?x?x?x?xf32>)
return %buf : memref<?x?x?x?x?xf32> return %buf : memref<?x?x?x?x?xf32>
} }
func.func @conv_3d_ndhwc_dhwcf(%arg0: memref<?x?x?x?x?xf32>, %arg1: memref<?x?x?x?x?xf32>, %arg2: memref<?x?x?x?x?xf32>) { func.func @conv_3d_ndhwc_dhwcf(%arg0: memref<?x?x?x?x?xf32>, %arg1: memref<?x?x?x?x?xf32>, %arg2: memref<?x?x?x?x?xf32>) {
linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>, linalg.conv_3d_ndhwc_dhwcf {dilations = dense<1> : tensor<3xi64>,
strides = dense<1> : tensor<3xi64>} strides = dense<1> : tensor<3xi64>}
ins (%arg0, %arg1: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>) ins (%arg0, %arg1: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
outs (%arg2: memref<?x?x?x?x?xf32>) inits (%arg2: memref<?x?x?x?x?xf32>)
return return
} }
transform.sequence failures(propagate) { transform.sequence failures(propagate) {
^bb0(%arg1: !pdl.operation): ^bb0(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["linalg.conv_3d_ndhwc_dhwcf"]} in %arg1 %0 = transform.structured.match ops{["linalg.conv_3d_ndhwc_dhwcf"]} in %arg1
%1, %loops:3 = transform.structured.tile %0 [0, 5, 5, 5] %1, %loops:3 = transform.structured.tile %0 [0, 5, 5, 5]
} }
▲ Show 20 Lines • Show All 153 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/Linalg/CPU/test-one-shot-bufferize.mlir

// RUN: mlir-opt %s -pass-pipeline="func.func(canonicalize,cse),one-shot-bufferize{bufferize-function-boundaries}" |\ // RUN: mlir-opt %s -pass-pipeline="func.func(canonicalize,cse),one-shot-bufferize{bufferize-function-boundaries}" |\
// RUN: mlir-opt -pass-pipeline="func.func(buffer-deallocation,convert-vector-to-scf,lower-affine,convert-linalg-to-loops)" |\ // RUN: mlir-opt -pass-pipeline="func.func(buffer-deallocation,convert-vector-to-scf,lower-affine,convert-linalg-to-loops)" |\
// RUN: mlir-opt -pass-pipeline="func.func(canonicalize,convert-scf-to-cf),convert-vector-to-llvm,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts" | \ // RUN: mlir-opt -pass-pipeline="func.func(canonicalize,convert-scf-to-cf),convert-vector-to-llvm,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts" | \
// RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \ // RUN: mlir-cpu-runner -O3 -e main -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext,%mlir_lib_dir/libmlir_c_runner_utils%shlibext |\ // RUN: -shared-libs=%mlir_lib_dir/libmlir_runner_utils%shlibext,%mlir_lib_dir/libmlir_c_runner_utils%shlibext |\
// RUN: FileCheck %s // RUN: FileCheck %s
#map0 = affine_map<(d0, d1)[s0] -> ((d1 - d0) ceildiv s0)> #map0 = affine_map<(d0, d1)[s0] -> ((d1 - d0) ceildiv s0)>
#map1 = affine_map<(d0, d1)[s0] -> ((d0 - d1) ceildiv s0)> #map1 = affine_map<(d0, d1)[s0] -> ((d0 - d1) ceildiv s0)>
func.func @init_and_dot(%arg0: tensor<64xf32>, %arg1: tensor<64xf32>, %arg2: tensor<f32>) -> tensor<f32> { func.func @init_and_dot(%arg0: tensor<64xf32>, %arg1: tensor<64xf32>, %arg2: tensor<f32>) -> tensor<f32> {
%c64 = arith.constant 64 : index %c64 = arith.constant 64 : index
%cst = arith.constant 0.000000e+00 : f32 %cst = arith.constant 0.000000e+00 : f32
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<f32>) -> tensor<f32> %0 = linalg.fill ins(%cst : f32) inits(%arg2 : tensor<f32>) -> tensor<f32>
%1 = affine.apply #map0(%c0, %c64)[%c2] %1 = affine.apply #map0(%c0, %c64)[%c2]
%2 = bufferization.alloc_tensor(%1) : tensor<?x2xf32> %2 = bufferization.alloc_tensor(%1) : tensor<?x2xf32>
%3 = scf.for %arg3 = %c0 to %c64 step %c2 iter_args(%arg4 = %2) -> (tensor<?x2xf32>) { %3 = scf.for %arg3 = %c0 to %c64 step %c2 iter_args(%arg4 = %2) -> (tensor<?x2xf32>) {
%8 = affine.apply #map1(%arg3, %c0)[%c2] %8 = affine.apply #map1(%arg3, %c0)[%c2]
%9 = tensor.extract_slice %arg1[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32> %9 = tensor.extract_slice %arg1[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32>
%10 = tensor.cast %9 : tensor<2xf32> to tensor<?xf32> %10 = tensor.cast %9 : tensor<2xf32> to tensor<?xf32>
%11 = tensor.pad %10 low[%c0] high[%c0] { %11 = tensor.pad %10 low[%c0] high[%c0] {
^bb0(%arg5: index): ^bb0(%arg5: index):
Show All 30 Lines %7 = scf.for %arg3 = %c0 to %c64 step %c2 iter_args(%arg4 = %0) -> (tensor<f32>) {
%8 = tensor.extract_slice %arg0[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32> %8 = tensor.extract_slice %arg0[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32>
%9 = tensor.cast %8 : tensor<2xf32> to tensor<?xf32> %9 = tensor.cast %8 : tensor<2xf32> to tensor<?xf32>
%10 = tensor.extract_slice %arg1[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32> %10 = tensor.extract_slice %arg1[%arg3] [2] [1] : tensor<64xf32> to tensor<2xf32>
%11 = tensor.cast %10 : tensor<2xf32> to tensor<?xf32> %11 = tensor.cast %10 : tensor<2xf32> to tensor<?xf32>
%12 = affine.apply #map1(%arg3, %c0)[%c2] %12 = affine.apply #map1(%arg3, %c0)[%c2]
%13 = tensor.extract_slice %6[%12, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32> %13 = tensor.extract_slice %6[%12, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
%14 = affine.apply #map1(%arg3, %c0)[%c2] %14 = affine.apply #map1(%arg3, %c0)[%c2]
%15 = tensor.extract_slice %3[%14, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32> %15 = tensor.extract_slice %3[%14, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
%16 = linalg.dot ins(%13, %15 : tensor<2xf32>, tensor<2xf32>) outs(%arg4 : tensor<f32>) -> tensor<f32> %16 = linalg.dot ins(%13, %15 : tensor<2xf32>, tensor<2xf32>) inits(%arg4 : tensor<f32>) -> tensor<f32>
// %AA = tensor.cast %13 : tensor<2xf32> to tensor<*xf32> // %AA = tensor.cast %13 : tensor<2xf32> to tensor<*xf32>
// call @printMemrefF32(%AA) : (tensor<*xf32>) -> () // call @printMemrefF32(%AA) : (tensor<*xf32>) -> ()
// %BB = tensor.cast %15 : tensor<2xf32> to tensor<*xf32> // %BB = tensor.cast %15 : tensor<2xf32> to tensor<*xf32>
// call @printMemrefF32(%BB) : (tensor<*xf32>) -> () // call @printMemrefF32(%BB) : (tensor<*xf32>) -> ()
// %CC = tensor.cast %16 : tensor<f32> to tensor<*xf32> // %CC = tensor.cast %16 : tensor<f32> to tensor<*xf32>
// call @printMemrefF32(%CC) : (tensor<*xf32>) -> () // call @printMemrefF32(%CC) : (tensor<*xf32>) -> ()
scf.yield %16 : tensor<f32> scf.yield %16 : tensor<f32>
} }
return %7 : tensor<f32> return %7 : tensor<f32>
} }
func.func @main() { func.func @main() {
%v0 = arith.constant 0.0 : f32 %v0 = arith.constant 0.0 : f32
%v1 = arith.constant 1.0 : f32 %v1 = arith.constant 1.0 : f32
%v2 = arith.constant 2.0 : f32 %v2 = arith.constant 2.0 : f32
%A = bufferization.alloc_tensor() : tensor<64xf32> %A = bufferization.alloc_tensor() : tensor<64xf32>
%B = bufferization.alloc_tensor() : tensor<64xf32> %B = bufferization.alloc_tensor() : tensor<64xf32>
%C = bufferization.alloc_tensor() : tensor<f32> %C = bufferization.alloc_tensor() : tensor<f32>
%AA = linalg.fill ins(%v1 : f32) outs(%A : tensor<64xf32>) -> tensor<64xf32> %AA = linalg.fill ins(%v1 : f32) inits(%A : tensor<64xf32>) -> tensor<64xf32>
%BB = linalg.fill ins(%v2 : f32) outs(%B : tensor<64xf32>) -> tensor<64xf32> %BB = linalg.fill ins(%v2 : f32) inits(%B : tensor<64xf32>) -> tensor<64xf32>
%CC = linalg.fill ins(%v0 : f32) outs(%C : tensor<f32>) -> tensor<f32> %CC = linalg.fill ins(%v0 : f32) inits(%C : tensor<f32>) -> tensor<f32>
%res = call @init_and_dot(%AA, %BB, %CC) : %res = call @init_and_dot(%AA, %BB, %CC) :
(tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32> (tensor<64xf32>, tensor<64xf32>, tensor<f32>) -> tensor<f32>
%res2 = tensor.cast %res: tensor<f32> to tensor<*xf32> %res2 = tensor.cast %res: tensor<f32> to tensor<*xf32>
// CHECK: Unranked Memref base@ = {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data = // CHECK: Unranked Memref base@ = {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =
// CHECK-NEXT: [128] // CHECK-NEXT: [128]
call @printMemrefF32(%res2) : (tensor<*xf32>) -> () call @printMemrefF32(%res2) : (tensor<*xf32>) -> ()
return return
} }
func.func private @printMemrefF32(tensor<*xf32>) attributes { llvm.emit_c_interface } func.func private @printMemrefF32(tensor<*xf32>) attributes { llvm.emit_c_interface }

mlir/test/Integration/Dialect/Linalg/CPU/test-tensor-matmul.mlir

Show All 17 Lines
func.func @main() { func.func @main() {
%A = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32> %A = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]> : tensor<2x3xf32>
%B = arith.constant dense<[[1.0, 2.0, 3.0, 4.0], %B = arith.constant dense<[[1.0, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0], [5.0, 6.0, 7.0, 8.0],
[9.0, 10.0, 11.0, 12.0]]> : tensor<3x4xf32> [9.0, 10.0, 11.0, 12.0]]> : tensor<3x4xf32>
%C = arith.constant dense<1000.0> : tensor<2x4xf32> %C = arith.constant dense<1000.0> : tensor<2x4xf32>
%D = linalg.matmul ins(%A, %B: tensor<2x3xf32>, tensor<3x4xf32>) %D = linalg.matmul ins(%A, %B: tensor<2x3xf32>, tensor<3x4xf32>)
outs(%C: tensor<2x4xf32>) -> tensor<2x4xf32> inits(%C: tensor<2x4xf32>) -> tensor<2x4xf32>
%unranked = tensor.cast %D : tensor<2x4xf32> to tensor<*xf32> %unranked = tensor.cast %D : tensor<2x4xf32> to tensor<*xf32>
call @printMemrefF32(%unranked) : (tensor<*xf32>) -> () call @printMemrefF32(%unranked) : (tensor<*xf32>) -> ()
// CHECK: Unranked Memref base@ = {{0x[-9a-f]*}} // CHECK: Unranked Memref base@ = {{0x[-9a-f]*}}
// CHECK-SAME: rank = 2 offset = 0 sizes = [2, 4] strides = [4, 1] data = // CHECK-SAME: rank = 2 offset = 0 sizes = [2, 4] strides = [4, 1] data =
// CHECK-NEXT: [1038, 1044, 1050, 1056] // CHECK-NEXT: [1038, 1044, 1050, 1056]
// CHECK-NEXT: [1083, 1098, 1113, 1128] // CHECK-NEXT: [1083, 1098, 1113, 1128]
Show All 11 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir

Show First 20 Lines • Show All 44 Lines • ▼ Show 20 Lines func.func @dense_output(%arga: tensor<?x?xf64, #SparseMatrix>) -> tensor<?x?xf64, #DenseMatrix> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2.0 : f64 %c2 = arith.constant 2.0 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #SparseMatrix> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #SparseMatrix>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #SparseMatrix> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #SparseMatrix>
%init = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DenseMatrix> %init = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DenseMatrix>
%0 = linalg.generic #trait_assign %0 = linalg.generic #trait_assign
ins(%arga: tensor<?x?xf64, #SparseMatrix>) ins(%arga: tensor<?x?xf64, #SparseMatrix>)
outs(%init: tensor<?x?xf64, #DenseMatrix>) { inits(%init: tensor<?x?xf64, #DenseMatrix>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.mulf %a, %c2 : f64 %0 = arith.mulf %a, %c2 : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<?x?xf64, #DenseMatrix> } -> tensor<?x?xf64, #DenseMatrix>
return %0 : tensor<?x?xf64, #DenseMatrix> return %0 : tensor<?x?xf64, #DenseMatrix>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 35 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir

Show All 22 Linesmodule {
// When there is no overlap, include the present value in the output. // When there is no overlap, include the present value in the output.
func.func @vector_min(%arga: tensor<?xbf16, #SparseVector>, func.func @vector_min(%arga: tensor<?xbf16, #SparseVector>,
%argb: tensor<?xbf16, #SparseVector>) -> tensor<?xbf16, #DenseVector> { %argb: tensor<?xbf16, #SparseVector>) -> tensor<?xbf16, #DenseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xbf16, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xbf16, #SparseVector>
%xv = bufferization.alloc_tensor (%d) : tensor<?xbf16, #DenseVector> %xv = bufferization.alloc_tensor (%d) : tensor<?xbf16, #DenseVector>
%0 = linalg.generic #trait_vec_op %0 = linalg.generic #trait_vec_op
ins(%arga, %argb: tensor<?xbf16, #SparseVector>, tensor<?xbf16, #SparseVector>) ins(%arga, %argb: tensor<?xbf16, #SparseVector>, tensor<?xbf16, #SparseVector>)
outs(%xv: tensor<?xbf16, #DenseVector>) { inits(%xv: tensor<?xbf16, #DenseVector>) {
^bb(%a: bf16, %b: bf16, %x: bf16): ^bb(%a: bf16, %b: bf16, %x: bf16):
%1 = sparse_tensor.binary %a, %b : bf16, bf16 to bf16 %1 = sparse_tensor.binary %a, %b : bf16, bf16 to bf16
overlap={ overlap={
^bb0(%a0: bf16, %b0: bf16): ^bb0(%a0: bf16, %b0: bf16):
%cmp = arith.cmpf "olt", %a0, %b0 : bf16 %cmp = arith.cmpf "olt", %a0, %b0 : bf16
%2 = arith.select %cmp, %a0, %b0: bf16 %2 = arith.select %cmp, %a0, %b0: bf16
sparse_tensor.yield %2 : bf16 sparse_tensor.yield %2 : bf16
} }
▲ Show 20 Lines • Show All 53 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir

Show All 20 Linesmodule {
// When there is no overlap, include the present value in the output. // When there is no overlap, include the present value in the output.
func.func @vector_min(%arga: tensor<?xf16, #SparseVector>, func.func @vector_min(%arga: tensor<?xf16, #SparseVector>,
%argb: tensor<?xf16, #SparseVector>) -> tensor<?xf16, #DenseVector> { %argb: tensor<?xf16, #SparseVector>) -> tensor<?xf16, #DenseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf16, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf16, #SparseVector>
%xv = bufferization.alloc_tensor (%d) : tensor<?xf16, #DenseVector> %xv = bufferization.alloc_tensor (%d) : tensor<?xf16, #DenseVector>
%0 = linalg.generic #trait_vec_op %0 = linalg.generic #trait_vec_op
ins(%arga, %argb: tensor<?xf16, #SparseVector>, tensor<?xf16, #SparseVector>) ins(%arga, %argb: tensor<?xf16, #SparseVector>, tensor<?xf16, #SparseVector>)
outs(%xv: tensor<?xf16, #DenseVector>) { inits(%xv: tensor<?xf16, #DenseVector>) {
^bb(%a: f16, %b: f16, %x: f16): ^bb(%a: f16, %b: f16, %x: f16):
%1 = sparse_tensor.binary %a, %b : f16, f16 to f16 %1 = sparse_tensor.binary %a, %b : f16, f16 to f16
overlap={ overlap={
^bb0(%a0: f16, %b0: f16): ^bb0(%a0: f16, %b0: f16):
%cmp = arith.cmpf "olt", %a0, %b0 : f16 %cmp = arith.cmpf "olt", %a0, %b0 : f16
%2 = arith.select %cmp, %a0, %b0: f16 %2 = arith.select %cmp, %a0, %b0: f16
sparse_tensor.yield %2 : f16 sparse_tensor.yield %2 : f16
} }
▲ Show 20 Lines • Show All 53 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_abs.mlir

Show All 17 Lines
module { module {
func.func @sparse_absf(%arg0: tensor<?xf64, #SparseVector>) func.func @sparse_absf(%arg0: tensor<?xf64, #SparseVector>)
-> tensor<?xf64, #SparseVector> { -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector> %d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector>
%xin = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xin = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arg0: tensor<?xf64, #SparseVector>) ins(%arg0: tensor<?xf64, #SparseVector>)
outs(%xin: tensor<?xf64, #SparseVector>) { inits(%xin: tensor<?xf64, #SparseVector>) {
^bb0(%a: f64, %x: f64) : ^bb0(%a: f64, %x: f64) :
%result = math.absf %a : f64 %result = math.absf %a : f64
linalg.yield %result : f64 linalg.yield %result : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
func.func @sparse_absi(%arg0: tensor<?xi32, #SparseVector>) func.func @sparse_absi(%arg0: tensor<?xi32, #SparseVector>)
-> tensor<?xi32, #SparseVector> { -> tensor<?xi32, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arg0, %c0 : tensor<?xi32, #SparseVector> %d = tensor.dim %arg0, %c0 : tensor<?xi32, #SparseVector>
%xin = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector> %xin = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arg0: tensor<?xi32, #SparseVector>) ins(%arg0: tensor<?xi32, #SparseVector>)
outs(%xin: tensor<?xi32, #SparseVector>) { inits(%xin: tensor<?xi32, #SparseVector>) {
^bb0(%a: i32, %x: i32) : ^bb0(%a: i32, %x: i32) :
%result = math.absi %a : i32 %result = math.absi %a : i32
linalg.yield %result : i32 linalg.yield %result : i32
} -> tensor<?xi32, #SparseVector> } -> tensor<?xi32, #SparseVector>
return %0 : tensor<?xi32, #SparseVector> return %0 : tensor<?xi32, #SparseVector>
} }
// Driver method to call and verify sign kernel. // Driver method to call and verify sign kernel.
▲ Show 20 Lines • Show All 61 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_binary.mlir

Show First 20 Lines • Show All 44 Lines • ▼ Show 20 Linesmodule {
// When there is no overlap, include the present value in the output. // When there is no overlap, include the present value in the output.
func.func @vector_min(%arga: tensor<?xi32, #SparseVector>, func.func @vector_min(%arga: tensor<?xi32, #SparseVector>,
%argb: tensor<?xi32, #SparseVector>) -> tensor<?xi32, #SparseVector> { %argb: tensor<?xi32, #SparseVector>) -> tensor<?xi32, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xi32, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xi32, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
%0 = linalg.generic #trait_vec_op %0 = linalg.generic #trait_vec_op
ins(%arga, %argb: tensor<?xi32, #SparseVector>, tensor<?xi32, #SparseVector>) ins(%arga, %argb: tensor<?xi32, #SparseVector>, tensor<?xi32, #SparseVector>)
outs(%xv: tensor<?xi32, #SparseVector>) { inits(%xv: tensor<?xi32, #SparseVector>) {
^bb(%a: i32, %b: i32, %x: i32): ^bb(%a: i32, %b: i32, %x: i32):
%1 = sparse_tensor.binary %a, %b : i32, i32 to i32 %1 = sparse_tensor.binary %a, %b : i32, i32 to i32
overlap={ overlap={
^bb0(%a0: i32, %b0: i32): ^bb0(%a0: i32, %b0: i32):
%2 = arith.minsi %a0, %b0: i32 %2 = arith.minsi %a0, %b0: i32
sparse_tensor.yield %2 : i32 sparse_tensor.yield %2 : i32
} }
left=identity left=identity
right=identity right=identity
linalg.yield %1 : i32 linalg.yield %1 : i32
} -> tensor<?xi32, #SparseVector> } -> tensor<?xi32, #SparseVector>
return %0 : tensor<?xi32, #SparseVector> return %0 : tensor<?xi32, #SparseVector>
} }
// Creates a new sparse vector by multiplying a sparse vector with a dense vector. // Creates a new sparse vector by multiplying a sparse vector with a dense vector.
// When there is no overlap, leave the result empty. // When there is no overlap, leave the result empty.
func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>, func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64>) -> tensor<?xf64, #SparseVector> { %argb: tensor<?xf64>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_vec_op %0 = linalg.generic #trait_vec_op
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = sparse_tensor.binary %a, %b : f64, f64 to f64 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={ overlap={
^bb0(%a0: f64, %b0: f64): ^bb0(%a0: f64, %b0: f64):
%ret = arith.mulf %a0, %b0 : f64 %ret = arith.mulf %a0, %b0 : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
} }
left={} left={}
right={} right={}
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Take a set difference of two sparse vectors. The result will include only those // Take a set difference of two sparse vectors. The result will include only those
// sparse elements present in the first, but not the second vector. // sparse elements present in the first, but not the second vector.
func.func @vector_setdiff(%arga: tensor<?xf64, #SparseVector>, func.func @vector_setdiff(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_vec_op %0 = linalg.generic #trait_vec_op
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = sparse_tensor.binary %a, %b : f64, f64 to f64 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={} overlap={}
left=identity left=identity
right={} right={}
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Return the index of each entry // Return the index of each entry
func.func @vector_index(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> { func.func @vector_index(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
%0 = linalg.generic #trait_vec_scale %0 = linalg.generic #trait_vec_scale
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xi32, #SparseVector>) { inits(%xv: tensor<?xi32, #SparseVector>) {
^bb(%a: f64, %x: i32): ^bb(%a: f64, %x: i32):
%idx = linalg.index 0 : index %idx = linalg.index 0 : index
%1 = sparse_tensor.binary %a, %idx : f64, index to i32 %1 = sparse_tensor.binary %a, %idx : f64, index to i32
overlap={ overlap={
^bb0(%x0: f64, %i: index): ^bb0(%x0: f64, %i: index):
%ret = arith.index_cast %i : index to i32 %ret = arith.index_cast %i : index to i32
sparse_tensor.yield %ret : i32 sparse_tensor.yield %ret : i32
} }
Show All 10 Lines func.func @matrix_intersect(%arga: tensor<?x?xf64, #DCSR>,
%argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
%xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>)
outs(%xv: tensor<?x?xf64, #DCSR>) { inits(%xv: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = sparse_tensor.binary %a, %b: f64, f64 to f64 %1 = sparse_tensor.binary %a, %b: f64, f64 to f64
overlap={ overlap={
^bb0(%x0: f64, %y0: f64): ^bb0(%x0: f64, %y0: f64):
%ret = arith.addf %x0, %y0 : f64 %ret = arith.addf %x0, %y0 : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
} }
left={} left={}
Show All 9 Linesmodule {
// Tensor addition (use semi-ring binary operation). // Tensor addition (use semi-ring binary operation).
func.func @add_tensor_1(%A: tensor<4x4xf64, #DCSR>, func.func @add_tensor_1(%A: tensor<4x4xf64, #DCSR>,
%B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) { inits(%C: tensor<4x4xf64, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: f64) : ^bb0(%a: f64, %b: f64, %c: f64) :
%result = sparse_tensor.binary %a, %b : f64, f64 to f64 %result = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={ overlap={
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%ret = arith.addf %x, %y : f64 %ret = arith.addf %x, %y : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
} }
left=identity left=identity
right=identity right=identity
linalg.yield %result : f64 linalg.yield %result : f64
} -> tensor<4x4xf64, #DCSR> } -> tensor<4x4xf64, #DCSR>
return %0 : tensor<4x4xf64, #DCSR> return %0 : tensor<4x4xf64, #DCSR>
} }
// Same as @add_tensor_1, but use sparse_tensor.yield instead of identity to yield value. // Same as @add_tensor_1, but use sparse_tensor.yield instead of identity to yield value.
func.func @add_tensor_2(%A: tensor<4x4xf64, #DCSR>, func.func @add_tensor_2(%A: tensor<4x4xf64, #DCSR>,
%B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) { inits(%C: tensor<4x4xf64, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: f64) : ^bb0(%a: f64, %b: f64, %c: f64) :
%result = sparse_tensor.binary %a, %b : f64, f64 to f64 %result = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={ overlap={
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%ret = arith.addf %x, %y : f64 %ret = arith.addf %x, %y : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
} }
left={ left={
Show All 11 Linesmodule {
// Performs triangular add/sub operation (using semi-ring binary op). // Performs triangular add/sub operation (using semi-ring binary op).
func.func @triangular(%A: tensor<4x4xf64, #DCSR>, func.func @triangular(%A: tensor<4x4xf64, #DCSR>,
%B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) { inits(%C: tensor<4x4xf64, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: f64) : ^bb0(%a: f64, %b: f64, %c: f64) :
%row = linalg.index 0 : index %row = linalg.index 0 : index
%col = linalg.index 1 : index %col = linalg.index 1 : index
%result = sparse_tensor.binary %a, %b : f64, f64 to f64 %result = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={ overlap={
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%cmp = arith.cmpi "uge", %col, %row : index %cmp = arith.cmpi "uge", %col, %row : index
%upperTriangleResult = arith.addf %x, %y : f64 %upperTriangleResult = arith.addf %x, %y : f64
Show All 20 Lines func.func @sub_with_thres(%A: tensor<4x4xf64, #DCSR>,
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
// Defines out-block constant bounds. // Defines out-block constant bounds.
%thres_out_up = arith.constant 2.0 : f64 %thres_out_up = arith.constant 2.0 : f64
%thres_out_lo = arith.constant -2.0 : f64 %thres_out_lo = arith.constant -2.0 : f64
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) { inits(%C: tensor<4x4xf64, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: f64) : ^bb0(%a: f64, %b: f64, %c: f64) :
%result = sparse_tensor.binary %a, %b : f64, f64 to f64 %result = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={ overlap={
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
// Defines in-block constant bounds. // Defines in-block constant bounds.
%thres_up = arith.constant 1.0 : f64 %thres_up = arith.constant 1.0 : f64
%thres_lo = arith.constant -1.0 : f64 %thres_lo = arith.constant -1.0 : f64
%result = arith.subf %x, %y : f64 %result = arith.subf %x, %y : f64
Show All 28 Linesmodule {
// Performs isEqual only on intersecting elements. // Performs isEqual only on intersecting elements.
func.func @intersect_equal(%A: tensor<4x4xf64, #DCSR>, func.func @intersect_equal(%A: tensor<4x4xf64, #DCSR>,
%B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xi8, #DCSR> { %B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xi8, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xi8, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xi8, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xi8, #DCSR>) { inits(%C: tensor<4x4xi8, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: i8) : ^bb0(%a: f64, %b: f64, %c: i8) :
%result = sparse_tensor.binary %a, %b : f64, f64 to i8 %result = sparse_tensor.binary %a, %b : f64, f64 to i8
overlap={ overlap={
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%cmp = arith.cmpf "oeq", %x, %y : f64 %cmp = arith.cmpf "oeq", %x, %y : f64
%ret = arith.extui %cmp : i1 to i8 %ret = arith.extui %cmp : i1 to i8
sparse_tensor.yield %ret : i8 sparse_tensor.yield %ret : i8
} }
left={} left={}
right={} right={}
linalg.yield %result : i8 linalg.yield %result : i8
} -> tensor<4x4xi8, #DCSR> } -> tensor<4x4xi8, #DCSR>
return %0 : tensor<4x4xi8, #DCSR> return %0 : tensor<4x4xi8, #DCSR>
} }
// Keeps values on left, negate value on right, ignore value when overlapping. // Keeps values on left, negate value on right, ignore value when overlapping.
func.func @only_left_right(%A: tensor<4x4xf64, #DCSR>, func.func @only_left_right(%A: tensor<4x4xf64, #DCSR>,
%B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%0 = linalg.generic #trait_mat_op %0 = linalg.generic #trait_mat_op
ins(%A, %B: tensor<4x4xf64, #DCSR>, ins(%A, %B: tensor<4x4xf64, #DCSR>,
tensor<4x4xf64, #DCSR>) tensor<4x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) { inits(%C: tensor<4x4xf64, #DCSR>) {
^bb0(%a: f64, %b: f64, %c: f64) : ^bb0(%a: f64, %b: f64, %c: f64) :
%result = sparse_tensor.binary %a, %b : f64, f64 to f64 %result = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap={} overlap={}
left=identity left=identity
right={ right={
^bb0(%y: f64): ^bb0(%y: f64):
%ret = arith.negf %y : f64 %ret = arith.negf %y : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
▲ Show 20 Lines • Show All 224 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

Show All 35 Linesmodule {
// bitcast // bitcast
// Since all casts are "zero preserving" unary operations, lattice computation // Since all casts are "zero preserving" unary operations, lattice computation
// and conversion to sparse code is straightforward. // and conversion to sparse code is straightforward.
// //
func.func @sparse_cast_s32_to_f32(%arga: tensor<10xi32, #SV>, func.func @sparse_cast_s32_to_f32(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xf32>) -> tensor<10xf32> { %argb: tensor<10xf32>) -> tensor<10xf32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: i32, %x : f32): ^bb(%a: i32, %x : f32):
%cst = arith.sitofp %a : i32 to f32 %cst = arith.sitofp %a : i32 to f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
func.func @sparse_cast_u32_to_f32(%arga: tensor<10xi32, #SV>, func.func @sparse_cast_u32_to_f32(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xf32>) -> tensor<10xf32> { %argb: tensor<10xf32>) -> tensor<10xf32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: i32, %x : f32): ^bb(%a: i32, %x : f32):
%cst = arith.uitofp %a : i32 to f32 %cst = arith.uitofp %a : i32 to f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
func.func @sparse_cast_f32_to_s32(%arga: tensor<10xf32, #SV>, func.func @sparse_cast_f32_to_s32(%arga: tensor<10xf32, #SV>,
%argb: tensor<10xi32>) -> tensor<10xi32> { %argb: tensor<10xi32>) -> tensor<10xi32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xf32, #SV>) ins(%arga: tensor<10xf32, #SV>)
outs(%argb: tensor<10xi32>) { inits(%argb: tensor<10xi32>) {
^bb(%a: f32, %x : i32): ^bb(%a: f32, %x : i32):
%cst = arith.fptosi %a : f32 to i32 %cst = arith.fptosi %a : f32 to i32
linalg.yield %cst : i32 linalg.yield %cst : i32
} -> tensor<10xi32> } -> tensor<10xi32>
return %0 : tensor<10xi32> return %0 : tensor<10xi32>
} }
func.func @sparse_cast_f64_to_u32(%arga: tensor<10xf64, #SV>, func.func @sparse_cast_f64_to_u32(%arga: tensor<10xf64, #SV>,
%argb: tensor<10xi32>) -> tensor<10xi32> { %argb: tensor<10xi32>) -> tensor<10xi32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xf64, #SV>) ins(%arga: tensor<10xf64, #SV>)
outs(%argb: tensor<10xi32>) { inits(%argb: tensor<10xi32>) {
^bb(%a: f64, %x : i32): ^bb(%a: f64, %x : i32):
%cst = arith.fptoui %a : f64 to i32 %cst = arith.fptoui %a : f64 to i32
linalg.yield %cst : i32 linalg.yield %cst : i32
} -> tensor<10xi32> } -> tensor<10xi32>
return %0 : tensor<10xi32> return %0 : tensor<10xi32>
} }
func.func @sparse_cast_f32_to_f64(%arga: tensor<10xf32, #SV>, func.func @sparse_cast_f32_to_f64(%arga: tensor<10xf32, #SV>,
%argb: tensor<10xf64>) -> tensor<10xf64> { %argb: tensor<10xf64>) -> tensor<10xf64> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xf32, #SV>) ins(%arga: tensor<10xf32, #SV>)
outs(%argb: tensor<10xf64>) { inits(%argb: tensor<10xf64>) {
^bb(%a: f32, %x : f64): ^bb(%a: f32, %x : f64):
%cst = arith.extf %a : f32 to f64 %cst = arith.extf %a : f32 to f64
linalg.yield %cst : f64 linalg.yield %cst : f64
} -> tensor<10xf64> } -> tensor<10xf64>
return %0 : tensor<10xf64> return %0 : tensor<10xf64>
} }
func.func @sparse_cast_f64_to_f32(%arga: tensor<10xf64, #SV>, func.func @sparse_cast_f64_to_f32(%arga: tensor<10xf64, #SV>,
%argb: tensor<10xf32>) -> tensor<10xf32> { %argb: tensor<10xf32>) -> tensor<10xf32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xf64, #SV>) ins(%arga: tensor<10xf64, #SV>)
outs(%argb: tensor<10xf32>) { inits(%argb: tensor<10xf32>) {
^bb(%a: f64, %x : f32): ^bb(%a: f64, %x : f32):
%cst = arith.truncf %a : f64 to f32 %cst = arith.truncf %a : f64 to f32
linalg.yield %cst : f32 linalg.yield %cst : f32
} -> tensor<10xf32> } -> tensor<10xf32>
return %0 : tensor<10xf32> return %0 : tensor<10xf32>
} }
func.func @sparse_cast_s32_to_u64(%arga: tensor<10xi32, #SV>, func.func @sparse_cast_s32_to_u64(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xi64>) -> tensor<10xi64> { %argb: tensor<10xi64>) -> tensor<10xi64> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xi64>) { inits(%argb: tensor<10xi64>) {
^bb(%a: i32, %x : i64): ^bb(%a: i32, %x : i64):
%cst = arith.extsi %a : i32 to i64 %cst = arith.extsi %a : i32 to i64
linalg.yield %cst : i64 linalg.yield %cst : i64
} -> tensor<10xi64> } -> tensor<10xi64>
return %0 : tensor<10xi64> return %0 : tensor<10xi64>
} }
func.func @sparse_cast_u32_to_s64(%arga: tensor<10xi32, #SV>, func.func @sparse_cast_u32_to_s64(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xi64>) -> tensor<10xi64> { %argb: tensor<10xi64>) -> tensor<10xi64> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xi64>) { inits(%argb: tensor<10xi64>) {
^bb(%a: i32, %x : i64): ^bb(%a: i32, %x : i64):
%cst = arith.extui %a : i32 to i64 %cst = arith.extui %a : i32 to i64
linalg.yield %cst : i64 linalg.yield %cst : i64
} -> tensor<10xi64> } -> tensor<10xi64>
return %0 : tensor<10xi64> return %0 : tensor<10xi64>
} }
func.func @sparse_cast_i32_to_i8(%arga: tensor<10xi32, #SV>, func.func @sparse_cast_i32_to_i8(%arga: tensor<10xi32, #SV>,
%argb: tensor<10xi8>) -> tensor<10xi8> { %argb: tensor<10xi8>) -> tensor<10xi8> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xi32, #SV>) ins(%arga: tensor<10xi32, #SV>)
outs(%argb: tensor<10xi8>) { inits(%argb: tensor<10xi8>) {
^bb(%a: i32, %x : i8): ^bb(%a: i32, %x : i8):
%cst = arith.trunci %a : i32 to i8 %cst = arith.trunci %a : i32 to i8
linalg.yield %cst : i8 linalg.yield %cst : i8
} -> tensor<10xi8> } -> tensor<10xi8>
return %0 : tensor<10xi8> return %0 : tensor<10xi8>
} }
func.func @sparse_cast_f32_as_s32(%arga: tensor<10xf32, #SV>, func.func @sparse_cast_f32_as_s32(%arga: tensor<10xf32, #SV>,
%argb: tensor<10xi32>) -> tensor<10xi32> { %argb: tensor<10xi32>) -> tensor<10xi32> {
%0 = linalg.generic #trait_cast %0 = linalg.generic #trait_cast
ins(%arga: tensor<10xf32, #SV>) ins(%arga: tensor<10xf32, #SV>)
outs(%argb: tensor<10xi32>) { inits(%argb: tensor<10xi32>) {
^bb(%a: f32, %x : i32): ^bb(%a: f32, %x : i32):
%cst = arith.bitcast %a : f32 to i32 %cst = arith.bitcast %a : f32 to i32
linalg.yield %cst : i32 linalg.yield %cst : i32
} -> tensor<10xi32> } -> tensor<10xi32>
return %0 : tensor<10xi32> return %0 : tensor<10xi32>
} }
// //
▲ Show 20 Lines • Show All 108 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex32.mlir

Show All 20 Lines func.func @cadd(%arga: tensor<?xcomplex<f32>, #SparseVector>,
%argb: tensor<?xcomplex<f32>, #SparseVector>) %argb: tensor<?xcomplex<f32>, #SparseVector>)
-> tensor<?xcomplex<f32>, #SparseVector> { -> tensor<?xcomplex<f32>, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f32>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f32>, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>, ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>,
tensor<?xcomplex<f32>, #SparseVector>) tensor<?xcomplex<f32>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f32>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f32>, #SparseVector>) {
^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>): ^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>):
%1 = complex.add %a, %b : complex<f32> %1 = complex.add %a, %b : complex<f32>
linalg.yield %1 : complex<f32> linalg.yield %1 : complex<f32>
} -> tensor<?xcomplex<f32>, #SparseVector> } -> tensor<?xcomplex<f32>, #SparseVector>
return %0 : tensor<?xcomplex<f32>, #SparseVector> return %0 : tensor<?xcomplex<f32>, #SparseVector>
} }
func.func @cmul(%arga: tensor<?xcomplex<f32>, #SparseVector>, func.func @cmul(%arga: tensor<?xcomplex<f32>, #SparseVector>,
%argb: tensor<?xcomplex<f32>, #SparseVector>) %argb: tensor<?xcomplex<f32>, #SparseVector>)
-> tensor<?xcomplex<f32>, #SparseVector> { -> tensor<?xcomplex<f32>, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f32>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f32>, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>, ins(%arga, %argb: tensor<?xcomplex<f32>, #SparseVector>,
tensor<?xcomplex<f32>, #SparseVector>) tensor<?xcomplex<f32>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f32>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f32>, #SparseVector>) {
^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>): ^bb(%a: complex<f32>, %b: complex<f32>, %x: complex<f32>):
%1 = complex.mul %a, %b : complex<f32> %1 = complex.mul %a, %b : complex<f32>
linalg.yield %1 : complex<f32> linalg.yield %1 : complex<f32>
} -> tensor<?xcomplex<f32>, #SparseVector> } -> tensor<?xcomplex<f32>, #SparseVector>
return %0 : tensor<?xcomplex<f32>, #SparseVector> return %0 : tensor<?xcomplex<f32>, #SparseVector>
} }
func.func @dump(%arg0: tensor<?xcomplex<f32>, #SparseVector>, %d: index) { func.func @dump(%arg0: tensor<?xcomplex<f32>, #SparseVector>, %d: index) {
▲ Show 20 Lines • Show All 62 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex64.mlir

Show All 20 Lines func.func @cadd(%arga: tensor<?xcomplex<f64>, #SparseVector>,
%argb: tensor<?xcomplex<f64>, #SparseVector>) %argb: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>, ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
tensor<?xcomplex<f64>, #SparseVector>) tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
%1 = complex.add %a, %b : complex<f64> %1 = complex.add %a, %b : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @cmul(%arga: tensor<?xcomplex<f64>, #SparseVector>, func.func @cmul(%arga: tensor<?xcomplex<f64>, #SparseVector>,
%argb: tensor<?xcomplex<f64>, #SparseVector>) %argb: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>, ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
tensor<?xcomplex<f64>, #SparseVector>) tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
%1 = complex.mul %a, %b : complex<f64> %1 = complex.mul %a, %b : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @dump(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) { func.func @dump(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) {
▲ Show 20 Lines • Show All 62 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_complex_ops.mlir

Show All 29 Lines func.func @cops(%arga: tensor<?xcomplex<f64>, #SparseVector>,
%argb: tensor<?xcomplex<f64>, #SparseVector>) %argb: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op2 %0 = linalg.generic #trait_op2
ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>, ins(%arga, %argb: tensor<?xcomplex<f64>, #SparseVector>,
tensor<?xcomplex<f64>, #SparseVector>) tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %b: complex<f64>, %x: complex<f64>):
%1 = complex.neg %b : complex<f64> %1 = complex.neg %b : complex<f64>
%2 = complex.sub %a, %1 : complex<f64> %2 = complex.sub %a, %1 : complex<f64>
linalg.yield %2 : complex<f64> linalg.yield %2 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @csin(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @csin(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%1 = complex.sin %a : complex<f64> %1 = complex.sin %a : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @complex_sqrt(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @complex_sqrt(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%1 = complex.sqrt %a : complex<f64> %1 = complex.sqrt %a : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @complex_tanh(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @complex_tanh(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%1 = complex.tanh %a : complex<f64> %1 = complex.tanh %a : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @clog1p_expm1(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @clog1p_expm1(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%1 = complex.log1p %a : complex<f64> %1 = complex.log1p %a : complex<f64>
%2 = complex.expm1 %1 : complex<f64> %2 = complex.expm1 %1 : complex<f64>
linalg.yield %2 : complex<f64> linalg.yield %2 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @cdiv(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @cdiv(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xcomplex<f64>, #SparseVector> { -> tensor<?xcomplex<f64>, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xcomplex<f64>, #SparseVector>
%c = complex.constant [2.0 : f64, 0.0 : f64] : complex<f64> %c = complex.constant [2.0 : f64, 0.0 : f64] : complex<f64>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xcomplex<f64>, #SparseVector>) { inits(%xv: tensor<?xcomplex<f64>, #SparseVector>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%1 = complex.div %a, %c : complex<f64> %1 = complex.div %a, %c : complex<f64>
linalg.yield %1 : complex<f64> linalg.yield %1 : complex<f64>
} -> tensor<?xcomplex<f64>, #SparseVector> } -> tensor<?xcomplex<f64>, #SparseVector>
return %0 : tensor<?xcomplex<f64>, #SparseVector> return %0 : tensor<?xcomplex<f64>, #SparseVector>
} }
func.func @cabs(%arga: tensor<?xcomplex<f64>, #SparseVector>) func.func @cabs(%arga: tensor<?xcomplex<f64>, #SparseVector>)
-> tensor<?xf64, #SparseVector> { -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector> %d = tensor.dim %arga, %c0 : tensor<?xcomplex<f64>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_op1 %0 = linalg.generic #trait_op1
ins(%arga: tensor<?xcomplex<f64>, #SparseVector>) ins(%arga: tensor<?xcomplex<f64>, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: complex<f64>, %x: f64): ^bb(%a: complex<f64>, %x: f64):
%1 = complex.abs %a : complex<f64> %1 = complex.abs %a : complex<f64>
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
func.func @dumpc(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) { func.func @dumpc(%arg0: tensor<?xcomplex<f64>, #SparseVector>, %d: index) {
▲ Show 20 Lines • Show All 116 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_dot.mlir

Show All 9 Linesmodule {
// //
// Sparse kernel. // Sparse kernel.
// //
func.func @sparse_dot(%a: tensor<1024xf32, #SparseVector>, func.func @sparse_dot(%a: tensor<1024xf32, #SparseVector>,
%b: tensor<1024xf32, #SparseVector>, %b: tensor<1024xf32, #SparseVector>,
%x: tensor<f32>) -> tensor<f32> { %x: tensor<f32>) -> tensor<f32> {
%dot = linalg.dot ins(%a, %b: tensor<1024xf32, #SparseVector>, %dot = linalg.dot ins(%a, %b: tensor<1024xf32, #SparseVector>,
tensor<1024xf32, #SparseVector>) tensor<1024xf32, #SparseVector>)
outs(%x: tensor<f32>) -> tensor<f32> inits(%x: tensor<f32>) -> tensor<f32>
return %dot : tensor<f32> return %dot : tensor<f32>
} }
// //
// Main driver. // Main driver.
// //
func.func @entry() { func.func @entry() {
// Setup two sparse vectors. // Setup two sparse vectors.
Show All 39 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_expand.mlir

Show All 13 Linesmodule {
// so that access pattern expansion (workspace) needs to be // so that access pattern expansion (workspace) needs to be
// done along dimension with size 8. // done along dimension with size 8.
// //
func.func @matmul(%A: tensor<8x2xf64, #CSC>, func.func @matmul(%A: tensor<8x2xf64, #CSC>,
%B: tensor<2x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> { %B: tensor<2x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> {
%C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSC> %C = bufferization.alloc_tensor() : tensor<8x4xf64, #CSC>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<8x2xf64, #CSC>, tensor<2x4xf64, #CSC>) ins(%A, %B: tensor<8x2xf64, #CSC>, tensor<2x4xf64, #CSC>)
outs(%C: tensor<8x4xf64, #CSC>) -> tensor<8x4xf64, #CSC> inits(%C: tensor<8x4xf64, #CSC>) -> tensor<8x4xf64, #CSC>
return %D: tensor<8x4xf64, #CSC> return %D: tensor<8x4xf64, #CSC>
} }
// //
// Main driver. // Main driver.
// //
func.func @entry() { func.func @entry() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
▲ Show 20 Lines • Show All 49 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

// RUN: mlir-opt %s --sparse-compiler | \ // RUN: mlir-opt %s --sparse-compiler | \
// RUN: mlir-cpu-runner -e entry -entry-point-result=void \ // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \ // RUN: -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s // RUN: FileCheck %s
#DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }> #DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
// An example of a 2D convolution with a sparse filter. // An example of a 2D convolution with a sparse filter.
module { module {
func.func @conv2d(%input: tensor<8x8xi32>, func.func @conv2d(%input: tensor<8x8xi32>,
%filter: tensor<3x3xi32, #DCSR>, %filter: tensor<3x3xi32, #DCSR>,
%output: tensor<6x6xi32>) -> tensor<6x6xi32> { %output: tensor<6x6xi32>) -> tensor<6x6xi32> {
%0 = linalg.conv_2d %0 = linalg.conv_2d
ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>) ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>)
outs (%output: tensor<6x6xi32>) -> tensor<6x6xi32> inits (%output: tensor<6x6xi32>) -> tensor<6x6xi32>
return %0 : tensor<6x6xi32> return %0 : tensor<6x6xi32>
} }
func.func @conv2d_sparse_out(%input: tensor<8x8xi32>, func.func @conv2d_sparse_out(%input: tensor<8x8xi32>,
%filter: tensor<3x3xi32, #DCSR>) -> tensor<6x6xi32, #DCSR> { %filter: tensor<3x3xi32, #DCSR>) -> tensor<6x6xi32, #DCSR> {
%s = bufferization.alloc_tensor() : tensor<6x6xi32, #DCSR> %s = bufferization.alloc_tensor() : tensor<6x6xi32, #DCSR>
%0 = linalg.conv_2d %0 = linalg.conv_2d
ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>) ins (%input, %filter: tensor<8x8xi32>, tensor<3x3xi32, #DCSR>)
outs (%s: tensor<6x6xi32, #DCSR>) -> tensor<6x6xi32, #DCSR> inits (%s: tensor<6x6xi32, #DCSR>) -> tensor<6x6xi32, #DCSR>
return %0 : tensor<6x6xi32, #DCSR> return %0 : tensor<6x6xi32, #DCSR>
} }
func.func @entry() { func.func @entry() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%i0 = arith.constant 0 : i32 %i0 = arith.constant 0 : i32
// A typical edge detection filter. // A typical edge detection filter.
▲ Show 20 Lines • Show All 62 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

Show All 34 Linesmodule {
// //
// A kernel that flattens a rank 8 tensor into a dense matrix. // A kernel that flattens a rank 8 tensor into a dense matrix.
// //
func.func @kernel_flatten(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>, func.func @kernel_flatten(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>,
%argx: tensor<7x3xf64>) %argx: tensor<7x3xf64>)
-> tensor<7x3xf64> { -> tensor<7x3xf64> {
%0 = linalg.generic #trait_flatten %0 = linalg.generic #trait_flatten
ins(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>) ins(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>)
outs(%argx: tensor<7x3xf64>) { inits(%argx: tensor<7x3xf64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.addf %x, %a : f64 %0 = arith.addf %x, %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<7x3xf64> } -> tensor<7x3xf64>
return %0 : tensor<7x3xf64> return %0 : tensor<7x3xf64>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index.mlir

Show All 37 Linesmodule {
// //
// Kernel that uses index in the index notation (conjunction). // Kernel that uses index in the index notation (conjunction).
// //
func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>) func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>)
-> tensor<8xi64, #SparseVector> { -> tensor<8xi64, #SparseVector> {
%init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector> %init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector>
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%init: tensor<8xi64, #SparseVector>) { inits(%init: tensor<8xi64, #SparseVector>) {
^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, #SparseVector> } -> tensor<8xi64, #SparseVector>
return %r : tensor<8xi64, #SparseVector> return %r : tensor<8xi64, #SparseVector>
} }
// //
// Kernel that uses index in the index notation (disjunction). // Kernel that uses index in the index notation (disjunction).
// //
func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>) func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>)
-> tensor<8xi64, #SparseVector> { -> tensor<8xi64, #SparseVector> {
%init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector> %init = bufferization.alloc_tensor() : tensor<8xi64, #SparseVector>
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%init: tensor<8xi64, #SparseVector>) { inits(%init: tensor<8xi64, #SparseVector>) {
^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, #SparseVector> } -> tensor<8xi64, #SparseVector>
return %r : tensor<8xi64, #SparseVector> return %r : tensor<8xi64, #SparseVector>
} }
// //
// Kernel that uses indices in the index notation (conjunction). // Kernel that uses indices in the index notation (conjunction).
// //
func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>) func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>)
-> tensor<3x4xi64, #SparseMatrix> { -> tensor<3x4xi64, #SparseMatrix> {
%init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix> %init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix>
%r = linalg.generic #trait_2d %r = linalg.generic #trait_2d
ins(%arga: tensor<3x4xi64, #SparseMatrix>) ins(%arga: tensor<3x4xi64, #SparseMatrix>)
outs(%init: tensor<3x4xi64, #SparseMatrix>) { inits(%init: tensor<3x4xi64, #SparseMatrix>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.muli %ii, %a : i64 %m1 = arith.muli %ii, %a : i64
%m2 = arith.muli %jj, %m1 : i64 %m2 = arith.muli %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
} -> tensor<3x4xi64, #SparseMatrix> } -> tensor<3x4xi64, #SparseMatrix>
return %r : tensor<3x4xi64, #SparseMatrix> return %r : tensor<3x4xi64, #SparseMatrix>
} }
// //
// Kernel that uses indices in the index notation (disjunction). // Kernel that uses indices in the index notation (disjunction).
// //
func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>) func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>)
-> tensor<3x4xi64, #SparseMatrix> { -> tensor<3x4xi64, #SparseMatrix> {
%init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix> %init = bufferization.alloc_tensor() : tensor<3x4xi64, #SparseMatrix>
%r = linalg.generic #trait_2d %r = linalg.generic #trait_2d
ins(%arga: tensor<3x4xi64, #SparseMatrix>) ins(%arga: tensor<3x4xi64, #SparseMatrix>)
outs(%init: tensor<3x4xi64, #SparseMatrix>) { inits(%init: tensor<3x4xi64, #SparseMatrix>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.addi %ii, %a : i64 %m1 = arith.addi %ii, %a : i64
%m2 = arith.addi %jj, %m1 : i64 %m2 = arith.addi %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
} -> tensor<3x4xi64, #SparseMatrix> } -> tensor<3x4xi64, #SparseMatrix>
return %r : tensor<3x4xi64, #SparseMatrix> return %r : tensor<3x4xi64, #SparseMatrix>
} }
func.func @add_outer_2d(%arg0: tensor<2x3xf32, #SparseMatrix>) func.func @add_outer_2d(%arg0: tensor<2x3xf32, #SparseMatrix>)
-> tensor<2x3xf32, #SparseMatrix> { -> tensor<2x3xf32, #SparseMatrix> {
%0 = bufferization.alloc_tensor() : tensor<2x3xf32, #SparseMatrix> %0 = bufferization.alloc_tensor() : tensor<2x3xf32, #SparseMatrix>
%1 = linalg.generic #trait_2d %1 = linalg.generic #trait_2d
ins(%arg0 : tensor<2x3xf32, #SparseMatrix>) ins(%arg0 : tensor<2x3xf32, #SparseMatrix>)
outs(%0 : tensor<2x3xf32, #SparseMatrix>) { inits(%0 : tensor<2x3xf32, #SparseMatrix>) {
^bb0(%arg1: f32, %arg2: f32): ^bb0(%arg1: f32, %arg2: f32):
%2 = linalg.index 0 : index %2 = linalg.index 0 : index
%3 = arith.index_cast %2 : index to i64 %3 = arith.index_cast %2 : index to i64
%4 = arith.uitofp %3 : i64 to f32 %4 = arith.uitofp %3 : i64 to f32
%5 = arith.addf %arg1, %4 : f32 %5 = arith.addf %arg1, %4 : f32
linalg.yield %5 : f32 linalg.yield %5 : f32
} -> tensor<2x3xf32, #SparseMatrix> } -> tensor<2x3xf32, #SparseMatrix>
return %1 : tensor<2x3xf32, #SparseMatrix> return %1 : tensor<2x3xf32, #SparseMatrix>
▲ Show 20 Lines • Show All 117 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir

Show All 35 Linesmodule {
// //
// Kernel that uses index in the index notation (conjunction). // Kernel that uses index in the index notation (conjunction).
// //
func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>, func.func @sparse_index_1d_conj(%arga: tensor<8xi64, #SparseVector>,
%out: tensor<8xi64>) -> tensor<8xi64> { %out: tensor<8xi64>) -> tensor<8xi64> {
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%out: tensor<8xi64>) { inits(%out: 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>
} }
// //
// Kernel that uses index in the index notation (disjunction). // Kernel that uses index in the index notation (disjunction).
// //
func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>, func.func @sparse_index_1d_disj(%arga: tensor<8xi64, #SparseVector>,
%out: tensor<8xi64>) -> tensor<8xi64> { %out: tensor<8xi64>) -> tensor<8xi64> {
%r = linalg.generic #trait_1d %r = linalg.generic #trait_1d
ins(%arga: tensor<8xi64, #SparseVector>) ins(%arga: tensor<8xi64, #SparseVector>)
outs(%out: tensor<8xi64>) { inits(%out: 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>
} }
// //
// Kernel that uses indices in the index notation (conjunction). // Kernel that uses indices in the index notation (conjunction).
// //
func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>, func.func @sparse_index_2d_conj(%arga: tensor<3x4xi64, #SparseMatrix>,
%out: tensor<3x4xi64>) -> tensor<3x4xi64> { %out: tensor<3x4xi64>) -> tensor<3x4xi64> {
%r = linalg.generic #trait_2d %r = linalg.generic #trait_2d
ins(%arga: tensor<3x4xi64, #SparseMatrix>) ins(%arga: tensor<3x4xi64, #SparseMatrix>)
outs(%out: tensor<3x4xi64>) { inits(%out: tensor<3x4xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.muli %ii, %a : i64 %m1 = arith.muli %ii, %a : i64
%m2 = arith.muli %jj, %m1 : i64 %m2 = arith.muli %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
} -> tensor<3x4xi64> } -> tensor<3x4xi64>
return %r : tensor<3x4xi64> return %r : tensor<3x4xi64>
} }
// //
// Kernel that uses indices in the index notation (disjunction). // Kernel that uses indices in the index notation (disjunction).
// //
func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>, func.func @sparse_index_2d_disj(%arga: tensor<3x4xi64, #SparseMatrix>,
%out: tensor<3x4xi64>) -> tensor<3x4xi64> { %out: tensor<3x4xi64>) -> tensor<3x4xi64> {
%r = linalg.generic #trait_2d %r = linalg.generic #trait_2d
ins(%arga: tensor<3x4xi64, #SparseMatrix>) ins(%arga: tensor<3x4xi64, #SparseMatrix>)
outs(%out: tensor<3x4xi64>) { inits(%out: tensor<3x4xi64>) {
^bb(%a: i64, %x: i64): ^bb(%a: i64, %x: i64):
%i = linalg.index 0 : index %i = linalg.index 0 : index
%j = linalg.index 1 : index %j = linalg.index 1 : index
%ii = arith.index_cast %i : index to i64 %ii = arith.index_cast %i : index to i64
%jj = arith.index_cast %j : index to i64 %jj = arith.index_cast %j : index to i64
%m1 = arith.addi %ii, %a : i64 %m1 = arith.addi %ii, %a : i64
%m2 = arith.addi %jj, %m1 : i64 %m2 = arith.addi %jj, %m1 : i64
linalg.yield %m2 : i64 linalg.yield %m2 : i64
▲ Show 20 Lines • Show All 81 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir

Show All 15 Lines
module { module {
// //
// Computes C = A x B with all matrices dense. // Computes C = A x B with all matrices dense.
// //
func.func @matmul1(%A: tensor<4x8xf64>, %B: tensor<8x4xf64>, func.func @matmul1(%A: tensor<4x8xf64>, %B: tensor<8x4xf64>,
%C: tensor<4x4xf64>) -> tensor<4x4xf64> { %C: tensor<4x4xf64>) -> tensor<4x4xf64> {
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<4x8xf64>, tensor<8x4xf64>) ins(%A, %B: tensor<4x8xf64>, tensor<8x4xf64>)
outs(%C: tensor<4x4xf64>) -> tensor<4x4xf64> inits(%C: tensor<4x4xf64>) -> tensor<4x4xf64>
return %D: tensor<4x4xf64> return %D: tensor<4x4xf64>
} }
// //
// Computes C = A x B with all matrices sparse (SpMSpM) in CSR. // Computes C = A x B with all matrices sparse (SpMSpM) in CSR.
// //
func.func @matmul2(%A: tensor<4x8xf64, #CSR>, func.func @matmul2(%A: tensor<4x8xf64, #CSR>,
%B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> { %B: tensor<8x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #CSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #CSR>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>) ins(%A, %B: tensor<4x8xf64, #CSR>, tensor<8x4xf64, #CSR>)
outs(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR> inits(%C: tensor<4x4xf64, #CSR>) -> tensor<4x4xf64, #CSR>
return %D: tensor<4x4xf64, #CSR> return %D: tensor<4x4xf64, #CSR>
} }
// //
// Computes C = A x B with all matrices sparse (SpMSpM) in DCSR. // Computes C = A x B with all matrices sparse (SpMSpM) in DCSR.
// //
func.func @matmul3(%A: tensor<4x8xf64, #DCSR>, func.func @matmul3(%A: tensor<4x8xf64, #DCSR>,
%B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> { %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
%C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR> %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
%D = linalg.matmul %D = linalg.matmul
ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>) ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>)
outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> inits(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>
return %D: tensor<4x4xf64, #DCSR> return %D: tensor<4x4xf64, #DCSR>
} }
// //
// Main driver. // Main driver.
// //
func.func @entry() { func.func @entry() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
▲ Show 20 Lines • Show All 188 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir

Show All 39 Lines func.func @matrix_scale(%arga: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%s = arith.constant 2.0 : f64 %s = arith.constant 2.0 : f64
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
%xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
ins(%arga: tensor<?x?xf64, #DCSR>) ins(%arga: tensor<?x?xf64, #DCSR>)
outs(%xm: tensor<?x?xf64, #DCSR>) { inits(%xm: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = arith.mulf %a, %s : f64 %1 = arith.mulf %a, %s : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64, #DCSR> } -> tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR> return %0 : tensor<?x?xf64, #DCSR>
} }
// Scales a sparse matrix in place. // Scales a sparse matrix in place.
func.func @matrix_scale_inplace(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { func.func @matrix_scale_inplace(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%s = arith.constant 2.0 : f64 %s = arith.constant 2.0 : f64
%0 = linalg.generic #trait_scale_inpl %0 = linalg.generic #trait_scale_inpl
outs(%argx: tensor<?x?xf64, #DCSR>) { inits(%argx: tensor<?x?xf64, #DCSR>) {
^bb(%x: f64): ^bb(%x: f64):
%1 = arith.mulf %x, %s : f64 %1 = arith.mulf %x, %s : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64, #DCSR> } -> tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR> return %0 : tensor<?x?xf64, #DCSR>
} }
// Adds two sparse matrices element-wise into a new sparse matrix. // Adds two sparse matrices element-wise into a new sparse matrix.
func.func @matrix_add(%arga: tensor<?x?xf64, #DCSR>, func.func @matrix_add(%arga: tensor<?x?xf64, #DCSR>,
%argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
%xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>)
outs(%xv: tensor<?x?xf64, #DCSR>) { inits(%xv: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.addf %a, %b : f64 %1 = arith.addf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64, #DCSR> } -> tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR> return %0 : tensor<?x?xf64, #DCSR>
} }
// Multiplies two sparse matrices element-wise into a new sparse matrix. // Multiplies two sparse matrices element-wise into a new sparse matrix.
func.func @matrix_mul(%arga: tensor<?x?xf64, #DCSR>, func.func @matrix_mul(%arga: tensor<?x?xf64, #DCSR>,
%argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { %argb: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #DCSR>
%xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>) ins(%arga, %argb: tensor<?x?xf64, #DCSR>, tensor<?x?xf64, #DCSR>)
outs(%xv: tensor<?x?xf64, #DCSR>) { inits(%xv: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.mulf %a, %b : f64 %1 = arith.mulf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64, #DCSR> } -> tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR> return %0 : tensor<?x?xf64, #DCSR>
} }
// Dump a sparse matrix. // Dump a sparse matrix.
▲ Show 20 Lines • Show All 65 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

Show All 33 Linesmodule {
// into a dense vector x. // into a dense vector x.
// //
func.func @kernel_matvec(%arga: tensor<?x?xi32, #SparseMatrix>, func.func @kernel_matvec(%arga: tensor<?x?xi32, #SparseMatrix>,
%argb: tensor<?xi32>, %argb: tensor<?xi32>,
%argx: tensor<?xi32>) %argx: tensor<?xi32>)
-> tensor<?xi32> { -> tensor<?xi32> {
%0 = linalg.generic #matvec %0 = linalg.generic #matvec
ins(%arga, %argb: tensor<?x?xi32, #SparseMatrix>, tensor<?xi32>) ins(%arga, %argb: tensor<?x?xi32, #SparseMatrix>, tensor<?xi32>)
outs(%argx: tensor<?xi32>) { inits(%argx: tensor<?xi32>) {
^bb(%a: i32, %b: i32, %x: i32): ^bb(%a: i32, %b: i32, %x: i32):
%0 = arith.muli %a, %b : i32 %0 = arith.muli %a, %b : i32
%1 = arith.addi %x, %0 : i32 %1 = arith.addi %x, %0 : i32
linalg.yield %1 : i32 linalg.yield %1 : i32
} -> tensor<?xi32> } -> tensor<?xi32>
return %0 : tensor<?xi32> return %0 : tensor<?xi32>
} }
▲ Show 20 Lines • Show All 50 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir

Show All 34 Linesmodule {
func.func @kernel_mttkrp(%argb: tensor<?x?x?xf64, #SparseTensor>, func.func @kernel_mttkrp(%argb: tensor<?x?x?xf64, #SparseTensor>,
%argc: tensor<?x?xf64>, %argc: tensor<?x?xf64>,
%argd: tensor<?x?xf64>, %argd: tensor<?x?xf64>,
%arga: tensor<?x?xf64>) %arga: tensor<?x?xf64>)
-> tensor<?x?xf64> { -> tensor<?x?xf64> {
%0 = linalg.generic #mttkrp %0 = linalg.generic #mttkrp
ins(%argb, %argc, %argd: ins(%argb, %argc, %argd:
tensor<?x?x?xf64, #SparseTensor>, tensor<?x?xf64>, tensor<?x?xf64>) tensor<?x?x?xf64, #SparseTensor>, tensor<?x?xf64>, tensor<?x?xf64>)
outs(%arga: tensor<?x?xf64>) { inits(%arga: tensor<?x?xf64>) {
^bb(%b: f64, %c: f64, %d: f64, %a: f64): ^bb(%b: f64, %c: f64, %d: f64, %a: f64):
%0 = arith.mulf %b, %c : f64 %0 = arith.mulf %b, %c : f64
%1 = arith.mulf %d, %0 : f64 %1 = arith.mulf %d, %0 : f64
%2 = arith.addf %a, %1 : f64 %2 = arith.addf %a, %1 : f64
linalg.yield %2 : f64 linalg.yield %2 : f64
} -> tensor<?x?xf64> } -> tensor<?x?xf64>
return %0 : tensor<?x?xf64> return %0 : tensor<?x?xf64>
} }
▲ Show 20 Lines • Show All 74 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_mult_elt.mlir

Show All 19 Lines
module { module {
// Sparse kernel. // Sparse kernel.
func.func @sparse_mult_elt( func.func @sparse_mult_elt(
%arga: tensor<32x16xf32, #DCSR>, %argb: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> { %arga: tensor<32x16xf32, #DCSR>, %argb: tensor<32x16xf32, #DCSR>) -> tensor<32x16xf32, #DCSR> {
%argx = bufferization.alloc_tensor() : tensor<32x16xf32, #DCSR> %argx = bufferization.alloc_tensor() : tensor<32x16xf32, #DCSR>
%0 = linalg.generic #trait_mult_elt %0 = linalg.generic #trait_mult_elt
ins(%arga, %argb: tensor<32x16xf32, #DCSR>, tensor<32x16xf32, #DCSR>) ins(%arga, %argb: tensor<32x16xf32, #DCSR>, tensor<32x16xf32, #DCSR>)
outs(%argx: tensor<32x16xf32, #DCSR>) { inits(%argx: tensor<32x16xf32, #DCSR>) {
^bb(%a: f32, %b: f32, %x: f32): ^bb(%a: f32, %b: f32, %x: f32):
%1 = arith.mulf %a, %b : f32 %1 = arith.mulf %a, %b : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<32x16xf32, #DCSR> } -> tensor<32x16xf32, #DCSR>
return %0 : tensor<32x16xf32, #DCSR> return %0 : tensor<32x16xf32, #DCSR>
} }
// Driver method to call and verify kernel. // Driver method to call and verify kernel.
Show All 37 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir

Show All 28 Lines func.func @redsum(%arga: tensor<?x?x?xi32, #SparseTensor>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?x?xi32, #SparseTensor> %d0 = tensor.dim %arga, %c0 : tensor<?x?x?xi32, #SparseTensor>
%d1 = tensor.dim %arga, %c1 : tensor<?x?x?xi32, #SparseTensor> %d1 = tensor.dim %arga, %c1 : tensor<?x?x?xi32, #SparseTensor>
%xinit = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xi32, #SparseMatrix> %xinit = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xi32, #SparseMatrix>
%0 = linalg.generic #redsum %0 = linalg.generic #redsum
ins(%arga, %argb: tensor<?x?x?xi32, #SparseTensor>, ins(%arga, %argb: tensor<?x?x?xi32, #SparseTensor>,
tensor<?x?x?xi32, #SparseTensor>) tensor<?x?x?xi32, #SparseTensor>)
outs(%xinit: tensor<?x?xi32, #SparseMatrix>) { inits(%xinit: tensor<?x?xi32, #SparseMatrix>) {
^bb(%a: i32, %b: i32, %x: i32): ^bb(%a: i32, %b: i32, %x: i32):
%0 = arith.muli %a, %b : i32 %0 = arith.muli %a, %b : i32
%1 = arith.addi %x, %0 : i32 %1 = arith.addi %x, %0 : i32
linalg.yield %1 : i32 linalg.yield %1 : i32
} -> tensor<?x?xi32, #SparseMatrix> } -> tensor<?x?xi32, #SparseMatrix>
return %0 : tensor<?x?xi32, #SparseMatrix> return %0 : tensor<?x?xi32, #SparseMatrix>
} }
▲ Show 20 Lines • Show All 44 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir

Show All 29 Linesmodule {
// A kernel that multiplies a sparse matrix A with itself // A kernel that multiplies a sparse matrix A with itself
// in an element-wise fashion. In this operation, we have // in an element-wise fashion. In this operation, we have
// a sparse tensor as output, but although the values of the // a sparse tensor as output, but although the values of the
// sparse tensor change, its nonzero structure remains the same. // sparse tensor change, its nonzero structure remains the same.
// //
func.func @kernel_eltwise_mult(%argx: tensor<?x?xf64, #DCSR>) func.func @kernel_eltwise_mult(%argx: tensor<?x?xf64, #DCSR>)
-> tensor<?x?xf64, #DCSR> { -> tensor<?x?xf64, #DCSR> {
%0 = linalg.generic #eltwise_mult %0 = linalg.generic #eltwise_mult
outs(%argx: tensor<?x?xf64, #DCSR>) { inits(%argx: tensor<?x?xf64, #DCSR>) {
^bb(%x: f64): ^bb(%x: f64):
%0 = arith.mulf %x, %x : f64 %0 = arith.mulf %x, %x : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<?x?xf64, #DCSR> } -> tensor<?x?xf64, #DCSR>
return %0 : tensor<?x?xf64, #DCSR> return %0 : tensor<?x?xf64, #DCSR>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 29 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir

Show All 12 Linesmodule {
func.func @quantized_matmul(%input1: tensor<5x3xi8>, func.func @quantized_matmul(%input1: tensor<5x3xi8>,
%input2: tensor<3x6xi8, #DCSR>, %input2: tensor<3x6xi8, #DCSR>,
%output: tensor<5x6xi32>) -> tensor<5x6xi32> { %output: tensor<5x6xi32>) -> tensor<5x6xi32> {
%c0 = arith.constant 0 : i32 %c0 = arith.constant 0 : i32
%c2 = arith.constant 2 : i32 %c2 = arith.constant 2 : i32
%0 = linalg.quantized_matmul %0 = linalg.quantized_matmul
ins(%input1, %input2, %c2, %c0 : tensor<5x3xi8>, tensor<3x6xi8, #DCSR>, i32, i32) ins(%input1, %input2, %c2, %c0 : tensor<5x3xi8>, tensor<3x6xi8, #DCSR>, i32, i32)
outs(%output : tensor<5x6xi32>) -> tensor<5x6xi32> inits(%output : tensor<5x6xi32>) -> tensor<5x6xi32>
return %0: tensor<5x6xi32> return %0: tensor<5x6xi32>
} }
func.func @entry() { func.func @entry() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%i0 = arith.constant 0 : i32 %i0 = arith.constant 0 : i32
%input1 = arith.constant dense<[ %input1 = arith.constant dense<[
▲ Show 20 Lines • Show All 41 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir

Show All 17 Lines
module { module {
func.func @cre(%arga: tensor<?xcomplex<f32>, #SparseVector>) func.func @cre(%arga: tensor<?xcomplex<f32>, #SparseVector>)
-> tensor<?xf32, #SparseVector> { -> tensor<?xf32, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga: tensor<?xcomplex<f32>, #SparseVector>) ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
outs(%xv: tensor<?xf32, #SparseVector>) { inits(%xv: tensor<?xf32, #SparseVector>) {
^bb(%a: complex<f32>, %x: f32): ^bb(%a: complex<f32>, %x: f32):
%1 = complex.re %a : complex<f32> %1 = complex.re %a : complex<f32>
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?xf32, #SparseVector> } -> tensor<?xf32, #SparseVector>
return %0 : tensor<?xf32, #SparseVector> return %0 : tensor<?xf32, #SparseVector>
} }
func.func @cim(%arga: tensor<?xcomplex<f32>, #SparseVector>) func.func @cim(%arga: tensor<?xcomplex<f32>, #SparseVector>)
-> tensor<?xf32, #SparseVector> { -> tensor<?xf32, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xcomplex<f32>, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf32, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga: tensor<?xcomplex<f32>, #SparseVector>) ins(%arga: tensor<?xcomplex<f32>, #SparseVector>)
outs(%xv: tensor<?xf32, #SparseVector>) { inits(%xv: tensor<?xf32, #SparseVector>) {
^bb(%a: complex<f32>, %x: f32): ^bb(%a: complex<f32>, %x: f32):
%1 = complex.im %a : complex<f32> %1 = complex.im %a : complex<f32>
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<?xf32, #SparseVector> } -> tensor<?xf32, #SparseVector>
return %0 : tensor<?xf32, #SparseVector> return %0 : tensor<?xf32, #SparseVector>
} }
func.func @dump(%arg0: tensor<?xf32, #SparseVector>) { func.func @dump(%arg0: tensor<?xf32, #SparseVector>) {
▲ Show 20 Lines • Show All 44 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir

Show First 20 Lines • Show All 44 Lines • ▼ Show 20 Lines
module { module {
func.func @redProdLex(%arga: tensor<?x?xf64, #CSR>) -> tensor<?xf64, #SparseVector> { func.func @redProdLex(%arga: tensor<?x?xf64, #CSR>) -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cf1 = arith.constant 1.0 : f64 %cf1 = arith.constant 1.0 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
%xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_mat_reduce_rowwise %0 = linalg.generic #trait_mat_reduce_rowwise
ins(%arga: tensor<?x?xf64, #CSR>) ins(%arga: tensor<?x?xf64, #CSR>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64): ^bb(%a: f64, %b: f64):
%1 = sparse_tensor.reduce %a, %b, %cf1 : f64 { %1 = sparse_tensor.reduce %a, %b, %cf1 : f64 {
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%2 = arith.mulf %x, %y : f64 %2 = arith.mulf %x, %y : f64
sparse_tensor.yield %2 : f64 sparse_tensor.yield %2 : f64
} }
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
func.func @redProdExpand(%arga: tensor<?x?xf64, #CSC>) -> tensor<?xf64, #SparseVector> { func.func @redProdExpand(%arga: tensor<?x?xf64, #CSC>) -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cf1 = arith.constant 1.0 : f64 %cf1 = arith.constant 1.0 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSC> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSC>
%xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_mat_reduce_rowwise %0 = linalg.generic #trait_mat_reduce_rowwise
ins(%arga: tensor<?x?xf64, #CSC>) ins(%arga: tensor<?x?xf64, #CSC>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64): ^bb(%a: f64, %b: f64):
%1 = sparse_tensor.reduce %a, %b, %cf1 : f64 { %1 = sparse_tensor.reduce %a, %b, %cf1 : f64 {
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%2 = arith.mulf %x, %y : f64 %2 = arith.mulf %x, %y : f64
sparse_tensor.yield %2 : f64 sparse_tensor.yield %2 : f64
} }
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
func.func @min_plus_csrcsr(%arga: tensor<?x?xf64, #CSR>, func.func @min_plus_csrcsr(%arga: tensor<?x?xf64, #CSR>,
%argb: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> { %argb: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%maxf = arith.constant 1.0e999 : f64 %maxf = arith.constant 1.0e999 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
%d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSR> %d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSR>
%xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR> %xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR>
%0 = linalg.generic #trait_matmul %0 = linalg.generic #trait_matmul
ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSR>) ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSR>)
outs(%xm: tensor<?x?xf64, #CSR>) { inits(%xm: tensor<?x?xf64, #CSR>) {
^bb(%a: f64, %b: f64, %output: f64): ^bb(%a: f64, %b: f64, %output: f64):
%1 = sparse_tensor.binary %a, %b : f64, f64 to f64 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap = { overlap = {
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%3 = arith.addf %x, %y : f64 %3 = arith.addf %x, %y : f64
sparse_tensor.yield %3 : f64 sparse_tensor.yield %3 : f64
} }
left={} left={}
Show All 14 Lines func.func @min_plus_csrcsc(%arga: tensor<?x?xf64, #CSR>,
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%maxf = arith.constant 1.0e999 : f64 %maxf = arith.constant 1.0e999 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
%d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSC> %d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSC>
%xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR> %xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR>
%0 = linalg.generic #trait_matmul %0 = linalg.generic #trait_matmul
ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSC>) ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSC>)
outs(%xm: tensor<?x?xf64, #CSR>) { inits(%xm: tensor<?x?xf64, #CSR>) {
^bb(%a: f64, %b: f64, %output: f64): ^bb(%a: f64, %b: f64, %output: f64):
%1 = sparse_tensor.binary %a, %b : f64, f64 to f64 %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
overlap = { overlap = {
^bb0(%x: f64, %y: f64): ^bb0(%x: f64, %y: f64):
%3 = arith.addf %x, %y : f64 %3 = arith.addf %x, %y : f64
sparse_tensor.yield %3 : f64 sparse_tensor.yield %3 : f64
} }
left={} left={}
▲ Show 20 Lines • Show All 101 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir

Show All 16 Lines
// An example of vector reductions. // An example of vector reductions.
module { module {
func.func @sum_reduction_i32(%arga: tensor<32xi32, #SV>, func.func @sum_reduction_i32(%arga: tensor<32xi32, #SV>,
%argx: tensor<i32>) -> tensor<i32> { %argx: tensor<i32>) -> tensor<i32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #SV>) ins(%arga: tensor<32xi32, #SV>)
outs(%argx: tensor<i32>) { inits(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%0 = arith.addi %x, %a : i32 %0 = arith.addi %x, %a : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
func.func @sum_reduction_f32(%arga: tensor<32xf32, #SV>, func.func @sum_reduction_f32(%arga: tensor<32xf32, #SV>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xf32, #SV>) ins(%arga: tensor<32xf32, #SV>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.addf %x, %a : f32 %0 = arith.addf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
func.func @prod_reduction_i32(%arga: tensor<32xi32, #DV>, func.func @prod_reduction_i32(%arga: tensor<32xi32, #DV>,
%argx: tensor<i32>) -> tensor<i32> { %argx: tensor<i32>) -> tensor<i32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #DV>) ins(%arga: tensor<32xi32, #DV>)
outs(%argx: tensor<i32>) { inits(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%0 = arith.muli %x, %a : i32 %0 = arith.muli %x, %a : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
func.func @prod_reduction_f32(%arga: tensor<32xf32, #DV>, func.func @prod_reduction_f32(%arga: tensor<32xf32, #DV>,
%argx: tensor<f32>) -> tensor<f32> { %argx: tensor<f32>) -> tensor<f32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xf32, #DV>) ins(%arga: tensor<32xf32, #DV>)
outs(%argx: tensor<f32>) { inits(%argx: tensor<f32>) {
^bb(%a: f32, %x: f32): ^bb(%a: f32, %x: f32):
%0 = arith.mulf %x, %a : f32 %0 = arith.mulf %x, %a : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} -> tensor<f32> } -> tensor<f32>
return %0 : tensor<f32> return %0 : tensor<f32>
} }
func.func @and_reduction_i32(%arga: tensor<32xi32, #DV>, func.func @and_reduction_i32(%arga: tensor<32xi32, #DV>,
%argx: tensor<i32>) -> tensor<i32> { %argx: tensor<i32>) -> tensor<i32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #DV>) ins(%arga: tensor<32xi32, #DV>)
outs(%argx: tensor<i32>) { inits(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%0 = arith.andi %x, %a : i32 %0 = arith.andi %x, %a : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
func.func @or_reduction_i32(%arga: tensor<32xi32, #SV>, func.func @or_reduction_i32(%arga: tensor<32xi32, #SV>,
%argx: tensor<i32>) -> tensor<i32> { %argx: tensor<i32>) -> tensor<i32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #SV>) ins(%arga: tensor<32xi32, #SV>)
outs(%argx: tensor<i32>) { inits(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%0 = arith.ori %x, %a : i32 %0 = arith.ori %x, %a : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
func.func @xor_reduction_i32(%arga: tensor<32xi32, #SV>, func.func @xor_reduction_i32(%arga: tensor<32xi32, #SV>,
%argx: tensor<i32>) -> tensor<i32> { %argx: tensor<i32>) -> tensor<i32> {
%0 = linalg.generic #trait_reduction %0 = linalg.generic #trait_reduction
ins(%arga: tensor<32xi32, #SV>) ins(%arga: tensor<32xi32, #SV>)
outs(%argx: tensor<i32>) { inits(%argx: tensor<i32>) {
^bb(%a: i32, %x: i32): ^bb(%a: i32, %x: i32):
%0 = arith.xori %x, %a : i32 %0 = arith.xori %x, %a : i32
linalg.yield %0 : i32 linalg.yield %0 : i32
} -> tensor<i32> } -> tensor<i32>
return %0 : tensor<i32> return %0 : tensor<i32>
} }
func.func @dump_i32(%arg0 : tensor<i32>) { func.func @dump_i32(%arg0 : tensor<i32>) {
▲ Show 20 Lines • Show All 92 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir

Show All 33 Linesmodule {
// A kernel that computes a sampled matrix matrix multiplication. // A kernel that computes a sampled matrix matrix multiplication.
// //
func.func @sampled_dense_dense(%args: tensor<?x?xf32, #SparseMatrix>, func.func @sampled_dense_dense(%args: tensor<?x?xf32, #SparseMatrix>,
%arga: tensor<?x?xf32>, %arga: tensor<?x?xf32>,
%argb: tensor<?x?xf32>, %argb: tensor<?x?xf32>,
%argx: tensor<?x?xf32>) -> tensor<?x?xf32> { %argx: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.generic #trait_sampled_dense_dense %0 = linalg.generic #trait_sampled_dense_dense
ins(%args, %arga, %argb: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32>, tensor<?x?xf32>) ins(%args, %arga, %argb: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32>, tensor<?x?xf32>)
outs(%argx: tensor<?x?xf32>) { inits(%argx: tensor<?x?xf32>) {
^bb(%s: f32, %a: f32, %b: f32, %x: f32): ^bb(%s: f32, %a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32 %0 = arith.mulf %a, %b : f32
%1 = arith.mulf %s, %0 : f32 %1 = arith.mulf %s, %0 : f32
%2 = arith.addf %x, %1 : f32 %2 = arith.addf %x, %1 : f32
linalg.yield %2 : f32 linalg.yield %2 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
▲ Show 20 Lines • Show All 68 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Linesmodule {
// //
func.func @sampled_dd(%args: tensor<8x8xf64, #SM>, func.func @sampled_dd(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64> {
%1 = arith.constant dense<0.0> : tensor<8x8xf64> %1 = arith.constant dense<0.0> : tensor<8x8xf64>
%2 = linalg.generic #trait_sampled_dense_dense %2 = linalg.generic #trait_sampled_dense_dense
ins(%args, %arga, %argb: tensor<8x8xf64, #SM>, ins(%args, %arga, %argb: tensor<8x8xf64, #SM>,
tensor<8x8xf64>, tensor<8x8xf64>) tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1: tensor<8x8xf64>) { inits(%1: tensor<8x8xf64>) {
^bb(%s: f64, %a: f64, %b: f64, %x: f64): ^bb(%s: f64, %a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.mulf %s, %p : f64 %q = arith.mulf %s, %p : f64
%r = arith.addf %x, %q : f64 %r = arith.addf %x, %q : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
return %2 : tensor<8x8xf64> return %2 : tensor<8x8xf64>
} }
// //
// A kernel that computes an unfused sampled matrix matrix multiplication // A kernel that computes an unfused sampled matrix matrix multiplication
// (with dense result). // (with dense result).
// //
func.func @sampled_dd_unfused(%args: tensor<8x8xf64, #SM>, func.func @sampled_dd_unfused(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64> {
// Perform dense-dense matrix matrix multiplication. // Perform dense-dense matrix matrix multiplication.
%1 = arith.constant dense<0.0> : tensor<8x8xf64> %1 = arith.constant dense<0.0> : tensor<8x8xf64>
%2 = linalg.generic #trait_matmul %2 = linalg.generic #trait_matmul
ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>) ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%a: f64, %b: f64, %x: f64): ^bb0(%a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.addf %x, %p : f64 %q = arith.addf %x, %p : f64
linalg.yield %q : f64 linalg.yield %q : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
// Sample the result with elements-wise multiplication with sparse matrix. // Sample the result with elements-wise multiplication with sparse matrix.
%3 = linalg.generic #trait_scale %3 = linalg.generic #trait_scale
ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>) ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%t: f64, %s: f64, %x: f64): ^bb0(%t: f64, %s: f64, %x: f64):
%r = arith.mulf %t, %s : f64 %r = arith.mulf %t, %s : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
return %3 : tensor<8x8xf64> return %3 : tensor<8x8xf64>
} }
// //
// A kernel that computes a direct sampled matrix matrix multiplication // A kernel that computes a direct sampled matrix matrix multiplication
// (with sparse result). // (with sparse result).
// //
func.func @sparse_sampled_dd(%args: tensor<8x8xf64, #SM>, func.func @sparse_sampled_dd(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> {
%1 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM> %1 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM>
%2 = linalg.generic #trait_sampled_dense_dense %2 = linalg.generic #trait_sampled_dense_dense
ins(%args, %arga, %argb: tensor<8x8xf64, #SM>, ins(%args, %arga, %argb: tensor<8x8xf64, #SM>,
tensor<8x8xf64>, tensor<8x8xf64>) tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1: tensor<8x8xf64, #SM>) { inits(%1: tensor<8x8xf64, #SM>) {
^bb(%s: f64, %a: f64, %b: f64, %x: f64): ^bb(%s: f64, %a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.mulf %s, %p : f64 %q = arith.mulf %s, %p : f64
%r = arith.addf %x, %q : f64 %r = arith.addf %x, %q : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64, #SM> } -> tensor<8x8xf64, #SM>
return %2 : tensor<8x8xf64, #SM> return %2 : tensor<8x8xf64, #SM>
} }
// //
// A kernel that computes an unfused sampled matrix matrix multiplication // A kernel that computes an unfused sampled matrix matrix multiplication
// (with sparse result). // (with sparse result).
// //
func.func @sparse_sampled_dd_unfused( func.func @sparse_sampled_dd_unfused(
%args: tensor<8x8xf64, #SM>, %args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>, %arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> { %argb: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> {
// Perform dense-dense matrix matrix multiplication. // Perform dense-dense matrix matrix multiplication.
%1 = arith.constant dense<0.0> : tensor<8x8xf64> %1 = arith.constant dense<0.0> : tensor<8x8xf64>
%2 = linalg.generic #trait_matmul %2 = linalg.generic #trait_matmul
ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>) ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1 : tensor<8x8xf64>) { inits(%1 : tensor<8x8xf64>) {
^bb0(%a: f64, %b: f64, %x: f64): ^bb0(%a: f64, %b: f64, %x: f64):
%p = arith.mulf %a, %b : f64 %p = arith.mulf %a, %b : f64
%q = arith.addf %x, %p : f64 %q = arith.addf %x, %p : f64
linalg.yield %q : f64 linalg.yield %q : f64
} -> tensor<8x8xf64> } -> tensor<8x8xf64>
// Sample the result with elements-wise multiplication with sparse matrix. // Sample the result with elements-wise multiplication with sparse matrix.
%3 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM> %3 = bufferization.alloc_tensor() : tensor<8x8xf64, #SM>
%4 = linalg.generic #trait_scale %4 = linalg.generic #trait_scale
ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>) ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>)
outs(%3 : tensor<8x8xf64, #SM>) { inits(%3 : tensor<8x8xf64, #SM>) {
^bb0(%t: f64, %s: f64, %x: f64): ^bb0(%t: f64, %s: f64, %x: f64):
%r = arith.mulf %t, %s : f64 %r = arith.mulf %t, %s : f64
linalg.yield %r : f64 linalg.yield %r : f64
} -> tensor<8x8xf64, #SM> } -> tensor<8x8xf64, #SM>
return %4 : tensor<8x8xf64, #SM> return %4 : tensor<8x8xf64, #SM>
} }
// //
▲ Show 20 Lines • Show All 67 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir

Show All 20 Lines
// //
module { module {
// //
// A kernel that scales a sparse matrix A by a factor of 2.0. // A kernel that scales a sparse matrix A by a factor of 2.0.
// //
func.func @sparse_scale(%argx: tensor<8x8xf32, #CSR>) -> tensor<8x8xf32, #CSR> { func.func @sparse_scale(%argx: tensor<8x8xf32, #CSR>) -> tensor<8x8xf32, #CSR> {
%c = arith.constant 2.0 : f32 %c = arith.constant 2.0 : f32
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
outs(%argx: tensor<8x8xf32, #CSR>) { inits(%argx: tensor<8x8xf32, #CSR>) {
^bb(%x: f32): ^bb(%x: f32):
%1 = arith.mulf %x, %c : f32 %1 = arith.mulf %x, %c : f32
linalg.yield %1 : f32 linalg.yield %1 : f32
} -> tensor<8x8xf32, #CSR> } -> tensor<8x8xf32, #CSR>
return %0 : tensor<8x8xf32, #CSR> return %0 : tensor<8x8xf32, #CSR>
} }
// //
Show All 39 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scf_nested.mlir

Show All 12 Lines func.func public @condition(%cond: i1, %arg0: tensor<2x3x4xf64, #SparseMatrix>) -> tensor<2x3x4xf64, #SparseMatrix> {
%1 = scf.if %cond -> (tensor<2x3x4xf64, #SparseMatrix>) { %1 = scf.if %cond -> (tensor<2x3x4xf64, #SparseMatrix>) {
%cst_2 = arith.constant dense<1.000000e+00> : tensor<f64> %cst_2 = arith.constant dense<1.000000e+00> : tensor<f64>
%cst_3 = arith.constant dense<1.000000e+00> : tensor<2x3x4xf64> %cst_3 = arith.constant dense<1.000000e+00> : tensor<2x3x4xf64>
%2 = bufferization.alloc_tensor() : tensor<2x3x4xf64, #SparseMatrix> %2 = bufferization.alloc_tensor() : tensor<2x3x4xf64, #SparseMatrix>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0, %cst_3 : tensor<2x3x4xf64, #SparseMatrix>, tensor<2x3x4xf64>) ins(%arg0, %cst_3 : tensor<2x3x4xf64, #SparseMatrix>, tensor<2x3x4xf64>)
outs(%2 : tensor<2x3x4xf64, #SparseMatrix>) { inits(%2 : tensor<2x3x4xf64, #SparseMatrix>) {
^bb0(%arg1: f64, %arg2: f64, %arg3: f64): ^bb0(%arg1: f64, %arg2: f64, %arg3: f64):
%4 = arith.subf %arg1, %arg2 : f64 %4 = arith.subf %arg1, %arg2 : f64
linalg.yield %4 : f64 linalg.yield %4 : f64
} -> tensor<2x3x4xf64, #SparseMatrix> } -> tensor<2x3x4xf64, #SparseMatrix>
scf.yield %3 : tensor<2x3x4xf64, #SparseMatrix> scf.yield %3 : tensor<2x3x4xf64, #SparseMatrix>
} else { } else {
%cst_2 = arith.constant dense<1.000000e+00> : tensor<f64> %cst_2 = arith.constant dense<1.000000e+00> : tensor<f64>
%cst_3 = arith.constant dense<1.000000e+00> : tensor<2x3x4xf64> %cst_3 = arith.constant dense<1.000000e+00> : tensor<2x3x4xf64>
%2 = bufferization.alloc_tensor() : tensor<2x3x4xf64, #SparseMatrix> %2 = bufferization.alloc_tensor() : tensor<2x3x4xf64, #SparseMatrix>
%3 = linalg.generic { %3 = linalg.generic {
indexing_maps = [#map, #map, #map], indexing_maps = [#map, #map, #map],
iterator_types = ["parallel", "parallel", "parallel"]} iterator_types = ["parallel", "parallel", "parallel"]}
ins(%arg0, %cst_3 : tensor<2x3x4xf64, #SparseMatrix>, tensor<2x3x4xf64>) ins(%arg0, %cst_3 : tensor<2x3x4xf64, #SparseMatrix>, tensor<2x3x4xf64>)
outs(%2 : tensor<2x3x4xf64, #SparseMatrix>) { inits(%2 : tensor<2x3x4xf64, #SparseMatrix>) {
^bb0(%arg1: f64, %arg2: f64, %arg3: f64): ^bb0(%arg1: f64, %arg2: f64, %arg3: f64):
%4 = arith.addf %arg1, %arg2 : f64 %4 = arith.addf %arg1, %arg2 : f64
linalg.yield %4 : f64 linalg.yield %4 : f64
} -> tensor<2x3x4xf64, #SparseMatrix> } -> tensor<2x3x4xf64, #SparseMatrix>
scf.yield %3 : tensor<2x3x4xf64, #SparseMatrix> scf.yield %3 : tensor<2x3x4xf64, #SparseMatrix>
} }
return %1 : tensor<2x3x4xf64, #SparseMatrix> return %1 : tensor<2x3x4xf64, #SparseMatrix>
} }
Show All 31 Lines func.func public @entry() {
call @dump(%sm_f) : (tensor<2x3x4xf64, #SparseMatrix>) -> () call @dump(%sm_f) : (tensor<2x3x4xf64, #SparseMatrix>) -> ()
bufferization.dealloc_tensor %sm : tensor<2x3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %sm : tensor<2x3x4xf64, #SparseMatrix>
bufferization.dealloc_tensor %sm_t : tensor<2x3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %sm_t : tensor<2x3x4xf64, #SparseMatrix>
bufferization.dealloc_tensor %sm_f : tensor<2x3x4xf64, #SparseMatrix> bufferization.dealloc_tensor %sm_f : tensor<2x3x4xf64, #SparseMatrix>
return return
} }
} }
No newline at end of file No newline at end of file

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_select.mlir

Show All 32 Lines
module { module {
func.func @vecSelect(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { func.func @vecSelect(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%cf1 = arith.constant 1.0 : f64 %cf1 = arith.constant 1.0 : f64
%d0 = tensor.dim %arga, %c0 : tensor<?xf64, #SparseVector> %d0 = tensor.dim %arga, %c0 : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d0): tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_vec_select %0 = linalg.generic #trait_vec_select
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64): ^bb(%a: f64, %b: f64):
%1 = sparse_tensor.select %a : f64 { %1 = sparse_tensor.select %a : f64 {
^bb0(%x: f64): ^bb0(%x: f64):
%keep = arith.cmpf "oge", %x, %cf1 : f64 %keep = arith.cmpf "oge", %x, %cf1 : f64
sparse_tensor.yield %keep : i1 sparse_tensor.yield %keep : i1
} }
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
func.func @matUpperTriangle(%arga: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> { func.func @matUpperTriangle(%arga: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR> %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
%d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #CSR> %d1 = tensor.dim %arga, %c1 : tensor<?x?xf64, #CSR>
%xv = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xf64, #CSR> %xv = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xf64, #CSR>
%0 = linalg.generic #trait_mat_select %0 = linalg.generic #trait_mat_select
ins(%arga: tensor<?x?xf64, #CSR>) ins(%arga: tensor<?x?xf64, #CSR>)
outs(%xv: tensor<?x?xf64, #CSR>) { inits(%xv: tensor<?x?xf64, #CSR>) {
^bb(%a: f64, %b: f64): ^bb(%a: f64, %b: f64):
%row = linalg.index 0 : index %row = linalg.index 0 : index
%col = linalg.index 1 : index %col = linalg.index 1 : index
%1 = sparse_tensor.select %a : f64 { %1 = sparse_tensor.select %a : f64 {
^bb0(%x: f64): ^bb0(%x: f64):
%keep = arith.cmpi "ugt", %col, %row : index %keep = arith.cmpi "ugt", %col, %row : index
sparse_tensor.yield %keep : i1 sparse_tensor.yield %keep : i1
} }
▲ Show 20 Lines • Show All 79 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sign.mlir

Show All 27 Linesmodule {
// -0.0: -0.0 // -0.0: -0.0
func.func @sparse_sign(%arg0: tensor<?xf64, #SparseVector>) func.func @sparse_sign(%arg0: tensor<?xf64, #SparseVector>)
-> tensor<?xf64, #SparseVector> { -> tensor<?xf64, #SparseVector> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector> %d = tensor.dim %arg0, %c0 : tensor<?xf64, #SparseVector>
%xin = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xin = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arg0: tensor<?xf64, #SparseVector>) ins(%arg0: tensor<?xf64, #SparseVector>)
outs(%xin: tensor<?xf64, #SparseVector>) { inits(%xin: tensor<?xf64, #SparseVector>) {
^bb0(%a: f64, %x: f64) : ^bb0(%a: f64, %x: f64) :
%result = sparse_tensor.unary %a : f64 to f64 %result = sparse_tensor.unary %a : f64 to f64
present={ present={
^bb1(%s: f64): ^bb1(%s: f64):
%z = arith.constant 0.0 : f64 %z = arith.constant 0.0 : f64
%1 = arith.cmpf one, %s, %z : f64 %1 = arith.cmpf one, %s, %z : f64
%2 = arith.uitofp %1 : i1 to f64 %2 = arith.uitofp %1 : i1 to f64
%3 = math.copysign %2, %s : f64 %3 = math.copysign %2, %s : f64
▲ Show 20 Lines • Show All 56 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sorted_coo.mlir

Show First 20 Lines • Show All 43 Lines • ▼ Show 20 Linesmodule {
// //
// A kernel that scales a sparse matrix A by a factor of 2.0. // A kernel that scales a sparse matrix A by a factor of 2.0.
// //
func.func @sparse_scale(%argx: tensor<?x?xf64, #SortedCOO>) func.func @sparse_scale(%argx: tensor<?x?xf64, #SortedCOO>)
-> tensor<?x?xf64, #SortedCOO> { -> tensor<?x?xf64, #SortedCOO> {
%c = arith.constant 2.0 : f64 %c = arith.constant 2.0 : f64
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
outs(%argx: tensor<?x?xf64, #SortedCOO>) { inits(%argx: tensor<?x?xf64, #SortedCOO>) {
^bb(%x: f64): ^bb(%x: f64):
%1 = arith.mulf %x, %c : f64 %1 = arith.mulf %x, %c : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64, #SortedCOO> } -> tensor<?x?xf64, #SortedCOO>
return %0 : tensor<?x?xf64, #SortedCOO> return %0 : tensor<?x?xf64, #SortedCOO>
} }
func.func @dumpi(%arg0: memref<?xindex>) { func.func @dumpi(%arg0: memref<?xindex>) {
▲ Show 20 Lines • Show All 170 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir

Show All 30 Linesmodule {
// A kernel that multiplies a sparse matrix A with a dense matrix B // A kernel that multiplies a sparse matrix A with a dense matrix B
// into a dense matrix X. // into a dense matrix X.
// //
func.func @kernel_spmm(%arga: tensor<?x?xf64, #SparseMatrix>, func.func @kernel_spmm(%arga: tensor<?x?xf64, #SparseMatrix>,
%argb: tensor<?x?xf64>, %argb: tensor<?x?xf64>,
%argx: tensor<?x?xf64>) -> tensor<?x?xf64> { %argx: tensor<?x?xf64>) -> tensor<?x?xf64> {
%0 = linalg.generic #spmm %0 = linalg.generic #spmm
ins(%arga, %argb: tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>) ins(%arga, %argb: tensor<?x?xf64, #SparseMatrix>, tensor<?x?xf64>)
outs(%argx: tensor<?x?xf64>) { inits(%argx: tensor<?x?xf64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%0 = arith.mulf %a, %b : f64 %0 = arith.mulf %a, %b : f64
%1 = arith.addf %x, %0 : f64 %1 = arith.addf %x, %0 : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?xf64> } -> tensor<?x?xf64>
return %0 : tensor<?x?xf64> return %0 : tensor<?x?xf64>
} }
▲ Show 20 Lines • Show All 52 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

Show All 27 Lines
module { module {
// //
// A kernel that sum-reduces a matrix to a single scalar. // A kernel that sum-reduces a matrix to a single scalar.
// //
func.func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix>, func.func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix>,
%argx: tensor<f64>) -> tensor<f64> { %argx: tensor<f64>) -> tensor<f64> {
%0 = linalg.generic #trait_sum_reduce %0 = linalg.generic #trait_sum_reduce
ins(%arga: tensor<?x?xf64, #SparseMatrix>) ins(%arga: tensor<?x?xf64, #SparseMatrix>)
outs(%argx: tensor<f64>) { inits(%argx: tensor<f64>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%0 = arith.addf %x, %a : f64 %0 = arith.addf %x, %a : f64
linalg.yield %0 : f64 linalg.yield %0 : f64
} -> tensor<f64> } -> tensor<f64>
return %0 : tensor<f64> return %0 : tensor<f64>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 33 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_bf16.mlir

Show All 23 Lines
module { module {
// //
// A kernel that sum-reduces a matrix to a single scalar. // A kernel that sum-reduces a matrix to a single scalar.
// //
func.func @kernel_sum_reduce(%arga: tensor<?x?xbf16, #SparseMatrix>, func.func @kernel_sum_reduce(%arga: tensor<?x?xbf16, #SparseMatrix>,
%argx: tensor<bf16>) -> tensor<bf16> { %argx: tensor<bf16>) -> tensor<bf16> {
%0 = linalg.generic #trait_sum_reduce %0 = linalg.generic #trait_sum_reduce
ins(%arga: tensor<?x?xbf16, #SparseMatrix>) ins(%arga: tensor<?x?xbf16, #SparseMatrix>)
outs(%argx: tensor<bf16>) { inits(%argx: tensor<bf16>) {
^bb(%a: bf16, %x: bf16): ^bb(%a: bf16, %x: bf16):
%0 = arith.addf %x, %a : bf16 %0 = arith.addf %x, %a : bf16
linalg.yield %0 : bf16 linalg.yield %0 : bf16
} -> tensor<bf16> } -> tensor<bf16>
return %0 : tensor<bf16> return %0 : tensor<bf16>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 36 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_c32.mlir

Show All 27 Lines
module { module {
// //
// A kernel that sum-reduces a matrix to a single scalar. // A kernel that sum-reduces a matrix to a single scalar.
// //
func.func @kernel_sum_reduce(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>, func.func @kernel_sum_reduce(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>,
%argx: tensor<complex<f64>>) -> tensor<complex<f64>> { %argx: tensor<complex<f64>>) -> tensor<complex<f64>> {
%0 = linalg.generic #trait_sum_reduce %0 = linalg.generic #trait_sum_reduce
ins(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>) ins(%arga: tensor<?x?xcomplex<f64>, #SparseMatrix>)
outs(%argx: tensor<complex<f64>>) { inits(%argx: tensor<complex<f64>>) {
^bb(%a: complex<f64>, %x: complex<f64>): ^bb(%a: complex<f64>, %x: complex<f64>):
%0 = complex.add %x, %a : complex<f64> %0 = complex.add %x, %a : complex<f64>
linalg.yield %0 : complex<f64> linalg.yield %0 : complex<f64>
} -> tensor<complex<f64>> } -> tensor<complex<f64>>
return %0 : tensor<complex<f64>> return %0 : tensor<complex<f64>>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 40 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum_f16.mlir

Show All 21 Lines
module { module {
// //
// A kernel that sum-reduces a matrix to a single scalar. // A kernel that sum-reduces a matrix to a single scalar.
// //
func.func @kernel_sum_reduce(%arga: tensor<?x?xf16, #SparseMatrix>, func.func @kernel_sum_reduce(%arga: tensor<?x?xf16, #SparseMatrix>,
%argx: tensor<f16>) -> tensor<f16> { %argx: tensor<f16>) -> tensor<f16> {
%0 = linalg.generic #trait_sum_reduce %0 = linalg.generic #trait_sum_reduce
ins(%arga: tensor<?x?xf16, #SparseMatrix>) ins(%arga: tensor<?x?xf16, #SparseMatrix>)
outs(%argx: tensor<f16>) { inits(%argx: tensor<f16>) {
^bb(%a: f16, %x: f16): ^bb(%a: f16, %x: f16):
%0 = arith.addf %x, %a : f16 %0 = arith.addf %x, %a : f16
linalg.yield %0 : f16 linalg.yield %0 : f16
} -> tensor<f16> } -> tensor<f16>
return %0 : tensor<f16> return %0 : tensor<f16>
} }
func.func private @getTensorFilename(index) -> (!Filename) func.func private @getTensorFilename(index) -> (!Filename)
Show All 36 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tanh.mlir

Show All 13 Lines#trait_op = {
doc = "X(i) = OP X(i)" doc = "X(i) = OP X(i)"
} }
module { module {
// Performs zero-preserving math to sparse vector. // Performs zero-preserving math to sparse vector.
func.func @sparse_tanh(%vec: tensor<?xf64, #SparseVector>) func.func @sparse_tanh(%vec: tensor<?xf64, #SparseVector>)
-> tensor<?xf64, #SparseVector> { -> tensor<?xf64, #SparseVector> {
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
outs(%vec: tensor<?xf64, #SparseVector>) { inits(%vec: tensor<?xf64, #SparseVector>) {
^bb(%x: f64): ^bb(%x: f64):
%1 = math.tanh %x : f64 %1 = math.tanh %x : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Dumps a sparse vector of type f64. // Dumps a sparse vector of type f64.
▲ Show 20 Lines • Show All 43 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_mul.mlir

Show All 26 Lines func.func @tensor_mul(%arga: tensor<?x?x?xf64, #ST>,
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?x?xf64, #ST> %d0 = tensor.dim %arga, %c0 : tensor<?x?x?xf64, #ST>
%d1 = tensor.dim %arga, %c1 : tensor<?x?x?xf64, #ST> %d1 = tensor.dim %arga, %c1 : tensor<?x?x?xf64, #ST>
%d2 = tensor.dim %arga, %c2 : tensor<?x?x?xf64, #ST> %d2 = tensor.dim %arga, %c2 : tensor<?x?x?xf64, #ST>
%xt = bufferization.alloc_tensor(%d0, %d1, %d2) : tensor<?x?x?xf64, #ST> %xt = bufferization.alloc_tensor(%d0, %d1, %d2) : tensor<?x?x?xf64, #ST>
%0 = linalg.generic #trait_mul %0 = linalg.generic #trait_mul
ins(%arga, %argb: tensor<?x?x?xf64, #ST>, tensor<?x?x?xf64, #ST>) ins(%arga, %argb: tensor<?x?x?xf64, #ST>, tensor<?x?x?xf64, #ST>)
outs(%xt: tensor<?x?x?xf64, #ST>) { inits(%xt: tensor<?x?x?xf64, #ST>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.mulf %a, %b : f64 %1 = arith.mulf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?x?xf64, #ST> } -> tensor<?x?x?xf64, #ST>
return %0 : tensor<?x?x?xf64, #ST> return %0 : tensor<?x?x?xf64, #ST>
} }
// Driver method to call and verify tensor multiplication kernel. // Driver method to call and verify tensor multiplication kernel.
▲ Show 20 Lines • Show All 59 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_tensor_ops.mlir

Show All 26 Lines func.func @tensor_scale(%arga: tensor<?x?x?xf64, #ST1>) -> tensor<?x?x?xf64, #ST2> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
%d0 = tensor.dim %arga, %c0 : tensor<?x?x?xf64, #ST1> %d0 = tensor.dim %arga, %c0 : tensor<?x?x?xf64, #ST1>
%d1 = tensor.dim %arga, %c1 : tensor<?x?x?xf64, #ST1> %d1 = tensor.dim %arga, %c1 : tensor<?x?x?xf64, #ST1>
%d2 = tensor.dim %arga, %c2 : tensor<?x?x?xf64, #ST1> %d2 = tensor.dim %arga, %c2 : tensor<?x?x?xf64, #ST1>
%xm = bufferization.alloc_tensor(%d0, %d1, %d2) : tensor<?x?x?xf64, #ST2> %xm = bufferization.alloc_tensor(%d0, %d1, %d2) : tensor<?x?x?xf64, #ST2>
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
ins(%arga: tensor<?x?x?xf64, #ST1>) ins(%arga: tensor<?x?x?xf64, #ST1>)
outs(%xm: tensor<?x?x?xf64, #ST2>) { inits(%xm: tensor<?x?x?xf64, #ST2>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = arith.mulf %a, %s : f64 %1 = arith.mulf %a, %s : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?x?x?xf64, #ST2> } -> tensor<?x?x?xf64, #ST2>
return %0 : tensor<?x?x?xf64, #ST2> return %0 : tensor<?x?x?xf64, #ST2>
} }
// Driver method to call and verify tensor kernel. // Driver method to call and verify tensor kernel.
Show All 40 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_transpose.mlir

Show All 31 Linesmodule {
func.func @sparse_transpose(%arga: tensor<3x4xf64, #DCSR>) func.func @sparse_transpose(%arga: tensor<3x4xf64, #DCSR>)
-> tensor<4x3xf64, #DCSR> { -> tensor<4x3xf64, #DCSR> {
%t = sparse_tensor.convert %arga %t = sparse_tensor.convert %arga
: tensor<3x4xf64, #DCSR> to tensor<3x4xf64, #DCSC> : tensor<3x4xf64, #DCSR> to tensor<3x4xf64, #DCSC>
%i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR> %i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR>
%0 = linalg.generic #transpose_trait %0 = linalg.generic #transpose_trait
ins(%t: tensor<3x4xf64, #DCSC>) ins(%t: tensor<3x4xf64, #DCSC>)
outs(%i: tensor<4x3xf64, #DCSR>) { inits(%i: tensor<4x3xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
linalg.yield %a : f64 linalg.yield %a : f64
} -> tensor<4x3xf64, #DCSR> } -> tensor<4x3xf64, #DCSR>
bufferization.dealloc_tensor %t : tensor<3x4xf64, #DCSC> bufferization.dealloc_tensor %t : tensor<3x4xf64, #DCSC>
return %0 : tensor<4x3xf64, #DCSR> return %0 : tensor<4x3xf64, #DCSR>
} }
// //
// However, even better, the sparse compiler is able to insert such a // However, even better, the sparse compiler is able to insert such a
// conversion automatically to resolve a cycle in the iteration graph! // conversion automatically to resolve a cycle in the iteration graph!
// //
func.func @sparse_transpose_auto(%arga: tensor<3x4xf64, #DCSR>) func.func @sparse_transpose_auto(%arga: tensor<3x4xf64, #DCSR>)
-> tensor<4x3xf64, #DCSR> { -> tensor<4x3xf64, #DCSR> {
%i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR> %i = bufferization.alloc_tensor() : tensor<4x3xf64, #DCSR>
%0 = linalg.generic #transpose_trait %0 = linalg.generic #transpose_trait
ins(%arga: tensor<3x4xf64, #DCSR>) ins(%arga: tensor<3x4xf64, #DCSR>)
outs(%i: tensor<4x3xf64, #DCSR>) { inits(%i: tensor<4x3xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
linalg.yield %a : f64 linalg.yield %a : f64
} -> tensor<4x3xf64, #DCSR> } -> tensor<4x3xf64, #DCSR>
return %0 : tensor<4x3xf64, #DCSR> return %0 : tensor<4x3xf64, #DCSR>
} }
// //
// Main driver. // Main driver.
▲ Show 20 Lines • Show All 53 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_unary.mlir

Show All 29 Linesmodule {
// Missing values are filled with 1 (i32). // Missing values are filled with 1 (i32).
func.func @vector_complement(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> { func.func @vector_complement(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xi32, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%ci1 = arith.constant 1 : i32 %ci1 = arith.constant 1 : i32
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xi32, #SparseVector>
%0 = linalg.generic #trait_vec_scale %0 = linalg.generic #trait_vec_scale
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xi32, #SparseVector>) { inits(%xv: tensor<?xi32, #SparseVector>) {
^bb(%a: f64, %x: i32): ^bb(%a: f64, %x: i32):
%1 = sparse_tensor.unary %a : f64 to i32 %1 = sparse_tensor.unary %a : f64 to i32
present={} present={}
absent={ absent={
sparse_tensor.yield %ci1 : i32 sparse_tensor.yield %ci1 : i32
} }
linalg.yield %1 : i32 linalg.yield %1 : i32
} -> tensor<?xi32, #SparseVector> } -> tensor<?xi32, #SparseVector>
return %0 : tensor<?xi32, #SparseVector> return %0 : tensor<?xi32, #SparseVector>
} }
// Negate existing values. Fill missing ones with +1. // Negate existing values. Fill missing ones with +1.
func.func @vector_negation(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { func.func @vector_negation(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%cf1 = arith.constant 1.0 : f64 %cf1 = arith.constant 1.0 : f64
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_vec_scale %0 = linalg.generic #trait_vec_scale
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = sparse_tensor.unary %a : f64 to f64 %1 = sparse_tensor.unary %a : f64 to f64
present={ present={
^bb0(%x0: f64): ^bb0(%x0: f64):
%ret = arith.negf %x0 : f64 %ret = arith.negf %x0 : f64
sparse_tensor.yield %ret : f64 sparse_tensor.yield %ret : f64
} }
absent={ absent={
sparse_tensor.yield %cf1 : f64 sparse_tensor.yield %cf1 : f64
} }
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Performs B[i] = i * A[i]. // Performs B[i] = i * A[i].
func.func @vector_magnify(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { func.func @vector_magnify(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_vec_scale %0 = linalg.generic #trait_vec_scale
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%idx = linalg.index 0 : index %idx = linalg.index 0 : index
%1 = sparse_tensor.unary %a : f64 to f64 %1 = sparse_tensor.unary %a : f64 to f64
present={ present={
^bb0(%x0: f64): ^bb0(%x0: f64):
%tmp = arith.index_cast %idx : index to i64 %tmp = arith.index_cast %idx : index to i64
%idxf = arith.uitofp %tmp : i64 to f64 %idxf = arith.uitofp %tmp : i64 to f64
%ret = arith.mulf %x0, %idxf : f64 %ret = arith.mulf %x0, %idxf : f64
Show All 11 Lines func.func @matrix_clip(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%cfmin = arith.constant 3.0 : f64 %cfmin = arith.constant 3.0 : f64
%cfmax = arith.constant 7.0 : f64 %cfmax = arith.constant 7.0 : f64
%d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR>
%xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_mat_scale %0 = linalg.generic #trait_mat_scale
ins(%argx: tensor<?x?xf64, #DCSR>) ins(%argx: tensor<?x?xf64, #DCSR>)
outs(%xv: tensor<?x?xf64, #DCSR>) { inits(%xv: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = sparse_tensor.unary %a: f64 to f64 %1 = sparse_tensor.unary %a: f64 to f64
present={ present={
^bb0(%x0: f64): ^bb0(%x0: f64):
%mincmp = arith.cmpf "ogt", %x0, %cfmin : f64 %mincmp = arith.cmpf "ogt", %x0, %cfmin : f64
%x1 = arith.select %mincmp, %x0, %cfmin : f64 %x1 = arith.select %mincmp, %x0, %cfmin : f64
%maxcmp = arith.cmpf "olt", %x1, %cfmax : f64 %maxcmp = arith.cmpf "olt", %x1, %cfmax : f64
%x2 = arith.select %maxcmp, %x1, %cfmax : f64 %x2 = arith.select %maxcmp, %x1, %cfmax : f64
Show All 10 Linesmodule {
func.func @matrix_slice(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> { func.func @matrix_slice(%argx: tensor<?x?xf64, #DCSR>) -> tensor<?x?xf64, #DCSR> {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index %c1 = arith.constant 1 : index
%d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR> %d0 = tensor.dim %argx, %c0 : tensor<?x?xf64, #DCSR>
%d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR> %d1 = tensor.dim %argx, %c1 : tensor<?x?xf64, #DCSR>
%xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR> %xv = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #DCSR>
%0 = linalg.generic #trait_mat_scale %0 = linalg.generic #trait_mat_scale
ins(%argx: tensor<?x?xf64, #DCSR>) ins(%argx: tensor<?x?xf64, #DCSR>)
outs(%xv: tensor<?x?xf64, #DCSR>) { inits(%xv: tensor<?x?xf64, #DCSR>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%row = linalg.index 0 : index %row = linalg.index 0 : index
%col = linalg.index 1 : index %col = linalg.index 1 : index
%1 = sparse_tensor.unary %a: f64 to f64 %1 = sparse_tensor.unary %a: f64 to f64
present={ present={
^bb0(%x0: f64): ^bb0(%x0: f64):
%v = arith.constant 99.0 : f64 %v = arith.constant 99.0 : f64
%two = arith.constant 2 : index %two = arith.constant 2 : index
▲ Show 20 Lines • Show All 123 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_vector_ops.mlir

Show First 20 Lines • Show All 47 Lines • ▼ Show 20 Linesmodule {
// Scales a sparse vector into a new sparse vector. // Scales a sparse vector into a new sparse vector.
func.func @vector_scale(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { func.func @vector_scale(%arga: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%s = arith.constant 2.0 : f64 %s = arith.constant 2.0 : f64
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_scale %0 = linalg.generic #trait_scale
ins(%arga: tensor<?xf64, #SparseVector>) ins(%arga: tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %x: f64): ^bb(%a: f64, %x: f64):
%1 = arith.mulf %a, %s : f64 %1 = arith.mulf %a, %s : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Scales a sparse vector in place. // Scales a sparse vector in place.
func.func @vector_scale_inplace(%argx: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { func.func @vector_scale_inplace(%argx: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%s = arith.constant 2.0 : f64 %s = arith.constant 2.0 : f64
%0 = linalg.generic #trait_scale_inpl %0 = linalg.generic #trait_scale_inpl
outs(%argx: tensor<?xf64, #SparseVector>) { inits(%argx: tensor<?xf64, #SparseVector>) {
^bb(%x: f64): ^bb(%x: f64):
%1 = arith.mulf %x, %s : f64 %1 = arith.mulf %x, %s : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Adds two sparse vectors into a new sparse vector. // Adds two sparse vectors into a new sparse vector.
func.func @vector_add(%arga: tensor<?xf64, #SparseVector>, func.func @vector_add(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.addf %a, %b : f64 %1 = arith.addf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Multiplies two sparse vectors into a new sparse vector. // Multiplies two sparse vectors into a new sparse vector.
func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>, func.func @vector_mul(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> { %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #SparseVector>) { inits(%xv: tensor<?xf64, #SparseVector>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.mulf %a, %b : f64 %1 = arith.mulf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #SparseVector> } -> tensor<?xf64, #SparseVector>
return %0 : tensor<?xf64, #SparseVector> return %0 : tensor<?xf64, #SparseVector>
} }
// Multiplies two sparse vectors into a new "annotated" dense vector. // Multiplies two sparse vectors into a new "annotated" dense vector.
func.func @vector_mul_d(%arga: tensor<?xf64, #SparseVector>, func.func @vector_mul_d(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #DenseVector> { %argb: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #DenseVector> {
%c = arith.constant 0 : index %c = arith.constant 0 : index
%d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector> %d = tensor.dim %arga, %c : tensor<?xf64, #SparseVector>
%xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #DenseVector> %xv = bufferization.alloc_tensor(%d) : tensor<?xf64, #DenseVector>
%0 = linalg.generic #trait_op %0 = linalg.generic #trait_op
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
outs(%xv: tensor<?xf64, #DenseVector>) { inits(%xv: tensor<?xf64, #DenseVector>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.mulf %a, %b : f64 %1 = arith.mulf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<?xf64, #DenseVector> } -> tensor<?xf64, #DenseVector>
return %0 : tensor<?xf64, #DenseVector> return %0 : tensor<?xf64, #DenseVector>
} }
// Sum reduces dot product of two sparse vectors. // Sum reduces dot product of two sparse vectors.
func.func @vector_dotprod(%arga: tensor<?xf64, #SparseVector>, func.func @vector_dotprod(%arga: tensor<?xf64, #SparseVector>,
%argb: tensor<?xf64, #SparseVector>, %argb: tensor<?xf64, #SparseVector>,
%argx: tensor<f64>) -> tensor<f64> { %argx: tensor<f64>) -> tensor<f64> {
%0 = linalg.generic #trait_dot %0 = linalg.generic #trait_dot
ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>) ins(%arga, %argb: tensor<?xf64, #SparseVector>, tensor<?xf64, #SparseVector>)
outs(%argx: tensor<f64>) { inits(%argx: tensor<f64>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.mulf %a, %b : f64 %1 = arith.mulf %a, %b : f64
%2 = arith.addf %x, %1 : f64 %2 = arith.addf %x, %1 : f64
linalg.yield %2 : f64 linalg.yield %2 : f64
} -> tensor<f64> } -> tensor<f64>
return %0 : tensor<f64> return %0 : tensor<f64>
} }
▲ Show 20 Lines • Show All 97 Lines • Show Last 20 Lines

mlir/test/Integration/Dialect/SparseTensor/python/test_elementwise_add_sparse_output.py

Show All 32 Lines#trait_add_elt = {
doc = "X(i,j) = A(i,j) + B(i,j)" doc = "X(i,j) = A(i,j) + B(i,j)"
} }
func.func @sparse_add_elt( func.func @sparse_add_elt(
%arga: tensor<3x4xf64, #DCSR>, %argb: tensor<3x4xf64, #DCSR>) -> tensor<3x4xf64, #DCSR> { %arga: tensor<3x4xf64, #DCSR>, %argb: tensor<3x4xf64, #DCSR>) -> tensor<3x4xf64, #DCSR> {
%argx = bufferization.alloc_tensor() : tensor<3x4xf64, #DCSR> %argx = bufferization.alloc_tensor() : tensor<3x4xf64, #DCSR>
%0 = linalg.generic #trait_add_elt %0 = linalg.generic #trait_add_elt
ins(%arga, %argb: tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>) ins(%arga, %argb: tensor<3x4xf64, #DCSR>, tensor<3x4xf64, #DCSR>)
outs(%argx: tensor<3x4xf64, #DCSR>) { inits(%argx: tensor<3x4xf64, #DCSR>) {
^bb(%a: f64, %b: f64, %x: f64): ^bb(%a: f64, %b: f64, %x: f64):
%1 = arith.addf %a, %b : f64 %1 = arith.addf %a, %b : f64
linalg.yield %1 : f64 linalg.yield %1 : f64
} -> tensor<3x4xf64, #DCSR> } -> tensor<3x4xf64, #DCSR>
return %0 : tensor<3x4xf64, #DCSR> return %0 : tensor<3x4xf64, #DCSR>
} }
func.func @main(%ad: tensor<3x4xf64>, %bd: tensor<3x4xf64>) -> tensor<3x4xf64, #DCSR> func.func @main(%ad: tensor<3x4xf64>, %bd: tensor<3x4xf64>) -> tensor<3x4xf64, #DCSR>
▲ Show 20 Lines • Show All 56 Lines • Show Last 20 Lines

mlir/test/Interfaces/TilingInterface/lower-to-loops-using-interface.mlir

// RUN: mlir-opt -test-tiling-interface=lower-to-scalar-using-scf-for -split-input-file %s | FileCheck %s // RUN: mlir-opt -test-tiling-interface=lower-to-scalar-using-scf-for -split-input-file %s | FileCheck %s
func.func @gemm(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, func.func @gemm(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>,
%arg2 : memref<?x?xf32>) { %arg2 : memref<?x?xf32>) {
linalg.matmul ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) linalg.matmul ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>)
outs(%arg2 : memref<?x?xf32>) inits(%arg2 : memref<?x?xf32>)
return return
} }
// CHECK-LABEL: func @gemm // CHECK-LABEL: func @gemm
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?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
Show All 14 Lines
func.func @indexed_generic(%arg0 : memref<200x300xi32>, %arg1 : memref<300xi16>, func.func @indexed_generic(%arg0 : memref<200x300xi32>, %arg1 : memref<300xi16>,
%arg2 : memref<200xi8>, %arg3 : memref<300x200xi64>) { %arg2 : memref<200xi8>, %arg3 : memref<300x200xi64>) {
linalg.generic { linalg.generic {
indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>, indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1)>,
affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>], affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
ins(%arg0, %arg1, %arg2 : memref<200x300xi32>, memref<300xi16>, memref<200xi8>) ins(%arg0, %arg1, %arg2 : memref<200x300xi32>, memref<300xi16>, memref<200xi8>)
outs(%arg3 : memref<300x200xi64>) { inits(%arg3 : memref<300x200xi64>) {
^bb0(%b0 : i32, %b1 : i16, %b2 : i8, %b3 : i64): ^bb0(%b0 : i32, %b1 : i16, %b2 : i8, %b3 : i64):
%0 = linalg.index 0 : index %0 = linalg.index 0 : index
%1 = arith.index_cast %0 : index to i16 %1 = arith.index_cast %0 : index to i16
%2 = arith.muli %b1, %1 : i16 %2 = arith.muli %b1, %1 : i16
%3 = linalg.index 1 : index %3 = linalg.index 1 : index
%4 = arith.index_cast %3 : index to i8 %4 = arith.index_cast %3 : index to i8
%5 = arith.muli %b2, %4 : i8 %5 = arith.muli %b2, %4 : i8
%6 = arith.extsi %2 : i16 to i32 %6 = arith.extsi %2 : i16 to i32
Show All 33 Lines
// ----- // -----
func.func @conv_strides_and_dilation(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>, func.func @conv_strides_and_dilation(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>,
%arg2 : memref<?x?x?x?xf32>) { %arg2 : memref<?x?x?x?xf32>) {
linalg.conv_2d_nhwc_hwcf { linalg.conv_2d_nhwc_hwcf {
strides = dense<[1, 2]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>,
dilations = dense<[3, 4]> : tensor<2xi64>} dilations = dense<[3, 4]> : tensor<2xi64>}
ins(%arg0, %arg1 : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>) ins(%arg0, %arg1 : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
outs(%arg2 : memref<?x?x?x?xf32>) inits(%arg2 : memref<?x?x?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1 + d4 * 3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1 + d4 * 3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d2 * 2 + d5 * 4)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d2 * 2 + d5 * 4)>
// CHECK: func @conv_strides_and_dilation( // CHECK: func @conv_strides_and_dilation(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
Show All 27 Lines
// ----- // -----
func.func @pool_strides_and_dilation(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?xf32>, func.func @pool_strides_and_dilation(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?xf32>,
%arg2 : memref<?x?x?x?xf32>) { %arg2 : memref<?x?x?x?xf32>) {
linalg.pooling_nhwc_max { linalg.pooling_nhwc_max {
strides = dense<[1, 2]> : tensor<2xi64>, strides = dense<[1, 2]> : tensor<2xi64>,
dilations = dense<[3, 4]> : tensor<2xi64>} dilations = dense<[3, 4]> : tensor<2xi64>}
ins(%arg0, %arg1 : memref<?x?x?x?xf32>, memref<?x?xf32>) ins(%arg0, %arg1 : memref<?x?x?x?xf32>, memref<?x?xf32>)
outs(%arg2 : memref<?x?x?x?xf32>) inits(%arg2 : memref<?x?x?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1 + d4 * 3)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d1 + d4 * 3)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2 * 2 + d5 * 4)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2 * 2 + d5 * 4)>
// CHECK: func @pool_strides_and_dilation // CHECK: func @pool_strides_and_dilation
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: memref<?x?x?x?xf32>
Show All 20 Lines
// CHECK: %[[T10:.+]] = arith.maxf %[[T9]], %[[T8]] // CHECK: %[[T10:.+]] = arith.maxf %[[T9]], %[[T8]]
// CHECK: memref.store %[[T10]], %[[ARG2]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]] // CHECK: memref.store %[[T10]], %[[ARG2]][%[[IV0]], %[[IV1]], %[[IV2]], %[[IV3]]]
// ----- // -----
func.func @map(%lhs: memref<64xf32>, func.func @map(%lhs: memref<64xf32>,
%rhs: memref<64xf32>, %out: memref<64xf32>) { %rhs: memref<64xf32>, %out: memref<64xf32>) {
linalg.map ins(%lhs, %rhs : memref<64xf32>, memref<64xf32>) linalg.map ins(%lhs, %rhs : memref<64xf32>, memref<64xf32>)
outs(%out : memref<64xf32>) inits(%out : memref<64xf32>)
(%in: f32, %in_0: f32) { (%in: f32, %in_0: f32) {
%0 = arith.addf %in, %in_0 : f32 %0 = arith.addf %in, %in_0 : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
// CHECK-LABEL: func.func @map( // CHECK-LABEL: func.func @map(
// CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<64xf32>, // CHECK-SAME: %[[LHS:[a-zA-Z0-9]+]]: memref<64xf32>,
Show All 10 Lines
// CHECK: %[[ADD:.*]] = arith.addf %[[LHS_ELEM]], %[[RHS_ELEM]] // CHECK: %[[ADD:.*]] = arith.addf %[[LHS_ELEM]], %[[RHS_ELEM]]
// CHECK: memref.store %[[ADD]], %[[OUT]][%[[I]]] // CHECK: memref.store %[[ADD]], %[[OUT]][%[[I]]]
// ----- // -----
func.func @transpose(%arg0: memref<16x32x64xf32>, func.func @transpose(%arg0: memref<16x32x64xf32>,
%arg1: memref<32x64x16xf32>) { %arg1: memref<32x64x16xf32>) {
linalg.transpose ins(%arg0 : memref<16x32x64xf32>) linalg.transpose ins(%arg0 : memref<16x32x64xf32>)
outs(%arg1 : memref<32x64x16xf32>) permutation = [1, 2, 0] inits(%arg1 : memref<32x64x16xf32>) permutation = [1, 2, 0]
return return
} }
// CHECK-LABEL: func.func @transpose( // CHECK-LABEL: func.func @transpose(
// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: memref<16x32x64xf32>, // CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: memref<16x32x64xf32>,
// CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: memref<32x64x16xf32>) // CHECK-SAME: %[[OUT:[a-zA-Z0-9]+]]: memref<32x64x16xf32>)
// 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: %[[C16:.*]] = arith.constant 16 : index // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
// CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index // CHECK-DAG: %[[C32:.*]] = arith.constant 32 : index
// CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index // CHECK-DAG: %[[C64:.*]] = arith.constant 64 : index
// CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C16]] step %[[C1]] { // CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C16]] step %[[C1]] {
// CHECK: scf.for %[[J:.*]] = %[[C0]] to %[[C32]] step %[[C1]] { // CHECK: scf.for %[[J:.*]] = %[[C0]] to %[[C32]] step %[[C1]] {
// CHECK: scf.for %[[K:.*]] = %[[C0]] to %[[C64]] step %[[C1]] { // CHECK: scf.for %[[K:.*]] = %[[C0]] to %[[C64]] step %[[C1]] {
// CHECK: %[[ELEM:.*]] = memref.load %[[OUT]][%[[J]], %[[K]], %[[I]]] // CHECK: %[[ELEM:.*]] = memref.load %[[OUT]][%[[J]], %[[K]], %[[I]]]
// CHECK: memref.store %[[ELEM]], %[[OUT]][%[[J]], %[[K]], %[[I]]] // CHECK: memref.store %[[ELEM]], %[[OUT]][%[[J]], %[[K]], %[[I]]]
// ----- // -----
func.func @reduce(%arg0: memref<16x32x64xf32>, func.func @reduce(%arg0: memref<16x32x64xf32>,
%arg1: memref<16x64xf32>) { %arg1: memref<16x64xf32>) {
linalg.reduce ins(%arg0 : memref<16x32x64xf32>) linalg.reduce ins(%arg0 : memref<16x32x64xf32>)
outs(%arg1 : memref<16x64xf32>) dimensions = [1] inits(%arg1 : memref<16x64xf32>) dimensions = [1]
(%in: f32, %init: f32) { (%in: f32, %init: f32) {
%0 = arith.addf %in, %init : f32 %0 = arith.addf %in, %init : f32
linalg.yield %0 : f32 linalg.yield %0 : f32
} }
return return
} }
// CHECK-LABEL: func.func @reduce( // CHECK-LABEL: func.func @reduce(
// CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: memref<16x32x64xf32>, // CHECK-SAME: %[[IN:[a-zA-Z0-9]+]]: memref<16x32x64xf32>,
Show All 15 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 = tensor.empty(%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) inits(%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> inits(%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:.+]] = tensor.empty // 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: inits(%[[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: inits(%[[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 = tensor.empty(%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) inits(%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> inits(%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>) inits(%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:.+]] = tensor.empty // 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
// CHECK-SAME: outs(%[[INIT_TILE]] : // CHECK-SAME: inits(%[[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: inits(%[[FILL_TILE]] :
// CHECK-DAG: %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]] // CHECK-DAG: %[[BIAS_TILE:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]]]
// CHECK-DAG: %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV0]], %[[IV1]]]
// CHECK: %[[GENERIC_TILE:.+]] = linalg.generic // CHECK: %[[GENERIC_TILE:.+]] = linalg.generic
// CHECK-SAME: ins(%[[GEMM_TILE]], %[[BIAS_TILE]] : // CHECK-SAME: ins(%[[GEMM_TILE]], %[[BIAS_TILE]] :
// CHECK-SAME: outs(%[[OUTS_TILE]] : // CHECK-SAME: inits(%[[OUTS_TILE]] :
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]] // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV0]], %[[IV1]]]
// CHECK: scf.yield %[[INSERT]] // CHECK: scf.yield %[[INSERT]]
// ----- // -----
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 = tensor.empty(%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) inits(%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>) inits(%fill0 : tensor<?x?xf32>) -> tensor<?x?xf32>
%d2 = tensor.dim %rhs1, %c1 : tensor<?x?xf32> %d2 = tensor.dim %rhs1, %c1 : tensor<?x?xf32>
%init1 = tensor.empty(%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) inits(%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>) inits(%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:.+]] = tensor.empty(%[[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:.+]] = tensor.empty(%[[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: inits(%[[INIT0_TILE]] :
// CHECK: %[[GEMM0_TILE:.+]] = linalg.matmul // CHECK: %[[GEMM0_TILE:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[LHS0_TILE]], %[[RHS0_TILE]] : // CHECK-SAME: ins(%[[LHS0_TILE]], %[[RHS0_TILE]] :
// CHECK-SAME: outs(%[[FILL0_TILE]] : // CHECK-SAME: inits(%[[FILL0_TILE]] :
// CHECK-DAG: %[[RHS1_TILE:.+]] = tensor.extract_slice %[[RHS1]][0, 0] // CHECK-DAG: %[[RHS1_TILE:.+]] = tensor.extract_slice %[[RHS1]][0, 0]
// CHECK-DAG: %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ITERARG]][%[[IV]], 0] // CHECK-DAG: %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ITERARG]][%[[IV]], 0]
// CHECK: %[[FILL1_TILE:.+]] = linalg.fill // CHECK: %[[FILL1_TILE:.+]] = linalg.fill
// CHECK-SAME: outs(%[[INIT1_TILE]] : // CHECK-SAME: inits(%[[INIT1_TILE]] :
// CHECK: %[[GEMM1_TILE:.+]] = linalg.matmul // CHECK: %[[GEMM1_TILE:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[GEMM0_TILE]], %[[RHS1_TILE]] : // CHECK-SAME: ins(%[[GEMM0_TILE]], %[[RHS1_TILE]] :
// CHECK-SAME: outs(%[[FILL1_TILE]] : // CHECK-SAME: inits(%[[FILL1_TILE]] :
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITERARG]][%[[IV]], 0] // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GEMM1_TILE]] into %[[ITERARG]][%[[IV]], 0]
// CHECK: scf.yield %[[INSERT]] // CHECK: scf.yield %[[INSERT]]
// ----- // -----
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 = tensor.empty(%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) inits(%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> inits(%fill : tensor<?x?xf32>) -> tensor<?x?xf32>
%init1 = tensor.empty(%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>) inits(%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:.+]] = tensor.empty(%[[D0]], %[[D1]]) // CHECK-DAG: %[[INIT0:.+]] = tensor.empty(%[[D0]], %[[D1]])
// CHECK-DAG: %[[INIT1:.+]] = tensor.empty(%[[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
// CHECK-SAME: outs(%[[INIT0_TILE]] : // CHECK-SAME: inits(%[[INIT0_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: inits(%[[FILL_TILE]] :
// CHECK-DAG: %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]] // CHECK-DAG: %[[OUTS_TILE:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
// CHECK: %[[GENERIC_TILE:.+]] = linalg.generic // CHECK: %[[GENERIC_TILE:.+]] = linalg.generic
// CHECK-SAME: ins(%[[GEMM_TILE]] : // CHECK-SAME: ins(%[[GEMM_TILE]] :
// CHECK-SAME: outs(%[[OUTS_TILE]] : // CHECK-SAME: inits(%[[OUTS_TILE]] :
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]] // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]]
// CHECK: scf.yield %[[INSERT]] // CHECK: scf.yield %[[INSERT]]
// ----- // -----
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 = tensor.empty(%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) inits(%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> inits(%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>) inits(%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:.+]] = tensor.empty // 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
// CHECK-SAME: outs(%[[INIT_TILE]] : // CHECK-SAME: inits(%[[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: inits(%[[FILL_TILE]] :
// CHECK: %[[INIT_TILE_2:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]] // CHECK: %[[INIT_TILE_2:.+]] = tensor.extract_slice %[[ITERARG1]][%[[IV1]], %[[IV0]]]
// CHECK: %[[GENERIC_TILE:.+]] = linalg.generic // CHECK: %[[GENERIC_TILE:.+]] = linalg.generic
// CHECK-SAME: ins(%[[GEMM_TILE]] : // CHECK-SAME: ins(%[[GEMM_TILE]] :
// CHECK-SAME: outs(%[[INIT_TILE_2]] : // CHECK-SAME: inits(%[[INIT_TILE_2]] :
// CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]] // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[GENERIC_TILE]] into %[[ITERARG1]][%[[IV1]], %[[IV0]]]
// CHECK: scf.yield %[[INSERT]] // CHECK: scf.yield %[[INSERT]]
// ----- // -----
func.func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, func.func @matmul_plus_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
%arg2: tensor<?x?xf32>) -> 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 %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> inits(%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 = tensor.empty(%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>) { inits(%5 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) : ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
%7 = arith.addf %arg3, %arg4 : f32 %7 = arith.addf %arg3, %arg4 : f32
linalg.yield %7 : f32 linalg.yield %7 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %6 : tensor<?x?xf32> return %6 : tensor<?x?xf32>
} }
// This fuses as expected but the gemm operation is inlined twice. It should be CSE-d but isnt today. // This fuses as expected but the gemm operation is inlined twice. It should be CSE-d but isnt today.
// CHECK: func @matmul_plus_matmul // CHECK: func @matmul_plus_matmul
// 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<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
// CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]] // CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}}) // CHECK-SAME: iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}})
// CHECK: %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]] // CHECK: %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]]) // CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]])
// CHECK-DAG: %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[ST_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[ST_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]]
// CHECK: %[[LHS:.+]] = linalg.matmul // CHECK: %[[LHS:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[ST_ARG0]], %[[ST_ARG1]] : // CHECK-SAME: ins(%[[ST_ARG0]], %[[ST_ARG1]] :
// CHECK-SAME: outs(%[[ST_ARG2]] : // CHECK-SAME: inits(%[[ST_ARG2]] :
// CHECK-DAG: %[[ST_ARG0_1:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[ST_ARG0_1:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[ST_ARG1_1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[ST_ARG1_1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[ST_ARG2_1:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[ST_ARG2_1:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]]
// CHECK: %[[RHS:.+]] = linalg.matmul // CHECK: %[[RHS:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[ST_ARG0_1]], %[[ST_ARG1_1]] : // CHECK-SAME: ins(%[[ST_ARG0_1]], %[[ST_ARG1_1]] :
// CHECK-SAME: outs(%[[ST_ARG2_1]] : // CHECK-SAME: inits(%[[ST_ARG2_1]] :
// CHECK: %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]]] // CHECK: %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: %[[ST_RESULT:.+]] = linalg.generic // CHECK: %[[ST_RESULT:.+]] = linalg.generic
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : // CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[ST_ARG6]] : // CHECK-SAME: inits(%[[ST_ARG6]] :
// CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[ST_RESULT]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[ST_RESULT]]
// CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]] // CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[UPDATE]]
// CHECK: scf.yield %[[YIELD]] // CHECK: scf.yield %[[YIELD]]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
func.func @matmul_plus_transpose_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, func.func @matmul_plus_transpose_matmul(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
%arg2: tensor<?x?xf32>) -> 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 %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> inits(%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 = tensor.empty(%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>) { inits(%5 : tensor<?x?xf32>) {
^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) : ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32) :
%7 = arith.addf %arg3, %arg4 : f32 %7 = arith.addf %arg3, %arg4 : f32
linalg.yield %7 : f32 linalg.yield %7 : f32
} -> tensor<?x?xf32> } -> tensor<?x?xf32>
return %6 : tensor<?x?xf32> return %6 : tensor<?x?xf32>
} }
// CHECK: func @matmul_plus_transpose_matmul // CHECK: func @matmul_plus_transpose_matmul
// 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<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
// CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]] // CHECK: %[[RESULT:.+]] = scf.for %[[IV0:[a-zA-Z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}}) // CHECK-SAME: iter_args(%[[ARG4:.+]] = %{{[a-zA-Z0-9_]+}})
// CHECK: %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]] // CHECK: %[[YIELD:.+]] = scf.for %[[IV1:[a-zA-Z0-9_]+]]
// CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]]) // CHECK-SAME: iter_args(%[[ARG6:.+]] = %[[ARG4]])
// CHECK-DAG: %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0] // CHECK-DAG: %[[ST_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV0]], 0]
// CHECK-DAG: %[[ST_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]] // CHECK-DAG: %[[ST_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV1]]]
// CHECK-DAG: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]] // CHECK-DAG: %[[ST_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV0]], %[[IV1]]]
// CHECK: %[[LHS:.+]] = linalg.matmul // CHECK: %[[LHS:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[ST_ARG0]], %[[ST_ARG1]] // CHECK-SAME: ins(%[[ST_ARG0]], %[[ST_ARG1]]
// CHECK-SAME: : tensor<?x?xf32>, tensor<?x?xf32>) // CHECK-SAME: : tensor<?x?xf32>, tensor<?x?xf32>)
// CHECK-SAME: outs(%[[ST_ARG2]] : tensor<?x?xf32>) // CHECK-SAME: inits(%[[ST_ARG2]] : tensor<?x?xf32>)
// CHECK-DAG: %[[STR_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0] // CHECK-DAG: %[[STR_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV1]], 0]
// CHECK-DAG: %[[STR_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV0]]] // CHECK-DAG: %[[STR_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, %[[IV0]]]
// CHECK-DAG: %[[STR_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]], %[[IV0]]] // CHECK-DAG: %[[STR_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV1]], %[[IV0]]]
// CHECK: %[[RHS:.+]] = linalg.matmul // CHECK: %[[RHS:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[STR_ARG0]], %[[STR_ARG1]] : // CHECK-SAME: ins(%[[STR_ARG0]], %[[STR_ARG1]] :
// CHECK-SAME: outs(%[[STR_ARG2]] : // CHECK-SAME: inits(%[[STR_ARG2]] :
// CHECK: %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]]] // CHECK: %[[ST_ARG6:.+]] = tensor.extract_slice %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: %[[ST_RESULT:.+]] = linalg.generic // CHECK: %[[ST_RESULT:.+]] = linalg.generic
// CHECK-SAME: ins(%[[LHS]], %[[RHS]] : // CHECK-SAME: ins(%[[LHS]], %[[RHS]] :
// CHECK-SAME: outs(%[[ST_ARG6]] : // CHECK-SAME: inits(%[[ST_ARG6]] :
// CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[ST_RESULT]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[ST_RESULT]]
// CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]] // CHECK-SAME: into %[[ARG6]][%[[IV0]], %[[IV1]]]
// CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[UPDATE]]
// CHECK: scf.yield %[[YIELD]] // CHECK: scf.yield %[[YIELD]]
// CHECK: return %[[RESULT]] // CHECK: return %[[RESULT]]
// ----- // -----
func.func @matmul_sequence_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, func.func @matmul_sequence_fusion(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>,
%arg2: tensor<?x?xf32>, %arg3: tensor<?x?xf32>, %arg4: tensor<?x?xf32>, %arg2: tensor<?x?xf32>, %arg3: tensor<?x?xf32>, %arg4: tensor<?x?xf32>,
%arg5: tensor<?x?xf32>, %arg6: tensor<?x?xf32>) -> tensor<?x?xf32> { %arg5: tensor<?x?xf32>, %arg6: tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N0] * [N0, N1] inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N0] * [N0, N1]
%1 = linalg.matmul ins(%0, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>) %1 = linalg.matmul ins(%0, %arg3 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg4 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N1] * [N1, N2] inits(%arg4 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N1] * [N1, N2]
%2 = linalg.matmul %2 = linalg.matmul
{__internal_linalg_transform__ = "gemm_sequence_fusion"} {__internal_linalg_transform__ = "gemm_sequence_fusion"}
ins(%1, %arg5 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%1, %arg5 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg6 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N2] * [N2, N3] inits(%arg6 : tensor<?x?xf32>) -> tensor<?x?xf32> // [M, N2] * [N2, N3]
return %2 : tensor<?x?xf32> return %2 : tensor<?x?xf32>
} }
// CHECK: #[[MAP:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK: #[[MAP:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK: func @matmul_sequence_fusion( // CHECK: func @matmul_sequence_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<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
Show All 12 Lines
// CHECK-DAG: %[[N3:.+]] = tensor.dim %[[ARG5]], %[[C1]] // CHECK-DAG: %[[N3:.+]] = tensor.dim %[[ARG5]], %[[C1]]
// CHECK: %[[R0:.+]] = scf.for %[[IV:[a-zA-Z0-9_]+]] = // CHECK: %[[R0:.+]] = scf.for %[[IV:[a-zA-Z0-9_]+]] =
// CHECK-SAME: iter_args(%[[ARG8:.+]] = %[[ARG6]]) -> (tensor<?x?xf32>) { // CHECK-SAME: iter_args(%[[ARG8:.+]] = %[[ARG6]]) -> (tensor<?x?xf32>) {
// CHECK-DAG: %[[TILE_M:.+]] = affine.min #[[MAP]](%[[IV]])[%{{.+}}, %[[M]]] // CHECK-DAG: %[[TILE_M:.+]] = affine.min #[[MAP]](%[[IV]])[%{{.+}}, %[[M]]]
// CHECK-DAG: %[[SLICE_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV]], 0] [%[[TILE_M]], %[[N0]]] // CHECK-DAG: %[[SLICE_ARG0:.+]] = tensor.extract_slice %[[ARG0]][%[[IV]], 0] [%[[TILE_M]], %[[N0]]]
// CHECK-DAG: %[[SLICE_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, 0] [%[[N0]], %[[N1]]] // CHECK-DAG: %[[SLICE_ARG1:.+]] = tensor.extract_slice %[[ARG1]][0, 0] [%[[N0]], %[[N1]]]
// CHECK-DAG: %[[SLICE_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV]], 0] [%[[TILE_M]], %[[N1]]] // CHECK-DAG: %[[SLICE_ARG2:.+]] = tensor.extract_slice %[[ARG2]][%[[IV]], 0] [%[[TILE_M]], %[[N1]]]
// CHECK-DAG: %[[TILE_GEMM1:.+]] = linalg.matmul ins(%[[SLICE_ARG0]], %[[SLICE_ARG1]] : // CHECK-DAG: %[[TILE_GEMM1:.+]] = linalg.matmul ins(%[[SLICE_ARG0]], %[[SLICE_ARG1]] :
// CHECK-SAME: outs(%[[SLICE_ARG2]] : // CHECK-SAME: inits(%[[SLICE_ARG2]] :
// CHECK-DAG: %[[SLICE_ARG3:.+]] = tensor.extract_slice %[[ARG3]][0, 0] [%[[N1]], %[[N2]]] // CHECK-DAG: %[[SLICE_ARG3:.+]] = tensor.extract_slice %[[ARG3]][0, 0] [%[[N1]], %[[N2]]]
// CHECK-DAG: %[[SLICE_ARG4:.+]] = tensor.extract_slice %[[ARG4]][%[[IV]], 0] [%[[TILE_M]], %[[N2]]] // CHECK-DAG: %[[SLICE_ARG4:.+]] = tensor.extract_slice %[[ARG4]][%[[IV]], 0] [%[[TILE_M]], %[[N2]]]
// CHECK-DAG: %[[TILE_GEMM2:.+]] = linalg.matmul ins(%[[TILE_GEMM1]], %[[SLICE_ARG3]] : // CHECK-DAG: %[[TILE_GEMM2:.+]] = linalg.matmul ins(%[[TILE_GEMM1]], %[[SLICE_ARG3]] :
// CHECK-SAME: outs(%[[SLICE_ARG4]] : // CHECK-SAME: inits(%[[SLICE_ARG4]] :
// CHECK-DAG: %[[SLICE_ARG5:.+]] = tensor.extract_slice %[[ARG5]][0, 0] [%[[N2]], %[[N3]]] // CHECK-DAG: %[[SLICE_ARG5:.+]] = tensor.extract_slice %[[ARG5]][0, 0] [%[[N2]], %[[N3]]]
// CHECK-DAG: %[[SLICE_ARG6:.+]] = tensor.extract_slice %[[ARG8]][%[[IV]], 0] [%[[TILE_M]], %[[N3]]] // CHECK-DAG: %[[SLICE_ARG6:.+]] = tensor.extract_slice %[[ARG8]][%[[IV]], 0] [%[[TILE_M]], %[[N3]]]
// CHECK-DAG: %[[TILE_GEMM3:.+]] = linalg.matmul // CHECK-DAG: %[[TILE_GEMM3:.+]] = linalg.matmul
// CHECK-SAME: ins(%[[TILE_GEMM2]], %[[SLICE_ARG5]] : // CHECK-SAME: ins(%[[TILE_GEMM2]], %[[SLICE_ARG5]] :
// CHECK-SAME: outs(%[[SLICE_ARG6]] : // CHECK-SAME: inits(%[[SLICE_ARG6]] :
// CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[TILE_GEMM3]] into %[[ARG8]][%[[IV]], 0] [%[[TILE_M]], %[[N3]]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[TILE_GEMM3]] into %[[ARG8]][%[[IV]], 0] [%[[TILE_M]], %[[N3]]]
// CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[UPDATE]]

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

// RUN: mlir-opt -test-tiling-interface=tile-using-scf-for -split-input-file %s | FileCheck %s // RUN: mlir-opt -test-tiling-interface=tile-using-scf-for -split-input-file %s | FileCheck %s
func.func @simple_matmul(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, func.func @simple_matmul(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> { %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.matmul {__internal_linalg_transform__ = "simple_gemm"} %0 = linalg.matmul {__internal_linalg_transform__ = "simple_gemm"}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
// CHECK: func.func @simple_matmul( // CHECK: func.func @simple_matmul(
// 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<?x?xf32> // CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<?x?xf32>
Show All 13 Lines
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]]
// CHECK-SAME: [%[[IV0]], 0] [%[[TS_Y]], %[[K]]] [1, 1] // CHECK-SAME: [%[[IV0]], 0] [%[[TS_Y]], %[[K]]] [1, 1]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]]
// CHECK-SAME: [0, %[[IV1]]] [%[[K]], %[[TS_X]]] [1, 1] // CHECK-SAME: [0, %[[IV1]]] [%[[K]], %[[TS_X]]] [1, 1]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT1]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT1]]
// CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_Y]], %[[TS_X]]] [1, 1] // CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_Y]], %[[TS_X]]] [1, 1]
// 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(%[[INIT_TILE]] : // CHECK-SAME: inits(%[[INIT_TILE]] :
// CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[INIT1]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[INIT1]]
// CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_Y]], %[[TS_X]]] [1, 1] // CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_Y]], %[[TS_X]]] [1, 1]
// CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[UPDATE]]
// CHECK: scf.yield %[[INNER]] // CHECK: scf.yield %[[INNER]]
// CHECK: return %[[OUTER]] // CHECK: return %[[OUTER]]
// ----- // -----
func.func @simple_matmul_memref(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, func.func @simple_matmul_memref(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>,
%arg2 : memref<?x?xf32>) { %arg2 : memref<?x?xf32>) {
linalg.matmul {__internal_linalg_transform__ = "simple_gemm_memref"} linalg.matmul {__internal_linalg_transform__ = "simple_gemm_memref"}
ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>)
outs(%arg2 : memref<?x?xf32>) inits(%arg2 : memref<?x?xf32>)
return return
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)>
// CHECK: func.func @simple_matmul_memref( // CHECK: func.func @simple_matmul_memref(
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<?x?xf32>
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32> // CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: memref<?x?xf32>
Show All 15 Lines
// CHECK-DAG: %[[LHS_TILE:.+]] = memref.subview %[[ARG0]] // CHECK-DAG: %[[LHS_TILE:.+]] = memref.subview %[[ARG0]]
// CHECK-SAME: [%[[IV0]], %[[IV2]]] [%[[TS_M]], %[[TS_K]]] [1, 1] // CHECK-SAME: [%[[IV0]], %[[IV2]]] [%[[TS_M]], %[[TS_K]]] [1, 1]
// CHECK-DAG: %[[RHS_TILE:.+]] = memref.subview %[[ARG1]] // CHECK-DAG: %[[RHS_TILE:.+]] = memref.subview %[[ARG1]]
// CHECK-SAME: [%[[IV2]], %[[IV1]]] [%[[TS_K]], %[[TS_N]]] [1, 1] // CHECK-SAME: [%[[IV2]], %[[IV1]]] [%[[TS_K]], %[[TS_N]]] [1, 1]
// CHECK-DAG: %[[OUT_TILE:.+]] = memref.subview %[[ARG2]] // CHECK-DAG: %[[OUT_TILE:.+]] = memref.subview %[[ARG2]]
// CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_M]], %[[TS_N]]] [1, 1] // CHECK-SAME: [%[[IV0]], %[[IV1]]] [%[[TS_M]], %[[TS_N]]] [1, 1]
// CHECK: linalg.matmul // CHECK: linalg.matmul
// CHECK-SAME: ins(%[[LHS_TILE]], %[[RHS_TILE]] : // CHECK-SAME: ins(%[[LHS_TILE]], %[[RHS_TILE]] :
// CHECK-SAME: outs(%[[OUT_TILE]] : // CHECK-SAME: inits(%[[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 = tensor.empty() : tensor<128x300x200xf32> %init0 = tensor.empty() : tensor<128x300x200xf32>
%init1 = tensor.empty() : 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>) { inits(%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>)
Show All 12 Lines
// 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]]
// CHECK-SAME: [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1] // CHECK-SAME: [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1]
// CHECK-DAG: %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ARG4]] // CHECK-DAG: %[[INIT1_TILE:.+]] = tensor.extract_slice %[[ARG4]]
// CHECK-SAME: [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1] // CHECK-SAME: [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1]
// CHECK: %[[RESULT_TILE:.+]]:2 = linalg.generic // CHECK: %[[RESULT_TILE:.+]]:2 = linalg.generic
// CHECK-SAME: ins(%[[ARG_TILE]] : // CHECK-SAME: ins(%[[ARG_TILE]] :
// CHECK-SAME: outs(%[[INIT0_TILE]], %[[INIT1_TILE]] : // CHECK-SAME: inits(%[[INIT0_TILE]], %[[INIT1_TILE]] :
// CHECK: %[[UPDATE0:.+]] = tensor.insert_slice %[[RESULT_TILE]]#0 into %[[ARG3]] // CHECK: %[[UPDATE0:.+]] = tensor.insert_slice %[[RESULT_TILE]]#0 into %[[ARG3]]
// CHECK-SAME: [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1] // CHECK-SAME: [%[[IV0]], %[[IV1]], 0] [%[[TS_Y]], 20, 200] [1, 1, 1]
// CHECK: %[[UPDATE1:.+]] = tensor.insert_slice %[[RESULT_TILE]]#1 into %[[ARG4]] // CHECK: %[[UPDATE1:.+]] = tensor.insert_slice %[[RESULT_TILE]]#1 into %[[ARG4]]
// CHECK-SAME: [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1] // CHECK-SAME: [%[[IV1]], %[[IV0]], 0] [20, %[[TS_Y]], 200] [1, 1, 1]
// CHECK: scf.yield %[[UPDATE0]], %[[UPDATE1]] // CHECK: scf.yield %[[UPDATE0]], %[[UPDATE1]]
// CHECK: scf.yield %[[INNER]]#0, %[[INNER]]#1 // CHECK: scf.yield %[[INNER]]#0, %[[INNER]]#1
// CHECK: return %[[OUTER]]#0, %[[OUTER]]#1 // CHECK: return %[[OUTER]]#0, %[[OUTER]]#1
// ----- // -----
func.func @conv2D(%arg0 : tensor<?x?x?x?xf32>, %arg1 : tensor<?x?x?x?xf32>, func.func @conv2D(%arg0 : tensor<?x?x?x?xf32>, %arg1 : tensor<?x?x?x?xf32>,
%arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> { %arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
%0 = linalg.conv_2d_nhwc_hwcf { %0 = linalg.conv_2d_nhwc_hwcf {
strides = dense<[2, 3]> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>,
dilation = dense<[4, 5]> : tensor<2xi64>, dilation = dense<[4, 5]> : tensor<2xi64>,
__internal_linalg_transform__ = "simple_conv"} __internal_linalg_transform__ = "simple_conv"}
ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) ins(%arg0, %arg1 : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
outs(%arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> inits(%arg2 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
return %0 : tensor<?x?x?x?xf32> return %0 : tensor<?x?x?x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)>
// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0)[s0] -> (d0 + s0 * 2 - 2)> // CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0)[s0] -> (d0 + s0 * 2 - 2)>
// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0)[s0] -> (d0 + s0 * 3 - 3)> // CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0)[s0] -> (d0 + s0 * 3 - 3)>
// CHECK: func.func @conv2D( // CHECK: func.func @conv2D(
Show All 29 Lines
// CHECK-SAME: [0, %[[IV0]], %[[IV1]], %[[IV2]]] [%[[N]], %[[TS_H]], %[[TS_W]], %[[TS_C]]] // CHECK-SAME: [0, %[[IV0]], %[[IV1]], %[[IV2]]] [%[[N]], %[[TS_H]], %[[TS_W]], %[[TS_C]]]
// CHECK-DAG: %[[FILTER_TILE:.+]] = tensor.extract_slice %[[FILTER]] // CHECK-DAG: %[[FILTER_TILE:.+]] = tensor.extract_slice %[[FILTER]]
// CHECK-SAME: [%[[IV0]], %[[IV1]], %[[IV2]], 0] [%[[TS_P]], %[[TS_Q]], %[[TS_C]], %[[F]]] // CHECK-SAME: [%[[IV0]], %[[IV1]], %[[IV2]], 0] [%[[TS_P]], %[[TS_Q]], %[[TS_C]], %[[F]]]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT2]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT2]]
// CHECK-SAME: [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]] // CHECK-SAME: [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]]
// CHECK: %[[CONV_TILE:.+]] = linalg.conv_2d_nhwc_hwcf // CHECK: %[[CONV_TILE:.+]] = linalg.conv_2d_nhwc_hwcf
// CHECK-SAME: dilation = dense<[4, 5]> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64> // CHECK-SAME: dilation = dense<[4, 5]> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>
// CHECK-SAME: ins(%[[INPUT_TILE]], %[[FILTER_TILE]] : // CHECK-SAME: ins(%[[INPUT_TILE]], %[[FILTER_TILE]] :
// CHECK-SAME: outs(%[[INIT_TILE]] : // CHECK-SAME: inits(%[[INIT_TILE]] :
// CHECK: tensor.insert_slice %[[CONV_TILE]] into %[[INIT2]] // CHECK: tensor.insert_slice %[[CONV_TILE]] into %[[INIT2]]
// CHECK-SAME: [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]] // CHECK-SAME: [0, 0, 0, 0] [%[[N]], %[[R]], %[[S]], %[[F]]]
// ----- // -----
// CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0, d1) -> (d0 + d1)> // CHECK: #[[$MAP_ADD:.+]] = affine_map<(d0, d1) -> (d0 + d1)>
// CHECK-LABEL: @indexed_semantics // CHECK-LABEL: @indexed_semantics
func.func @indexed_semantics(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> { func.func @indexed_semantics(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>) -> tensor<?x?xf32> {
// Check that we correctly amend "linalg.index" results. // Check that we correctly amend "linalg.index" results.
// CHECK: scf.for %[[I0:.+]] = %{{.*}} to %{{.*}} step %{{.*}} // CHECK: scf.for %[[I0:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
// CHECK: scf.for %[[I1:.+]] = %{{.*}} to %{{.*}} step %{{.*}} // CHECK: scf.for %[[I1:.+]] = %{{.*}} to %{{.*}} step %{{.*}}
%0 = linalg.generic { %0 = 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)>],
iterator_types = ["parallel", "parallel"]} iterator_types = ["parallel", "parallel"]}
{__internal_linalg_transform__ = "indexed_semantics"} {__internal_linalg_transform__ = "indexed_semantics"}
ins(%arg0: tensor<?x?xf32>) ins(%arg0: tensor<?x?xf32>)
outs(%arg1: tensor<?x?xf32>) { inits(%arg1: tensor<?x?xf32>) {
^bb0(%arg2: f32, %arg3: f32): ^bb0(%arg2: f32, %arg3: f32):
// CHECK: %[[INDEX0:.+]] = linalg.index 0 // CHECK: %[[INDEX0:.+]] = linalg.index 0
// CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX0]], %[[I0]]) // CHECK: %[[INDEX0_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX0]], %[[I0]])
%1 = linalg.index 0 : index %1 = linalg.index 0 : index
// CHECK: %[[INDEX1:.+]] = linalg.index 1 // CHECK: %[[INDEX1:.+]] = linalg.index 1
// CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX1]], %[[I1]]) // CHECK: %[[INDEX1_AMENDED:.+]] = affine.apply #[[$MAP_ADD]](%[[INDEX1]], %[[I1]])
%2 = linalg.index 1 : index %2 = linalg.index 1 : index
// CHECK: arith.addi %[[INDEX0_AMENDED]], %[[INDEX1_AMENDED]] // CHECK: arith.addi %[[INDEX0_AMENDED]], %[[INDEX1_AMENDED]]
%3 = arith.addi %1, %2 : index %3 = arith.addi %1, %2 : index
%4 = arith.index_cast %3 : index to i64 %4 = arith.index_cast %3 : index to i64
%5 = arith.uitofp %4 : i64 to f32 %5 = arith.uitofp %4 : i64 to f32
%6 = arith.addf %5, %arg2 : f32 %6 = arith.addf %5, %arg2 : f32
linalg.yield %6 : f32 linalg.yield %6 : f32
} -> (tensor<?x?xf32>) } -> (tensor<?x?xf32>)
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// ----- // -----
func.func @interchange_matmul(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, func.func @interchange_matmul(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> { %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%0 = linalg.matmul {__internal_linalg_transform__ = "gemm_interchange"} %0 = linalg.matmul {__internal_linalg_transform__ = "gemm_interchange"}
ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32> inits(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
return %0 : tensor<?x?xf32> return %0 : tensor<?x?xf32>
} }
// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)> // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0)[s0, s1] -> (20, -d0 + s1)>
// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)> // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0)[s0, s1] -> (30, -d0 + s1)>
// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)> // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0, s1] -> (10, -d0 + s1)>
// CHECK: func.func @interchange_matmul( // CHECK: func.func @interchange_matmul(
// 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>
Show All 18 Lines
// CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]] // CHECK-DAG: %[[LHS_TILE:.+]] = tensor.extract_slice %[[ARG0]]
// CHECK-SAME: [%[[IV2]], %[[IV1]]] [%[[TS_M]], %[[TS_K]]] [1, 1] // CHECK-SAME: [%[[IV2]], %[[IV1]]] [%[[TS_M]], %[[TS_K]]] [1, 1]
// CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]] // CHECK-DAG: %[[RHS_TILE:.+]] = tensor.extract_slice %[[ARG1]]
// CHECK-SAME: [%[[IV1]], %[[IV0]]] [%[[TS_K]], %[[TS_N]]] [1, 1] // CHECK-SAME: [%[[IV1]], %[[IV0]]] [%[[TS_K]], %[[TS_N]]] [1, 1]
// CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT2]] // CHECK-DAG: %[[INIT_TILE:.+]] = tensor.extract_slice %[[INIT2]]
// CHECK-SAME: [%[[IV2]], %[[IV0]]] [%[[TS_M]], %[[TS_N]]] [1, 1] // CHECK-SAME: [%[[IV2]], %[[IV0]]] [%[[TS_M]], %[[TS_N]]] [1, 1]
// 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(%[[INIT_TILE]] : // CHECK-SAME: inits(%[[INIT_TILE]] :
// CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[INIT2]] // CHECK: %[[UPDATE:.+]] = tensor.insert_slice %[[GEMM_TILE]] into %[[INIT2]]
// CHECK-SAME: [%[[IV2]], %[[IV0]]] [%[[TS_M]], %[[TS_N]]] [1, 1] // CHECK-SAME: [%[[IV2]], %[[IV0]]] [%[[TS_M]], %[[TS_N]]] [1, 1]
// CHECK: scf.yield %[[UPDATE]] // CHECK: scf.yield %[[UPDATE]]
// CHECK: scf.yield %[[INNER2]] // CHECK: scf.yield %[[INNER2]]
// CHECK: scf.yield %[[INNER1]] // CHECK: scf.yield %[[INNER1]]
// CHECK: return %[[OUTER]] // CHECK: return %[[OUTER]]

mlir/test/lib/Dialect/Test/TestOps.td

Show First 20 Lines • Show All 2,800 Lines • ▼ Show 20 Linesdef TestLinalgConvOp :
let arguments = (ins Variadic<AnyType>:$inputs, let arguments = (ins Variadic<AnyType>:$inputs,
Variadic<AnyType>:$outputs); Variadic<AnyType>:$outputs);
let results = (outs Variadic<AnyType>:$results); let results = (outs Variadic<AnyType>:$results);
let regions = (region AnyRegion:$region); let regions = (region AnyRegion:$region);
let assemblyFormat = [{ let assemblyFormat = [{
attr-dict (`ins` `(` $inputs^ `:` type($inputs) `)`)? attr-dict (`ins` `(` $inputs^ `:` type($inputs) `)`)?
`outs` `(` $outputs `:` type($outputs) `)` `inits` `(` $outputs `:` type($outputs) `)`
$region (`->` type($results)^)? $region (`->` type($results)^)?
}]; }];
let extraClassDeclaration = [{ let extraClassDeclaration = [{
bool hasIndexSemantics() { return false; } bool hasIndexSemantics() { return false; }
static void regionBuilder(mlir::ImplicitLocOpBuilder &b, mlir::Block &block, static void regionBuilder(mlir::ImplicitLocOpBuilder &b, mlir::Block &block,
mlir::ArrayRef<mlir::NamedAttribute> attrs) { mlir::ArrayRef<mlir::NamedAttribute> attrs) {
▲ Show 20 Lines • Show All 43 Lines • ▼ Show 20 Linesdef TestLinalgFillOp :
let arguments = (ins Variadic<AnyType>:$inputs, let arguments = (ins Variadic<AnyType>:$inputs,
Variadic<AnyType>:$outputs); Variadic<AnyType>:$outputs);
let results = (outs Variadic<AnyType>:$results); let results = (outs Variadic<AnyType>:$results);
let regions = (region AnyRegion:$region); let regions = (region AnyRegion:$region);
let assemblyFormat = [{ let assemblyFormat = [{
attr-dict (`ins` `(` $inputs^ `:` type($inputs) `)`)? attr-dict (`ins` `(` $inputs^ `:` type($inputs) `)`)?
`outs` `(` $outputs `:` type($outputs) `)` `inits` `(` $outputs `:` type($outputs) `)`
$region (`->` type($results)^)? $region (`->` type($results)^)?
}]; }];
let extraClassDeclaration = [{ let extraClassDeclaration = [{
bool hasIndexSemantics() { return false; } bool hasIndexSemantics() { return false; }
static void regionBuilder(mlir::ImplicitLocOpBuilder &b, mlir::Block &block, static void regionBuilder(mlir::ImplicitLocOpBuilder &b, mlir::Block &block,
mlir::ArrayRef<mlir::NamedAttribute> attrs) { mlir::ArrayRef<mlir::NamedAttribute> attrs) {
▲ Show 20 Lines • Show All 230 Lines • Show Last 20 Lines

mlir/test/mlir-cpu-runner/async.mlir

Show All 16 Linesfunc.func @main() {
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
%c1 = arith.constant 1.0 : f32 %c1 = arith.constant 1.0 : f32
%c2 = arith.constant 2.0 : f32 %c2 = arith.constant 2.0 : f32
%c3 = arith.constant 3.0 : f32 %c3 = arith.constant 3.0 : f32
%c4 = arith.constant 4.0 : f32 %c4 = arith.constant 4.0 : f32
%A = memref.alloc() : memref<4xf32> %A = memref.alloc() : memref<4xf32>
linalg.fill ins(%c0 : f32) outs(%A : memref<4xf32>) linalg.fill ins(%c0 : f32) inits(%A : memref<4xf32>)
// CHECK: [0, 0, 0, 0] // CHECK: [0, 0, 0, 0]
%U = memref.cast %A : memref<4xf32> to memref<*xf32> %U = memref.cast %A : memref<4xf32> to memref<*xf32>
call @printMemrefF32(%U): (memref<*xf32>) -> () call @printMemrefF32(%U): (memref<*xf32>) -> ()
// CHECK: Current thread id: [[MAIN:.*]] // CHECK: Current thread id: [[MAIN:.*]]
// CHECK: [1, 0, 0, 0] // CHECK: [1, 0, 0, 0]
memref.store %c1, %A[%i0]: memref<4xf32> memref.store %c1, %A[%i0]: memref<4xf32>
▲ Show 20 Lines • Show All 51 Lines • Show Last 20 Lines

mlir/test/mlir-cpu-runner/sgemm-naive-codegen.mlir

// RUN: mlir-opt -pass-pipeline="func.func(convert-linalg-to-loops,lower-affine,convert-scf-to-cf,convert-arith-to-llvm),convert-vector-to-llvm,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts" %s | mlir-cpu-runner -O3 -e main -entry-point-result=void -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | FileCheck %s // RUN: mlir-opt -pass-pipeline="func.func(convert-linalg-to-loops,lower-affine,convert-scf-to-cf,convert-arith-to-llvm),convert-vector-to-llvm,convert-memref-to-llvm,convert-func-to-llvm,reconcile-unrealized-casts" %s | mlir-cpu-runner -O3 -e main -entry-point-result=void -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | FileCheck %s
func.func @main() { func.func @main() {
%A = memref.alloc() : memref<16x16xf32> %A = memref.alloc() : memref<16x16xf32>
%B = memref.alloc() : memref<16x16xf32> %B = memref.alloc() : memref<16x16xf32>
%C = memref.alloc() : memref<16x16xf32> %C = memref.alloc() : memref<16x16xf32>
%cf1 = arith.constant 1.00000e+00 : f32 %cf1 = arith.constant 1.00000e+00 : f32
linalg.fill ins(%cf1 : f32) outs(%A : memref<16x16xf32>) linalg.fill ins(%cf1 : f32) inits(%A : memref<16x16xf32>)
linalg.fill ins(%cf1 : f32) outs(%B : memref<16x16xf32>) linalg.fill ins(%cf1 : f32) inits(%B : memref<16x16xf32>)
%reps = arith.constant 1 : index %reps = arith.constant 1 : index
%t_start = call @rtclock() : () -> f64 %t_start = call @rtclock() : () -> f64
affine.for %arg0 = 0 to 5 { affine.for %arg0 = 0 to 5 {
linalg.fill ins(%cf1 : f32) outs(%C : memref<16x16xf32>) linalg.fill ins(%cf1 : f32) inits(%C : memref<16x16xf32>)
func.call @sgemm_naive(%A, %B, %C) : (memref<16x16xf32>, memref<16x16xf32>, memref<16x16xf32>) -> () func.call @sgemm_naive(%A, %B, %C) : (memref<16x16xf32>, memref<16x16xf32>, memref<16x16xf32>) -> ()
} }
%t_end = call @rtclock() : () -> f64 %t_end = call @rtclock() : () -> f64
%t = arith.subf %t_end, %t_start : f64 %t = arith.subf %t_end, %t_start : f64
%res = affine.load %C[0, 0]: memref<16x16xf32> %res = affine.load %C[0, 0]: memref<16x16xf32>
vector.print %res: f32 vector.print %res: f32
▲ Show 20 Lines • Show All 51 Lines • Show Last 20 Lines

mlir/test/mlir-cpu-runner/unranked-memref.mlir

Show All 36 Lines
// CHECK-COUNT-4: [1, 1, 1] // CHECK-COUNT-4: [1, 1, 1]
func.func @main() -> () { func.func @main() -> () {
%A = memref.alloc() : memref<10x3xf32, 0> %A = memref.alloc() : memref<10x3xf32, 0>
%f2 = arith.constant 2.00000e+00 : f32 %f2 = arith.constant 2.00000e+00 : f32
%f5 = arith.constant 5.00000e+00 : f32 %f5 = arith.constant 5.00000e+00 : f32
%f10 = arith.constant 10.00000e+00 : f32 %f10 = arith.constant 10.00000e+00 : f32
%V = memref.cast %A : memref<10x3xf32, 0> to memref<?x?xf32> %V = memref.cast %A : memref<10x3xf32, 0> to memref<?x?xf32>
linalg.fill ins(%f10 : f32) outs(%V : memref<?x?xf32, 0>) linalg.fill ins(%f10 : f32) inits(%V : memref<?x?xf32, 0>)
%U = memref.cast %A : memref<10x3xf32, 0> to memref<*xf32> %U = memref.cast %A : memref<10x3xf32, 0> to memref<*xf32>
call @printMemrefF32(%U) : (memref<*xf32>) -> () call @printMemrefF32(%U) : (memref<*xf32>) -> ()
%V2 = memref.cast %U : memref<*xf32> to memref<?x?xf32> %V2 = memref.cast %U : memref<*xf32> to memref<?x?xf32>
linalg.fill ins(%f5 : f32) outs(%V2 : memref<?x?xf32, 0>) linalg.fill ins(%f5 : f32) inits(%V2 : memref<?x?xf32, 0>)
%U2 = memref.cast %V2 : memref<?x?xf32, 0> to memref<*xf32> %U2 = memref.cast %V2 : memref<?x?xf32, 0> to memref<*xf32>
call @printMemrefF32(%U2) : (memref<*xf32>) -> () call @printMemrefF32(%U2) : (memref<*xf32>) -> ()
%V3 = memref.cast %V2 : memref<?x?xf32> to memref<*xf32> %V3 = memref.cast %V2 : memref<?x?xf32> to memref<*xf32>
%V4 = memref.cast %V3 : memref<*xf32> to memref<?x?xf32> %V4 = memref.cast %V3 : memref<*xf32> to memref<?x?xf32>
linalg.fill ins(%f2 : f32) outs(%V4 : memref<?x?xf32, 0>) linalg.fill ins(%f2 : f32) inits(%V4 : memref<?x?xf32, 0>)
%U3 = memref.cast %V2 : memref<?x?xf32> to memref<*xf32> %U3 = memref.cast %V2 : memref<?x?xf32> to memref<*xf32>
call @printMemrefF32(%U3) : (memref<*xf32>) -> () call @printMemrefF32(%U3) : (memref<*xf32>) -> ()
// 122 is ASCII for 'z'. // 122 is ASCII for 'z'.
%i8_z = arith.constant 122 : i8 %i8_z = arith.constant 122 : i8
%I8 = memref.alloc() : memref<i8> %I8 = memref.alloc() : memref<i8>
memref.store %i8_z, %I8[]: memref<i8> memref.store %i8_z, %I8[]: memref<i8>
%U4 = memref.cast %I8 : memref<i8> to memref<*xi8> %U4 = memref.cast %I8 : memref<i8> to memref<*xi8>
Show All 9 Lines
} }
func.func private @printMemrefI8(memref<*xi8>) attributes { llvm.emit_c_interface } func.func private @printMemrefI8(memref<*xi8>) attributes { llvm.emit_c_interface }
func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface } func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }
func.func @return_two_var_memref_caller() { func.func @return_two_var_memref_caller() {
%0 = memref.alloca() : memref<4x3xf32> %0 = memref.alloca() : memref<4x3xf32>
%c0f32 = arith.constant 1.0 : f32 %c0f32 = arith.constant 1.0 : f32
linalg.fill ins(%c0f32 : f32) outs(%0 : memref<4x3xf32>) linalg.fill ins(%c0f32 : f32) inits(%0 : memref<4x3xf32>)
%1:2 = call @return_two_var_memref(%0) : (memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) %1:2 = call @return_two_var_memref(%0) : (memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>)
call @printMemrefF32(%1#0) : (memref<*xf32>) -> () call @printMemrefF32(%1#0) : (memref<*xf32>) -> ()
call @printMemrefF32(%1#1) : (memref<*xf32>) -> () call @printMemrefF32(%1#1) : (memref<*xf32>) -> ()
return return
} }
func.func @return_two_var_memref(%arg0: memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) { func.func @return_two_var_memref(%arg0: memref<4x3xf32>) -> (memref<*xf32>, memref<*xf32>) {
%0 = memref.cast %arg0 : memref<4x3xf32> to memref<*xf32> %0 = memref.cast %arg0 : memref<4x3xf32> to memref<*xf32>
return %0, %0 : memref<*xf32>, memref<*xf32> return %0, %0 : memref<*xf32>, memref<*xf32>
} }
func.func @return_var_memref_caller() { func.func @return_var_memref_caller() {
%0 = memref.alloca() : memref<4x3xf32> %0 = memref.alloca() : memref<4x3xf32>
%c0f32 = arith.constant 1.0 : f32 %c0f32 = arith.constant 1.0 : f32
linalg.fill ins(%c0f32 : f32) outs(%0 : memref<4x3xf32>) linalg.fill ins(%c0f32 : f32) inits(%0 : memref<4x3xf32>)
%1 = call @return_var_memref(%0) : (memref<4x3xf32>) -> memref<*xf32> %1 = call @return_var_memref(%0) : (memref<4x3xf32>) -> memref<*xf32>
call @printMemrefF32(%1) : (memref<*xf32>) -> () call @printMemrefF32(%1) : (memref<*xf32>) -> ()
return return
} }
func.func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> { func.func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> {
%0 = memref.cast %arg0: memref<4x3xf32> to memref<*xf32> %0 = memref.cast %arg0: memref<4x3xf32> to memref<*xf32>
return %0 : memref<*xf32> return %0 : memref<*xf32>
Show All 23 Lines

mlir/test/mlir-cpu-runner/utils.mlir

Show All 13 Lines
} }
// PRINT-0D: Unranked Memref base@ = {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data = // PRINT-0D: Unranked Memref base@ = {{.*}} rank = 0 offset = 0 sizes = [] strides = [] data =
// PRINT-0D: [2] // PRINT-0D: [2]
func.func @print_1d() { func.func @print_1d() {
%f = arith.constant 2.00000e+00 : f32 %f = arith.constant 2.00000e+00 : f32
%A = memref.alloc() : memref<16xf32> %A = memref.alloc() : memref<16xf32>
%B = memref.cast %A: memref<16xf32> to memref<?xf32> %B = memref.cast %A: memref<16xf32> to memref<?xf32>
linalg.fill ins(%f : f32) outs(%B : memref<?xf32>) linalg.fill ins(%f : f32) inits(%B : memref<?xf32>)
%U = memref.cast %B : memref<?xf32> to memref<*xf32> %U = memref.cast %B : memref<?xf32> to memref<*xf32>
call @printMemrefF32(%U): (memref<*xf32>) -> () call @printMemrefF32(%U): (memref<*xf32>) -> ()
memref.dealloc %A : memref<16xf32> memref.dealloc %A : memref<16xf32>
return return
} }
// PRINT-1D: Unranked Memref base@ = {{.*}} rank = 1 offset = 0 sizes = [16] strides = [1] data = // PRINT-1D: Unranked Memref base@ = {{.*}} rank = 1 offset = 0 sizes = [16] strides = [1] data =
// PRINT-1D-NEXT: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2] // PRINT-1D-NEXT: [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
func.func @print_3d() { func.func @print_3d() {
%f = arith.constant 2.00000e+00 : f32 %f = arith.constant 2.00000e+00 : f32
%f4 = arith.constant 4.00000e+00 : f32 %f4 = arith.constant 4.00000e+00 : f32
%A = memref.alloc() : memref<3x4x5xf32> %A = memref.alloc() : memref<3x4x5xf32>
%B = memref.cast %A: memref<3x4x5xf32> to memref<?x?x?xf32> %B = memref.cast %A: memref<3x4x5xf32> to memref<?x?x?xf32>
linalg.fill ins(%f : f32) outs(%B : memref<?x?x?xf32>) linalg.fill ins(%f : f32) inits(%B : memref<?x?x?xf32>)
%c2 = arith.constant 2 : index %c2 = arith.constant 2 : index
memref.store %f4, %B[%c2, %c2, %c2]: memref<?x?x?xf32> memref.store %f4, %B[%c2, %c2, %c2]: memref<?x?x?xf32>
%U = memref.cast %B : memref<?x?x?xf32> to memref<*xf32> %U = memref.cast %B : memref<?x?x?xf32> to memref<*xf32>
call @printMemrefF32(%U): (memref<*xf32>) -> () call @printMemrefF32(%U): (memref<*xf32>) -> ()
memref.dealloc %A : memref<3x4x5xf32> memref.dealloc %A : memref<3x4x5xf32>
return return
} }
Show All 29 Lines

mlir/test/mlir-opt/async.mlir

Show All 14 Linesfunc.func @main() {
%c0 = arith.constant 0.0 : f32 %c0 = arith.constant 0.0 : f32
%c1 = arith.constant 1.0 : f32 %c1 = arith.constant 1.0 : f32
%c2 = arith.constant 2.0 : f32 %c2 = arith.constant 2.0 : f32
%c3 = arith.constant 3.0 : f32 %c3 = arith.constant 3.0 : f32
%c4 = arith.constant 4.0 : f32 %c4 = arith.constant 4.0 : f32
%A = memref.alloc() : memref<4xf32> %A = memref.alloc() : memref<4xf32>
linalg.fill ins(%c0 : f32) outs(%A : memref<4xf32>) linalg.fill ins(%c0 : f32) inits(%A : memref<4xf32>)
%U = memref.cast %A : memref<4xf32> to memref<*xf32> %U = memref.cast %A : memref<4xf32> to memref<*xf32>
call @printMemrefF32(%U): (memref<*xf32>) -> () call @printMemrefF32(%U): (memref<*xf32>) -> ()
memref.store %c1, %A[%i0]: memref<4xf32> memref.store %c1, %A[%i0]: memref<4xf32>
call @mlirAsyncRuntimePrintCurrentThreadId(): () -> () call @mlirAsyncRuntimePrintCurrentThreadId(): () -> ()
call @printMemrefF32(%U): (memref<*xf32>) -> () call @printMemrefF32(%U): (memref<*xf32>) -> ()
Show All 39 Lines

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

Show First 20 Lines • Show All 50 Lines • ▼ Show 20 Lines with InsertionPoint(module.body):
# 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:.+]] = tensor.empty() : 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: inits(%[[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 = tensor.EmptyOp([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)
▲ Show 20 Lines • Show All 104 Lines • Show Last 20 Lines

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

Show All 15 Lines
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
# CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<12x?xf32> # CHECK-SAME: %[[OUT:[0-9a-z]+]]: tensor<12x?xf32>
# CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32 # CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
# CHECK-NEXT: %[[RES:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[OUT]] : tensor<12x?xf32>) -> tensor<12x?xf32> # CHECK-NEXT: %[[RES:.*]] = linalg.fill ins(%[[CST]] : f32) inits(%[[OUT]] : tensor<12x?xf32>) -> tensor<12x?xf32>
# CHECK-NEXT: return %[[RES]] : tensor<12x?xf32> # CHECK-NEXT: return %[[RES]] : tensor<12x?xf32>
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32)) RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32))
def fill_tensor(out): def fill_tensor(out):
zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
return linalg.fill(zero, outs=[out]) return linalg.fill(zero, outs=[out])
# CHECK-LABEL: func @fill_buffer # CHECK-LABEL: func @fill_buffer
# CHECK-SAME: %[[OUT:[0-9a-z]+]]: memref<12x?xf32> # CHECK-SAME: %[[OUT:[0-9a-z]+]]: memref<12x?xf32>
# CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32 # CHECK-NEXT: %[[CST:.*]] = arith.constant 0.0{{.*}} : f32
# CHECK-NEXT: linalg.fill ins(%[[CST]] : f32) outs(%[[OUT]] : memref<12x?xf32>) # CHECK-NEXT: linalg.fill ins(%[[CST]] : f32) inits(%[[OUT]] : memref<12x?xf32>)
# CHECK-NEXT: return # CHECK-NEXT: return
@func.FuncOp.from_py_func( @func.FuncOp.from_py_func(
MemRefType.get((12, ShapedType.get_dynamic_size()), f32)) MemRefType.get((12, ShapedType.get_dynamic_size()), f32))
def fill_buffer(out): def fill_buffer(out):
zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result zero = arith.ConstantOp(value=FloatAttr.get(f32, 0.), result=f32).result
linalg.fill(zero, outs=[out]) linalg.fill(zero, outs=[out])
print(module) print(module)
Show All 10 Lines 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 = tensor.EmptyOp([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>) inits(%{{.*}} : 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>
# CHECK-SAME: fun = #linalg.binary_fn<mul> # CHECK-SAME: fun = #linalg.binary_fn<mul>
# CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<4x8xf32>, tensor<4x8xf32>) outs(%{{.*}} : tensor<4x8xf32>) # CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<4x8xf32>, tensor<4x8xf32>) inits(%{{.*}} : tensor<4x8xf32>)
# CHECK: return # CHECK: return
binary_result = linalg.elemwise_binary( binary_result = linalg.elemwise_binary(
lhs, lhs,
rhs, rhs,
outs=[init_result.result], outs=[init_result.result],
fun=BinaryFn.mul, fun=BinaryFn.mul,
cast=TypeFn.cast_unsigned) cast=TypeFn.cast_unsigned)
return unary_result, binary_result return unary_result, binary_result
▲ Show 20 Lines • Show All 63 Lines • ▼ Show 20 Lines with InsertionPoint(module.body):
# 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:.*]] = tensor.empty # CHECK: %[[INIT:.*]] = tensor.empty
init = tensor.EmptyOp([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: inits(%[[INIT]]
return linalg.matmul(lhs, rhs, outs=init) return linalg.matmul(lhs, rhs, outs=init)
print(module) print(module)

mlir/test/python/integration/dialects/linalg/opsrun.py

Show All 23 Linesfunc.func @main() -> f32 attributes {llvm.emit_c_interface} {
%v0 = arith.constant 0.0 : f32 %v0 = arith.constant 0.0 : f32
%v1 = arith.constant 1.0 : f32 %v1 = arith.constant 1.0 : f32
%v2 = arith.constant 2.0 : f32 %v2 = arith.constant 2.0 : f32
%lhs = memref.alloc() : memref<f32> %lhs = memref.alloc() : memref<f32>
%rhs = memref.alloc() : memref<4x8xf32> %rhs = memref.alloc() : memref<4x8xf32>
%O0 = memref.alloc() : memref<4x8xf32> %O0 = memref.alloc() : memref<4x8xf32>
%O1 = memref.alloc() : memref<4x8xf32> %O1 = memref.alloc() : memref<4x8xf32>
linalg.fill ins(%v1 : f32) outs(%lhs : memref<f32>) linalg.fill ins(%v1 : f32) inits(%lhs : memref<f32>)
linalg.fill ins(%v2 : f32) outs(%rhs : memref<4x8xf32>) linalg.fill ins(%v2 : f32) inits(%rhs : memref<4x8xf32>)
linalg.fill ins(%v0 : f32) outs(%O0 : memref<4x8xf32>) linalg.fill ins(%v0 : f32) inits(%O0 : memref<4x8xf32>)
linalg.fill ins(%v0 : f32) outs(%O1 : memref<4x8xf32>) linalg.fill ins(%v0 : f32) inits(%O1 : memref<4x8xf32>)
call @elemwise_exp_add_on_buffers(%lhs, %rhs, %O0) : call @elemwise_exp_add_on_buffers(%lhs, %rhs, %O0) :
(memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> () (memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> ()
call @elemwise_log_mul_on_buffers(%lhs, %rhs, %O1) : call @elemwise_log_mul_on_buffers(%lhs, %rhs, %O1) :
(memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> () (memref<f32>, memref<4x8xf32>, memref<4x8xf32>) -> ()
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%res0 = memref.load %O0[%c0, %c0] : memref<4x8xf32> %res0 = memref.load %O0[%c0, %c0] : memref<4x8xf32>
Show All 11 Linesfunc.func @main() -> f32 attributes {llvm.emit_c_interface} {
%v0 = arith.constant 0.0 : f32 %v0 = arith.constant 0.0 : f32
%v1 = arith.constant -1 : i8 %v1 = arith.constant -1 : i8
%v2 = arith.constant 2.0 : f32 %v2 = arith.constant 2.0 : f32
%A = memref.alloc() : memref<4x16xi8> %A = memref.alloc() : memref<4x16xi8>
%B = memref.alloc() : memref<16x8xf32> %B = memref.alloc() : memref<16x8xf32>
%C0 = memref.alloc() : memref<4x8xf32> %C0 = memref.alloc() : memref<4x8xf32>
%C1 = memref.alloc() : memref<4x8xf32> %C1 = memref.alloc() : memref<4x8xf32>
linalg.fill ins(%v1 : i8) outs(%A : memref<4x16xi8>) linalg.fill ins(%v1 : i8) inits(%A : memref<4x16xi8>)
linalg.fill ins(%v2 : f32) outs(%B : memref<16x8xf32>) linalg.fill ins(%v2 : f32) inits(%B : memref<16x8xf32>)
linalg.fill ins(%v0 : f32) outs(%C0 : memref<4x8xf32>) linalg.fill ins(%v0 : f32) inits(%C0 : memref<4x8xf32>)
linalg.fill ins(%v0 : f32) outs(%C1 : memref<4x8xf32>) linalg.fill ins(%v0 : f32) inits(%C1 : memref<4x8xf32>)
call @matmul_signed_on_buffers(%A, %B, %C0) : call @matmul_signed_on_buffers(%A, %B, %C0) :
(memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> () (memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
call @matmul_unsigned_on_buffers(%A, %B, %C1) : call @matmul_unsigned_on_buffers(%A, %B, %C1) :
(memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> () (memref<4x16xi8>, memref<16x8xf32>, memref<4x8xf32>) -> ()
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%res0 = memref.load %C0[%c0, %c0] : memref<4x8xf32> %res0 = memref.load %C0[%c0, %c0] : memref<4x8xf32>
▲ Show 20 Lines • Show All 57 Lines • ▼ Show 20 Lines
func.func @main() -> i32 attributes {llvm.emit_c_interface} { func.func @main() -> i32 attributes {llvm.emit_c_interface} {
%v0 = arith.constant 0 : i32 %v0 = arith.constant 0 : i32
%v1 = arith.constant 1.0 : f64 %v1 = arith.constant 1.0 : f64
%v2 = arith.constant 2.0 : f64 %v2 = arith.constant 2.0 : f64
%input = memref.alloc() : memref<1x4x16x1xf64> %input = memref.alloc() : memref<1x4x16x1xf64>
%filter = memref.alloc() : memref<2x2x1xf64> %filter = memref.alloc() : memref<2x2x1xf64>
%output = memref.alloc() : memref<1x2x4x1xi32> %output = memref.alloc() : memref<1x2x4x1xi32>
linalg.fill ins(%v1 : f64) outs(%input : memref<1x4x16x1xf64>) linalg.fill ins(%v1 : f64) inits(%input : memref<1x4x16x1xf64>)
linalg.fill ins(%v2 : f64) outs(%filter : memref<2x2x1xf64>) linalg.fill ins(%v2 : f64) inits(%filter : memref<2x2x1xf64>)
linalg.fill ins(%v0 : i32) outs(%output : memref<1x2x4x1xi32>) linalg.fill ins(%v0 : i32) inits(%output : memref<1x2x4x1xi32>)
call @conv_on_buffers(%input, %filter, %output) : call @conv_on_buffers(%input, %filter, %output) :
(memref<1x4x16x1xf64>, memref<2x2x1xf64>, memref<1x2x4x1xi32>) -> () (memref<1x4x16x1xf64>, memref<2x2x1xf64>, memref<1x2x4x1xi32>) -> ()
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32> %0 = memref.load %output[%c0, %c0, %c0, %c0] : memref<1x2x4x1xi32>
// TODO: FFI-based solution to allow testing and printing with python code. // TODO: FFI-based solution to allow testing and printing with python code.
return %0 : i32 return %0 : i32
} }
""" """
pooling_boiler = """ pooling_boiler = """
func.func @main() -> i32 attributes {llvm.emit_c_interface} { func.func @main() -> i32 attributes {llvm.emit_c_interface} {
%v0 = arith.constant 0 : i32 %v0 = arith.constant 0 : i32
%v42 = arith.constant 42.0 : f64 %v42 = arith.constant 42.0 : f64
%v77 = arith.constant 77.0 : f64 %v77 = arith.constant 77.0 : f64
%v-13 = arith.constant -13.0 : f64 %v-13 = arith.constant -13.0 : f64
%v1 = arith.constant 1.0 : f64 %v1 = arith.constant 1.0 : f64
%input = memref.alloc() : memref<1x4x16x1xf64> %input = memref.alloc() : memref<1x4x16x1xf64>
%shape = memref.alloc() : memref<2x2xf64> %shape = memref.alloc() : memref<2x2xf64>
%output = memref.alloc() : memref<1x2x4x1xi32> %output = memref.alloc() : memref<1x2x4x1xi32>
linalg.fill ins(%v1 : f64) outs(%input : memref<1x4x16x1xf64>) linalg.fill ins(%v1 : f64) inits(%input : memref<1x4x16x1xf64>)
linalg.fill ins(%v1 : f64) outs(%shape : memref<2x2xf64>) linalg.fill ins(%v1 : f64) inits(%shape : memref<2x2xf64>)
linalg.fill ins(%v0 : i32) outs(%output : memref<1x2x4x1xi32>) linalg.fill ins(%v0 : i32) inits(%output : memref<1x2x4x1xi32>)
%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
memref.store %v42, %input[%c0, %c0, %c0, %c0] : memref<1x4x16x1xf64> memref.store %v42, %input[%c0, %c0, %c0, %c0] : memref<1x4x16x1xf64>
memref.store %v77, %input[%c0, %c0, %c1, %c0] : memref<1x4x16x1xf64> memref.store %v77, %input[%c0, %c0, %c1, %c0] : memref<1x4x16x1xf64>
memref.store %v-13, %input[%c0, %c1, %c0, %c0] : memref<1x4x16x1xf64> memref.store %v-13, %input[%c0, %c1, %c0, %c0] : memref<1x4x16x1xf64>
▲ Show 20 Lines • Show All 431 Lines • Show Last 20 Lines