[Mlir-commits] [mlir] [mlir][tosa] Add folder for multiply like reduce_prod operation (PR #128067)

[Jack Frankland]

================
@@ -963,10 +963,38 @@ REDUCE_FOLDER(ReduceAllOp)
 REDUCE_FOLDER(ReduceAnyOp)
 REDUCE_FOLDER(ReduceMaxOp)
 REDUCE_FOLDER(ReduceMinOp)
-REDUCE_FOLDER(ReduceProdOp)
 REDUCE_FOLDER(ReduceSumOp)
 #undef REDUCE_FOLDER
 
+OpFoldResult ReduceProdOp::fold(FoldAdaptor adaptor) {
+  ShapedType inputTy = llvm::cast<ShapedType>(getInput().getType());
+  if (!inputTy.hasRank())
+    return {};
+  if (inputTy == getType() &&
+      (inputTy.getRank() == 0 || inputTy.getDimSize(getAxis()) == 1))
+    return getInput();
+
+  // Fold multiply like reduce_prod operators using mulBinaryFolder
+  if (inputTy.getRank() == 1 && inputTy.getDimSize(0) == 2) {
+    const auto resultTy = llvm::dyn_cast<RankedTensorType>(getType());
----------------
FranklandJack wrote:

I don't think `Type` and its derived classes are const correct.

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