This class assists with generating IR required to materialize an arbitrary-sized slice from the result of a CollapseShapeOp. More...

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

Public Member Functions
	ExtractSliceFromCollapseHelper (tensor::CollapseShapeOp collapseShapeOp, ArrayRef< OpFoldResult > collapseShapeInputShape, ArrayRef< OpFoldResult > collapseShapeOutputShape, ArrayRef< Range > extractSliceParams, const llvm::SmallBitVector &linearizedDimensions, const llvm::SmallBitVector &slicedDimensions, ArrayRef< Value > tiledSizes)

const SmallVector< Value > &	getIterationSpaceSizes ()
	Return the upper bounds of the iteration space (with 0 offset and stride 1) required to create the desired slice. More...

std::pair< Value, SmallVector< Range > >	emitLoopNestBody (OpBuilder &builder, Location loc, ValueRange tileInductionVars)
	Generates the IR inside of the caller's loop nest for 1) inverting the index mappings of the ExtractSliceOp->CollapseShapeOp chain and 2) extracting the CollapseShapeOp source tensor tile for this specified iteration space point `tileInductionVars` and 3) calculating where to insert the extracted tile. More...

Static Public Member Functions
static FailureOr< ExtractSliceFromCollapseHelper >	create (OpBuilder &b, tensor::CollapseShapeOp op, ArrayRef< Range > sliceParams)
	Given a CollapseShapeOp and a set of ranges describing the desired slice of its result, emits IR to materialize the shapes of the input and output tensors, and returns an instance of the initialized class. More...

static FailureOr< ExtractSliceFromCollapseHelper >	create (OpBuilder &b, tensor::CollapseShapeOp collapseOp, tensor::ExtractSliceOp extractOp)
	Given a CollapseShapeOp and an ExtractSliceOp acting on its result, emits IR to materialize the shapes of the input and output tensors of the CollapseShapeOp, and returns an instance of the initialized class. More...

Detailed Description

This class assists with generating IR required to materialize an arbitrary-sized slice from the result of a CollapseShapeOp.

In order to accomplish this, a loop nest or similar operation must be created by the caller. The purpose of the loop nest is to generate a "tiling by 1" of all sliced dimensions. The "tiling by 1" assembles all elements of the result tile over dimensions that would have been impossible to directly slice.

The class provides three methods:

ExtractSliceFromCollapseHelper::create: emits IR that should appear before the loop nest and populates the internal state.
ExtractSliceFromCollapseHelper::getIterationSpaceSizes: returns parameters used by the caller to construct the loop nest.
ExtractSliceFromCollapseHelper::emitLoopNestBody: emits IR to construct a "size-1 tile" of the desired result and returns a set of ranges where the tile should be inserted into the destination tensor.

Intended usage:

The caller should first call ExtractSliceFromCollapseHelper::create and then create a destination tensor that is the same size as the desired slice. The caller then creates a loop nest that iterates over the multi-dimensional iteration space defined by [0, ub[0]) x [0, ub[1]] x ... x [0, ub[N-1]] where ub is the upper bound given by ExtractSliceFromCollapseHelper::getIterationSpaceSizes. Inside the body of the loop nest, the caller should call ExtractSliceFromCollapseHelper::emitLoopNestBody and provide the induction variables. This returns a sub-tile and a set of ranges that describe where this tile should be inserted into the result by the caller. For a complete example of usage, see the examples in the TestTensorTransforms pass.

Example:

Consider the following IR:

%0 = linalg.generic ... -> tensor<3x?x?x11x?xf32>
%1 = tensor.collapse_shape %0 [[0, 1, 2], [3, 4]]
       : tensor<3x?x?x11x?xf32> into tensor<?x?xf32>
%2 = tensor.extract_slice %1 [%offt0, %offt1][%size0, %size1][1, 1]
       : tensor<?x?xf32> to tensor<?x?xf32>

We can construct %2 by generating the following, which only uses %0:

%dest = tensor.empty(%size0, %size1) : tensor<?x?xf32>
%1 = tensor.dim %0, %c1 : tensor<3x?x?x11x?xf32>
%2 = tensor.dim %0, %c2 : tensor<3x?x?x11x?xf32>
%3 = tensor.dim %0, %c4 : tensor<3x?x?x11x?xf32>
 
