Commit Graph

91 Commits

Author SHA1 Message Date
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
Matthias Springer
4cd7362083 [mlir][SCF] foreach_thread: Capture shared output tensors explicitly
This change refines the semantics of scf.foreach_thread. Tensors that are inserted into in the terminator must now be passed to the region explicitly via `shared_outs`. Inside of the body of the op, those tensors are then accessed via block arguments.

The body of a scf.foreach_thread is now treated as a repetitive region. I.e., op dominance can no longer be used in conflict detection when using a value that is defined outside of the body. Such uses may now be considered as conflicts (if there is at least one read and one write in the body), effectively privatizing the tensor. Shared outputs are not privatized when they are used via their corresponding block arguments.

As part of this change, it was also necessary to update the "tiling to scf.foreach_thread", such that the generated tensor.extract_slice ops use the scf.foreach_thread's block arguments. This is implemented by cloning the TilingInterface op inside the scf.foreach_thread, rewriting all of its outputs with block arguments and then calling the tiling implementation. Afterwards, the cloned op is deleted again.

Differential Revision: https://reviews.llvm.org/D133114
2022-09-02 14:54:04 +02:00
Matthias Springer
547942841f [mlir][interfaces] Drop dest/tileDestOperands from TilingInterface
`getTiledImplementation`/`generateResultTileValue` only computes the tiled operation, but does not insert the result into any tensor.

Differential Revision: https://reviews.llvm.org/D133015
2022-09-01 08:53:53 +02:00
Michele Scuttari
67d0d7ac0a [MLIR] Update pass declarations to new autogenerated files
The patch introduces the required changes to update the pass declarations and definitions to use the new autogenerated files and allow dropping the old infrastructure.

Reviewed By: mehdi_amini, rriddle

Differential Review: https://reviews.llvm.org/D132838
2022-08-31 12:28:45 +02:00
Michele Scuttari
039b969b32 Revert "[MLIR] Update pass declarations to new autogenerated files"
This reverts commit 2be8af8f0e.
2022-08-30 22:21:55 +02:00
Michele Scuttari
2be8af8f0e [MLIR] Update pass declarations to new autogenerated files
The patch introduces the required changes to update the pass declarations and definitions to use the new autogenerated files and allow dropping the old infrastructure.

Reviewed By: mehdi_amini, rriddle

Differential Review: https://reviews.llvm.org/D132838
2022-08-30 21:56:31 +02:00
Matthias Springer
123c4b0251 [mlir][SCF][bufferize] Support different iter_arg/init_arg types (scf.for)
Even though iter_arg and init_arg of an scf.for loop may have the same tensor type, their bufferized memref types are not necessarily equal. It is sometimes necessary to insert a cast in case of differing layout maps.

Differential Revision: https://reviews.llvm.org/D132860
2022-08-30 16:35:32 +02:00
Matthias Springer
111c919665 [mlir][bufferization] Generalize getBufferType
This change generalizes getBufferType. This function can be used to predict the buffer type of any tensor value (not just BlockArguments) without changing any IR. It also subsumes getMemorySpace. This is useful for loop bufferization, where the precise buffer type of an iter_arg cannot be known without examining the loop body.

Differential Revision: https://reviews.llvm.org/D132859
2022-08-30 16:26:44 +02: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
Matthias Springer
c37ed7762e [tensor][bufferize] Use affine.apply instead of arith.addi in PadOp lowering
Affine exprs compose better than arith ops.

Differential Revision: https://reviews.llvm.org/D132456
2022-08-23 11:46:11 +02:00
Matthias Springer
9ee12f4778 [mlir][tensor][bufferize] Bufferize tensor.pad
tensor.pad is lowered to tensor.generate + tensor.insert_slice during bufferization. For best performance with constant padding values, users should vectorize the IR before bufferizing it.

This change also relaxes tje restriction that no new ops that bufferize to a memory write should be added during bufferization. Since bufferization has been split into two steps a while ago (tensor copy insertion + bufferization), it is reasonable to allow this now.

