| |
| |
| |
| #include <array> |
| #include <cfenv> |
| #include <cmath> |
| #include <cstdint> |
| #include <vector> |
| #include "paddle/extension.h" |
|
|
| namespace { |
| constexpr int kWidth = 384; |
| constexpr int kAttention = 578; |
| struct Moments { |
| int count = 0; |
| std::array<float, 8> mean{}; |
| std::array<float, 8> m2{}; |
| }; |
| void Merge(Moments& dst, const Moments& src) { |
| const int old = dst.count; |
| const float factor = float(src.count) / float(old + src.count); |
| for (int i = 0; i < 8; ++i) { |
| const float delta = src.mean[i] - dst.mean[i]; |
| const float scaled = factor * delta; |
| dst.mean[i] = dst.mean[i] + scaled; |
| dst.m2[i] = std::fma(delta * float(old), scaled, dst.m2[i] + src.m2[i]); |
| } |
| dst.count += src.count; |
| } |
| void CheckCPU(const paddle::Tensor& value) { |
| PD_CHECK(value.is_cpu(), "Source-order research operators require CPU inputs."); |
| PD_CHECK(value.dtype() == paddle::DataType::FLOAT32, "Require FP32 inputs."); |
| PD_CHECK(std::fegetround() == FE_TONEAREST, "Require round-to-nearest floating point."); |
| } |
| float NegativeExp(float d) { |
| if (d < -104.f) return 0.f; |
| const int q = int(std::nearbyint(d * 1.4426950408889634073599246810018921374f)); |
| float s = std::fma(float(q), -0.693145751953125f, d); |
| s = std::fma(float(q), -1.428606765330187045e-6f, s); |
| float u = 0.000198527617612853646278381f; |
| u = std::fma(u, s, 0.00139304355252534151077271f); |
| u = std::fma(u, s, 0.00833336077630519866943359f); |
| u = std::fma(u, s, 0.0416664853692054748535156f); |
| u = std::fma(u, s, 0.166666671633720397949219f); |
| u = std::fma(u, s, 0.5f); |
| u = 1.f + std::fma(s * s, u, s); |
| const int half = int(std::floor(float(q) * .5f)); |
| return std::ldexp(std::ldexp(u, half), q - half); |
| } |
| } |
|
|
| std::vector<paddle::Tensor> SourceNorm384(const paddle::Tensor& x, |
| const paddle::Tensor& gamma, |
| const paddle::Tensor& beta) { |
| CheckCPU(x); CheckCPU(gamma); CheckCPU(beta); |
| PD_CHECK(!x.shape().empty() && x.shape().back() == kWidth, "Require 384 channels."); |
| PD_CHECK(gamma.numel() == kWidth && beta.numel() == kWidth, "Require 384 affine values."); |
| auto out = paddle::empty(x.shape(), x.dtype(), x.place()); |
| const float* input = x.data<float>(); |
| const float* weight = gamma.data<float>(); |
| const float* bias = beta.data<float>(); |
| float* result = out.data<float>(); |
| for (int64_t row = 0; row < x.numel() / kWidth; ++row) { |
| const float* values = input + row * kWidth; |
| Moments first, second; |
| for (int chunk = 0; chunk < 3; ++chunk) { |
| Moments local; local.count = 16; |
| for (int j = 0; j < 16; ++j) { |
| const float factor = 1.f / float(j + 1); |
| for (int lane = 0; lane < 8; ++lane) { |
| const float value = values[chunk * 128 + j * 8 + lane]; |
| const float delta = value - local.mean[lane]; |
| local.mean[lane] = std::fma(factor, delta, local.mean[lane]); |
| local.m2[lane] = std::fma(delta, value - local.mean[lane], local.m2[lane]); |
| } |
| } |
| Merge(first, local); |
| if (chunk == 1) { Merge(second, first); first = Moments{}; } |
| } |
| Merge(first, second); |
| float mean = 0.f, m2 = 0.f; |
| for (int lane = 0; lane < 8; ++lane) { |
| const float factor = 48.f / float((lane + 1) * 48); |
| const float delta = first.mean[lane] - mean; |
| mean = std::fma(factor, delta, mean); |
| const float correction = (delta * delta) * factor; |
| m2 = m2 + std::fma(correction, float(lane * 48), first.m2[lane]); |
| } |
| const float variance = m2 / float(kWidth); |
| const float scale = 1.f / std::sqrt(variance + 1e-12f); |
| for (int i = 0; i < kWidth; ++i) { |
| const float normalized = (values[i] - mean) * scale; |
| result[row * kWidth + i] = std::fma(normalized, weight[i], bias[i]); |
| } |
| } |
| return {out}; |
| } |
|
|
| std::vector<paddle::Tensor> SourceSoftmax578(const paddle::Tensor& x) { |
| CheckCPU(x); |
| PD_CHECK(!x.shape().empty() && x.shape().back() == kAttention, "Require 578 attention scores."); |
| auto out = paddle::empty(x.shape(), x.dtype(), x.place()); |
| const float* input = x.data<float>(); float* output = out.data<float>(); |
| for (int64_t row = 0; row < x.numel() / kAttention; ++row) { |
| const float* values = input + row * kAttention; |
| float* result = output + row * kAttention; |
| float maximum = values[0]; |
| for (int i = 0; i < kAttention; ++i) { |
| PD_CHECK(std::isfinite(values[i]), "Require finite scores."); |
| if (values[i] > maximum) maximum = values[i]; |
| } |
| for (int i = 0; i < kAttention; ++i) result[i] = NegativeExp(values[i] - maximum); |
| std::array<float, 16> sum; |
| for (int i = 0; i < 16; ++i) sum[i] = result[i]; |
| for (int block = 1; block < 36; ++block) |
| for (int lane = 0; lane < 16; ++lane) sum[lane] = sum[lane] + result[block * 16 + lane]; |
| for (int lane = 0; lane < 2; ++lane) sum[lane] = sum[lane] + result[576 + lane]; |
| for (int width : {8, 4, 2, 1}) |
| for (int lane = 0; lane < width; ++lane) sum[lane] = sum[lane] + sum[lane + width]; |
| const float inverse = 1.f / sum[0]; |
| for (int i = 0; i < kAttention; ++i) result[i] = result[i] * inverse; |
| } |
| return {out}; |
| } |
|
|
| std::vector<std::vector<int64_t>> NormShape(const std::vector<int64_t>& x, |
| const std::vector<int64_t>& gamma, |
| const std::vector<int64_t>& beta) { return {x}; } |
| std::vector<paddle::DataType> NormType(paddle::DataType x, paddle::DataType gamma, |
| paddle::DataType beta) { return {x}; } |
| std::vector<std::vector<int64_t>> SoftmaxShape(const std::vector<int64_t>& x) { return {x}; } |
| std::vector<paddle::DataType> SoftmaxType(paddle::DataType x) { return {x}; } |
|
|
| PD_BUILD_OP(ocr_source_norm384_v1) |
| .Inputs({"X", "Gamma", "Beta"}).Outputs({"Out"}) |
| .SetKernelFn(PD_KERNEL(SourceNorm384)) |
| .SetInferShapeFn(PD_INFER_SHAPE(NormShape)).SetInferDtypeFn(PD_INFER_DTYPE(NormType)); |
| PD_BUILD_OP(ocr_source_softmax578_v1) |
| .Inputs({"X"}).Outputs({"Out"}) |
| .SetKernelFn(PD_KERNEL(SourceSoftmax578)) |
| .SetInferShapeFn(PD_INFER_SHAPE(SoftmaxShape)).SetInferDtypeFn(PD_INFER_DTYPE(SoftmaxType)); |
|
|
