Revert "[MLIR] Generalize expand_shape to take shape as explicit input" #89540
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
joker-eph
added
the
skip-precommit-approval
joker-eph
requested review from
rupprecht,
aartbik,
PeimingLiu,
yinying-lisa-li,
matthias-springer,
hanhanW,
nicolasvasilache and
MaheshRavishankar
as code owners
llvmbot
added
mlir:linalg
mlir:sparse
|
@llvm/pr-subscribers-mlir-linalg @llvm/pr-subscribers-mlir-tensor Author: Mehdi Amini (joker-eph) ChangesReverts llvm/llvm-project#69267 this broke some bots. Patch is 259.47 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/89540.diff 52 Files Affected:
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 14b8d95ea15b41..39e66cd9e6e5ab 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -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 = [{
@@ -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;
@@ -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
@@ -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;
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index a403e89a39f98c..cf7f3e89079c1c 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -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"; }
@@ -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;
@@ -1097,75 +1102,43 @@ 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;
@@ -1173,7 +1146,6 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
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.
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 8a41a0a18b0ab3..ae9824f728da4d 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -30,27 +30,6 @@ using ReassociationExprs = SmallVector<AffineExpr, 2>;
/// Attribute name for the ArrayAttr which encodes reassociation indices.
constexpr StringRef getReassociationAttrName() { return "reassociation"; }
-// 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);
-
/// Compose reassociation maps that are used in pair of reshape ops where one
/// is a producer and other is the consumer. Only valid to use this method when
/// both the producer and consumer are collapsing dimensions or both are
@@ -83,7 +62,7 @@ getReassociationIndicesAttribute(OpBuilder &b,
/// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
- ArrayRef<ReassociationExprs> reassociationExprs);
+ OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
/// Return the reassociations maps to use to reshape given the source type and
/// the target type when possible. Return std::nullopt when this computation
@@ -161,11 +140,14 @@ static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
op.getReassociationIndices(), isExpansion);
}
-/// Verify that shapes of the reshaped types using following rule:
-/// if a dimension in the collapsed type is static, then the corresponding
-/// dimensions in the expanded shape should be
+/// Verify that shapes of the reshaped types using following rules
+/// 1) if a dimension in the collapsed type is static, then the corresponding
+/// dimensions in the expanded shape should be
/// a) static
/// b) the product should be same as the collaped shape.
+/// 2) if a dimension in the collaped type is dynamic, one and only one of the
+/// corresponding dimensions in the expanded type should be dynamic. This
+/// rule is only needed with reshape operations that are expanding.
LogicalResult reshapeLikeShapesAreCompatible(
function_ref<LogicalResult(const Twine &)> emitError,
ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
@@ -174,11 +156,9 @@ LogicalResult reshapeLikeShapesAreCompatible(
/// Returns true iff the type is a MemRefType and has a non-identity layout.
bool hasNonIdentityLayout(Type type);
-enum class ReshapeOpKind { kExpand, kCollapse };
-
/// Pattern to collapse producer/consumer reshape ops that are both collapsing
/// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy, ReshapeOpKind opKind>
+template <typename ReshapeOpTy>
struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
@@ -201,18 +181,8 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
rewriter.getContext());
if (!reassociationIndices)
return failure();
-
- if constexpr (opKind == ReshapeOpKind::kExpand) {
- SmallVector<OpFoldResult> outputShape(
- getMixedValues(reshapeOp.getStaticOutputShape(),
- reshapeOp.getOutputShape(), rewriter));
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
- outputShape);
- } else {
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
- }
+ rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+ reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
return success();
}
};
@@ -245,8 +215,7 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
//
/// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
/// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy,
- typename DimOpTy, typename TensorTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
@@ -353,11 +322,8 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
if (!composedReassociation)
return failure();
- SmallVector<OpFoldResult> outputShape(getMixedValues(
- expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter));
rewriter.replaceOpWithNewOp<ExpandOpTy>(
- expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
- outputShape);
+ expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
return success();
}
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index 594bcf5dbb399a..20f019666a2e6a 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -125,8 +125,9 @@ SmallVect...
[truncated]
|
|
@llvm/pr-subscribers-mlir-memref Author: Mehdi Amini (joker-eph) ChangesReverts llvm/llvm-project#69267 this broke some bots. Patch is 259.47 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/89540.diff 52 Files Affected:
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 14b8d95ea15b41..39e66cd9e6e5ab 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -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 = [{
@@ -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;
@@ -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
@@ -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;
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index a403e89a39f98c..cf7f3e89079c1c 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -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"; }
@@ -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;
@@ -1097,75 +1102,43 @@ 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;
@@ -1173,7 +1146,6 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
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.
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 8a41a0a18b0ab3..ae9824f728da4d 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -30,27 +30,6 @@ using ReassociationExprs = SmallVector<AffineExpr, 2>;
/// Attribute name for the ArrayAttr which encodes reassociation indices.
constexpr StringRef getReassociationAttrName() { return "reassociation"; }
-// 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);
-
/// Compose reassociation maps that are used in pair of reshape ops where one
/// is a producer and other is the consumer. Only valid to use this method when
/// both the producer and consumer are collapsing dimensions or both are
@@ -83,7 +62,7 @@ getReassociationIndicesAttribute(OpBuilder &b,
/// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
- ArrayRef<ReassociationExprs> reassociationExprs);
+ OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
/// Return the reassociations maps to use to reshape given the source type and
/// the target type when possible. Return std::nullopt when this computation
@@ -161,11 +140,14 @@ static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
op.getReassociationIndices(), isExpansion);
}
-/// Verify that shapes of the reshaped types using following rule:
-/// if a dimension in the collapsed type is static, then the corresponding
-/// dimensions in the expanded shape should be
+/// Verify that shapes of the reshaped types using following rules
+/// 1) if a dimension in the collapsed type is static, then the corresponding
+/// dimensions in the expanded shape should be
/// a) static
/// b) the product should be same as the collaped shape.
+/// 2) if a dimension in the collaped type is dynamic, one and only one of the
+/// corresponding dimensions in the expanded type should be dynamic. This
+/// rule is only needed with reshape operations that are expanding.
LogicalResult reshapeLikeShapesAreCompatible(
function_ref<LogicalResult(const Twine &)> emitError,
ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
@@ -174,11 +156,9 @@ LogicalResult reshapeLikeShapesAreCompatible(
/// Returns true iff the type is a MemRefType and has a non-identity layout.
bool hasNonIdentityLayout(Type type);
-enum class ReshapeOpKind { kExpand, kCollapse };
-
/// Pattern to collapse producer/consumer reshape ops that are both collapsing
/// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy, ReshapeOpKind opKind>
+template <typename ReshapeOpTy>
struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
@@ -201,18 +181,8 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
rewriter.getContext());
if (!reassociationIndices)
return failure();
-
- if constexpr (opKind == ReshapeOpKind::kExpand) {
- SmallVector<OpFoldResult> outputShape(
- getMixedValues(reshapeOp.getStaticOutputShape(),
- reshapeOp.getOutputShape(), rewriter));
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
- outputShape);
- } else {
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
- }
+ rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+ reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
return success();
}
};
@@ -245,8 +215,7 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
//
/// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
/// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy,
- typename DimOpTy, typename TensorTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
@@ -353,11 +322,8 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
if (!composedReassociation)
return failure();
- SmallVector<OpFoldResult> outputShape(getMixedValues(
- expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter));
rewriter.replaceOpWithNewOp<ExpandOpTy>(
- expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
- outputShape);
+ expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
return success();
}
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index 594bcf5dbb399a..20f019666a2e6a 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -125,8 +125,9 @@ SmallVect...
[truncated]
|
|
@llvm/pr-subscribers-mlir-bufferization Author: Mehdi Amini (joker-eph) ChangesReverts llvm/llvm-project#69267 this broke some bots. Patch is 259.47 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/89540.diff 52 Files Affected:
diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
index 14b8d95ea15b41..39e66cd9e6e5ab 100644
--- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
+++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td
@@ -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 = [{
@@ -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;
@@ -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
@@ -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;
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index a403e89a39f98c..cf7f3e89079c1c 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -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"; }
@@ -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;
@@ -1097,75 +1102,43 @@ 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;
@@ -1173,7 +1146,6 @@ def Tensor_ExpandShapeOp : Tensor_ReassociativeReshapeOp<"expand_shape"> {
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.
@@ -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`.
@@ -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);
}]>,
@@ -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);
}]>
];
diff --git a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
index 8a41a0a18b0ab3..ae9824f728da4d 100644
--- a/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/ReshapeOpsUtils.h
@@ -30,27 +30,6 @@ using ReassociationExprs = SmallVector<AffineExpr, 2>;
/// Attribute name for the ArrayAttr which encodes reassociation indices.
constexpr StringRef getReassociationAttrName() { return "reassociation"; }
-// 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);
-
/// Compose reassociation maps that are used in pair of reshape ops where one
/// is a producer and other is the consumer. Only valid to use this method when
/// both the producer and consumer are collapsing dimensions or both are
@@ -83,7 +62,7 @@ getReassociationIndicesAttribute(OpBuilder &b,
/// Convert Array<Array<AffineExpr>> to Array<Array<int64_t>>.
SmallVector<ReassociationIndices, 2> convertReassociationMapsToIndices(
- ArrayRef<ReassociationExprs> reassociationExprs);
+ OpBuilder &b, ArrayRef<ReassociationExprs> reassociationExprs);
/// Return the reassociations maps to use to reshape given the source type and
/// the target type when possible. Return std::nullopt when this computation
@@ -161,11 +140,14 @@ static LogicalResult verifyReshapeLikeTypes(Op op, T expandedType,
op.getReassociationIndices(), isExpansion);
}
-/// Verify that shapes of the reshaped types using following rule:
-/// if a dimension in the collapsed type is static, then the corresponding
-/// dimensions in the expanded shape should be
+/// Verify that shapes of the reshaped types using following rules
+/// 1) if a dimension in the collapsed type is static, then the corresponding
+/// dimensions in the expanded shape should be
/// a) static
/// b) the product should be same as the collaped shape.
+/// 2) if a dimension in the collaped type is dynamic, one and only one of the
+/// corresponding dimensions in the expanded type should be dynamic. This
+/// rule is only needed with reshape operations that are expanding.
LogicalResult reshapeLikeShapesAreCompatible(
function_ref<LogicalResult(const Twine &)> emitError,
ArrayRef<int64_t> collapsedShape, ArrayRef<int64_t> expandedShape,
@@ -174,11 +156,9 @@ LogicalResult reshapeLikeShapesAreCompatible(
/// Returns true iff the type is a MemRefType and has a non-identity layout.
bool hasNonIdentityLayout(Type type);
-enum class ReshapeOpKind { kExpand, kCollapse };
-
/// Pattern to collapse producer/consumer reshape ops that are both collapsing
/// dimensions or are both expanding dimensions.
-template <typename ReshapeOpTy, ReshapeOpKind opKind>
+template <typename ReshapeOpTy>
struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
using OpRewritePattern<ReshapeOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(ReshapeOpTy reshapeOp,
@@ -201,18 +181,8 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
rewriter.getContext());
if (!reassociationIndices)
return failure();
-
- if constexpr (opKind == ReshapeOpKind::kExpand) {
- SmallVector<OpFoldResult> outputShape(
- getMixedValues(reshapeOp.getStaticOutputShape(),
- reshapeOp.getOutputShape(), rewriter));
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices,
- outputShape);
- } else {
- rewriter.replaceOpWithNewOp<ReshapeOpTy>(
- reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
- }
+ rewriter.replaceOpWithNewOp<ReshapeOpTy>(
+ reshapeOp, resultType, srcReshapeOp.getSrc(), *reassociationIndices);
return success();
}
};
@@ -245,8 +215,7 @@ struct ComposeReassociativeReshapeOps : public OpRewritePattern<ReshapeOpTy> {
//
/// When `rank(srcType) < rank(resultType)`, then we just swap `reassociation_1`
/// `reassociation_2` and produce `expand_shape`.
-template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy,
- typename DimOpTy, typename TensorTy>
+template <typename CollapseOpTy, typename ExpandOpTy, typename CastOpTy>
struct ComposeCollapseOfExpandOp : public OpRewritePattern<CollapseOpTy> {
using OpRewritePattern<CollapseOpTy>::OpRewritePattern;
LogicalResult matchAndRewrite(CollapseOpTy collapseOp,
@@ -353,11 +322,8 @@ struct ComposeExpandOfCollapseOp : public OpRewritePattern<ExpandOpTy> {
if (!composedReassociation)
return failure();
- SmallVector<OpFoldResult> outputShape(getMixedValues(
- expandOp.getStaticOutputShape(), expandOp.getOutputShape(), rewriter));
rewriter.replaceOpWithNewOp<ExpandOpTy>(
- expandOp, resultType, collapseOp.getSrc(), *composedReassociation,
- outputShape);
+ expandOp, resultType, collapseOp.getSrc(), *composedReassociation);
return success();
}
diff --git a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
index 594bcf5dbb399a..20f019666a2e6a 100644
--- a/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/StaticValueUtils.h
@@ -125,8 +125,9 @@ SmallVect...
[truncated]
|
aniplcc
pushed a commit
to aniplcc/llvm-project
that referenced
this pull request
Apr 21, 2024
…t" (llvm#89540) Reverts llvm#69267 this broke some bots.
This was referenced Apr 22, 2024
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Reverts #69267
this broke some bots.