Skip to content

feat: support strided bilinear upsampling #1531

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
PanZezhong1725 wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat: support strided bilinear upsampling #1531

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 20 additions & 11 deletions src/infiniop/ops/upsample_bilinear/cpu/upsample_bilinear_cpu.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#include "../../../devices/cpu/common_cpu.h"
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <omp.h>
#include <vector>

Expand Down Expand Up @@ -104,6 +105,14 @@ void calculate_cpu_impl(
size_t out_h = info.h_out();
size_t out_w = info.w_out();
bool align_corners = info.align_corners();
const ptrdiff_t input_stride_n = info.input_stride(0);
const ptrdiff_t input_stride_c = info.input_stride(1);
const ptrdiff_t input_stride_h = info.input_stride(2);
const ptrdiff_t input_stride_w = info.input_stride(3);
const ptrdiff_t output_stride_n = info.output_stride(0);
const ptrdiff_t output_stride_c = info.output_stride(1);
const ptrdiff_t output_stride_h = info.output_stride(2);
const ptrdiff_t output_stride_w = info.output_stride(3);

auto out_ptr = reinterpret_cast<T *>(output);
auto in_ptr = reinterpret_cast<const T *>(input);
Expand All @@ -116,32 +125,32 @@ void calculate_cpu_impl(

#pragma omp parallel for schedule(static)
for (ptrdiff_t nc = 0; nc < (ptrdiff_t)n_c; ++nc) {
// 当前 channel 的输入输出起始指针
const T *src_base = in_ptr + nc * in_h * in_w;
T *dst_base = out_ptr + nc * out_h * out_w;
const size_t n = static_cast<size_t>(nc) / C;
const size_t c = static_cast<size_t>(nc) % C;
const ptrdiff_t input_base = n * input_stride_n + c * input_stride_c;
const ptrdiff_t output_base = n * output_stride_n + c * output_stride_c;

for (size_t h = 0; h < out_h; ++h) {
const auto &hp = h_params[h];
// 缓存行指针，避免内层循环重复计算乘法
const T *src_row0 = src_base + hp.idx0 * in_w;
const T *src_row1 = src_base + hp.idx1 * in_w;
const ptrdiff_t row0 = input_base + hp.idx0 * input_stride_h;
const ptrdiff_t row1 = input_base + hp.idx1 * input_stride_h;

for (size_t w = 0; w < out_w; ++w) {
const auto &wp = w_params[w];

// 获取四个采样点的值
float val00 = utils::cast<float>(src_row0[wp.idx0]);
float val01 = utils::cast<float>(src_row0[wp.idx1]);
float val10 = utils::cast<float>(src_row1[wp.idx0]);
float val11 = utils::cast<float>(src_row1[wp.idx1]);
float val00 = utils::cast<float>(in_ptr[row0 + wp.idx0 * input_stride_w]);
float val01 = utils::cast<float>(in_ptr[row0 + wp.idx1 * input_stride_w]);
float val10 = utils::cast<float>(in_ptr[row1 + wp.idx0 * input_stride_w]);
float val11 = utils::cast<float>(in_ptr[row1 + wp.idx1 * input_stride_w]);

// 双线性插值计算
// interpolation = (val00 * w0 + val01 * w1) * h_w0 + (val10 * w0 + val11 * w1) * h_w1
float val_h0 = val00 * wp.w0 + val01 * wp.w1;
float val_h1 = val10 * wp.w0 + val11 * wp.w1;
float result = val_h0 * hp.w0 + val_h1 * hp.w1;

dst_base[h * out_w + w] = utils::cast<T>(result);
out_ptr[output_base + h * output_stride_h + w * output_stride_w] = utils::cast<T>(result);
}
}
}
Expand Down
26 changes: 17 additions & 9 deletions src/infiniop/ops/upsample_bilinear/cuda/kernel.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#define __UPSAMPLE_BILINEAR_CUDA_CUH__

#include <cmath>
#include <cstddef>
#include <cstdio>

namespace op::upsample_bilinear::cuda {
Expand Down Expand Up @@ -39,6 +40,14 @@ __global__ void upsample_bilinear_kernel(
size_t W_in,
size_t H_out,
size_t W_out,
ptrdiff_t input_stride_n,
ptrdiff_t input_stride_c,
ptrdiff_t input_stride_h,
ptrdiff_t input_stride_w,
ptrdiff_t output_stride_n,
ptrdiff_t output_stride_c,
ptrdiff_t output_stride_h,
ptrdiff_t output_stride_w,
float scale_h, // 预计算的缩放比例
float scale_w, // 预计算的缩放比例
bool align_corners) {
Expand Down Expand Up @@ -82,20 +91,19 @@ __global__ void upsample_bilinear_kernel(
w0 = clamp(w0, 0, static_cast<int>(W_in) - 1);
w1 = clamp(w1, 0, static_cast<int>(W_in) - 1);

// 6. 读取数据
// 计算当前 Batch 和 Channel 的 Input 基地址
const T *img_base = input + (n_idx * C + c_idx) * H_in * W_in;

float val00 = static_cast<float>(img_base[h0 * W_in + w0]);
float val01 = static_cast<float>(img_base[h0 * W_in + w1]);
float val10 = static_cast<float>(img_base[h1 * W_in + w0]);
float val11 = static_cast<float>(img_base[h1 * W_in + w1]);
const ptrdiff_t input_base = n_idx * input_stride_n + c_idx * input_stride_c;
float val00 = static_cast<float>(input[input_base + h0 * input_stride_h + w0 * input_stride_w]);
float val01 = static_cast<float>(input[input_base + h0 * input_stride_h + w1 * input_stride_w]);
float val10 = static_cast<float>(input[input_base + h1 * input_stride_h + w0 * input_stride_w]);
float val11 = static_cast<float>(input[input_base + h1 * input_stride_h + w1 * input_stride_w]);

// 7. 双线性插值计算
// result = (val00 * w0 + val01 * w1) * h0 + (val10 * w0 + val11 * w1) * h1
float val = h0_lambda * (w0_lambda * val00 + w1_lambda * val01) + h1_lambda * (w0_lambda * val10 + w1_lambda * val11);

output[i] = static_cast<T>(val);
const ptrdiff_t output_offset = n_idx * output_stride_n + c_idx * output_stride_c
+ h_out_idx * output_stride_h + w_out_idx * output_stride_w;
output[output_offset] = static_cast<T>(val);
}
}

Expand Down
56 changes: 28 additions & 28 deletions src/infiniop/ops/upsample_bilinear/info.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

#include "../../../utils.h"
#include "../../tensor.h"
#include <array>
#include <cstddef>
#include <vector>

namespace op::upsample_bilinear {
Expand All @@ -23,6 +25,8 @@ class UpsampleBilinearInfo {
size_t _w_in; // Input Width
size_t _h_out; // Output Height
size_t _w_out; // Output Width
std::array<ptrdiff_t, 4> _input_strides;
std::array<ptrdiff_t, 4> _output_strides;

int dtype() const { return _dtype; }
bool align_corners() const { return _align_corners; }
Expand All @@ -32,27 +36,32 @@ class UpsampleBilinearInfo {
size_t w_in() const { return _w_in; }
size_t h_out() const { return _h_out; }
size_t w_out() const { return _w_out; }
ptrdiff_t input_stride(size_t dim) const { return _input_strides[dim]; }
ptrdiff_t output_stride(size_t dim) const { return _output_strides[dim]; }

// 构造函数
UpsampleBilinearInfo(int dtype, bool align_corners,
size_t n, size_t c,
size_t h_in, size_t w_in,
size_t h_out, size_t w_out)
size_t h_out, size_t w_out,
std::array<ptrdiff_t, 4> input_strides,
std::array<ptrdiff_t, 4> output_strides)
: _dtype(dtype), _align_corners(align_corners),
_n(n), _c(c),
_h_in(h_in), _w_in(w_in),
_h_out(h_out), _w_out(w_out) {}
_h_out(h_out), _w_out(w_out),
_input_strides(input_strides),
_output_strides(output_strides) {}

static utils::Result<UpsampleBilinearInfo> create(
infiniopTensorDescriptor_t out_desc,
infiniopTensorDescriptor_t input_desc,
int align_corners) { // C 接口通常传入 int 替代 bool

// 1. 检查维度数量
// 至少需要 2 维 (H, W)
// 修复: 使用 size_t 避免与 ndim() 返回值比较时的 signed/unsigned 警告
// Normalize [H, W], [C, H, W], and [N, C, H, W] to NCHW.
size_t ndim = input_desc->ndim();
if (ndim < 2) {
if (ndim < 2 || ndim > 4) {
return INFINI_STATUS_BAD_TENSOR_SHAPE;
}
if (out_desc->ndim() != ndim) {
Expand All @@ -67,32 +76,21 @@ class UpsampleBilinearInfo {

// 3. 检查 Batch/Channel 维度一致性
// 除了最后两维 (H, W)，前面的维度必须完全匹配
size_t n = 1;
size_t c = 1;

// 解析 N 和 C 用于 Info 缓存
// 逻辑：
// ndim = 4: [N, C, H, W] -> n=dims[0], c=dims[1]
// ndim = 3: [C, H, W] -> n=1, c=dims[0]
// ndim = 2: [H, W] -> n=1, c=1
// 其他情况将所有非 spatial 维度累乘到 c 中 (视为 flattened channels)

for (size_t i = 0; i < ndim - 2; ++i) { // 循环变量 i 也建议改为 size_t
for (size_t i = 0; i < ndim - 2; ++i) {
if (input_desc->shape()[i] != out_desc->shape()[i]) {
return INFINI_STATUS_BAD_TENSOR_SHAPE;
}
}

// 简单 heuristic 来填充 n 和 c
if (ndim == 4 && i == 0) {
n = input_desc->shape()[i];
} else if (ndim == 4 && i == 1) {
c = input_desc->shape()[i];
} else if (ndim == 3 && i == 0) {
c = input_desc->shape()[i];
} else {
// 对于 >4 维的情况，简单地归约为 c
c *= input_desc->shape()[i];
}
size_t n = ndim == 4 ? input_desc->shape()[0] : 1;
size_t c = ndim == 4 ? input_desc->shape()[1]
: (ndim == 3 ? input_desc->shape()[0] : 1);
std::array<ptrdiff_t, 4> input_strides{0, 0, 0, 0};
std::array<ptrdiff_t, 4> output_strides{0, 0, 0, 0};
const size_t stride_offset = 4 - ndim;
for (size_t i = 0; i < ndim; ++i) {
input_strides[stride_offset + i] = input_desc->strides()[i];
output_strides[stride_offset + i] = out_desc->strides()[i];
}

// 4. 获取空间维度
Expand All @@ -114,7 +112,9 @@ class UpsampleBilinearInfo {
h_in,
w_in,
h_out,
w_out});
w_out,
input_strides,
output_strides});
}
};

Expand Down
29 changes: 21 additions & 8 deletions src/infiniop/ops/upsample_bilinear/metax/upsample_bilinear_metax.maca
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include "upsample_bilinear_metax.h"
#include <algorithm>
#include <cmath>
#include <cstddef>
#include <cstdio>

namespace op::upsample_bilinear::metax {
Expand Down Expand Up @@ -60,6 +61,14 @@ __global__ void upsample_bilinear_kernel(
size_t W_in,
size_t H_out,
size_t W_out,
ptrdiff_t input_stride_n,
ptrdiff_t input_stride_c,
ptrdiff_t input_stride_h,
ptrdiff_t input_stride_w,
ptrdiff_t output_stride_n,
ptrdiff_t output_stride_c,
ptrdiff_t output_stride_h,
ptrdiff_t output_stride_w,
float scale_h, // 预计算的缩放比例
float scale_w, // 预计算的缩放比例
bool align_corners) {
Expand Down Expand Up @@ -101,18 +110,18 @@ __global__ void upsample_bilinear_kernel(
w0 = clamp(w0, 0, static_cast<int>(W_in) - 1);
w1 = clamp(w1, 0, static_cast<int>(W_in) - 1);

// 6. 读取数据并转换为 float
const T *img_base = input + (n_idx * C + c_idx) * H_in * W_in;

float val00 = to_float(img_base[h0 * W_in + w0]);
float val01 = to_float(img_base[h0 * W_in + w1]);
float val10 = to_float(img_base[h1 * W_in + w0]);
float val11 = to_float(img_base[h1 * W_in + w1]);
const ptrdiff_t input_base = n_idx * input_stride_n + c_idx * input_stride_c;
float val00 = to_float(input[input_base + h0 * input_stride_h + w0 * input_stride_w]);
float val01 = to_float(input[input_base + h0 * input_stride_h + w1 * input_stride_w]);
float val10 = to_float(input[input_base + h1 * input_stride_h + w0 * input_stride_w]);
float val11 = to_float(input[input_base + h1 * input_stride_h + w1 * input_stride_w]);

// 7. 双线性插值计算
float val = h0_lambda * (w0_lambda * val00 + w1_lambda * val01) + h1_lambda * (w0_lambda * val10 + w1_lambda * val11);

output[i] = static_cast<T>(val);
const ptrdiff_t output_offset = n_idx * output_stride_n + c_idx * output_stride_c
+ h_out_idx * output_stride_h + w_out_idx * output_stride_w;
output[output_offset] = static_cast<T>(val);
}
}

Expand Down Expand Up @@ -165,6 +174,10 @@ void launch_kernel(
out_ptr,
in_ptr,
N, C, H_in, W_in, H_out, W_out,
info.input_stride(0), info.input_stride(1),
info.input_stride(2), info.input_stride(3),
info.output_stride(0), info.output_stride(1),
info.output_stride(2), info.output_stride(3),
scale_h, scale_w,
align_corners);
}
Expand Down
4 changes: 4 additions & 0 deletions src/infiniop/ops/upsample_bilinear/moore/upsample_bilinear_moore.mu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,10 @@ void launch_kernel(
out_ptr,
in_ptr,
N, C, H_in, W_in, H_out, W_out,
info.input_stride(0), info.input_stride(1),
info.input_stride(2), info.input_stride(3),
info.output_stride(0), info.output_stride(1),
info.output_stride(2), info.output_stride(3),
scale_h, scale_w,
align_corners);
}
Expand Down
24 changes: 17 additions & 7 deletions src/infiniop/ops/upsample_bilinear/moore/upsample_bilinear_moore_kernel.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#define __UPSAMPLE_BILINEAR_MOORE_H__

#include <cmath>
#include <cstddef>
#include <cstdio>
#include <musa_bf16.h>
#include <musa_fp16.h>
Expand Down Expand Up @@ -33,6 +34,14 @@ __global__ void upsample_bilinear_kernel(
size_t W_in,
size_t H_out,
size_t W_out,
ptrdiff_t input_stride_n,
ptrdiff_t input_stride_c,
ptrdiff_t input_stride_h,
ptrdiff_t input_stride_w,
ptrdiff_t output_stride_n,
ptrdiff_t output_stride_c,
ptrdiff_t output_stride_h,
ptrdiff_t output_stride_w,
float scale_h,
float scale_w,
bool align_corners) {
Expand Down Expand Up @@ -67,16 +76,17 @@ __global__ void upsample_bilinear_kernel(
w0 = clamp(w0, 0, static_cast<int>(W_in) - 1);
w1 = clamp(w1, 0, static_cast<int>(W_in) - 1);

const T *img_base = input + (n_idx * C + c_idx) * H_in * W_in;

float val00 = static_cast<float>(img_base[h0 * W_in + w0]);
float val01 = static_cast<float>(img_base[h0 * W_in + w1]);
float val10 = static_cast<float>(img_base[h1 * W_in + w0]);
float val11 = static_cast<float>(img_base[h1 * W_in + w1]);
const ptrdiff_t input_base = n_idx * input_stride_n + c_idx * input_stride_c;
float val00 = static_cast<float>(input[input_base + h0 * input_stride_h + w0 * input_stride_w]);
float val01 = static_cast<float>(input[input_base + h0 * input_stride_h + w1 * input_stride_w]);
float val10 = static_cast<float>(input[input_base + h1 * input_stride_h + w0 * input_stride_w]);
float val11 = static_cast<float>(input[input_base + h1 * input_stride_h + w1 * input_stride_w]);

float val = h0_lambda * (w0_lambda * val00 + w1_lambda * val01) + h1_lambda * (w0_lambda * val10 + w1_lambda * val11);

output[i] = static_cast<T>(val);
const ptrdiff_t output_offset = n_idx * output_stride_n + c_idx * output_stride_c
+ h_out_idx * output_stride_h + w_out_idx * output_stride_w;
output[output_offset] = static_cast<T>(val);
}
}

Expand Down
4 changes: 4 additions & 0 deletions src/infiniop/ops/upsample_bilinear/nvidia/upsample_bilinear_nvidia.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,10 @@ void launch_kernel(
out_ptr,
in_ptr,
N, C, H_in, W_in, H_out, W_out,
info.input_stride(0), info.input_stride(1),
info.input_stride(2), info.input_stride(3),
info.output_stride(0), info.output_stride(1),
info.output_stride(2), info.output_stride(3),
scale_h, scale_w,
align_corners);
}
Expand Down
2 changes: 2 additions & 0 deletions test/infinicore/ops/upsample_bilinear.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@
((2, 3, 6, 6), (12, 12), None, None),
((4, 3, 7, 7), 2.0, False, None),
((3, 3, 5, 5), (10, 10), True, None),
# Channel-first view of a contiguous NHWC tensor, as used by Qwen3-VL.
((1, 8, 4, 4), (6, 8), True, (128, 1, 32, 8)),
]

_TOLERANCE_MAP = {
Expand Down
Loading