File size: 6,396 Bytes
5f70891 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | // Experimental Paddle CPU operators for the fixed TrOCR encoder profile.
// Welford/reduction order follows pinned PyTorch; exp follows pinned SLEEF.
// Retain trocr_paddle_norm_probe.LICENSE and trocr_paddle_sleef_probe.LICENSE.
#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);
}
} // namespace
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));
|