Differential Revision: https://reviews.llvm.org/D132355
2022-08-22 17:00:33 +02:00
Matthias Springer
1defec8730 [mlir][tensor][bufferize][NFC] Remove duplicate code
InsertSliceOp and ParallelInsertSliceOp are very similar and can share some of the bufferization analysis code.

Differential Revision: https://reviews.llvm.org/D130465
2022-07-25 12:34:16 +02:00
Matthias Springer
664ffa46bb [mlir][tensor][bufferize] Fix deallocation of GenerateOp/FromElementsOp
Both ops allocate a buffer. There were cases in which the buffer was not deallocated.

Differential Revision: https://reviews.llvm.org/D130469
2022-07-25 12:25:06 +02:00
Matthias Springer
5f5f71e737 [mlir][tensor][bufferize] Load dependent dialects
Load dialects that will be generated by the extension. (Except for BufferizationDialect and MemrefDialect which are loaded already.)

Differential Revision: https://reviews.llvm.org/D130463
2022-07-25 11:36:10 +02:00
Kazu Hirata
c27d815249 [mlir] Use value instead of getValue (NFC) 2022-07-14 00:19:59 -07:00
Jacques Pienaar
136d746ec7 [mlir] Flip accessors to prefixed form (NFC)
Another mechanical sweep to keep diff small for flip to _Prefixed.
2022-07-10 21:19:11 -07:00
Matthias Springer
606f7c8f7a [mlir][bufferization][NFC] Move more unknown type conversion logic into BufferizationOptions
The `unknownTypeConversion` bufferization option (enum) is now a type converter function option. Some logic of `getMemRefType` is now handled by that function.

This change makes type conversion more controllable. Previously, there were only two options when generating memref types for non-bufferizable ops: Static identity layout or fully dynamic layout. With this change, users of One-Shot Bufferize can provide a function with custom logic.

Differential Revision: https://reviews.llvm.org/D129273
2022-07-07 13:36:28 +02:00
Matthias Springer
6c3c5f8069 [mlir][memref] Improve type inference for rank-reducing subviews
The result shape of a rank-reducing subview cannot be inferred in the general case. Just the result rank is not enough. The only thing that we can infer is the layout map.

This change also improves the bufferization patterns of tensor.extract_slice and tensor.insert_slice to fully support rank-reducing operations.

Differential Revision: https://reviews.llvm.org/D129144
2022-07-05 16:49:07 +02: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
Matthias Springer
c0b0b6a00a [mlir][bufferize] Infer memory space in all bufferization patterns
This change updates all remaining bufferization patterns (except for scf.while) and the remaining bufferization infrastructure to infer the memory space whenever possible instead of falling back to "0". (If a default memory space is set in the bufferization options, we still fall back to that value if the memory space could not be inferred.)

Differential Revision: https://reviews.llvm.org/D128423
2022-06-27 16:32:52 +02:00
Matthias Springer
45b995cda4 [mlir][bufferize][NFC] Change signature of allocateTensorForShapedValue
Add a failure return value and bufferization options argument. This is to keep a subsequent change smaller.

Differential Revision: https://reviews.llvm.org/D128278
2022-06-27 16:00:06 +02:00
Matthias Springer
5d50f51c97 [mlir][bufferization][NFC] Add error handling to getBuffer
This is in preparation of adding memory space support.

Differential Revision: https://reviews.llvm.org/D128277
2022-06-27 13:48:01 +02:00
Matthias Springer
b06614e2e8 [mlir][bufferization][NFC] Change signature of getMemRefType
These functions now accep unsigned attributes for address spaces instead of Attributes.

Differential Revision: https://reviews.llvm.org/D128275
2022-06-27 10:41:40 +02: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
Kazu Hirata
3b7c3a654c Revert "Don't use Optional::hasValue (NFC)"
This reverts commit aa8feeefd3.
2022-06-25 11:56:50 -07:00
Kazu Hirata
aa8feeefd3 Don't use Optional::hasValue (NFC) 2022-06-25 11:55:57 -07:00
Mahesh Ravishankar
2f637fe730 [mlir][TilingInterface] Enable tile and fuse using TilingInterface.
This patch implements tile and fuse transformation for ops that
implement the tiling interface. To do so,
- `TilingInterface` needs a new method that generates a tiled
  implementation of the operation based on the tile of the result
  needed.
