Export csv graph for gephi

2026-07-02 17:01:26 +02:00
parent 8d3eb929f6
commit c4dd28a607
14 changed files with 926 additions and 10 deletions
@@ -0,0 +1,112 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/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 FailureOr<int64_t> normalizeFlattenAxis(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = axis < 0 ? rank + axis : axis;
+  if (normalizedAxis < 0 || normalizedAxis > rank)
+    return failure();
+  return normalizedAxis;
+}
+
+static int64_t product(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values)
+    result *= value;
+  return result;
+}
+
+static SmallVector<ReassociationIndices> getCollapseTo1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          ConversionPatternRewriter& rewriter,
+                          Location loc) {
+  if (sourceType == resultType)
+    return input;
+
+  if (axis > 0 && axis < sourceType.getRank()) {
+    SmallVector<ReassociationIndices> reassociation(2);
+    for (int64_t dim = 0; dim < axis; ++dim)
+      reassociation[0].push_back(dim);
+    for (int64_t dim = axis; dim < sourceType.getRank(); ++dim)
+      reassociation[1].push_back(dim);
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
+  }
+
+  Value flattened = input;
+  if (sourceType.getRank() != 1) {
+    auto flatType = RankedTensorType::get({sourceType.getNumElements()}, sourceType.getElementType());
+    flattened = tensor::CollapseShapeOp::create(
+      rewriter, loc, flatType, flattened, getCollapseTo1DReassociation(sourceType.getRank()));
+  }
+  return tensor::ExpandShapeOp::create(
+    rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
+}
+
+struct Flatten : OpConversionPattern<ONNXFlattenOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXFlattenOp flattenOp,
+                                ONNXFlattenOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto sourceType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOperation()->getResult(0).getType());
+    if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+    if (!hasStaticPositiveShape(sourceType) || !hasStaticPositiveShape(resultType) || resultType.getRank() != 2)
+      return failure();
+
+    auto axis = normalizeFlattenAxis(flattenOp.getAxis(), sourceType.getRank());
+    if (failed(axis))
+      return failure();
+
+    int64_t outerDim = product(sourceType.getShape().take_front(*axis));
+    int64_t innerDim = product(sourceType.getShape().drop_front(*axis));
+    if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
+      return failure();
+
+    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
+      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
+      rewriter.replaceOp(flattenOp, flattened);
+      return success();
+    };
+
+    return replaceWithFlatten([&](Value input) {
+      return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc());
+    });
+  }
+};
+
+} // namespace
+
+void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
+
+} // namespace onnx_mlir