ggml : add bilinear upscale support (ggml/1185)

This commit is contained in:
Diego Devesa
2025-04-09 12:32:13 +02:00
committed by Georgi Gerganov
parent 459895c326
commit fe92821ea9
9 changed files with 119 additions and 40 deletions

View File

@ -6351,24 +6351,72 @@ static void ggml_compute_forward_upscale_f32(
const float sf2 = (float)ne2/src0->ne[2];
const float sf3 = (float)ne3/src0->ne[3];
// TODO: optimize
const ggml_scale_mode mode = (ggml_scale_mode) ggml_get_op_params_i32(dst, 0);
for (int64_t i3 = 0; i3 < ne3; i3++) {
const int64_t i03 = i3 / sf3;
for (int64_t i2 = ith; i2 < ne2; i2 += nth) {
const int64_t i02 = i2 / sf2;
for (int64_t i1 = 0; i1 < ne1; i1++) {
const int64_t i01 = i1 / sf1;
for (int64_t i0 = 0; i0 < ne0; i0++) {
const int64_t i00 = i0 / sf0;
if (mode == GGML_SCALE_MODE_NEAREST) {
for (int64_t i3 = 0; i3 < ne3; i3++) {
const int64_t i03 = i3 / sf3;
for (int64_t i2 = ith; i2 < ne2; i2 += nth) {
const int64_t i02 = i2 / sf2;
for (int64_t i1 = 0; i1 < ne1; i1++) {
const int64_t i01 = i1 / sf1;
for (int64_t i0 = 0; i0 < ne0; i0++) {
const int64_t i00 = i0 / sf0;
const float * x = (float *)((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03);
float * y = (float *)((char *) dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
const float * x = (float *)((char *) src0->data + i00*nb00 + i01*nb01 + i02*nb02 + i03*nb03);
float * y = (float *)((char *) dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
*y = *x;
*y = *x;
}
}
}
}
} else if (mode == GGML_SCALE_MODE_BILINEAR) {
// setting a pixel offset of 0 would replicate the behavior of pytorch interpolate with align_corners=True
const float pixel_offset = 0.5f;
for (int64_t i3 = 0; i3 < ne3; i3++) {
const int64_t i03 = i3 / sf3;
for (int64_t i2 = ith; i2 < ne2; i2 += nth) {
const int64_t i02 = i2 / sf2;
for (int64_t i1 = 0; i1 < ne1; i1++) {
const float y = ((float)i1 + pixel_offset) / sf1 - pixel_offset;
int64_t y0 = (int64_t)floorf(y);
int64_t y1 = y0 + 1;
y0 = std::max(int64_t(0), std::min(y0, ne01 - 1));
y1 = std::max(int64_t(0), std::min(y1, ne01 - 1));
float dy = y - (float)y0;
dy = std::max(0.0f, std::min(dy, 1.0f));
for (int64_t i0 = 0; i0 < ne0; i0++) {
const float x = ((float)i0 + pixel_offset) / sf0 - pixel_offset;
int64_t x0 = (int64_t)floorf(x);
int64_t x1 = x0 + 1;
x0 = std::max(int64_t(0), std::min(x0, ne00 - 1));
x1 = std::max(int64_t(0), std::min(x1, ne00 - 1));
float dx = x - (float)x0;
dx = std::max(0.0f, std::min(dx, 1.0f));
// fetch the four surrounding pixel values and interpolate
const float a = *(const float *)((const char *)src0->data + x0*nb00 + y0*nb01 + i02*nb02 + i03*nb03);
const float b = *(const float *)((const char *)src0->data + x1*nb00 + y0*nb01 + i02*nb02 + i03*nb03);
const float c = *(const float *)((const char *)src0->data + x0*nb00 + y1*nb01 + i02*nb02 + i03*nb03);
const float d = *(const float *)((const char *)src0->data + x1*nb00 + y1*nb01 + i02*nb02 + i03*nb03);
const float val = a*(1 - dx)*(1 - dy) + b*dx*(1 - dy) + c*(1 - dx)*dy + d*dx*dy;
float * y_dst = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3);
*y_dst = val;
}
}
}
}
} else {
GGML_ABORT("unsupported upscale mode");
}
}