Commit Graph

119 Commits

Author SHA1 Message Date
liqinweng
5c18ae3135 [MLIR][Tensor] Canonicalize expand/collapse_shape of splat to splat
Collapsing / expanding a splatted value can be replaced with a single `tensor.splat` operation. Replace
these cases with a simple `tensor.splat` operation.

Reviewed By: rsuderman

Differential Revision: https://reviews.llvm.org/D140552
2023-01-04 13:07:55 -08:00
Mehdi Amini
1211af761f Apply clang-tidy fixes for llvm-else-after-return in TensorOps.cpp (NFC) 2022-12-22 15:33:01 +00:00
Fangrui Song
cbb0981388 [mlir] llvm::Optional::value => operator*/operator->
std::optional::value() has undesired exception checking semantics and is
unavailable in older Xcode (see _LIBCPP_AVAILABILITY_BAD_OPTIONAL_ACCESS). The
call sites block std::optional migration.
2022-12-17 19:07:38 +00:00
Ramkumar Ramachandra
22426110c5 mlir/tblgen: use std::optional in generation
This is part of an effort to migrate from llvm::Optional to
std::optional. This patch changes the way mlir-tblgen generates .inc
files, and modifies tests and documentation appropriately. It is a "no
compromises" patch, and doesn't leave the user with an unpleasant mix of
llvm::Optional and std::optional.

A non-trivial change has been made to ControlFlowInterfaces to split one
constructor into two, relating to a build failure on Windows.

See also: https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716

Signed-off-by: Ramkumar Ramachandra <r@artagnon.com>

Differential Revision: https://reviews.llvm.org/D138934
2022-12-17 11:13:26 +01:00
Aliia Khasanova
ded75a282a Remove sentinel argument from dispatchIndexOpFoldResults.
Post clean-up after merger of kDynamicSize and kDynamicStrideOrOffset.

Differential Revision: https://reviews.llvm.org/D139929
2022-12-13 14:04:46 +01:00
Nicolas Vasilache
93bbcffc7e [mlir][Transform] Make FuseIntoContainingOp support rank-reducing extract slices
This fixes an issue where rank-reducing + fusion would not interop properly.

Differential Revision: https://reviews.llvm.org/D139844
2022-12-12 12:55:08 -08:00
Alexander Belyaev
f6fb0a4f35 [mlir] Make patterns for folding tensor.empty optional.
At the moment, they are a part of EmptyOp::getCanonicalizationPatterns. When
extract_slice(tensor.empty) is rewritten as a new tensor.empty, it could
happen that we end up with two tensor.empty ops, since the original
tensor.empty can have two users. After bufferization such cases result in two
allocations.

Differential Revision: https://reviews.llvm.org/D139308
2022-12-07 23:01:34 +01:00
Lorenzo Chelini
87ecf9d155 [MLIR][Tensor] Add custom builder for unpack op
Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D139344
2022-12-07 12:40:45 +01:00
Hanhan Wang
0f297cad4d [mlir][tensor][linalg] Introduce DataLayoutPropagation pass.
It introduces a pattern that swaps `linalg.generic + tensor.pack` to
`tensor.pack + linalg.generic`. It requires all the iteration types
being parallel; the indexing map of output operand is identiy. They can
all be relaxed in the future.

The user can decide whether the propagation should be applied or not by
passing a control function.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D138882
2022-12-06 15:00:07 -08:00
Hanhan Wang
193cefd1b1 [mlir][tensor] Adapt FoldTensorCastProducerOp pattern on DPS interface.
This revision adapts the pattern in LinAlg to work on DPS interface, and
adds it to canonicalization patterns of tensor dialect. The
InsertSliceOp is skipped in the pattern because it has its own logic
about folding tensor.cast ops.

Reviewed By: pifon2a