- A pattern is added that replaces a `tensor.extract_slice` whose
  source is defined by an operation that implements the
  `TilingInterface` with a tiled implementation that produces the
  extracted slice in-place (using the method added to
  `TilingInterface`).
- A pattern is added that takes a sequence of operations that
  implement the `TilingInterface` (for now `LinalgOp`s), tiles the
  consumer, and greedily fuses its producers iteratively.

Differential Revision: https://reviews.llvm.org/D127809
2022-06-21 17:22:58 +00:00
Mahesh Ravishankar
c584771f54 Revert "[mlir][TilingInterface] Enable tile and fuse using TilingInterface."
This reverts commit ea75511319 due to build failures.
2022-06-21 16:56:59 +00:00
Mahesh Ravishankar
ea75511319 [mlir][TilingInterface] Enable tile and fuse using TilingInterface.
This patch implements tile and fuse transformation for ops that
implement the tiling interface. To do so,
- `TilingInterface` needs a new method that generates a tiled
  implementation of the operation based on the tile of the result
  needed.
- A pattern is added that replaces a `tensor.extract_slice` whose
  source is defined by an operation that implements the
  `TilingInterface` with a tiled implementation that produces the
  extracted slice in-place (using the method added to
  `TilingInterface`).
- A pattern is added that takes a sequence of operations that
  implement the `TilingInterface` (for now `LinalgOp`s), tiles the
  consumer, and greedily fuses its producers iteratively.

Differential Revision: https://reviews.llvm.org/D127809
2022-06-21 16:47:14 +00:00
Kazu Hirata
6d5fc1e3d5 [mlir] Don't use Optional::getValue (NFC) 2022-06-20 23:20:25 -07:00
Alex Zinenko
8b68da2c7d [mlir] move SCF headers to SCF/{IR,Transforms} respectively
This aligns the SCF dialect file layout with the majority of the dialects.

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D128049
2022-06-20 10:18:01 +02:00
Jacques Pienaar
8df54a6a03 [mlir] Update accessors to prefixed form (NFC)
Follow up from flipping dialects to both, flip accessor used to prefixed
variant ahead to flipping from _Both to _Prefixed. This just flips to
the accessors introduced in the preceding change which are just prefixed
forms of the existing accessor changed from.

Mechanical change using helper script
https://github.com/jpienaar/llvm-project/blob/main/clang-tools-extra/clang-tidy/misc/AddGetterCheck.cpp and clang-format.
2022-06-18 17:53:22 -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
Matthias Springer
b55d55ecd9 [mlir][bufferize][NFC] Remove BufferizationState
With the recent refactorings, this class is no longer needed. We can use BufferizationOptions in all places were BufferizationState was used.

Differential Revision: https://reviews.llvm.org/D127653
2022-06-17 14:04:11 +02:00
Matthias Springer
b3ebe3beed [mlir][bufferize] Bufferize after TensorCopyInsertion
This change changes the bufferization so that it utilizes the new TensorCopyInsertion pass. One-Shot Bufferize no longer calls the One-Shot Analysis. Instead, it relies on the TensorCopyInsertion pass to make the entire IR fully inplacable. The `bufferize` implementations of all ops are simplified; they no longer have to account for out-of-place bufferization decisions. These were already materialized in the IR in the form of `bufferization.alloc_tensor` ops during the TensorCopyInsertion pass.

Differential Revision: https://reviews.llvm.org/D127652
2022-06-17 13:29:52 +02:00
Mogball
e16d13322b [mlir] (NFC) Clean up bazel and CMake target names
All dialect targets in bazel have been named *Dialect and all dialect
targets in CMake have been named MLIR*Dialect.
2022-06-13 16:24:15 +00:00
Matthias Springer
1534177f8f [mlir][bufferization][NFC] Move OpFilter out of BufferizationOptions
Differential Revision: https://reviews.llvm.org/D126568
2022-05-28 01:47:39 +02:00
Matthias Springer
996834e681 [mlir][SCF] Fix scf.while bufferization
Before this fix, the bufferization implementation made the incorrect assumption that the values yielded from the "before" region must match with the values yielded from the "after" region.

