NiccoloN
263 lines
10 KiB
C++
263 lines
10 KiB
C++
#include "llvm/ADT/DenseSet.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/SetVector.h"
|
|
|
|
#include <string>
|
|
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
namespace spatial {
|
|
namespace {
|
|
|
|
std::optional<BlockArgument> getBlockArgument(Region& body, unsigned argIdx) {
|
|
if (body.empty())
|
|
return std::nullopt;
|
|
|
|
Block& block = body.front();
|
|
if (argIdx >= block.getNumArguments())
|
|
return std::nullopt;
|
|
return block.getArgument(argIdx);
|
|
}
|
|
|
|
std::optional<BlockArgument> insertBlockArgument(Region& body, unsigned argIdx, Type type, Location loc) {
|
|
if (body.empty())
|
|
return std::nullopt;
|
|
return body.insertArgument(argIdx, type, loc);
|
|
}
|
|
|
|
void setComputeOperandSegmentSizes(Operation* op, int32_t weightCount, int32_t inputCount) {
|
|
if (auto compute = dyn_cast<SpatCompute>(op)) {
|
|
compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
|
|
return;
|
|
}
|
|
cast<SpatComputeBatch>(op).getProperties().setOperandSegmentSizes({weightCount, inputCount});
|
|
}
|
|
|
|
using CrossbarWeightSet = llvm::SetVector<Value, llvm::SmallVector<Value, 4>, llvm::SmallDenseSet<Value, 4>>;
|
|
|
|
CrossbarWeightSet collectCrossbarWeights(Region& body) {
|
|
CrossbarWeightSet weights;
|
|
body.walk([&](SpatVMMOp vmmOp) {
|
|
Value weight = vmmOp.getWeight();
|
|
weights.insert(weight);
|
|
});
|
|
return weights;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
std::optional<BlockArgument> SpatCompute::getWeightArgument(unsigned idx) { return getBlockArgument(getBody(), idx); }
|
|
|
|
std::optional<BlockArgument> SpatCompute::getInputArgument(unsigned idx) {
|
|
return getBlockArgument(getBody(), getWeights().size() + idx);
|
|
}
|
|
|
|
std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertWeight(unsigned idx, Value weight, Location loc) {
|
|
if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
|
|
auto index = std::distance(getWeights().begin(), existing);
|
|
return {
|
|
{*existing, *getWeightArgument(index)}
|
|
};
|
|
}
|
|
|
|
unsigned weightCount = getWeights().size();
|
|
unsigned inputCount = getInputs().size();
|
|
getOperation()->insertOperands(idx, ValueRange {weight});
|
|
setComputeOperandSegmentSizes(
|
|
getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
|
|
auto blockArg = insertBlockArgument(getBody(), idx, weight.getType(), loc);
|
|
if (!blockArg)
|
|
return std::nullopt;
|
|
return std::make_tuple(getOperation()->getOperand(idx), *blockArg);
|
|
}
|
|
|
|
std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertInput(unsigned idx, Value input, Location loc) {
|
|
unsigned weightCount = getWeights().size();
|
|
unsigned inputCount = getInputs().size();
|
|
getOperation()->insertOperands(weightCount + idx, ValueRange {input});
|
|
setComputeOperandSegmentSizes(
|
|
getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
|
|
auto blockArg = insertBlockArgument(getBody(), weightCount + idx, input.getType(), loc);
|
|
if (!blockArg)
|
|
return std::nullopt;
|
|
return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
|
|
}
|
|
|
|
CrossbarWeightSet SpatCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
|
|
|
|
FailureOr<std::tuple<OpResult, SpatCompute>>
|
|
SpatCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
|
|
if (idx > getNumResults())
|
|
return failure();
|
|
|
|
rewriter.setInsertionPoint(getOperation());
|
|
SmallVector<Type> resultTypes(getResultTypes().begin(), getResultTypes().end());
|
|
resultTypes.insert(resultTypes.begin() + idx, type);
|
|
auto newCompute = SpatCompute::create(rewriter, getLoc(), TypeRange(resultTypes), getWeights(), getInputs());
|
|
newCompute->setAttrs((*this)->getAttrs());
|
|
setComputeOperandSegmentSizes(newCompute.getOperation(),
|
|
static_cast<int32_t>(newCompute.getWeights().size()),
|
|
static_cast<int32_t>(newCompute.getInputs().size()));
|
|
rewriter.inlineRegionBefore(getBody(), newCompute.getBody(), newCompute.getBody().end());
|
|
for (unsigned oldResultIdx = 0; oldResultIdx < getNumResults(); ++oldResultIdx)
|
|
getResult(oldResultIdx)
|
|
.replaceAllUsesWith(newCompute.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
|
|
rewriter.eraseOp(getOperation());
|
|
return std::make_tuple(cast<OpResult>(newCompute.getResult(idx)), newCompute);
|
|
}
|
|
|
|
void SpatCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
|
|
if (region.empty())
|
|
return;
|
|
|
|
for (unsigned index = 0; index < getWeights().size(); ++index)
|
|
if (auto weightArg = getWeightArgument(index))
|
|
setNameFn(*weightArg, ("w" + std::to_string(index)).c_str());
|
|
|
|
for (unsigned index = 0; index < getInputs().size(); ++index)
|
|
if (auto inputArg = getInputArgument(index))
|
|
setNameFn(*inputArg, ("in" + std::to_string(index)).c_str());
|
|
}
|
|
|
|
std::optional<BlockArgument> SpatComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
|
|
|
|
std::optional<BlockArgument> SpatComputeBatch::getWeightArgument(unsigned idx) {
|
|
return getBlockArgument(getBody(), 1 + idx);
|
|
}
|
|
|
|
std::optional<BlockArgument> SpatComputeBatch::getInputArgument(unsigned idx) {
|
|
return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
|
|
}
|
|
|
|
std::optional<BlockArgument> SpatComputeBatch::getOutputArgument(unsigned idx) {
|
|
return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
|
|
}
|
|
|
|
std::optional<std::tuple<Value, BlockArgument>>
|
|
SpatComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
|
|
if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
|
|
auto index = std::distance(getWeights().begin(), existing);
|
|
return {
|
|
{*existing, *getWeightArgument(index)}
|
|
};
|
|
}
|
|
|
|
unsigned weightCount = getWeights().size();
|
|
unsigned inputCount = getInputs().size();
|
|
getOperation()->insertOperands(idx, ValueRange {weight});
|
|
setComputeOperandSegmentSizes(
|
|
getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
|
|
auto blockArg = insertBlockArgument(getBody(), 1 + idx, weight.getType(), loc);
|
|
if (!blockArg)
|
|
return std::nullopt;
|
|
return std::make_tuple(getOperation()->getOperand(idx), *blockArg);
|
|
}
|
|
|
|
std::optional<std::tuple<Value, BlockArgument>> SpatComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
|
|
unsigned weightCount = getWeights().size();
|
|
unsigned inputCount = getInputs().size();
|
|
getOperation()->insertOperands(weightCount + idx, ValueRange {input});
|
|
setComputeOperandSegmentSizes(
|
|
getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
|
|
auto blockArg = insertBlockArgument(getBody(), 1 + weightCount + idx, input.getType(), loc);
|
|
if (!blockArg)
|
|
return std::nullopt;
|
|
return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
|
|
}
|
|
|
|
CrossbarWeightSet SpatComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
|
|
|
|
FailureOr<std::tuple<OpResult, BlockArgument, SpatComputeBatch>>
|
|
SpatComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
|
|
if (idx > getNumResults())
|
|
return failure();
|
|
|
|
rewriter.setInsertionPoint(getOperation());
|
|
SmallVector<Type> resultTypes(getResultTypes().begin(), getResultTypes().end());
|
|
resultTypes.insert(resultTypes.begin() + idx, type);
|
|
auto newBatch =
|
|
SpatComputeBatch::create(rewriter, getLoc(), TypeRange(resultTypes), getLaneCountAttr(), getWeights(), getInputs());
|
|
newBatch->setAttrs((*this)->getAttrs());
|
|
setComputeOperandSegmentSizes(newBatch.getOperation(),
|
|
static_cast<int32_t>(newBatch.getWeights().size()),
|
|
static_cast<int32_t>(newBatch.getInputs().size()));
|
|
rewriter.inlineRegionBefore(getBody(), newBatch.getBody(), newBatch.getBody().end());
|
|
if (newBatch.getBody().empty()) {
|
|
rewriter.eraseOp(newBatch);
|
|
return failure();
|
|
}
|
|
auto blockArg = newBatch.getBody().front().insertArgument(
|
|
1 + newBatch.getWeights().size() + newBatch.getInputs().size() + idx, type, loc);
|
|
for (unsigned oldResultIdx = 0; oldResultIdx < getNumResults(); ++oldResultIdx)
|
|
getResult(oldResultIdx)
|
|
.replaceAllUsesWith(newBatch.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
|
|
rewriter.eraseOp(getOperation());
|
|
return std::make_tuple(cast<OpResult>(newBatch.getResult(idx)), blockArg, newBatch);
|
|
}
|
|
|
|
void SpatComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
|
|
if (region.empty())
|
|
return;
|
|
|
|
if (auto laneArg = getLaneArgument())
|
|
setNameFn(*laneArg, "lane");
|
|
|
|
for (unsigned index = 0; index < getWeights().size(); ++index)
|
|
if (auto weightArg = getWeightArgument(index))
|
|
setNameFn(*weightArg, ("w" + std::to_string(index)).c_str());
|
|
|
|
for (unsigned index = 0; index < getInputs().size(); ++index)
|
|
if (auto inputArg = getInputArgument(index))
|
|
setNameFn(*inputArg, ("in" + std::to_string(index)).c_str());
|
|
|
|
for (unsigned index = 0; index < getNumResults(); ++index) {
|
|
auto outputArg = getOutputArgument(index);
|
|
if (!outputArg)
|
|
continue;
|
|
if (index == 0) {
|
|
setNameFn(*outputArg, "out");
|
|
continue;
|
|
}
|
|
setNameFn(*outputArg, ("out" + std::to_string(index)).c_str());
|
|
}
|
|
}
|
|
|
|
void SpatInParallelOp::build(OpBuilder& builder, OperationState& result) {
|
|
OpBuilder::InsertionGuard guard(builder);
|
|
Region* bodyRegion = result.addRegion();
|
|
builder.createBlock(bodyRegion);
|
|
}
|
|
|
|
OpResult SpatInParallelOp::getParentResult(int64_t idx) { return getOperation()->getParentOp()->getResult(idx); }
|
|
|
|
llvm::iterator_range<Block::iterator> SpatInParallelOp::getYieldingOps() { return getRegion().front().getOperations(); }
|
|
|
|
void SpatialDialect::initialize() {
|
|
addTypes<
|
|
#define GET_TYPEDEF_LIST
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"
|
|
|
|
>();
|
|
addOperations<
|
|
#define GET_OP_LIST
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.cpp.inc"
|
|
|
|
>();
|
|
}
|
|
} // namespace spatial
|
|
} // namespace onnx_mlir
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// TableGen'd op method definitions
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#define GET_OP_CLASSES
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.cpp.inc"
|
|
|
|
#define GET_TYPEDEF_CLASSES
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.cpp.inc"
|
|
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.cpp.inc"
|