Differential Revision: https://reviews.llvm.org/D139375
2022-12-06 12:13:37 -08:00
Adrian Kuegel
94d3df2015 [mlir][Tensor] Apply ClangTidy performance finding (NFC) 2022-12-05 11:22:20 +01:00
Kazu Hirata
1a36588ec6 [mlir] Use std::nullopt instead of None (NFC)
This patch mechanically replaces None with std::nullopt where the
compiler would warn if None were deprecated.  The intent is to reduce
the amount of manual work required in migrating from Optional to
std::optional.

This is part of an effort to migrate from llvm::Optional to
std::optional:

https://discourse.llvm.org/t/deprecating-llvm-optional-x-hasvalue-getvalue-getvalueor/63716
2022-12-03 18:50:27 -08:00
Lorenzo Chelini
44f7356005 [MLIR][Tensor] Add canonicalization for UnpackOp
pack(unpack(x)) -> x

Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D138917
2022-12-01 15:17:50 +01:00
Hanhan Wang
a971d51932 [mlir][tensor] Enhance the verifier of pack and unpack op.
The outer_dims_perm must be a permutation or empty.

Reviewed By: chelini

Differential Revision: https://reviews.llvm.org/D138936
2022-11-29 15:47:52 -08:00
Hanhan Wang
e86169f090 [mlir][tensor] Add a custom builder for pack op.
The `paddingValue` and `outerDimsPerm` are optional to the op;
`innerTiles` can be variadic in terms of static sizes and dynamic sizes.
Add a custom builder for building pack op easier.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D138860
2022-11-28 15:18:42 -08:00
Lorenzo Chelini
a9733b8a5e [MLIR] Adopt DenseI64ArrayAttr in tensor, memref and linalg transform
This commit is a first step toward removing inconsistencies between dynamic
and static attributes (i64 v. index) by dropping `I64ArrayAttr` and
using `DenseI64ArrayAttr` in Tensor, Memref and Linalg Transform ops.
In Linalg Transform ops only `TileToScfForOp` and `TileOp` have been updated.

See related discussion: https://discourse.llvm.org/t/rfc-inconsistency-between-dynamic-and-static-attributes-i64-v-index/66612/1

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D138567
2022-11-25 09:43:30 +01:00
Matthias Springer
f2d91a7ae1 [mlir][utils] Fix invalid reshapes in ComposeCollapseOfExpandOp
Do not generate CollapseShapeOps/ExpandShapeOps that have the same source and result shape. Generate casts instead. Such reshapes became invalid with D138498.

Differential Revision: https://reviews.llvm.org/D138557
2022-11-23 13:52:00 +01:00
Matthias Springer
b9745ad812 [mlir][tensor/memref] Disallow Collapse/ExpandShapeOps that do not reduce/increase the rank
CollapseShapeOp/ExpandShapeOp that do not change the rank (or increase/reduce it) are invalid.

Differential Revision: https://reviews.llvm.org/D138498
2022-11-23 09:19:35 +01:00
Matthias Springer
6052b17aab [mlir][tensor] Add dim(expand_shape/collapse_shape) folding
Differential Revision: https://reviews.llvm.org/D138487
2022-11-22 17:34:49 +01:00
Lorenzo Chelini
85e38e5292 [MLIR][Tensor] Use the existing helper function applyPermutationToVector (NFC)
Avoid duplicate code by using an existing helper function to interchange
a vector based on a permutation. Address comments emerged after landing
D138119.

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D138480
2022-11-22 11:34:44 +01:00
Lorenzo Chelini
9aa505a28d Introduce tensor.pack and tensor.unpack operations
Pack and Unpack return new tensors within which the individual elements
are reshuffled according to the packing specification. This has the
consequence of modifying the canonical order in which a given operator
(i.e., Matmul) accesses the individual elements. After bufferization,
this typically translates to increased access locality and cache
behavior improvement, e.g., eliminating cache line splitting.

Co-authored-by: Mahesh Ravishankar <ravishankarm@google.com>
Co-authored-by: Han-Chung Wang <hanchung@google.com>

