fixes

slightly faster codegen

src/PIM/Compiler/PimCoreProgram.hpp

@@ -0,0 +1,74 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <memory>
+#include <optional>
+
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
+namespace onnx_mlir {
+
+struct CompiledTransposePlan {
+  mlir::Value source;
+  mlir::Value destination;
+  size_t elementBytes = 0;
+  size_t totalElements = 0;
+  size_t totalBytes = 0;
+  llvm::SmallVector<size_t> sourceShape;
+  llvm::SmallVector<size_t> destinationStrides;
+  llvm::SmallVector<unsigned> destinationDimensionForSource;
+  llvm::SmallVector<size_t> destinationRewinds;
+  bool storagePreserving = false;
+};
+
+enum class CompiledCoreOpKind : uint8_t {
+  Load,
+  Store,
+  Lmv,
+  Receive,
+  Send,
+  Concat,
+  Vmm,
+  Transpose,
+  VVAdd,
+  VVSub,
+  VVMul,
+  VVMax,
+  VVDMul,
+  VAvg,
+  VRelu,
+  VTanh,
+  VSigm,
+  VSoftmax
+};
+
+struct CompiledCoreNode {
+  enum class Kind : uint8_t { Op, Loop, If, IndexSwitch };
+
+  Kind kind = Kind::Op;
+  mlir::Operation* op = nullptr;
+  CompiledCoreOpKind opKind = CompiledCoreOpKind::Load;
+  CompiledIndexExpr lowerBound;
+  CompiledIndexExpr upperBound;
+  CompiledIndexExpr step;
+  CompiledIndexExpr condition;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
+  llvm::SmallVector<int64_t> caseValues;
+  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
+  std::optional<CompiledTransposePlan> transposePlan;
+};
+
+struct CompiledCoreProgram {
+  llvm::SmallVector<CompiledCoreNode, 32> nodes;
+};
+
+mlir::LogicalResult compileCoreProgram(mlir::Operation* coreLikeOp, CompiledCoreProgram& program);
+
+} // namespace onnx_mlir