NiccoloN
53 lines
2.2 KiB
C++
53 lines
2.2 KiB
C++
#pragma once
|
|
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/IR/ValueRange.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
|
|
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
|
|
|
|
namespace onnx_mlir {
|
|
|
|
struct MixedSliceGeometry {
|
|
llvm::SmallVector<mlir::OpFoldResult> offsets;
|
|
llvm::SmallVector<mlir::OpFoldResult> sizes;
|
|
llvm::SmallVector<mlir::OpFoldResult> strides;
|
|
};
|
|
|
|
mlir::Value extractAxisSlice(
|
|
mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
|
|
|
|
mlir::Value extractStaticSliceOrIdentity(mlir::OpBuilder& rewriter,
|
|
mlir::Location loc,
|
|
mlir::Value source,
|
|
mlir::RankedTensorType resultType,
|
|
llvm::ArrayRef<mlir::OpFoldResult> offsets,
|
|
llvm::ArrayRef<mlir::OpFoldResult> sizes,
|
|
llvm::ArrayRef<mlir::OpFoldResult> strides);
|
|
|
|
mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
|
|
mlir::Location loc,
|
|
mlir::Value source,
|
|
mlir::Value dest,
|
|
llvm::ArrayRef<mlir::OpFoldResult> offsets);
|
|
|
|
mlir::Value extractMixedSliceOrIdentity(mlir::OpBuilder &rewriter,
|
|
mlir::Location loc,
|
|
mlir::Value source,
|
|
mlir::RankedTensorType resultType,
|
|
const MixedSliceGeometry &geometry);
|
|
|
|
mlir::Value insertMixedSlice(mlir::OpBuilder &builder,
|
|
mlir::Location loc,
|
|
mlir::Value source,
|
|
mlir::Value dest,
|
|
const MixedSliceGeometry &geometry);
|
|
|
|
mlir::FailureOr<mlir::Value>
|
|
addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);
|
|
|
|
mlir::FailureOr<mlir::Value> removeLeadingUnitTensorDimension(
|
|
mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType);
|
|
|
|
} // namespace onnx_mlir
|