Skip to content

Commit

Permalink
Revert "[MLIR] Generalize expand_shape to take shape as explicit input (
Browse files Browse the repository at this point in the history 
#69267)"

This reverts commit e095d97.
  • Loading branch information
@joker-eph
joker-eph authored Apr 21, 2024
1 parent e095d97 commit b8bffad
Show file tree
Hide file tree
Showing 52 changed files with 634 additions and 1,193 deletions.
80 changes: 22 additions & 58 deletions mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1548,6 +1548,7 @@ 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 @@ -1572,6 +1573,10 @@ 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 @@ -1593,10 +1598,14 @@ def MemRef_ExpandShapeOp : MemRef_ReassociativeReshapeOp<"expand_shape", [
Example:

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

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 @@ -1613,80 +1622,41 @@ 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,
"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),
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationIndices =
convertReassociationMapsToIndices(reassociation);
build($_builder, $_state, resultType, src, reassociationIndices);
build($_builder, $_state, resultType, src, attrs);
$_state.addAttribute("reassociation",
getReassociationIndicesAttribute($_builder, reassociation));
}]>,

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

// 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<OpFoldResult>":$outputShape)>
"ArrayRef<ReassociationIndices>":$reassociation)>
];

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 @@ -1737,12 +1707,6 @@ 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 @@ -1754,7 +1718,7 @@ def MemRef_CollapseShapeOp : MemRef_ReassociativeReshapeOp<"collapse_shape", [
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices(reassociation);
convertReassociationMapsToIndices($_builder, reassociation);
build($_builder, $_state, src, reassociationMaps, attrs);
}]>,

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

code commonExtraClassDeclaration = [{
static StringRef getReassociationAttrStrName() { return "reassociation"; }
Expand All @@ -1085,6 +1086,10 @@ 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 @@ -1097,83 +1102,50 @@ 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 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.
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.

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`.
The reassociation maps are applied to the result shape to obtain the operand
shape.

Example:

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

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

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 @@ -1191,11 +1163,6 @@ 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 @@ -1207,7 +1174,7 @@ def Tensor_CollapseShapeOp : Tensor_ReassociativeReshapeOp<"collapse_shape"> {
CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs),
[{
auto reassociationMaps =
convertReassociationMapsToIndices(reassociation);
convertReassociationMapsToIndices($_builder, reassociation);
build($_builder, $_state, src, reassociationMaps, attrs);
}]>,

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

0 comments on commit b8bffad

Please sign in to comment.