Refactor + ReduceMean batched

2026-05-29 15:57:13 +02:00
parent 832bd7f1f7
commit 819d8af0f7
27 changed files with 929 additions and 568 deletions
@@ -0,0 +1,45 @@
+#pragma once
+
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/FoldInterfaces.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <optional>
+
+namespace onnx_mlir {
+
+int64_t normalizeAxis(int64_t axis, int64_t rank);
+
+mlir::FailureOr<int64_t> normalizeAxisChecked(int64_t axis, int64_t rank);
+
+int64_t normalizeIndex(int64_t index, int64_t dimSize);
+
+llvm::SmallVector<int64_t> normalizeAxes(mlir::ArrayAttr axesAttr, int64_t rank);
+
+llvm::SmallVector<int64_t> normalizeAxes(std::optional<mlir::ArrayAttr> axesAttr, int64_t rank);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> normalizeAxesChecked(mlir::ArrayAttr axesAttr, int64_t rank);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> normalizeAxesChecked(std::optional<mlir::ArrayAttr> axesAttr, int64_t rank);
+
+mlir::Value createAffineApplyOrConstant(mlir::PatternRewriter& rewriter,
+                                        mlir::Location loc,
+                                        mlir::AffineExpr expr,
+                                        mlir::ValueRange operands);
+
+mlir::Value
+multiplyIndexByConstant(mlir::PatternRewriter& rewriter, mlir::Operation* anchorOp, mlir::Value value, int64_t multiplier);
+
+mlir::Value modIndexByConstant(mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value value, int64_t divisor);
+
+mlir::Value
+floorDivIndexByConstant(mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value value, int64_t divisor);
+
+mlir::Value getOrMaterializeIndexValue(mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::OpFoldResult value);
+
+} // namespace onnx_mlir