add support for softmax, resize, split, gather

2026-04-09 14:25:00 +02:00
parent 86916a8fa0
commit 1a0192d1f9
16 changed files with 560 additions and 8 deletions
@@ -0,0 +1,111 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static int64_t normalizeAxis(int64_t axis, int64_t rank) { return axis >= 0 ? axis : rank + axis; }
+
+static SmallVector<int64_t> permuteShape(ArrayRef<int64_t> shape, ArrayRef<int64_t> permutation) {
+  SmallVector<int64_t> permutedShape;
+  permutedShape.reserve(permutation.size());
+  for (int64_t axis : permutation)
+    permutedShape.push_back(shape[axis]);
+  return permutedShape;
+}
+
+static Value createSoftmaxCompute(Value input, ConversionPatternRewriter& rewriter, Location loc) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  constexpr size_t numInputs = 1;
+  auto computeOp = createSpatCompute<numInputs>(rewriter, loc, TypeRange {inputType}, {}, ValueRange {input}, [&](Value x) {
+    auto softmaxOp = spatial::SpatSoftmaxOp::create(rewriter, loc, inputType, x);
+    spatial::SpatYieldOp::create(rewriter, loc, softmaxOp.getResult());
+  });
+  return computeOp.getResult(0);
+}
+
+static Value buildSoftmax(Value input,
+                          int64_t softmaxAxis,
+                          int64_t axis,
+                          ConversionPatternRewriter& rewriter,
+                          Location loc) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  if (axis == inputType.getRank())
+    return createSoftmaxCompute(input, rewriter, loc);
+
+  if (axis == softmaxAxis)
+    return buildSoftmax(input, softmaxAxis, axis + 1, rewriter, loc);
+
+  SmallVector<Value> slices = sliceTensor(input, axis, /*sliceSize=*/1, rewriter, loc);
+  SmallVector<Value> rebuiltSlices;
+  rebuiltSlices.reserve(slices.size());
+  for (Value slice : slices)
+    rebuiltSlices.push_back(buildSoftmax(slice, softmaxAxis, axis + 1, rewriter, loc));
+
+  return rebuiltSlices.size() == 1 ? rebuiltSlices.front()
+                                   : tensor::ConcatOp::create(rewriter, loc, axis, rebuiltSlices).getResult();
+}
+
+struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXSoftmaxOp softmaxOp,
+                                ONNXSoftmaxOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto inputType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    if (!inputType || !inputType.hasStaticShape())
+      return failure();
+
+    int64_t axis = normalizeAxis(softmaxOp.getAxis(), inputType.getRank());
+    if (axis < 0 || axis >= inputType.getRank())
+      return failure();
+
+    Value input = adaptor.getInput();
+    Value result;
+    if (axis == inputType.getRank() - 1) {
+      result = buildSoftmax(input, axis, /*axis=*/0, rewriter, softmaxOp.getLoc());
+    } else {
+      SmallVector<int64_t> permutation;
+      permutation.reserve(inputType.getRank());
+      for (int64_t dim = 0; dim < inputType.getRank(); ++dim)
+        if (dim != axis)
+          permutation.push_back(dim);
+      permutation.push_back(axis);
+
+      SmallVector<int64_t> inversePermutation(inputType.getRank());
+      for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
+        inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
+
+      auto transposedType = RankedTensorType::get(
+        permuteShape(inputType.getShape(), permutation), inputType.getElementType(), inputType.getEncoding());
+      auto preTransposeCompute = createSpatCompute<1>(
+        rewriter, softmaxOp.getLoc(), TypeRange {transposedType}, {}, input, [&](Value x) {
+          Value transposed =
+            ONNXTransposeOp::create(rewriter, softmaxOp.getLoc(), transposedType, x, rewriter.getI64ArrayAttr(permutation));
+          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
+        });
+      Value transposedInput = preTransposeCompute.getResult(0);
+      Value transposedResult = buildSoftmax(transposedInput, /*softmaxAxis=*/inputType.getRank() - 1, /*axis=*/0, rewriter, softmaxOp.getLoc());
+      result = ONNXTransposeOp::create(
+        rewriter, softmaxOp.getLoc(), inputType, transposedResult, rewriter.getI64ArrayAttr(inversePermutation));
+    }
+
+    rewriter.replaceOp(softmaxOp, result);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateSoftmaxPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<SoftmaxToSpatialCompute>(ctx);
+}
+
+} // namespace onnx_mlir