RFC: https://discourse.llvm.org/t/rfc-tensor-pack-and-tensor-unpack/66408/1

Reviewed By: nicolasvasilache, rengolin, hanchung

Differential Revision: https://reviews.llvm.org/D138119
2022-11-22 09:11:59 +01:00
Aliia Khasanova
399638f98c Merge kDynamicSize and kDynamicSentinel into one constant.
resolve conflicts

Differential Revision: https://reviews.llvm.org/D138282
2022-11-21 13:01:26 +00:00
Mahesh Ravishankar
24f9293de8 [mlir][Tensor] Allow builders of tensor.empty to accept encoding attribute.
The `RankedTensorType` can have an optional encoding
attribute. Allowing the builders of `tensor.empty` to accept the
encoding attribute (optionally), allows building empty tensors with
the type having the encoding attribute.

Reviewed By: nicolasvasilache, hanchung, springerm

Differential Revision: https://reviews.llvm.org/D137297
2022-11-03 20:30:12 +00:00
Matthias Springer
cfaf3292df [mlir][tensor] Disallow unranked tensors for tensor.extract/insert
When writing a tensor.extract/tensor.insert, the rank of the tensor is implied by the number of specified indices. When extracting from/inserting into an unranked tensor, it should first be casted to a ranked version.

Differential Revision: https://reviews.llvm.org/D136756
2022-10-27 10:09:31 +02:00
Mahesh Ravishankar
94b8469a88 [mlir][Tensor] Add a helper build method for pad operations with constant padding.
Drop the `createPadScalarOp` from Utils.h since it is a duplicate of
the `build` method added here.

Differential Revision: https://reviews.llvm.org/D136493
2022-10-24 18:11:53 +00:00
Matthias Springer
b169643f3a [mlir][interfaces] Remove getDestinationOperands from TilingInterface
`getDestinationOperands` was almost a duplicate of `DestinationStyleOpInterface::getOutputOperands`. Now that the interface has been moved to mlir/Interfaces, it is no longer needed.

Differential Revision: https://reviews.llvm.org/D136240
2022-10-24 09:26:19 +02:00
Lorenzo Chelini
9bcac22be5 [MLIR][Tensor] Remove assert in PadOp builder
The assert is misplaced as the result type is allowed to be null. A few
lines below the result type is inferred if it is passed a nullptr.
Besides, this behavior is described in the documentation of the builder.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D136262
2022-10-19 18:02:50 +02:00
Sanjoy Das
adabce4118 Correctly model undefined behavior in {tensor|memref}.dim
These operations have undefined behavior if the index is not less than the rank of the source tensor / memref, so they cannot be freely speculated like they were before this patch.  After this patch we speculate them only if we can prove that they don't have UB.

Depends on D135505.

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D135748
2022-10-12 17:30:13 -07:00
Nicolas Vasilache
54a4e9685d [mlir][Tensor] NFC - Add result pretty printing to TensorOps
Differential Revision: https://reviews.llvm.org/D135135
2022-10-04 09:16:51 -07:00
Matthias Springer
81ca5aa452 [mlir][tensor][NFC] Rename linalg.init_tensor to tensor.empty
tensor.empty/linalg.init_tensor produces an uninititalized tensor that can be used as a destination operand for destination-style ops (ops that implement `DestinationStyleOpInterface`).

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

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

Differential Revision: https://reviews.llvm.org/D135129
2022-10-04 17:25:35 +09:00
Jakub Kuderski
abc362a107 [mlir][arith] Change dialect name from Arithmetic to Arith
Suggested by @lattner in https://discourse.llvm.org/t/rfc-define-precise-arith-semantics/65507/22.

Tested with:
`ninja check-mlir check-mlir-integration check-mlir-mlir-spirv-cpu-runner check-mlir-mlir-vulkan-runner check-mlir-examples`

and `bazel build --config=generic_clang @llvm-project//mlir:all`.

