Files
Raptor/src/PIM/PimAccelerator.cpp
NiccoloN 6933804003 constant fold linalg.map (generated from tensor.pad for padding)
refactor pim helpers in PimCommon
2026-03-20 20:51:20 +01:00

111 lines
4.5 KiB
C++

#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Func/Transforms/Passes.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "llvm/Support/Debug.h"
#include "src/Accelerators/PIM/Compiler/PimCompilerUtils.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
#include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/SpatialBufferizableOpInterface.hpp"
#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
#include "src/Accelerators/PIM/PimAccelerator.hpp"
#define DEBUG_TYPE "PimAccelerator"
namespace onnx_mlir {
namespace accel {
Accelerator* createPIM() { return PimAccelerator::getInstance(); }
PimAccelerator* PimAccelerator::instance = nullptr;
PimAccelerator* PimAccelerator::getInstance() {
if (instance == nullptr)
instance = new PimAccelerator();
return instance;
}
PimAccelerator::PimAccelerator()
: Accelerator(Kind::PIM) {
LLVM_DEBUG(llvm::dbgs() << "Creating a PIM accelerator\n");
acceleratorTargets.push_back(this);
}
uint64_t PimAccelerator::getVersionNumber() const { return 0x000001; }
void PimAccelerator::addPasses(mlir::OwningOpRef<mlir::ModuleOp>& module,
mlir::PassManager& pm,
EmissionTargetType& emissionTarget,
std::string outputNameNoExt) const {
LLVM_DEBUG(llvm::dbgs() << "Adding passes for PIM accelerator\n");
addPassesPim(module, pm, emissionTarget, outputNameNoExt);
}
void PimAccelerator::registerDialects(mlir::DialectRegistry& registry) const {
LLVM_DEBUG(llvm::dbgs() << "Registering dialects for PIM accelerator\n");
registry.insert<mlir::tensor::TensorDialect>();
registry.insert<mlir::tosa::TosaDialect>();
registry.insert<mlir::bufferization::BufferizationDialect>();
registry.insert<pim::PimDialect>();
registry.insert<spatial::SpatialDialect>();
mlir::tensor::registerBufferizableOpInterfaceExternalModels(registry);
mlir::tensor::registerInferTypeOpInterfaceExternalModels(registry);
mlir::arith::registerBufferizableOpInterfaceExternalModels(registry);
mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(registry);
spatial::registerBufferizableOpInterfaceExternalModels(registry);
spatial::registerONNXBufferizableOpInterfaceExternalModels(registry);
pim::registerOpBufferizationInterfaces(registry);
}
void PimAccelerator::registerPasses(int optLevel) const {
LLVM_DEBUG(llvm::dbgs() << "Registering passes for PIM accelerator\n");
registerPass(createONNXToSpatialPass);
registerPass(createSpatialToGraphvizPass);
registerPass(createSpatialToPimPass);
registerPass(createBufferizePimPass);
registerPass(createPimConstantFoldingPass);
registerPass(createPimHostVerificationPass);
registerPass(createEmitPimJsonPass);
}
void PimAccelerator::configurePasses() const {
LLVM_DEBUG(llvm::dbgs() << "Configuring passes for PIM accelerator\n");
// TODO: This does nothing for now.
}
mlir::MemRefType PimAccelerator::convertTensorTypeToMemRefType(const mlir::TensorType tensorType) const {
// Do not convert tensor types to memref types.
return nullptr;
}
void PimAccelerator::conversionTargetONNXToKrnl(mlir::ConversionTarget& target) const {
target.addLegalDialect<pim::PimDialect>();
}
void PimAccelerator::rewritePatternONNXToKrnl(mlir::RewritePatternSet& patterns,
mlir::TypeConverter& typeConverter,
mlir::MLIRContext* ctx) const {
// TODO: Add patterns for conversion
}
void PimAccelerator::conversionTargetKrnlToLLVM(mlir::ConversionTarget& target) const {}
void PimAccelerator::rewritePatternKrnlToLLVM(mlir::RewritePatternSet& patterns,
mlir::LLVMTypeConverter& typeConverter,
mlir::MLIRContext* ctx) const {
// We should not need this, since we offload it all to PIM.
}
} // namespace accel
} // namespace onnx_mlir