Skip to content

[MLIR] Generalize expand_shape to take shape as explicit input #90040

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 30, 2024
Merged

[MLIR] Generalize expand_shape to take shape as explicit input #90040

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e74c1d1
[MLIR] Generalize expand_shape to take shape as explicit input
ramiro050 Oct 17, 2023
bab64ea
[MLIR] Add test cases for illegal output shape argument
Shukla-Gaurav Apr 26, 2024
0a15b0c
[MLIR] Remove ArithDialect dependency from Dialect/Utils
Shukla-Gaurav Apr 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions mlir/include/mlir/Dialect/Arith/Utils/Utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,29 @@

namespace mlir {

using ReassociationIndices = SmallVector<int64_t, 2>;

/// Infer the output shape for a {memref|tensor}.expand_shape when it is
/// possible to do so.
///
/// Note: This should *only* be used to implement
/// `ExpandShapeOp::inferOutputShape` in both the memref and tensor namespaces.
/// If you need to infer the output shape you should use the static method of
/// `ExpandShapeOp` instead of calling this.
///
/// `inputShape` is the shape of the tensor or memref being expanded as a
/// sequence of SSA values or constants. `expandedType` is the output shape of
/// the expand_shape operation. `reassociation` is the reassociation denoting
/// the output dims each input dim is mapped to.
///
/// Returns the output shape in `outputShape` and `staticOutputShape`, following
/// the conventions for the output_shape and static_output_shape inputs to the
/// expand_shape ops.
std::optional<SmallVector<OpFoldResult>>
inferExpandShapeOutputShape(OpBuilder &b, Location loc, ShapedType expandedType,
ArrayRef<ReassociationIndices> reassociation,
ArrayRef<OpFoldResult> inputShape);

/// Matches a ConstantIndexOp.
detail::op_matcher<arith::ConstantIndexOp> matchConstantIndex();

Expand Down
80 changes: 58 additions & 22 deletions mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1548,7 +1548,6 @@ def MemRef_ReshapeOp: MemRef_Op<"reshape", [
class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
MemRef_Op<mnemonic, !listconcat(traits,
[Pure, ViewLikeOpInterface])>,
Arguments<(ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation)>,
Results<(outs AnyStridedMemRef:$result)>{

code commonExtraClassDeclaration = [{
Expand All @@ -1573,10 +1572,6 @@ class MemRef_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
Value getViewSource() { return getSrc(); }
}];

let assemblyFormat = [{
$src $reassociation attr-dict `:` type($src) `into` type($result)
}];

let hasFolder = 1;
let hasCanonicalizer = 1;
let hasVerifier = 1;
Expand All @@ -1598,14 +1593,10 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
Example:

```mlir
%r = memref.expand_shape %0 [[0, 1], [2]]
: memref<?x?xf32> into memref<?x5x?xf32>
%r = memref.expand_shape %0 [[0, 1], [2]] output_shape [%sz0, %sz1, 32]
: memref<?x32xf32> into memref<?x?x32xf32>
```

At most one dimension of a reassociation group (e.g., [0, 1] above) may be
dynamic in the result type. Otherwise, the op would be ambiguous, as it
would not be clear how the source dimension is extended.

If an op can be statically proven to be invalid (e.g, an expansion from
`memref<10xf32>` to `memref<2x6xf32>`), it is rejected by the verifier. If
it cannot statically be proven invalid (e.g., the full example above; it is
Expand All @@ -1622,41 +1613,80 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
there must be a dynamic result dimension in the corresponding reassociation
group. Same for strides.

The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
argument. A special sentinel value `ShapedType::kDynamic` encodes that the
corresponding entry has a dynamic value. There must be exactly as many SSA
inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
`static_output_shape`.

Note: This op currently assumes that the inner strides are of the
source/result layout map are the faster-varying ones.
}];

let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation,
Variadic<Index>:$output_shape,
DenseI64ArrayAttr:$static_output_shape);

let assemblyFormat = [{
$src $reassociation `output_shape`
custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
type($src) `into` type($result)
}];

let builders = [
// Builders using ReassociationIndices.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
"ArrayRef<OpFoldResult>":$outputShape)>,

// It will infer output shape using inferOutputShape() method.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation)>,

// Builder using ReassociationExprs.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationExprs>":$reassociation),
[{
build($_builder, $_state, resultType, src, attrs);
$_state.addAttribute("reassociation",
getReassociationIndicesAttribute($_builder, reassociation));
auto reassociationIndices =
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationIndices);
}]>,

// Builder using ReassociationExprs.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationExprs>":$reassociation,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
"ArrayRef<OpFoldResult>":$outputShape),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
build($_builder, $_state, resultType, src, reassociationMaps, attrs);
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationMaps,
outputShape);
}]>,

// Builder that infers the result layout map. The result shape must be
// specified. Otherwise, the op may be ambiguous. The output shape for
// the op will be inferred using the inferOutputShape() method.
OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation)>,

// Builder that infers the result layout map. The result shape must be
// specified. Otherwise, the op may be ambiguous.
OpBuilder<(ins "ArrayRef<int64_t>":$resultShape, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation)>
"ArrayRef<ReassociationIndices>":$reassociation,
"ArrayRef<OpFoldResult>":$outputShape)>
];

let extraClassDeclaration = commonExtraClassDeclaration # [{
static FailureOr<MemRefType> computeExpandedType(
MemRefType srcType, ArrayRef<int64_t> resultShape,
ArrayRef<ReassociationIndices> reassociation);

