08cf6ae537
This pass reifies the shapes of a subset of
`ReifyRankedShapedTypeOpInterface` ops with `tensor` results.
The pass currently only supports result shape type reification for:
- tensor::PadOp
- tensor::ConcatOp
It addresses a representation gap where implicit op semantics are needed
to infer static result types from dynamic
operands. But it does so by using `ReifyRankedShapedTypeOpInterface` as
the source of truth rather than the op itself.
As a consequence, this cannot generalize today.
TODO: in the future, we should consider coupling this information with
op "transfer functions" (e.g.
`IndexingMapOpInterface`) to provide a source of truth that can work
across result shape inference, canonicalization and
op verifiers.
The pass replaces the operations with their reified versions, when more
static information can be derived, and inserts
casts when results shapes are updated.
Example:
```mlir
#map = affine_map<(d0) -> (-d0 + 256)>
func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
%0 = affine.apply #map(%arg1)
%extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
%padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %arg0 : f32
} : tensor<1x?x64xf32> to tensor<1x?x64xf32>
return %padded : tensor<1x?x64xf32>
}
// mlir-opt --reify-result-shapes
#map = affine_map<()[s0] -> (-s0 + 256)>
func.func @func(%arg0: f32, %arg1: index, %arg2: tensor<64x?x64xf32>) -> tensor<1x?x64xf32> {
%0 = affine.apply #map()[%arg1]
%extracted_slice = tensor.extract_slice %arg2[0, 0, 0] [1, %arg1, 64] [1, 1, 1] : tensor<64x?x64xf32> to tensor<1x?x64xf32>
%padded = tensor.pad %extracted_slice low[0, 0, 0] high[0, %0, 0] {
^bb0(%arg3: index, %arg4: index, %arg5: index):
tensor.yield %arg0 : f32
} : tensor<1x?x64xf32> to tensor<1x256x64xf32>
%cast = tensor.cast %padded : tensor<1x256x64xf32> to tensor<1x?x64xf32>
return %cast : tensor<1x?x64xf32>
}
```
---------
Co-authored-by: Fabian Mora <fabian.mora-cordero@amd.com>
46 lines
950 B
CMake
46 lines
950 B
CMake
add_mlir_dialect_library(MLIRMemRefTransforms
|
|
AllocationOpInterfaceImpl.cpp
|
|
BufferViewFlowOpInterfaceImpl.cpp
|
|
ComposeSubView.cpp
|
|
ExpandOps.cpp
|
|
ExpandRealloc.cpp
|
|
ExpandStridedMetadata.cpp
|
|
EmulateWideInt.cpp
|
|
EmulateNarrowType.cpp
|
|
ExtractAddressComputations.cpp
|
|
FlattenMemRefs.cpp
|
|
FoldMemRefAliasOps.cpp
|
|
IndependenceTransforms.cpp
|
|
MultiBuffer.cpp
|
|
NormalizeMemRefs.cpp
|
|
ReifyResultShapes.cpp
|
|
ResolveShapedTypeResultDims.cpp
|
|
RuntimeOpVerification.cpp
|
|
|
|
ADDITIONAL_HEADER_DIRS
|
|
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/MemRef
|
|
|
|
DEPENDS
|
|
MLIRMemRefPassIncGen
|
|
|
|
LINK_LIBS PUBLIC
|
|
MLIRAffineTransforms
|
|
MLIRAffineUtils
|
|
MLIRArithDialect
|
|
MLIRArithTransforms
|
|
MLIRDialectUtils
|
|
MLIRFuncDialect
|
|
MLIRGPUDialect
|
|
MLIRInferTypeOpInterface
|
|
MLIRLoopLikeInterface
|
|
MLIRMemRefDialect
|
|
MLIRMemRefUtils
|
|
MLIRNVGPUDialect
|
|
MLIRPass
|
|
MLIRTensorDialect
|
|
MLIRTransforms
|
|
MLIRValueBoundsOpInterface
|
|
MLIRVectorDialect
|
|
)
|
|
|