40 lines
1.2 KiB
C++
40 lines
1.2 KiB
C++
#include "mlir/Dialect/Linalg/IR/Linalg.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Interfaces/SideEffectInterfaces.h"
|
|
|
|
#include "ShapingUtils.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
bool isShapingOnlyOp(Operation *op) {
|
|
return isa<tensor::CastOp,
|
|
tensor::CollapseShapeOp,
|
|
tensor::ExpandShapeOp,
|
|
tensor::ExtractSliceOp,
|
|
tensor::InsertSliceOp,
|
|
tensor::ConcatOp,
|
|
tensor::EmptyOp,
|
|
tensor::ExtractOp,
|
|
tensor::InsertOp,
|
|
tensor::SplatOp,
|
|
linalg::TransposeOp,
|
|
ONNXTransposeOp,
|
|
spatial::SpatConcatOp,
|
|
spatial::SpatExtractRowsOp>(op);
|
|
}
|
|
|
|
bool isPureIndexComputationOp(Operation *op) {
|
|
if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait<OpTrait::IsTerminator>()
|
|
|| !isMemoryEffectFree(op))
|
|
return false;
|
|
auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa<IntegerType>(type); };
|
|
return llvm::all_of(op->getOperandTypes(), isIndexOrInteger)
|
|
&& llvm::all_of(op->getResultTypes(), isIndexOrInteger);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|