NiccoloN
dbb66be93e
Validate Operations / validate-operations (push) Has been cancelled
slightly faster codegen
75 lines
1.8 KiB
C++
75 lines
1.8 KiB
C++
#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
|