[Mlir-commits] [mlir] [mlir][tosa] Add constant folding for tosa.add_shape operation (PR #173112)

[Luke Hutton]

================
@@ -890,32 +890,47 @@ void SliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
 //===----------------------------------------------------------------------===//
 
 template <typename Folder>
-static DenseElementsAttr binaryFolder(DenseElementsAttr lhs,
-                                      DenseElementsAttr rhs,
-                                      RankedTensorType returnTy) {
-  if (rhs && lhs && rhs.isSplat() && lhs.isSplat()) {
-    const auto lETy = llvm::cast<ShapedType>(lhs.getType()).getElementType();
-    const auto rETy = llvm::cast<ShapedType>(rhs.getType()).getElementType();
-    if (lETy != rETy)
-      return {};
+static DenseElementsAttr
+binaryFolder(DenseElementsAttr lhs, DenseElementsAttr rhs, ShapedType returnTy,
+             bool foldDenseValues = false) {
+  if (!lhs || !rhs)
+    return {};
 
-    if (const auto lIntTy = dyn_cast<IntegerType>(lETy)) {
+  const auto lETy = llvm::cast<ShapedType>(lhs.getType()).getElementType();
+  const auto rETy = llvm::cast<ShapedType>(rhs.getType()).getElementType();
+  if (lETy != rETy)
+    return {};
+
+  if (lhs.isSplat() && rhs.isSplat()) {
+    if (isa<FloatType>(lETy)) {
+      const APFloat l = lhs.getSplatValue<APFloat>();
+      const APFloat r = rhs.getSplatValue<APFloat>();
+      const auto maybeResult = Folder::fold(l, r);
+      if (failed(maybeResult))
+        return {};
+      return DenseElementsAttr::get(returnTy, maybeResult.value());
+    }
+
+    if (const auto lIntTy = llvm::dyn_cast<IntegerType>(lETy)) {
       const APInt l = lhs.getSplatValue<APInt>();
       const APInt r = rhs.getSplatValue<APInt>();
       const auto maybeResult = Folder::fold(l, r, lIntTy.isUnsigned());
       if (failed(maybeResult))
         return {};
       return DenseElementsAttr::get(returnTy, maybeResult.value());
     }
+  }
 
-    if (llvm::isa<FloatType>(lETy)) {
-      const APFloat l = lhs.getSplatValue<APFloat>();
-      const APFloat r = rhs.getSplatValue<APFloat>();
-      const auto maybeResult = Folder::fold(l, r);
+  if (foldDenseValues) {
+    SmallVector<APInt> resultValues;
----------------
lhutton1 wrote:

Thanks, yes it should. I've added an assert

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