This commit is contained in:
NiccoloN
@@ -0,0 +1,239 @@
|
||||
#include "llvm/ADT/SmallVector.h"
|
||||
|
||||
#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
|
||||
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
|
||||
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
|
||||
#include "src/Dialect/ONNX/ONNXOps.hpp"
|
||||
|
||||
using namespace mlir;
|
||||
|
||||
namespace onnx_mlir {
|
||||
|
||||
RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
|
||||
return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
|
||||
logicalType.getElementType(),
|
||||
logicalType.getEncoding());
|
||||
}
|
||||
|
||||
RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
|
||||
return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
|
||||
logicalType.getElementType(),
|
||||
logicalType.getEncoding());
|
||||
}
|
||||
|
||||
std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
|
||||
SmallVector<int64_t> offsets;
|
||||
SmallVector<int64_t> sizes;
|
||||
const int64_t channels = type.getDimSize(1);
|
||||
const int64_t height = type.getDimSize(2);
|
||||
const int64_t width = type.getDimSize(3);
|
||||
offsets.reserve(height * 4);
|
||||
sizes.reserve(height * 4);
|
||||
for (int64_t row = 0; row < height; ++row) {
|
||||
offsets.append({0, 0, row, 0});
|
||||
sizes.append({1, channels, 1, width});
|
||||
}
|
||||
return {offsets, sizes};
|
||||
}
|
||||
|
||||
SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
|
||||
return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
|
||||
}
|
||||
|
||||
SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
|
||||
return {rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(logicalType.getDimSize(1)),
|
||||
rewriter.getIndexAttr(1),
|
||||
rewriter.getIndexAttr(logicalType.getDimSize(3))};
|
||||
}
|
||||
|
||||
Value extractRowStripFragment(Value storage,
|
||||
RankedTensorType logicalType,
|
||||
OpFoldResult row,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
return tensor::ExtractSliceOp::create(rewriter,
|
||||
loc,
|
||||
getRowStripFragmentType(logicalType),
|
||||
storage,
|
||||
buildRowStripFragmentOffsets(rewriter, row),
|
||||
buildRowStripFragmentSizes(rewriter, logicalType),
|
||||
getUnitStrides(rewriter, 4));
|
||||
}
|
||||
|
||||
void insertRowStripFragment(Value fragment,
|
||||
Value output,
|
||||
RankedTensorType logicalType,
|
||||
OpFoldResult row,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
createParallelInsertSliceIntoBatchOutput(rewriter,
|
||||
loc,
|
||||
fragment,
|
||||
output,
|
||||
buildRowStripFragmentOffsets(rewriter, row),
|
||||
buildRowStripFragmentSizes(rewriter, logicalType),
|
||||
getUnitStrides(rewriter, 4));
|
||||
}
|
||||
|
||||
FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
|
||||
RankedTensorType fragmentType,
|
||||
PatternRewriter& rewriter) {
|
||||
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
|
||||
if (failed(channelValues))
|
||||
return failure();
|
||||
|
||||
SmallVector<Attribute> values;
|
||||
values.reserve(fragmentType.getNumElements());
|
||||
for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
|
||||
for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
|
||||
for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
|
||||
for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
|
||||
values.push_back((*channelValues)[channel]);
|
||||
|
||||
auto attr = DenseElementsAttr::get(fragmentType, values);
|
||||
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
|
||||
}
|
||||
|
||||
FailureOr<Value> createRowStripStorageFromRows(Value rows,
|
||||
RankedTensorType logicalType,
|
||||
PatternRewriter& rewriter,
|
||||
Location loc) {
|
||||
auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
|
||||
if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
|
||||
return failure();
|
||||
if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
|
||||
return failure();
|
||||
if (logicalType.getDimSize(0) != 1)
|
||||
return failure();
|
||||
if (rowsType.getElementType() != logicalType.getElementType())
|
||||
return failure();
|
||||
|
||||
const int64_t channels = logicalType.getDimSize(1);
|
||||
const int64_t height = logicalType.getDimSize(2);
|
||||
const int64_t width = logicalType.getDimSize(3);
|
||||
if (rowsType.getDimSize(0) != height * width)
|
||||
return failure();
|
||||
if (rowsType.getDimSize(1) != channels)
|
||||
return failure();
|
||||
|
||||
auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
|
||||
auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
|
||||
auto fragmentType = getRowStripFragmentType(logicalType);
|
||||
auto storageType = getRowStripStorageType(logicalType);
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
||||
Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
|
||||
SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
|
||||
SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
|
||||
Value rowSlice = tensor::ExtractSliceOp::create(
|
||||
rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
|
||||
Value channelWidth = ONNXTransposeOp::create(
|
||||
rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
|
||||
Value fragment = tensor::ExpandShapeOp::create(
|
||||
rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
|
||||
insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
return success();
|
||||
});
|
||||
if (failed(batchOp))
|
||||
return failure();
|
||||
return batchOp->getResult(0);
|
||||
}
|
||||
|
||||
FailureOr<Value>
|
||||
materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
|
||||
auto storageType = dyn_cast<RankedTensorType>(storage.getType());
|
||||
if (!storageType || storageType != getRowStripStorageType(logicalType))
|
||||
return failure();
|
||||
|
||||
auto batchOp = createSpatComputeBatch(
|
||||
rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
createParallelInsertSliceIntoBatchOutput(rewriter,
|
||||
loc,
|
||||
fragment,
|
||||
args.outputs.front(),
|
||||
SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
|
||||
rewriter.getIndexAttr(0),
|
||||
args.lane,
|
||||
rewriter.getIndexAttr(0)},
|
||||
buildRowStripFragmentSizes(rewriter, logicalType),
|
||||
getUnitStrides(rewriter, 4));
|
||||
return success();
|
||||
});
|
||||
if (failed(batchOp))
|
||||
return failure();
|
||||
return batchOp->getResult(0);
|
||||
}
|
||||
|
||||
FailureOr<Value>
|
||||
applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
|
||||
auto fragmentType = getRowStripFragmentType(logicalType);
|
||||
auto storageType = getRowStripStorageType(logicalType);
|
||||
auto batchOp = createSpatComputeBatch(rewriter,
|
||||
loc,
|
||||
TypeRange {storageType},
|
||||
logicalType.getDimSize(2),
|
||||
{},
|
||||
ValueRange {storage},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Value fragment =
|
||||
extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
|
||||
insertRowStripFragment(
|
||||
fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
return success();
|
||||
});
|
||||
if (failed(batchOp))
|
||||
return failure();
|
||||
return batchOp->getResult(0);
|
||||
}
|
||||
|
||||
FailureOr<Value>
|
||||
applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
|
||||
DenseElementsAttr denseAttr;
|
||||
if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
|
||||
return failure();
|
||||
auto fragmentType = getRowStripFragmentType(logicalType);
|
||||
auto storageType = getRowStripStorageType(logicalType);
|
||||
auto batchOp = createSpatComputeBatch(rewriter,
|
||||
loc,
|
||||
TypeRange {storageType},
|
||||
logicalType.getDimSize(2),
|
||||
{},
|
||||
ValueRange {storage},
|
||||
[&](detail::SpatComputeBatchBodyArgs args) {
|
||||
Value fragment =
|
||||
extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
Value constant;
|
||||
if (denseAttr.isSplat()) {
|
||||
constant = getOrCreateConstant(
|
||||
rewriter,
|
||||
rewriter.getInsertionBlock()->getParentOp(),
|
||||
DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
|
||||
fragmentType);
|
||||
}
|
||||
else {
|
||||
FailureOr<Value> perChannel =
|
||||
createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
|
||||
if (failed(perChannel))
|
||||
return failure();
|
||||
constant = *perChannel;
|
||||
}
|
||||
fragment =
|
||||
spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
|
||||
insertRowStripFragment(
|
||||
fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
|
||||
return success();
|
||||
});
|
||||
if (failed(batchOp))
|
||||
return failure();
|
||||
return batchOp->getResult(0);
|
||||
}
|
||||
|
||||
} // namespace onnx_mlir
|
||||
Reference in New Issue
Block a user