slightly faster codegen
This commit is contained in:
NiccoloN
@@ -0,0 +1,209 @@
|
||||
#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
|
||||
Reference in New Issue
Block a user