diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt index a99f6d4..a665c0a 100644 --- a/core/CMakeLists.txt +++ b/core/CMakeLists.txt @@ -19,6 +19,7 @@ set(sources_headers SOFIE/RModel.hxx SOFIE/RModelProfiler.hxx SOFIE/RModelProfilerGPU.hxx + SOFIE/RWeightFile.hxx SOFIE/ROperator.hxx SOFIE/ROperator_BasicUnary.hxx SOFIE/ROperator_BasicBinary.hxx @@ -27,6 +28,16 @@ set(sources_headers SOFIE/ROperator_Conv.hxx SOFIE/ROperator_ConvTranspose.hxx SOFIE/ROperator_Gemm.hxx + SOFIE/ROperator_QuantizedGemm.hxx + SOFIE/ROperator_QuantizedMatMul.hxx + SOFIE/ROperator_QuantizedMatrix.hxx + SOFIE/SOFIE_Quantized.hxx + SOFIE/SOFIE_QuantizedAlpaka.hxx + SOFIE/RQuantization.hxx + SOFIE/RQuantization_Analysis.hxx + SOFIE/RQuantization_DenseLinear.hxx + SOFIE/RQuantization_Parameters.hxx + SOFIE/RQuantization_Storage.hxx SOFIE/ROperator_Relu.hxx SOFIE/ROperator_Tanh.hxx SOFIE/ROperator_LeakyRelu.hxx @@ -57,6 +68,8 @@ set(sources_headers SOFIE/ROperator_Split.hxx SOFIE/ROperator_SubGraph.hxx SOFIE/ROperator_Pad.hxx + SOFIE/ROperator_QONNXQuant.hxx + SOFIE/ROperator_ONNXQuantizeLinear.hxx SOFIE/ROperator_Where.hxx SOFIE/ROperator_Einsum.hxx SOFIE/ROperator_Random.hxx @@ -78,6 +91,12 @@ list(TRANSFORM sources_headers PREPEND "inc/") set(sources_cxx src/RModel_Base.cxx src/RModel.cxx + src/RModel_Quantization.cxx + src/RWeightFile.cxx + src/RQuantization_Analysis.cxx + src/RQuantization_DenseLinear.cxx + src/RQuantization_Parameters.cxx + src/RQuantization_Storage.cxx src/RModelProfiler.cxx src/RModelProfilerGPU.cxx src/RModel_ALPAKA.cxx @@ -120,4 +139,3 @@ install(TARGETS SOFIE_core install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/inc/" DESTINATION ${CMAKE_INSTALL_INCLUDEDIR} ) - diff --git a/core/inc/SOFIE/OperatorList.hxx b/core/inc/SOFIE/OperatorList.hxx index 02ae60d..4e268e9 100644 --- a/core/inc/SOFIE/OperatorList.hxx +++ b/core/inc/SOFIE/OperatorList.hxx @@ -1,5 +1,8 @@ #include "SOFIE/ROperator_Transpose.hxx" #include "SOFIE/ROperator_Gemm.hxx" +#include "SOFIE/ROperator_QuantizedGemm.hxx" +#include "SOFIE/ROperator_QuantizedMatMul.hxx" +#include "SOFIE/ROperator_ONNXQuantizeLinear.hxx" #include "SOFIE/ROperator_Relu.hxx" #include "SOFIE/ROperator_Tanh.hxx" #include "SOFIE/ROperator_LeakyRelu.hxx" diff --git a/core/inc/SOFIE/RModel.hxx b/core/inc/SOFIE/RModel.hxx index 8153408..2a7d7c5 100644 --- a/core/inc/SOFIE/RModel.hxx +++ b/core/inc/SOFIE/RModel.hxx @@ -3,6 +3,7 @@ #include "SOFIE/RModel_Base.hxx" #include "SOFIE/SOFIE_common.hxx" +#include "SOFIE/RQuantization.hxx" #include "SOFIE/ROperator.hxx" @@ -32,6 +33,7 @@ private: std::unordered_map fInitializedTensors; std::unordered_map fIntermediateTensorInfos; std::unordered_map fDynamicTensorInfos; + QuantizationModelState fQuantizationState; // quantization metadata, regions, storage, and backend lowering plans std::unordered_map, bool>> fShapeTensors; // constant tensors describing a shape std::unordered_map fAliasTensors; // alias tensors (name -> original tensor name) std::unordered_map @@ -42,6 +44,11 @@ private: std::vector> fOperators; + // transient lowered operator view while the parsed fOperators graph remains intact. + // let code generation use synthetic operators for proven regions. + std::unordered_map> fLoweredOperators; + std::set fLoweredConsumedOperatorIndices; + std::vector> fSubGraphs; /// GetDynamicTensorShape(const std::string & name) const ; + void AddQuantizationInfo(const std::string & tensor_name, QuantizationInfo info); + bool HasQuantizationInfo(const std::string & tensor_name) const; + const QuantizationInfo & GetQuantizationInfo(const std::string & tensor_name) const; + + void RegisterQuantizedTensorStorage(QuantizedTensorStorage storage); + bool HasQuantizedTensorStorage(const std::string & storage_tensor_name) const; + const QuantizedTensorStorage & GetQuantizedTensorStorage(const std::string & storage_tensor_name) const; + + void AnalyzeQuantizedRegions(); + const QuantizationModelState & GetQuantizationState() const { return fQuantizationState; } + // get the values for the tensor representing a shape const std::vector & GetShapeTensorValues(const std::string & tensor_name) const; @@ -145,6 +167,18 @@ public: AddInitializedTensor(tensor_name, GetTemplatedType(T()), shape, data); } + template + void AddInitializedTensor(const std::string & tensor_name, const std::vector & shape, const std::vector &values) + { + size_t size = ConvertShapeToLength(shape); + if (values.size() != size) { + throw std::runtime_error("sofie: initialized tensor " + tensor_name + " data length does not match shape"); + } + std::shared_ptr data(malloc(size * sizeof(T)), free); + std::copy(values.begin(), values.end(), static_cast(data.get())); + AddInitializedTensor(tensor_name, GetTemplatedType(T()), shape, data); + } + void AddShapeTensor(const std::string & name, const std::vector & shapeValues, bool scalar = false); void AddAliasTensor(const std::string & name, const std::string & origin); bool IsAliasTensor(const std::string & tensor_name) const; diff --git a/core/inc/SOFIE/RModel_Base.hxx b/core/inc/SOFIE/RModel_Base.hxx index b598652..179377d 100644 --- a/core/inc/SOFIE/RModel_Base.hxx +++ b/core/inc/SOFIE/RModel_Base.hxx @@ -28,6 +28,8 @@ enum class Options { kGNN = 0x8, kGNNComponent = 0x10, kProfile = 0x20, + kBinaryWeightFile = 0x40, + kTextWeightFile = 0x80, }; // Optimization levels inspired by ONNXRuntime. @@ -39,7 +41,7 @@ enum class OptimizationLevel { kExtended = 0x1, }; -enum class WeightFileType { None, RootBinary, Text }; +enum class WeightFileType { None, RootBinary, Text, Binary }; inline std::underlying_type_t operator|(Options opA, Options opB) { diff --git a/core/inc/SOFIE/ROperator.hxx b/core/inc/SOFIE/ROperator.hxx index c24fd70..84439d2 100644 --- a/core/inc/SOFIE/ROperator.hxx +++ b/core/inc/SOFIE/ROperator.hxx @@ -38,12 +38,16 @@ enum class OperatorKind { UNARY_COS=22, UNARY_ABS=23, CLIP=24, - NOT=25 + NOT=25, + QUANTIZED_GEMM=26, + QUANTIZED_MATMUL=27 }; inline const char* toString(OperatorKind kind) { switch (kind) { case OperatorKind::GEMM: return "GEMM"; + case OperatorKind::QUANTIZED_GEMM: return "QUANTIZED_GEMM"; + case OperatorKind::QUANTIZED_MATMUL: return "QUANTIZED_MATMUL"; case OperatorKind::LAYERNORM: return "LAYERNORM"; case OperatorKind::RELU: return "RELU"; case OperatorKind::CONSTANT: return "CONSTANT"; @@ -82,6 +86,15 @@ public: virtual std::string GetBlasConfig() { return ""; } virtual void UpdateFusableTensorName(std::string, const std::function& removal_func){ return;}; + // Semantic graph-analysis hooks + virtual bool IsQuantizationBoundary() const { return false; } + virtual std::string GetQuantizationSourceTensor() const + { + if (fInputTensorNames.empty()) + return {}; + return std::string(fInputTensorNames.front()); + } + // Elementwise kernel fusion interface virtual bool IsElementwise() const { return false; } // Returns the C++ expression applying this op to inputVar (a local T variable) for fused kernel generation diff --git a/core/inc/SOFIE/ROperator_Gemm.hxx b/core/inc/SOFIE/ROperator_Gemm.hxx index eecb33b..1f58266 100644 --- a/core/inc/SOFIE/ROperator_Gemm.hxx +++ b/core/inc/SOFIE/ROperator_Gemm.hxx @@ -12,6 +12,7 @@ #include #include #include +#include namespace SOFIE{ @@ -72,6 +73,20 @@ namespace SOFIE{ fOutputTensorNames = { fNY }; fKind = OperatorKind::GEMM; } + // Getters for further quantized GEMM lowering-elimination. + float GetAlpha() const { return fAttrAlpha; } + float GetBeta() const { return fAttrBeta; } + int_t GetTransA() const { return fAttrTransA; } + int_t GetTransB() const { return fAttrTransB; } + bool HasBias() const { return !fNC.empty(); } + const std::string &GetInputTensorName() const { return fNA; } + const std::string &GetWeightTensorName() const { return fNB; } + const std::string &GetBiasTensorName() const { return fNC; } + const std::string &GetOutputTensorName() const { return fNY; } + const std::vector &GetInputShape() const { return fShapeA; } + const std::vector &GetWeightShape() const { return fShapeB; } + const std::vector &GetOutputShape() const { return fShapeY; } + std::vector GetStdLibs() override { return { std::string("cmath"), std::string("cstdint"), std::string("vector") }; } std::vector TypeInference(std::vector input) override { ETensorType out = input[0]; diff --git a/core/inc/SOFIE/ROperator_ONNXQuantizeLinear.hxx b/core/inc/SOFIE/ROperator_ONNXQuantizeLinear.hxx new file mode 100644 index 0000000..3a12f3f --- /dev/null +++ b/core/inc/SOFIE/ROperator_ONNXQuantizeLinear.hxx @@ -0,0 +1,334 @@ +#ifndef SOFIE_ROPERATOR_ONNXQUANTIZELINEAR +#define SOFIE_ROPERATOR_ONNXQUANTIZELINEAR + +#include "SOFIE/RModel.hxx" +#include "SOFIE/ROperator.hxx" +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/RQuantization_Parameters.hxx" +#include "SOFIE/SOFIE_common.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +namespace DETAIL { + +inline unsigned BitWidthForONNXQuantizedType(ETensorType type) +{ + switch (type) { + case ETensorType::INT8: + case ETensorType::UINT8: + return 8; + case ETensorType::INT16: + case ETensorType::UINT16: + return 16; + case ETensorType::INT32: + case ETensorType::UINT32: + return 32; + default: + throw std::runtime_error("SOFIE ONNX Q/DQ supports integer zero-point carrier types only"); + } +} + +inline bool IsSignedONNXQuantizedType(ETensorType type) +{ + switch (type) { + case ETensorType::INT8: + case ETensorType::INT16: + case ETensorType::INT32: + return true; + case ETensorType::UINT8: + case ETensorType::UINT16: + case ETensorType::UINT32: + return false; + default: + throw std::runtime_error("SOFIE ONNX Q/DQ supports integer zero-point carrier types only"); + } +} + +inline std::vector GetFloatScaleInitializer(RModel &model, const std::string &tensorName, + const std::string &opName) +{ + auto values = model.GetTensorData(tensorName); + if (values.empty()) { + throw std::runtime_error("SOFIE " + opName + " expected non-empty FLOAT scale initializer " + tensorName); + } + return values; +} + +inline std::vector GetIntegerZeroPointInitializer(RModel &model, const std::string &tensorName, + ETensorType type, const std::string &opName) +{ + if (tensorName.empty()) + return {0}; + + switch (type) { + case ETensorType::INT8: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + case ETensorType::UINT8: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + case ETensorType::INT16: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + case ETensorType::UINT16: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + case ETensorType::INT32: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + case ETensorType::UINT32: { + const auto values = model.GetTensorData(tensorName); + return {values.begin(), values.end()}; + } + default: + throw std::runtime_error("SOFIE " + opName + " zero-point tensor has unsupported integer carrier type"); + } +} + +inline QuantizationInfo MakeONNXQDQInfo(RModel &model, const std::string &scaleTensor, + const std::string &zeroPointTensor, ETensorType carrierType, + const std::vector &tensorShape, int explicitAxis, + const std::string &opName) +{ + if (!model.IsInitializedTensor(scaleTensor)) { + throw std::runtime_error("SOFIE " + opName + " scale must be an initialized tensor"); + } + if (!zeroPointTensor.empty() && !model.IsInitializedTensor(zeroPointTensor)) { + throw std::runtime_error("SOFIE " + opName + " zero-point must be an initialized tensor when provided"); + } + + const auto scales = GetFloatScaleInitializer(model, scaleTensor, opName); + const auto zeroPoints = GetIntegerZeroPointInitializer(model, zeroPointTensor, carrierType, opName); + QuantizationParameterSpec spec; + spec.scales.assign(scales.begin(), scales.end()); + spec.zeroPoints = zeroPoints; + spec.bitWidth = BitWidthForONNXQuantizedType(carrierType); + spec.isSigned = IsSignedONNXQuantizedType(carrierType); + spec.narrow = false; + spec.rounding = EQuantizationRoundingMode::ROUND; + spec.overflow = EQuantizationOverflowMode::SAT; + spec.scaleTensor = scaleTensor; + spec.zeroPointTensor = zeroPointTensor; + spec.tensorShape = tensorShape; + spec.explicitAxis = explicitAxis; + spec.context = "SOFIE " + opName; + return MakeValidatedQuantizationInfo(spec); +} + +inline double ScaleForElement(const QuantizationInfo &info, RModel &model, const std::vector &shape, + std::size_t linearIndex) +{ + if (info.granularity != EQuantizationGranularity::PerChannel || info.axis < 0) + return info.scale; + const auto scales = model.GetTensorData(info.scaleTensor); + if (scales.empty()) + return info.scale; + std::size_t stride = 1; + for (std::size_t i = static_cast(info.axis) + 1; i < shape.size(); ++i) + stride *= shape[i]; + const std::size_t channel = (linearIndex / stride) % scales.size(); + return static_cast(scales[channel]); +} + +inline std::int64_t ZeroPointForElement(const QuantizationInfo &info, RModel &model, const std::vector &shape, + ETensorType carrierType, std::size_t linearIndex) +{ + if (info.zeroPointTensor.empty()) + return info.zeroPoint; + if (info.granularity != EQuantizationGranularity::PerChannel || info.axis < 0) + return info.zeroPoint; + const auto zeroPoints = GetIntegerZeroPointInitializer(model, info.zeroPointTensor, carrierType, "ONNX Q/DQ"); + if (zeroPoints.empty()) + return info.zeroPoint; + std::size_t stride = 1; + for (std::size_t i = static_cast(info.axis) + 1; i < shape.size(); ++i) + stride *= shape[i]; + const std::size_t channel = (linearIndex / stride) % zeroPoints.size(); + return zeroPoints[channel]; +} + +} // namespace DETAIL + +class ROperator_ONNXQuantizeLinear final : public ROperator { +private: + std::string fNX; + std::string fNScale; + std::string fNZeroPoint; + std::string fNY; + ETensorType fOutputType = ETensorType::UNDEFINED; + int fAxis = -1; + std::vector fShape; + QuantizationInfo fInfo; + +public: + ROperator_ONNXQuantizeLinear() = default; + + ROperator_ONNXQuantizeLinear(std::string nameX, std::string nameScale, std::string nameZeroPoint, + std::string nameY, ETensorType outputType, int axis) + : fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)), + fNZeroPoint(UTILITY::Clean_name(nameZeroPoint)), fNY(UTILITY::Clean_name(nameY)), + fOutputType(outputType), fAxis(axis) + { + fInputTensorNames = fNZeroPoint.empty() ? std::vector{fNX, fNScale} + : std::vector{fNX, fNScale, fNZeroPoint}; + fOutputTensorNames = {fNY}; + } + + bool IsQuantizationBoundary() const override { return true; } + std::string GetQuantizationSourceTensor() const override { return fNX; } + + std::vector TypeInference(std::vector) override { return {fOutputType}; } + std::vector> ShapeInference(std::vector> input) override + { + if (input.empty()) return {}; + return {input.front()}; + } + + void Initialize(RModel &model) override + { + if (!model.CheckIfTensorAlreadyExist(fNX)) + throw std::runtime_error("SOFIE ONNX QuantizeLinear input tensor " + fNX + " is not found in model"); + fShape = model.GetTensorShape(fNX); + fInfo = DETAIL::MakeONNXQDQInfo(model, fNScale, fNZeroPoint, fOutputType, fShape, fAxis, + "ONNX QuantizeLinear"); + model.AddQuantizationInfo(fNY, fInfo); + model.AddIntermediateTensor(fNY, fOutputType, fShape); + model.AddNeededStdLib("cmath"); + model.AddNeededStdLib("cstdint"); + } + + std::string Generate(std::string OpName) override + { + OpName = "op_" + OpName; + if (fInfo.granularity != EQuantizationGranularity::PerTensor) { + throw std::runtime_error("SOFIE ONNX QuantizeLinear literal code generation supports scalar parameters only; vector parameters require fused lowering"); + } + const auto [qMin, qMax] = QuantizedIntegerRange(fInfo); + const auto length = ConvertShapeToLength(fShape); + std::stringstream out; + out << "\n//------ ONNX QUANTIZELINEAR " << OpName << "\n"; + out << SP << "for (size_t id = 0; id < " << length << "; ++id) {\n"; + out << SP << SP << "double q = std::nearbyint((static_cast(tensor_" << fNX << "[id]) / " + << fInfo.scale << ") + " << fInfo.zeroPoint << ");\n"; + out << SP << SP << "q = (q < " << qMin << ") ? " << qMin << " : ((q > " << qMax << ") ? " << qMax << " : q);\n"; + out << SP << SP << "tensor_" << fNY << "[id] = static_cast<" << ConvertTypeToString(fOutputType) << ">(q);\n"; + out << SP << "}\n"; + return out.str(); + } +}; + +class ROperator_ONNXDequantizeLinear final : public ROperator { +private: + std::string fNX; + std::string fNScale; + std::string fNZeroPoint; + std::string fNY; + ETensorType fInputType = ETensorType::UNDEFINED; + int fAxis = -1; + std::vector fShape; + QuantizationInfo fInfo; + + template + std::vector DequantizeInitializedTensor(RModel &model) const + { + const auto values = model.GetTensorData(fNX); + std::vector output(values.size()); + for (std::size_t i = 0; i < values.size(); ++i) { + const double scale = DETAIL::ScaleForElement(fInfo, model, fShape, i); + const auto zeroPoint = DETAIL::ZeroPointForElement(fInfo, model, fShape, fInputType, i); + output[i] = static_cast((static_cast(values[i]) - zeroPoint) * scale); + } + return output; + } + +public: + ROperator_ONNXDequantizeLinear() = default; + + ROperator_ONNXDequantizeLinear(std::string nameX, std::string nameScale, std::string nameZeroPoint, + std::string nameY, ETensorType inputType, int axis) + : fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)), + fNZeroPoint(UTILITY::Clean_name(nameZeroPoint)), fNY(UTILITY::Clean_name(nameY)), + fInputType(inputType), fAxis(axis) + { + fInputTensorNames = fNZeroPoint.empty() ? std::vector{fNX, fNScale} + : std::vector{fNX, fNScale, fNZeroPoint}; + fOutputTensorNames = {fNY}; + } + + bool IsQuantizationBoundary() const override { return true; } + std::string GetQuantizationSourceTensor() const override { return fNX; } + + std::vector TypeInference(std::vector) override { return {ETensorType::FLOAT}; } + std::vector> ShapeInference(std::vector> input) override + { + if (input.empty()) return {}; + return {input.front()}; + } + + void Initialize(RModel &model) override + { + if (!model.CheckIfTensorAlreadyExist(fNX)) + throw std::runtime_error("SOFIE ONNX DequantizeLinear input tensor " + fNX + " is not found in model"); + fShape = model.GetTensorShape(fNX); + fInfo = DETAIL::MakeONNXQDQInfo(model, fNScale, fNZeroPoint, fInputType, fShape, fAxis, + "ONNX DequantizeLinear"); + model.AddQuantizationInfo(fNY, fInfo); + if (model.IsInitializedTensor(fNX)) { + std::vector values; + switch (fInputType) { + case ETensorType::INT8: values = DequantizeInitializedTensor(model); break; + case ETensorType::UINT8: values = DequantizeInitializedTensor(model); break; + case ETensorType::INT16: values = DequantizeInitializedTensor(model); break; + case ETensorType::UINT16: values = DequantizeInitializedTensor(model); break; + case ETensorType::INT32: values = DequantizeInitializedTensor(model); break; + case ETensorType::UINT32: values = DequantizeInitializedTensor(model); break; + default: + throw std::runtime_error("SOFIE ONNX DequantizeLinear supports initialized integer carriers only"); + } + model.AddConstantTensor(fNY, fShape, values); + } else { + model.AddIntermediateTensor(fNY, ETensorType::FLOAT, fShape); + } + model.AddNeededStdLib("cmath"); + model.AddNeededStdLib("cstdint"); + } + + std::string Generate(std::string OpName) override + { + OpName = "op_" + OpName; + if (fInfo.granularity != EQuantizationGranularity::PerTensor) { + throw std::runtime_error("SOFIE ONNX DequantizeLinear literal code generation supports scalar parameters only; vector parameters require fused lowering"); + } + if (fInputType != ETensorType::INT8 && fInputType != ETensorType::UINT8 && fInputType != ETensorType::INT16 && + fInputType != ETensorType::UINT16 && fInputType != ETensorType::INT32 && fInputType != ETensorType::UINT32) { + throw std::runtime_error("SOFIE ONNX DequantizeLinear literal code generation supports integer carriers only"); + } + const auto length = ConvertShapeToLength(fShape); + std::stringstream out; + out << "\n//------ ONNX DEQUANTIZELINEAR " << OpName << "\n"; + out << SP << "for (size_t id = 0; id < " << length << "; ++id) {\n"; + out << SP << SP << "tensor_" << fNY << "[id] = (static_cast(tensor_" << fNX << "[id]) - " + << fInfo.zeroPoint << ") * " << fInfo.scale << ";\n"; + out << SP << "}\n"; + return out.str(); + } +}; + +} // namespace SOFIE + +#endif // SOFIE_ROPERATOR_ONNXQUANTIZELINEAR diff --git a/core/inc/SOFIE/ROperator_QONNXQuant.hxx b/core/inc/SOFIE/ROperator_QONNXQuant.hxx new file mode 100644 index 0000000..f9c504f --- /dev/null +++ b/core/inc/SOFIE/ROperator_QONNXQuant.hxx @@ -0,0 +1,176 @@ +#ifndef SOFIE_ROPERATOR_QONNXQUANT +#define SOFIE_ROPERATOR_QONNXQUANT + +#include "SOFIE/RModel.hxx" +#include "SOFIE/ROperator.hxx" +#include "SOFIE/SOFIE_common.hxx" +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/RQuantization_Parameters.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +// Fallback for qonnx.custom_op.general::Quant. +// Operator implements fake quantization: values are rounded and clamped to +// the quantized integer grid, then converted back to the floating carrier type. +class ROperator_QONNXQuant final : public ROperator { +private: + std::string fNX; + std::string fNScale; + std::string fNZeroPoint; + std::string fNBitWidth; + std::string fNY; + bool fIsSigned = false; + bool fNarrow = false; + EQuantizationRoundingMode fRounding = EQuantizationRoundingMode::UNDEFINED; + EQuantizationOverflowMode fOverflow = EQuantizationOverflowMode::UNDEFINED; + std::vector fShape; + double fScale = 1.0; + std::int64_t fZeroPoint = 0; + unsigned fBitWidth = 0; + bool fHasVectorParameters = false; + double fQMin = 0.0; + double fQMax = 0.0; + + static std::vector GetFloatInitializer(RModel &model, const std::string &tensorName) + { + auto values = model.GetTensorData(tensorName); + if (values.empty()) { + throw std::runtime_error("SOFIE QONNX Quant expected non-empty FLOAT initializer " + tensorName); + } + return values; + } + + static float GetScalarFloat(RModel &model, const std::string &tensorName) + { + auto values = GetFloatInitializer(model, tensorName); + if (values.size() != 1) { + throw std::runtime_error("SOFIE QONNX Quant expected scalar FLOAT initializer " + tensorName); + } + return values.front(); + } + + std::string RoundingExpression(const std::string &value) const + { + switch (fRounding) { + case EQuantizationRoundingMode::ROUND: return "std::nearbyint(" + value + ")"; + case EQuantizationRoundingMode::FLOOR: return "std::floor(" + value + ")"; + case EQuantizationRoundingMode::TRUNCATE: return "std::trunc(" + value + ")"; + default: throw std::runtime_error("SOFIE QONNX Quant has unsupported rounding mode"); + } + } + +public: + ROperator_QONNXQuant() = default; + + ROperator_QONNXQuant(std::string nameX, std::string nameScale, std::string nameZeroPoint, + std::string nameBitWidth, std::string nameY, bool isSigned, bool narrow, + EQuantizationRoundingMode rounding, EQuantizationOverflowMode overflow) + : fNX(UTILITY::Clean_name(nameX)), fNScale(UTILITY::Clean_name(nameScale)), + fNZeroPoint(UTILITY::Clean_name(nameZeroPoint)), fNBitWidth(UTILITY::Clean_name(nameBitWidth)), + fNY(UTILITY::Clean_name(nameY)), fIsSigned(isSigned), fNarrow(narrow), fRounding(rounding), + fOverflow(overflow) + { + fInputTensorNames = { fNX, fNScale, fNZeroPoint, fNBitWidth }; + fOutputTensorNames = { fNY }; + } + + bool IsQuantizationBoundary() const override { return true; } + std::string GetQuantizationSourceTensor() const override { return fNX; } + + std::vector TypeInference(std::vector input) override + { + if (input.empty()) return {}; + return { input.front() }; + } + + std::vector> ShapeInference(std::vector> input) override + { + if (input.empty()) return {}; + return { input.front() }; + } + + void Initialize(RModel &model) override + { + if (!model.CheckIfTensorAlreadyExist(fNX)) { + throw std::runtime_error("SOFIE QONNX Quant input tensor " + fNX + " is not found in model"); + } + if (!model.IsInitializedTensor(fNScale) || !model.IsInitializedTensor(fNZeroPoint) || + !model.IsInitializedTensor(fNBitWidth)) { + throw std::runtime_error("SOFIE QONNX Quant scale, zero-point, and bit-width must be initialized tensors"); + } + + const auto scaleValues = GetFloatInitializer(model, fNScale); + const auto zeroPointValues = GetFloatInitializer(model, fNZeroPoint); + const double bitWidthFloat = static_cast(GetScalarFloat(model, fNBitWidth)); + + if (std::round(bitWidthFloat) != bitWidthFloat || bitWidthFloat <= 0.0 || bitWidthFloat > 32.0) { + throw std::runtime_error("SOFIE QONNX Quant bit-width must be an integer in [1, 32] for tensor " + fNY); + } + + fBitWidth = static_cast(bitWidthFloat); + fShape = model.GetTensorShape(fNX); + QuantizationParameterSpec spec; + spec.scales.assign(scaleValues.begin(), scaleValues.end()); + spec.zeroPoints = ValidateIntegralZeroPoints(zeroPointValues, "SOFIE QONNX Quant " + fNY); + spec.bitWidth = fBitWidth; + spec.isSigned = fIsSigned; + spec.narrow = fNarrow; + spec.rounding = fRounding; + spec.overflow = fOverflow; + spec.scaleTensor = fNScale; + spec.zeroPointTensor = fNZeroPoint; + spec.tensorShape = fShape; + spec.context = "SOFIE QONNX Quant " + fNY; + auto info = MakeValidatedQuantizationInfo(spec); + fScale = info.scale; + fZeroPoint = info.zeroPoint; + fHasVectorParameters = info.granularity == EQuantizationGranularity::PerChannel; + + auto [qMin, qMax] = QuantizedIntegerRange(info); + fQMin = static_cast(qMin); + fQMax = static_cast(qMax); + model.AddQuantizationInfo(fNY, std::move(info)); + + model.AddIntermediateTensor(fNY, model.GetTensorType(fNX), fShape); + model.AddNeededStdLib("cmath"); + } + + std::string Generate(std::string OpName) override + { + OpName = "op_" + OpName; + if (fBitWidth == 0) { + throw std::runtime_error("SOFIE QONNX Quant called to Generate without being initialized first"); + } + if (fHasVectorParameters) { + throw std::runtime_error("SOFIE QONNX Quant literal code generation supports scalar parameters only; vector parameters require a fused quantized lowering"); + } + const auto length = ConvertShapeToLength(fShape); + const std::string scaledValue = "((static_cast(tensor_" + fNX + "[id]) / " + std::to_string(fScale) + ") + " + std::to_string(fZeroPoint) + ")"; + + std::stringstream out; + out << "\n//------ QONNX QUANT FAKE-QUANT " << OpName << "\n"; + out << SP << "for (size_t id = 0; id < " << length << "; ++id) {\n"; + out << SP << SP << "double q = " << RoundingExpression(scaledValue) << ";\n"; + out << SP << SP << "q = (q < " << fQMin << ") ? " << fQMin << " : ((q > " << fQMax << ") ? " << fQMax << " : q);\n"; + out << SP << SP << "tensor_" << fNY << "[id] = (q - " << fZeroPoint << ") * " << fScale << ";\n"; + out << SP << "}\n"; + return out.str(); + } + + std::string Generate_GPU_ALPAKA(std::string /*OpName*/) override + { + throw std::runtime_error("SOFIE QONNX Quant Alpaka code generation is not available"); + } +}; + +} // namespace SOFIE + +#endif // SOFIE_ROPERATOR_QONNXQUANT diff --git a/core/inc/SOFIE/ROperator_QuantizedGemm.hxx b/core/inc/SOFIE/ROperator_QuantizedGemm.hxx new file mode 100644 index 0000000..3b04954 --- /dev/null +++ b/core/inc/SOFIE/ROperator_QuantizedGemm.hxx @@ -0,0 +1,420 @@ +#ifndef SOFIE_ROPERATOR_QUANTIZED_GEMM +#define SOFIE_ROPERATOR_QUANTIZED_GEMM + +#include "SOFIE/ROperator.hxx" +#include "SOFIE/SOFIE_common.hxx" +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/ROperator_QuantizedMatrix.hxx" +#include "SOFIE/SOFIE_Quantized.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +struct QuantizedGemmCodegenContext : QuantizedMatrixCodegenContext { + float alpha = 1.0f; + float beta = 1.0f; + std::int64_t transA = 0; + std::int64_t transB = 0; + EActivationType activation = EActivationType::UNDEFINED; +}; + +namespace INTERNAL { + +inline void ValidateQuantizedGemmContext(const QuantizedGemmCodegenContext &context, + const std::string &pathName) +{ + ValidateQuantizedMatrixContext(context, "Gemm", pathName); + if (context.transA != 0 || context.transB != 1) { + throw std::runtime_error("SOFIE " + pathName + " supports transA=0 and transB=1"); + } + if (context.alpha != 1.0f || context.beta != 1.0f) { + throw std::runtime_error("SOFIE " + pathName + " supports alpha=1 and beta=1"); + } + if (context.activation != EActivationType::UNDEFINED && context.activation != EActivationType::RELU) { + throw std::runtime_error("SOFIE " + pathName + " supports only no fused activation or fused ReLU"); + } +} + +inline bool HasQuantizedGemmBias(const QuantizedGemmRegion ®ion) +{ + return !region.biasSourceTensor.empty(); +} + +} // namespace INTERNAL + +inline std::string GenerateFusedQuantizedGemmCallCPUCode(std::string opName, const QuantizedGemmCodegenContext &context, + const QuantizedGemmRegion ®ion, + const QuantizedLoweringPlan &plan) +{ + opName = "op_" + opName; + INTERNAL::ValidateQuantizedGemmContext(context, "fused Quantized Gemm call CPU"); + + if (!QuantizedPlanUsesPrequantizedWeights(plan)) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path requires pre-quantized weight storage"); + } + if (plan.weightStorageTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path is missing a weight storage tensor"); + } + if (plan.weightLayout != EQuantizedLayout::PackedCPU) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path requires PackedCPU weight layout"); + } + if (plan.weightStorage != EQuantizedStorageType::Int8 && plan.weightStorage != EQuantizedStorageType::UInt8) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path requires int8/uint8 weight storage"); + } + if (region.inputSourceTensor.empty() || region.outputTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path is missing quantization source/output tensors"); + } + if (INTERNAL::HasQuantizedGemmBias(region) && !region.biasQuant.has_value()) { + throw std::runtime_error("SOFIE fused Quantized Gemm call CPU path is missing quantized bias metadata"); + } + + const auto dimA = context.inputShape.size(); + const auto dimB = context.weightShape.size(); + const auto m = context.inputShape[dimA - 2].GetVal(); + const auto k = context.inputShape[dimA - 1].GetVal(); + const auto n = context.weightShape[dimB - 2].GetVal(); + const std::string SP = " "; + constexpr std::size_t tileN = 4; + + const auto inputRange = QuantizedIntegerRange(region.inputQuant); + const auto outputRange = QuantizedIntegerRange(region.outputQuant); + + std::stringstream out; + out << "\n//--------- Fused Quantized Gemm call CPU " << opName << " " + << ConvertDimShapeToString(context.inputShape) << " * " << ConvertDimShapeToString(context.weightShape) + << " -> " << ConvertDimShapeToString(context.outputShape) << "\n"; + out << SP << "// Call-boundary path: generated code delegates quantized GEMM to SOFIE::QuantizedGemm_Call(...).\n"; + out << SP << "// The runtime call consumes quantized parameters and packed constant weights.\n"; + + out << SP << "SOFIE::QuantizedGemmParams " << opName << "_params{};\n"; + out << SP << opName << "_params.m = static_cast(" << m << ");\n"; + out << SP << opName << "_params.n = static_cast(" << n << ");\n"; + out << SP << opName << "_params.k = static_cast(" << k << ");\n"; + out << SP << opName << "_params.tileN = " << tileN << ";\n"; + + out << SP << opName << "_params.scaleX = " + << std::setprecision(std::numeric_limits::max_digits10) << region.inputQuant.scale << ";\n"; + out << SP << opName << "_params.scaleW = " + << std::setprecision(std::numeric_limits::max_digits10) << region.weightQuant.scale << ";\n"; + out << SP << opName << "_params.scaleY = " + << std::setprecision(std::numeric_limits::max_digits10) << region.outputQuant.scale << ";\n"; + out << SP << opName << "_params.zeroX = " << region.inputQuant.zeroPoint << ";\n"; + out << SP << opName << "_params.zeroW = " << region.weightQuant.zeroPoint << ";\n"; + out << SP << opName << "_params.zeroY = " << region.outputQuant.zeroPoint << ";\n"; + out << SP << opName << "_params.qminX = " << inputRange.first << ";\n"; + out << SP << opName << "_params.qmaxX = " << inputRange.second << ";\n"; + out << SP << opName << "_params.qminY = " << outputRange.first << ";\n"; + out << SP << opName << "_params.qmaxY = " << outputRange.second << ";\n"; + + const auto accumulatorScale = region.inputQuant.scale * region.weightQuant.scale; + const auto requantScale = accumulatorScale / region.outputQuant.scale; + std::int64_t requantMultiplier = 0; + int requantShift = 0; + const bool hasFixedPointRequantization = + MakeQuantizedFixedPointMultiplier(requantScale, requantMultiplier, requantShift); + + const bool hasBiasTensor = INTERNAL::HasQuantizedGemmBias(region); + const bool hasAccumulatorBias = + !hasBiasTensor || + (region.biasQuant.has_value() && NearlyEqualQuantizedScale(region.biasQuant->scale, accumulatorScale)); + std::int64_t biasRequantMultiplier = 0; + int biasRequantShift = 0; + const bool hasExactIntegerBias = + !hasBiasTensor || + (region.biasQuant.has_value() && MakeExactIntegerScaleMultiplier( + region.biasQuant->scale / region.outputQuant.scale, biasRequantMultiplier, + biasRequantShift)); + const bool useIntegerEpilogue = hasFixedPointRequantization && (!hasBiasTensor || hasAccumulatorBias || hasExactIntegerBias); + if (useIntegerEpilogue) { + out << SP << opName << "_params.useIntegerEpilogue = true;\n"; + out << SP << opName << "_params.requantMultiplier = static_cast(" << requantMultiplier << ");\n"; + out << SP << opName << "_params.requantShift = " << requantShift << ";\n"; + if (hasBiasTensor && hasAccumulatorBias) { + out << SP << opName << "_params.useAccumulatorBias = true;\n"; + } else if (hasBiasTensor && hasExactIntegerBias) { + out << SP << opName << "_params.useIntegerBias = true;\n"; + out << SP << opName << "_params.biasRequantMultiplier = static_cast(" << biasRequantMultiplier << ");\n"; + out << SP << opName << "_params.biasRequantShift = " << biasRequantShift << ";\n"; + } + out << SP << "// Integer epilogue: fixed-point requantization"; + if (hasBiasTensor && hasAccumulatorBias) { + out << " with accumulator-domain bias"; + } else if (hasBiasTensor && hasExactIntegerBias) { + out << " with exact integer bias contribution"; + } + out << ".\n"; + } else { + out << SP << "// Integer epilogue disabled: QuantizedGemm_Call uses the exact float fallback.\n"; + } + + if (context.activation == EActivationType::RELU) { + out << SP << opName << "_params.activation = SOFIE::EQuantizedGemmActivation::Relu;\n"; + } else { + out << SP << opName << "_params.activation = SOFIE::EQuantizedGemmActivation::None;\n"; + } + + if (INTERNAL::HasQuantizedGemmBias(region)) { + const auto biasRange = QuantizedIntegerRange(*region.biasQuant); + out << SP << opName << "_params.hasBias = true;\n"; + out << SP << opName << "_params.scaleB = " + << std::setprecision(std::numeric_limits::max_digits10) << region.biasQuant->scale << ";\n"; + out << SP << opName << "_params.zeroB = " << region.biasQuant->zeroPoint << ";\n"; + out << SP << opName << "_params.qminB = " << biasRange.first << ";\n"; + out << SP << opName << "_params.qmaxB = " << biasRange.second << ";\n"; + } else { + out << SP << opName << "_params.hasBias = false;\n"; + } + + out << SP << "SOFIE::QuantizedGemm_Call(tensor_" << region.outputTensor << ", tensor_" + << region.inputSourceTensor << ", tensor_" << plan.weightStorageTensor << ", "; + if (INTERNAL::HasQuantizedGemmBias(region)) { + out << "tensor_" << region.biasSourceTensor; + } else { + out << "nullptr"; + } + out << ", " << opName << "_params);\n"; + + return out.str(); +} + + +inline std::string GenerateFusedQuantizedGemmAlpakaFakeQuantKernel(std::string opName, + const QuantizedGemmCodegenContext &context) +{ + INTERNAL::ValidateQuantizedGemmContext(context, "fused Quantized Gemm Alpaka fake-quant kernel"); + + std::stringstream out; + out << "\n//--------- ROperator_QuantizedGemm Alpaka fake-quant kernel " << opName << "\n"; + out << " struct QuantizedGemmAlpakaFakeQuantKernel_" << opName << " {\n"; + out << " template\n"; + out << " ALPAKA_FN_ACC void operator()(TAcc const & acc, const float * input, const float * weight, const float * bias, float * output,\n"; + out << " std::size_t elements, std::size_t n, std::size_t k, bool hasBias, bool hasRelu,\n"; + out << " double scaleX, double scaleW, double scaleB, double scaleY,\n"; + out << " std::int32_t zeroX, std::int32_t zeroW, std::int32_t zeroB, std::int32_t zeroY,\n"; + out << " std::int32_t qminX, std::int32_t qmaxX, std::int32_t qminW, std::int32_t qmaxW,\n"; + out << " std::int32_t qminB, std::int32_t qmaxB, std::int32_t qminY, std::int32_t qmaxY) const {\n"; + out << " const auto idx = alpaka::getIdx(acc)[0];\n"; + out << " if (idx >= elements) return;\n"; + out << " const std::size_t row = static_cast(idx) / n;\n"; + out << " const std::size_t col = static_cast(idx) % n;\n"; + out << " std::int32_t dot = 0;\n"; + out << " for (std::size_t kk = 0; kk < k; ++kk) {\n"; + out << " auto xReal = static_cast(input[row * k + kk]) / scaleX + static_cast(zeroX);\n"; + out << " auto wReal = static_cast(weight[col * k + kk]) / scaleW + static_cast(zeroW);\n"; + out << " auto xq = static_cast((xReal >= 0.0) ? (xReal + 0.5) : (xReal - 0.5));\n"; + out << " auto wq = static_cast((wReal >= 0.0) ? (wReal + 0.5) : (wReal - 0.5));\n"; + out << " xq = (xq < qminX) ? qminX : ((xq > qmaxX) ? qmaxX : xq);\n"; + out << " wq = (wq < qminW) ? qminW : ((wq > qmaxW) ? qmaxW : wq);\n"; + out << " dot += (xq - zeroX) * (wq - zeroW);\n"; + out << " }\n"; + out << " double real = static_cast(dot) * scaleX * scaleW;\n"; + out << " if (hasBias) {\n"; + out << " auto bReal = static_cast(bias[col]) / scaleB + static_cast(zeroB);\n"; + out << " auto bq = static_cast((bReal >= 0.0) ? (bReal + 0.5) : (bReal - 0.5));\n"; + out << " bq = (bq < qminB) ? qminB : ((bq > qmaxB) ? qmaxB : bq);\n"; + out << " real += static_cast(bq - zeroB) * scaleB;\n"; + out << " }\n"; + out << " auto yReal = real / scaleY + static_cast(zeroY);\n"; + out << " auto yq = static_cast((yReal >= 0.0) ? (yReal + 0.5) : (yReal - 0.5));\n"; + out << " if (hasRelu && yq < zeroY) yq = zeroY;\n"; + out << " yq = (yq < qminY) ? qminY : ((yq > qmaxY) ? qmaxY : yq);\n"; + out << " output[row * n + col] = static_cast(static_cast(yq - zeroY) * scaleY);\n"; + out << " }\n"; + out << " };\n"; + return out.str(); +} + +inline std::string GenerateFusedQuantizedGemmAlpakaFakeQuantDefinition(std::string opName, + const QuantizedGemmCodegenContext &context) +{ + INTERNAL::ValidateQuantizedGemmContext(context, "fused Quantized Gemm Alpaka fake-quant definition"); + return " QuantizedGemmAlpakaFakeQuantKernel_" + opName + " quantizedGemmAlpakaFakeQuantKernel_" + opName + ";\n"; +} + +inline std::string GenerateFusedQuantizedGemmCublasLtCoreLaunch(std::string opName, + const QuantizedGemmCodegenContext &context, + const QuantizedGemmRegion ®ion, + const QuantizedLoweringPlan &plan) +{ + INTERNAL::ValidateQuantizedGemmContext(context, "fused Quantized Gemm cuBLASLt core launch"); + if (plan.backend != EQuantizedBackend::ALPAKA || plan.status != EQuantizedLoweringStatus::Optimized || + !QuantizedPlanUsesPrequantizedWeights(plan) || plan.weightLayout != EQuantizedLayout::PlainDevice || + plan.weightStorageTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm cuBLASLt core launch requires an optimized Alpaka PlainDevice plan"); + } + if (region.inputSourceTensor.empty() || region.outputTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm cuBLASLt core launch is missing input/output tensors"); + } + if (plan.weightScaleMode == EQuantizedParameterMode::PerOutputChannel && plan.weightScaleTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm cuBLASLt per-channel launch is missing a weight scale tensor"); + } + + const auto dimA = context.inputShape.size(); + const auto dimB = context.weightShape.size(); + + INTERNAL::QuantizedCudaLtMatMulCall call; + call.boundaryName = "ROperator_QuantizedGemm cuBLASLt int8 GEMM core boundary " + opName; + call.stateName = "quantizedGemmCudaLtState_" + opName; + call.paramsName = "params_quantizedGemm_" + opName; + call.outputTensor = region.outputTensor; + call.inputTensor = region.inputSourceTensor; + call.weightStorageTensor = plan.weightStorageTensor; + call.biasTensor = region.biasSourceTensor; + call.weightScaleTensor = plan.weightScaleTensor; + call.m = context.inputShape[dimA - 2].GetVal(); + call.k = context.inputShape[dimA - 1].GetVal(); + call.n = context.weightShape[dimB - 2].GetVal(); + call.inputQuant = region.inputQuant; + call.weightQuant = region.weightQuant; + call.biasQuant = region.biasQuant; + call.outputQuant = region.outputQuant; + call.outputMode = plan.outputMode; + call.inputCarrierMode = plan.inputCarrierMode; + call.weightScaleMode = plan.weightScaleMode; + call.shapePolicy = plan.shapePolicy; + call.capabilityTag = plan.capabilityTag; + call.reason = plan.reason; + call.hasBias = INTERNAL::HasQuantizedGemmBias(region); + call.hasRelu = context.activation == EActivationType::RELU; + call.weightIsSigned = region.weightQuant.isSigned; + return INTERNAL::GenerateQuantizedCudaLtMatMulCall(call); +} + +inline std::string GenerateFusedQuantizedGemmAlpakaFakeQuantLaunch(std::string opName, + const QuantizedGemmCodegenContext &context, + const QuantizedGemmRegion ®ion, + const QuantizedLoweringPlan &plan) +{ + INTERNAL::ValidateQuantizedGemmContext(context, "fused Quantized Gemm Alpaka fake-quant launch"); + if (plan.backend != EQuantizedBackend::ALPAKA || !IsQuantizedLoweringAvailable(plan.status)) { + throw std::runtime_error("SOFIE fused Quantized Gemm Alpaka fake-quant launch requires an available Alpaka plan"); + } + if (region.inputSourceTensor.empty() || region.weightSourceTensor.empty() || region.outputTensor.empty()) { + throw std::runtime_error("SOFIE fused Quantized Gemm Alpaka fake-quant launch is missing quantization source/output tensors"); + } + if (INTERNAL::HasQuantizedGemmBias(region) && !region.biasQuant.has_value()) { + throw std::runtime_error("SOFIE fused Quantized Gemm Alpaka fake-quant launch is missing quantized bias metadata"); + } + + const auto dimA = context.inputShape.size(); + const auto dimB = context.weightShape.size(); + const auto m = context.inputShape[dimA - 2].GetVal(); + const auto k = context.inputShape[dimA - 1].GetVal(); + const auto n = context.weightShape[dimB - 2].GetVal(); + const auto elements = m + " * " + n; + const auto inputRange = QuantizedIntegerRange(region.inputQuant); + const auto weightRange = QuantizedIntegerRange(region.weightQuant); + const auto outputRange = QuantizedIntegerRange(region.outputQuant); + const auto biasRange = region.biasQuant ? QuantizedIntegerRange(*region.biasQuant) : std::pair{0, 0}; + + std::stringstream out; + out << "\n//--------- ROperator_QuantizedGemm Alpaka fake-quant launch " << opName << "\n"; + out << " // Fake-quant GPU path: one Alpaka thread computes one quantized GEMM output element.\n"; + out << " auto const elementsPerThread_quantizedGemm_" << opName << " = Vec::all(static_cast(1));\n"; + out << " auto const elementsPerGrid_quantizedGemm_" << opName << " = Vec::all(Idx{" << elements << "});\n"; + out << " auto const workDiv_quantizedGemm_" << opName << " = sofie_workdiv(elementsPerGrid_quantizedGemm_" << opName << ");\n"; + out << " auto task_op_" << opName << " = alpaka::createTaskKernel(workDiv_quantizedGemm_" << opName + << ", quantizedGemmAlpakaFakeQuantKernel_" << opName + << ", alpaka::getPtrNative(deviceBuf_" << region.inputSourceTensor << ")" + << ", alpaka::getPtrNative(deviceBuf_" << region.weightSourceTensor << ")"; + if (INTERNAL::HasQuantizedGemmBias(region)) { + out << ", alpaka::getPtrNative(deviceBuf_" << region.biasSourceTensor << ")"; + } else { + out << ", static_cast(nullptr)"; + } + out << ", alpaka::getPtrNative(deviceBuf_" << region.outputTensor << ")" + << ", static_cast(" << elements << ")" + << ", static_cast(" << n << ")" + << ", static_cast(" << k << ")" + << ", " << (INTERNAL::HasQuantizedGemmBias(region) ? "true" : "false") + << ", " << (context.activation == EActivationType::RELU ? "true" : "false") + << ", " << std::setprecision(std::numeric_limits::max_digits10) << region.inputQuant.scale + << ", " << std::setprecision(std::numeric_limits::max_digits10) << region.weightQuant.scale + << ", " << std::setprecision(std::numeric_limits::max_digits10) << (region.biasQuant ? region.biasQuant->scale : 1.0) + << ", " << std::setprecision(std::numeric_limits::max_digits10) << region.outputQuant.scale + << ", static_cast(" << region.inputQuant.zeroPoint << ")" + << ", static_cast(" << region.weightQuant.zeroPoint << ")" + << ", static_cast(" << (region.biasQuant ? region.biasQuant->zeroPoint : 0) << ")" + << ", static_cast(" << region.outputQuant.zeroPoint << ")" + << ", static_cast(" << inputRange.first << ")" + << ", static_cast(" << inputRange.second << ")" + << ", static_cast(" << weightRange.first << ")" + << ", static_cast(" << weightRange.second << ")" + << ", static_cast(" << biasRange.first << ")" + << ", static_cast(" << biasRange.second << ")" + << ", static_cast(" << outputRange.first << ")" + << ", static_cast(" << outputRange.second << "));\n"; + out << " alpaka::enqueue(queue, task_op_" << opName << ");\n"; + return out.str(); +} + +class ROperator_QuantizedGemm final : public ROperator { +private: + QuantizedGemmRegion fRegion; + QuantizedLoweringPlan fPlan; + QuantizedGemmCodegenContext fContext; + +public: + ROperator_QuantizedGemm(QuantizedGemmRegion region, QuantizedLoweringPlan plan, + QuantizedGemmCodegenContext context) + : fRegion(std::move(region)), fPlan(std::move(plan)), fContext(std::move(context)) + { + fKind = OperatorKind::QUANTIZED_GEMM; + fName = "QuantizedGemm"; + fInputTensorNames = { fRegion.inputSourceTensor, fRegion.weightSourceTensor }; + if (!fRegion.biasSourceTensor.empty()) { + fInputTensorNames.emplace_back(fRegion.biasSourceTensor); + } + fOutputTensorNames = { fRegion.outputTensor }; + } + + std::vector GetStdLibs() override { return { "cmath", "cstdint", "vector" }; } + + void Initialize(RModel &) override {} + + std::string Generate(std::string opName) override + { + if (fPlan.backend != EQuantizedBackend::CPU || !IsQuantizedLoweringAvailable(fPlan.status) || + !fPlan.suppressesGraphOperators || !QuantizedPlanUsesPrequantizedWeights(fPlan) || + fPlan.weightLayout != EQuantizedLayout::PackedCPU) { + throw std::runtime_error("SOFIE ROperator_QuantizedGemm CPU code generation requires an available packed-weight CPU lowering plan"); + } + std::string code = "\n//--------- ROperator_QuantizedGemm synthetic fused CPU operator " + opName + "\n"; + return code + GenerateFusedQuantizedGemmCallCPUCode(std::move(opName), fContext, fRegion, fPlan); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string opName) override + { + if (IsOptimizedQuantizedAlpakaPlainDevicePlan(fPlan)) + return ""; + return GenerateFusedQuantizedGemmAlpakaFakeQuantKernel(std::move(opName), fContext); + } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override + { + if (IsOptimizedQuantizedAlpakaPlainDevicePlan(fPlan)) + return " SOFIE::QuantizedGemmCudaLtState quantizedGemmCudaLtState_" + opName + "; // owns cuBLASLt state and CUDA temporaries\n"; + return GenerateFusedQuantizedGemmAlpakaFakeQuantDefinition(std::move(opName), fContext); + } + + std::string Generate_GPU_ALPAKA(std::string opName) override + { + if (IsOptimizedQuantizedAlpakaPlainDevicePlan(fPlan)) + return GenerateFusedQuantizedGemmCublasLtCoreLaunch(std::move(opName), fContext, fRegion, fPlan); + return GenerateFusedQuantizedGemmAlpakaFakeQuantLaunch(std::move(opName), fContext, fRegion, fPlan); + } +}; + +} // namespace SOFIE + +#endif // SOFIE_ROPERATOR_QUANTIZED_GEMM diff --git a/core/inc/SOFIE/ROperator_QuantizedMatMul.hxx b/core/inc/SOFIE/ROperator_QuantizedMatMul.hxx new file mode 100644 index 0000000..80c5a1a --- /dev/null +++ b/core/inc/SOFIE/ROperator_QuantizedMatMul.hxx @@ -0,0 +1,120 @@ +#ifndef SOFIE_ROPERATOR_QUANTIZED_MATMUL +#define SOFIE_ROPERATOR_QUANTIZED_MATMUL + +#include "SOFIE/ROperator.hxx" +#include "SOFIE/SOFIE_common.hxx" +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/ROperator_QuantizedMatrix.hxx" + +#include +#include +#include +#include +#include + +namespace SOFIE { + +namespace INTERNAL { + +inline void ValidateQuantizedMatMulContext(const QuantizedMatrixCodegenContext &context, + const QuantizedMatMulRegion ®ion, + const QuantizedLoweringPlan &plan, + const std::string &pathName) +{ + ValidateQuantizedMatrixContext(context, "MatMul", pathName); + if (context.inputShape[1].GetVal() != context.weightShape[0].GetVal() || + context.inputShape[0].GetVal() != context.outputShape[0].GetVal() || + context.weightShape[1].GetVal() != context.outputShape[1].GetVal()) { + throw std::runtime_error("SOFIE " + pathName + " requires X[M,K] @ W[K,N] -> Y[M,N]"); + } + if (region.inputSourceTensor.empty() || region.outputTensor.empty()) { + throw std::runtime_error("SOFIE " + pathName + " is missing input/output tensors"); + } + ValidateQuantizedCudaLtMatrixPlan(plan, pathName); +} + +} // namespace INTERNAL + +inline std::string GenerateFusedQuantizedMatMulCublasLtLaunch(std::string opName, + const QuantizedMatrixCodegenContext &context, + const QuantizedMatMulRegion ®ion, + const QuantizedLoweringPlan &plan) +{ + INTERNAL::ValidateQuantizedMatMulContext(context, region, plan, "fused Quantized MatMul cuBLASLt launch"); + + INTERNAL::QuantizedCudaLtMatMulCall call; + call.boundaryName = "ROperator_QuantizedMatMul cuBLASLt int8 MatMul boundary " + opName; + call.stateName = "quantizedMatMulCudaLtState_" + opName; + call.paramsName = "params_quantizedMatMul_" + opName; + call.outputTensor = region.outputTensor; + call.inputTensor = region.inputSourceTensor; + call.weightStorageTensor = plan.weightStorageTensor; + call.biasTensor = region.epilogue.biasSourceTensor; + call.weightScaleTensor = plan.weightScaleTensor; + call.m = context.inputShape[0].GetVal(); + call.k = context.inputShape[1].GetVal(); + call.n = context.weightShape[1].GetVal(); + call.inputQuant = region.inputQuant; + call.weightQuant = region.weightQuant; + call.biasQuant = region.epilogue.biasQuant; + call.outputQuant = region.outputQuant; + call.outputMode = plan.outputMode; + call.inputCarrierMode = plan.inputCarrierMode; + call.weightScaleMode = plan.weightScaleMode; + call.shapePolicy = plan.shapePolicy; + call.capabilityTag = plan.capabilityTag; + call.reason = plan.reason; + call.hasBias = QuantizedEpilogueHasBias(region.epilogue.kind); + call.hasRelu = QuantizedEpilogueHasRelu(region.epilogue.kind); + call.weightIsSigned = region.weightQuant.isSigned; + return INTERNAL::GenerateQuantizedCudaLtMatMulCall(call); +} + +class ROperator_QuantizedMatMul final : public ROperator { +private: + QuantizedMatMulRegion fRegion; + QuantizedLoweringPlan fPlan; + QuantizedMatrixCodegenContext fContext; + +public: + ROperator_QuantizedMatMul(QuantizedMatMulRegion region, QuantizedLoweringPlan plan, + QuantizedMatrixCodegenContext context) + : fRegion(std::move(region)), fPlan(std::move(plan)), fContext(std::move(context)) + { + fKind = OperatorKind::QUANTIZED_MATMUL; + fName = "QuantizedMatMul"; + fInputTensorNames = { fRegion.inputSourceTensor, fRegion.weightSourceTensor }; + if (QuantizedEpilogueHasBias(fRegion.epilogue.kind)) { + fInputTensorNames.emplace_back(fRegion.epilogue.biasSourceTensor); + } + fOutputTensorNames = { fRegion.outputTensor }; + } + + std::vector GetStdLibs() override { return { "cstdint", "vector" }; } + + void Initialize(RModel &) override {} + + std::string Generate(std::string) override + { + throw std::runtime_error("SOFIE ROperator_QuantizedMatMul CPU code generation is not implemented"); + } + + std::string Generate_GPU_Kernel_ALPAKA(std::string) override { return ""; } + + std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string opName) override + { + if (!IsOptimizedQuantizedAlpakaPlainDevicePlan(fPlan)) { + throw std::runtime_error("SOFIE ROperator_QuantizedMatMul Alpaka code generation requires an optimized PlainDevice plan"); + } + return " SOFIE::QuantizedGemmCudaLtState quantizedMatMulCudaLtState_" + opName + "; // owns cuBLASLt state and CUDA temporaries\n"; + } + + std::string Generate_GPU_ALPAKA(std::string opName) override + { + return GenerateFusedQuantizedMatMulCublasLtLaunch(std::move(opName), fContext, fRegion, fPlan); + } +}; + +} // namespace SOFIE + +#endif // SOFIE_ROPERATOR_QUANTIZED_MATMUL diff --git a/core/inc/SOFIE/ROperator_QuantizedMatrix.hxx b/core/inc/SOFIE/ROperator_QuantizedMatrix.hxx new file mode 100644 index 0000000..fb4a080 --- /dev/null +++ b/core/inc/SOFIE/ROperator_QuantizedMatrix.hxx @@ -0,0 +1,202 @@ +#ifndef SOFIE_ROPERATOR_QUANTIZED_MATRIX +#define SOFIE_ROPERATOR_QUANTIZED_MATRIX + +#include "SOFIE/RQuantization.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +struct QuantizedMatrixCodegenContext { + std::vector inputShape; + std::vector weightShape; + std::vector outputShape; +}; + +namespace INTERNAL { + +inline void ValidateQuantizedMatrixContext(const QuantizedMatrixCodegenContext &context, + const std::string &operatorName, + const std::string &pathName) +{ + if (context.inputShape.empty() || context.weightShape.empty() || context.outputShape.empty()) { + throw std::runtime_error("SOFIE " + pathName + " called before " + operatorName + " initialization"); + } + if (context.inputShape.size() != 2 || context.weightShape.size() != 2 || context.outputShape.size() != 2) { + throw std::runtime_error("SOFIE " + pathName + " supports rank-2 " + operatorName + " only"); + } +} + +inline void ValidateQuantizedCudaLtMatrixPlan(const QuantizedLoweringPlan &plan, + const std::string &pathName) +{ + if (plan.backend != EQuantizedBackend::ALPAKA || plan.status != EQuantizedLoweringStatus::Optimized || + !QuantizedPlanUsesPrequantizedWeights(plan) || plan.weightLayout != EQuantizedLayout::PlainDevice || + plan.weightStorageTensor.empty()) { + throw std::runtime_error("SOFIE " + pathName + " requires an optimized Alpaka PlainDevice lowering plan"); + } + if (plan.weightStorage != EQuantizedStorageType::Int8) { + throw std::runtime_error("SOFIE " + pathName + " currently supports signed int8 weight storage only"); + } +} + +inline const char *QuantizedCudaInputCarrierName(EQuantizedCarrierMode mode, const std::string &pathName) +{ + switch (mode) { + case EQuantizedCarrierMode::Int8: + return "Int8"; + case EQuantizedCarrierMode::Float: + return "Float"; + case EQuantizedCarrierMode::UInt8: + throw std::runtime_error("SOFIE " + pathName + " currently supports signed int8 input carriers only"); + default: + throw std::runtime_error("SOFIE " + pathName + " received unsupported input carrier mode"); + } +} + +inline const char *QuantizedCudaEpilogueModeName(EQuantizedOutputMode mode, const std::string &pathName) +{ + switch (mode) { + case EQuantizedOutputMode::ExactFakeQuantFloat: + return "ExactFakeQuant"; + case EQuantizedOutputMode::Quantized: + return "Quantized"; + default: + throw std::runtime_error("SOFIE " + pathName + " received unsupported output mode"); + } +} + +inline const char *QuantizedCudaOutputCarrierName(const QuantizationInfo &info, const std::string &pathName) +{ + if (info.bitWidth != 8) { + throw std::runtime_error("SOFIE " + pathName + " currently supports only 8-bit output carriers"); + } + return info.isSigned ? "Int8" : "UInt8"; +} + +struct QuantizedCudaLtMatMulCall { + std::string boundaryName; + std::string stateName; + std::string paramsName; + std::string outputTensor; + std::string inputTensor; + std::string weightStorageTensor; + std::string biasTensor; + std::string weightScaleTensor; + std::string m; + std::string n; + std::string k; + QuantizedDenseLinearShapePolicy shapePolicy; + QuantizationInfo inputQuant; + QuantizationInfo weightQuant; + std::optional biasQuant; + QuantizationInfo outputQuant; + EQuantizedOutputMode outputMode = EQuantizedOutputMode::UNDEFINED; + EQuantizedCarrierMode inputCarrierMode = EQuantizedCarrierMode::UNDEFINED; + EQuantizedParameterMode weightScaleMode = EQuantizedParameterMode::Scalar; + std::string capabilityTag; + std::string reason; + bool hasBias = false; + bool hasRelu = false; + bool weightIsSigned = true; +}; + +inline std::string GenerateQuantizedCudaLtMatMulCall(const QuantizedCudaLtMatMulCall &call) +{ + if (call.outputTensor.empty() || call.inputTensor.empty() || call.weightStorageTensor.empty()) { + throw std::runtime_error("SOFIE " + call.boundaryName + " is missing input/output/weight-storage tensors"); + } + if (call.hasBias && (!call.biasQuant.has_value() || call.biasTensor.empty())) { + throw std::runtime_error("SOFIE " + call.boundaryName + " is missing bias tensor or quantization metadata"); + } + if (call.weightScaleMode == EQuantizedParameterMode::PerOutputChannel && call.weightScaleTensor.empty()) { + throw std::runtime_error("SOFIE " + call.boundaryName + " per-channel launch is missing a weight scale tensor"); + } + + const auto inputRange = QuantizedIntegerRange(call.inputQuant); + const auto biasRange = call.biasQuant ? QuantizedIntegerRange(*call.biasQuant) : std::pair{0, 0}; + const auto outputRange = QuantizedIntegerRange(call.outputQuant); + const bool paddedExecution = call.shapePolicy.policy == EQuantizedShapePolicy::Padded; + + std::stringstream out; + out << "\n//--------- " << call.boundaryName << "\n"; + out << " // Optimized GPU boundary: stream-ordered cuBLASLt int8 matrix multiply selected by the lowering plan.\n"; + out << " {\n"; + out << " // Quantized lowering capability: " << call.capabilityTag << "\n"; + out << " // Quantized lowering reason: " << call.reason << "\n"; + out << " SOFIE::QuantizedGemmCudaLtParams " << call.paramsName << "{};\n"; + out << " " << call.paramsName << ".logicalM = static_cast(" << call.m << ");\n"; + out << " " << call.paramsName << ".logicalN = static_cast(" << call.n << ");\n"; + out << " " << call.paramsName << ".logicalK = static_cast(" << call.k << ");\n"; + if (paddedExecution) { + out << " " << call.paramsName << ".m = " << call.shapePolicy.physicalM << ";\n"; + out << " " << call.paramsName << ".n = " << call.shapePolicy.physicalN << ";\n"; + out << " " << call.paramsName << ".k = " << call.shapePolicy.physicalK << ";\n"; + out << " " << call.paramsName << ".paddedExecution = true;\n"; + } else { + out << " " << call.paramsName << ".m = " << call.paramsName << ".logicalM;\n"; + out << " " << call.paramsName << ".n = " << call.paramsName << ".logicalN;\n"; + out << " " << call.paramsName << ".k = " << call.paramsName << ".logicalK;\n"; + } + out << std::setprecision(std::numeric_limits::max_digits10); + out << " " << call.paramsName << ".inputScale = " << call.inputQuant.scale << ";\n"; + out << " " << call.paramsName << ".weightScale = " << call.weightQuant.scale << ";\n"; + out << " " << call.paramsName << ".biasScale = " << (call.biasQuant ? call.biasQuant->scale : 1.0) << ";\n"; + out << " " << call.paramsName << ".outputScale = " << call.outputQuant.scale << ";\n"; + out << " " << call.paramsName << ".inputZeroPoint = " << call.inputQuant.zeroPoint << ";\n"; + out << " " << call.paramsName << ".weightZeroPoint = " << call.weightQuant.zeroPoint << ";\n"; + out << " " << call.paramsName << ".biasZeroPoint = " << (call.biasQuant ? call.biasQuant->zeroPoint : 0) << ";\n"; + out << " " << call.paramsName << ".outputZeroPoint = " << call.outputQuant.zeroPoint << ";\n"; + out << " " << call.paramsName << ".inputQMin = static_cast(" << inputRange.first << ");\n"; + out << " " << call.paramsName << ".inputQMax = static_cast(" << inputRange.second << ");\n"; + out << " " << call.paramsName << ".biasQMin = static_cast(" << biasRange.first << ");\n"; + out << " " << call.paramsName << ".biasQMax = static_cast(" << biasRange.second << ");\n"; + out << " " << call.paramsName << ".outputQMin = static_cast(" << outputRange.first << ");\n"; + out << " " << call.paramsName << ".outputQMax = static_cast(" << outputRange.second << ");\n"; + out << " " << call.paramsName << ".hasBias = " << (call.hasBias ? "true" : "false") << ";\n"; + out << " " << call.paramsName << ".hasRelu = " << (call.hasRelu ? "true" : "false") << ";\n"; + out << " " << call.paramsName << ".maxWorkspaceBytes = 32ULL * 1024ULL * 1024ULL;\n"; + out << " " << call.paramsName << ".epilogueMode = SOFIE::EQuantizedCudaEpilogueMode::" + << QuantizedCudaEpilogueModeName(call.outputMode, call.boundaryName) << ";\n"; + out << " " << call.paramsName << ".inputCarrier = SOFIE::EQuantizedCudaInputCarrier::" + << QuantizedCudaInputCarrierName(call.inputCarrierMode, call.boundaryName) << ";\n"; + out << " " << call.paramsName << ".outputCarrier = SOFIE::EQuantizedCudaOutputCarrier::" + << QuantizedCudaOutputCarrierName(call.outputQuant, call.boundaryName) << ";\n"; + out << " " << call.paramsName << ".weightType = SOFIE::EQuantizedCudaWeightType::" + << (call.weightIsSigned ? "Int8" : "UInt8") << ";\n"; + out << " " << call.paramsName << ".weightScaleMode = SOFIE::EQuantizedCudaScaleMode::" + << (call.weightScaleMode == EQuantizedParameterMode::PerOutputChannel ? "PerOutputChannel" : "PerTensor") << ";\n"; + + out << " SOFIE::QuantizedGemmCudaLt_Call(" << call.stateName + << ", alpaka::getNativeHandle(queue)" + << ", alpaka::getPtrNative(deviceBuf_" << call.outputTensor << ")" + << ", alpaka::getPtrNative(deviceBuf_" << call.inputTensor << ")" + << ", alpaka::getPtrNative(deviceBuf_" << call.weightStorageTensor << ")"; + if (call.hasBias) { + out << ", alpaka::getPtrNative(deviceBuf_" << call.biasTensor << ")"; + } else { + out << ", static_cast(nullptr)"; + } + if (call.weightScaleMode == EQuantizedParameterMode::PerOutputChannel) { + out << ", alpaka::getPtrNative(deviceBuf_" << call.weightScaleTensor << ")"; + } else { + out << ", static_cast(nullptr)"; + } + out << ", " << call.paramsName << ");\n"; + out << " }\n"; + return out.str(); +} + +} // namespace INTERNAL + +} // namespace SOFIE + +#endif // SOFIE_ROPERATOR_QUANTIZED_MATRIX diff --git a/core/inc/SOFIE/RQuantization.hxx b/core/inc/SOFIE/RQuantization.hxx new file mode 100644 index 0000000..ce07e2e --- /dev/null +++ b/core/inc/SOFIE/RQuantization.hxx @@ -0,0 +1,465 @@ +#ifndef SOFIE_RQUANTIZATION +#define SOFIE_RQUANTIZATION + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +enum class EQuantizationRoundingMode { + UNDEFINED = 0, ROUND = 1, FLOOR = 2, TRUNCATE = 3 +}; + +enum class EQuantizationOverflowMode { + UNDEFINED = 0, SAT = 1, SAT_SYM = 2 +}; + +enum class EQuantizationGranularity { + UNDEFINED = 0, PerTensor = 1, PerChannel = 2, PerWeight = 3 +}; + +struct QuantizationInfo { + unsigned bitWidth = 0; + bool isSigned = false; + bool narrow = false; + double scale = 1.0; + std::int64_t zeroPoint = 0; + std::string scaleTensor; + std::string zeroPointTensor; + EQuantizationRoundingMode rounding = EQuantizationRoundingMode::UNDEFINED; + EQuantizationOverflowMode overflow = EQuantizationOverflowMode::UNDEFINED; + EQuantizationGranularity granularity = EQuantizationGranularity::PerTensor; + int axis = -1; +}; + +inline std::pair QuantizedIntegerRange(const QuantizationInfo &info) +{ + if (info.bitWidth == 0 || info.bitWidth >= 63) { + throw std::runtime_error("SOFIE quantization metadata received unsupported bit width"); + } + if (info.isSigned) { + const std::int64_t qmax = (std::int64_t{1} << (info.bitWidth - 1)) - 1; + const std::int64_t qmin = info.narrow ? -qmax : -(std::int64_t{1} << (info.bitWidth - 1)); + return {qmin, qmax}; + } + const std::int64_t qmax = (std::int64_t{1} << info.bitWidth) - 1; + const std::int64_t qmin = info.narrow ? 1 : 0; + return {qmin, qmax}; +} + +inline std::int64_t QuantizeScalarToIntegerGrid(float value, const QuantizationInfo &info) +{ + const auto range = QuantizedIntegerRange(info); + auto quantized = static_cast(std::llround((static_cast(value) / info.scale) + info.zeroPoint)); + if (quantized < range.first) + quantized = range.first; + if (quantized > range.second) + quantized = range.second; + return quantized; +} + +enum class EQuantizedLoweringStatus { + UNDEFINED = 0, + Optimized = 1, + Baseline = 2, + BackendUnsupported = 3, + SemanticUnsupported = 4, + SemanticRecognized = 5 +}; + +inline bool IsQuantizedLoweringAvailable(EQuantizedLoweringStatus status) +{ + return status == EQuantizedLoweringStatus::Optimized || status == EQuantizedLoweringStatus::Baseline; +} + +inline bool IsQuantizedLoweringUnsupported(EQuantizedLoweringStatus status) +{ + return status == EQuantizedLoweringStatus::BackendUnsupported || + status == EQuantizedLoweringStatus::SemanticUnsupported; +} + +inline bool IsQuantizedLoweringOptimized(EQuantizedLoweringStatus status) +{ + return status == EQuantizedLoweringStatus::Optimized; +} + +enum class EQuantizedBackend { + UNDEFINED = 0, CPU = 1, ALPAKA = 2 +}; + +enum class EQuantizedStorageType { + UNDEFINED = 0, FloatCarrier = 1, Int8 = 2, UInt8 = 3, Int32Accumulator = 4, MetadataOnly = 5 +}; + +enum class EQuantizedCarrierMode { + UNDEFINED = 0, + Float = 1, + Int8 = 2, + UInt8 = 3, + Int32Accumulator = 4 +}; + +enum class EQuantizedOutputMode { + UNDEFINED = 0, + ExactFakeQuantFloat = 1, + Quantized = 2, + Int32Accumulator = 3 +}; + +enum class EQuantizedComputeProfile { + UNDEFINED = 0, + GenericRecognized = 1, + SignedInt8SymmetricPerTensorRank2 = 2, + SignedInt8PerTensorActivationPerChannelWeightRank2 = 3, + UnsignedInt8ActivationSignedInt8WeightRank2 = 4, + UnsignedInt8SymmetricRank2 = 5, + AsymmetricZeroPointRank2 = 6, + UnsupportedDenseLinearRank2 = 7 +}; + +enum class EQuantizedLayout { + UNDEFINED = 0, Plain = 1, Transposed = 2, PackedCPU = 3, PlainDevice = 4, TiledAlpaka = 5 +}; + + +enum class EQuantizedParameterMode { + UNDEFINED = 0, + Scalar = 1, + PerOutputChannel = 2 +}; + + +enum class EQuantizedEpilogueKind { + None = 0, + Bias = 1, + Relu = 2, + BiasRelu = 3 +}; + +inline bool QuantizedEpilogueHasBias(EQuantizedEpilogueKind kind) +{ + return kind == EQuantizedEpilogueKind::Bias || kind == EQuantizedEpilogueKind::BiasRelu; +} + +inline bool QuantizedEpilogueHasRelu(EQuantizedEpilogueKind kind) +{ + return kind == EQuantizedEpilogueKind::Relu || kind == EQuantizedEpilogueKind::BiasRelu; +} + +struct QuantizedEpilogue { + EQuantizedEpilogueKind kind = EQuantizedEpilogueKind::None; + std::string biasSourceTensor; + std::optional biasQuant; + std::optional addOpIndex; +}; + +enum class EQuantizedShapePolicy { + UNDEFINED = 0, + Exact = 1, + ExactTooSmall = 2, + PaddedCandidate = 3, + Padded = 4, + Fallback = 5, + Unsupported = 6 +}; + +inline bool QuantizedShapePolicyUsesPadding(EQuantizedShapePolicy policy) +{ + return policy == EQuantizedShapePolicy::PaddedCandidate || policy == EQuantizedShapePolicy::Padded; +} + +inline bool QuantizedShapePolicyIsExecutable(EQuantizedShapePolicy policy) +{ + return policy == EQuantizedShapePolicy::Exact || policy == EQuantizedShapePolicy::Padded; +} + +struct QuantizedDenseLinearShapePolicy { + EQuantizedShapePolicy policy = EQuantizedShapePolicy::UNDEFINED; + std::size_t logicalM = 0; + std::size_t logicalK = 0; + std::size_t logicalN = 0; + std::size_t physicalM = 0; + std::size_t physicalK = 0; + std::size_t physicalN = 0; + std::size_t logicalMacs = 0; + std::size_t physicalMacs = 0; + std::size_t minimumOptimizedMacs = 0; + bool belowMinimumWork = false; + double paddingWorkRatio = 1.0; + std::string reason; +}; + +enum class EQuantizedMatMulShapeKind { + UNDEFINED = 0, + Unsupported = 1, + Rank2 = 2, + FlattenableProjection = 3, + TrueBatched = 4 +}; + +struct QuantizedMatMulShapeAssessment { + EQuantizedMatMulShapeKind kind = EQuantizedMatMulShapeKind::UNDEFINED; + std::size_t logicalM = 0; + std::size_t logicalK = 0; + std::size_t logicalN = 0; + std::vector flattenedInputShape; + std::vector flattenedOutputShape; + std::string reason; + std::vector unsupportedReasons; +}; + +inline bool QuantizedMatMulShapeIsRecognized(const QuantizedMatMulShapeAssessment &assessment) +{ + return assessment.kind == EQuantizedMatMulShapeKind::Rank2 || + assessment.kind == EQuantizedMatMulShapeKind::FlattenableProjection || + assessment.kind == EQuantizedMatMulShapeKind::TrueBatched; +} + +inline bool QuantizedMatMulShapeIsRank2Executable(const QuantizedMatMulShapeAssessment &assessment) +{ + return assessment.kind == EQuantizedMatMulShapeKind::Rank2; +} + +struct QuantizedTensorStorage { + std::string logicalTensor; + std::string sourceTensor; + std::string storageTensor; + EQuantizedStorageType storageType = EQuantizedStorageType::UNDEFINED; + EQuantizedLayout layout = EQuantizedLayout::UNDEFINED; + QuantizationInfo quantization; + std::vector shape; + + EQuantizedBackend residentBackend = EQuantizedBackend::UNDEFINED; +}; + +inline std::size_t QuantizedStorageElementSize(EQuantizedStorageType type) +{ + switch (type) { + case EQuantizedStorageType::FloatCarrier: + return sizeof(float); + case EQuantizedStorageType::Int8: + return sizeof(std::int8_t); + case EQuantizedStorageType::UInt8: + return sizeof(std::uint8_t); + case EQuantizedStorageType::Int32Accumulator: + return sizeof(std::int32_t); + default: + return 0; + } +} + +inline std::size_t QuantizedStorageElementCount(const std::vector &shape) +{ + if (shape.empty()) + return 0; + std::size_t count = 1; + for (auto dim : shape) + count *= dim; + return count; +} + +inline std::size_t QuantizedStorageByteSize(EQuantizedStorageType type, const std::vector &shape) +{ + return QuantizedStorageElementSize(type) * QuantizedStorageElementCount(shape); +} + +inline bool IsPhysicalQuantizedStorage(EQuantizedStorageType type) +{ + return type == EQuantizedStorageType::Int8 || type == EQuantizedStorageType::UInt8 || + type == EQuantizedStorageType::Int32Accumulator; +} + +inline EQuantizedStorageType QuantizedStorageTypeForCarrier(const QuantizationInfo &info) +{ + return info.isSigned ? EQuantizedStorageType::Int8 : EQuantizedStorageType::UInt8; +} + +inline EQuantizedCarrierMode QuantizedCarrierModeForStorage(EQuantizedStorageType type) +{ + switch (type) { + case EQuantizedStorageType::FloatCarrier: + return EQuantizedCarrierMode::Float; + case EQuantizedStorageType::Int8: + return EQuantizedCarrierMode::Int8; + case EQuantizedStorageType::UInt8: + return EQuantizedCarrierMode::UInt8; + case EQuantizedStorageType::Int32Accumulator: + return EQuantizedCarrierMode::Int32Accumulator; + default: + return EQuantizedCarrierMode::UNDEFINED; + } +} + +struct QuantizedLoweringPlan { + EQuantizedBackend backend = EQuantizedBackend::UNDEFINED; + EQuantizedLoweringStatus status = EQuantizedLoweringStatus::UNDEFINED; + std::string reason; + + EQuantizedStorageType inputStorage = EQuantizedStorageType::UNDEFINED; + EQuantizedStorageType weightStorage = EQuantizedStorageType::UNDEFINED; + EQuantizedStorageType biasStorage = EQuantizedStorageType::UNDEFINED; + EQuantizedStorageType accumulatorStorage = EQuantizedStorageType::UNDEFINED; + EQuantizedStorageType outputStorage = EQuantizedStorageType::UNDEFINED; + + EQuantizedCarrierMode inputCarrierMode = EQuantizedCarrierMode::UNDEFINED; + EQuantizedOutputMode outputMode = EQuantizedOutputMode::UNDEFINED; + EQuantizedComputeProfile computeProfile = EQuantizedComputeProfile::UNDEFINED; + std::string capabilityTag; + QuantizedDenseLinearShapePolicy shapePolicy; + + std::string weightStorageTensor; + EQuantizedLayout weightLayout = EQuantizedLayout::UNDEFINED; + EQuantizedParameterMode weightScaleMode = EQuantizedParameterMode::Scalar; + std::string weightScaleTensor; + std::string weightZeroPointTensor; + + std::vector consumedOperatorIndices; + bool preservesQuantizationSemantics = false; + bool isMetadataOnly = false; + bool suppressesGraphOperators = false; +}; + +inline bool QuantizedPlanUsesInt32Accumulator(const QuantizedLoweringPlan &plan) +{ + return plan.accumulatorStorage == EQuantizedStorageType::Int32Accumulator; +} + +inline bool QuantizedPlanUsesPrequantizedWeights(const QuantizedLoweringPlan &plan) +{ + return !plan.weightStorageTensor.empty(); +} + +inline bool QuantizedPlanExposesQuantizedInputCarrier(const QuantizedLoweringPlan &plan) +{ + return plan.inputCarrierMode == EQuantizedCarrierMode::Int8 || + plan.inputCarrierMode == EQuantizedCarrierMode::UInt8; +} + +inline bool QuantizedPlanExposesQuantizedOutputCarrier(const QuantizedLoweringPlan &plan) +{ + return plan.outputMode == EQuantizedOutputMode::Quantized && + (plan.outputStorage == EQuantizedStorageType::Int8 || + plan.outputStorage == EQuantizedStorageType::UInt8); +} + +inline bool IsOptimizedQuantizedPlainDevicePlan(const QuantizedLoweringPlan &plan) +{ + return IsQuantizedLoweringOptimized(plan.status) && QuantizedPlanUsesPrequantizedWeights(plan) && + plan.weightLayout == EQuantizedLayout::PlainDevice; +} + +inline bool IsOptimizedQuantizedAlpakaPlainDevicePlan(const QuantizedLoweringPlan &plan) +{ + return plan.backend == EQuantizedBackend::ALPAKA && IsOptimizedQuantizedPlainDevicePlan(plan); +} + +struct QuantizedMatMulRegion { + // Quantized carrier tensors, i.e. outputs of quantization boundaries. + std::string inputTensor; + std::string weightTensor; + std::string matmulOutputTensor; + std::string outputTensor; + + // Source tensors consumed by the quantization boundaries. + std::string inputSourceTensor; + std::string weightSourceTensor; + + std::size_t inputQuantOpIndex = static_cast(-1); + std::size_t weightQuantOpIndex = static_cast(-1); + std::size_t matmulOpIndex = static_cast(-1); + std::size_t outputQuantOpIndex = static_cast(-1); + + QuantizedEpilogue epilogue; + QuantizedMatMulShapeAssessment shape; + + QuantizationInfo inputQuant; + QuantizationInfo weightQuant; + QuantizationInfo outputQuant; + + EQuantizedLoweringStatus status = EQuantizedLoweringStatus::UNDEFINED; + std::string reason; +}; + +struct QuantizedGemmRegion { + // Quantized carrier tensors, i.e. outputs of quantization boundaries. + std::string inputTensor; + std::string weightTensor; + std::string biasTensor; + std::string gemmOutputTensor; + std::string outputTensor; + + // Source tensors consumed by the quantization boundaries. Fused region code + // uses these names when it suppresses the literal Quant nodes. + std::string inputSourceTensor; + std::string weightSourceTensor; + std::string biasSourceTensor; + + std::size_t inputQuantOpIndex = static_cast(-1); + std::size_t weightQuantOpIndex = static_cast(-1); + std::optional biasQuantOpIndex; + std::size_t gemmOpIndex = static_cast(-1); + std::size_t outputQuantOpIndex = static_cast(-1); + + QuantizationInfo inputQuant; + QuantizationInfo weightQuant; + std::optional biasQuant; + QuantizationInfo outputQuant; + + float alpha = 1.0f; + float beta = 1.0f; + std::int64_t transA = 0; + std::int64_t transB = 0; + + EQuantizedLoweringStatus status = EQuantizedLoweringStatus::UNDEFINED; + std::string reason; +}; + +struct QuantizationModelState { + std::unordered_map tensorInfos; + std::unordered_map tensorStorages; + std::unordered_map gemmRegions; + std::unordered_map matmulRegions; + std::unordered_map> loweringPlans; + + void ClearDerivedAnalysis() + { + tensorStorages.clear(); + gemmRegions.clear(); + matmulRegions.clear(); + loweringPlans.clear(); + } +}; + +template +std::vector SortedQuantizedRegionOperatorIndices(const RegionMap ®ions) +{ + std::vector indices; + indices.reserve(regions.size()); + for (const auto &entry : regions) + indices.push_back(entry.first); + std::sort(indices.begin(), indices.end()); + return indices; +} + +inline const QuantizedLoweringPlan *FindQuantizedLoweringPlan(const QuantizationModelState &state, + std::size_t opIndex, + EQuantizedBackend backend) +{ + auto opIt = state.loweringPlans.find(opIndex); + if (opIt == state.loweringPlans.end()) + return nullptr; + auto backendIt = opIt->second.find(backend); + return backendIt == opIt->second.end() ? nullptr : &backendIt->second; +} + +} // namespace SOFIE + +#endif // SOFIE_RQUANTIZATION diff --git a/core/inc/SOFIE/RQuantization_Analysis.hxx b/core/inc/SOFIE/RQuantization_Analysis.hxx new file mode 100644 index 0000000..e5be74e --- /dev/null +++ b/core/inc/SOFIE/RQuantization_Analysis.hxx @@ -0,0 +1,71 @@ +#ifndef SOFIE_RQUANTIZATION_ANALYSIS +#define SOFIE_RQUANTIZATION_ANALYSIS + +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/ROperator.hxx" +#include "SOFIE/ROperator_Gemm.hxx" + +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +struct QuantizationGraphIndex { + std::unordered_map producerByTensor; + std::unordered_map> consumersByTensor; +}; + +struct QuantizedDenseLinearPatternMatch { + QuantizedGemmRegion region; + QuantizedMatMulShapeAssessment matmulShape; + std::vector reasons; + bool isMatMul = false; + bool hasInlineMatMulBias = false; +}; + +QuantizedDenseLinearPatternMatch MatchQuantizedDenseLinearPattern( + const ROperator_Gemm &gemm, std::size_t opIndex, + const std::function(const std::string &)> &tensorShape); + +QuantizationGraphIndex BuildQuantizationGraphIndex(const std::vector> &operators); + +std::optional MatchQuantizationBoundaryProducer( + const QuantizationGraphIndex &graph, const std::vector> &operators, + const std::string &tensor, const std::string &role, std::vector &reasons); + +std::optional MatchSingleTensorConsumer(const QuantizationGraphIndex &graph, + const std::string &tensor, + const std::string &role, + std::vector &reasons); + +bool IsFloatAddOperator(const ROperator &op); +void CheckQuantizationInfo(const QuantizationInfo &info, const std::string &role, + std::vector &reasons); +void CheckQuantizedGemmAttributes(const QuantizedGemmRegion ®ion, + std::vector &reasons); +void CheckQuantizedGemmRank2Shape(const std::vector &shape, + const std::string &role, + std::vector &reasons); +std::vector QuantizedGemmLoweringUnsupportedReasons(const QuantizedGemmRegion ®ion); + +std::string JoinQuantizationReasons(const std::vector &reasons); + +std::vector QuantizedGemmConsumedOperatorIndices(const QuantizedGemmRegion ®ion); +std::vector QuantizedMatMulConsumedOperatorIndices(const QuantizedMatMulRegion ®ion); + +QuantizedMatMulRegion MakeQuantizedMatMulRegionFromGemmLikeRegion(const QuantizedGemmRegion ®ion); + +bool IsDenseLinearBiasLikeShape(const std::vector &biasShape, + const std::vector &outputShape); + +QuantizationInfo MakeAccumulatorBiasQuantization(const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant); + +} // namespace SOFIE + +#endif // SOFIE_RQUANTIZATION_ANALYSIS diff --git a/core/inc/SOFIE/RQuantization_DenseLinear.hxx b/core/inc/SOFIE/RQuantization_DenseLinear.hxx new file mode 100644 index 0000000..90fffba --- /dev/null +++ b/core/inc/SOFIE/RQuantization_DenseLinear.hxx @@ -0,0 +1,123 @@ +#ifndef SOFIE_RQUANTIZATION_DENSELINEAR +#define SOFIE_RQUANTIZATION_DENSELINEAR + +#include "SOFIE/RQuantization.hxx" + +#include +#include +#include +#include +#include + +namespace SOFIE { + +struct QuantizedDenseLinearProfileAssessment { + EQuantizedComputeProfile profile = EQuantizedComputeProfile::UNDEFINED; + bool cublasLtOptimizedCandidate = false; + std::vector reasons; +}; + +struct QuantizedDenseLinearOperands { + QuantizationInfo inputQuant; + QuantizationInfo weightQuant; + QuantizationInfo outputQuant; + std::optional biasQuant; + + std::vector inputShape; + std::vector weightShape; + std::vector outputShape; + + int weightOutputChannelAxis = -1; + std::string operatorName; +}; + +struct QuantizedDenseLinearCublasLtCapability { + bool optimized = false; + EQuantizedComputeProfile profile = EQuantizedComputeProfile::GenericRecognized; + std::string tag = "recognized_not_cublaslt_optimized"; + std::string reason; + QuantizedDenseLinearShapePolicy shapePolicy; +}; + +bool IsScalarPerTensor(const QuantizationInfo &info); +bool IsPerChannelAxis(const QuantizationInfo &info, int axis); + +std::vector DenseLinearQuantizationParameterUnsupportedReasons( + const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant, + const QuantizationInfo &outputQuant, + const std::optional &biasQuant, + int expectedWeightPerChannelAxis, + const std::string &operatorName); + +std::vector QuantizeTensorToInt8(const float *data, std::size_t length, const QuantizationInfo &info); +std::vector QuantizeTensorToUInt8(const float *data, std::size_t length, const QuantizationInfo &info); + +std::vector QuantizeGemmWeightTensorToInt8(const float *data, std::size_t n, std::size_t k, + const QuantizationInfo &info, + const std::vector &perChannelScale); +std::vector QuantizeMatMulWeightTensorToInt8Transposed(const float *data, std::size_t k, std::size_t n, + const QuantizationInfo &info, + const std::vector &perChannelScale); + +std::vector QuantizeGemmWeightTensorToInt8Padded(const float *data, std::size_t n, std::size_t k, + std::size_t physicalN, std::size_t physicalK, + const QuantizationInfo &info, + const std::vector &perChannelScale); +std::vector QuantizeMatMulWeightTensorToInt8TransposedPadded(const float *data, std::size_t k, std::size_t n, + std::size_t physicalK, std::size_t physicalN, + const QuantizationInfo &info, + const std::vector &perChannelScale); + +QuantizedDenseLinearProfileAssessment AssessDenseLinearComputeProfile( + const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant, + const QuantizationInfo &outputQuant, + int expectedWeightPerChannelAxis, + const std::string &operatorName); + +QuantizedDenseLinearShapePolicy MakeCublasLtShapePolicy(std::size_t m, std::size_t k, std::size_t n); + +bool IsProfitableCublasLtPaddedDenseLinearPolicy(const QuantizedDenseLinearShapePolicy &policy); +std::string ExplainCublasLtPaddedDenseLinearProfitability(const QuantizedDenseLinearShapePolicy &policy); + +QuantizedMatMulShapeAssessment AssessQuantizedMatMulShape( + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape); + +QuantizedDenseLinearCublasLtCapability AssessCublasLtDenseLinearCapability( + const QuantizedDenseLinearOperands &operands); + +QuantizedDenseLinearCublasLtCapability +SelectExecutableDenseLinearCapability(QuantizedDenseLinearCublasLtCapability capability); + +QuantizedDenseLinearOperands MakeDenseLinearOperands(const QuantizedGemmRegion ®ion, + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape); +QuantizedDenseLinearOperands MakeDenseLinearOperands(const QuantizedMatMulRegion ®ion, + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape); + +QuantizedLoweringPlan MakeUnsupportedQuantizedMatMulPlan(const QuantizedMatMulRegion ®ion, + EQuantizedBackend backend, + std::string reason, + bool preservesSemantics); +QuantizedLoweringPlan MakeMatMulAlpakaTransposedWeightStoragePlan( + const QuantizedMatMulRegion ®ion, const std::string &weightStorageTensor, + const QuantizedDenseLinearShapePolicy &shapePolicy); +QuantizedLoweringPlan MakeCPUPackedWeightBaselinePlan(const QuantizedGemmRegion ®ion, + const std::string &weightStorageTensor); +QuantizedLoweringPlan MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend backend, + std::string reason, + bool preservesSemantics); +QuantizedLoweringPlan MakeAlpakaFakeQuantPlan(const QuantizedGemmRegion ®ion); +QuantizedLoweringPlan MakeAlpakaCublasLtCorePlan( + const QuantizedGemmRegion ®ion, const std::string &weightStorageTensor, + const QuantizedDenseLinearCublasLtCapability &capability); + +} // namespace SOFIE + +#endif // SOFIE_RQUANTIZATION_DENSELINEAR diff --git a/core/inc/SOFIE/RQuantization_Parameters.hxx b/core/inc/SOFIE/RQuantization_Parameters.hxx new file mode 100644 index 0000000..671a78b --- /dev/null +++ b/core/inc/SOFIE/RQuantization_Parameters.hxx @@ -0,0 +1,40 @@ +#ifndef SOFIE_RQUANTIZATION_PARAMETERS +#define SOFIE_RQUANTIZATION_PARAMETERS + +#include "SOFIE/RQuantization.hxx" + +#include +#include +#include +#include + +namespace SOFIE { + +struct QuantizationParameterSpec { + std::vector scales; + std::vector zeroPoints; + unsigned bitWidth = 0; + bool isSigned = false; + bool narrow = false; + EQuantizationRoundingMode rounding = EQuantizationRoundingMode::UNDEFINED; + EQuantizationOverflowMode overflow = EQuantizationOverflowMode::UNDEFINED; + std::string scaleTensor; + std::string zeroPointTensor; + std::vector tensorShape; + std::optional explicitAxis; + std::string context; +}; + +std::vector ValidateIntegralZeroPoints(const std::vector &values, + const std::string &context); + +int InferQuantizationParameterAxis(const std::vector &tensorShape, + std::size_t parameterCount, + std::optional explicitAxis, + const std::string &context); + +QuantizationInfo MakeValidatedQuantizationInfo(const QuantizationParameterSpec &spec); + +} // namespace SOFIE + +#endif // SOFIE_RQUANTIZATION_PARAMETERS diff --git a/core/inc/SOFIE/RQuantization_Storage.hxx b/core/inc/SOFIE/RQuantization_Storage.hxx new file mode 100644 index 0000000..f3ca00f --- /dev/null +++ b/core/inc/SOFIE/RQuantization_Storage.hxx @@ -0,0 +1,58 @@ +#ifndef SOFIE_RQUANTIZATION_STORAGE +#define SOFIE_RQUANTIZATION_STORAGE + +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/SOFIE_common.hxx" + +#include +#include +#include +#include +#include + +namespace SOFIE { + +ETensorType TensorTypeForQuantizedStorage(EQuantizedStorageType storage); + +QuantizedTensorStorage MakeQuantizedTensorStorage(std::string logicalTensor, + std::string sourceTensor, + std::string storageTensor, + const QuantizationInfo &quantization, + EQuantizedLayout layout, + std::vector shape, + EQuantizedBackend backend); + +std::vector PackQuantizedGemmWeightsInt8(const float *data, + std::size_t n, + std::size_t k, + std::size_t tileN, + const QuantizationInfo &info); + +std::vector PackQuantizedGemmWeightsUInt8(const float *data, + std::size_t n, + std::size_t k, + std::size_t tileN, + const QuantizationInfo &info); + +using QuantizedWeightBuffer = std::variant, std::vector>; + +struct MaterializedQuantizedWeight { + QuantizedTensorStorage storage; + QuantizedWeightBuffer buffer; +}; + +MaterializedQuantizedWeight MaterializeQuantizedGemmWeight( + const QuantizedGemmRegion ®ion, const QuantizedLoweringPlan &plan, + EQuantizedBackend backend, const float *sourceData, + const std::vector &sourceShape, + const std::vector &perChannelScales); + +MaterializedQuantizedWeight MaterializeQuantizedMatMulWeight( + const QuantizedMatMulRegion ®ion, const QuantizedLoweringPlan &plan, + EQuantizedBackend backend, const float *sourceData, + const std::vector &sourceShape, + const std::vector &perChannelScales); + +} // namespace SOFIE + +#endif // SOFIE_RQUANTIZATION_STORAGE diff --git a/core/inc/SOFIE/RWeightFile.hxx b/core/inc/SOFIE/RWeightFile.hxx new file mode 100644 index 0000000..5202cc0 --- /dev/null +++ b/core/inc/SOFIE/RWeightFile.hxx @@ -0,0 +1,98 @@ +#ifndef SOFIE_RWEIGHTFILE +#define SOFIE_RWEIGHTFILE + +#include "SOFIE/SOFIE_common.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +inline constexpr char kBinaryWeightFileMagic[8] = {'S', 'O', 'F', 'I', 'E', 'W', 'B', '1'}; +inline constexpr std::uint32_t kBinaryWeightFileVersion = 1; +inline constexpr std::uint32_t kBinaryWeightFileEndianMarker = 0x01020304u; + +using BinaryWeightTensorEntry = std::pair; + +bool IsBinaryWeightTensorType(ETensorType type); + +std::vector +CollectBinaryWeightTensors(const std::unordered_map &tensors); + +long WriteBinaryWeightFile(std::ostream &stream, + const std::unordered_map &tensors); + +template +void ReadBinaryWeightValue(std::istream &stream, T &value) +{ + static_assert(std::is_trivially_copyable_v); + stream.read(reinterpret_cast(&value), sizeof(T)); + if (!stream) + throw std::runtime_error("sofie binary weight file is truncated"); +} + +inline void ReadBinaryWeightFileHeader(std::istream &stream, std::uint64_t expectedTensorCount) +{ + char magic[8]{}; + stream.read(magic, sizeof(magic)); + std::uint32_t version = 0; + std::uint32_t endianMarker = 0; + std::uint64_t tensorCount = 0; + ReadBinaryWeightValue(stream, version); + ReadBinaryWeightValue(stream, endianMarker); + ReadBinaryWeightValue(stream, tensorCount); + if (!std::equal(std::begin(magic), std::end(magic), std::begin(kBinaryWeightFileMagic))) + throw std::runtime_error("sofie binary weight file has an invalid magic header"); + if (version != kBinaryWeightFileVersion) + throw std::runtime_error("sofie binary weight file has an unsupported version"); + if (endianMarker != kBinaryWeightFileEndianMarker) + throw std::runtime_error("sofie binary weight file uses an incompatible byte order"); + if (tensorCount != expectedTensorCount) + throw std::runtime_error("sofie binary weight file tensor count does not match the generated model"); +} + +template +void ReadTensorFromBinaryStream(std::istream &stream, T &target, + const std::string &expectedName, ETensorType expectedType, + const std::vector &expectedShape) +{ + std::uint32_t nameLength = 0; + ReadBinaryWeightValue(stream, nameLength); + std::string name(nameLength, 0); + stream.read(name.data(), name.size()); + std::uint32_t type = 0; + std::uint32_t rank = 0; + ReadBinaryWeightValue(stream, type); + ReadBinaryWeightValue(stream, rank); + std::vector shape(rank); + for (auto &dim : shape) { + std::uint64_t value = 0; + ReadBinaryWeightValue(stream, value); + dim = static_cast(value); + } + std::uint64_t elementCount = 0; + std::uint64_t byteCount = 0; + ReadBinaryWeightValue(stream, elementCount); + ReadBinaryWeightValue(stream, byteCount); + std::size_t expectedElementCount = 1; + for (auto dim : expectedShape) + expectedElementCount *= dim; + using Element = std::remove_cv_t>; + if (name != expectedName || type != static_cast(expectedType) || shape != expectedShape || + elementCount != expectedElementCount || byteCount != expectedElementCount * sizeof(Element)) + throw std::runtime_error("sofie binary weight tensor metadata does not match generated tensor " + expectedName); + stream.read(reinterpret_cast(&target[0]), static_cast(byteCount)); + if (!stream) + throw std::runtime_error("sofie binary weight data is truncated for tensor " + expectedName); +} + +} // namespace SOFIE + +#endif // SOFIE_RWEIGHTFILE diff --git a/core/inc/SOFIE/SOFIE_Quantized.hxx b/core/inc/SOFIE/SOFIE_Quantized.hxx new file mode 100644 index 0000000..464ed9f --- /dev/null +++ b/core/inc/SOFIE/SOFIE_Quantized.hxx @@ -0,0 +1,314 @@ +#ifndef SOFIE_QUANTIZED +#define SOFIE_QUANTIZED + +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +// Portable quantized runtime helpers used by generated lowered quantized operators. + +inline bool NearlyEqualQuantizedScale(double lhs, double rhs) +{ + const auto scale = std::max({1.0, std::fabs(lhs), std::fabs(rhs)}); + return std::fabs(lhs - rhs) <= (1.0e-12 * scale); +} + +inline bool MakeQuantizedFixedPointMultiplier(double realMultiplier, std::int64_t &multiplier, int &shift) +{ + if (!std::isfinite(realMultiplier) || realMultiplier <= 0.0) { + return false; + } + + int exponent = 0; + const double significand = std::frexp(realMultiplier, &exponent); + auto q31 = static_cast(std::llround(significand * 2147483648.0)); + if (q31 == (std::int64_t{1} << 31)) { + q31 /= 2; + ++exponent; + } + + const int candidateShift = 31 - exponent; + if (q31 <= 0 || candidateShift < 0 || candidateShift >= 62) { + return false; + } + + multiplier = q31; + shift = candidateShift; + return true; +} + +inline bool MakeExactIntegerScaleMultiplier(double scale, std::int64_t &multiplier, int &shift) +{ + if (!std::isfinite(scale) || scale < 0.0) { + return false; + } + + const auto rounded = std::llround(scale); + if (!NearlyEqualQuantizedScale(scale, static_cast(rounded))) { + return false; + } + + multiplier = rounded; + shift = 0; + return true; +} + +enum class EQuantizedGemmActivation { + None = 0, + Relu = 1 +}; + +struct QuantizedGemmParams { + std::size_t m = 0; + std::size_t n = 0; + std::size_t k = 0; + std::size_t tileN = 4; + + double scaleX = 1.0; + double scaleW = 1.0; + double scaleY = 1.0; + double scaleB = 1.0; + + std::int32_t zeroX = 0; + std::int32_t zeroW = 0; + std::int32_t zeroY = 0; + std::int32_t zeroB = 0; + + std::int32_t qminX = 0; + std::int32_t qmaxX = 0; + std::int32_t qminY = 0; + std::int32_t qmaxY = 0; + std::int32_t qminB = 0; + std::int32_t qmaxB = 0; + + bool hasBias = false; + + bool useIntegerEpilogue = false; + bool useAccumulatorBias = false; + bool useIntegerBias = false; + std::int64_t requantMultiplier = 0; + int requantShift = 0; + std::int64_t biasRequantMultiplier = 0; + int biasRequantShift = 0; + + EQuantizedGemmActivation activation = EQuantizedGemmActivation::None; +}; + +namespace INTERNAL { + +enum class EQuantizedGemmCPUBackend { + Portable = 0, + X86 = 1, + ARM = 2 +}; + +inline constexpr bool IsQuantizedGemmX86Target() +{ +#if defined(__x86_64__) || defined(_M_X64) || defined(__i386) || defined(_M_IX86) + return true; +#else + return false; +#endif +} + +inline constexpr bool IsQuantizedGemmARMTarget() +{ +#if defined(__aarch64__) || defined(__arm__) || defined(_M_ARM64) || defined(_M_ARM) + return true; +#else + return false; +#endif +} + +inline constexpr EQuantizedGemmCPUBackend SelectQuantizedGemmCPUBackend() +{ + // Selects the CPU backend family visible at compile time. + // Portable remains the fallback for every target family. + if constexpr (IsQuantizedGemmX86Target()) { + return EQuantizedGemmCPUBackend::X86; + } else if constexpr (IsQuantizedGemmARMTarget()) { + return EQuantizedGemmCPUBackend::ARM; + } else { + return EQuantizedGemmCPUBackend::Portable; + } +} + +inline std::int32_t QuantizedGemmClamp(std::int32_t value, std::int32_t qmin, std::int32_t qmax) +{ + return std::max(qmin, std::min(qmax, value)); +} + +inline std::int32_t QuantizedGemmQuantizeClamp(double value, double scale, std::int32_t zero, std::int32_t qmin, + std::int32_t qmax) +{ + auto quantized = static_cast(std::nearbyint((value / scale) + zero)); + return QuantizedGemmClamp(quantized, qmin, qmax); +} + +inline std::int64_t QuantizedGemmRoundDivideByPowerOfTwo(std::int64_t value, int shift) +{ + if (shift <= 0) { + const auto leftShift = -shift; + if (leftShift >= 62) { + return value >= 0 ? std::numeric_limits::max() : std::numeric_limits::min(); + } + return value << leftShift; + } + + if (shift >= 62) { + return 0; + } + + const std::uint64_t magnitude = value < 0 ? static_cast(-value) : static_cast(value); + const std::uint64_t quotient = magnitude >> shift; + const std::uint64_t remainder = magnitude & ((std::uint64_t{1} << shift) - 1); + const std::uint64_t half = std::uint64_t{1} << (shift - 1); + const bool roundUp = (remainder > half) || (remainder == half && (quotient & std::uint64_t{1}) != 0); + const auto rounded = static_cast(quotient + (roundUp ? 1 : 0)); + return value < 0 ? -rounded : rounded; +} + +inline std::int64_t QuantizedGemmApplyFixedPointMultiplier(std::int32_t accumulator, std::int64_t multiplier, int shift) +{ + return QuantizedGemmRoundDivideByPowerOfTwo(static_cast(accumulator) * multiplier, shift); +} + +inline std::int32_t QuantizedGemmFinalizeY(std::int64_t yq, const QuantizedGemmParams ¶ms) +{ + if (params.activation == EQuantizedGemmActivation::Relu) { + yq = std::max(yq, params.zeroY); + } + yq = std::max(params.qminY, std::min(params.qmaxY, yq)); + return static_cast(yq); +} + +inline float QuantizedGemmFinalizeAccumulator(std::int32_t accumulator, std::size_t col, const float *bias, + const QuantizedGemmParams ¶ms) +{ + const bool hasRuntimeBias = params.hasBias && bias != nullptr; + const bool canUseIntegerEpilogue = + params.useIntegerEpilogue && (!hasRuntimeBias || params.useAccumulatorBias || params.useIntegerBias); + + if (canUseIntegerEpilogue) { + auto acc = accumulator; + std::int32_t biasCentered = 0; + if (hasRuntimeBias) { + const auto bq = QuantizedGemmQuantizeClamp(static_cast(bias[col]), params.scaleB, params.zeroB, + params.qminB, params.qmaxB); + biasCentered = bq - params.zeroB; + if (params.useAccumulatorBias) { + acc += biasCentered; + } + } + + auto scaled = QuantizedGemmApplyFixedPointMultiplier(acc, params.requantMultiplier, params.requantShift); + if (hasRuntimeBias && params.useIntegerBias && !params.useAccumulatorBias) { + scaled += QuantizedGemmApplyFixedPointMultiplier(biasCentered, params.biasRequantMultiplier, + params.biasRequantShift); + } + const auto yq = QuantizedGemmFinalizeY(scaled + params.zeroY, params); + return static_cast(static_cast(yq - params.zeroY) * params.scaleY); + } + + double real = static_cast(accumulator) * params.scaleX * params.scaleW; + if (hasRuntimeBias) { + const auto bq = QuantizedGemmQuantizeClamp(static_cast(bias[col]), params.scaleB, params.zeroB, + params.qminB, params.qmaxB); + real += static_cast(bq - params.zeroB) * params.scaleB; + } + + const auto yq = QuantizedGemmFinalizeY( + static_cast(std::nearbyint((real / params.scaleY) + params.zeroY)), params); + return static_cast(static_cast(yq - params.zeroY) * params.scaleY); +} + +template +inline void QuantizedGemmCPU_Portable(float *output, const float *input, const WeightT *packedWeight, const float *bias, + const QuantizedGemmParams ¶ms) +{ + if (params.tileN == 0) { + return; + } + + thread_local std::vector xqScratch; + xqScratch.resize(params.k); + auto *xqRow = xqScratch.data(); + + for (std::size_t row = 0; row < params.m; ++row) { + for (std::size_t kk = 0; kk < params.k; ++kk) { + xqRow[kk] = QuantizedGemmQuantizeClamp(static_cast(input[row * params.k + kk]), params.scaleX, + params.zeroX, params.qminX, params.qmaxX); + } + + for (std::size_t colBase = 0; colBase < params.n; colBase += params.tileN) { + const std::size_t block = colBase / params.tileN; + const std::size_t columns = std::min(params.tileN, params.n - colBase); + + if (params.tileN == 4) { + std::int32_t acc[4] = {0, 0, 0, 0}; + + for (std::size_t kk = 0; kk < params.k; ++kk) { + const auto xCentered = xqRow[kk] - params.zeroX; + const auto weightBase = (block * params.k + kk) * params.tileN; + for (std::size_t ji = 0; ji < columns; ++ji) { + const auto wq = static_cast(packedWeight[weightBase + ji]); + acc[ji] += xCentered * (wq - params.zeroW); + } + } + + for (std::size_t ji = 0; ji < columns; ++ji) { + const std::size_t col = colBase + ji; + output[row * params.n + col] = QuantizedGemmFinalizeAccumulator(acc[ji], col, bias, params); + } + } else { + std::vector acc(params.tileN); + + for (std::size_t kk = 0; kk < params.k; ++kk) { + const auto xCentered = xqRow[kk] - params.zeroX; + const auto weightBase = (block * params.k + kk) * params.tileN; + for (std::size_t ji = 0; ji < columns; ++ji) { + const auto wq = static_cast(packedWeight[weightBase + ji]); + acc[ji] += xCentered * (wq - params.zeroW); + } + } + + for (std::size_t ji = 0; ji < columns; ++ji) { + const std::size_t col = colBase + ji; + output[row * params.n + col] = QuantizedGemmFinalizeAccumulator(acc[ji], col, bias, params); + } + } + } + } +} + +template +inline void QuantizedGemmCPU_Dispatch(float *output, const float *input, const WeightT *packedWeight, const float *bias, + const QuantizedGemmParams ¶ms) +{ + [[maybe_unused]] constexpr auto backend = SelectQuantizedGemmCPUBackend(); + // The dispatch boundary is explicit even when the portable kernel is selected. + // Architecture-specific kernels can replace this call without changing generated code. + QuantizedGemmCPU_Portable(output, input, packedWeight, bias, params); +} + +} // namespace INTERNAL + +template +inline void QuantizedGemm_Call(float *output, const float *input, const WeightT *packedWeight, const float *bias, + const QuantizedGemmParams ¶ms) +{ + static_assert(std::is_same_v || std::is_same_v, + "SOFIE::QuantizedGemm_Call supports int8/uint8 packed weights"); + + INTERNAL::QuantizedGemmCPU_Dispatch(output, input, packedWeight, bias, params); +} + +} // namespace SOFIE + +#endif // SOFIE_QUANTIZED diff --git a/core/inc/SOFIE/SOFIE_QuantizedAlpaka.hxx b/core/inc/SOFIE/SOFIE_QuantizedAlpaka.hxx new file mode 100644 index 0000000..d377e47 --- /dev/null +++ b/core/inc/SOFIE/SOFIE_QuantizedAlpaka.hxx @@ -0,0 +1,1113 @@ +#ifndef SOFIE_QUANTIZED_ALPAKA +#define SOFIE_QUANTIZED_ALPAKA + +#include +#include +#include +#include + +#ifdef SOFIE_USE_CUBLASLT +#include +#include +#endif + +namespace SOFIE { + +enum class EQuantizedCudaWeightType { + Int8, + UInt8 +}; + +enum class EQuantizedCudaEpilogueMode { + ExactFakeQuant, + Quantized +}; + +enum class EQuantizedCudaInputCarrier { + Float, + Int8 +}; + +enum class EQuantizedCudaOutputCarrier { + Float, + Int8, + UInt8 +}; + +enum class EQuantizedCudaScaleMode { + PerTensor, + PerOutputChannel +}; + +struct QuantizedGemmCudaLtParams { + std::size_t m = 0; + std::size_t n = 0; + std::size_t k = 0; + std::size_t logicalM = 0; + std::size_t logicalN = 0; + std::size_t logicalK = 0; + bool paddedExecution = false; + double inputScale = 1.0; + double weightScale = 1.0; + double biasScale = 1.0; + double outputScale = 1.0; + std::int32_t inputZeroPoint = 0; + std::int32_t weightZeroPoint = 0; + std::int32_t biasZeroPoint = 0; + std::int32_t outputZeroPoint = 0; + std::int32_t inputQMin = -128; + std::int32_t inputQMax = 127; + std::int32_t biasQMin = -2147483648; + std::int32_t biasQMax = 2147483647; + std::int32_t outputQMin = -128; + std::int32_t outputQMax = 127; + bool hasBias = false; + bool hasRelu = false; + std::size_t maxWorkspaceBytes = 32ULL * 1024ULL * 1024ULL; + EQuantizedCudaEpilogueMode epilogueMode = EQuantizedCudaEpilogueMode::ExactFakeQuant; + EQuantizedCudaInputCarrier inputCarrier = EQuantizedCudaInputCarrier::Float; + EQuantizedCudaOutputCarrier outputCarrier = EQuantizedCudaOutputCarrier::Float; + EQuantizedCudaWeightType weightType = EQuantizedCudaWeightType::Int8; + EQuantizedCudaScaleMode weightScaleMode = EQuantizedCudaScaleMode::PerTensor; + bool enableAutotuning = true; + int autotuneIterations = 3; + double accumulatorToOutputScale = 0.0; +}; + +#ifndef SOFIE_USE_CUBLASLT + +using QuantizedGemmCudaStream = void *; + +struct QuantizedGemmCudaLtState {}; + +struct QuantizedGemmCudaLtProfile { + float inputQuantizeMs = 0.0f; + float gemmMs = 0.0f; + float epilogueMs = 0.0f; + float totalMs = 0.0f; + std::size_t workspaceSize = 0; + int heuristicResultCount = 0; + int selectedHeuristicIndex = 0; + float autotuneMs = 0.0f; + int autotunedCandidateCount = 0; + float selectedCandidateMs = 0.0f; + std::size_t accumulatorBytes = 0; + std::size_t outputBytes = 0; + std::size_t epilogueTrafficBytes = 0; + bool directOutputCarrier = false; +}; + +inline void QuantizedGemmCudaLt_SetProfiling(bool) {} +inline bool QuantizedGemmCudaLt_ProfilingEnabled() { return false; } +inline QuantizedGemmCudaLtProfile QuantizedGemmCudaLt_GetLastProfile() { return {}; } + +inline void QuantizedGemmCudaPadInt8Matrix(QuantizedGemmCudaStream, const std::int8_t *, std::int8_t *, + std::size_t, std::size_t, std::size_t, std::size_t, std::int8_t = 0) +{ + throw std::runtime_error("SOFIE quantized padded CUDA input path was selected, but SOFIE_USE_CUBLASLT is not enabled"); +} + +inline void QuantizedGemmCudaUnpadInt8Matrix(QuantizedGemmCudaStream, const std::int8_t *, std::int8_t *, + std::size_t, std::size_t, std::size_t, std::size_t) +{ + throw std::runtime_error("SOFIE quantized padded CUDA output path was selected, but SOFIE_USE_CUBLASLT is not enabled"); +} + +inline void QuantizedGemmCudaUnpadUInt8Matrix(QuantizedGemmCudaStream, const std::uint8_t *, std::uint8_t *, + std::size_t, std::size_t, std::size_t, std::size_t) +{ + throw std::runtime_error("SOFIE quantized padded CUDA output path was selected, but SOFIE_USE_CUBLASLT is not enabled"); +} + +inline void QuantizedGemmCudaLt_Call(QuantizedGemmCudaLtState &, QuantizedGemmCudaStream, void *, const void *, + const void *, const float *, const float *, const QuantizedGemmCudaLtParams &) +{ + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM path was selected, but SOFIE_USE_CUBLASLT is not enabled"); +} + + +#else + +using QuantizedGemmCudaStream = cudaStream_t; + +namespace INTERNAL { + +inline void CheckCudaStatus(cudaError_t status, const char *where) +{ + if (status != cudaSuccess) { + throw std::runtime_error(std::string("SOFIE CUDA quantized GEMM failure in ") + where + ": " + + cudaGetErrorString(status)); + } +} + +inline void CheckCublasLtStatus(cublasStatus_t status, const char *where) +{ + if (status != CUBLAS_STATUS_SUCCESS) { + throw std::runtime_error(std::string("SOFIE cuBLASLt quantized GEMM failure in ") + where + + ": status " + std::to_string(static_cast(status))); + } +} + +__device__ inline std::int32_t QuantizedCudaClamp(std::int32_t value, std::int32_t qmin, std::int32_t qmax) +{ + return value < qmin ? qmin : (value > qmax ? qmax : value); +} + +__device__ inline std::int32_t QuantizedCudaQuantizeClamp(double value, double scale, std::int32_t zero, + std::int32_t qmin, std::int32_t qmax) +{ + const auto quantized = static_cast(nearbyint((value / scale) + static_cast(zero))); + return QuantizedCudaClamp(quantized, qmin, qmax); +} + +__global__ void QuantizedGemmCudaQuantizeInputKernel(const float *input, std::int8_t *inputQuantized, + std::size_t elements, double scale, std::int32_t zero, + std::int32_t qmin, std::int32_t qmax) +{ + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elements) + return; + inputQuantized[idx] = static_cast(QuantizedCudaQuantizeClamp(static_cast(input[idx]), scale, + zero, qmin, qmax)); +} + +__global__ void QuantizedGemmCudaPadInt8MatrixKernel(const std::int8_t *__restrict__ input, + std::int8_t *__restrict__ padded, + std::size_t logicalRows, std::size_t logicalCols, + std::size_t physicalRows, std::size_t physicalCols, + std::int8_t padValue) +{ + const std::size_t elements = physicalRows * physicalCols; + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elements) + return; + const std::size_t row = idx / physicalCols; + const std::size_t col = idx % physicalCols; + padded[idx] = (row < logicalRows && col < logicalCols) ? input[row * logicalCols + col] : padValue; +} + +template +__global__ void QuantizedGemmCudaUnpadMatrixKernel(const T *__restrict__ padded, + T *__restrict__ output, + std::size_t logicalRows, std::size_t logicalCols, + std::size_t physicalRows, std::size_t physicalCols) +{ + const std::size_t elements = logicalRows * logicalCols; + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elements) + return; + const std::size_t row = idx / logicalCols; + const std::size_t col = idx % logicalCols; + output[idx] = row < physicalRows && col < physicalCols ? padded[row * physicalCols + col] : T{}; +} + +__global__ void QuantizedGemmCudaBiasOutputOffsetKernel(float *__restrict__ biasOutputOffset, + const float *__restrict__ bias, + const float *__restrict__ weightScaleVector, + QuantizedGemmCudaLtParams params) +{ + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= params.n) + return; + + const float biasScale = (params.weightScaleMode == EQuantizedCudaScaleMode::PerOutputChannel) + ? static_cast(params.inputScale * static_cast(weightScaleVector[idx])) + : static_cast(params.biasScale); + int bq = __float2int_rn((bias[idx] / biasScale) + static_cast(params.biasZeroPoint)); + bq = QuantizedCudaClamp(bq, params.biasQMin, params.biasQMax); + biasOutputOffset[idx] = + (static_cast(bq - params.biasZeroPoint) * biasScale / static_cast(params.outputScale)) + + static_cast(params.outputZeroPoint); +} +template +__global__ void QuantizedGemmCudaQuantizedEpilogueKernel(OutputT *__restrict__ output, + const std::int32_t *__restrict__ accumulator, + const float *__restrict__ biasOutputOffset, + const float *__restrict__ weightScaleVector, + QuantizedGemmCudaLtParams params) +{ + const std::size_t elements = params.m * params.n; + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elements) + return; + + const std::size_t col = idx % params.n; + const float scale = (params.weightScaleMode == EQuantizedCudaScaleMode::PerOutputChannel) + ? static_cast((params.inputScale * static_cast(weightScaleVector[col])) / params.outputScale) + : static_cast(params.accumulatorToOutputScale); + const float offset = HasBias ? biasOutputOffset[col] : static_cast(params.outputZeroPoint); + int yq = __float2int_rn( + __fmaf_rn(static_cast(accumulator[idx]), scale, offset)); + if constexpr (HasRelu) { + if (yq < params.outputZeroPoint) + yq = params.outputZeroPoint; + } + yq = QuantizedCudaClamp(yq, params.outputQMin, params.outputQMax); + output[idx] = static_cast(yq); +} + +__global__ void QuantizedGemmCudaEpilogueKernel(float *output, const std::int32_t *accumulator, const float *bias, + QuantizedGemmCudaLtParams params) +{ + const std::size_t elements = params.m * params.n; + const std::size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= elements) + return; + + const std::size_t col = idx % params.n; + double real = static_cast(accumulator[idx]) * params.inputScale * params.weightScale; + if (params.hasBias && bias != nullptr) { + const auto bq = QuantizedCudaQuantizeClamp(static_cast(bias[col]), params.biasScale, + params.biasZeroPoint, params.biasQMin, params.biasQMax); + real += static_cast(bq - params.biasZeroPoint) * params.biasScale; + } + + auto yq = QuantizedCudaQuantizeClamp(real, params.outputScale, params.outputZeroPoint, + params.outputQMin, params.outputQMax); + if (params.hasRelu && yq < params.outputZeroPoint) + yq = params.outputZeroPoint; + yq = QuantizedCudaClamp(yq, params.outputQMin, params.outputQMax); + output[idx] = static_cast(static_cast(yq - params.outputZeroPoint) * params.outputScale); +} + +inline void SetRowMajorLayout(cublasLtMatrixLayout_t layout) +{ + const cublasLtOrder_t order = CUBLASLT_ORDER_ROW; + CheckCublasLtStatus(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_ORDER, &order, sizeof(order)), + "cublasLtMatrixLayoutSetAttribute(row-major)"); +} + +inline cublasLtMatrixLayout_t CreateRowMajorLayout(cudaDataType_t type, std::uint64_t rows, std::uint64_t cols, + std::int64_t leadingDimension) +{ + cublasLtMatrixLayout_t layout = nullptr; + CheckCublasLtStatus(cublasLtMatrixLayoutCreate(&layout, type, rows, cols, leadingDimension), + "cublasLtMatrixLayoutCreate"); + SetRowMajorLayout(layout); + return layout; +} + +inline void DestroyLayout(cublasLtMatrixLayout_t &layout) +{ + if (layout != nullptr) { + cublasLtMatrixLayoutDestroy(layout); + layout = nullptr; + } +} + +} // namespace INTERNAL + +struct QuantizedGemmCudaLtProfile { + float inputQuantizeMs = 0.0f; + float gemmMs = 0.0f; + float epilogueMs = 0.0f; + float totalMs = 0.0f; + std::size_t workspaceSize = 0; + int heuristicResultCount = 0; + int selectedHeuristicIndex = 0; + float autotuneMs = 0.0f; + int autotunedCandidateCount = 0; + float selectedCandidateMs = 0.0f; + std::size_t accumulatorBytes = 0; + std::size_t outputBytes = 0; + std::size_t epilogueTrafficBytes = 0; + bool directOutputCarrier = false; +}; + +inline bool &QuantizedGemmCudaLt_ProfileEnabledStorage() +{ + static bool enabled = false; + return enabled; +} + +inline QuantizedGemmCudaLtProfile &QuantizedGemmCudaLt_LastProfileStorage() +{ + static QuantizedGemmCudaLtProfile profile{}; + return profile; +} + +inline void QuantizedGemmCudaLt_SetProfiling(bool enabled) +{ + QuantizedGemmCudaLt_ProfileEnabledStorage() = enabled; +} + +inline bool QuantizedGemmCudaLt_ProfilingEnabled() +{ + return QuantizedGemmCudaLt_ProfileEnabledStorage(); +} + +inline QuantizedGemmCudaLtProfile QuantizedGemmCudaLt_GetLastProfile() +{ + return QuantizedGemmCudaLt_LastProfileStorage(); +} + +inline void QuantizedGemmCudaPadInt8Matrix(QuantizedGemmCudaStream stream, const std::int8_t *input, + std::int8_t *padded, std::size_t logicalRows, + std::size_t logicalCols, std::size_t physicalRows, + std::size_t physicalCols, std::int8_t padValue = 0) +{ + const std::size_t elements = physicalRows * physicalCols; + if (elements == 0) + return; + constexpr int blockSize = 256; + const int gridSize = static_cast((elements + blockSize - 1) / blockSize); + INTERNAL::QuantizedGemmCudaPadInt8MatrixKernel<<>>( + input, padded, logicalRows, logicalCols, physicalRows, physicalCols, padValue); + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaPadInt8MatrixKernel"); +} + +template +inline void QuantizedGemmCudaUnpadMatrix(QuantizedGemmCudaStream stream, const T *padded, T *output, + std::size_t logicalRows, std::size_t logicalCols, + std::size_t physicalRows, std::size_t physicalCols) +{ + const std::size_t elements = logicalRows * logicalCols; + if (elements == 0) + return; + constexpr int blockSize = 256; + const int gridSize = static_cast((elements + blockSize - 1) / blockSize); + INTERNAL::QuantizedGemmCudaUnpadMatrixKernel<<>>( + padded, output, logicalRows, logicalCols, physicalRows, physicalCols); + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaUnpadMatrixKernel"); +} + +inline void QuantizedGemmCudaUnpadInt8Matrix(QuantizedGemmCudaStream stream, const std::int8_t *padded, + std::int8_t *output, std::size_t logicalRows, + std::size_t logicalCols, std::size_t physicalRows, + std::size_t physicalCols) +{ + QuantizedGemmCudaUnpadMatrix(stream, padded, output, logicalRows, logicalCols, physicalRows, physicalCols); +} + +inline void QuantizedGemmCudaUnpadUInt8Matrix(QuantizedGemmCudaStream stream, const std::uint8_t *padded, + std::uint8_t *output, std::size_t logicalRows, + std::size_t logicalCols, std::size_t physicalRows, + std::size_t physicalCols) +{ + QuantizedGemmCudaUnpadMatrix(stream, padded, output, logicalRows, logicalCols, physicalRows, physicalCols); +} + +struct QuantizedGemmCudaLtEventTimer { + cudaEvent_t totalStart = nullptr; + cudaEvent_t inputStop = nullptr; + cudaEvent_t gemmStop = nullptr; + cudaEvent_t epilogueStop = nullptr; + bool active = false; + + explicit QuantizedGemmCudaLtEventTimer(bool enabled) + : active(enabled) + { + if (!active) + return; + INTERNAL::CheckCudaStatus(cudaEventCreate(&totalStart), "cudaEventCreate(totalStart)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&inputStop), "cudaEventCreate(inputStop)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&gemmStop), "cudaEventCreate(gemmStop)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&epilogueStop), "cudaEventCreate(epilogueStop)"); + } + + QuantizedGemmCudaLtEventTimer(const QuantizedGemmCudaLtEventTimer &) = delete; + QuantizedGemmCudaLtEventTimer &operator=(const QuantizedGemmCudaLtEventTimer &) = delete; + + ~QuantizedGemmCudaLtEventTimer() + { + if (epilogueStop != nullptr) + cudaEventDestroy(epilogueStop); + if (gemmStop != nullptr) + cudaEventDestroy(gemmStop); + if (inputStop != nullptr) + cudaEventDestroy(inputStop); + if (totalStart != nullptr) + cudaEventDestroy(totalStart); + } + + void RecordStart(QuantizedGemmCudaStream stream) + { + if (active) + INTERNAL::CheckCudaStatus(cudaEventRecord(totalStart, stream), "cudaEventRecord(totalStart)"); + } + + void RecordInputStop(QuantizedGemmCudaStream stream) + { + if (active) + INTERNAL::CheckCudaStatus(cudaEventRecord(inputStop, stream), "cudaEventRecord(inputStop)"); + } + + void RecordGemmStop(QuantizedGemmCudaStream stream) + { + if (active) + INTERNAL::CheckCudaStatus(cudaEventRecord(gemmStop, stream), "cudaEventRecord(gemmStop)"); + } + + void RecordEpilogueStop(QuantizedGemmCudaStream stream, std::size_t workspaceSize, + int heuristicResultCount, int selectedHeuristicIndex, + float autotuneMs, int autotunedCandidateCount, float selectedCandidateMs, + std::size_t accumulatorBytes, std::size_t outputBytes, bool directOutputCarrier) + { + if (!active) + return; + INTERNAL::CheckCudaStatus(cudaEventRecord(epilogueStop, stream), "cudaEventRecord(epilogueStop)"); + INTERNAL::CheckCudaStatus(cudaEventSynchronize(epilogueStop), "cudaEventSynchronize(epilogueStop)"); + + QuantizedGemmCudaLtProfile profile{}; + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&profile.inputQuantizeMs, totalStart, inputStop), + "cudaEventElapsedTime(inputQuantize)"); + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&profile.gemmMs, inputStop, gemmStop), + "cudaEventElapsedTime(gemm)"); + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&profile.epilogueMs, gemmStop, epilogueStop), + "cudaEventElapsedTime(epilogue)"); + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&profile.totalMs, totalStart, epilogueStop), + "cudaEventElapsedTime(total)"); + profile.workspaceSize = workspaceSize; + profile.heuristicResultCount = heuristicResultCount; + profile.selectedHeuristicIndex = selectedHeuristicIndex; + profile.autotuneMs = autotuneMs; + profile.autotunedCandidateCount = autotunedCandidateCount; + profile.selectedCandidateMs = selectedCandidateMs; + profile.accumulatorBytes = accumulatorBytes; + profile.outputBytes = outputBytes; + profile.directOutputCarrier = directOutputCarrier; + profile.epilogueTrafficBytes = directOutputCarrier ? outputBytes : (accumulatorBytes + outputBytes); + QuantizedGemmCudaLt_LastProfileStorage() = profile; + } +}; + +struct QuantizedGemmCudaLtState { + cublasLtHandle_t fHandle = nullptr; + cublasLtMatmulDesc_t fOperation = nullptr; + cublasLtMatrixLayout_t fALayout = nullptr; + cublasLtMatrixLayout_t fBLayout = nullptr; + cublasLtMatrixLayout_t fCLayout = nullptr; + cublasLtMatmulPreference_t fPreference = nullptr; + static constexpr int kMaxHeuristicResults = 8; + cublasLtMatmulHeuristicResult_t fHeuristicResults[kMaxHeuristicResults]{}; + cublasLtMatmulHeuristicResult_t fHeuristic{}; + int fHeuristicResultCount = 0; + int fSelectedHeuristicIndex = 0; + std::size_t fWorkspaceSize = 0; + std::size_t fWorkspaceAllocatedBytes = 0; + std::size_t fWorkspaceLimitBytes = 0; + void *fWorkspace = nullptr; + bool fAutotuned = false; + float fAutotuneMs = 0.0f; + int fAutotunedCandidateCount = 0; + float fSelectedCandidateMs = 0.0f; + std::int8_t *fInputQuantized = nullptr; + std::int32_t *fAccumulator = nullptr; + void *fOutputQuantized = nullptr; + float *fBiasOutputOffset = nullptr; + std::size_t fInputQuantizedBytes = 0; + std::size_t fAccumulatorBytes = 0; + std::size_t fOutputQuantizedBytes = 0; + std::size_t fBiasOutputOffsetBytes = 0; + const float *fBiasOutputOffsetSource = nullptr; + double fBiasOutputOffsetBiasScale = 0.0; + double fBiasOutputOffsetOutputScale = 0.0; + std::int32_t fBiasOutputOffsetBiasZeroPoint = 0; + std::int32_t fBiasOutputOffsetOutputZeroPoint = 0; + std::int32_t fBiasOutputOffsetBiasQMin = 0; + std::int32_t fBiasOutputOffsetBiasQMax = 0; + std::size_t fBiasOutputOffsetN = 0; + const float *fBiasOutputOffsetWeightScaleVector = nullptr; + EQuantizedCudaScaleMode fBiasOutputOffsetWeightScaleMode = EQuantizedCudaScaleMode::PerTensor; + std::size_t fM = 0; + std::size_t fN = 0; + std::size_t fK = 0; + bool fInitialized = false; + + QuantizedGemmCudaLtState() = default; + QuantizedGemmCudaLtState(const QuantizedGemmCudaLtState &) = delete; + QuantizedGemmCudaLtState &operator=(const QuantizedGemmCudaLtState &) = delete; + + QuantizedGemmCudaLtState(QuantizedGemmCudaLtState &&other) noexcept + { + MoveFrom(other); + } + + QuantizedGemmCudaLtState &operator=(QuantizedGemmCudaLtState &&other) noexcept + { + if (this != &other) { + Reset(); + MoveFrom(other); + } + return *this; + } + + ~QuantizedGemmCudaLtState() { Reset(); } + + void Reset() noexcept + { + if (fPreference != nullptr) { + cublasLtMatmulPreferenceDestroy(fPreference); + fPreference = nullptr; + } + INTERNAL::DestroyLayout(fCLayout); + INTERNAL::DestroyLayout(fBLayout); + INTERNAL::DestroyLayout(fALayout); + if (fOperation != nullptr) { + cublasLtMatmulDescDestroy(fOperation); + fOperation = nullptr; + } + if (fHandle != nullptr) { + cublasLtDestroy(fHandle); + fHandle = nullptr; + } + if (fWorkspace != nullptr) { + cudaFree(fWorkspace); + fWorkspace = nullptr; + } + if (fInputQuantized != nullptr) { + cudaFree(fInputQuantized); + fInputQuantized = nullptr; + } + if (fAccumulator != nullptr) { + cudaFree(fAccumulator); + fAccumulator = nullptr; + } + if (fOutputQuantized != nullptr) { + cudaFree(fOutputQuantized); + fOutputQuantized = nullptr; + } + if (fBiasOutputOffset != nullptr) { + cudaFree(fBiasOutputOffset); + fBiasOutputOffset = nullptr; + } + for (auto &heuristic : fHeuristicResults) + heuristic = cublasLtMatmulHeuristicResult_t{}; + fHeuristic = cublasLtMatmulHeuristicResult_t{}; + fHeuristicResultCount = 0; + fSelectedHeuristicIndex = 0; + fWorkspaceSize = 0; + fWorkspaceAllocatedBytes = 0; + fWorkspaceLimitBytes = 0; + fWorkspace = nullptr; + fAutotuned = false; + fAutotuneMs = 0.0f; + fAutotunedCandidateCount = 0; + fSelectedCandidateMs = 0.0f; + fInputQuantizedBytes = 0; + fAccumulatorBytes = 0; + fOutputQuantizedBytes = 0; + fBiasOutputOffsetBytes = 0; + fBiasOutputOffsetSource = nullptr; + fBiasOutputOffsetBiasScale = 0.0; + fBiasOutputOffsetOutputScale = 0.0; + fBiasOutputOffsetBiasZeroPoint = 0; + fBiasOutputOffsetOutputZeroPoint = 0; + fBiasOutputOffsetBiasQMin = 0; + fBiasOutputOffsetBiasQMax = 0; + fBiasOutputOffsetN = 0; + fBiasOutputOffsetWeightScaleVector = nullptr; + fBiasOutputOffsetWeightScaleMode = EQuantizedCudaScaleMode::PerTensor; + fM = 0; + fN = 0; + fK = 0; + fInitialized = false; + } + + void Initialize(const QuantizedGemmCudaLtParams ¶ms) + { + if (fInitialized && fM == params.m && fN == params.n && fK == params.k && + fWorkspaceLimitBytes == params.maxWorkspaceBytes) + return; + + Reset(); + try { + INTERNAL::CheckCublasLtStatus(cublasLtCreate(&fHandle), "cublasLtCreate"); + INTERNAL::CheckCublasLtStatus(cublasLtMatmulDescCreate(&fOperation, CUBLAS_COMPUTE_32I, CUDA_R_32I), + "cublasLtMatmulDescCreate"); + + const cublasOperation_t transA = CUBLAS_OP_N; + const cublasOperation_t transB = CUBLAS_OP_T; + INTERNAL::CheckCublasLtStatus(cublasLtMatmulDescSetAttribute(fOperation, CUBLASLT_MATMUL_DESC_TRANSA, + &transA, sizeof(transA)), + "cublasLtMatmulDescSetAttribute(transA)"); + INTERNAL::CheckCublasLtStatus(cublasLtMatmulDescSetAttribute(fOperation, CUBLASLT_MATMUL_DESC_TRANSB, + &transB, sizeof(transB)), + "cublasLtMatmulDescSetAttribute(transB)"); + + fALayout = INTERNAL::CreateRowMajorLayout(CUDA_R_8I, params.m, params.k, + static_cast(params.k)); + fBLayout = INTERNAL::CreateRowMajorLayout(CUDA_R_8I, params.n, params.k, + static_cast(params.k)); + fCLayout = INTERNAL::CreateRowMajorLayout(CUDA_R_32I, params.m, params.n, + static_cast(params.n)); + + INTERNAL::CheckCublasLtStatus(cublasLtMatmulPreferenceCreate(&fPreference), + "cublasLtMatmulPreferenceCreate"); + fWorkspaceLimitBytes = params.maxWorkspaceBytes; + INTERNAL::CheckCublasLtStatus(cublasLtMatmulPreferenceSetAttribute( + fPreference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &fWorkspaceLimitBytes, sizeof(fWorkspaceLimitBytes)), + "cublasLtMatmulPreferenceSetAttribute(workspace)"); + + INTERNAL::CheckCublasLtStatus(cublasLtMatmulAlgoGetHeuristic(fHandle, fOperation, fALayout, fBLayout, + fCLayout, fCLayout, fPreference, + kMaxHeuristicResults, fHeuristicResults, + &fHeuristicResultCount), + "cublasLtMatmulAlgoGetHeuristic"); + if (fHeuristicResultCount == 0) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM found no algorithm for the selected int8 shape"); + } + + fSelectedHeuristicIndex = 0; + fHeuristic = fHeuristicResults[fSelectedHeuristicIndex]; + fWorkspaceSize = fHeuristic.workspaceSize; + for (int i = 0; i < fHeuristicResultCount; ++i) { + if (fHeuristicResults[i].workspaceSize > fWorkspaceAllocatedBytes) + fWorkspaceAllocatedBytes = fHeuristicResults[i].workspaceSize; + } + if (fWorkspaceAllocatedBytes > 0) { + INTERNAL::CheckCudaStatus(cudaMalloc(&fWorkspace, fWorkspaceAllocatedBytes), "cudaMalloc(cuBLASLt workspace)"); + } + + fM = params.m; + fN = params.n; + fK = params.k; + fInitialized = true; + } catch (...) { + Reset(); + throw; + } + } + + void EnsureTemporaryBuffers(const QuantizedGemmCudaLtParams ¶ms) + { + const std::size_t requiredInputBytes = params.m * params.k * sizeof(std::int8_t); + if (requiredInputBytes > fInputQuantizedBytes) { + if (fInputQuantized != nullptr) { + INTERNAL::CheckCudaStatus(cudaFree(fInputQuantized), "cudaFree(inputQuantized cache)"); + fInputQuantized = nullptr; + fInputQuantizedBytes = 0; + } + INTERNAL::CheckCudaStatus(cudaMalloc(&fInputQuantized, requiredInputBytes), + "cudaMalloc(inputQuantized cache)"); + fInputQuantizedBytes = requiredInputBytes; + } + + const std::size_t requiredAccumulatorBytes = params.m * params.n * sizeof(std::int32_t); + if (requiredAccumulatorBytes > fAccumulatorBytes) { + if (fAccumulator != nullptr) { + INTERNAL::CheckCudaStatus(cudaFree(fAccumulator), "cudaFree(accumulator cache)"); + fAccumulator = nullptr; + fAccumulatorBytes = 0; + } + INTERNAL::CheckCudaStatus(cudaMalloc(&fAccumulator, requiredAccumulatorBytes), + "cudaMalloc(accumulator cache)"); + fAccumulatorBytes = requiredAccumulatorBytes; + } + + if (params.paddedExecution && params.epilogueMode == EQuantizedCudaEpilogueMode::Quantized) { + const std::size_t outputElementSize = params.outputCarrier == EQuantizedCudaOutputCarrier::UInt8 + ? sizeof(std::uint8_t) : sizeof(std::int8_t); + const std::size_t requiredOutputBytes = params.m * params.n * outputElementSize; + if (requiredOutputBytes > fOutputQuantizedBytes) { + if (fOutputQuantized != nullptr) { + INTERNAL::CheckCudaStatus(cudaFree(fOutputQuantized), "cudaFree(padded output cache)"); + fOutputQuantized = nullptr; + fOutputQuantizedBytes = 0; + } + INTERNAL::CheckCudaStatus(cudaMalloc(&fOutputQuantized, requiredOutputBytes), + "cudaMalloc(padded output cache)"); + fOutputQuantizedBytes = requiredOutputBytes; + } + } + } + + const float *EnsureBiasOutputOffsetBuffer(const QuantizedGemmCudaLtParams ¶ms, const float *bias, + const float *weightScaleVector, QuantizedGemmCudaStream stream) + { + if (bias == nullptr || !params.hasBias) + return nullptr; + + const std::size_t requiredBytes = params.n * sizeof(float); + if (requiredBytes > fBiasOutputOffsetBytes) { + if (fBiasOutputOffset != nullptr) { + INTERNAL::CheckCudaStatus(cudaFree(fBiasOutputOffset), "cudaFree(bias output offset cache)"); + fBiasOutputOffset = nullptr; + fBiasOutputOffsetBytes = 0; + } + INTERNAL::CheckCudaStatus(cudaMalloc(&fBiasOutputOffset, requiredBytes), + "cudaMalloc(bias output offset cache)"); + fBiasOutputOffsetBytes = requiredBytes; + fBiasOutputOffsetSource = nullptr; + } + + const bool cacheValid = fBiasOutputOffsetSource == bias && fBiasOutputOffsetN == params.n && + fBiasOutputOffsetBiasScale == params.biasScale && + fBiasOutputOffsetOutputScale == params.outputScale && + fBiasOutputOffsetBiasZeroPoint == params.biasZeroPoint && + fBiasOutputOffsetOutputZeroPoint == params.outputZeroPoint && + fBiasOutputOffsetBiasQMin == params.biasQMin && + fBiasOutputOffsetBiasQMax == params.biasQMax && + fBiasOutputOffsetWeightScaleVector == weightScaleVector && + fBiasOutputOffsetWeightScaleMode == params.weightScaleMode; + if (!cacheValid) { + constexpr int threads = 256; + const int blocks = static_cast((params.n + threads - 1) / threads); + INTERNAL::QuantizedGemmCudaBiasOutputOffsetKernel<<>>( + fBiasOutputOffset, bias, weightScaleVector, params); + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaBiasOutputOffsetKernel launch"); + fBiasOutputOffsetSource = bias; + fBiasOutputOffsetN = params.n; + fBiasOutputOffsetBiasScale = params.biasScale; + fBiasOutputOffsetOutputScale = params.outputScale; + fBiasOutputOffsetBiasZeroPoint = params.biasZeroPoint; + fBiasOutputOffsetOutputZeroPoint = params.outputZeroPoint; + fBiasOutputOffsetBiasQMin = params.biasQMin; + fBiasOutputOffsetBiasQMax = params.biasQMax; + fBiasOutputOffsetWeightScaleVector = weightScaleVector; + fBiasOutputOffsetWeightScaleMode = params.weightScaleMode; + } + + return fBiasOutputOffset; + } + + std::int8_t *InputQuantizedBuffer() const { return fInputQuantized; } + std::int32_t *AccumulatorBuffer() const { return fAccumulator; } + void *OutputQuantizedBuffer() const { return fOutputQuantized; } + std::size_t AccumulatorBytes() const { return fAccumulatorBytes; } + std::size_t WorkspaceSize() const { return fWorkspaceSize; } + int HeuristicResultCount() const { return fHeuristicResultCount; } + int SelectedHeuristicIndex() const { return fSelectedHeuristicIndex; } + float AutotuneMs() const { return fAutotuneMs; } + int AutotunedCandidateCount() const { return fAutotunedCandidateCount; } + float SelectedCandidateMs() const { return fSelectedCandidateMs; } + + void Autotune(std::int32_t *accumulator, const std::int8_t *inputQuantized, const std::int8_t *weightQuantized, + const QuantizedGemmCudaLtParams ¶ms, QuantizedGemmCudaStream stream) + { + if (fAutotuned || !params.enableAutotuning || fHeuristicResultCount <= 1) { + fAutotuned = true; + return; + } + + cudaEvent_t totalStart = nullptr; + cudaEvent_t totalStop = nullptr; + cudaEvent_t candidateStart = nullptr; + cudaEvent_t candidateStop = nullptr; + INTERNAL::CheckCudaStatus(cudaEventCreate(&totalStart), "cudaEventCreate(autotuneTotalStart)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&totalStop), "cudaEventCreate(autotuneTotalStop)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&candidateStart), "cudaEventCreate(autotuneCandidateStart)"); + INTERNAL::CheckCudaStatus(cudaEventCreate(&candidateStop), "cudaEventCreate(autotuneCandidateStop)"); + + const std::int32_t alpha = 1; + const std::int32_t beta = 0; + const int iterations = params.autotuneIterations > 0 ? params.autotuneIterations : 1; + float bestMs = 0.0f; + int bestIndex = fSelectedHeuristicIndex; + int measuredCandidates = 0; + INTERNAL::CheckCudaStatus(cudaEventRecord(totalStart, stream), "cudaEventRecord(autotuneTotalStart)"); + for (int i = 0; i < fHeuristicResultCount; ++i) { + const auto warmupStatus = cublasLtMatmul(fHandle, fOperation, &alpha, inputQuantized, fALayout, + weightQuantized, fBLayout, &beta, accumulator, fCLayout, + accumulator, fCLayout, &fHeuristicResults[i].algo, fWorkspace, + fWorkspaceAllocatedBytes, stream); + if (warmupStatus != CUBLAS_STATUS_SUCCESS) + continue; + INTERNAL::CheckCudaStatus(cudaEventRecord(candidateStart, stream), "cudaEventRecord(autotuneCandidateStart)"); + bool candidateOk = true; + for (int iteration = 0; iteration < iterations; ++iteration) { + const auto status = cublasLtMatmul(fHandle, fOperation, &alpha, inputQuantized, fALayout, + weightQuantized, fBLayout, &beta, accumulator, fCLayout, + accumulator, fCLayout, &fHeuristicResults[i].algo, fWorkspace, + fWorkspaceAllocatedBytes, stream); + if (status != CUBLAS_STATUS_SUCCESS) { + candidateOk = false; + break; + } + } + if (!candidateOk) + continue; + INTERNAL::CheckCudaStatus(cudaEventRecord(candidateStop, stream), "cudaEventRecord(autotuneCandidateStop)"); + INTERNAL::CheckCudaStatus(cudaEventSynchronize(candidateStop), "cudaEventSynchronize(autotuneCandidateStop)"); + float candidateMs = 0.0f; + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&candidateMs, candidateStart, candidateStop), + "cudaEventElapsedTime(autotuneCandidate)"); + candidateMs /= static_cast(iterations); + ++measuredCandidates; + if (measuredCandidates == 1 || candidateMs < bestMs) { + bestMs = candidateMs; + bestIndex = i; + } + } + INTERNAL::CheckCudaStatus(cudaEventRecord(totalStop, stream), "cudaEventRecord(autotuneTotalStop)"); + INTERNAL::CheckCudaStatus(cudaEventSynchronize(totalStop), "cudaEventSynchronize(autotuneTotalStop)"); + INTERNAL::CheckCudaStatus(cudaEventElapsedTime(&fAutotuneMs, totalStart, totalStop), + "cudaEventElapsedTime(autotuneTotal)"); + + cudaEventDestroy(candidateStop); + cudaEventDestroy(candidateStart); + cudaEventDestroy(totalStop); + cudaEventDestroy(totalStart); + + if (measuredCandidates > 0) { + fSelectedHeuristicIndex = bestIndex; + fHeuristic = fHeuristicResults[fSelectedHeuristicIndex]; + fWorkspaceSize = fHeuristic.workspaceSize; + fSelectedCandidateMs = bestMs; + fAutotunedCandidateCount = measuredCandidates; + } + fAutotuned = true; + } + + void Execute(std::int32_t *accumulator, const std::int8_t *inputQuantized, const std::int8_t *weightQuantized, + const QuantizedGemmCudaLtParams ¶ms, QuantizedGemmCudaStream stream) + { + Initialize(params); + Autotune(accumulator, inputQuantized, weightQuantized, params, stream); + const std::int32_t alpha = 1; + const std::int32_t beta = 0; + INTERNAL::CheckCublasLtStatus(cublasLtMatmul(fHandle, fOperation, &alpha, inputQuantized, fALayout, + weightQuantized, fBLayout, &beta, accumulator, fCLayout, + accumulator, fCLayout, &fHeuristic.algo, fWorkspace, + fWorkspaceAllocatedBytes, stream), + "cublasLtMatmul"); + } + +private: + void MoveFrom(QuantizedGemmCudaLtState &other) noexcept + { + fHandle = other.fHandle; + fOperation = other.fOperation; + fALayout = other.fALayout; + fBLayout = other.fBLayout; + fCLayout = other.fCLayout; + fPreference = other.fPreference; + for (int i = 0; i < kMaxHeuristicResults; ++i) + fHeuristicResults[i] = other.fHeuristicResults[i]; + fHeuristic = other.fHeuristic; + fHeuristicResultCount = other.fHeuristicResultCount; + fSelectedHeuristicIndex = other.fSelectedHeuristicIndex; + fWorkspaceSize = other.fWorkspaceSize; + fWorkspaceAllocatedBytes = other.fWorkspaceAllocatedBytes; + fWorkspaceLimitBytes = other.fWorkspaceLimitBytes; + fWorkspace = other.fWorkspace; + fAutotuned = other.fAutotuned; + fAutotuneMs = other.fAutotuneMs; + fAutotunedCandidateCount = other.fAutotunedCandidateCount; + fSelectedCandidateMs = other.fSelectedCandidateMs; + fInputQuantized = other.fInputQuantized; + fAccumulator = other.fAccumulator; + fOutputQuantized = other.fOutputQuantized; + fBiasOutputOffset = other.fBiasOutputOffset; + fInputQuantizedBytes = other.fInputQuantizedBytes; + fAccumulatorBytes = other.fAccumulatorBytes; + fOutputQuantizedBytes = other.fOutputQuantizedBytes; + fBiasOutputOffsetBytes = other.fBiasOutputOffsetBytes; + fBiasOutputOffsetSource = other.fBiasOutputOffsetSource; + fBiasOutputOffsetBiasScale = other.fBiasOutputOffsetBiasScale; + fBiasOutputOffsetOutputScale = other.fBiasOutputOffsetOutputScale; + fBiasOutputOffsetBiasZeroPoint = other.fBiasOutputOffsetBiasZeroPoint; + fBiasOutputOffsetOutputZeroPoint = other.fBiasOutputOffsetOutputZeroPoint; + fBiasOutputOffsetBiasQMin = other.fBiasOutputOffsetBiasQMin; + fBiasOutputOffsetBiasQMax = other.fBiasOutputOffsetBiasQMax; + fBiasOutputOffsetN = other.fBiasOutputOffsetN; + fBiasOutputOffsetWeightScaleVector = other.fBiasOutputOffsetWeightScaleVector; + fBiasOutputOffsetWeightScaleMode = other.fBiasOutputOffsetWeightScaleMode; + fM = other.fM; + fN = other.fN; + fK = other.fK; + fInitialized = other.fInitialized; + + other.fHandle = nullptr; + other.fOperation = nullptr; + other.fALayout = nullptr; + other.fBLayout = nullptr; + other.fCLayout = nullptr; + other.fPreference = nullptr; + for (auto &heuristic : other.fHeuristicResults) + heuristic = cublasLtMatmulHeuristicResult_t{}; + other.fHeuristic = cublasLtMatmulHeuristicResult_t{}; + other.fHeuristicResultCount = 0; + other.fSelectedHeuristicIndex = 0; + other.fWorkspaceSize = 0; + other.fWorkspaceAllocatedBytes = 0; + other.fWorkspaceLimitBytes = 0; + other.fWorkspace = nullptr; + other.fAutotuned = false; + other.fAutotuneMs = 0.0f; + other.fAutotunedCandidateCount = 0; + other.fSelectedCandidateMs = 0.0f; + other.fInputQuantized = nullptr; + other.fAccumulator = nullptr; + other.fOutputQuantized = nullptr; + other.fBiasOutputOffset = nullptr; + other.fInputQuantizedBytes = 0; + other.fAccumulatorBytes = 0; + other.fOutputQuantizedBytes = 0; + other.fBiasOutputOffsetBytes = 0; + other.fBiasOutputOffsetSource = nullptr; + other.fBiasOutputOffsetBiasScale = 0.0; + other.fBiasOutputOffsetOutputScale = 0.0; + other.fBiasOutputOffsetBiasZeroPoint = 0; + other.fBiasOutputOffsetOutputZeroPoint = 0; + other.fBiasOutputOffsetBiasQMin = 0; + other.fBiasOutputOffsetBiasQMax = 0; + other.fBiasOutputOffsetN = 0; + other.fBiasOutputOffsetWeightScaleVector = nullptr; + other.fBiasOutputOffsetWeightScaleMode = EQuantizedCudaScaleMode::PerTensor; + other.fM = 0; + other.fN = 0; + other.fK = 0; + other.fInitialized = false; + } +}; + +inline void QuantizedGemmCudaLt_Call(QuantizedGemmCudaLtState &state, QuantizedGemmCudaStream stream, + void *output, const void *input, const void *weight, const float *bias, + const float *weightScaleVector, + const QuantizedGemmCudaLtParams ¶ms) +{ + if (output == nullptr || input == nullptr || weight == nullptr) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM received a null required pointer"); + } + if (params.m == 0 || params.n == 0 || params.k == 0) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM requires nonzero M, N, and K"); + } + if (params.weightType != EQuantizedCudaWeightType::Int8) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM currently supports signed int8 weights only"); + } + if (params.inputZeroPoint != 0 || params.weightZeroPoint != 0) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM currently requires input and weight zero points to be 0"); + } + if (params.weightScaleMode == EQuantizedCudaScaleMode::PerOutputChannel && weightScaleVector == nullptr) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM per-channel weight scale mode requires a scale vector"); + } + if (params.epilogueMode == EQuantizedCudaEpilogueMode::Quantized && + params.outputCarrier == EQuantizedCudaOutputCarrier::Float) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM quantized epilogue requires an integer output carrier"); + } + if (params.epilogueMode != EQuantizedCudaEpilogueMode::Quantized && + params.outputCarrier != EQuantizedCudaOutputCarrier::Float) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM float epilogues require a float output carrier"); + } + + QuantizedGemmCudaLtParams effectiveParams = params; + if (effectiveParams.logicalM == 0) + effectiveParams.logicalM = effectiveParams.m; + if (effectiveParams.logicalN == 0) + effectiveParams.logicalN = effectiveParams.n; + if (effectiveParams.logicalK == 0) + effectiveParams.logicalK = effectiveParams.k; + if (effectiveParams.paddedExecution && + (effectiveParams.logicalM > effectiveParams.m || effectiveParams.logicalN > effectiveParams.n || + effectiveParams.logicalK > effectiveParams.k)) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM padded execution received logical dimensions larger than physical dimensions"); + } + if (effectiveParams.accumulatorToOutputScale == 0.0) + effectiveParams.accumulatorToOutputScale = + (effectiveParams.inputScale * effectiveParams.weightScale) / effectiveParams.outputScale; + + const std::size_t inputElements = effectiveParams.m * effectiveParams.k; + const std::size_t outputElements = effectiveParams.m * effectiveParams.n; + const std::size_t logicalOutputElements = effectiveParams.logicalM * effectiveParams.logicalN; + state.Initialize(effectiveParams); + state.EnsureTemporaryBuffers(effectiveParams); + QuantizedGemmCudaLtEventTimer profileTimer(QuantizedGemmCudaLt_ProfilingEnabled()); + profileTimer.RecordStart(stream); + + constexpr int threads = 256; + const std::int8_t *inputQuantized = nullptr; + if (effectiveParams.inputCarrier == EQuantizedCudaInputCarrier::Int8) { + inputQuantized = static_cast(input); + if (effectiveParams.paddedExecution) { + QuantizedGemmCudaPadInt8Matrix(stream, inputQuantized, state.InputQuantizedBuffer(), + effectiveParams.logicalM, effectiveParams.logicalK, + effectiveParams.m, effectiveParams.k, 0); + inputQuantized = state.InputQuantizedBuffer(); + } + } else { + if (effectiveParams.paddedExecution) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM padded execution currently requires int8 input carriers"); + } + const int inputBlocks = static_cast((inputElements + threads - 1) / threads); + const float *inputFloat = static_cast(input); + INTERNAL::QuantizedGemmCudaQuantizeInputKernel<<>>( + inputFloat, state.InputQuantizedBuffer(), inputElements, effectiveParams.inputScale, effectiveParams.inputZeroPoint, + effectiveParams.inputQMin, effectiveParams.inputQMax); + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaQuantizeInputKernel launch"); + inputQuantized = state.InputQuantizedBuffer(); + } + profileTimer.RecordInputStop(stream); + + state.Execute(state.AccumulatorBuffer(), inputQuantized, static_cast(weight), effectiveParams, + stream); + profileTimer.RecordGemmStop(stream); + + const int outputBlocks = static_cast((outputElements + threads - 1) / threads); + std::size_t outputBytes = outputElements * sizeof(float); + bool directOutputCarrier = false; + if (effectiveParams.epilogueMode == EQuantizedCudaEpilogueMode::Quantized) { + directOutputCarrier = true; + const float *biasOutputOffset = state.EnsureBiasOutputOffsetBuffer(effectiveParams, bias, weightScaleVector, stream); + + outputBytes = logicalOutputElements * (effectiveParams.outputCarrier == EQuantizedCudaOutputCarrier::UInt8 ? sizeof(std::uint8_t) : sizeof(std::int8_t)); + void *epilogueOutput = effectiveParams.paddedExecution ? state.OutputQuantizedBuffer() : output; + if (epilogueOutput == nullptr) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM padded execution did not allocate an output scratch buffer"); + } + if (effectiveParams.outputCarrier == EQuantizedCudaOutputCarrier::UInt8) { + auto *quantizedOutput = static_cast(epilogueOutput); + if (effectiveParams.hasBias && bias != nullptr) { + if (effectiveParams.hasRelu) { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), biasOutputOffset, weightScaleVector, effectiveParams); + } else { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), biasOutputOffset, weightScaleVector, effectiveParams); + } + } else { + if (effectiveParams.hasRelu) { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), nullptr, weightScaleVector, effectiveParams); + } else { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), nullptr, weightScaleVector, effectiveParams); + } + } + } else { + auto *quantizedOutput = static_cast(epilogueOutput); + if (effectiveParams.hasBias && bias != nullptr) { + if (effectiveParams.hasRelu) { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), biasOutputOffset, weightScaleVector, effectiveParams); + } else { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), biasOutputOffset, weightScaleVector, effectiveParams); + } + } else { + if (effectiveParams.hasRelu) { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), nullptr, weightScaleVector, effectiveParams); + } else { + INTERNAL::QuantizedGemmCudaQuantizedEpilogueKernel<<>>(quantizedOutput, state.AccumulatorBuffer(), nullptr, weightScaleVector, effectiveParams); + } + } + } + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaQuantizedEpilogueKernel launch"); + if (effectiveParams.paddedExecution) { + if (effectiveParams.outputCarrier == EQuantizedCudaOutputCarrier::UInt8) { + QuantizedGemmCudaUnpadUInt8Matrix(stream, static_cast(state.OutputQuantizedBuffer()), + static_cast(output), + effectiveParams.logicalM, effectiveParams.logicalN, + effectiveParams.m, effectiveParams.n); + } else { + QuantizedGemmCudaUnpadInt8Matrix(stream, static_cast(state.OutputQuantizedBuffer()), + static_cast(output), + effectiveParams.logicalM, effectiveParams.logicalN, + effectiveParams.m, effectiveParams.n); + } + } + } else { + if (effectiveParams.paddedExecution) { + throw std::runtime_error("SOFIE cuBLASLt quantized GEMM padded execution currently requires quantized output mode"); + } + auto *floatOutput = static_cast(output); + INTERNAL::QuantizedGemmCudaEpilogueKernel<<>>(floatOutput, state.AccumulatorBuffer(), bias, effectiveParams); + INTERNAL::CheckCudaStatus(cudaGetLastError(), "QuantizedGemmCudaEpilogueKernel launch"); + } + profileTimer.RecordEpilogueStop(stream, state.WorkspaceSize(), state.HeuristicResultCount(), + state.SelectedHeuristicIndex(), state.AutotuneMs(), + state.AutotunedCandidateCount(), state.SelectedCandidateMs(), + state.AccumulatorBytes(), outputBytes, directOutputCarrier); +} + + +#endif // SOFIE_USE_CUBLASLT + +} // namespace SOFIE + +#endif // SOFIE_QUANTIZED_ALPAKA diff --git a/core/inc/SOFIE/SOFIE_common.hxx b/core/inc/SOFIE/SOFIE_common.hxx index e36df0a..f5d1332 100644 --- a/core/inc/SOFIE/SOFIE_common.hxx +++ b/core/inc/SOFIE/SOFIE_common.hxx @@ -20,6 +20,7 @@ #include #include #include +#include namespace SOFIE { @@ -302,7 +303,7 @@ public: void SetNotWritable() { fIsNotWritable = true;} // set writable initialized tensors - i.e. tensor that must be written in a file void SetWritable() { fIsNotWritable = false;} - // set as constant (needed for non-float initialized tensors) + // embed an initialized tensor in generated code instead of the weight file void SetConstant() { fConstant = true;} template diff --git a/core/src/RModel.cxx b/core/src/RModel.cxx index 377171c..f41728c 100644 --- a/core/src/RModel.cxx +++ b/core/src/RModel.cxx @@ -3,6 +3,8 @@ #include #include #include +#include +#include #ifdef SOFIE_SUPPORT_ROOT_BINARY #include "TFile.h" @@ -10,6 +12,7 @@ #include "SOFIE/RModel.hxx" #include "SOFIE/RModelProfiler.hxx" +#include "SOFIE/RWeightFile.hxx" #include "SOFIE/SOFIE_common.hxx" namespace SOFIE { @@ -52,6 +55,14 @@ std::string TensorMember(std::string const &name) } // namespace +void RModel::BuildLoweredOperatorView(EQuantizedBackend backend) +{ + fLoweredOperators.clear(); + fLoweredConsumedOperatorIndices.clear(); + PrepareQuantizedTensorStorage(backend); + AddLoweredQuantizedOperators(backend); +} + std::vector RModel::GetTensorShape(const std::string & name) const { auto f = fReadyInputTensorInfos.find(name); if (f != fReadyInputTensorInfos.end()) { @@ -677,9 +688,11 @@ void RModel::Initialize(const std::map & inputParams, bool i++; } - // loop on initialized tensors and make the integers as constant to be - // not written in a weight file and check if the tensors flagged as not writable are really not writable, - // i.e. are not used by non constant operators + AnalyzeQuantizedRegions(); + + // Check whether tensors flagged as not writable are used by runtime operators. + // Integer initializers remain embedded for compatibility, except for physical + // quantized storage explicitly registered as file-backed model weights. for (auto &it : fInitializedTensors) { // check if not-writable tensors are really not writable, i.e. are not used by non constant operators if (it.second.IsNotWritable() && runtimeInitializedInputs.find(it.first) != runtimeInitializedInputs.end()) { @@ -688,9 +701,8 @@ void RModel::Initialize(const std::map & inputParams, bool std::cout << "Initialized tensor " << it.first << " is flagged as not writable but is used by non constant operators, set it as writable \n"; } } - // if the tensor is an integer we can flag it as constant since it will not be written in a weight file and it is considered equivalent as being created from a Constant operator - // only FLOAT tensors are written in a weight file - if (it.second.type() != ETensorType::FLOAT) { + const bool isFileBackedQuantizedStorage = HasQuantizedTensorStorage(it.first); + if (it.second.type() != ETensorType::FLOAT && !isFileBackedQuantizedStorage) { it.second.SetConstant(); } } @@ -800,7 +812,7 @@ void RModel::GenerateInitializedTensorInfo() for (auto &i : fInitializedTensors) { if (i.second.IsNotWritable()) continue; size_t length = ConvertShapeToLength(i.second.shape()); - if (!fUseWeightFile || i.second.IsConstantTensor() || !i.second.IsWeightTensor() || i.second.type() != ETensorType::FLOAT ) { + if (!fUseWeightFile || i.second.IsConstantTensor() || !i.second.IsWeightTensor()) { if (i.second.type() == ETensorType::FLOAT) { // check if NaN of Inf are inside tensor data bool hasInfOrNaN = false; @@ -823,19 +835,24 @@ void RModel::GenerateInitializedTensorInfo() } else if (i.second.type() == ETensorType::INT32) { fGC += GenerateConstantTensorCode(i); fConstantTensorSize += length * sizeof(int32_t); - } else if (i.second.type() == ETensorType::BOOL || i.second.type() == ETensorType::UINT8 ) { + } else if (i.second.type() == ETensorType::INT8) { + fGC += GenerateConstantTensorCode(i); + fConstantTensorSize += length * sizeof(int8_t); + } else if (i.second.type() == ETensorType::BOOL || i.second.type() == ETensorType::UINT8 ) { fGC += GenerateConstantTensorCode(i); fConstantTensorSize += length * sizeof(uint8_t); } } else { - // case of tensors which are read from a file - if (i.second.type() == ETensorType::FLOAT) { - fGC += "std::vector fTensor_" + i.first + " = std::vector(" + std::to_string(length) + ");\n"; - fGC += "float * " + TensorMember(i.first) + " = fTensor_" + i.first + ".data();\n"; - fWeightsTensorSize += length * sizeof(float); - } + const auto typeName = ConvertTypeToString(i.second.type()); + if (i.second.type() != ETensorType::FLOAT && i.second.type() != ETensorType::INT8 && + i.second.type() != ETensorType::UINT8 && i.second.type() != ETensorType::INT32 && + i.second.type() != ETensorType::INT64) + throw std::runtime_error("sofie initialized tensor " + i.first + " has a type unsupported by weight files"); + fGC += "std::vector<" + typeName + "> fTensor_" + i.first + " = std::vector<" + typeName + ">(" + std::to_string(length) + ");\n"; + fGC += typeName + " * " + TensorMember(i.first) + " = fTensor_" + i.first + ".data();\n"; + fWeightsTensorSize += length * GetTypeSize(i.second.type()); } } } @@ -1359,6 +1376,9 @@ void RModel::GenerateSessionCode() if (fWeightFile == WeightFileType::RootBinary) { fileName += ".root"; } + if (fWeightFile == WeightFileType::Binary) { + fileName += ".bin"; + } fGC += sessionName + "(std::string filename =\"" + fileName + "\""; } else { // no need to pass weight file since it is not used @@ -1439,8 +1459,13 @@ void RModel::GenerateSessionCode() for (size_t op_idx = 0; op_idx < fOperators.size(); ++op_idx) { if (fVerbose) std::cout << "Generating code for operator .... " << op_idx << std::endl; + if (!fProfile && fLoweredConsumedOperatorIndices.count(op_idx) != 0) { + continue; + } if (fProfile) { allOperatorCode += RModelProfiler::GenerateOperatorCode(*fOperators[op_idx], op_idx); + } else if (auto lowered = fLoweredOperators.find(op_idx); lowered != fLoweredOperators.end()) { + allOperatorCode += lowered->second->Generate(std::to_string(op_idx)); } else { allOperatorCode += fOperators[op_idx]->Generate(std::to_string(op_idx)); } @@ -1504,6 +1529,18 @@ void RModel::Generate(std::underlying_type_t options, int batchSize, lo fUseWeightFile = true; fWeightFile = WeightFileType::RootBinary; } + if (static_cast>(Options::kBinaryWeightFile) & options) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Binary; + } + if (static_cast>(Options::kTextWeightFile) & options) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Text; + } + const auto explicitWeightPolicy = static_cast>(Options::kNoWeightFile) | + static_cast>(Options::kRootBinaryWeightFile) | + static_cast>(Options::kBinaryWeightFile) | + static_cast>(Options::kTextWeightFile); if (fUseWeightFile && !fUseSession) { throw std::runtime_error( "sofie: RModel::Generate: cannot use a separate weight file without generating a Session class"); @@ -1519,6 +1556,15 @@ void RModel::Generate(std::underlying_type_t options, int batchSize, lo // initialize the model including all operators and sub-graphs Initialize(batchSize, verbose); + if ((options & explicitWeightPolicy) == 0 && !fQuantizationState.tensorInfos.empty()) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Binary; + } + if (fWeightFile == WeightFileType::Binary) + AddNeededCustomHeader("SOFIE/RWeightFile.hxx"); + BuildLoweredOperatorView(EQuantizedBackend::CPU); + + AddQuantizedGeneratedHeaders(EQuantizedBackend::CPU); // if having dynamic tensor we need to have a Session if (!fDynamicTensorInfos.empty()) { @@ -1577,7 +1623,9 @@ void RModel::ReadInitializedTensorsFromFile(long pos) { // skip Constant and shape tensors (not written in a file) if (!i.second.IsWeightTensor()) continue; std::string tensor_name = "tensor_" + i.first; - if (i.second.type() == ETensorType::FLOAT) { + if (i.second.type() == ETensorType::FLOAT || i.second.type() == ETensorType::INT8 || + i.second.type() == ETensorType::UINT8 || i.second.type() == ETensorType::INT32 || + i.second.type() == ETensorType::INT64) { std::string length = std::to_string(ConvertShapeToLength(i.second.shape())); fGC += " ReadTensorFromStream(f, " + tensor_name + ", \"" + tensor_name + "\", " + length + ");\n"; } else { @@ -1587,6 +1635,29 @@ void RModel::ReadInitializedTensorsFromFile(long pos) { fGC += " f.close();\n"; } + if (fWeightFile == WeightFileType::Binary) { + if (!fUseWeightFile) return; + if (fIsGNNComponent || pos != 0) + throw std::runtime_error("SOFIE binary weight files do not support appended GNN component offsets"); + + const auto binaryWeights = CollectBinaryWeightTensors(fInitializedTensors); + fGC += " std::ifstream f(filename, std::ios::binary);\n"; + fGC += " if (!f.is_open()) throw std::runtime_error(\"sofie failed to open binary weight file \" + filename);\n"; + fGC += " SOFIE::ReadBinaryWeightFileHeader(f, " + std::to_string(binaryWeights.size()) + ");\n"; + for (const auto &[name, tensor] : binaryWeights) { + const auto type = tensor->type(); + if (!IsBinaryWeightTensorType(type)) + throw std::runtime_error("sofie tensor " + name + " has a type unsupported by binary weights"); + std::string shape = "{"; + for (auto dim : tensor->shape()) shape += std::to_string(dim) + ","; + shape += "}"; + fGC += " SOFIE::ReadTensorFromBinaryStream(f, tensor_" + name + ", \"tensor_" + name + + "\", static_cast(" + std::to_string(static_cast(type)) + + "), std::vector" + shape + ");\n"; + } + fGC += " f.close();\n"; + } + // generate the code to read initialized tensors from a ROOT data file if(fWeightFile == WeightFileType::RootBinary) { #ifdef SOFIE_SUPPORT_ROOT_BINARY @@ -1640,6 +1711,9 @@ long RModel::WriteInitializedTensorsToFile(std::string filename) { case WeightFileType::Text: fileExtension = ".dat"; break; + case WeightFileType::Binary: + fileExtension = ".bin"; + break; } // If filename is empty, use the model name as the base filename @@ -1648,7 +1722,16 @@ long RModel::WriteInitializedTensorsToFile(std::string filename) { } // Write the initialized tensors to the file - if (fWeightFile == WeightFileType::RootBinary) { + if (fWeightFile == WeightFileType::Binary) { + if (fIsGNNComponent || fIsGNN) + throw std::runtime_error("SOFIE binary weight files do not support appended GNN components"); + std::ofstream f(filename, std::ios::binary | std::ios::trunc); + if (!f.is_open()) + throw std::runtime_error("sofie failed to open binary weight file " + filename); + const auto position = WriteBinaryWeightFile(f, fInitializedTensors); + f.close(); + return position; + } else if (fWeightFile == WeightFileType::RootBinary) { #ifdef SOFIE_SUPPORT_ROOT_BINARY if(fIsGNNComponent || fIsGNN) { throw std::runtime_error("SOFIE-GNN yet not supports writing to a ROOT file."); @@ -1731,6 +1814,30 @@ long RModel::WriteInitializedTensorsToFile(std::string filename) { f << ( (idx < length-1) ? " " : "\n" ); } } + else if (i.second.type() == ETensorType::INT8) { + const int8_t * data = i.second.data(); + for (size_t idx = 0; idx < length; idx++) { + f << static_cast(data[idx]) << ( (idx < length-1) ? " " : "\n" ); + } + } + else if (i.second.type() == ETensorType::UINT8) { + const uint8_t * data = i.second.data(); + for (size_t idx = 0; idx < length; idx++) { + f << static_cast(data[idx]) << ( (idx < length-1) ? " " : "\n" ); + } + } + else if (i.second.type() == ETensorType::INT32) { + const int32_t * data = i.second.data(); + for (size_t idx = 0; idx < length; idx++) { + f << data[idx] << ( (idx < length-1) ? " " : "\n" ); + } + } + else if (i.second.type() == ETensorType::INT64) { + const int64_t * data = i.second.data(); + for (size_t idx = 0; idx < length; idx++) { + f << data[idx] << ( (idx < length-1) ? " " : "\n" ); + } + } else { throw std::runtime_error("sofie tensor " + tensor_name + " with type " + ConvertTypeToString(i.second.type()) + " cannot be written to a file"); } @@ -1984,9 +2091,13 @@ void RModel::OutputGenerated(std::string filename, bool append) { filename = filename.erase(pos, 4); filename += ".root"; } + if (fWeightFile == WeightFileType::Binary) + filename.replace(pos, 4, ".bin"); } else { filename = fName; - filename += fWeightFile == WeightFileType::Text ? ".dat" : ".root"; + if (fWeightFile == WeightFileType::Text) filename += ".dat"; + else if (fWeightFile == WeightFileType::RootBinary) filename += ".root"; + else if (fWeightFile == WeightFileType::Binary) filename += ".bin"; } WriteInitializedTensorsToFile(filename); } diff --git a/core/src/RModel_ALPAKA.cxx b/core/src/RModel_ALPAKA.cxx index 9e0e84c..c9c0c2b 100644 --- a/core/src/RModel_ALPAKA.cxx +++ b/core/src/RModel_ALPAKA.cxx @@ -148,6 +148,8 @@ void RModel::GenerateInitializedTensorInfo_GPU_ALPAKA() { } for (auto &i : fInitializedTensors) { + if (i.second.IsNotWritable()) + continue; if (!fUseWeightFile || i.second.IsConstantTensor()) { if (i.second.type() == ETensorType::FLOAT) fGC += GenerateConstantTensorCode(i); @@ -155,6 +157,8 @@ void RModel::GenerateInitializedTensorInfo_GPU_ALPAKA() { fGC += GenerateConstantTensorCode(i); else if (i.second.type() == ETensorType::INT32) fGC += GenerateConstantTensorCode(i); + else if (i.second.type() == ETensorType::INT8) + fGC += GenerateConstantTensorCode(i); else if (i.second.type() == ETensorType::BOOL || i.second.type() == ETensorType::UINT8) @@ -170,6 +174,10 @@ void RModel::GenerateInitializedTensorInfo_GPU_ALPAKA() { fGC += "BufI321D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + std::to_string(length) + "}));\n"; + } else if (i.second.type() == ETensorType::INT8) { + fGC += "BufI81D deviceBuf_" + i.first + + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + + std::to_string(length) + "}));\n"; } else if (i.second.type() == ETensorType::INT64) { fGC += "BufI641D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + @@ -190,6 +198,8 @@ void RModel::GenerateTemporaryInitializedTensorContainers_GPU_ALPAKA() fGC += "// temporary initialized tensors for loading weights\n"; for (auto &i : fInitializedTensors) { + if (i.second.IsNotWritable()) + continue; if (fUseWeightFile && !i.second.IsConstantTensor()) { // case of tensors which are read from a file size_t length = ConvertShapeToLength(i.second.shape()); @@ -197,6 +207,8 @@ void RModel::GenerateTemporaryInitializedTensorContainers_GPU_ALPAKA() fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; } else if (i.second.type() == ETensorType::INT32) { fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; + } else if (i.second.type() == ETensorType::INT8) { + fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; } else if (i.second.type() == ETensorType::INT64) { fGC += "std::vector tensor_" + i.first + "(" + std::to_string(length) + ");\n"; } else if (i.second.type() == ETensorType::BOOL || @@ -229,11 +241,15 @@ void RModel::GenerateGPU_ALPAKA_Buffers() { tensor_declaration_block += "BufI321D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + std::to_string(length) + "}));\n"; + } else if (i.second.type == ETensorType::INT8) { + tensor_declaration_block += "BufI81D deviceBuf_" + i.first + + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + + std::to_string(length) + "}));\n"; } else if (i.second.type == ETensorType::INT64) { tensor_declaration_block += "BufI641D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + std::to_string(length) + "}));\n"; - } else if (i.second.type == ETensorType::BOOL) { + } else if (i.second.type == ETensorType::BOOL || i.second.type == ETensorType::UINT8) { tensor_declaration_block += "BufUI81D deviceBuf_" + i.first + " = alpaka::allocBuf(devAcc, Ext1D::all(Idx{" + std::to_string(length) + "}));\n"; @@ -286,7 +302,8 @@ std::string RModel::GenerateInferSignature_GPU_ALPAKA(bool isdecl) { if (type == ETensorType::DOUBLE) return "BufD1D"; if (type == ETensorType::INT32) return "BufI321D"; if (type == ETensorType::INT64) return "BufI641D"; - if (type == ETensorType::BOOL) return "BufUI81D"; + if (type == ETensorType::INT8) return "BufI81D"; + if (type == ETensorType::BOOL || type == ETensorType::UINT8) return "BufUI81D"; throw std::runtime_error("sofie: input tensor " + name + " is of a data type which is not yet supported."); }; @@ -326,7 +343,8 @@ std::string RModel::GenerateImplSignature_GPU_ALPAKA(bool isdecl) { if (type == ETensorType::DOUBLE) return "ViewConstD1D"; if (type == ETensorType::INT32) return "ViewConstI321D"; if (type == ETensorType::INT64) return "ViewConstI641D"; - if (type == ETensorType::BOOL) return "ViewConstUI81D"; + if (type == ETensorType::INT8) return "ViewConstI81D"; + if (type == ETensorType::BOOL || type == ETensorType::UINT8) return "ViewConstUI81D"; throw std::runtime_error("sofie: input tensor " + name + " is of a data type which is not yet supported."); }; @@ -378,7 +396,8 @@ void RModel::GenerateOutput_GPU_ALPAKA() { if (type == ETensorType::DOUBLE) return "ViewConstD1D"; if (type == ETensorType::INT32) return "ViewConstI321D"; if (type == ETensorType::INT64) return "ViewConstI641D"; - if (type == ETensorType::BOOL) return "ViewConstUI81D"; + if (type == ETensorType::INT8) return "ViewConstI81D"; + if (type == ETensorType::BOOL || type == ETensorType::UINT8) return "ViewConstUI81D"; throw std::runtime_error("sofie: input tensor " + name + " is of an unsupported data type."); }; @@ -416,7 +435,9 @@ void RModel::GenerateOutput_GPU_ALPAKA() { if (fVerbose) std::cout << "Generating code for operator .... " << op_idx << std::endl; - if (fSkipOperators.count(op_idx)) continue; + if (fSkipOperators.count(op_idx) || fLoweredConsumedOperatorIndices.count(op_idx) != 0) continue; + auto loweredIt = fLoweredOperators.find(op_idx); + ROperator *op = (loweredIt != fLoweredOperators.end()) ? loweredIt->second.get() : fOperators[op_idx].get(); auto gIt = fOpToFusionGroupIdx.find(op_idx); size_t gIdx = (gIt != fOpToFusionGroupIdx.end()) ? gIt->second : SIZE_MAX; @@ -457,9 +478,9 @@ void RModel::GenerateOutput_GPU_ALPAKA() { // Chain followers: skip — their logic is inside the fused kernel } else { if (fProfile) { - fGC += RModelProfilerGPU::GenerateOperatorCode(*fOperators[op_idx], op_idx); + fGC += RModelProfilerGPU::GenerateOperatorCode(*op, op_idx); } else { - fGC += fOperators[op_idx]->Generate_GPU_ALPAKA(std::to_string(op_idx)); + fGC += op->Generate_GPU_ALPAKA(std::to_string(op_idx)); } } } @@ -477,31 +498,34 @@ void RModel::GenerateOutput_GPU_ALPAKA() { for (auto &p : dynParamNames) spanDynDecl += ", size_t " + p; - fGC += "void infer(std::span inputs, std::span outputs" + spanDynDecl + "){\n"; - - { - fGC += SP + "_infer_impl("; - bool first = true; - for (auto &p : dynParamNames) { - if (!first) fGC += ", "; - fGC += p; - first = false; - } - for (size_t i = 0; i < fInputTensorNames.size(); i++) { - if (!first) fGC += ", "; - fGC += "inputs[" + std::to_string(i) + "]"; - first = false; + const bool allFloatOutputs = sameOutputTypes && eFirstOutputType == ETensorType::FLOAT; + if (allFloatOutputs) { + fGC += "void infer(std::span inputs, std::span outputs" + spanDynDecl + "){\n"; + + { + fGC += SP + "_infer_impl("; + bool first = true; + for (auto &p : dynParamNames) { + if (!first) fGC += ", "; + fGC += p; + first = false; + } + for (size_t i = 0; i < fInputTensorNames.size(); i++) { + if (!first) fGC += ", "; + fGC += "inputs[" + std::to_string(i) + "]"; + first = false; + } + fGC += ");\n"; } - fGC += ");\n"; - } - // Copy member output buffers into caller-provided output views - for (size_t i = 0; i < outputSize; i++) { - std::string tensorName = *(fOutputTensorNames.begin() + i); - fGC += SP + "alpaka::memcpy(queue, outputs[" + std::to_string(i) + "], deviceBuf_" + tensorName + ");\n"; + // Copy member output buffers into caller-provided output views + for (size_t i = 0; i < outputSize; i++) { + std::string tensorName = *(fOutputTensorNames.begin() + i); + fGC += SP + "alpaka::memcpy(queue, outputs[" + std::to_string(i) + "], deviceBuf_" + tensorName + ");\n"; + } + fGC += SP + "alpaka::wait(queue);\n"; + fGC += "}\n\n"; } - fGC += SP + "alpaka::wait(queue);\n"; - fGC += "}\n\n"; std::string returnType; @@ -583,7 +607,10 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fGC += "\n//--- ALPAKA Kernels\n"; for (size_t id = 0; id < fOperators.size(); id++) { - if(fOperators[id]->GetKind() == OperatorKind::GEMM || fOperators[id]->GetKind() == OperatorKind::CONV) { + if (fLoweredConsumedOperatorIndices.count(id) != 0) continue; + auto loweredIt = fLoweredOperators.find(id); + ROperator *op = (loweredIt != fLoweredOperators.end()) ? loweredIt->second.get() : fOperators[id].get(); + if(op->GetKind() == OperatorKind::GEMM || op->GetKind() == OperatorKind::CONV) { OpNeedsBlas = true; } @@ -614,17 +641,17 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { // Chain followers: skip (their logic is inside the fused kernel) } else { // Unfused op: generate individual kernel struct (with dedup for single_initialized_operators) - if (single_initialized_operators.find(fOperators[id]->GetKind()) != single_initialized_operators.end()) { - if (registered_operators.find(fOperators[id]->GetKind()) == registered_operators.end()) { + if (single_initialized_operators.find(op->GetKind()) != single_initialized_operators.end()) { + if (registered_operators.find(op->GetKind()) == registered_operators.end()) { if (fVerbose) - std::cout << "Generating ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); - registered_operators.insert(fOperators[id]->GetKind()); + std::cout << "Generating ALPAKA kernel for operator " << toString(op->GetKind()) << std::endl; + fGC += op->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); + registered_operators.insert(op->GetKind()); } } else { if (fVerbose) - std::cout << "Generating ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); + std::cout << "Generating ALPAKA kernel for operator " << toString(op->GetKind()) << std::endl; + fGC += op->Generate_GPU_Kernel_ALPAKA(std::to_string(id)); } } } @@ -659,6 +686,7 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fGC += "using QueueAcc = alpaka::Queue;\n\n"; fGC += "using BufF1D = alpaka::Buf;\n"; fGC += "using BufD1D = alpaka::Buf;\n"; + fGC += "using BufI81D = alpaka::Buf;\n"; fGC += "using BufI321D = alpaka::Buf;\n"; fGC += "using BufI641D = alpaka::Buf;\n"; fGC += "using BufUI81D = alpaka::Buf;\n\n"; @@ -667,6 +695,8 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fGC += "using ViewConstF1D = alpaka::ViewPlainPtr;\n"; fGC += "using ViewD1D = alpaka::ViewPlainPtr;\n"; fGC += "using ViewConstD1D = alpaka::ViewPlainPtr;\n"; + fGC += "using ViewI81D = alpaka::ViewPlainPtr;\n"; + fGC += "using ViewConstI81D = alpaka::ViewPlainPtr;\n"; fGC += "using ViewI321D = alpaka::ViewPlainPtr;\n"; fGC += "using ViewConstI321D = alpaka::ViewPlainPtr;\n"; fGC += "using ViewI641D = alpaka::ViewPlainPtr;\n"; @@ -707,6 +737,8 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fileName += ".dat"; if (fWeightFile == WeightFileType::RootBinary) fileName += ".root"; + if (fWeightFile == WeightFileType::Binary) + fileName += ".bin"; fGC += sessionName + "(std::string filename =\"" + fileName + "\""; } else { @@ -732,12 +764,14 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { GenerateDynamicTensorInfo_GPU_ALPAKA(); for (size_t id = 0; id < fOperators.size(); id++) { - if (fSkipOperators.count(id)) continue; - fGC += fOperators[id]->GenerateInitCode_GPU_ALPAKA(); - if (fOperators[id]->GetKind() == OperatorKind::GEMM || fOperators[id]->GetKind() == OperatorKind::CONV) { + if (fSkipOperators.count(id) || fLoweredConsumedOperatorIndices.count(id) != 0) continue; + auto loweredIt = fLoweredOperators.find(id); + ROperator *op = (loweredIt != fLoweredOperators.end()) ? loweredIt->second.get() : fOperators[id].get(); + fGC += op->GenerateInitCode_GPU_ALPAKA(); + if (op->GetKind() == OperatorKind::GEMM || op->GetKind() == OperatorKind::CONV) { // GetBlasConfig() returns "" for ops that use gemmStridedBatched - // (legacy cuBLAS path, no cuBLASLt layout registration needed). - auto blasCfg = fOperators[id]->GetBlasConfig(); + // (legacy cuBLAS path, no cuBLASLt layout needed). + auto blasCfg = op->GetBlasConfig(); if (!blasCfg.empty()) fGC += "\nblas.addLayoutConfig("+blasCfg+");\n"; } @@ -751,6 +785,9 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fusedGroupsEmitted.clear(); for (size_t id = 0; id < fOperators.size(); id++) { + if (fLoweredConsumedOperatorIndices.count(id) != 0) continue; + auto loweredIt = fLoweredOperators.find(id); + ROperator *op = (loweredIt != fLoweredOperators.end()) ? loweredIt->second.get() : fOperators[id].get(); // Same as the kernel-struct loop above: fused activation ops must still // declare their member variable (e.g. `leakyReluKernel`) even though // their Generate_GPU_ALPAKA call is skipped in the infer-body loop. @@ -766,17 +803,17 @@ void RModel::GenerateSessionCode_GPU_ALPAKA() { fusedGroupsEmitted.insert(gIdx); } } else { - if (single_initialized_operators.find(fOperators[id]->GetKind()) != single_initialized_operators.end()) { - if (registered_operators.find(fOperators[id]->GetKind()) == registered_operators.end()) { + if (single_initialized_operators.find(op->GetKind()) != single_initialized_operators.end()) { + if (registered_operators.find(op->GetKind()) == registered_operators.end()) { if (fVerbose) - std::cout << "Declaring ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_Definitions_ALPAKA(std::to_string(id)); - registered_operators.insert(fOperators[id]->GetKind()); + std::cout << "Declaring ALPAKA kernel for operator " << toString(op->GetKind()) << std::endl; + fGC += op->Generate_GPU_Kernel_Definitions_ALPAKA(std::to_string(id)); + registered_operators.insert(op->GetKind()); } } else { if (fVerbose) - std::cout << "Declaring ALPAKA kernel for operator " << toString(fOperators[id]->GetKind()) << std::endl; - fGC += fOperators[id]->Generate_GPU_Kernel_Definitions_ALPAKA(std::to_string(id)); + std::cout << "Declaring ALPAKA kernel for operator " << toString(op->GetKind()) << std::endl; + fGC += op->Generate_GPU_Kernel_Definitions_ALPAKA(std::to_string(id)); } } } @@ -815,6 +852,18 @@ void RModel::GenerateGPU_ALPAKA(std::underlying_type_t options, int bat fUseWeightFile = true; fWeightFile = WeightFileType::RootBinary; } + if (static_cast>(Options::kBinaryWeightFile) & options) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Binary; + } + if (static_cast>(Options::kTextWeightFile) & options) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Text; + } + const auto explicitWeightPolicy = static_cast>(Options::kNoWeightFile) | + static_cast>(Options::kRootBinaryWeightFile) | + static_cast>(Options::kBinaryWeightFile) | + static_cast>(Options::kTextWeightFile); if (fUseWeightFile && !fUseSession) { throw std::runtime_error( "sofie: RModel::Generate: cannot use a separate weight file without generating a Session class"); @@ -825,6 +874,14 @@ void RModel::GenerateGPU_ALPAKA(std::underlying_type_t options, int bat throw std::runtime_error("SOFIE GPU does not yet supports GNN Inference."); Initialize(batchSize, verbose); + if ((options & explicitWeightPolicy) == 0 && !fQuantizationState.tensorInfos.empty()) { + fUseWeightFile = true; + fWeightFile = WeightFileType::Binary; + } + if (fWeightFile == WeightFileType::Binary) + AddNeededCustomHeader("SOFIE/RWeightFile.hxx"); + BuildLoweredOperatorView(EQuantizedBackend::ALPAKA); + AddQuantizedGeneratedHeaders(EQuantizedBackend::ALPAKA); FuseGemmActivations_GPU(); // must run before elementwise fusion (redirects tensors) ComputeEltwiseFusionGroups(); @@ -848,6 +905,16 @@ void RModel::GenerateGPU_ALPAKA(std::underlying_type_t options, int bat void RModel::MoveInitializedTensorsToBuffers_ALPAKA(){ for (auto &i : fInitializedTensors) { if (i.second.IsNotWritable()) continue; + if (HasQuantizedTensorStorage(i.first)) { + const auto &storage = GetQuantizedTensorStorage(i.first); + if (storage.layout == EQuantizedLayout::PackedCPU || + storage.residentBackend != EQuantizedBackend::ALPAKA) + continue; + } + const auto type = i.second.type(); + if (type != ETensorType::FLOAT && type != ETensorType::INT32 && type != ETensorType::INT64 && + type != ETensorType::BOOL && type != ETensorType::INT8 && type != ETensorType::UINT8) + continue; std::string tensor_name = "tensor_" + i.first; auto length = ConvertShapeToLength(i.second.shape()); std::string slength = std::to_string(length); diff --git a/core/src/RModel_Quantization.cxx b/core/src/RModel_Quantization.cxx new file mode 100644 index 0000000..98556be --- /dev/null +++ b/core/src/RModel_Quantization.cxx @@ -0,0 +1,745 @@ +#include "SOFIE/RModel.hxx" +#include "SOFIE/ROperator_Gemm.hxx" +#include "SOFIE/ROperator_QuantizedGemm.hxx" +#include "SOFIE/ROperator_QuantizedMatMul.hxx" +#include "SOFIE/RQuantization.hxx" +#include "SOFIE/RQuantization_Analysis.hxx" +#include "SOFIE/RQuantization_DenseLinear.hxx" +#include "SOFIE/RQuantization_Storage.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace SOFIE { + +namespace { + +std::string GetQuantBoundarySourceTensor(const ROperator &op) +{ + return op.GetQuantizationSourceTensor(); +} + +QuantizedGemmCodegenContext MakeQuantizedGemmCodegenContext(const ROperator_Gemm &gemm) +{ + QuantizedGemmCodegenContext context; + context.inputShape = gemm.GetInputShape(); + context.weightShape = gemm.GetWeightShape(); + context.outputShape = gemm.GetOutputShape(); + context.alpha = gemm.GetAlpha(); + context.beta = gemm.GetBeta(); + context.transA = gemm.GetTransA(); + context.transB = gemm.GetTransB(); + context.activation = gemm.GetActivationType(); + return context; +} + +QuantizedMatrixCodegenContext MakeQuantizedMatMulCodegenContext(const ROperator_Gemm &gemm) +{ + QuantizedMatrixCodegenContext context; + context.inputShape = gemm.GetInputShape(); + context.weightShape = gemm.GetWeightShape(); + context.outputShape = gemm.GetOutputShape(); + return context; +} + + +} // namespace + +void RModel::AddQuantizationInfo(const std::string & tensor_name, QuantizationInfo info) +{ + fQuantizationState.tensorInfos[UTILITY::Clean_name(tensor_name)] = std::move(info); +} + +bool RModel::HasQuantizationInfo(const std::string & tensor_name) const +{ + auto clean_name = UTILITY::Clean_name(tensor_name); + if (fQuantizationState.tensorInfos.find(clean_name) != fQuantizationState.tensorInfos.end()) + return true; + if (fIsSubGraph && fParentGraph) + return fParentGraph->HasQuantizationInfo(clean_name); + return false; +} + +const QuantizationInfo & RModel::GetQuantizationInfo(const std::string & tensor_name) const +{ + auto clean_name = UTILITY::Clean_name(tensor_name); + auto f = fQuantizationState.tensorInfos.find(clean_name); + if (f != fQuantizationState.tensorInfos.end()) + return f->second; + if (fIsSubGraph && fParentGraph) + return fParentGraph->GetQuantizationInfo(clean_name); + throw std::runtime_error("SOFIE tensor [" + clean_name + "] has no quantization information"); +} + +void RModel::RegisterQuantizedTensorStorage(QuantizedTensorStorage storage) +{ + storage.storageTensor = UTILITY::Clean_name(storage.storageTensor); + storage.logicalTensor = UTILITY::Clean_name(storage.logicalTensor); + storage.sourceTensor = UTILITY::Clean_name(storage.sourceTensor); + if (storage.storageTensor.empty()) { + throw std::runtime_error("SOFIE quantized tensor storage registration requires a storage tensor name"); + } + fQuantizationState.tensorStorages[storage.storageTensor] = std::move(storage); +} + +bool RModel::HasQuantizedTensorStorage(const std::string & storage_tensor_name) const +{ + auto clean_name = UTILITY::Clean_name(storage_tensor_name); + return fQuantizationState.tensorStorages.find(clean_name) != fQuantizationState.tensorStorages.end(); +} + +const QuantizedTensorStorage & RModel::GetQuantizedTensorStorage(const std::string & storage_tensor_name) const +{ + auto clean_name = UTILITY::Clean_name(storage_tensor_name); + auto it = fQuantizationState.tensorStorages.find(clean_name); + if (it == fQuantizationState.tensorStorages.end()) { + throw std::runtime_error("SOFIE tensor [" + clean_name + "] has no quantized storage information"); + } + return it->second; +} + +void RModel::AnalyzeQuantizedRegions() +{ + for (const auto &[name, storage] : fQuantizationState.tensorStorages) + fInitializedTensors.erase(name); + fQuantizationState.ClearDerivedAnalysis(); + + const auto graph = BuildQuantizationGraphIndex(fOperators); + auto readZeroPointTensor = [this](const std::string &tensorName) { + std::vector values; + auto appendValues = [&values](const auto &typedValues) { + values.reserve(typedValues.size()); + for (auto value : typedValues) + values.push_back(static_cast(value)); + }; + + switch (GetTensorType(tensorName)) { + case ETensorType::FLOAT: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::DOUBLE: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::INT8: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::UINT8: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::INT16: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::UINT16: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::INT32: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::UINT32: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::INT64: + appendValues(GetTensorData(tensorName)); + break; + case ETensorType::UINT64: + appendValues(GetTensorData(tensorName)); + break; + default: + throw std::runtime_error("SOFIE quantized lowering expects numeric zero-point tensor [" + tensorName + "]"); + } + return values; + }; + + for (std::size_t opIndex = 0; opIndex < fOperators.size(); ++opIndex) { + if (fOperators[opIndex]->GetKind() != OperatorKind::GEMM) + continue; + + auto *gemm = dynamic_cast *>(fOperators[opIndex].get()); + if (!gemm) + continue; + + auto pattern = MatchQuantizedDenseLinearPattern( + *gemm, opIndex, [this](const std::string &tensor) { return GetTensorShape(tensor); }); + auto info = std::move(pattern.region); + auto reasons = std::move(pattern.reasons); + const bool isQuantizedMatMulSpelling = pattern.isMatMul; + const bool isMatMulAddSpelling = pattern.hasInlineMatMulBias; + const bool isMatMulSpelling = pattern.isMatMul && !pattern.hasInlineMatMulBias; + auto matmulShape = std::move(pattern.matmulShape); + + if (!info.inputTensor.empty()) { + if (auto producer = MatchQuantizationBoundaryProducer(graph, fOperators, info.inputTensor, "input", reasons)) { + info.inputQuantOpIndex = *producer; + info.inputSourceTensor = GetQuantBoundarySourceTensor(*fOperators[*producer]); + } + if (HasQuantizationInfo(info.inputTensor)) { + info.inputQuant = GetQuantizationInfo(info.inputTensor); + CheckQuantizationInfo(info.inputQuant, "input", reasons); + } else { + reasons.push_back("input tensor has no QuantizationInfo"); + } + } + + if (!info.weightTensor.empty()) { + if (auto producer = MatchQuantizationBoundaryProducer(graph, fOperators, info.weightTensor, "weight", reasons)) { + info.weightQuantOpIndex = *producer; + info.weightSourceTensor = GetQuantBoundarySourceTensor(*fOperators[*producer]); + if (IsInitializedTensor(info.weightTensor) && GetTensorType(info.weightTensor) == ETensorType::FLOAT) + info.weightSourceTensor = info.weightTensor; + } + if (HasQuantizationInfo(info.weightTensor)) { + info.weightQuant = GetQuantizationInfo(info.weightTensor); + CheckQuantizationInfo(info.weightQuant, "weight", reasons); + } else { + reasons.push_back("weight tensor has no QuantizationInfo"); + } + } + + if (!info.biasTensor.empty()) { + if (isMatMulAddSpelling) { + info.biasSourceTensor = info.biasTensor; + if (!IsInitializedTensor(info.biasSourceTensor)) { + reasons.push_back("MatMul fused Add bias must be an initialized constant tensor"); + } else if (!info.gemmOutputTensor.empty() && + !IsDenseLinearBiasLikeShape(GetTensorShape(info.biasSourceTensor), GetTensorShape(info.gemmOutputTensor))) { + reasons.push_back("MatMul fused Add bias is not a dense-linear projection bias broadcast shape"); + } else { + info.biasQuant = MakeAccumulatorBiasQuantization(info.inputQuant, info.weightQuant); + } + } else { + if (auto producer = MatchQuantizationBoundaryProducer(graph, fOperators, info.biasTensor, "bias", reasons)) { + info.biasQuantOpIndex = *producer; + info.biasSourceTensor = GetQuantBoundarySourceTensor(*fOperators[*producer]); + if (IsInitializedTensor(info.biasTensor) && GetTensorType(info.biasTensor) == ETensorType::FLOAT) + info.biasSourceTensor = info.biasTensor; + } + if (HasQuantizationInfo(info.biasTensor)) { + info.biasQuant = GetQuantizationInfo(info.biasTensor); + CheckQuantizationInfo(*info.biasQuant, "bias", reasons); + } else { + reasons.push_back("bias tensor has no QuantizationInfo"); + } + } + } + + QuantizedEpilogue matmulEpilogue; + if (isMatMulAddSpelling && !info.biasSourceTensor.empty() && info.biasQuant.has_value()) { + matmulEpilogue.kind = EQuantizedEpilogueKind::Bias; + matmulEpilogue.biasSourceTensor = info.biasSourceTensor; + matmulEpilogue.biasQuant = info.biasQuant; + } + + if (!info.gemmOutputTensor.empty()) { + auto consumers = graph.consumersByTensor.find(info.gemmOutputTensor); + if (consumers == graph.consumersByTensor.end() || consumers->second.empty()) { + reasons.push_back("Gemm output has no output quantization consumer"); + } else if (consumers->second.size() != 1) { + reasons.push_back("Gemm output has multiple consumers"); + } else { + auto consumerIndex = consumers->second.front(); + auto setOutputQuantFromBoundary = [&](std::size_t quantIndex) { + info.outputQuantOpIndex = quantIndex; + auto quantOutputs = fOperators[quantIndex]->GetOpOutputTensors(); + if (quantOutputs.size() != 1) { + reasons.push_back("output quantization boundary does not have exactly one output"); + } else { + info.outputTensor = std::string(quantOutputs[0]); + if (HasQuantizationInfo(info.outputTensor)) { + info.outputQuant = GetQuantizationInfo(info.outputTensor); + CheckQuantizationInfo(info.outputQuant, "output", reasons); + } else { + reasons.push_back("output tensor has no QuantizationInfo"); + } + } + }; + + if (fOperators[consumerIndex]->IsQuantizationBoundary()) { + setOutputQuantFromBoundary(consumerIndex); + } else if (isMatMulSpelling && IsFloatAddOperator(*fOperators[consumerIndex])) { + const auto addInputs = fOperators[consumerIndex]->GetOpInputTensors(); + const auto addOutputs = fOperators[consumerIndex]->GetOpOutputTensors(); + if (addInputs.size() != 2 || addOutputs.size() != 1) { + reasons.push_back("MatMul Add epilogue does not have two inputs and one output"); + } else { + const std::string addInputA = std::string(addInputs[0]); + const std::string addInputB = std::string(addInputs[1]); + const std::string biasCandidate = addInputA == info.gemmOutputTensor ? addInputB : + (addInputB == info.gemmOutputTensor ? addInputA : std::string{}); + if (biasCandidate.empty()) { + reasons.push_back("MatMul Add epilogue does not consume the MatMul output"); + } else if (!IsInitializedTensor(biasCandidate)) { + reasons.push_back("MatMul Add epilogue bias must be an initialized constant tensor"); + } else if (!IsDenseLinearBiasLikeShape(GetTensorShape(biasCandidate), GetTensorShape(info.gemmOutputTensor))) { + reasons.push_back("MatMul Add epilogue constant is not a dense-linear projection bias broadcast shape"); + } else { + const std::string addOutput = std::string(addOutputs[0]); + auto addOutputConsumers = graph.consumersByTensor.find(addOutput); + if (addOutputConsumers == graph.consumersByTensor.end() || addOutputConsumers->second.empty()) { + reasons.push_back("MatMul Add epilogue output has no output quantization consumer"); + } else if (addOutputConsumers->second.size() != 1) { + reasons.push_back("MatMul Add epilogue output has multiple consumers"); + } else if (!fOperators[addOutputConsumers->second.front()]->IsQuantizationBoundary()) { + reasons.push_back("MatMul Add epilogue output consumer is not a quantization boundary"); + } else { + matmulEpilogue.kind = EQuantizedEpilogueKind::Bias; + matmulEpilogue.biasSourceTensor = biasCandidate; + matmulEpilogue.biasQuant = MakeAccumulatorBiasQuantization(info.inputQuant, info.weightQuant); + matmulEpilogue.addOpIndex = consumerIndex; + setOutputQuantFromBoundary(addOutputConsumers->second.front()); + } + } + } + } else { + reasons.push_back("Gemm output consumer is not a quantization boundary"); + } + } + } + + const bool hasQuantizationEvidence = + info.inputQuantOpIndex != static_cast(-1) || + info.weightQuantOpIndex != static_cast(-1) || + info.biasQuantOpIndex.has_value() || info.outputQuantOpIndex != static_cast(-1) || + (!info.inputTensor.empty() && HasQuantizationInfo(info.inputTensor)) || + (!info.weightTensor.empty() && HasQuantizationInfo(info.weightTensor)) || + (!info.biasTensor.empty() && HasQuantizationInfo(info.biasTensor)) || + (!info.outputTensor.empty() && HasQuantizationInfo(info.outputTensor)); + + if (isQuantizedMatMulSpelling) { + if (hasQuantizationEvidence) { + auto matmul = MakeQuantizedMatMulRegionFromGemmLikeRegion(info); + matmul.epilogue = matmulEpilogue; + matmul.shape = matmulShape; + auto &plans = fQuantizationState.loweringPlans[opIndex]; + if (reasons.empty()) { + matmul.status = EQuantizedLoweringStatus::SemanticRecognized; + matmul.reason = QuantizedEpilogueHasBias(matmul.epilogue.kind) + ? "recognized quantized MatMul+Add bias region" + : "recognized quantized MatMul region"; + if (!matmul.shape.reason.empty()) + matmul.reason += "; " + matmul.shape.reason; + + auto cpuReason = matmul.reason + "; CPU QuantizedMatMul lowering is not implemented"; + plans[EQuantizedBackend::CPU] = MakeUnsupportedQuantizedMatMulPlan(matmul, EQuantizedBackend::CPU, cpuReason, true); + + std::vector storageReasons; + std::vector perChannelWeightScales; + const bool perChannelWeight = IsPerChannelAxis(matmul.weightQuant, 1); + + std::vector inputShape; + std::vector weightShape; + std::vector outputShape; + if (!matmul.inputSourceTensor.empty()) + inputShape = GetTensorShape(matmul.inputSourceTensor); + if (!matmul.weightSourceTensor.empty() && IsInitializedTensor(matmul.weightSourceTensor)) + weightShape = GetTensorShape(matmul.weightSourceTensor); + else + storageReasons.push_back("MatMul weight source tensor must be initialized for transposed quantized storage"); + if (!matmul.outputTensor.empty()) + outputShape = GetTensorShape(matmul.outputTensor); + + const auto capability = AssessCublasLtDenseLinearCapability( + MakeDenseLinearOperands(matmul, inputShape, weightShape, outputShape)); + const auto selectedCapability = SelectExecutableDenseLinearCapability(capability); + if (!selectedCapability.optimized) { + storageReasons.push_back("MatMul cuBLASLt optimized profile unavailable: " + selectedCapability.reason); + } + + if (perChannelWeight) { + if (!IsInitializedTensor(matmul.weightQuant.scaleTensor)) { + storageReasons.push_back("MatMul per-channel weight scale tensor is not initialized"); + } else if (weightShape.size() == 2) { + perChannelWeightScales = GetTensorData(matmul.weightQuant.scaleTensor); + if (perChannelWeightScales.size() != weightShape[1]) { + storageReasons.push_back("MatMul per-channel weight scale length does not match output channels N"); + } + } + if (!IsInitializedTensor(matmul.weightQuant.zeroPointTensor)) { + storageReasons.push_back("MatMul per-channel weight zero-point tensor is not initialized"); + } else { + const auto zeroPoints = readZeroPointTensor(matmul.weightQuant.zeroPointTensor); + if (weightShape.size() == 2 && zeroPoints.size() != weightShape[1]) { + storageReasons.push_back("MatMul per-channel weight zero-point length does not match output channels N"); + } + for (std::int64_t zeroPoint : zeroPoints) { + if (zeroPoint != 0) { + storageReasons.push_back("MatMul per-channel weight zero-points must all be 0"); + break; + } + } + } + } + + if (storageReasons.empty()) { + const bool paddedStorage = selectedCapability.shapePolicy.policy == EQuantizedShapePolicy::Padded; + const auto deviceStorageTensor = matmul.weightSourceTensor + + (paddedStorage ? "_s22_matmul_transposed_padded_plain_device_storage" + : "_s19_matmul_transposed_plain_device_storage"); + auto alpakaPlan = MakeMatMulAlpakaTransposedWeightStoragePlan(matmul, deviceStorageTensor, selectedCapability.shapePolicy); + alpakaPlan.computeProfile = selectedCapability.profile; + alpakaPlan.capabilityTag = selectedCapability.tag; + alpakaPlan.reason = matmul.reason + "; " + selectedCapability.reason; + plans[EQuantizedBackend::ALPAKA] = std::move(alpakaPlan); + matmul.reason += "; transposed pre-quantized ALPAKA weight storage selected"; + } else { + auto alpakaReason = matmul.reason + "; " + JoinQuantizationReasons(storageReasons); + auto unsupportedPlan = MakeUnsupportedQuantizedMatMulPlan(matmul, EQuantizedBackend::ALPAKA, alpakaReason, true); + unsupportedPlan.capabilityTag = capability.tag; + unsupportedPlan.computeProfile = capability.profile; + unsupportedPlan.shapePolicy = capability.shapePolicy; + plans[EQuantizedBackend::ALPAKA] = std::move(unsupportedPlan); + matmul.reason = alpakaReason; + } + } else { + matmul.status = EQuantizedLoweringStatus::SemanticUnsupported; + matmul.reason = JoinQuantizationReasons(reasons); + plans[EQuantizedBackend::CPU] = MakeUnsupportedQuantizedMatMulPlan(matmul, EQuantizedBackend::CPU, matmul.reason, false); + plans[EQuantizedBackend::ALPAKA] = MakeUnsupportedQuantizedMatMulPlan(matmul, EQuantizedBackend::ALPAKA, matmul.reason, false); + } + fQuantizationState.matmulRegions[opIndex] = std::move(matmul); + if (fVerbose > 0) { + std::cout << "SOFIE quantized MatMul candidate at operator " << opIndex << ": " + << fQuantizationState.matmulRegions[opIndex].reason << std::endl; + } + } + continue; + } + + if (reasons.empty()) { + info.status = EQuantizedLoweringStatus::SemanticRecognized; + info.reason = "recognized quantized Gemm region"; + + auto currentLoweringUnsupportedReasons = QuantizedGemmLoweringUnsupportedReasons(info); + std::vector perChannelWeightScales; + if (IsPerChannelAxis(info.weightQuant, 0)) { + if (!IsInitializedTensor(info.weightQuant.scaleTensor)) { + currentLoweringUnsupportedReasons.push_back("per-channel weight scale tensor is not initialized"); + } else { + perChannelWeightScales = GetTensorData(info.weightQuant.scaleTensor); + const auto weightShape = GetTensorShape(info.weightSourceTensor); + if (weightShape.size() != 2 || perChannelWeightScales.size() != weightShape[0]) { + currentLoweringUnsupportedReasons.push_back("per-channel weight scale length does not match GEMM output channels"); + } + } + if (!IsInitializedTensor(info.weightQuant.zeroPointTensor)) { + currentLoweringUnsupportedReasons.push_back("per-channel weight zero-point tensor is not initialized"); + } else { + const auto zeroPoints = readZeroPointTensor(info.weightQuant.zeroPointTensor); + for (std::int64_t zeroPoint : zeroPoints) { + if (zeroPoint != 0) { + currentLoweringUnsupportedReasons.push_back("per-channel weight zero-points must all be 0"); + break; + } + } + } + } + if (!currentLoweringUnsupportedReasons.empty()) { + info.reason += "; " + JoinQuantizationReasons(currentLoweringUnsupportedReasons); + auto &plans = fQuantizationState.loweringPlans[opIndex]; + plans[EQuantizedBackend::CPU] = MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend::CPU, info.reason, true); + plans[EQuantizedBackend::ALPAKA] = MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend::ALPAKA, info.reason, true); + fQuantizationState.gemmRegions[opIndex] = std::move(info); + if (fVerbose > 0) { + std::cout << "SOFIE quantized Gemm candidate recognized but not lowered at operator " << opIndex << ": " + << fQuantizationState.gemmRegions[opIndex].reason << std::endl; + } + continue; + } + + auto cpuPlan = MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend::CPU, "CPU quantized Gemm lowering requires constant pre-quantized weight storage", true); + auto alpakaPlan = MakeAlpakaFakeQuantPlan(info); + + if (!IsPerChannelAxis(info.weightQuant, 0) && !info.weightSourceTensor.empty() && IsInitializedTensor(info.weightSourceTensor)) { + const auto storageTensor = info.weightSourceTensor + "_s11_packed_cpu_storage"; + const auto weightShape = GetTensorShape(info.weightSourceTensor); + if (weightShape.size() != 2) { + reasons.push_back("weight tensor is not rank-2 for packed CPU storage"); + } else { + cpuPlan = MakeCPUPackedWeightBaselinePlan(info, storageTensor); + } + } + + if (!info.weightSourceTensor.empty() && IsInitializedTensor(info.weightSourceTensor)) { + const auto deviceStorageTensor = info.weightSourceTensor + "_s17g3_plain_device_storage"; + const auto weightShape = GetTensorShape(info.weightSourceTensor); + + try { + const auto inputShape = GetTensorShape(info.inputSourceTensor); + const auto outputShape = GetTensorShape(info.outputTensor); + auto capability = AssessCublasLtDenseLinearCapability( + MakeDenseLinearOperands(info, inputShape, weightShape, outputShape)); + auto selectedCapability = SelectExecutableDenseLinearCapability(capability); + if (selectedCapability.optimized) { + std::string selectedStorageTensor = deviceStorageTensor; + if (selectedCapability.shapePolicy.policy == EQuantizedShapePolicy::Padded) { + const auto paddedStorageTensor = info.weightSourceTensor + "_s22_gemm_padded_plain_device_storage"; + selectedStorageTensor = paddedStorageTensor; + } + alpakaPlan = MakeAlpakaCublasLtCorePlan(info, selectedStorageTensor, selectedCapability); + } else { + alpakaPlan.reason += "; cuBLASLt optimized profile unavailable: " + capability.reason; + alpakaPlan.capabilityTag = capability.tag; + alpakaPlan.computeProfile = capability.profile; + alpakaPlan.shapePolicy = capability.shapePolicy; + } + } catch (const std::exception &e) { + alpakaPlan.reason += "; cuBLASLt optimized profile unavailable: " + std::string(e.what()); + alpakaPlan.capabilityTag = "cublaslt_shape_unavailable"; + } + } + + auto &plans = fQuantizationState.loweringPlans[opIndex]; + plans[EQuantizedBackend::CPU] = std::move(cpuPlan); + plans[EQuantizedBackend::ALPAKA] = std::move(alpakaPlan); + fQuantizationState.gemmRegions[opIndex] = std::move(info); + } else { + info.reason = JoinQuantizationReasons(reasons); + if (hasQuantizationEvidence) { + auto &plans = fQuantizationState.loweringPlans[opIndex]; + plans[EQuantizedBackend::CPU] = MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend::CPU, info.reason, true); + plans[EQuantizedBackend::ALPAKA] = MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend::ALPAKA, info.reason, true); + fQuantizationState.gemmRegions[opIndex] = std::move(info); + } + if (fVerbose > 0) { + std::cout << "SOFIE quantized Gemm candidate rejected at operator " << opIndex << ": " << info.reason << std::endl; + } + } + } +} + +void RModel::PrepareQuantizedTensorStorage(EQuantizedBackend backend) +{ + for (const auto &[name, storage] : fQuantizationState.tensorStorages) + fInitializedTensors.erase(name); + fQuantizationState.tensorStorages.clear(); + + auto restoreSource = [this](const std::string &name) { + auto it = fInitializedTensors.find(name); + if (it != fInitializedTensors.end()) + it->second.SetWritable(); + }; + for (const auto &[index, region] : fQuantizationState.gemmRegions) + restoreSource(region.weightSourceTensor); + for (const auto &[index, region] : fQuantizationState.matmulRegions) + restoreSource(region.weightSourceTensor); + + auto installStorage = [this](MaterializedQuantizedWeight materialized) { + const auto name = materialized.storage.storageTensor; + const auto shape = materialized.storage.shape; + std::visit([this, &name, &shape](auto &&buffer) { + AddInitializedTensor(name, shape, std::forward(buffer)); + }, std::move(materialized.buffer)); + RegisterQuantizedTensorStorage(std::move(materialized.storage)); + }; + + for (auto opIndex : SortedQuantizedRegionOperatorIndices(fQuantizationState.gemmRegions)) { + const auto *plan = FindQuantizedLoweringPlan(fQuantizationState, opIndex, backend); + if (plan == nullptr || !IsQuantizedLoweringAvailable(plan->status) || plan->weightStorageTensor.empty()) + continue; + + const auto regionIt = fQuantizationState.gemmRegions.find(opIndex); + if (regionIt == fQuantizationState.gemmRegions.end()) + throw std::runtime_error("SOFIE quantized Gemm storage plan has no matching region"); + const auto ®ion = regionIt->second; + const auto weightShape = GetTensorShape(region.weightSourceTensor); + if (weightShape.size() != 2 || !IsInitializedTensor(region.weightSourceTensor)) + throw std::runtime_error("SOFIE quantized Gemm storage requires an initialized rank-2 weight tensor"); + const auto *weightData = fInitializedTensors.at(region.weightSourceTensor).data(); + + std::vector perChannelScales; + if (backend == EQuantizedBackend::ALPAKA && IsPerChannelAxis(region.weightQuant, 0)) + perChannelScales = GetTensorData(region.weightQuant.scaleTensor); + installStorage(MaterializeQuantizedGemmWeight(region, *plan, backend, weightData, + weightShape, perChannelScales)); + } + + for (auto opIndex : SortedQuantizedRegionOperatorIndices(fQuantizationState.matmulRegions)) { + const auto *plan = FindQuantizedLoweringPlan(fQuantizationState, opIndex, backend); + if (plan == nullptr || !IsQuantizedLoweringAvailable(plan->status) || plan->weightStorageTensor.empty()) + continue; + if (backend != EQuantizedBackend::ALPAKA) + continue; + + const auto regionIt = fQuantizationState.matmulRegions.find(opIndex); + if (regionIt == fQuantizationState.matmulRegions.end()) + throw std::runtime_error("SOFIE quantized MatMul storage plan has no matching region"); + const auto ®ion = regionIt->second; + const auto weightShape = GetTensorShape(region.weightSourceTensor); + if (weightShape.size() != 2 || !IsInitializedTensor(region.weightSourceTensor)) + throw std::runtime_error("SOFIE quantized MatMul storage requires an initialized rank-2 weight tensor"); + + const auto *weightData = fInitializedTensors.at(region.weightSourceTensor).data(); + std::vector perChannelScales; + if (IsPerChannelAxis(region.weightQuant, 1)) + perChannelScales = GetTensorData(region.weightQuant.scaleTensor); + + installStorage(MaterializeQuantizedMatMulWeight(region, *plan, backend, weightData, + weightShape, perChannelScales)); + } + + std::unordered_set consumedOperators; + std::unordered_set pruneCandidates; + std::unordered_set protectedTensors; + for (const auto &[opIndex, backendPlans] : fQuantizationState.loweringPlans) { + auto planIt = backendPlans.find(backend); + if (planIt == backendPlans.end() || !IsQuantizedLoweringAvailable(planIt->second.status) || + planIt->second.weightStorageTensor.empty()) + continue; + consumedOperators.insert(planIt->second.consumedOperatorIndices.begin(), + planIt->second.consumedOperatorIndices.end()); + protectedTensors.insert(planIt->second.weightStorageTensor); + if (!planIt->second.weightScaleTensor.empty()) + protectedTensors.insert(planIt->second.weightScaleTensor); + if (!planIt->second.weightZeroPointTensor.empty()) + protectedTensors.insert(planIt->second.weightZeroPointTensor); + if (auto gemm = fQuantizationState.gemmRegions.find(opIndex); gemm != fQuantizationState.gemmRegions.end()) { + pruneCandidates.insert(gemm->second.weightSourceTensor); + if (!gemm->second.biasSourceTensor.empty()) + protectedTensors.insert(gemm->second.biasSourceTensor); + } + if (auto matmul = fQuantizationState.matmulRegions.find(opIndex); matmul != fQuantizationState.matmulRegions.end()) { + pruneCandidates.insert(matmul->second.weightSourceTensor); + if (!matmul->second.epilogue.biasSourceTensor.empty()) + protectedTensors.insert(matmul->second.epilogue.biasSourceTensor); + } + } + + for (auto opIndex : consumedOperators) { + if (opIndex >= fOperators.size()) + continue; + for (const auto &input : fOperators[opIndex]->GetOpInputTensors()) { + const auto name = UTILITY::Clean_name(std::string(input)); + if (fInitializedTensors.find(name) != fInitializedTensors.end()) + pruneCandidates.insert(name); + } + } + + for (const auto &source : pruneCandidates) { + auto tensor = fInitializedTensors.find(source); + if (tensor == fInitializedTensors.end() || protectedTensors.count(source) != 0) + continue; + + bool hasLiveConsumer = false; + for (std::size_t opIndex = 0; opIndex < fOperators.size() && !hasLiveConsumer; ++opIndex) { + for (const auto &input : fOperators[opIndex]->GetOpInputTensors()) { + if (UTILITY::Clean_name(std::string(input)) == source && consumedOperators.count(opIndex) == 0) { + hasLiveConsumer = true; + break; + } + } + } + if (std::find(fOutputTensorNames.begin(), fOutputTensorNames.end(), source) != fOutputTensorNames.end()) + hasLiveConsumer = true; + + if (!hasLiveConsumer) + tensor->second.SetNotWritable(); + } +} + +void RModel::AddLoweredQuantizedOperators(EQuantizedBackend backend) +{ + auto setKnownTensorType = [this](const std::string &tensorName, ETensorType type) { + if (auto it = fIntermediateTensorInfos.find(tensorName); it != fIntermediateTensorInfos.end()) { + it->second.type = type; + return; + } + if (auto it = fReadyInputTensorInfos.find(tensorName); it != fReadyInputTensorInfos.end()) { + it->second.type = type; + return; + } + if (auto it = fInputTensorInfos.find(tensorName); it != fInputTensorInfos.end()) { + it->second.type = type; + return; + } + if (auto it = fDynamicTensorInfos.find(tensorName); it != fDynamicTensorInfos.end()) { + it->second.type = type; + } + }; + + auto installLoweredOperator = [this, &setKnownTensorType](std::size_t opIndex, + const QuantizedLoweringPlan &plan, + const std::string &inputSourceTensor, + const std::string &outputTensor, + std::unique_ptr lowered) { + if (QuantizedPlanExposesQuantizedInputCarrier(plan)) + setKnownTensorType(inputSourceTensor, TensorTypeForQuantizedStorage(plan.inputStorage)); + if (QuantizedPlanExposesQuantizedOutputCarrier(plan)) + setKnownTensorType(outputTensor, TensorTypeForQuantizedStorage(plan.outputStorage)); + + fLoweredOperators[opIndex] = std::move(lowered); + if (!plan.suppressesGraphOperators) + return; + for (auto consumedOpIndex : plan.consumedOperatorIndices) { + if (consumedOpIndex != opIndex) + fLoweredConsumedOperatorIndices.insert(consumedOpIndex); + } + }; + + for (auto op_idx : SortedQuantizedRegionOperatorIndices(fQuantizationState.gemmRegions)) { + const auto *planPtr = FindQuantizedLoweringPlan(fQuantizationState, op_idx, backend); + if (planPtr == nullptr) + continue; + + auto *gemm = dynamic_cast *>(fOperators[op_idx].get()); + if (!gemm) + throw std::runtime_error("SOFIE quantized Gemm region is attached to a non-float Gemm operator"); + + const auto &plan = *planPtr; + if (!IsQuantizedLoweringAvailable(plan.status)) + continue; + + const auto regionIt = fQuantizationState.gemmRegions.find(op_idx); + if (regionIt == fQuantizationState.gemmRegions.end()) + throw std::runtime_error("SOFIE quantized Gemm lowering plan has no matching region"); + const auto ®ion = regionIt->second; + installLoweredOperator(op_idx, plan, region.inputSourceTensor, region.outputTensor, + std::make_unique(region, plan, MakeQuantizedGemmCodegenContext(*gemm))); + } + + for (auto op_idx : SortedQuantizedRegionOperatorIndices(fQuantizationState.matmulRegions)) { + const auto *planPtr = FindQuantizedLoweringPlan(fQuantizationState, op_idx, backend); + if (planPtr == nullptr) + continue; + + auto *gemm = dynamic_cast *>(fOperators[op_idx].get()); + if (!gemm) + throw std::runtime_error("SOFIE quantized MatMul region is attached to a non-float Gemm-spelled MatMul operator"); + + const auto &plan = *planPtr; + if (!IsQuantizedLoweringAvailable(plan.status)) + continue; + + const auto regionIt = fQuantizationState.matmulRegions.find(op_idx); + if (regionIt == fQuantizationState.matmulRegions.end()) + throw std::runtime_error("SOFIE quantized MatMul lowering plan has no matching region"); + const auto ®ion = regionIt->second; + installLoweredOperator(op_idx, plan, region.inputSourceTensor, region.outputTensor, + std::make_unique(region, plan, MakeQuantizedMatMulCodegenContext(*gemm))); + } +} + +void RModel::AddQuantizedGeneratedHeaders(EQuantizedBackend backend) +{ + for (const auto &[opIndex, backendPlans] : fQuantizationState.loweringPlans) { + auto planIt = backendPlans.find(backend); + if (planIt == backendPlans.end()) + continue; + const auto &plan = planIt->second; + if (!IsQuantizedLoweringAvailable(plan.status) || !plan.suppressesGraphOperators) + continue; + if (backend == EQuantizedBackend::CPU && QuantizedPlanUsesPrequantizedWeights(plan) && + plan.weightLayout == EQuantizedLayout::PackedCPU) + AddNeededCustomHeader("SOFIE/SOFIE_Quantized.hxx"); + if (backend == EQuantizedBackend::ALPAKA && IsOptimizedQuantizedAlpakaPlainDevicePlan(plan)) { + AddNeededCustomHeader("SOFIE/SOFIE_QuantizedAlpaka.hxx"); + } + } +} + +} // namespace SOFIE diff --git a/core/src/RQuantization_Analysis.cxx b/core/src/RQuantization_Analysis.cxx new file mode 100644 index 0000000..f9ba65d --- /dev/null +++ b/core/src/RQuantization_Analysis.cxx @@ -0,0 +1,241 @@ +#include "SOFIE/RQuantization_Analysis.hxx" +#include "SOFIE/ROperator_BasicBinary.hxx" +#include "SOFIE/RQuantization_DenseLinear.hxx" + +#include +#include +#include + +namespace SOFIE { + +QuantizedDenseLinearPatternMatch MatchQuantizedDenseLinearPattern( + const ROperator_Gemm &gemm, std::size_t opIndex, + const std::function(const std::string &)> &tensorShape) +{ + QuantizedDenseLinearPatternMatch match; + auto ®ion = match.region; + region.status = EQuantizedLoweringStatus::SemanticUnsupported; + region.alpha = gemm.GetAlpha(); + region.beta = gemm.GetBeta(); + region.transA = gemm.GetTransA(); + region.transB = gemm.GetTransB(); + region.gemmOpIndex = opIndex; + + const auto inputs = gemm.GetOpInputTensors(); + const auto outputs = gemm.GetOpOutputTensors(); + if (inputs.size() < 2 || outputs.size() != 1) { + match.reasons.push_back("Gemm does not have the expected input/output arity"); + } else { + region.inputTensor = std::string(inputs[0]); + region.weightTensor = std::string(inputs[1]); + if (inputs.size() >= 3) + region.biasTensor = std::string(inputs[2]); + region.gemmOutputTensor = std::string(outputs[0]); + } + + const bool matmul = inputs.size() == 2 && region.alpha == 1.0f && region.beta == 0.0f && + region.transA == 0 && region.transB == 0; + match.hasInlineMatMulBias = inputs.size() == 3 && region.alpha == 1.0f && region.beta == 1.0f && + region.transA == 0 && region.transB == 0; + match.isMatMul = matmul || match.hasInlineMatMulBias; + if (match.isMatMul) { + const auto inputShape = region.inputTensor.empty() ? std::vector{} : tensorShape(region.inputTensor); + const auto weightShape = region.weightTensor.empty() ? std::vector{} : tensorShape(region.weightTensor); + const auto outputShape = region.gemmOutputTensor.empty() ? std::vector{} : tensorShape(region.gemmOutputTensor); + match.matmulShape = AssessQuantizedMatMulShape(inputShape, weightShape, outputShape); + if (!QuantizedMatMulShapeIsRecognized(match.matmulShape)) { + match.reasons.push_back(match.matmulShape.reason.empty() + ? "MatMul shape is not recognized by quantized dense-linear analysis" + : match.matmulShape.reason); + } + } else { + CheckQuantizedGemmAttributes(region, match.reasons); + if (!region.inputTensor.empty()) + CheckQuantizedGemmRank2Shape(tensorShape(region.inputTensor), "input", match.reasons); + if (!region.weightTensor.empty()) + CheckQuantizedGemmRank2Shape(tensorShape(region.weightTensor), "weight", match.reasons); + if (!region.gemmOutputTensor.empty()) + CheckQuantizedGemmRank2Shape(tensorShape(region.gemmOutputTensor), "Gemm output", match.reasons); + } + return match; +} + +QuantizationGraphIndex BuildQuantizationGraphIndex(const std::vector> &operators) +{ + QuantizationGraphIndex graph; + for (std::size_t opIndex = 0; opIndex < operators.size(); ++opIndex) { + for (const auto &output : operators[opIndex]->GetOpOutputTensors()) + graph.producerByTensor[std::string(output)] = opIndex; + for (const auto &input : operators[opIndex]->GetOpInputTensors()) + graph.consumersByTensor[std::string(input)].push_back(opIndex); + } + return graph; +} + +std::optional MatchQuantizationBoundaryProducer( + const QuantizationGraphIndex &graph, const std::vector> &operators, + const std::string &tensor, const std::string &role, std::vector &reasons) +{ + auto producer = graph.producerByTensor.find(tensor); + if (producer == graph.producerByTensor.end()) { + reasons.push_back(role + " tensor has no producer quantization boundary"); + return std::nullopt; + } + if (!operators[producer->second]->IsQuantizationBoundary()) { + reasons.push_back(role + " tensor producer is not a quantization boundary"); + return std::nullopt; + } + return producer->second; +} + +std::optional MatchSingleTensorConsumer(const QuantizationGraphIndex &graph, + const std::string &tensor, + const std::string &role, + std::vector &reasons) +{ + auto consumers = graph.consumersByTensor.find(tensor); + if (consumers == graph.consumersByTensor.end() || consumers->second.empty()) { + reasons.push_back(role + " has no consumer"); + return std::nullopt; + } + if (consumers->second.size() != 1) { + reasons.push_back(role + " has multiple consumers"); + return std::nullopt; + } + return consumers->second.front(); +} + +bool IsFloatAddOperator(const ROperator &op) +{ + return dynamic_cast *>(&op) != nullptr; +} + +void CheckQuantizationInfo(const QuantizationInfo &info, const std::string &role, + std::vector &reasons) +{ + if (info.bitWidth == 0 || info.bitWidth > 8) + reasons.push_back(role + " bit width is not in the supported quantized integer range [1, 8]"); + if (info.scale <= 0.0 || !std::isfinite(info.scale)) + reasons.push_back(role + " scale is not positive and finite"); + if (info.rounding != EQuantizationRoundingMode::ROUND) + reasons.push_back(role + " rounding mode is not ROUND"); + if (info.overflow != EQuantizationOverflowMode::SAT && info.overflow != EQuantizationOverflowMode::SAT_SYM) + reasons.push_back(role + " overflow mode is unsupported"); +} + +void CheckQuantizedGemmAttributes(const QuantizedGemmRegion ®ion, + std::vector &reasons) +{ + if (std::fabs(region.alpha - 1.0f) > 0.0f) + reasons.push_back("Gemm alpha is not 1"); + if (std::fabs(region.beta - 1.0f) > 0.0f) + reasons.push_back("Gemm beta is not 1"); + if (region.transA != 0) + reasons.push_back("Gemm transA is not 0"); + if (region.transB != 1) + reasons.push_back("Gemm transB is not 1"); +} + +void CheckQuantizedGemmRank2Shape(const std::vector &shape, + const std::string &role, + std::vector &reasons) +{ + if (shape.size() != 2) + reasons.push_back(role + " tensor is not rank-2 for quantized Gemm lowering"); +} + +std::vector QuantizedGemmLoweringUnsupportedReasons(const QuantizedGemmRegion ®ion) +{ + return DenseLinearQuantizationParameterUnsupportedReasons(region.inputQuant, region.weightQuant, + region.outputQuant, region.biasQuant, 0, + "quantized Gemm lowering"); +} + +std::string JoinQuantizationReasons(const std::vector &reasons) +{ + std::ostringstream out; + for (std::size_t i = 0; i < reasons.size(); ++i) { + if (i != 0) + out << "; "; + out << reasons[i]; + } + return out.str(); +} + +std::vector QuantizedGemmConsumedOperatorIndices(const QuantizedGemmRegion ®ion) +{ + std::vector indices = {region.inputQuantOpIndex, region.weightQuantOpIndex, + region.gemmOpIndex, region.outputQuantOpIndex}; + if (region.biasQuantOpIndex) + indices.push_back(*region.biasQuantOpIndex); + std::sort(indices.begin(), indices.end()); + return indices; +} + +std::vector QuantizedMatMulConsumedOperatorIndices(const QuantizedMatMulRegion ®ion) +{ + std::vector indices = {region.inputQuantOpIndex, region.weightQuantOpIndex, + region.matmulOpIndex, region.outputQuantOpIndex}; + if (region.epilogue.addOpIndex) + indices.push_back(*region.epilogue.addOpIndex); + std::sort(indices.begin(), indices.end()); + return indices; +} + +QuantizedMatMulRegion MakeQuantizedMatMulRegionFromGemmLikeRegion(const QuantizedGemmRegion ®ion) +{ + QuantizedMatMulRegion matmul; + matmul.inputTensor = region.inputTensor; + matmul.weightTensor = region.weightTensor; + matmul.matmulOutputTensor = region.gemmOutputTensor; + matmul.outputTensor = region.outputTensor; + matmul.inputSourceTensor = region.inputSourceTensor; + matmul.weightSourceTensor = region.weightSourceTensor; + matmul.inputQuantOpIndex = region.inputQuantOpIndex; + matmul.weightQuantOpIndex = region.weightQuantOpIndex; + matmul.matmulOpIndex = region.gemmOpIndex; + matmul.outputQuantOpIndex = region.outputQuantOpIndex; + matmul.inputQuant = region.inputQuant; + matmul.weightQuant = region.weightQuant; + matmul.outputQuant = region.outputQuant; + return matmul; +} + +bool IsDenseLinearBiasLikeShape(const std::vector &biasShape, + const std::vector &outputShape) +{ + if (outputShape.empty() || biasShape.empty()) + return false; + const auto n = outputShape.back(); + if (biasShape.size() == 1) + return biasShape[0] == n; + if (biasShape.back() != n) + return false; + for (std::size_t i = 0; i + 1 < biasShape.size(); ++i) { + if (biasShape[i] != 1) + return false; + } + return true; +} + +QuantizationInfo MakeAccumulatorBiasQuantization(const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant) +{ + QuantizationInfo biasQuant; + biasQuant.bitWidth = 32; + biasQuant.isSigned = true; + biasQuant.narrow = false; + biasQuant.scale = inputQuant.scale * weightQuant.scale; + biasQuant.zeroPoint = 0; + biasQuant.rounding = inputQuant.rounding; + biasQuant.overflow = EQuantizationOverflowMode::SAT; + biasQuant.granularity = weightQuant.granularity == EQuantizationGranularity::PerChannel + ? EQuantizationGranularity::PerChannel + : EQuantizationGranularity::PerTensor; + biasQuant.axis = weightQuant.granularity == EQuantizationGranularity::PerChannel ? 0 : -1; + biasQuant.scaleTensor = weightQuant.scaleTensor; + biasQuant.zeroPointTensor = weightQuant.zeroPointTensor; + return biasQuant; +} + +} // namespace SOFIE diff --git a/core/src/RQuantization_DenseLinear.cxx b/core/src/RQuantization_DenseLinear.cxx new file mode 100644 index 0000000..35896f1 --- /dev/null +++ b/core/src/RQuantization_DenseLinear.cxx @@ -0,0 +1,782 @@ +#include "SOFIE/RQuantization_DenseLinear.hxx" +#include "SOFIE/RQuantization_Analysis.hxx" + +#include +#include +#include + +namespace SOFIE { + +namespace { + +constexpr std::size_t kCublasLtInt8Alignment = 16; +constexpr std::size_t kCublasLtMinOptimizedMacs = 1'000'000; +constexpr double kCublasLtPaddingCandidateMaxWorkRatio = 1.50; +constexpr std::size_t kCublasLtMinProfitablePaddedMacs = 100'000; +constexpr double kCublasLtProfitablePaddedMaxWorkRatio = 1.10; +constexpr std::size_t kCublasLtMinProfitablePaddedK = 64; +constexpr std::size_t kCublasLtMinProfitablePaddedN = 64; + +std::size_t RoundUpToMultiple(std::size_t value, std::size_t multiple) +{ + if (multiple == 0 || value == 0) + return value; + return ((value + multiple - 1) / multiple) * multiple; +} + +bool IsAlignedTo(std::size_t value, std::size_t multiple) +{ + return multiple != 0 && (value % multiple) == 0; +} + +bool IsScalarZeroPointZero(const QuantizationInfo &info) +{ + return info.zeroPoint == 0; +} + +} // namespace + +std::vector QuantizeTensorToInt8(const float *data, std::size_t length, const QuantizationInfo &info) +{ + std::vector quantized(length); + for (std::size_t i = 0; i < length; ++i) { + quantized[i] = static_cast(QuantizeScalarToIntegerGrid(data[i], info)); + } + return quantized; +} + +std::vector QuantizeTensorToUInt8(const float *data, std::size_t length, const QuantizationInfo &info) +{ + std::vector quantized(length); + for (std::size_t i = 0; i < length; ++i) { + quantized[i] = static_cast(QuantizeScalarToIntegerGrid(data[i], info)); + } + return quantized; +} + +std::vector QuantizeGemmWeightTensorToInt8(const float *data, std::size_t n, std::size_t k, + const QuantizationInfo &info, + const std::vector &perChannelScale) +{ + std::vector quantized(n * k); + for (std::size_t col = 0; col < n; ++col) { + QuantizationInfo channelInfo = info; + channelInfo.granularity = EQuantizationGranularity::PerTensor; + channelInfo.axis = -1; + channelInfo.scale = static_cast(perChannelScale[col]); + channelInfo.zeroPoint = 0; + for (std::size_t kk = 0; kk < k; ++kk) { + quantized[col * k + kk] = static_cast(QuantizeScalarToIntegerGrid(data[col * k + kk], channelInfo)); + } + } + return quantized; +} + +std::vector QuantizeMatMulWeightTensorToInt8Transposed(const float *data, std::size_t k, std::size_t n, + const QuantizationInfo &info, + const std::vector &perChannelScale) +{ + std::vector quantized(n * k); + for (std::size_t col = 0; col < n; ++col) { + QuantizationInfo channelInfo = info; + if (!perChannelScale.empty()) { + channelInfo.granularity = EQuantizationGranularity::PerTensor; + channelInfo.axis = -1; + channelInfo.scale = static_cast(perChannelScale[col]); + channelInfo.zeroPoint = 0; + } + for (std::size_t kk = 0; kk < k; ++kk) { + quantized[col * k + kk] = static_cast(QuantizeScalarToIntegerGrid(data[kk * n + col], channelInfo)); + } + } + return quantized; +} + +std::vector QuantizeGemmWeightTensorToInt8Padded(const float *data, std::size_t n, std::size_t k, + std::size_t physicalN, std::size_t physicalK, + const QuantizationInfo &info, + const std::vector &perChannelScale) +{ + std::vector quantized(physicalN * physicalK, 0); + for (std::size_t col = 0; col < n; ++col) { + QuantizationInfo channelInfo = info; + if (!perChannelScale.empty()) { + channelInfo.granularity = EQuantizationGranularity::PerTensor; + channelInfo.axis = -1; + channelInfo.scale = static_cast(perChannelScale[col]); + channelInfo.zeroPoint = 0; + } + for (std::size_t kk = 0; kk < k; ++kk) { + quantized[col * physicalK + kk] = static_cast(QuantizeScalarToIntegerGrid(data[col * k + kk], channelInfo)); + } + } + return quantized; +} + +std::vector QuantizeMatMulWeightTensorToInt8TransposedPadded(const float *data, std::size_t k, std::size_t n, + std::size_t physicalK, std::size_t physicalN, + const QuantizationInfo &info, + const std::vector &perChannelScale) +{ + std::vector quantized(physicalN * physicalK, 0); + for (std::size_t col = 0; col < n; ++col) { + QuantizationInfo channelInfo = info; + if (!perChannelScale.empty()) { + channelInfo.granularity = EQuantizationGranularity::PerTensor; + channelInfo.axis = -1; + channelInfo.scale = static_cast(perChannelScale[col]); + channelInfo.zeroPoint = 0; + } + for (std::size_t kk = 0; kk < k; ++kk) { + quantized[col * physicalK + kk] = static_cast(QuantizeScalarToIntegerGrid(data[kk * n + col], channelInfo)); + } + } + return quantized; +} + +bool IsScalarPerTensor(const QuantizationInfo &info) +{ + return info.granularity == EQuantizationGranularity::PerTensor && info.axis == -1; +} + +bool IsPerChannelAxis(const QuantizationInfo &info, int axis) +{ + return info.granularity == EQuantizationGranularity::PerChannel && info.axis == axis; +} + +std::vector DenseLinearQuantizationParameterUnsupportedReasons( + const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant, + const QuantizationInfo &outputQuant, + const std::optional &biasQuant, + int expectedWeightPerChannelAxis, + const std::string &operatorName) +{ + std::vector reasons; + + const bool inputPerTensor = IsScalarPerTensor(inputQuant); + const bool outputPerTensor = IsScalarPerTensor(outputQuant); + const bool weightPerTensor = IsScalarPerTensor(weightQuant); + const bool weightPerChannel = IsPerChannelAxis(weightQuant, expectedWeightPerChannelAxis); + + if (!inputPerTensor) + reasons.push_back(operatorName + " input quantization requires per-tensor scalar parameters"); + if (!outputPerTensor) + reasons.push_back(operatorName + " output quantization requires per-tensor scalar parameters"); + if (!weightPerTensor && !weightPerChannel) { + reasons.push_back(operatorName + " weight quantization supports per-tensor scalar parameters or per-output-channel axis " + + std::to_string(expectedWeightPerChannelAxis)); + } + + if (biasQuant.has_value()) { + const bool biasPerOutputChannel = IsPerChannelAxis(*biasQuant, 0) || + IsPerChannelAxis(*biasQuant, expectedWeightPerChannelAxis); + if (weightPerChannel && !biasPerOutputChannel) { + reasons.push_back(operatorName + " per-channel weight quantization requires per-output-channel bias parameters"); + } else if (weightPerTensor && !IsScalarPerTensor(*biasQuant)) { + reasons.push_back(operatorName + " per-tensor weight quantization requires per-tensor bias parameters"); + } + } + + return reasons; +} + +QuantizedDenseLinearProfileAssessment AssessDenseLinearComputeProfile( + const QuantizationInfo &inputQuant, + const QuantizationInfo &weightQuant, + const QuantizationInfo &outputQuant, + int expectedWeightPerChannelAxis, + const std::string &operatorName) +{ + QuantizedDenseLinearProfileAssessment assessment; + + const bool input8 = inputQuant.bitWidth == 8; + const bool weight8 = weightQuant.bitWidth == 8; + const bool output8 = outputQuant.bitWidth == 8; + if (!input8) + assessment.reasons.push_back(operatorName + " input bit width is not 8"); + if (!weight8) + assessment.reasons.push_back(operatorName + " weight bit width is not 8"); + if (!output8) + assessment.reasons.push_back(operatorName + " output bit width is not 8"); + + const bool inputPerTensor = IsScalarPerTensor(inputQuant); + const bool outputPerTensor = IsScalarPerTensor(outputQuant); + const bool weightPerTensor = IsScalarPerTensor(weightQuant); + const bool weightPerChannel = IsPerChannelAxis(weightQuant, expectedWeightPerChannelAxis); + auto parameterReasons = DenseLinearQuantizationParameterUnsupportedReasons( + inputQuant, weightQuant, outputQuant, std::nullopt, expectedWeightPerChannelAxis, operatorName); + assessment.reasons.insert(assessment.reasons.end(), parameterReasons.begin(), parameterReasons.end()); + + const bool hasAsymmetricInput = !IsScalarZeroPointZero(inputQuant); + const bool hasAsymmetricWeight = !IsScalarZeroPointZero(weightQuant); + const bool hasAsymmetricOutput = !IsScalarZeroPointZero(outputQuant); + if (hasAsymmetricInput) { + assessment.reasons.push_back(operatorName + " input zero point is nonzero; cuBLASLt int8 lowering requires row-sum zero-point correction"); + } + if (hasAsymmetricWeight) { + assessment.reasons.push_back(operatorName + " weight zero point is nonzero; cuBLASLt int8 lowering requires activation-sum/weight-sum zero-point correction"); + } + if (hasAsymmetricOutput) { + assessment.reasons.push_back(operatorName + " output zero point is nonzero; cuBLASLt int8 lowering emits zero-centered int8 output only"); + } + + if (!input8 || !weight8 || !output8 || !inputPerTensor || !outputPerTensor || (!weightPerTensor && !weightPerChannel)) { + assessment.profile = EQuantizedComputeProfile::UnsupportedDenseLinearRank2; + return assessment; + } + + if (hasAsymmetricInput || hasAsymmetricWeight || hasAsymmetricOutput) { + assessment.profile = EQuantizedComputeProfile::AsymmetricZeroPointRank2; + return assessment; + } + + if (inputQuant.isSigned && weightQuant.isSigned && outputQuant.isSigned) { + assessment.profile = weightPerChannel ? EQuantizedComputeProfile::SignedInt8PerTensorActivationPerChannelWeightRank2 + : EQuantizedComputeProfile::SignedInt8SymmetricPerTensorRank2; + assessment.cublasLtOptimizedCandidate = true; + return assessment; + } + + if (!inputQuant.isSigned && weightQuant.isSigned) { + assessment.profile = EQuantizedComputeProfile::UnsignedInt8ActivationSignedInt8WeightRank2; + assessment.reasons.push_back(operatorName + " unsigned activation profile is recognized but the cuBLASLt int8 lowering supports signed int8 carriers only"); + return assessment; + } + + if (!inputQuant.isSigned && !weightQuant.isSigned) { + assessment.profile = EQuantizedComputeProfile::UnsignedInt8SymmetricRank2; + assessment.reasons.push_back(operatorName + " unsigned activation/weight profile is recognized but the cuBLASLt int8 lowering supports signed int8 carriers only"); + return assessment; + } + + assessment.profile = EQuantizedComputeProfile::UnsupportedDenseLinearRank2; + assessment.reasons.push_back(operatorName + " signedness combination is not supported by dense-linear int8 lowering"); + return assessment; +} + +QuantizedDenseLinearShapePolicy MakeCublasLtShapePolicy(std::size_t m, std::size_t k, std::size_t n) +{ + QuantizedDenseLinearShapePolicy policy; + policy.logicalM = m; + policy.logicalK = k; + policy.logicalN = n; + policy.physicalM = RoundUpToMultiple(m, kCublasLtInt8Alignment); + policy.physicalK = RoundUpToMultiple(k, kCublasLtInt8Alignment); + policy.physicalN = RoundUpToMultiple(n, kCublasLtInt8Alignment); + + policy.logicalMacs = m * k * n; + policy.physicalMacs = policy.physicalM * policy.physicalK * policy.physicalN; + policy.minimumOptimizedMacs = kCublasLtMinOptimizedMacs; + policy.belowMinimumWork = policy.logicalMacs < policy.minimumOptimizedMacs; + policy.paddingWorkRatio = policy.logicalMacs > 0 ? static_cast(policy.physicalMacs) / + static_cast(policy.logicalMacs) : 1.0; + + std::ostringstream reason; + reason << "logical M/K/N=" << policy.logicalM << "/" << policy.logicalK << "/" << policy.logicalN + << ", physical M/K/N=" << policy.physicalM << "/" << policy.physicalK << "/" << policy.physicalN + << ", logical MACs=" << policy.logicalMacs + << ", physical MACs=" << policy.physicalMacs + << ", minimum optimized MACs=" << policy.minimumOptimizedMacs + << ", padding work ratio=" << policy.paddingWorkRatio; + + if (IsAlignedTo(m, kCublasLtInt8Alignment) && IsAlignedTo(k, kCublasLtInt8Alignment) && + IsAlignedTo(n, kCublasLtInt8Alignment)) { + if (policy.belowMinimumWork) { + policy.policy = EQuantizedShapePolicy::ExactTooSmall; + policy.reason = "exact cuBLASLt int8 shape below minimum optimized work threshold; " + reason.str(); + } else { + policy.policy = EQuantizedShapePolicy::Exact; + policy.reason = "exact cuBLASLt int8 shape; " + reason.str(); + } + } else if (policy.paddingWorkRatio <= kCublasLtPaddingCandidateMaxWorkRatio) { + policy.policy = EQuantizedShapePolicy::PaddedCandidate; + policy.reason = "padded cuBLASLt candidate; profitability policy selects executable padded lowering; " + reason.str(); + } else { + policy.policy = EQuantizedShapePolicy::Fallback; + policy.reason = "padding too expensive for cuBLASLt candidate; " + reason.str(); + } + return policy; +} + + +bool IsProfitableCublasLtPaddedDenseLinearPolicy(const QuantizedDenseLinearShapePolicy &policy) +{ + return policy.policy == EQuantizedShapePolicy::PaddedCandidate && + policy.logicalMacs >= kCublasLtMinProfitablePaddedMacs && + policy.paddingWorkRatio <= kCublasLtProfitablePaddedMaxWorkRatio && + policy.logicalK >= kCublasLtMinProfitablePaddedK && + policy.logicalN >= kCublasLtMinProfitablePaddedN; +} + +std::string ExplainCublasLtPaddedDenseLinearProfitability(const QuantizedDenseLinearShapePolicy &policy) +{ + std::ostringstream reason; + reason << "padded cuBLASLt profitability requires logical MACs >= " << kCublasLtMinProfitablePaddedMacs + << ", padding work ratio <= " << kCublasLtProfitablePaddedMaxWorkRatio + << ", logical K >= " << kCublasLtMinProfitablePaddedK + << ", and logical N >= " << kCublasLtMinProfitablePaddedN + << "; observed logical MACs=" << policy.logicalMacs + << ", padding work ratio=" << policy.paddingWorkRatio + << ", logical K=" << policy.logicalK + << ", logical N=" << policy.logicalN; + return reason.str(); +} + +namespace { + +std::string JoinCapabilityReasons(const std::vector &reasons) +{ + std::ostringstream out; + for (std::size_t i = 0; i < reasons.size(); ++i) { + if (i != 0) + out << "; "; + out << reasons[i]; + } + return out.str(); +} + +std::size_t DenseLinearLeadingElementCount(const std::vector &shape) +{ + if (shape.size() <= 1) + return 0; + std::size_t count = 1; + for (std::size_t i = 0; i + 1 < shape.size(); ++i) + count *= shape[i]; + return count; +} + +bool DenseLinearShapeMatches(const std::vector &lhs, const std::vector &rhs) +{ + return lhs.size() == rhs.size() && std::equal(lhs.begin(), lhs.end(), rhs.begin()); +} + +std::string DenseLinearShapeToString(const std::vector &shape) +{ + std::ostringstream out; + out << "["; + for (std::size_t i = 0; i < shape.size(); ++i) { + if (i != 0) + out << ","; + out << shape[i]; + } + out << "]"; + return out.str(); +} + +void AddCapabilityReason(std::vector &reasons, std::string reason) +{ + reasons.push_back(std::move(reason)); +} + +} // namespace + +QuantizedMatMulShapeAssessment AssessQuantizedMatMulShape( + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape) +{ + QuantizedMatMulShapeAssessment assessment; + + if (inputShape.size() < 2) + assessment.unsupportedReasons.push_back("MatMul input rank is less than 2"); + if (weightShape.size() < 2) + assessment.unsupportedReasons.push_back("MatMul weight rank is less than 2"); + if (!outputShape.empty() && outputShape.size() < 2) + assessment.unsupportedReasons.push_back("MatMul output rank is less than 2"); + + if (!assessment.unsupportedReasons.empty()) { + assessment.kind = EQuantizedMatMulShapeKind::Unsupported; + assessment.reason = JoinCapabilityReasons(assessment.unsupportedReasons); + return assessment; + } + + const auto inputK = inputShape.back(); + const auto weightK = weightShape[weightShape.size() - 2]; + const auto weightN = weightShape.back(); + + if (inputK == 0 || weightK == 0 || weightN == 0) + assessment.unsupportedReasons.push_back("MatMul K and N dimensions must be nonzero"); + if (inputK != weightK) + assessment.unsupportedReasons.push_back("MatMul input K does not match weight K"); + + std::vector expectedOutput = inputShape; + expectedOutput.back() = weightN; + + if (weightShape.size() == 2) { + if (!outputShape.empty() && !DenseLinearShapeMatches(outputShape, expectedOutput)) { + assessment.unsupportedReasons.push_back( + "MatMul output shape does not match X" + DenseLinearShapeToString(inputShape) + + " @ W" + DenseLinearShapeToString(weightShape) + + " -> Y" + DenseLinearShapeToString(expectedOutput)); + } + if (!assessment.unsupportedReasons.empty()) { + assessment.kind = EQuantizedMatMulShapeKind::Unsupported; + assessment.reason = JoinCapabilityReasons(assessment.unsupportedReasons); + return assessment; + } + + assessment.logicalM = inputShape.size() == 2 ? inputShape[0] : DenseLinearLeadingElementCount(inputShape); + assessment.logicalK = inputK; + assessment.logicalN = weightN; + assessment.flattenedInputShape = { assessment.logicalM, assessment.logicalK }; + assessment.flattenedOutputShape = { assessment.logicalM, assessment.logicalN }; + if (inputShape.size() == 2) { + assessment.kind = EQuantizedMatMulShapeKind::Rank2; + assessment.reason = "rank-2 MatMul shape X[M,K] @ W[K,N] -> Y[M,N]"; + } else { + assessment.kind = EQuantizedMatMulShapeKind::FlattenableProjection; + assessment.reason = "flattenable projection MatMul shape " + DenseLinearShapeToString(inputShape) + + " @ " + DenseLinearShapeToString(weightShape) + + " can be viewed as [prod(prefix),K] @ [K,N]"; + } + return assessment; + } + + if (inputShape.size() >= 3 && weightShape.size() >= 3) { + if (!assessment.unsupportedReasons.empty()) { + assessment.kind = EQuantizedMatMulShapeKind::Unsupported; + assessment.reason = JoinCapabilityReasons(assessment.unsupportedReasons); + return assessment; + } + assessment.logicalM = inputShape[inputShape.size() - 2]; + assessment.logicalK = inputK; + assessment.logicalN = weightN; + assessment.kind = EQuantizedMatMulShapeKind::TrueBatched; + assessment.reason = "true batched MatMul requires strided-batched quantized lowering"; + return assessment; + } + + assessment.kind = EQuantizedMatMulShapeKind::Unsupported; + assessment.unsupportedReasons.push_back("MatMul broadcasted shape family is not a dense projection or true batched MatMul"); + assessment.reason = JoinCapabilityReasons(assessment.unsupportedReasons); + return assessment; +} + +QuantizedDenseLinearCublasLtCapability AssessCublasLtDenseLinearCapability( + const QuantizedDenseLinearOperands &operands) +{ + QuantizedDenseLinearCublasLtCapability capability; + std::vector semanticReasons; + + if (operands.inputShape.size() != 2) + AddCapabilityReason(semanticReasons, operands.operatorName + " input rank is not 2"); + if (operands.weightShape.size() != 2) + AddCapabilityReason(semanticReasons, operands.operatorName + " weight rank is not 2"); + if (!operands.outputShape.empty() && operands.outputShape.size() != 2) + AddCapabilityReason(semanticReasons, operands.operatorName + " output rank is not 2"); + + if (operands.inputShape.size() == 2 && operands.weightShape.size() == 2) { + const auto m = operands.inputShape[0]; + const auto k = operands.inputShape[1]; + const auto n = operands.weightOutputChannelAxis == 0 ? operands.weightShape[0] : operands.weightShape[1]; + const auto weightK = operands.weightOutputChannelAxis == 0 ? operands.weightShape[1] : operands.weightShape[0]; + if (m == 0 || n == 0 || k == 0) + AddCapabilityReason(semanticReasons, operands.operatorName + " M, N, and K must be nonzero"); + if (k != weightK) + AddCapabilityReason(semanticReasons, operands.operatorName + " input K does not match weight K"); + if (!operands.outputShape.empty() && operands.outputShape.size() == 2 && + (operands.outputShape[0] != m || operands.outputShape[1] != n)) { + AddCapabilityReason(semanticReasons, operands.operatorName + " output shape does not match X[M,K] @ W -> Y[M,N]"); + } + if (m != 0 && n != 0 && k != 0 && k == weightK) + capability.shapePolicy = MakeCublasLtShapePolicy(m, k, n); + } + + const auto computeProfile = AssessDenseLinearComputeProfile(operands.inputQuant, operands.weightQuant, + operands.outputQuant, + operands.weightOutputChannelAxis, + operands.operatorName); + semanticReasons.insert(semanticReasons.end(), computeProfile.reasons.begin(), computeProfile.reasons.end()); + + const bool perTensorWeight = IsScalarPerTensor(operands.weightQuant); + const bool perChannelWeight = IsPerChannelAxis(operands.weightQuant, operands.weightOutputChannelAxis); + if (operands.biasQuant.has_value()) { + const bool biasPerOutputChannel = IsPerChannelAxis(*operands.biasQuant, 0) || + IsPerChannelAxis(*operands.biasQuant, operands.weightOutputChannelAxis); + if (perChannelWeight && !biasPerOutputChannel) { + AddCapabilityReason(semanticReasons, operands.operatorName + " per-channel weight requires per-output-channel bias parameters"); + } + if (!perChannelWeight && !IsScalarPerTensor(*operands.biasQuant)) { + AddCapabilityReason(semanticReasons, operands.operatorName + " per-tensor weight requires per-tensor bias parameters"); + } + } + + if (!semanticReasons.empty()) { + capability.shapePolicy.policy = EQuantizedShapePolicy::Unsupported; + capability.shapePolicy.reason = "cuBLASLt semantic requirements are not met"; + capability.profile = computeProfile.profile; + capability.reason = JoinCapabilityReasons(semanticReasons); + capability.tag = "cublaslt_dense_linear_profile_unsupported"; + return capability; + } + + capability.profile = computeProfile.profile; + if (capability.shapePolicy.policy == EQuantizedShapePolicy::Exact) { + capability.optimized = true; + capability.tag = perChannelWeight ? "cublaslt_i8i8_per_channel_weight_rank2_exact" + : "cublaslt_i8i8_symmetric_per_tensor_rank2_exact"; + capability.reason = perChannelWeight ? + "cuBLASLt optimized signed-int8 per-tensor activation/per-channel weight rank-2 exact-shape " + operands.operatorName + "; " + capability.shapePolicy.reason : + "cuBLASLt optimized signed-int8 symmetric per-tensor rank-2 exact-shape " + operands.operatorName + "; " + capability.shapePolicy.reason; + } else if (capability.shapePolicy.policy == EQuantizedShapePolicy::ExactTooSmall) { + capability.tag = perChannelWeight ? "cublaslt_i8i8_per_channel_weight_rank2_exact_too_small" + : "cublaslt_i8i8_symmetric_per_tensor_rank2_exact_too_small"; + capability.reason = "cuBLASLt exact-shape execution is legal but below the minimum optimized work threshold; " + + capability.shapePolicy.reason; + } else if (capability.shapePolicy.policy == EQuantizedShapePolicy::PaddedCandidate) { + capability.tag = perChannelWeight ? "cublaslt_i8i8_per_channel_weight_rank2_padded_candidate" + : "cublaslt_i8i8_symmetric_per_tensor_rank2_padded_candidate"; + capability.reason = "cuBLASLt padded execution is shape-compatible; profitability policy selects executable padded lowering; " + + capability.shapePolicy.reason; + } else { + capability.tag = perChannelWeight ? "cublaslt_i8i8_per_channel_weight_rank2_shape_fallback" + : "cublaslt_i8i8_symmetric_per_tensor_rank2_shape_fallback"; + capability.reason = capability.shapePolicy.reason.empty() ? + "cuBLASLt shape policy is unavailable" : capability.shapePolicy.reason; + } + return capability; +} + +QuantizedLoweringPlan MakeUnsupportedQuantizedMatMulPlan(const QuantizedMatMulRegion ®ion, + EQuantizedBackend backend, + std::string reason, + bool preservesSemantics) +{ + QuantizedLoweringPlan plan; + plan.backend = backend; + plan.status = preservesSemantics ? EQuantizedLoweringStatus::BackendUnsupported + : EQuantizedLoweringStatus::SemanticUnsupported; + plan.reason = std::move(reason); + plan.inputStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.weightStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.biasStorage = EQuantizedStorageType::UNDEFINED; + plan.accumulatorStorage = EQuantizedStorageType::UNDEFINED; + plan.outputStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.inputCarrierMode = preservesSemantics ? EQuantizedCarrierMode::Float : EQuantizedCarrierMode::UNDEFINED; + plan.outputMode = preservesSemantics ? EQuantizedOutputMode::ExactFakeQuantFloat : EQuantizedOutputMode::UNDEFINED; + plan.computeProfile = preservesSemantics ? EQuantizedComputeProfile::GenericRecognized : EQuantizedComputeProfile::UNDEFINED; + plan.capabilityTag = preservesSemantics ? "matmul_recognized_backend_unsupported" : "matmul_semantic_unsupported"; + plan.consumedOperatorIndices = QuantizedMatMulConsumedOperatorIndices(region); + plan.preservesQuantizationSemantics = preservesSemantics; + plan.isMetadataOnly = preservesSemantics; + plan.suppressesGraphOperators = false; + return plan; +} + + +QuantizedDenseLinearCublasLtCapability SelectExecutableDenseLinearCapability(QuantizedDenseLinearCublasLtCapability capability) +{ + if (capability.shapePolicy.policy == EQuantizedShapePolicy::PaddedCandidate) { + if (IsProfitableCublasLtPaddedDenseLinearPolicy(capability.shapePolicy)) { + capability.shapePolicy.policy = EQuantizedShapePolicy::Padded; + capability.shapePolicy.reason = "padded cuBLASLt int8 execution selected by profitability policy; " + capability.shapePolicy.reason; + capability.optimized = true; + auto pos = capability.tag.find("padded_candidate"); + if (pos != std::string::npos) { + capability.tag.replace(pos, std::string("padded_candidate").size(), "padded"); + } + capability.reason = "cuBLASLt optimized padded signed-int8 rank-2 dense-linear execution; " + capability.shapePolicy.reason; + } else { + capability.shapePolicy.policy = EQuantizedShapePolicy::Fallback; + capability.shapePolicy.reason = "padded cuBLASLt candidate rejected by profitability policy; " + + ExplainCublasLtPaddedDenseLinearProfitability(capability.shapePolicy) + + "; " + capability.shapePolicy.reason; + capability.optimized = false; + auto pos = capability.tag.find("padded_candidate"); + if (pos != std::string::npos) { + capability.tag.replace(pos, std::string("padded_candidate").size(), "padded_unprofitable"); + } + capability.reason = "cuBLASLt padded dense-linear execution is shape-compatible but not expected to beat the baseline; " + + capability.shapePolicy.reason; + } + } + return capability; +} + +QuantizedLoweringPlan MakeMatMulAlpakaTransposedWeightStoragePlan(const QuantizedMatMulRegion ®ion, + const std::string &weightStorageTensor, + const QuantizedDenseLinearShapePolicy &shapePolicy) +{ + QuantizedLoweringPlan plan; + plan.backend = EQuantizedBackend::ALPAKA; + if (QuantizedShapePolicyIsExecutable(shapePolicy.policy)) { + plan.status = EQuantizedLoweringStatus::Optimized; + plan.reason = shapePolicy.policy == EQuantizedShapePolicy::Padded + ? "ALPAKA cuBLASLt MatMul lowering with padded transposed pre-quantized weight storage" + : "ALPAKA cuBLASLt MatMul lowering with transposed pre-quantized weight storage"; + plan.capabilityTag = shapePolicy.policy == EQuantizedShapePolicy::Padded + ? "matmul_cublaslt_i8i8_transposed_weight_rank2_padded" + : "matmul_cublaslt_i8i8_transposed_weight_rank2_exact"; + plan.suppressesGraphOperators = true; + } else { + plan.status = EQuantizedLoweringStatus::BackendUnsupported; + plan.reason = "MatMul transposed pre-quantized weight storage is prepared but the shape is not executable by cuBLASLt int8 lowering: " + + shapePolicy.reason; + plan.capabilityTag = "matmul_transposed_weight_storage_shape_unsupported"; + plan.suppressesGraphOperators = false; + } + plan.inputStorage = QuantizedStorageTypeForCarrier(region.inputQuant); + plan.weightStorage = QuantizedStorageTypeForCarrier(region.weightQuant); + plan.biasStorage = EQuantizedStorageType::UNDEFINED; + plan.accumulatorStorage = EQuantizedStorageType::Int32Accumulator; + plan.outputStorage = QuantizedStorageTypeForCarrier(region.outputQuant); + plan.inputCarrierMode = QuantizedCarrierModeForStorage(plan.inputStorage); + plan.outputMode = EQuantizedOutputMode::Quantized; + plan.computeProfile = IsPerChannelAxis(region.weightQuant, 1) + ? EQuantizedComputeProfile::SignedInt8PerTensorActivationPerChannelWeightRank2 + : EQuantizedComputeProfile::SignedInt8SymmetricPerTensorRank2; + plan.shapePolicy = shapePolicy; + plan.weightStorageTensor = weightStorageTensor; + plan.weightLayout = EQuantizedLayout::PlainDevice; + if (IsPerChannelAxis(region.weightQuant, 1)) { + plan.weightScaleMode = EQuantizedParameterMode::PerOutputChannel; + plan.weightScaleTensor = region.weightQuant.scaleTensor; + plan.weightZeroPointTensor = region.weightQuant.zeroPointTensor; + } + plan.consumedOperatorIndices = QuantizedMatMulConsumedOperatorIndices(region); + plan.preservesQuantizationSemantics = true; + plan.isMetadataOnly = false; + return plan; +} + +QuantizedLoweringPlan MakeAvailableQuantizedGemmPlan(const QuantizedGemmRegion ®ion, + EQuantizedBackend backend, + EQuantizedLoweringStatus status, + std::string reason, + std::string capabilityTag) +{ + QuantizedLoweringPlan plan; + plan.backend = backend; + plan.status = status; + plan.reason = std::move(reason); + plan.capabilityTag = std::move(capabilityTag); + plan.consumedOperatorIndices = QuantizedGemmConsumedOperatorIndices(region); + plan.preservesQuantizationSemantics = true; + plan.isMetadataOnly = false; + plan.suppressesGraphOperators = true; + return plan; +} + +QuantizedLoweringPlan MakeCPUPackedWeightBaselinePlan(const QuantizedGemmRegion ®ion, + const std::string &weightStorageTensor) +{ + auto plan = MakeAvailableQuantizedGemmPlan(region, EQuantizedBackend::CPU, EQuantizedLoweringStatus::Baseline, + "CPU baseline lowering with packed pre-quantized weight storage", + "cpu_packed_weight_baseline"); + plan.inputStorage = QuantizedStorageTypeForCarrier(region.inputQuant); + plan.weightStorage = QuantizedStorageTypeForCarrier(region.weightQuant); + plan.biasStorage = EQuantizedStorageType::FloatCarrier; + plan.accumulatorStorage = EQuantizedStorageType::Int32Accumulator; + plan.outputStorage = EQuantizedStorageType::FloatCarrier; + plan.inputCarrierMode = EQuantizedCarrierMode::Float; + plan.outputMode = EQuantizedOutputMode::ExactFakeQuantFloat; + plan.computeProfile = EQuantizedComputeProfile::GenericRecognized; + plan.weightStorageTensor = weightStorageTensor; + plan.weightLayout = EQuantizedLayout::PackedCPU; + return plan; +} + +QuantizedLoweringPlan MakeUnsupportedQuantizedGemmPlan(EQuantizedBackend backend, std::string reason, bool preservesSemantics) +{ + QuantizedLoweringPlan plan; + plan.backend = backend; + plan.status = preservesSemantics ? EQuantizedLoweringStatus::BackendUnsupported + : EQuantizedLoweringStatus::SemanticUnsupported; + plan.reason = std::move(reason); + plan.inputStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.weightStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.biasStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.accumulatorStorage = EQuantizedStorageType::UNDEFINED; + plan.outputStorage = preservesSemantics ? EQuantizedStorageType::MetadataOnly : EQuantizedStorageType::UNDEFINED; + plan.inputCarrierMode = preservesSemantics ? EQuantizedCarrierMode::Float : EQuantizedCarrierMode::UNDEFINED; + plan.outputMode = preservesSemantics ? EQuantizedOutputMode::ExactFakeQuantFloat : EQuantizedOutputMode::UNDEFINED; + plan.computeProfile = preservesSemantics ? EQuantizedComputeProfile::GenericRecognized : EQuantizedComputeProfile::UNDEFINED; + plan.capabilityTag = preservesSemantics ? "recognized_backend_unsupported" : "semantic_unsupported"; + plan.preservesQuantizationSemantics = preservesSemantics; + plan.isMetadataOnly = preservesSemantics; + plan.suppressesGraphOperators = false; + return plan; +} + +QuantizedLoweringPlan MakeAlpakaFakeQuantPlan(const QuantizedGemmRegion ®ion) +{ + auto plan = MakeAvailableQuantizedGemmPlan(region, EQuantizedBackend::ALPAKA, EQuantizedLoweringStatus::Baseline, + "Alpaka fake-quant lowering over float carrier tensors", + "alpaka_fake_quant_baseline"); + plan.inputStorage = EQuantizedStorageType::FloatCarrier; + plan.weightStorage = EQuantizedStorageType::FloatCarrier; + plan.biasStorage = EQuantizedStorageType::FloatCarrier; + plan.accumulatorStorage = EQuantizedStorageType::Int32Accumulator; + plan.outputStorage = EQuantizedStorageType::FloatCarrier; + plan.inputCarrierMode = EQuantizedCarrierMode::Float; + plan.outputMode = EQuantizedOutputMode::ExactFakeQuantFloat; + plan.computeProfile = EQuantizedComputeProfile::GenericRecognized; + plan.weightLayout = EQuantizedLayout::Plain; + return plan; +} + + +QuantizedDenseLinearOperands MakeDenseLinearOperands(const QuantizedGemmRegion ®ion, + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape) +{ + QuantizedDenseLinearOperands operands; + operands.inputQuant = region.inputQuant; + operands.weightQuant = region.weightQuant; + operands.outputQuant = region.outputQuant; + operands.biasQuant = region.biasQuant; + operands.inputShape = inputShape; + operands.weightShape = weightShape; + operands.outputShape = outputShape; + operands.weightOutputChannelAxis = 0; + operands.operatorName = "Gemm"; + return operands; +} + +QuantizedDenseLinearOperands MakeDenseLinearOperands(const QuantizedMatMulRegion ®ion, + const std::vector &inputShape, + const std::vector &weightShape, + const std::vector &outputShape) +{ + QuantizedDenseLinearOperands operands; + operands.inputQuant = region.inputQuant; + operands.weightQuant = region.weightQuant; + operands.outputQuant = region.outputQuant; + operands.biasQuant = region.epilogue.biasQuant; + operands.inputShape = inputShape; + operands.weightShape = weightShape; + operands.outputShape = outputShape; + operands.weightOutputChannelAxis = 1; + operands.operatorName = "MatMul"; + return operands; +} + +QuantizedLoweringPlan MakeAlpakaCublasLtCorePlan(const QuantizedGemmRegion ®ion, + const std::string &weightStorageTensor, + const QuantizedDenseLinearCublasLtCapability &capability) +{ + auto plan = MakeAvailableQuantizedGemmPlan(region, EQuantizedBackend::ALPAKA, EQuantizedLoweringStatus::Optimized, + capability.reason, capability.tag); + plan.inputStorage = QuantizedStorageTypeForCarrier(region.inputQuant); + plan.weightStorage = QuantizedStorageTypeForCarrier(region.weightQuant); + plan.biasStorage = EQuantizedStorageType::FloatCarrier; + plan.accumulatorStorage = EQuantizedStorageType::Int32Accumulator; + plan.outputStorage = QuantizedStorageTypeForCarrier(region.outputQuant); + plan.inputCarrierMode = QuantizedCarrierModeForStorage(plan.inputStorage); + plan.outputMode = EQuantizedOutputMode::Quantized; + plan.computeProfile = capability.profile; + plan.shapePolicy = capability.shapePolicy; + plan.weightStorageTensor = weightStorageTensor; + plan.weightLayout = EQuantizedLayout::PlainDevice; + if (IsPerChannelAxis(region.weightQuant, 0)) { + plan.weightScaleMode = EQuantizedParameterMode::PerOutputChannel; + plan.weightScaleTensor = region.weightQuant.scaleTensor; + plan.weightZeroPointTensor = region.weightQuant.zeroPointTensor; + } + return plan; +} + + +} // namespace SOFIE diff --git a/core/src/RQuantization_Parameters.cxx b/core/src/RQuantization_Parameters.cxx new file mode 100644 index 0000000..5379efe --- /dev/null +++ b/core/src/RQuantization_Parameters.cxx @@ -0,0 +1,79 @@ +#include "SOFIE/RQuantization_Parameters.hxx" + +#include +#include + +namespace SOFIE { + +std::vector ValidateIntegralZeroPoints(const std::vector &values, + const std::string &context) +{ + if (values.empty()) + throw std::runtime_error(context + " expected a non-empty zero-point tensor"); + std::vector result; + result.reserve(values.size()); + for (float value : values) { + if (!std::isfinite(value) || std::round(static_cast(value)) != static_cast(value)) + throw std::runtime_error(context + " zero-point values must be finite integers"); + result.push_back(static_cast(std::llround(static_cast(value)))); + } + return result; +} + +int InferQuantizationParameterAxis(const std::vector &tensorShape, + std::size_t parameterCount, + std::optional explicitAxis, + const std::string &context) +{ + if (parameterCount == 1) + return -1; + if (explicitAxis) { + if (*explicitAxis < -static_cast(tensorShape.size()) || + *explicitAxis >= static_cast(tensorShape.size())) + throw std::runtime_error(context + " axis is outside the tensor rank"); + return *explicitAxis < 0 ? *explicitAxis + static_cast(tensorShape.size()) : *explicitAxis; + } + for (std::size_t i = 0; i < tensorShape.size(); ++i) { + if (tensorShape[i] == parameterCount) + return static_cast(i); + } + return -1; +} + +QuantizationInfo MakeValidatedQuantizationInfo(const QuantizationParameterSpec &spec) +{ + if (spec.scales.empty()) + throw std::runtime_error(spec.context + " expected a non-empty scale tensor"); + if (spec.zeroPoints.empty()) + throw std::runtime_error(spec.context + " expected a non-empty zero-point tensor"); + for (double scale : spec.scales) { + if (!(scale > 0.0) || !std::isfinite(scale)) + throw std::runtime_error(spec.context + " scale values must be positive and finite"); + } + if (spec.bitWidth == 0 || spec.bitWidth > 32) + throw std::runtime_error(spec.context + " bit width must be in [1, 32]"); + + const bool vectorScale = spec.scales.size() != 1; + const bool vectorZeroPoint = spec.zeroPoints.size() != 1; + if (vectorScale && vectorZeroPoint && spec.scales.size() != spec.zeroPoints.size()) + throw std::runtime_error(spec.context + " vector scale and zero-point sizes must match"); + const auto parameterCount = vectorScale ? spec.scales.size() : spec.zeroPoints.size(); + + QuantizationInfo info; + info.bitWidth = spec.bitWidth; + info.isSigned = spec.isSigned; + info.narrow = spec.narrow; + info.scale = spec.scales.front(); + info.zeroPoint = spec.zeroPoints.front(); + info.scaleTensor = spec.scaleTensor; + info.zeroPointTensor = spec.zeroPointTensor; + info.rounding = spec.rounding; + info.overflow = spec.overflow; + info.granularity = (vectorScale || vectorZeroPoint) ? EQuantizationGranularity::PerChannel + : EQuantizationGranularity::PerTensor; + info.axis = InferQuantizationParameterAxis(spec.tensorShape, parameterCount, + spec.explicitAxis, spec.context); + return info; +} + +} // namespace SOFIE diff --git a/core/src/RQuantization_Storage.cxx b/core/src/RQuantization_Storage.cxx new file mode 100644 index 0000000..e7d47a4 --- /dev/null +++ b/core/src/RQuantization_Storage.cxx @@ -0,0 +1,149 @@ +#include "SOFIE/RQuantization_Storage.hxx" +#include "SOFIE/RQuantization_DenseLinear.hxx" + +#include +#include + +namespace SOFIE { + +ETensorType TensorTypeForQuantizedStorage(EQuantizedStorageType storage) +{ + switch (storage) { + case EQuantizedStorageType::FloatCarrier: + return ETensorType::FLOAT; + case EQuantizedStorageType::Int8: + return ETensorType::INT8; + case EQuantizedStorageType::UInt8: + return ETensorType::UINT8; + case EQuantizedStorageType::Int32Accumulator: + return ETensorType::INT32; + default: + throw std::runtime_error("SOFIE quantized lowering plan has no physical tensor type for this storage"); + } +} + +QuantizedTensorStorage MakeQuantizedTensorStorage(std::string logicalTensor, + std::string sourceTensor, + std::string storageTensor, + const QuantizationInfo &quantization, + EQuantizedLayout layout, + std::vector shape, + EQuantizedBackend backend) +{ + QuantizedTensorStorage storage; + storage.logicalTensor = std::move(logicalTensor); + storage.sourceTensor = std::move(sourceTensor); + storage.storageTensor = std::move(storageTensor); + storage.storageType = QuantizedStorageTypeForCarrier(quantization); + storage.layout = layout; + storage.quantization = quantization; + storage.shape = std::move(shape); + storage.residentBackend = backend; + return storage; +} + +template +std::vector PackQuantizedGemmWeights(const float *data, std::size_t n, std::size_t k, + std::size_t tileN, const QuantizationInfo &info) +{ + const auto blocks = (n + tileN - 1) / tileN; + std::vector weights(blocks * k * tileN, 0); + for (std::size_t block = 0; block < blocks; ++block) + for (std::size_t kk = 0; kk < k; ++kk) + for (std::size_t lane = 0; lane < tileN; ++lane) { + const auto column = block * tileN + lane; + if (column < n) + weights[(block * k + kk) * tileN + lane] = + static_cast(QuantizeScalarToIntegerGrid(data[column * k + kk], info)); + } + return weights; +} + +std::vector PackQuantizedGemmWeightsInt8(const float *data, std::size_t n, std::size_t k, + std::size_t tileN, const QuantizationInfo &info) +{ + return PackQuantizedGemmWeights(data, n, k, tileN, info); +} + +std::vector PackQuantizedGemmWeightsUInt8(const float *data, std::size_t n, std::size_t k, + std::size_t tileN, const QuantizationInfo &info) +{ + return PackQuantizedGemmWeights(data, n, k, tileN, info); +} + +MaterializedQuantizedWeight MaterializeQuantizedGemmWeight( + const QuantizedGemmRegion ®ion, const QuantizedLoweringPlan &plan, + EQuantizedBackend backend, const float *sourceData, + const std::vector &sourceShape, + const std::vector &perChannelScales) +{ + if (sourceShape.size() != 2) + throw std::runtime_error("SOFIE quantized Gemm storage requires a rank-2 weight tensor"); + const auto n = sourceShape[0]; + const auto k = sourceShape[1]; + MaterializedQuantizedWeight result; + + if (backend == EQuantizedBackend::CPU) { + constexpr std::size_t tileN = 4; + const std::vector shape = {(n + tileN - 1) / tileN, k, tileN}; + result.storage = MakeQuantizedTensorStorage(region.weightTensor, region.weightSourceTensor, + plan.weightStorageTensor, region.weightQuant, + EQuantizedLayout::PackedCPU, shape, backend); + result.buffer = region.weightQuant.isSigned + ? QuantizedWeightBuffer{PackQuantizedGemmWeightsInt8(sourceData, n, k, tileN, + region.weightQuant)} + : QuantizedWeightBuffer{PackQuantizedGemmWeightsUInt8(sourceData, n, k, tileN, + region.weightQuant)}; + return result; + } + + if (backend != EQuantizedBackend::ALPAKA) + throw std::runtime_error("SOFIE quantized Gemm storage received an unsupported backend"); + + std::vector shape = sourceShape; + if (plan.shapePolicy.policy == EQuantizedShapePolicy::Padded) { + shape = {plan.shapePolicy.physicalN, plan.shapePolicy.physicalK}; + result.buffer = QuantizeGemmWeightTensorToInt8Padded(sourceData, n, k, shape[0], shape[1], + region.weightQuant, perChannelScales); + } else if (IsPerChannelAxis(region.weightQuant, 0)) { + result.buffer = QuantizeGemmWeightTensorToInt8(sourceData, n, k, region.weightQuant, perChannelScales); + } else if (region.weightQuant.isSigned) { + result.buffer = QuantizeTensorToInt8(sourceData, n * k, region.weightQuant); + } else { + result.buffer = QuantizeTensorToUInt8(sourceData, n * k, region.weightQuant); + } + result.storage = MakeQuantizedTensorStorage(region.weightTensor, region.weightSourceTensor, + plan.weightStorageTensor, region.weightQuant, + EQuantizedLayout::PlainDevice, shape, backend); + return result; +} + +MaterializedQuantizedWeight MaterializeQuantizedMatMulWeight( + const QuantizedMatMulRegion ®ion, const QuantizedLoweringPlan &plan, + EQuantizedBackend backend, const float *sourceData, + const std::vector &sourceShape, + const std::vector &perChannelScales) +{ + if (backend != EQuantizedBackend::ALPAKA) + throw std::runtime_error("SOFIE quantized MatMul storage currently requires the Alpaka backend"); + if (sourceShape.size() != 2) + throw std::runtime_error("SOFIE quantized MatMul storage requires a rank-2 weight tensor"); + const auto k = sourceShape[0]; + const auto n = sourceShape[1]; + std::vector shape = {n, k}; + MaterializedQuantizedWeight result; + if (plan.shapePolicy.policy == EQuantizedShapePolicy::Padded) { + shape = {plan.shapePolicy.physicalN, plan.shapePolicy.physicalK}; + result.buffer = QuantizeMatMulWeightTensorToInt8TransposedPadded( + sourceData, k, n, shape[1], shape[0], region.weightQuant, perChannelScales); + } else { + result.buffer = QuantizeMatMulWeightTensorToInt8Transposed( + sourceData, k, n, region.weightQuant, perChannelScales); + } + result.storage = MakeQuantizedTensorStorage(region.weightTensor, region.weightSourceTensor, + plan.weightStorageTensor, region.weightQuant, + EQuantizedLayout::PlainDevice, shape, backend); + return result; +} + +} // namespace SOFIE diff --git a/core/src/RWeightFile.cxx b/core/src/RWeightFile.cxx new file mode 100644 index 0000000..4bd5574 --- /dev/null +++ b/core/src/RWeightFile.cxx @@ -0,0 +1,64 @@ +#include "SOFIE/RWeightFile.hxx" + +#include +#include + +namespace SOFIE { + +bool IsBinaryWeightTensorType(ETensorType type) +{ + return type == ETensorType::FLOAT || type == ETensorType::INT8 || type == ETensorType::UINT8 || + type == ETensorType::INT32 || type == ETensorType::INT64; +} + +std::vector +CollectBinaryWeightTensors(const std::unordered_map &tensors) +{ + std::vector result; + result.reserve(tensors.size()); + for (const auto &[name, tensor] : tensors) { + if (tensor.IsWeightTensor()) + result.emplace_back(name, &tensor); + } + std::sort(result.begin(), result.end(), [](const auto &lhs, const auto &rhs) { return lhs.first < rhs.first; }); + return result; +} + +long WriteBinaryWeightFile(std::ostream &stream, + const std::unordered_map &tensors) +{ + auto writeValue = [&stream](const auto &value) { + stream.write(reinterpret_cast(&value), sizeof(value)); + }; + const auto weights = CollectBinaryWeightTensors(tensors); + stream.write(kBinaryWeightFileMagic, sizeof(kBinaryWeightFileMagic)); + writeValue(kBinaryWeightFileVersion); + writeValue(kBinaryWeightFileEndianMarker); + writeValue(static_cast(weights.size())); + + for (const auto &[sourceName, tensor] : weights) { + const auto type = tensor->type(); + if (!IsBinaryWeightTensorType(type)) + throw std::runtime_error("sofie tensor " + sourceName + " has a type unsupported by binary weights"); + const std::string name = "tensor_" + sourceName; + const auto nameLength = static_cast(name.size()); + const auto typeCode = static_cast(type); + const auto rank = static_cast(tensor->shape().size()); + const auto elementCount = static_cast(ConvertShapeToLength(tensor->shape())); + const auto byteCount = elementCount * GetTypeSize(type); + writeValue(nameLength); + stream.write(name.data(), name.size()); + writeValue(typeCode); + writeValue(rank); + for (auto dim : tensor->shape()) + writeValue(static_cast(dim)); + writeValue(elementCount); + writeValue(byteCount); + stream.write(static_cast(tensor->sharedptr().get()), static_cast(byteCount)); + if (!stream) + throw std::runtime_error("sofie failed to write binary tensor " + name); + } + return static_cast(stream.tellp()); +} + +} // namespace SOFIE diff --git a/parsers/CMakeLists.txt b/parsers/CMakeLists.txt index 7174e90..248bc2c 100644 --- a/parsers/CMakeLists.txt +++ b/parsers/CMakeLists.txt @@ -41,6 +41,8 @@ set(sources_cxx src/ParseConv.cxx src/ParseConvTranspose.cxx src/ParseGemm.cxx + src/ParseQONNXQuant.cxx + src/ParseONNXQuantization.cxx src/ParseGRU.cxx src/ParseIdentity.cxx src/ParseLeakyRelu.cxx diff --git a/parsers/inc/SOFIE/RModelParser_ONNX.hxx b/parsers/inc/SOFIE/RModelParser_ONNX.hxx index c1dc463..93d6fa1 100644 --- a/parsers/inc/SOFIE/RModelParser_ONNX.hxx +++ b/parsers/inc/SOFIE/RModelParser_ONNX.hxx @@ -35,17 +35,24 @@ private: std::unique_ptr fOperatorsMapImpl; // Type of the tensors std::unordered_map fTensorTypeMap; + std::unordered_map fOpsetVersionMap; // flag list of fused operators std::vector fFusedOperators; public: - // Register an ONNX operator + // Register an ONNX operator in the standard/default ONNX domain void RegisterOperator(const std::string &name, ParserFuncSignature func); - // Check if the operator is registered + // Register an ONNX operator in an explicit domain. The standard ONNX domain is "". + void RegisterOperator(const std::string &domain, const std::string &name, ParserFuncSignature func); + + // Check if the operator is registered in the standard/default ONNX domain bool IsRegisteredOperator(const std::string &name); + // Check if the operator is registered in an explicit domain + bool IsRegisteredOperator(const std::string &domain, const std::string &name); + // List of registered operators (in alphabetical order) std::vector GetRegisteredOperators(); @@ -63,6 +70,9 @@ public: // Get the type of the tensor ETensorType GetTensorType(const std::string &name); + // Get imported ONNX opset version for a domain. Returns -1 if absent. + int GetOpsetVersion(const std::string &domain) const; + // Parse the index'th node from the ONNX graph std::unique_ptr ParseOperator(const size_t /*index*/, const onnx::GraphProto & /*graphproto*/, const std::vector & /*nodes*/, const std::vector & /* children */); diff --git a/parsers/src/ParseONNXQuantization.cxx b/parsers/src/ParseONNXQuantization.cxx new file mode 100644 index 0000000..a908d48 --- /dev/null +++ b/parsers/src/ParseONNXQuantization.cxx @@ -0,0 +1,109 @@ +#include "SOFIE/RModelParser_ONNX.hxx" +#include "SOFIE/ROperator_ONNXQuantizeLinear.hxx" +#include "onnx_proto3.pb.h" + +#include +#include + +namespace SOFIE { + +namespace { + +int ParseAxisAttribute(const onnx::NodeProto &nodeproto) +{ + int axis = -1; + for (const auto &attribute : nodeproto.attribute()) { + if (attribute.name() == "axis") { + axis = static_cast(attribute.i()); + } else if (attribute.name() == "block_size") { + throw std::runtime_error("TMVA::SOFIE ONNX Q/DQ block quantization is not supported"); + } else if (attribute.name() == "output_dtype") { + throw std::runtime_error("TMVA::SOFIE ONNX Q/DQ output_dtype attribute is not supported; provide an explicit zero-point tensor"); + } else { + throw std::runtime_error("TMVA::SOFIE ONNX Q/DQ unsupported attribute " + attribute.name()); + } + } + return axis; +} + +void CheckQDQInputCount(const onnx::NodeProto &nodeproto, const std::string &opName) +{ + if (nodeproto.input_size() < 2 || nodeproto.input_size() > 3) { + throw std::runtime_error("TMVA::SOFIE ONNX " + opName + " expects two or three inputs: tensor, scale, optional zero_point"); + } + if (nodeproto.output_size() != 1) { + throw std::runtime_error("TMVA::SOFIE ONNX " + opName + " expects one output"); + } +} + +std::string OptionalInputName(const onnx::NodeProto &nodeproto, int index) +{ + if (nodeproto.input_size() <= index) + return {}; + return nodeproto.input(index); +} + +} // namespace + +ParserFuncSignature ParseQuantizeLinear = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) { + CheckQDQInputCount(nodeproto, "QuantizeLinear"); + + const std::string inputName = nodeproto.input(0); + const std::string scaleName = nodeproto.input(1); + const std::string zeroPointName = OptionalInputName(nodeproto, 2); + const std::string outputName = nodeproto.output(0); + const int axis = ParseAxisAttribute(nodeproto); + + if (!parser.IsRegisteredTensorType(inputName)) { + throw std::runtime_error("TMVA::SOFIE ONNX QuantizeLinear input tensor " + inputName + " has no registered type"); + } + if (!parser.IsRegisteredTensorType(scaleName)) { + throw std::runtime_error("TMVA::SOFIE ONNX QuantizeLinear scale tensor " + scaleName + " has no registered type"); + } + + ETensorType outputType = ETensorType::UINT8; + if (!zeroPointName.empty()) { + if (!parser.IsRegisteredTensorType(zeroPointName)) { + throw std::runtime_error("TMVA::SOFIE ONNX QuantizeLinear zero-point tensor " + zeroPointName + " has no registered type"); + } + outputType = parser.GetTensorType(zeroPointName); + } else if (parser.IsRegisteredTensorType(outputName)) { + outputType = parser.GetTensorType(outputName); + } + + if (!parser.IsRegisteredTensorType(outputName)) { + parser.RegisterTensorType(outputName, outputType); + } + + return std::make_unique(inputName, scaleName, zeroPointName, outputName, + outputType, axis); +}; + +ParserFuncSignature ParseDequantizeLinear = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) { + CheckQDQInputCount(nodeproto, "DequantizeLinear"); + + const std::string inputName = nodeproto.input(0); + const std::string scaleName = nodeproto.input(1); + const std::string zeroPointName = OptionalInputName(nodeproto, 2); + const std::string outputName = nodeproto.output(0); + const int axis = ParseAxisAttribute(nodeproto); + + if (!parser.IsRegisteredTensorType(inputName)) { + throw std::runtime_error("TMVA::SOFIE ONNX DequantizeLinear input tensor " + inputName + " has no registered type"); + } + if (!parser.IsRegisteredTensorType(scaleName)) { + throw std::runtime_error("TMVA::SOFIE ONNX DequantizeLinear scale tensor " + scaleName + " has no registered type"); + } + if (!zeroPointName.empty() && !parser.IsRegisteredTensorType(zeroPointName)) { + throw std::runtime_error("TMVA::SOFIE ONNX DequantizeLinear zero-point tensor " + zeroPointName + " has no registered type"); + } + + if (!parser.IsRegisteredTensorType(outputName)) { + parser.RegisterTensorType(outputName, ETensorType::FLOAT); + } + + return std::make_unique(inputName, scaleName, zeroPointName, outputName, + parser.GetTensorType(inputName), axis); +}; + +} // namespace SOFIE diff --git a/parsers/src/ParseQONNXQuant.cxx b/parsers/src/ParseQONNXQuant.cxx new file mode 100644 index 0000000..1fd5029 --- /dev/null +++ b/parsers/src/ParseQONNXQuant.cxx @@ -0,0 +1,88 @@ +#include "SOFIE/RModelParser_ONNX.hxx" +#include "SOFIE/ROperator_QONNXQuant.hxx" +#include "onnx_proto3.pb.h" + +#include +#include + +namespace SOFIE { + +namespace { + +EQuantizationRoundingMode ParseQONNXRoundingMode(const std::string &roundingMode) +{ + if (roundingMode == "ROUND") return EQuantizationRoundingMode::ROUND; + if (roundingMode == "FLOOR") return EQuantizationRoundingMode::FLOOR; + if (roundingMode == "TRUNCATE") return EQuantizationRoundingMode::TRUNCATE; + throw std::runtime_error("TMVA::SOFIE QONNX Quant unsupported rounding_mode " + roundingMode); +} + +} // namespace + +ParserFuncSignature ParseQONNXQuant = [](RModelParser_ONNX &parser, const onnx::NodeProto &nodeproto) { + if (nodeproto.domain() != "qonnx.custom_op.general") { + throw std::runtime_error("TMVA::SOFIE QONNX Quant parser received unexpected domain " + nodeproto.domain()); + } + if (parser.GetOpsetVersion("qonnx.custom_op.general") != 1) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant requires qonnx.custom_op.general opset version 1"); + } + if (nodeproto.input_size() != 4) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant expects four inputs: tensor, scale, zero_point, bit_width"); + } + if (nodeproto.output_size() != 1) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant expects one output"); + } + + const std::string inputName = nodeproto.input(0); + const std::string scaleName = nodeproto.input(1); + const std::string zeroPointName = nodeproto.input(2); + const std::string bitWidthName = nodeproto.input(3); + const std::string outputName = nodeproto.output(0); + + if (!parser.IsRegisteredTensorType(inputName)) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant input tensor " + inputName + " has no registered type"); + } + if (!parser.IsRegisteredTensorType(scaleName) || !parser.IsRegisteredTensorType(zeroPointName) || + !parser.IsRegisteredTensorType(bitWidthName)) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant scale, zero_point, and bit_width must be registered tensors"); + } + + bool isSigned = false; + bool narrow = false; + bool hasSigned = false; + bool hasRoundingMode = false; + EQuantizationRoundingMode rounding = EQuantizationRoundingMode::UNDEFINED; + + for (const auto &attribute : nodeproto.attribute()) { + if (attribute.name() == "signed") { + isSigned = attribute.i() != 0; + hasSigned = true; + } else if (attribute.name() == "narrow") { + narrow = attribute.i() != 0; + } else if (attribute.name() == "rounding_mode") { + rounding = ParseQONNXRoundingMode(attribute.s()); + hasRoundingMode = true; + } else { + throw std::runtime_error("TMVA::SOFIE QONNX Quant unsupported attribute " + attribute.name()); + } + } + + if (!hasSigned) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant missing required signed attribute"); + } + if (!hasRoundingMode) { + throw std::runtime_error("TMVA::SOFIE QONNX Quant missing required rounding_mode attribute"); + } + + // The initial PQuant/QONNX profile accepts saturating per-tensor affine semantics only. + const auto overflow = narrow ? EQuantizationOverflowMode::SAT_SYM : EQuantizationOverflowMode::SAT; + + if (!parser.IsRegisteredTensorType(outputName)) { + parser.RegisterTensorType(outputName, parser.GetTensorType(inputName)); + } + + return std::make_unique(inputName, scaleName, zeroPointName, bitWidthName, outputName, + isSigned, narrow, rounding, overflow); +}; + +} // namespace SOFIE diff --git a/parsers/src/RModelParser_ONNX.cxx b/parsers/src/RModelParser_ONNX.cxx index afb8b93..7ac128e 100644 --- a/parsers/src/RModelParser_ONNX.cxx +++ b/parsers/src/RModelParser_ONNX.cxx @@ -73,6 +73,9 @@ extern ParserFuncSignature ParseLeakyRelu; extern ParserFuncSignature ParseSelu; extern ParserFuncSignature ParseSigmoid; extern ParserFuncSignature ParseGemm; +extern ParserFuncSignature ParseQONNXQuant; +extern ParserFuncSignature ParseQuantizeLinear; +extern ParserFuncSignature ParseDequantizeLinear; extern ParserFuncSignature ParseRNN; extern ParserFuncSignature ParseLSTM; extern ParserFuncSignature ParsePool; @@ -117,10 +120,61 @@ extern ParserFuseFuncSignature ParseFuseBatchnormRelu; extern ParserFuseFuncSignature ParseFuseConvTransposeAdd; extern ParserFuseFuncSignature ParseFuseMatMulAdd; +namespace { + +std::string NormalizeONNXDomain(const std::string &domain) +{ + // ONNX's default operator set is normally encoded with an empty domain. + // Just treat the explicit ai.onnx spelling as the same standard domain if it appears. + return (domain.empty() || domain == "ai.onnx") ? std::string{} : domain; +} + +std::string FormatQualifiedOperatorName(const std::string &domain, const std::string &opType) +{ + const std::string normalizedDomain = NormalizeONNXDomain(domain); + return normalizedDomain.empty() ? opType : normalizedDomain + "::" + opType; +} + + +void PopulateOpsetVersionMap(const onnx::ModelProto &model, std::unordered_map &opsetVersionMap) +{ + opsetVersionMap.clear(); + for (int i = 0; i < model.opset_import_size(); ++i) { + const auto &opset = model.opset_import(i); + opsetVersionMap[NormalizeONNXDomain(opset.domain())] = static_cast(opset.version()); + } +} + +} // namespace + // Definition of RModelParser_ONNX::OperatorsMap struct RModelParser_ONNX::OperatorsMapImpl { - // Registered operators - std::unordered_map fOperatorsMap; + struct OperatorKey { + std::string fDomain; + std::string fOpType; + + bool operator==(const OperatorKey &other) const + { + return fDomain == other.fDomain && fOpType == other.fOpType; + } + }; + + struct OperatorKeyHash { + std::size_t operator()(const OperatorKey &key) const + { + const std::size_t h1 = std::hash{}(key.fDomain); + const std::size_t h2 = std::hash{}(key.fOpType); + return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2)); + } + }; + + static OperatorKey MakeKey(const std::string &domain, const std::string &opType) + { + return {NormalizeONNXDomain(domain), opType}; + } + + // Registered operators keyed by ONNX domain and operator type. + std::unordered_map fOperatorsMap; }; // helper function to get initialized tensor data @@ -150,6 +204,46 @@ struct ExtractDataFromTP { } }; template<> +struct ExtractDataFromTP { + static void Copy(onnx::TensorProto * tensor, void * data) { + auto *out = static_cast(data); + for (int i = 0; i < tensor->int32_data_size(); ++i) + out[i] = static_cast(tensor->int32_data(i)); + } +}; +template<> +struct ExtractDataFromTP { + static void Copy(onnx::TensorProto * tensor, void * data) { + auto *out = static_cast(data); + for (int i = 0; i < tensor->int32_data_size(); ++i) + out[i] = static_cast(tensor->int32_data(i)); + } +}; +template<> +struct ExtractDataFromTP { + static void Copy(onnx::TensorProto * tensor, void * data) { + auto *out = static_cast(data); + for (int i = 0; i < tensor->int32_data_size(); ++i) + out[i] = static_cast(tensor->int32_data(i)); + } +}; +template<> +struct ExtractDataFromTP { + static void Copy(onnx::TensorProto * tensor, void * data) { + auto *out = static_cast(data); + for (int i = 0; i < tensor->int32_data_size(); ++i) + out[i] = static_cast(tensor->int32_data(i)); + } +}; +template<> +struct ExtractDataFromTP { + static void Copy(onnx::TensorProto * tensor, void * data) { + auto *out = static_cast(data); + for (int i = 0; i < tensor->uint64_data_size(); ++i) + out[i] = static_cast(tensor->uint64_data(i)); + } +}; +template<> struct ExtractDataFromTP { static void Copy(onnx::TensorProto * tensor, void * data) { tensor->mutable_int64_data()->ExtractSubrange(0, tensor->int64_data_size(), @@ -241,6 +335,9 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un RegisterOperator("Conv", ParseConv); RegisterOperator("ConvTranspose", ParseConvTranspose); RegisterOperator("Gemm", ParseGemm); + RegisterOperator("QuantizeLinear", ParseQuantizeLinear); + RegisterOperator("DequantizeLinear", ParseDequantizeLinear); + RegisterOperator("qonnx.custom_op.general", "Quant", ParseQONNXQuant); RegisterOperator("GRU", ParseGRU); RegisterOperator("Identity", ParseIdentity); RegisterOperator("LeakyRelu", ParseLeakyRelu); @@ -299,12 +396,23 @@ RModelParser_ONNX::~RModelParser_ONNX() = default; void RModelParser_ONNX::RegisterOperator(const std::string &name, ParserFuncSignature func) { - fOperatorsMapImpl->fOperatorsMap[name] = func; + RegisterOperator("", name, std::move(func)); +} + +void RModelParser_ONNX::RegisterOperator(const std::string &domain, const std::string &name, ParserFuncSignature func) +{ + fOperatorsMapImpl->fOperatorsMap[RModelParser_ONNX::OperatorsMapImpl::MakeKey(domain, name)] = std::move(func); } bool RModelParser_ONNX::IsRegisteredOperator(const std::string &name) { - return fOperatorsMapImpl->fOperatorsMap.find(name) != fOperatorsMapImpl->fOperatorsMap.end(); + return IsRegisteredOperator("", name); +} + +bool RModelParser_ONNX::IsRegisteredOperator(const std::string &domain, const std::string &name) +{ + return fOperatorsMapImpl->fOperatorsMap.find(RModelParser_ONNX::OperatorsMapImpl::MakeKey(domain, name)) != + fOperatorsMapImpl->fOperatorsMap.end(); } std::vector RModelParser_ONNX::GetRegisteredOperators() @@ -312,7 +420,7 @@ std::vector RModelParser_ONNX::GetRegisteredOperators() std::vector ops; ops.reserve(fOperatorsMapImpl->fOperatorsMap.size()); for (auto &it : fOperatorsMapImpl->fOperatorsMap) { - ops.emplace_back(it.first); + ops.emplace_back(FormatQualifiedOperatorName(it.first.fDomain, it.first.fOpType)); } // return sorted list in alphabetical order std::sort(ops.begin(), ops.end()); @@ -334,6 +442,12 @@ ETensorType RModelParser_ONNX::GetTensorType(const std::string &name) return fTensorTypeMap[UTILITY::Clean_name(name)]; } +int RModelParser_ONNX::GetOpsetVersion(const std::string &domain) const +{ + auto it = fOpsetVersionMap.find(NormalizeONNXDomain(domain)); + return it == fOpsetVersionMap.end() ? -1 : it->second; +} + // Parse an operator std::unique_ptr RModelParser_ONNX::ParseOperator(const size_t i, const onnx::GraphProto &graphproto, const std::vector &nodes, const std::vector & children) @@ -343,14 +457,16 @@ RModelParser_ONNX::ParseOperator(const size_t i, const onnx::GraphProto &graphpr int idx = nodes[i]; const auto &nodeproto = graphproto.node(idx); const std::string op_type = nodeproto.op_type(); + const std::string op_domain = NormalizeONNXDomain(nodeproto.domain()); if (fVerbose) - std::cout << "Parsing operator " << op_type << std::endl; + std::cout << "Parsing operator " << FormatQualifiedOperatorName(op_domain, op_type) << std::endl; // skip already fused operators if (fFusedOperators[idx]) return nullptr; - // try to fuse with following operator in case it is not last one - if (children.size() == 1) { + // try to fuse with following operator in case it is not last one. + // Parser fusions are defined only for standard ONNX operators. + if (op_domain.empty() && children.size() == 1) { int idx2 = children.front(); if (op_type == "MatMul") { // Fuse MatMul and Add @@ -388,13 +504,15 @@ RModelParser_ONNX::ParseOperator(const size_t i, const onnx::GraphProto &graphpr - auto it = fOperatorsMapImpl->fOperatorsMap.find(op_type); + const auto op_key = RModelParser_ONNX::OperatorsMapImpl::MakeKey(op_domain, op_type); + auto it = fOperatorsMapImpl->fOperatorsMap.find(op_key); if (it == fOperatorsMapImpl->fOperatorsMap.end()) { - std::cout << "operator " << op_type << " is not supported" << std::endl; - throw std::runtime_error("TMVA::SOFIE Operator type " + op_type + " is not yet supported"); + const std::string qualifiedName = FormatQualifiedOperatorName(op_domain, op_type); + std::cout << "operator " << qualifiedName << " is not supported" << std::endl; + throw std::runtime_error("TMVA::SOFIE Operator type " + qualifiedName + " is not yet supported"); } if (fVerbose) { - std::cout << "\tCreating operator " << op_type << std::endl; + std::cout << "\tCreating operator " << FormatQualifiedOperatorName(op_domain, op_type) << std::endl; } return it->second(*this, nodeproto); } @@ -405,11 +523,14 @@ RModel RModelParser_ONNX::Parse(std::string filename, bool verbose) fVerbose = verbose; fTensorTypeMap.clear(); + fOpsetVersionMap.clear(); auto model = LoadModel(filename); if (!model) throw std::runtime_error("TMVA::SOFIE - Failed to load onnx file " + filename); + PopulateOpsetVersionMap(*model, fOpsetVersionMap); + const onnx::GraphProto &graph = model->graph(); // not a memory leak. model freed automatically at the end. @@ -459,8 +580,10 @@ void RModelParser_ONNX::CheckGraph(const onnx::GraphProto & graph, int & level, for (int i = 0; i < graph.node_size(); i++) { const auto & node = graph.node(i); const std::string opType = node.op_type(); + const std::string opDomain = NormalizeONNXDomain(node.domain()); + const std::string qualifiedName = FormatQualifiedOperatorName(opDomain, opType); if (fVerbose) { - std::cout << "\tOperator " << i << " : " << opType << " (" << node.name() << "), " << graph.node(i).input_size() + std::cout << "\tOperator " << i << " : " << qualifiedName << " (" << node.name() << "), " << graph.node(i).input_size() << " inputs : {"; for (int j = 0; j < graph.node(i).input_size(); j++) { std::cout << graph.node(i).input(j); @@ -470,8 +593,8 @@ void RModelParser_ONNX::CheckGraph(const onnx::GraphProto & graph, int & level, std::cout << " }" << std::endl; } // check if operator exists - if (!IsRegisteredOperator(opType)) - missingOperators[opType] = level; + if (!IsRegisteredOperator(opDomain, opType)) + missingOperators[qualifiedName] = level; // see if sub-graph exists as node attributes for (int j = 0; j < node.attribute_size(); j++) { const auto & attribute = node.attribute(j); @@ -487,9 +610,12 @@ void RModelParser_ONNX::CheckGraph(const onnx::GraphProto & graph, int & level, bool RModelParser_ONNX::CheckModel(std::string filename, bool verbose) { fVerbose = verbose; + fOpsetVersionMap.clear(); auto model = LoadModel(filename); if (!model) return false; + PopulateOpsetVersionMap(*model, fOpsetVersionMap); + const onnx::GraphProto &graph = model->graph(); // Initial operator order if (fVerbose) @@ -631,6 +757,34 @@ void RModelParser_ONNX::ParseONNXGraph(RModel & rmodel, const onnx::GraphProto & allInitializedTensors[input_name] = i; break; } + case ETensorType::INT8: { + std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); + if (verbose) std::cout << "add INT8 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; + rmodel.AddInitializedTensor(input_name, ETensorType::INT8, shape, data); + allInitializedTensors[input_name] = i; + break; + } + case ETensorType::UINT8: { + std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); + if (verbose) std::cout << "add UINT8 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; + rmodel.AddInitializedTensor(input_name, ETensorType::UINT8, shape, data); + allInitializedTensors[input_name] = i; + break; + } + case ETensorType::INT16: { + std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); + if (verbose) std::cout << "add INT16 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; + rmodel.AddInitializedTensor(input_name, ETensorType::INT16, shape, data); + allInitializedTensors[input_name] = i; + break; + } + case ETensorType::UINT16: { + std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); + if (verbose) std::cout << "add UINT16 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; + rmodel.AddInitializedTensor(input_name, ETensorType::UINT16, shape, data); + allInitializedTensors[input_name] = i; + break; + } case ETensorType::INT32: { std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); if (verbose) std::cout << "add INT32 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; @@ -638,6 +792,13 @@ void RModelParser_ONNX::ParseONNXGraph(RModel & rmodel, const onnx::GraphProto & allInitializedTensors[input_name] = i; break; } + case ETensorType::UINT32: { + std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); + if (verbose) std::cout << "add UINT32 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; + rmodel.AddInitializedTensor(input_name, ETensorType::UINT32, shape, data); + allInitializedTensors[input_name] = i; + break; + } case ETensorType::INT64: { std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); if (verbose) std::cout << "add INT64 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 12f19b1..56d24b6 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -103,6 +103,14 @@ if (ENABLE_ALPAKA_TESTS) FIXTURES_SETUP sofie-compile-models-onnx-alpaka ) + ROOTTEST_ADD_TEST( + SofieQuantizedStorageExternalized_ONNX_Alpaka + COMMAND ${CMAKE_COMMAND} + -DGENERATED_DIR=${CMAKE_CURRENT_BINARY_DIR} + -P ${CMAKE_CURRENT_SOURCE_DIR}/CheckQuantizedStorageExternalized.cmake + FIXTURES_REQUIRED sofie-compile-models-onnx-alpaka + ) + ########################################################################## # CUDA backend ########################################################################## @@ -148,6 +156,7 @@ if (ENABLE_ALPAKA_TESTS) TestCustomModelsFromONNXForAlpakaCuda PRIVATE ALPAKA_ACC_GPU_CUDA_ENABLED ALPAKA_HAS_STD_ATOMIC_REF + SOFIE_USE_CUBLASLT ) target_compile_options( diff --git a/test/TestCustomModelsFromONNXForAlpakaCuda.cxx b/test/TestCustomModelsFromONNXForAlpakaCuda.cxx index fccacbe..30ae974 100644 --- a/test/TestCustomModelsFromONNXForAlpakaCuda.cxx +++ b/test/TestCustomModelsFromONNXForAlpakaCuda.cxx @@ -1,5 +1,12 @@ +#include +#include #include #include +#include +#include +#include + +#include "SOFIE/RWeightFile.hxx" // ── Trilu ────────────────────────────────────────────────────────────────── #include "Trilu_upper_FromONNX_GPU_ALPAKA.hxx" @@ -44,6 +51,16 @@ #include "Linear_64_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/Linear_64.ref.hxx" +#include "QONNX_QuantGemm_Binary_FromONNX_GPU_ALPAKA.hxx" +#include "QONNX_QuantGemm_NoBias_FromONNX_GPU_ALPAKA.hxx" +#include "QONNX_QuantMatMul_FromONNX_GPU_ALPAKA.hxx" +#include "QONNX_QuantMatMul_Padded_FromONNX_GPU_ALPAKA.hxx" +#include "QONNX_QuantMatMul_Add_FromONNX_GPU_ALPAKA.hxx" +#include "ONNX_QDQ_QuantGemm_FromONNX_GPU_ALPAKA.hxx" +#include "ONNX_QDQ_QuantMatMul_FromONNX_GPU_ALPAKA.hxx" +#include "ONNX_QDQ_QuantGemm_PerChannelWeight_FromONNX_GPU_ALPAKA.hxx" +#include "ONNX_QDQ_QuantMatMul_PerChannelWeight_FromONNX_GPU_ALPAKA.hxx" + #include "AddBroadcast1_FromONNX_GPU_ALPAKA.hxx" #include "input_models/references/AddBroadcast1.ref.hxx" @@ -189,6 +206,57 @@ using Idx = std::size_t; using Dim = alpaka::DimInt<1>; using Ext1D = alpaka::Vec; +struct QuantizedLinearTest { + Idx m; + Idx k; + Idx n; + bool matMul; + bool hasBias; + bool perChannelWeight; +}; + +std::int8_t QuantizedLinearTestInputValue(Idx index) +{ + return static_cast(((index * 5 + index / 7) % 31) - 15); +} + +std::int8_t QuantizedLinearTestWeightValue(Idx index, Idx n) +{ + return static_cast(((index * 3 + index / 11 + n) % 29) - 14); +} + +std::vector MakeQuantizedLinearTestInput(const QuantizedLinearTest &test) +{ + std::vector input(test.m * test.k); + for (Idx i = 0; i < input.size(); ++i) + input[i] = QuantizedLinearTestInputValue(i); + return input; +} + +std::vector MakeQuantizedLinearTestExpected(const QuantizedLinearTest &test, + const std::vector &input) +{ + std::vector output(test.m * test.n); + constexpr int scaleNumerators[] = {3, 4, 5, 6}; + for (Idx row = 0; row < test.m; ++row) { + for (Idx column = 0; column < test.n; ++column) { + std::int32_t accumulator = 0; + for (Idx inner = 0; inner < test.k; ++inner) { + const Idx weightIndex = test.matMul ? inner * test.n + column : column * test.k + inner; + accumulator += static_cast(input[row * test.k + inner]) * + static_cast(QuantizedLinearTestWeightValue(weightIndex, test.n)); + } + if (test.hasBias) + accumulator += static_cast((column * 3) % 17) - 8; + const int scaleNumerator = test.perChannelWeight ? scaleNumerators[column % 4] : 4; + const auto quantized = static_cast(std::nearbyint( + static_cast(accumulator) * static_cast(scaleNumerator) / 128.0)); + output[row * test.n + column] = static_cast(std::clamp(quantized, -128L, 127L)); + } + } + return output; +} + class SofieAlpakaTest : public ::testing::Test { protected: // Shared devices and platforms @@ -198,6 +266,41 @@ class SofieAlpakaTest : public ::testing::Test { alpaka::DevCudaRt device; alpaka::Queue queue; + template + void RunQuantizedLinearInt8(const char *weightFile, const QuantizedLinearTest &test) + { + const auto input = MakeQuantizedLinearTestInput(test); + const auto expectedOutput = MakeQuantizedLinearTestExpected(test, input); + const Idx inputSize = input.size(); + const Idx outputSize = expectedOutput.size(); + auto input_h = alpaka::allocBuf(host, Ext1D::all(inputSize)); + std::int8_t *input_ptr = reinterpret_cast(alpaka::getPtrNative(input_h)); + for (Idx i = 0; i < inputSize; ++i) { + input_ptr[i] = input[i]; + } + + auto input_d = alpaka::allocBuf(device, Ext1D::all(inputSize)); + alpaka::memcpy(queue, input_d, input_h); + alpaka::wait(queue); + + auto result_h = alpaka::allocBuf(host, Ext1D::all(outputSize)); + + { + TModel model(weightFile); + auto result = model.infer(input_d); + alpaka::wait(queue); + cudaDeviceSynchronize(); + + alpaka::memcpy(queue, result_h, result); + alpaka::wait(queue); + } + + const auto *res_ptr = reinterpret_cast(alpaka::getPtrNative(result_h)); + for (Idx i = 0; i < outputSize; ++i) { + EXPECT_EQ(static_cast(res_ptr[i]), static_cast(expectedOutput[i])) << "i=" << i; + } + } + SofieAlpakaTest() : hostPlatform{} , host(alpaka::getDevByIdx(hostPlatform, 0u)) @@ -222,6 +325,72 @@ class SofieAlpakaTest : public ::testing::Test { }; +TEST_F(SofieAlpakaTest, QuantizedGemm) +{ + // Biased Gemm: Yq = QY(SX * SW_j * sum_k Xq_ik * Wq_jk + SX * SW_j * Bq_j). + RunQuantizedLinearInt8>( + "QONNX_QuantGemm_Binary_FromONNX_GPU_ALPAKA.bin", + QuantizedLinearTest{512, 64, 32, false, true, true}); +} + +TEST(SofieBinaryWeights, RejectsInvalidHeader) +{ + std::string bytes(24, 0); + std::stringstream stream(bytes, std::ios::in | std::ios::binary); + EXPECT_THROW(SOFIE::ReadBinaryWeightFileHeader(stream, 0), std::runtime_error); +} + +TEST_F(SofieAlpakaTest, QuantizedGemmFrontendsEquivalent) +{ + // QONNX Quant and standard Q/DQ encode the same no-bias Gemm semantics. + RunQuantizedLinearInt8>( + "QONNX_QuantGemm_NoBias_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, false, false, false}); + RunQuantizedLinearInt8>( + "ONNX_QDQ_QuantGemm_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, false, false, false}); +} + +TEST_F(SofieAlpakaTest, QuantizedMatMulFrontends) +{ + // MatMul uses W as [K,N]: Yq = QY(SX * SW_j * sum_k Xq_ik * Wq_kj). + RunQuantizedLinearInt8>( + "QONNX_QuantMatMul_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{512, 64, 32, true, false, true}); + + // Standard Q/DQ example currently uses the smaller shared 256x64 input. + RunQuantizedLinearInt8>( + "ONNX_QDQ_QuantMatMul_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, true, false, false}); +} + +TEST_F(SofieAlpakaTest, QuantizedMatMulPadded) +{ + // Padded MatMul has logical M=511; the backend may pad physical storage but returns logical Y. + RunQuantizedLinearInt8>( + "QONNX_QuantMatMul_Padded_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{511, 64, 80, true, false, true}); +} + +TEST_F(SofieAlpakaTest, QuantizedMatMulAdd) +{ + // Projection bias: Yq = QY(SX * SW_j * sum_k Xq_ik * Wq_kj + bias_j). + RunQuantizedLinearInt8>( + "QONNX_QuantMatMul_Add_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, true, true, false}); +} + +TEST_F(SofieAlpakaTest, StandardQDQPerChannelDenseLinear) +{ + // Gemm uses output-channel axis 0; MatMul uses output-channel axis 1. + RunQuantizedLinearInt8>( + "ONNX_QDQ_QuantGemm_PerChannelWeight_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, false, false, true}); + RunQuantizedLinearInt8>( + "ONNX_QDQ_QuantMatMul_PerChannelWeight_FromONNX_GPU_ALPAKA.dat", + QuantizedLinearTest{256, 64, 64, true, false, true}); +} + TEST_F(SofieAlpakaTest, Linear64) { constexpr float TOLERANCE = DEFAULT_TOLERANCE; diff --git a/test/input_models/ONNX_QDQ_QuantGemm.onnx b/test/input_models/ONNX_QDQ_QuantGemm.onnx new file mode 100644 index 0000000..41aef07 Binary files /dev/null and b/test/input_models/ONNX_QDQ_QuantGemm.onnx differ diff --git a/test/input_models/ONNX_QDQ_QuantGemm_PerChannelWeight.onnx b/test/input_models/ONNX_QDQ_QuantGemm_PerChannelWeight.onnx new file mode 100644 index 0000000..5ee5475 Binary files /dev/null and b/test/input_models/ONNX_QDQ_QuantGemm_PerChannelWeight.onnx differ diff --git a/test/input_models/ONNX_QDQ_QuantMatMul.onnx b/test/input_models/ONNX_QDQ_QuantMatMul.onnx new file mode 100644 index 0000000..56cde93 Binary files /dev/null and b/test/input_models/ONNX_QDQ_QuantMatMul.onnx differ diff --git a/test/input_models/ONNX_QDQ_QuantMatMul_PerChannelWeight.onnx b/test/input_models/ONNX_QDQ_QuantMatMul_PerChannelWeight.onnx new file mode 100644 index 0000000..4578420 Binary files /dev/null and b/test/input_models/ONNX_QDQ_QuantMatMul_PerChannelWeight.onnx differ diff --git a/test/input_models/QONNX_QuantGemm.onnx b/test/input_models/QONNX_QuantGemm.onnx new file mode 100644 index 0000000..f7ef7e8 Binary files /dev/null and b/test/input_models/QONNX_QuantGemm.onnx differ diff --git a/test/input_models/QONNX_QuantGemm_NoBias.onnx b/test/input_models/QONNX_QuantGemm_NoBias.onnx new file mode 100644 index 0000000..3c9d857 Binary files /dev/null and b/test/input_models/QONNX_QuantGemm_NoBias.onnx differ diff --git a/test/input_models/QONNX_QuantMatMul.onnx b/test/input_models/QONNX_QuantMatMul.onnx new file mode 100644 index 0000000..8d4930c Binary files /dev/null and b/test/input_models/QONNX_QuantMatMul.onnx differ diff --git a/test/input_models/QONNX_QuantMatMul_Add.onnx b/test/input_models/QONNX_QuantMatMul_Add.onnx new file mode 100644 index 0000000..371eb41 Binary files /dev/null and b/test/input_models/QONNX_QuantMatMul_Add.onnx differ diff --git a/test/input_models/QONNX_QuantMatMul_Padded.onnx b/test/input_models/QONNX_QuantMatMul_Padded.onnx new file mode 100644 index 0000000..a4d0aad Binary files /dev/null and b/test/input_models/QONNX_QuantMatMul_Padded.onnx differ