Refactor + ReduceMean batched

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Value.h"
+#include "mlir/IR/ValueRange.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "llvm/ADT/ArrayRef.h"
@@ -11,6 +12,7 @@
 
 #include <cassert>
 #include <cstddef>
+#include <optional>
 #include <type_traits>
 #include <utility>
 
@@ -109,6 +111,33 @@ inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
       && lhsType.getShape() == rhsType.getShape();
 }
 
+bool hasStaticPositiveShape(mlir::ArrayRef<int64_t> shape);
+
+bool hasStaticPositiveShape(mlir::RankedTensorType type);
+
+int64_t getStaticShapeElementCount(mlir::ArrayRef<int64_t> shape);
+
+int64_t getStaticShapeElementCount(mlir::RankedTensorType type);
+
+llvm::SmallVector<int64_t> permuteShape(mlir::ArrayRef<int64_t> shape, mlir::ArrayRef<int64_t> permutation);
+
+llvm::SmallVector<int64_t> invertPermutation(mlir::ArrayRef<int64_t> permutation);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr,
+                                                                           int64_t rank);
+
+mlir::Value transposeMaybeInCompute(mlir::Value value,
+                                    mlir::RankedTensorType resultType,
+                                    mlir::ArrayRef<int64_t> permutation,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, mlir::ArrayRef<int64_t> shape);
+
 /// Slices a statically shaped tensor along one axis into contiguous pieces of
 /// at most `sliceSize` elements.
 llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
@@ -148,4 +177,23 @@ mlir::Value materializeContiguousTensorSlice(mlir::Value source,
                                              mlir::ConversionPatternRewriter& rewriter,
                                              mlir::Location loc);
 
+mlir::Value extractStaticSlice(mlir::PatternRewriter& rewriter,
+                               mlir::Location loc,
+                               mlir::Value source,
+                               mlir::RankedTensorType resultType,
+                               llvm::ArrayRef<mlir::OpFoldResult> offsets);
+
+mlir::Value extractAxisSlice(mlir::PatternRewriter& rewriter,
+                             mlir::Location loc,
+                             mlir::Value source,
+                             int64_t axis,
+                             int64_t offset,
+                             int64_t size);
+
+mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
+                              mlir::Location loc,
+                              mlir::Value source,
+                              mlir::Value dest,
+                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
+
 } // namespace onnx_mlir