NiccoloN
dbb66be93e
Validate Operations / validate-operations (push) Has been cancelled
slightly faster codegen
210 lines
9.6 KiB
C++
210 lines
9.6 KiB
C++
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/SCF/IR/SCF.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
|
|
#include "llvm/ADT/STLExtras.h"
|
|
|
|
#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
|
|
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
|
|
#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
|
|
#include "src/Accelerators/PIM/Compiler/PimCoreProgram.hpp"
|
|
#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
|
|
|
|
using namespace llvm;
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
namespace {
|
|
|
|
static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
|
|
if (isa<pim::PimMemCopyHostToDevOp>(op)) return CompiledCoreOpKind::Load;
|
|
if (isa<pim::PimMemCopyDevToHostOp>(op)) return CompiledCoreOpKind::Store;
|
|
if (isa<pim::PimMemCopyOp>(op)) return CompiledCoreOpKind::Lmv;
|
|
if (isa<pim::PimReceiveOp>(op)) return CompiledCoreOpKind::Receive;
|
|
if (isa<pim::PimSendOp>(op)) return CompiledCoreOpKind::Send;
|
|
if (isa<pim::PimConcatOp>(op)) return CompiledCoreOpKind::Concat;
|
|
if (isa<pim::PimVMMOp>(op)) return CompiledCoreOpKind::Vmm;
|
|
if (isa<pim::PimTransposeOp>(op)) return CompiledCoreOpKind::Transpose;
|
|
if (isa<pim::PimVVAddOp>(op)) return CompiledCoreOpKind::VVAdd;
|
|
if (isa<pim::PimVVSubOp>(op)) return CompiledCoreOpKind::VVSub;
|
|
if (isa<pim::PimVVMulOp>(op)) return CompiledCoreOpKind::VVMul;
|
|
if (isa<pim::PimVVMaxOp>(op)) return CompiledCoreOpKind::VVMax;
|
|
if (isa<pim::PimVVDMulOp>(op)) return CompiledCoreOpKind::VVDMul;
|
|
if (isa<pim::PimVAvgOp>(op)) return CompiledCoreOpKind::VAvg;
|
|
if (isa<pim::PimVReluOp>(op)) return CompiledCoreOpKind::VRelu;
|
|
if (isa<pim::PimVTanhOp>(op)) return CompiledCoreOpKind::VTanh;
|
|
if (isa<pim::PimVSigmOp>(op)) return CompiledCoreOpKind::VSigm;
|
|
if (isa<pim::PimVSoftmaxOp>(op)) return CompiledCoreOpKind::VSoftmax;
|
|
return failure();
|
|
}
|
|
|
|
static bool isStoragePreservingTranspose(ArrayRef<size_t> sourceShape, ArrayRef<int64_t> permutation) {
|
|
SmallVector<unsigned> sourceNonUnitDims;
|
|
SmallVector<unsigned> destinationSourceNonUnitDims;
|
|
for (auto [dim, size] : llvm::enumerate(sourceShape))
|
|
if (size != 1)
|
|
sourceNonUnitDims.push_back(dim);
|
|
for (int64_t sourceDim : permutation)
|
|
if (sourceShape[sourceDim] != 1)
|
|
destinationSourceNonUnitDims.push_back(static_cast<unsigned>(sourceDim));
|
|
return sourceNonUnitDims == destinationSourceNonUnitDims;
|
|
}
|
|
|
|
static FailureOr<CompiledTransposePlan> compileTransposePlan(pim::PimTransposeOp transposeOp) {
|
|
auto sourceType = cast<ShapedType>(transposeOp.getInput().getType());
|
|
ArrayRef<int64_t> sourceShape = sourceType.getShape();
|
|
size_t rank = sourceShape.size();
|
|
CompiledTransposePlan plan;
|
|
plan.source = transposeOp.getInput();
|
|
plan.destination = transposeOp.getOutputBuffer();
|
|
plan.elementBytes = getElementTypeSizeInBytes(sourceType.getElementType());
|
|
auto totalElements = pim::checkedSize(sourceType.getNumElements(), transposeOp, "transpose elements");
|
|
if (failed(totalElements)) return failure();
|
|
plan.totalElements = *totalElements;
|
|
auto totalBytes = pim::checkedMul(plan.totalElements, plan.elementBytes, transposeOp, "transpose byte size");
|
|
if (failed(totalBytes)) return failure();
|
|
plan.totalBytes = *totalBytes;
|
|
|
|
SmallVector<int64_t> permutation = map_to_vector(transposeOp.getPermutation().getAsRange<IntegerAttr>(),
|
|
[](IntegerAttr attr) { return attr.getInt(); });
|
|
if (permutation.size() != rank) {
|
|
transposeOp.emitOpError("requires permutation rank to match source rank for PIM codegen");
|
|
return failure();
|
|
}
|
|
|
|
SmallVector<size_t> destinationShape(rank);
|
|
plan.destinationStrides.assign(rank, 1);
|
|
plan.destinationDimensionForSource.assign(rank, 0);
|
|
plan.destinationRewinds.assign(rank, 0);
|
|
SmallVector<bool> seenSourceDimensions(rank, false);
|
|
for (size_t dim = 0; dim < rank; ++dim) {
|
|
auto size = pim::checkedSize(sourceShape[dim], transposeOp, "transpose source dimension");
|
|
if (failed(size)) return failure();
|
|
plan.sourceShape.push_back(*size);
|
|
}
|
|
for (auto [destinationDim, sourceDim] : llvm::enumerate(permutation)) {
|
|
if (sourceDim < 0 || static_cast<size_t>(sourceDim) >= rank || seenSourceDimensions[sourceDim]) {
|
|
transposeOp.emitOpError("requires a valid permutation containing each source dimension exactly once");
|
|
return failure();
|
|
}
|
|
seenSourceDimensions[sourceDim] = true;
|
|
destinationShape[destinationDim] = plan.sourceShape[sourceDim];
|
|
plan.destinationDimensionForSource[sourceDim] = destinationDim;
|
|
}
|
|
for (size_t dim = rank; dim > 1; --dim) {
|
|
auto stride = pim::checkedMul(
|
|
plan.destinationStrides[dim - 1], destinationShape[dim - 1], transposeOp, "transpose destination stride");
|
|
if (failed(stride)) return failure();
|
|
plan.destinationStrides[dim - 2] = *stride;
|
|
}
|
|
for (size_t sourceDim = 0; sourceDim < rank; ++sourceDim) {
|
|
auto rewind = pim::checkedMul(plan.sourceShape[sourceDim],
|
|
plan.destinationStrides[plan.destinationDimensionForSource[sourceDim]],
|
|
transposeOp,
|
|
"transpose destination rewind");
|
|
if (failed(rewind)) return failure();
|
|
plan.destinationRewinds[sourceDim] = *rewind;
|
|
}
|
|
plan.storagePreserving = isStoragePreservingTranspose(plan.sourceShape, permutation);
|
|
return plan;
|
|
}
|
|
|
|
static LogicalResult compileCoreEmissionPlan(Block& block, SmallVectorImpl<CompiledCoreNode>& plan) {
|
|
for (Operation& op : block) {
|
|
if (isa<pim::PimHaltOp, scf::YieldOp, memref::GetGlobalOp>(op) || isCoreStaticAddressOp(&op))
|
|
continue;
|
|
if (auto loadOp = dyn_cast<memref::LoadOp>(op); loadOp && succeeded(compileIndexExpr(loadOp.getResult())))
|
|
continue;
|
|
|
|
if (auto forOp = dyn_cast<scf::ForOp>(op)) {
|
|
auto lower = compileIndexExpr(forOp.getLowerBound());
|
|
auto upper = compileIndexExpr(forOp.getUpperBound());
|
|
auto step = compileIndexExpr(forOp.getStep());
|
|
if (failed(lower) || failed(upper) || failed(step)) {
|
|
forOp.emitOpError("requires statically evaluable scf.for bounds for PIM codegen");
|
|
return failure();
|
|
}
|
|
CompiledCoreNode node;
|
|
node.kind = CompiledCoreNode::Kind::Loop;
|
|
node.op = forOp;
|
|
node.lowerBound = *lower;
|
|
node.upperBound = *upper;
|
|
node.step = *step;
|
|
node.loopBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
|
|
if (failed(compileCoreEmissionPlan(forOp.getRegion().front(), *node.loopBody))) return failure();
|
|
plan.push_back(std::move(node));
|
|
continue;
|
|
}
|
|
if (auto ifOp = dyn_cast<scf::IfOp>(op)) {
|
|
auto condition = compileIndexExpr(ifOp.getCondition());
|
|
if (failed(condition)) {
|
|
ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
|
|
return failure();
|
|
}
|
|
CompiledCoreNode node;
|
|
node.kind = CompiledCoreNode::Kind::If;
|
|
node.op = ifOp;
|
|
node.condition = *condition;
|
|
node.thenBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
|
|
node.elseBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
|
|
if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), *node.thenBody))) return failure();
|
|
if (!ifOp.getElseRegion().empty()
|
|
&& failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), *node.elseBody)))
|
|
return failure();
|
|
plan.push_back(std::move(node));
|
|
continue;
|
|
}
|
|
if (auto switchOp = dyn_cast<scf::IndexSwitchOp>(op)) {
|
|
auto selector = compileIndexExpr(switchOp.getArg());
|
|
if (failed(selector)) {
|
|
switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
|
|
return failure();
|
|
}
|
|
CompiledCoreNode node;
|
|
node.kind = CompiledCoreNode::Kind::IndexSwitch;
|
|
node.op = switchOp;
|
|
node.condition = *selector;
|
|
llvm::append_range(node.caseValues, switchOp.getCases());
|
|
for (Region& region : switchOp.getCaseRegions()) {
|
|
auto body = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
|
|
if (failed(compileCoreEmissionPlan(region.front(), *body))) return failure();
|
|
node.caseBodies.push_back(std::move(body));
|
|
}
|
|
node.defaultBody = std::make_unique<SmallVector<CompiledCoreNode, 8>>();
|
|
if (failed(compileCoreEmissionPlan(switchOp.getDefaultRegion().front(), *node.defaultBody))) return failure();
|
|
plan.push_back(std::move(node));
|
|
continue;
|
|
}
|
|
|
|
auto opKind = classifyCompiledCoreOpKind(op);
|
|
if (failed(opKind)) {
|
|
InFlightDiagnostic diagnostic = op.emitError() << "unsupported codegen for op '" << op.getName() << "'";
|
|
if (auto coreOp = op.getParentOfType<pim::PimCoreOp>())
|
|
diagnostic << " inside pim.core " << coreOp.getCoreId();
|
|
else if (auto batchOp = op.getParentOfType<pim::PimCoreBatchOp>())
|
|
diagnostic << " inside pim.core_batch with laneCount " << batchOp.getLaneCount();
|
|
return failure();
|
|
}
|
|
CompiledCoreNode node;
|
|
node.op = &op;
|
|
node.opKind = *opKind;
|
|
if (*opKind == CompiledCoreOpKind::Transpose) {
|
|
auto transposePlan = compileTransposePlan(cast<pim::PimTransposeOp>(op));
|
|
if (failed(transposePlan)) return failure();
|
|
node.transposePlan = *transposePlan;
|
|
}
|
|
plan.push_back(std::move(node));
|
|
}
|
|
return success();
|
|
}
|
|
|
|
} // namespace
|
|
|
|
LogicalResult compileCoreProgram(Operation* coreLikeOp, CompiledCoreProgram& program) {
|
|
Block& block = isa<pim::PimCoreOp>(coreLikeOp) ? cast<pim::PimCoreOp>(coreLikeOp).getBody().front()
|
|
: cast<pim::PimCoreBatchOp>(coreLikeOp).getBody().front();
|
|
return compileCoreEmissionPlan(block, program.nodes);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|