TOSA: add folder for tosa.sub

mlir/lib/Dialect/Tosa/Transforms/TosaFolders.cpp

@@ -1005,6 +1005,38 @@ struct TosaFoldConstantAdd : public TosaFoldConstantBinary<TosaFoldConstantAdd,
   }
 };
 
+struct TosaFoldConstantSub : public TosaFoldConstantBinary<TosaFoldConstantSub, SubOp> {
+  using TosaFoldConstantBinary<TosaFoldConstantSub, SubOp>::TosaFoldConstantBinary;
+
+  DenseElementsAttr computeInteger(DenseElementsAttr lhsValues,
+                                   DenseElementsAttr rhsValues,
+                                   PatternRewriter &rewriter, SubOp op) const {
+    bool addOverflowed = false;
+    auto intAdd = [&addOverflowed](const APInt &first, const APInt &second) {
+      bool didOverflow;
+      auto res = first.ssub_ov(second, didOverflow);
+      addOverflowed |= didOverflow;
+      return res;
+    };
+    auto newTensor = applyElementWise<APInt, APInt>(lhsValues, rhsValues,
+                                                    op.getType(), intAdd);
+    if (addOverflowed) {
+      op->emitWarning("Addition did overflow. The results are unspecified.");
+    }
+    return newTensor;
+  }
+
+  DenseElementsAttr computeFloat(DenseElementsAttr lhsValues,
+                                 DenseElementsAttr rhsValues,
+                                 PatternRewriter &rewriter, SubOp op) const {
+    auto floatAdd = [](const APFloat &first, const APFloat &second) {
+      return first - second;
+    };
+    return applyElementWise<APFloat, APFloat>(lhsValues, rhsValues,
+                                              op.getType(), floatAdd);
+  }
+};
+
 struct TosaFoldConstantGreater : public TosaFoldConstantBinary<TosaFoldConstantGreater, GreaterOp> {
   using TosaFoldConstantBinary<TosaFoldConstantGreater, GreaterOp>::TosaFoldConstantBinary;
 
@@ -1382,6 +1414,7 @@ void mlir::tosa::populateTosaFoldConstantPatterns(
     patterns.add<TosaFoldConstantFloatCasts>(ctx, foldSplatOrSingleUseOnly);
   }
   patterns.add<TosaFoldConstantAdd>(ctx, foldSplatOrSingleUseOnly);
+  patterns.add<TosaFoldConstantSub>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantGreater>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantBitwiseNot>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantCeil>(ctx, foldSplatOrSingleUseOnly);
@@ -1394,4 +1427,4 @@ void mlir::tosa::populateTosaFoldConstantPatterns(
   patterns.add<TosaFoldConstantMinimum>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantMaximum>(ctx, foldSplatOrSingleUseOnly);
   patterns.add<TosaFoldConstantPad>(ctx, foldSplatOrSingleUseOnly);
-}
+}