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

[Abhishek Varma]

================
@@ -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())) {
----------------
Abhishek-Varma wrote:

So, I did try that but with the test case I added it gives :-
```
Assertion `shapedType.isDynamicDim(dim) == reifiedReturnShapes[resultIdx][dim].is<Value>() && "incorrect implementation of ReifyRankedShapedTypeOpInterface"' failed.
```

This is something I faced while adding the test case the first time - it pointed out to me that I need to segregate `input` and `output/inits` type.
And this follows the same implementation currently being done for [rest of the LinalgOps](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp#L951-L967).

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