Reviewed By: lattner, Mogball, rriddle, jpienaar, mehdi_amini

Differential Revision: https://reviews.llvm.org/D134762
2022-09-29 11:23:28 -04:00
Lei Zhang
9d59705169 [mlir][tensor] Fold round-tripping extract/insert slice ops
Reviewed By: ThomasRaoux, nicolasvasilache

Differential Revision: https://reviews.llvm.org/D133909
2022-09-19 12:58:52 -04:00
Christopher Bate
f4a478cd01 [mlir][Tensor] Add rewrites to extract slices through tensor.collape_shape
This change adds a set of utilities to replace the result of a
`tensor.collapse_shape -> tensor.extract_slice` chain with the
equivalent result formed by aggregating slices of the
`tensor.collapse_shape` source. In general, it is not possible to
commute `extract_slice` and `collapse_shape` if linearized dimensions
are sliced. The i-th dimension of the `tensor.collapse_shape`
result is a "linearized sliced dimension" if:

1) Reassociation indices of tensor.collapse_shape in the i'th position
   is greater than size 1 (multiple dimensions of the input are collapsed)
2) The i-th dimension is sliced by `tensor.extract_slice`.

We can work around this by stitching together the result of
`tensor.extract_slice` by iterating over any linearized sliced dimensions.
This is equivalent to "tiling" the linearized-and-sliced dimensions of
the `tensor.collapse_shape` operation in order to manifest the result
tile (the result of the `tensor.extract_slice`). The user of the
utilities must provide the mechanism to create the tiling (e.g. a loop).
In the tests, it is demonstrated how to apply the utilities using either
`scf.for` or `scf.foreach_thread`.

The below example illustrates the pattern using `scf.for`:

```
%0 = linalg.generic ... -> tensor<3x7x11x10xf32>
%1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32>
%2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32>
```

We can construct %2 by generating the following IR:

```
%dest = linalg.init_tensor() : tensor<10x10xf32>
%2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> {
   // Step 1: Map this output idx (%iv) to a multi-index for the input (%3):
   %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv)
   %3:3 = arith.delinearize_index %iv into (3, 7, 11)
   // Step 2: Extract the slice from the input
   %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] :
         tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32>
   %5 = tensor.collapse_shape %4 [[0, 1, 2], [3]] :
         tensor<1x1x1x10xf32> into tensor<1x10xf32>
   // Step 3: Insert the slice into the destination
   %6 = tensor.insert_slice %5 into %arg0 [%iv, 0] [1, 10] [1, 1] :
         tensor<1x10xf32> into tensor<10x10xf32>
   scf.yield %6 : tensor<10x10xf32>
}
```

The pattern was discussed in the RFC here: https://discourse.llvm.org/t/rfc-tensor-extracting-slices-from-tensor-collapse-shape/64034

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D129699
2022-09-08 21:58:21 -06:00
Nicolas Vasilache
d2613d5bb5 [mlir][tensor] Add gather/scatter op definitions to the tensor dialect.
Gather/Scatter are examined from first principles in light of our recent progress on tensor-based codegen
and in-place bufferization.

In the future, lowering of these abstractions to operate **inplace** on buffers
will likely require a more powerful buffer representation than strided memref.

General context: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707
Relevant TL;DR parts of the proposal:
- gather: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-gatherop-and-friends-10
- need for more expressive types: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-bufferization-copy-view-and-the-need-for-more-expressive-types-12
- jagged buffer discussion: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707#proposal-first-class-jagged-buffer-13

Differential Revision: https://reviews.llvm.org/D130348
2022-09-05 02:02:22 -07:00
Mehdi Amini
0b1aee38bd Revert "[mlir][Tensor] Add rewrites to extract slices through tensor.collape_shape"
This reverts commit 5711957875.

A circular dependency is introduced here from Dialect/Utils/ to the
ViewLikeInterface, but it already depends on Dialect/Utils.