%result = scf.for %iv0 = %c0 to %arg2 step %c1 iter_args(%arg6 = %dest) ->
                                                 (tensor<?x?xf32>) {
  %5 = scf.for %iv1 = %c0 to %arg4 step %c1 iter_args(%arg8 = %arg6)
                                                 -> (tensor<?x?xf32>) {
    %lin0 = (affine.apply) %iv0 + %offt0
    %lin1 = (affine.apply) %iv1 + %offt1
 
    %mi0:3 = affine.delinearize_index %lin0 into (%c3, %1, %2)
    %mi1:2 = affine.delinearize_index %lin1 into (%c11, %3)
 
    %sub_tile = tensor.extract_slice %0
                   [%mi0#0, %mi0#1, %mi0#2, %mi1#0, %mi1#1]
                   [1, 1, 1, 1, 1]
                   [1, 1, 1, 1, 1]
           : tensor<3x?x?x11x?xf32> to tensor<1x1x1x1x1xf32>
    %sub_tile_collapsed = tensor.collapse_shape %sub_tile
            [[0, 1, 2], [3, 4]]
           : tensor<1x1x1x1x1xf32> into tensor<1x1xf3
 
    %12 = tensor.insert_slice %sub_tile_collapsed into
            %arg8[%iv0, %iv1] [1, 1] [1, 1]
            : tensor<1x1xf32> into tensor<?x?xf32>
    scf.yield %12 : tensor<?x?xf32>
  }
  scf.yield %5 : tensor<?x?xf32>
}

Explanation of example:

Each step above is explained below.

Step 0: Create %dest and materialization of shapes.

This step is self-explanatory and performed by the caller. It can be done before or after calling ExtractSliceFromCollapseHelper::create, which materializes the source shape (%0, %1, %2).

Step 1: Create loop nest.

The caller creates the loop nest (depicted here is scf.for, but any other similar op can be used). The iteration should start at zero and proceed with step size 1 to the upper bounds given by ExtractSliceFromCollapseHelper::getIterationSpaceSizes. This forms the basis for the "tiling by 1".

Step 2: Transform (%iv0, %iv1) from the index space of %3 to the index

space of %0.

This step is performed by ExtractSliceFromCollapseHelper::emitLoopNestBody.

The induction variables iv0 and iv1 live in the index space of %2 (for dimensions 0 and 1, respectively). lin0 and lin1 are the result of inverting or resolve the index space transformation represented by the slice operation, accounting for offset and stride. Subsequently, mi0 and mi1 are the result of applying the inverse index space transformation represented by tensor.collapse_shape. This is accomplished using affine.delinearize_index. Note that iv0 and iv1 now correspond to multi-indices mi0:3 and mi1:2.

Step 3: Extract a sub-tile slice from the source.

This step is also performed by ExtractSliceFromCollapseHelper::emitLoopNestBody.

The indices mi0 and mi1 are used to extract a slice from %0. This slice is then collapsed down to match the result rank.

Step 4: Insert sub-tile into the destination

This step is performed by the caller using the results of ExtractSliceFromCollapseHelper::emitLoopNestBody.

In the above example, the slice insertion parameters are straightforward, but in other possible situations, the slice parameters are more complicated, which is why this helper calculates them for the caller. These other situations correspond to:

The presence of linearized dimensions that are not sliced
The presence of non-linearized dimensions that are sliced.

Definition at line 149 of file TransformUtils.h.

Constructor & Destructor Documentation

◆ ExtractSliceFromCollapseHelper()

mlir::tensor::ExtractSliceFromCollapseHelper::ExtractSliceFromCollapseHelper	(	tensor::CollapseShapeOp	collapseShapeOp,
		ArrayRef< OpFoldResult >	collapseShapeInputShape,
		ArrayRef< OpFoldResult >	collapseShapeOutputShape,
		ArrayRef< Range >	extractSliceParams,
		const llvm::SmallBitVector &	linearizedDimensions,
		const llvm::SmallBitVector &	slicedDimensions,
		ArrayRef< Value >	tiledSizes
	)

inline

Definition at line 166 of file TransformUtils.h.

Member Function Documentation

◆ create() [1/2]

FailureOr< ExtractSliceFromCollapseHelper > tensor::ExtractSliceFromCollapseHelper::create	(	OpBuilder &	b,
		tensor::CollapseShapeOp	collapseOp,
		tensor::ExtractSliceOp	extractOp
	)

static

Given a CollapseShapeOp and an ExtractSliceOp acting on its result, emits IR to materialize the shapes of the input and output tensors of the CollapseShapeOp, and returns an instance of the initialized class.

Returns failure if the slice is rank-reducing.

Definition at line 63 of file ExtractSliceFromReshapeUtils.cpp.

◆ create() [2/2]

FailureOr< ExtractSliceFromCollapseHelper > tensor::ExtractSliceFromCollapseHelper::create	(	OpBuilder &	b,
		tensor::CollapseShapeOp	op,
		ArrayRef< Range >	sliceParams
	)

static

Given a CollapseShapeOp and a set of ranges describing the desired slice of its result, emits IR to materialize the shapes of the input and output tensors, and returns an instance of the initialized class.

Returns failure if the slice is rank-reducing.

Definition at line 80 of file ExtractSliceFromReshapeUtils.cpp.

References mlir::getLinearizedDimensions(), mlir::tensor::getMixedSizes(), mlir::getSlicedDimensions(), mlir::getValueOrCreateConstantIndexOp(), and mlir::reifyResultShapes().

◆ emitLoopNestBody()

std::pair< Value, SmallVector< Range > > tensor::ExtractSliceFromCollapseHelper::emitLoopNestBody	(	OpBuilder &	builder,
		Location	loc,
		ValueRange	tileInductionVars
	)

Generates the IR inside of the caller's loop nest for 1) inverting the index mappings of the ExtractSliceOp->CollapseShapeOp chain and 2) extracting the CollapseShapeOp source tensor tile for this specified iteration space point tileInductionVars and 3) calculating where to insert the extracted tile.

The returned pair consists of the results of (2) and (3) and should be used by the caller to insert into the destination tensor.

Definition at line 121 of file ExtractSliceFromReshapeUtils.cpp.

References mlir::OpBuilder::create(), mlir::Builder::getContext(), invertCollapseShapeIndexing(), and invertSliceIndexing().

◆ getIterationSpaceSizes()

const SmallVector<Value>& mlir::tensor::ExtractSliceFromCollapseHelper::getIterationSpaceSizes ( )