ilgeco
253 lines
14 KiB
C++
253 lines
14 KiB
C++
#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/ComputeRegionBuilder.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 {
|
|
|
|
FailureOr<RowStripPhysicalValue> describeRowStripPhysicalValue(Value storage, RankedTensorType logicalType) {
|
|
auto storageType = dyn_cast<RankedTensorType>(storage.getType());
|
|
if (!storageType || !storageType.hasStaticShape() || !logicalType || !logicalType.hasStaticShape()
|
|
|| storageType.getRank() != 5 || logicalType.getRank() != 4 || logicalType.getDimSize(0) != 1
|
|
|| storageType.getElementType() != logicalType.getElementType()
|
|
|| storageType.getDimSize(1) != 1 || storageType.getDimSize(2) != 1
|
|
|| storageType.getDimSize(3) != logicalType.getDimSize(3) || storageType.getDimSize(4) <= 0)
|
|
return failure();
|
|
const int64_t tilesPerRow = ceilIntegerDivide(logicalType.getDimSize(1), storageType.getDimSize(4));
|
|
if (storageType.getDimSize(0) != logicalType.getDimSize(2) * tilesPerRow)
|
|
return failure();
|
|
return RowStripPhysicalValue {storage, logicalType,
|
|
RankedTensorType::get(storageType.getShape().drop_front(), storageType.getElementType(), storageType.getEncoding()),
|
|
tilesPerRow};
|
|
}
|
|
|
|
RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
|
|
return RankedTensorType::get({logicalType.getDimSize(0), 1, logicalType.getDimSize(3),
|
|
logicalType.getDimSize(1)},
|
|
logicalType.getElementType(),
|
|
logicalType.getEncoding());
|
|
}
|
|
|
|
RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
|
|
return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
|
|
}
|
|
|
|
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};
|
|
}
|
|
|
|
Value extractRowStripFragment(Value storage,
|
|
RankedTensorType logicalType,
|
|
OpFoldResult row,
|
|
PatternRewriter& rewriter,
|
|
Location loc) {
|
|
return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
|
|
}
|
|
|
|
void insertRowStripFragment(Value fragment,
|
|
Value output,
|
|
RankedTensorType logicalType,
|
|
OpFoldResult row,
|
|
PatternRewriter& rewriter,
|
|
Location loc) {
|
|
assert(fragment.getType() == getRowStripFragmentType(logicalType));
|
|
assert(output.getType() == getRowStripStorageType(logicalType));
|
|
auto slot = dyn_cast<Value>(row);
|
|
assert(slot && "row-strip graph publication requires a dynamic physical slot");
|
|
publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
|
|
}
|
|
|
|
FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
|
|
RankedTensorType fragmentType,
|
|
PatternRewriter& rewriter) {
|
|
auto logicalType = RankedTensorType::get(
|
|
{1, fragmentType.getDimSize(3), 1, 1}, fragmentType.getElementType());
|
|
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, logicalType);
|
|
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 h = 0; h < fragmentType.getDimSize(1); ++h)
|
|
for (int64_t w = 0; w < fragmentType.getDimSize(2); ++w)
|
|
for (int64_t channel = 0; channel < fragmentType.getDimSize(3); ++channel)
|
|
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 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 fragment = tensor::ExpandShapeOp::create(
|
|
rewriter, loc, fragmentType, rowSlice, 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> createRowStripAssemblyBlueprint(const RowStripPhysicalValue& value,
|
|
PatternRewriter& rewriter,
|
|
Location loc) {
|
|
const int64_t laneCount = cast<RankedTensorType>(value.storage.getType()).getDimSize(0);
|
|
const int64_t tileChannels = value.fragmentType.getDimSize(3);
|
|
auto nchwFragmentType = RankedTensorType::get(
|
|
{1, tileChannels, 1, value.logicalType.getDimSize(3)}, value.fragmentType.getElementType(),
|
|
value.fragmentType.getEncoding());
|
|
auto nchwStorageType = spatial::getGraphBatchPhysicalResultType(laneCount, nchwFragmentType);
|
|
auto transposed = createSpatComputeBatch(
|
|
rewriter, loc, TypeRange {nchwStorageType}, laneCount, {}, ValueRange {value.storage},
|
|
[&](detail::SpatComputeBatchBodyArgs args) {
|
|
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
|
|
rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
|
|
if (failed(fragment))
|
|
return failure();
|
|
Value nchw = ONNXTransposeOp::create(
|
|
rewriter, loc, nchwFragmentType, *fragment, rewriter.getI64ArrayAttr({0, 3, 1, 2}));
|
|
publishGraphBatchPhysicalFragment(rewriter, loc, nchw, args.outputs.front(), args.lane);
|
|
return success();
|
|
});
|
|
if (failed(transposed))
|
|
return failure();
|
|
|
|
SmallVector<FragmentAssemblyEntry> entries;
|
|
for (int64_t row = 0; row < value.logicalType.getDimSize(2); ++row)
|
|
for (int64_t tile = 0; tile < value.tilesPerRow; ++tile) {
|
|
const int64_t channelOffset = tile * tileChannels;
|
|
entries.push_back({row * value.tilesPerRow + tile, 0, {0, channelOffset, row, 0},
|
|
{1, std::min(tileChannels, value.logicalType.getDimSize(1) - channelOffset), 1,
|
|
value.logicalType.getDimSize(3)}});
|
|
}
|
|
return createFragmentAssemblyBlueprint(transposed->getResult(0), value.logicalType, entries, "nhwc_row_strip",
|
|
kRowStripIndexMap, rewriter, loc);
|
|
}
|
|
|
|
FailureOr<Value> applyRowStripRelu(const RowStripPhysicalValue& value, PatternRewriter& rewriter, Location loc) {
|
|
auto storageType = cast<RankedTensorType>(value.storage.getType());
|
|
const int64_t laneCount = storageType.getDimSize(0);
|
|
auto batchOp = createSpatComputeBatch(rewriter,
|
|
loc,
|
|
TypeRange {storageType},
|
|
laneCount,
|
|
{},
|
|
ValueRange {value.storage},
|
|
[&](detail::SpatComputeBatchBodyArgs args) {
|
|
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
|
|
rewriter, loc, args.inputs.front(), args.lane, value.fragmentType);
|
|
if (failed(fragment)) return failure();
|
|
Value relu = spatial::SpatReluOp::create(
|
|
rewriter, loc, value.fragmentType, *fragment).getResult();
|
|
publishGraphBatchPhysicalFragment(
|
|
rewriter, loc, relu, args.outputs.front(), args.lane);
|
|
return success();
|
|
});
|
|
if (failed(batchOp))
|
|
return failure();
|
|
return batchOp->getResult(0);
|
|
}
|
|
|
|
FailureOr<Value> applyRowStripBiasAdd(const RowStripPhysicalValue& value,
|
|
Value bias,
|
|
PatternRewriter& rewriter,
|
|
Location loc) {
|
|
DenseElementsAttr denseAttr;
|
|
if (!isSupportedBiasAddValue(bias, value.logicalType, &denseAttr))
|
|
return failure();
|
|
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, value.logicalType);
|
|
if (failed(channelValues)) return failure();
|
|
auto storageType = cast<RankedTensorType>(value.storage.getType());
|
|
auto biasStorageType = spatial::getGraphBatchPhysicalResultType(value.tilesPerRow, value.fragmentType);
|
|
SmallVector<Attribute> biasValues(
|
|
biasStorageType.getNumElements(), cast<Attribute>(rewriter.getZeroAttr(value.fragmentType.getElementType())));
|
|
const int64_t tileChannels = value.fragmentType.getDimSize(3);
|
|
const int64_t width = value.fragmentType.getDimSize(2);
|
|
for (int64_t channel = 0; channel < value.logicalType.getDimSize(1); ++channel)
|
|
for (int64_t w = 0; w < width; ++w)
|
|
biasValues[(channel / tileChannels) * width * tileChannels + w * tileChannels
|
|
+ channel % tileChannels] =
|
|
(*channelValues)[channel];
|
|
Value biasStorage = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(),
|
|
DenseElementsAttr::get(biasStorageType, biasValues), biasStorageType);
|
|
const int64_t laneCount = storageType.getDimSize(0);
|
|
auto batchOp = createSpatComputeBatch(rewriter,
|
|
loc,
|
|
TypeRange {storageType},
|
|
laneCount,
|
|
{},
|
|
ValueRange {value.storage, biasStorage},
|
|
[&](detail::SpatComputeBatchBodyArgs args) {
|
|
Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
|
|
FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(
|
|
rewriter, loc, args.inputs[0], args.lane, value.fragmentType);
|
|
Value tile = affineModConst(rewriter, loc, args.lane,
|
|
value.tilesPerRow, anchorOp);
|
|
FailureOr<Value> constant = extractGraphBatchPhysicalFragment(
|
|
rewriter, loc, args.inputs[1], tile, value.fragmentType);
|
|
if (failed(fragment) || failed(constant)) return failure();
|
|
Value added = spatial::SpatVAddOp::create(
|
|
rewriter, loc, value.fragmentType, *fragment, *constant).getResult();
|
|
publishGraphBatchPhysicalFragment(
|
|
rewriter, loc, added, args.outputs.front(), args.lane);
|
|
return success();
|
|
});
|
|
if (failed(batchOp))
|
|
return failure();
|
|
return batchOp->getResult(0);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|