Also this introduces a dependency from lib/Dialect/Tensor to Linalg,
which isn't obviously correct from a layering point of view.
2022-09-02 23:34:52 +00:00
Christopher Bate
5711957875 [mlir][Tensor] Add rewrites to extract slices through tensor.collape_shape
This change adds a set of utilities to replace the result of a
`tensor.collapse_shape -> tensor.extract_slice` chain with the
equivalent result formed by aggregating slices of the
`tensor.collapse_shape` source. In general, it is not possible to
commute `extract_slice` and `collapse_shape` if linearized dimensions
are sliced. The i-th dimension of the `tensor.collapse_shape`
result is a "linearized sliced dimension" if:

1) Reassociation indices of tensor.collapse_shape in the i'th position
   is greater than size 1 (multiple dimensions of the input are collapsed)
2) The i-th dimension is sliced by `tensor.extract_slice`.

We can work around this by stitching together the result of
`tensor.extract_slice` by iterating over any linearized sliced dimensions.
This is equivalent to "tiling" the linearized-and-sliced dimensions of
the `tensor.collapse_shape` operation in order to manifest the result
tile (the result of the `tensor.extract_slice`). The user of the
utilities must provide the mechanism to create the tiling (e.g. a loop).
In the tests, it is demonstrated how to apply the utilities using either
`scf.for` or `scf.foreach_thread`.

The below example illustrates the pattern using `scf.for`:

```
%0 = linalg.generic ... -> tensor<3x7x11x10xf32>
%1 = tensor.collapse_shape %0 [[0, 1, 2], [3]] : ... to tensor<341x10xf32>
%2 = tensor.extract_slice %1 [13, 0] [10, 10] [2, 1] : .... tensor<10x10xf32>
```

We can construct %2 by generating the following IR:

```
%dest = linalg.init_tensor() : tensor<10x10xf32>
%2 = scf.for %iv = %c0 to %c10 step %c1 iter_args(%arg0) -> tensor<10x10xf32> {
   // Step 1: Map this output idx (%iv) to a multi-index for the input (%3):
   %linear_index = affine.apply affine_map<(d0)[]->(d0*2 + 11)>(%iv)
   %3:3 = arith.delinearize_index %iv into (3, 7, 11)
   // Step 2: Extract the slice from the input
   %4 = tensor.extract_slice %0 [%3#0, %3#1, %3#2, 0] [1, 1, 1, 10] [1, 1, 1, 1] :
         tensor<3x7x11x10xf32> to tensor<1x1x1x10xf32>
   %5 = tensor.collapse_shape %4 [[0, 1, 2], [3]] :
         tensor<1x1x1x10xf32> into tensor<1x10xf32>
   // Step 3: Insert the slice into the destination
   %6 = tensor.insert_slice %5 into %arg0 [%iv, 0] [1, 10] [1, 1] :
         tensor<1x10xf32> into tensor<10x10xf32>
   scf.yield %6 : tensor<10x10xf32>
}
```

The pattern was discussed in the RFC here: https://discourse.llvm.org/t/rfc-tensor-extracting-slices-from-tensor-collapse-shape/64034

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D129699
2022-09-02 11:29:04 -06:00
Thomas Raoux
1ee0d60a9b [mlir][tensor] Remove incorrect parallel_insert_slice folder
parallel_insert_slice doesn't return a value therefore we shouldn't try
to fold the result. The insert folding don't apply to this op.
The current folding would cause pattern rewrite to not be able to
converge.

Differential Revision: https://reviews.llvm.org/D132668
2022-08-26 15:27:54 +00:00
Matthias Springer
ba95bf765d [mlir][tensor] Add getMixedSizes helper
This helper function computes the dimensions of a tensor value as OpFoldResults.

