add PIM accelerator

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -0,0 +1,131 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <filesystem>
+#include <fstream>
+
+#include "Conversion/ONNXToSpatial/Utils/AnnotateReplication.hpp"
+#include "ONNXToSpatialPass.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/PIM/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
+#include "src/Compiler/CompilerOptions.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace spatial {
+
+void ONNXToSpatialPass::runOnOperation() {
+  llvm::dbgs() << "Running ONNXToSpatialLoweringPass\n";
+
+  ModuleOp module = getOperation();
+  MLIRContext* ctx = &getContext();
+
+  RewritePatternSet mergeActivationPatterns(ctx);
+  mergeActivationPatterns.add<onnxToArithConstantOp>(ctx);
+  mergeActivationPatterns.add<convAddToConvWithBiasPatternLeft>(ctx);
+  mergeActivationPatterns.add<convAddToConvWithBiasPatternRight>(ctx);
+  mergeActivationPatterns.add<matMulAddToGemmPattern>(ctx);
+  mergeActivationPatterns.add<matMulToGemmPattern>(ctx);
+  mergeActivationPatterns.add<removeFlattenSameShapePattern>(ctx);
+
+  if (failed(applyPatternsAndFoldGreedily(module, std::move(mergeActivationPatterns))))
+    llvm::dbgs() << "Failed to merge activation patterns, continuing...\n";
+
+  IRRewriter rewriter(module);
+  func::FuncOp funcOp = *module.getOps<func::FuncOp>().begin();
+  if (annotateReplication(funcOp, rewriter).failed()) {
+    llvm::dbgs() << "Failed during annotation for replication analysis\n";
+    signalPassFailure();
+    return;
+  }
+
+  ConversionTarget target(*ctx);
+  target.addLegalDialect<ONNXDialect, SpatialDialect, tensor::TensorDialect, arith::ArithDialect, tosa::TosaDialect>();
+  target.addIllegalOp<ONNXMatMulOp>();
+  target.addIllegalOp<ONNXGemmOp>();
+  target.addIllegalOp<ONNXConvOp>();
+  target.addIllegalOp<ONNXLRNOp>();
+  target.addIllegalOp<ONNXMaxPoolSingleOutOp>();
+  target.addIllegalOp<ONNXAveragePoolOp>();
+  target.addIllegalOp<ONNXConcatOp>();
+  target.addIllegalOp<ONNXSoftmaxOp>();
+  target.addIllegalOp<ONNXReduceMeanV13Op>();
+
+  RewritePatternSet patterns(ctx);
+  patterns.add<removeLRNPattern>(ctx);
+
+  if (useExperimentalConvImpl) {
+    populateExperimentalTilingConvOpPattern(patterns, ctx);
+    populateExperimentalPoolingTilingPattern(patterns, ctx);
+    populateGemmToConvConversionPattern(patterns, ctx);
+  }
+  else {
+    populateTilingConvOpPattern(patterns, ctx);
+    populatePoolingTilingPattern(patterns, ctx);
+    populateTilingGemmOpPattern(patterns, ctx);
+  }
+
+  populateONNXConcatToTensorConcatPattern(patterns, ctx);
+  populateReduceMeanConversionPattern(patterns, ctx);
+
+  if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
+    signalPassFailure();
+    return;
+  }
+
+  // Count the number of compute ops and check they do not exceed the core count
+  if (coresCount != -1) {
+    int computeOpsCount = 0;
+    for (auto& op : funcOp.getFunctionBody().front().getOperations())
+      if (isa<spatial::SpatWeightedCompute>(op))
+        computeOpsCount++;
+
+    if (computeOpsCount > coresCount) {
+      llvm::dbgs() << "Number of compute ops exceeds the core count\n";
+      signalPassFailure();
+      return;
+    }
+  }
+
+  // Remove trailing "helper ops" i.e. concat,img_concat,reshape.
+  RewritePatternSet removeUnusedHelperOpsPatterns(ctx);
+  populateRemoveUnusedHelperOpsPatterns(removeUnusedHelperOpsPatterns, ctx);
+
+  if (failed(applyPatternsAndFoldGreedily(module, std::move(removeUnusedHelperOpsPatterns))))
+    llvm::dbgs() << "Failed to remove unused helper ops, continuing...\n";
+
+  annotateWeightsConstants(funcOp);
+
+  // Dump to file for debug
+  std::string outputDir = outputBaseName.substr(0, outputBaseName.find_last_of('/')).append("/dialects");
+  std::filesystem::create_directory(outputDir);
+  std::fstream file(outputDir + "/spatial.mlir", std::ios::out);
+  llvm::raw_os_ostream os(file);
+  os << *module;
+  os.flush();
+  file.close();
+}
+
+void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
+  MLIRContext* ctx = funcOp.getContext();
+  funcOp.walk([&](arith::ConstantOp constantOp) {
+    bool isAlwaysWeight =
+      llvm::all_of(constantOp->getUsers(), [](auto user) -> bool { return isa<SpatWeightedCompute>(user); });
+    if (isAlwaysWeight)
+      constantOp->setAttr("weightAlways", UnitAttr::get(ctx));
+  });
+}
+
+} // namespace spatial
+
+} // namespace onnx_mlir