// Infer the output shape for a memref.expand_shape when it is possible
// to do so.
static FailureOr<SmallVector<OpFoldResult>> inferOutputShape(
OpBuilder &b, Location loc, MemRefType expandedType,
ArrayRef<ReassociationIndices> reassociation,
ArrayRef<OpFoldResult> inputShape);
}];

let hasVerifier = 1;
Expand Down Expand Up @@ -1707,6 +1737,12 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
source/result layout map are the faster-varying ones.
}];

let arguments = (ins AnyStridedMemRef:$src, IndexListArrayAttr:$reassociation);

let assemblyFormat = [{
$src $reassociation attr-dict `:` type($src) `into` type($result)
}];

let builders = [
// Builders for a contracting reshape whose result type is computed from
// `src` and `reassociation`.
Expand All @@ -1718,7 +1754,7 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, src, reassociationMaps, attrs);
}]>,

Expand All @@ -1736,7 +1772,7 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationMaps, attrs);
}]>
];
Expand Down
79 changes: 56 additions & 23 deletions mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1062,8 +1062,7 @@ class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
Tensor_Op<mnemonic, !listconcat(traits, [
DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
Pure])>,
Arguments<(ins AnyRankedTensor:$src, IndexListArrayAttr:$reassociation)>,
Results<(outs AnyRankedTensor:$result)> {
Results<(outs AnyTensor:$result)> {

code commonExtraClassDeclaration = [{
static StringRef getReassociationAttrStrName() { return "reassociation"; }
Expand All @@ -1086,10 +1085,6 @@ class Tensor_ReassociativeReshapeOp<string mnemonic, list<Trait> traits = []> :
}
}];

let assemblyFormat = [{
$src $reassociation attr-dict `:` type($src) `into` type($result)
}];

let hasFolder = 1;
let hasCanonicalizer = 1;
let hasVerifier = 1;
Expand All @@ -1102,50 +1097,83 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
rank than the operand `src` whose dimension sizes are a reassociation of
`src`.

A reassociation is defined as a continuous grouping of dimensions. It is
represented with an array of DenseI64ArrayAttr attribute. Entries in the
array are referred to as reassociation maps.
A reassociation is defined as a continuous grouping of dimensions and is
represented with an array of DenseI64ArrayAttr attribute. The reassociation
maps applied to the result tensor with the higher rank must result in the
operand tensor with the smaller rank.

The reassociation maps are applied to the result shape to obtain the operand
shape.
The representation for the output shape supports a partially-static
specification via attributes specified through the `static_output_shape`
argument. A special sentinel value `ShapedType::kDynamic` encodes that the
corresponding entry has a dynamic value. There must be exactly as many SSA
inputs in `output_shape` as there are `ShapedType::kDynamic` entries in
`static_output_shape`.

Example:

```mlir
// Dimension expansion i -> (i', j') and (k) -> (k')
%b = tensor.expand_shape %a [[0, 1], [2]]
: tensor<?x?xf32> into tensor<?x?x?xf32>
%b = tensor.expand_shape %a [[0, 1], [2]] output_shape [%sz0, %sz1, 32]
: tensor<?x32xf32> into tensor<?x?x32xf32>
```
}];

let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation,
Variadic<Index>:$output_shape,
DenseI64ArrayAttr:$static_output_shape);

let assemblyFormat = [{
$src $reassociation `output_shape`
custom<DynamicIndexList>($output_shape, $static_output_shape) attr-dict `:`
type($src) `into` type($result)
}];

let builders = [
// Builders using ReassociationIndices.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
"ArrayRef<OpFoldResult>":$outputShape)>,

// It will infer output shape using inferOutputShape() method.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationIndices>":$reassociation)>,

// Builder using ReassociationExprs.
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationExprs>":$reassociation),
[{
build($_builder, $_state, resultType, src, attrs);
$_state.addAttribute("reassociation",
getReassociationIndicesAttribute($_builder, reassociation));
auto reassociationIndices =
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationIndices);
}]>,
OpBuilder<(ins "Type":$resultType, "Value":$src,
"ArrayRef<ReassociationExprs>":$reassociation,
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
"ArrayRef<OpFoldResult>":$outputShape),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
build($_builder, $_state, resultType, src, reassociationMaps, attrs);
auto reassociationIndices =
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationIndices,
outputShape);
}]>
];

let extraClassDeclaration = commonExtraClassDeclaration # [{
int64_t getCorrespondingSourceDim(int64_t resultDim);

// Infer the output shape for a tensor.expand_shape when it is possible
// to do so.
static FailureOr<SmallVector<OpFoldResult>> inferOutputShape(
OpBuilder &b, Location loc, RankedTensorType expandedType,
ArrayRef<ReassociationIndices> reassociation,
ArrayRef<OpFoldResult> inputShape);
}];

let hasVerifier = 1;
}

def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
let summary = "operation to produce a tensor with a smaller rank";
let arguments = (ins AnyTensor:$src, IndexListArrayAttr:$reassociation);
let description = [{
The `tensor.collapse_shape` op produces a new tensor of lower (or equal)
rank whose dimension sizes are a reassociation of the original `src` dimensions.
Expand All @@ -1163,6 +1191,11 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
: tensor<?x?x?xf32> into tensor<?x?xf32>
```
}];

let assemblyFormat = [{
$src $reassociation attr-dict `:` type($src) `into` type($result)
}];

let builders = [
// Builders for a contracting reshape whose result type is computed from
// `src` and `reassociation`.
Expand All @@ -1174,7 +1207,7 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, src, reassociationMaps, attrs);
}]>,

Expand All @@ -1192,7 +1225,7 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices($_builder, reassociation);
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationMaps, attrs);
}]>
];
Expand Down
Loading