Differential Revision: https://reviews.llvm.org/D132475
2022-08-25 10:25:41 +02:00
Gaurav Shukla
7d6ef5caef [mlir][tensor] Fold tensor.cast into tensor.collapse_shape op
This commit folds a `tensor.cast` op into a `tensor.collapse_shape` op
when following two conditions meet:
1. the `tensor.collapse_shape` op consumes result of the `tensor.cast` op.
2. `tensor.cast` op casts to a more dynamic version of the source tensor.
This is added as a canonicalization pattern in `tensor.collapse_shape` op.

Signed-Off-By: Gaurav Shukla <gaurav@nod-labs.com>

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D130650
2022-07-28 13:11:43 +05:30
Nicolas Vasilache
c9fb3c6ea6 [mlir][Tensor] Update ParallelInsertSlicOp semantics to match that of InsertSliceOp
This revision updates the op semantics to also allow rank-reducing behavior as well
as updates the implementation to reuse code between the sequential and the parallel
version of the op.

Depends on D128920

Differential Revision: https://reviews.llvm.org/D128985
2022-07-04 02:37:46 -07:00
Nicolas Vasilache
7fbf55c927 [mlir][Tensor] Move ParallelInsertSlice to the tensor dialect
This is moslty NFC and will allow tensor.parallel_insert_slice to gain
rank-reducing semantics by reusing the vast majority of the tensor.insert_slice impl.

Depends on D128857

Differential Revision: https://reviews.llvm.org/D128920
2022-07-04 01:53:12 -07:00
Nicolas Vasilache
741f8f2bed [mlir][Tensor][NFC] Better document rank-reducing behavior of ExtractSliceOp and cleanup 2022-06-29 07:37:58 -07:00
Jacques Pienaar
2d70eff802 [mlir] Flip more uses to prefixed accessor form (NFC).
Try to keep the final flip small. Need to flip MemRef as there are many
templated cases with it and Tensor.
2022-06-26 19:12:38 -07:00
Aart Bik
e7d3ba1066 [mlir][sparse] accept sparse reshape (expand/collapse)
This revision makes sure we accept sparse tensors as arguments
of the expand/collapse reshaping operations in the tensor dialect.
Note that the actual lowering to runnable IR is still TBD.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D128311
2022-06-22 09:40:38 -07:00
Kazu Hirata
6d5fc1e3d5 [mlir] Don't use Optional::getValue (NFC) 2022-06-20 23:20:25 -07:00
Kazu Hirata
037f09959a [mlir] Don't use Optional::hasValue (NFC) 2022-06-20 11:22:37 -07:00
Jacques Pienaar
eca86cb2ed [mlir] Start migrating more dialects to prefixed form
Marked all dialects that could be (reasonably) easily flipped to _Both
prefix. Updating the accessors to prefixed form will happen in follow
up, this was to flush out conflicts and to mark all dialects explicitly
as I plan to flip OpBase default to _Prefixed to avoid needing to
migrate new dialects.

Except for Standalone example which got flipped to _Prefixed.

Differential Revision: https://reviews.llvm.org/D128027
2022-06-18 10:10:31 -07:00
Thomas Raoux
f2676b151d [mlir][tensor] Add canonicalization for tensor.cast from extract_slice
Propagate static size information into extract_slice producer if
possible.

Differential Revision: https://reviews.llvm.org/D125972
2022-05-19 17:34:59 +00:00
Chris Lattner
f21896f2c6 [DenseElementAttr] Simplify the public API for creating these.
Instead of requiring the client to compute the "isSplat" bit,
compute it internally.  This makes the logic more consistent
and defines away a lot of "elements.size()==1" in the clients.

This addresses Issue #55185

Differential Revision: https://reviews.llvm.org/D125447
2022-05-12 16:18:23 +01:00
River Riddle
eda6f907d2 [mlir][NFC] Shift a bunch of dialect includes from the .h to the .cpp
Now that dialect constructors are generated in the .cpp file, we can
drop all of the dependent dialect includes from the .h file.

Differential Revision: https://reviews.llvm.org/D124298
2022-04-23 01:09:29 -07:00