[Mlir-commits] [mlir] [Linalg] Fix SoftmaxOp's reify result shape calculation (PR #67790)

[llvmlistbot at llvm.org]

================
@@ -2387,8 +2387,23 @@ LogicalResult SoftmaxOp::fold(FoldAdaptor, SmallVectorImpl<OpFoldResult> &) {
 LogicalResult
 SoftmaxOp::reifyResultShapes(OpBuilder &b,
                              ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
-  return cast<LinalgOp>(getOperation())
-      .reifyResultShapes(b, reifiedReturnShapes);
+  SmallVector<OpFoldResult> shapes;
+  Location loc = getOperation()->getLoc();
+  IRRewriter rewriter(b);
+  auto inputShapedType = llvm::cast<ShapedType>(getInputOperandType());
+  auto outputShapedType = llvm::cast<ShapedType>(getOutputOperandType());
+  for (int64_t dim : llvm::seq<int64_t>(0, getOutputOperandRank())) {
----------------
MaheshRavishankar wrote:

Actually this might be made much simpler
```
Value input = getInput();
unsigned inputRank = getInputOperandType().getRank();
SmallVector<OpFoldResult> shapes = llvm::map_to_vector(llvm::seq<int64_t>(0, inputRank), [&](int64_t dim) {
  return createFoldedDimOp(rewriter, getLoc(), input, dim);
} )
return shapes.
```

https://github.com/llvm/llvm-project/pull/67790