Differential Revision: https://reviews.llvm.org/D125835
2022-05-18 00:35:50 +02:00
Matthias Springer
8f42939a07 [mlir][bufferize][NFC] Make getContiguousMemRefType a static function
No need to expose this as public API anymore.

Differential Revision: https://reviews.llvm.org/D125361
2022-05-13 11:27:43 +02:00
Matthias Springer
248e113e9f [mlir][bufferize][NFC] Move helper functions to BufferizationOptions
Move helper functions for creating allocs/deallocs/memcpys to BufferizationOptions.

Differential Revision: https://reviews.llvm.org/D125375
2022-05-11 16:23:22 +02:00
Ashay Rane
53ff0daa7e [mlir] Fail early if AnalysisState::getBuffer() returns failure
This patch updates calls to AnalysisState::getBuffer() so that we return
early with a failure if the call does not succeed.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D125251
2022-05-10 08:08:38 -07:00
Ashay Rane
e287d647c6 [mlir] Add translation from tensor.reshape to memref.reshape
This patch augments the `tensor-bufferize` pass by adding a conversion
rule to translate ReshapeOp from the `tensor` dialect to the `memref`
dialect, in addition to adding a unit test to validate the translation.

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D125031
2022-05-09 17:45:07 +02: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
Adrian Kuegel
a74e5a89b9 [mlir] Move isGuaranteedCollapsible to CollapseShapeOp (NFC).
It seems more natural than to have it as a static method of ExpandShapeOp.
Also fix a typo ("the the" -> "the").

Differential Revision: https://reviews.llvm.org/D124234
2022-04-22 10:31:25 +02:00
Matthias Springer
d7a9bf9143 [mlir][tensor] Fix verifier and bufferization of collapse_shape
Insert a buffer copy unless the dims are guaranteed to be collapsible. In the verifier, accept collapses unless they are guaranteed to be non-collapsible.

Differential Revision: https://reviews.llvm.org/D123316
2022-04-08 18:20:40 +09:00
Matthias Springer
73c0333dee [mlir][tensor][bufferize] Support 0-d collapse_shape with offset
Differential Revision: https://reviews.llvm.org/D122901
2022-04-01 22:30:37 +09:00
Matthias Springer
51df62388e [mlir][tensor] Fix bufferization of CollapseShapeOp / ExpandShapeOp
Infer a tighter MemRef type instead of always falling back to the most dynamic MemRef type. This is inefficient and caused op verification errors.

Differential Revision: https://reviews.llvm.org/D122649
2022-03-31 17:11:45 +09:00
River Riddle
77eee5795e [mlir] Refactor DialectRegistry delayed interface support into a general DialectExtension mechanism
The current dialect registry allows for attaching delayed interfaces, that are added to attrs/dialects/ops/etc.
when the owning dialect gets loaded. This is clunky for quite a few reasons, e.g. each interface type has a
separate tracking structure, and is also quite limiting. This commit refactors this delayed mutation of
dialect constructs into a more general DialectExtension mechanism. This mechanism is essentially a registration
callback that is invoked when a set of dialects have been loaded. This allows for attaching interfaces directly
on the loaded constructs, and also allows for loading new dependent dialects. The latter of which is
extremely useful as it will now enable dependent dialects to only apply in the contexts in which they
are necessary. For example, a dialect dependency can now be conditional on if a user actually needs the
interface that relies on it.

Differential Revision: https://reviews.llvm.org/D120367
2022-03-16 22:15:25 -07:00
River Riddle
3655069234 [mlir] Move the Builtin FuncOp to the Func dialect
This commit moves FuncOp out of the builtin dialect, and into the Func
dialect. This move has been planned in some capacity from the moment
we made FuncOp an operation (years ago). This commit handles the
functional aspects of the move, but various aspects are left untouched
to ease migration: func::FuncOp is re-exported into mlir to reduce
the actual API churn, the assembly format still accepts the unqualified
`func`. These temporary measures will remain for a little while to
simplify migration before being removed.

Differential Revision: https://reviews.llvm.org/D121266
2022-03-16 17:07:03 -07:00