add support for operations: reduceMean, add, mul, div, sigmoid

2026-03-30 15:41:12 +02:00
parent 5e7114f517
commit 39830be888
32 changed files with 1057 additions and 224 deletions
@@ -0,0 +1,36 @@
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+struct SigmoidToSpatialCompute : OpConversionPattern<ONNXSigmoidOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXSigmoidOp sigmoidOp,
+                                ONNXSigmoidOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    Location loc = sigmoidOp.getLoc();
+    Type resultType = sigmoidOp.getResult().getType();
+    constexpr size_t numInputs = 1;
+    auto computeOp = createSpatCompute<numInputs>(rewriter, loc, resultType, {}, adaptor.getX(), [&](Value x) {
+      auto spatSigmoidOp = spatial::SpatSigmoidOp::create(rewriter, loc, resultType, x);
+      spatial::SpatYieldOp::create(rewriter, loc, spatSigmoidOp.getResult());
+    });
+    rewriter.replaceOp(sigmoidOp, computeOp);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateSigmoidPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<SigmoidToSpatialCompute>(ctx);
+}
+
+} // namespace onnx_mlir