NiccoloN
113 lines
4.2 KiB
C++
113 lines
4.2 KiB
C++
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/BuiltinTypes.h"
|
|
|
|
#include "llvm/ADT/SmallVector.h"
|
|
|
|
#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
|
|
#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
|
|
#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace onnx_mlir {
|
|
|
|
LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
|
|
if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
|
|
return failure();
|
|
if (resultType.getRank() != 4)
|
|
return failure();
|
|
if (biasType.getElementType() != resultType.getElementType())
|
|
return failure();
|
|
|
|
const int64_t channels = resultType.getDimSize(1);
|
|
ArrayRef<int64_t> shape = biasType.getShape();
|
|
if (shape.empty())
|
|
return success();
|
|
if (shape.size() == 1)
|
|
return success(shape[0] == channels);
|
|
if (shape.size() == 2)
|
|
return success(shape[0] == 1 && shape[1] == channels);
|
|
if (shape.size() == 4)
|
|
return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
|
|
return failure();
|
|
}
|
|
|
|
FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
|
|
auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
|
|
if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
|
|
return failure();
|
|
|
|
const int64_t channels = resultType.getDimSize(1);
|
|
if (denseAttr.isSplat()) {
|
|
return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
|
|
}
|
|
|
|
SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
|
|
if (biasType.getRank() == 1)
|
|
return flattened;
|
|
if (biasType.getRank() == 2)
|
|
return flattened;
|
|
|
|
SmallVector<Attribute> channelValues;
|
|
channelValues.reserve(channels);
|
|
const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
|
|
for (int64_t channel = 0; channel < channels; ++channel)
|
|
channelValues.push_back(flattened[channel * channelStride]);
|
|
return channelValues;
|
|
}
|
|
|
|
bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
|
|
auto attr = getHostConstDenseElementsAttr(bias);
|
|
if (!attr)
|
|
return false;
|
|
auto biasType = dyn_cast<RankedTensorType>(attr.getType());
|
|
if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
|
|
return false;
|
|
if (failed(getBiasChannelValues(attr, resultType)))
|
|
return false;
|
|
if (denseAttr)
|
|
*denseAttr = attr;
|
|
return true;
|
|
}
|
|
|
|
FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
|
|
auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
|
|
auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
|
|
if (!lhsType || !rhsType)
|
|
return failure();
|
|
if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
|
|
return BiasAddPlanCandidate {lhs, rhs};
|
|
if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
|
|
return BiasAddPlanCandidate {rhs, lhs};
|
|
return failure();
|
|
}
|
|
|
|
FailureOr<Value>
|
|
materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
|
|
DenseElementsAttr denseAttr;
|
|
if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
|
|
return failure();
|
|
|
|
FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
|
|
if (failed(channelValues))
|
|
return failure();
|
|
|
|
SmallVector<Attribute> resultValues;
|
|
resultValues.reserve(resultType.getNumElements());
|
|
const int64_t batches = resultType.getDimSize(0);
|
|
const int64_t channels = resultType.getDimSize(1);
|
|
const int64_t height = resultType.getDimSize(2);
|
|
const int64_t width = resultType.getDimSize(3);
|
|
for (int64_t n = 0; n < batches; ++n)
|
|
for (int64_t c = 0; c < channels; ++c)
|
|
for (int64_t h = 0; h < height; ++h)
|
|
for (int64_t w = 0; w < width; ++w)
|
|
resultValues.push_back((*channelValues)[c]);
|
|
|
|
auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
|
|
return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
|
|
}
|
|
|
|
} // namespace onnx_mlir
|