Skip to content

feature/adm/cuda: clamp the scale 0 contrast masking neighbours at the border #1649

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

Open
lusoris wants to merge 2 commits into
base: master
Choose a base branch
from
Open

feature/adm/cuda: clamp the scale 0 contrast masking neighbours at the border #1649

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4cb62aa
feature/adm/cuda: round the scale 0 and scale 1 to 3 contrast masking…
lusoris Oct 1, 2026
d3df91e
feature/adm/cuda: clamp the scale 0 contrast masking neighbours at th…
lusoris Oct 1, 2026
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
118 changes: 39 additions & 79 deletions libvmaf/src/feature/cuda/integer_adm/adm_cm.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,31 +115,14 @@ __constant__ const int fixed_shift[3] = {4, 4, 3};
__constant__ const int32_t shift_xsq[3] = {29, 29, 30};
__constant__ const int32_t add_shift_xsq[3] = {268435456, 268435456, 536870912};

// HACK: the 256 byte alignment is required to ensure that the struct is not moved to lmem
struct WarpShift
{
uint32_t shift_cub[3];
uint32_t add_shift_cub[3];
uint32_t shift_sq[3];
uint32_t add_shift_sq[3];
};

template <int rows_per_thread>
__device__ __forceinline__ void adm_cm_line_kernel(AdmBufferCuda buf, int h, int w, int top,
int bottom, int left, int right,
int start_row, int end_row, int start_col,
int end_col, int src_stride,
int csf_a_stride, int buffer_h,
int buffer_stride, int32_t *accum_per_block,
AdmFixedParametersCuda params,
// reduce
int scale, int64_t* accum_global,

// shift warp
WarpShift ws,
// shift global
const uint32_t shift_inner_accum, const uint32_t add_shift_inner_accum
) {
AdmFixedParametersCuda params) {
const cuda_adm_dwt_band_t *src = &buf.decouple_r;
const cuda_adm_dwt_band_t *csf_f = &buf.csf_f;
const cuda_adm_dwt_band_t *csf_a = &buf.csf_a;
Expand All @@ -161,8 +144,11 @@ __device__ __forceinline__ void adm_cm_line_kernel(AdmBufferCuda buf, int h, int
if (y < end_row && x < end_col)
{
int pos_x[3] = {x - 1, x, x + 1};
pos_x[0] = abs(pos_x[0]);
pos_x[2] = pos_x[2] - max(0, 2*(x - w)+1);
// the left border mirrors {1, 0, 1} and the right border replicates
// {w - 2, w - 1, w - 1}, as the CPU's ADM_CM_THRESH_S_* macros do. Clamp the
// neighbour's position: x itself is inside the band.
pos_x[0] = min(abs(pos_x[0]), w - 1);
pos_x[2] = min(pos_x[2], w - 1);

#pragma unroll
for (int theta = 0; theta < 3; ++theta)
Expand All @@ -182,8 +168,7 @@ __device__ __forceinline__ void adm_cm_line_kernel(AdmBufferCuda buf, int h, int
for (int row = 0; row < total_rows;++row)
{
int pos_y = y - 1 + row;
pos_y = abs(pos_y);
pos_y = pos_y - max(0, 2*(y - h)+1);
pos_y = min(abs(pos_y), h - 1);

int16_t src = angles[theta][pos_y * src_stride + x];
int16_t *flt_ptr = flt_angles[theta] + pos_y*src_stride;
Expand Down Expand Up @@ -220,36 +205,23 @@ __device__ __forceinline__ void adm_cm_line_kernel(AdmBufferCuda buf, int h, int
accum_thread_reg[row] = max(0, sb);
}

const int band2 = blockIdx.z;
int64_t accum = 0;

// the compiler does not assume that parameters are constant, move them to local variables to give the compiler
// a hint that those values have to be loaded only once from constant memory.
int32_t add_shift_cub = ws.add_shift_cub[band2];
int32_t shift_cub = ws.shift_cub[band2];
int32_t add_shift_sq = ws.add_shift_sq[band2];
int32_t shift_sq = ws.shift_sq[band2];

// accumulate per thread
for (int row = 0;row < rows_per_thread;++row) {
int32_t accum_thread = accum_thread_reg[row];
const int32_t x_sq = (int32_t)((((int64_t)accum_thread * accum_thread) + add_shift_sq >> shift_sq));
accum += (((int64_t)x_sq * accum_thread) + add_shift_cub) >> shift_cub;
}

// accumulate warp
accum = warp_reduce(accum);

if (threadIdx.x % 32 == 0)
{
accum = (accum + add_shift_inner_accum) >> shift_inner_accum;
atomicAdd_int64(&accum_global[band2],
accum);
// Store the masked values. adm_cm_reduce_line_kernel cubes and sums them and
// applies the shift_inner_accum rounding once per row, as the CPU adm_cm does.
// Rounding the partial sum of every warp tile instead is not the same value.
if (cta_x < buffer_stride) {
for (int row = 0; row < rows_per_thread; ++row) {
const int buffer_row = cta_y + row;
if (buffer_row < buffer_h)
accum_per_block[(blockIdx.z * buffer_h + buffer_row) * buffer_stride + cta_x] =
accum_thread_reg[row];
}
}
}

template <int val_per_thread>
__device__ __forceinline__ void adm_cm_reduce_line_kernel(int h, int w, int scale, int buffer_h,
// One block per (band, row). The CPU sums a whole row before it applies the
// rounded shift_inner_accum shift, so the row is reduced within one block and
// thread 0 applies the shift once.
extern "C" __global__ void adm_cm_reduce_line_kernel(int h, int w, int scale, int buffer_h,
int buffer_stride,
const int32_t *buffer,
int64_t *accum) {
Expand All @@ -259,66 +231,54 @@ __device__ __forceinline__ void adm_cm_reduce_line_kernel(int h, int w, int scal
const int b_off = off + line * buffer_stride;

uint32_t shift_cub = __float2uint_ru(__log2f(w));
uint32_t add_shift_cub = 1 << (shift_cub - 1);
uint32_t add_shift_cub = shift_cub ? 1u << (shift_cub - 1) : 0;
int32_t shift_sq = 30;
int32_t add_shift_sq = 536870912; // 2^29
if (scale == 0) {
shift_cub = __float2uint_ru(__log2f(w) - fixed_shift[band]);
add_shift_cub = 1 << (shift_cub - 1);
add_shift_cub = shift_cub ? 1u << (shift_cub - 1) : 0;
shift_sq = shift_xsq[band];
add_shift_sq = add_shift_xsq[band];
}

int64_t temp_value = 0;
const int buffer_col = (blockDim.x * blockIdx.x + threadIdx.x) * val_per_thread;
const int32_t *buffer_loc = buffer + b_off + buffer_col;
for (int i = 0; i < val_per_thread; ++i) {
if ((buffer_col + i) < buffer_stride) {
const int32_t x = buffer_loc[i];
const int32_t x_sq =
(int32_t)((((int64_t)x * x) + add_shift_sq) >> shift_sq);
temp_value += (((int64_t)x_sq * x) + add_shift_cub) >> shift_cub;
}
for (int col = threadIdx.x; col < buffer_stride; col += blockDim.x) {
const int32_t x = buffer[b_off + col];
const int32_t x_sq = (int32_t)((((int64_t)x * x) + add_shift_sq) >> shift_sq);
temp_value += (((int64_t)x_sq * x) + add_shift_cub) >> shift_cub;
}
temp_value = warp_reduce(temp_value);

if ((threadIdx.x % VMAF_CUDA_THREADS_PER_WARP) == 0) {
__shared__ int64_t warp_sums[32]; // at most 1024 threads of 32 lanes
if ((threadIdx.x % VMAF_CUDA_THREADS_PER_WARP) == 0)
warp_sums[threadIdx.x / VMAF_CUDA_THREADS_PER_WARP] = temp_value;
__syncthreads();

if (threadIdx.x == 0) {
int64_t row_sum = 0;
for (int i = 0; i < blockDim.x / VMAF_CUDA_THREADS_PER_WARP; ++i)
row_sum += warp_sums[i];
const uint32_t shift_inner_accum = __float2uint_ru(__log2f(h));
const uint32_t add_shift_inner_accum = 1 << (shift_inner_accum - 1);
atomicAdd_int64(&accum[band],
(temp_value + add_shift_inner_accum) >> shift_inner_accum);
atomicAdd_int64(&accum[band], (row_sum + add_shift_inner_accum) >> shift_inner_accum);
}
}

#define ADM_CM_REDUCE_LINE(val_per_thread) \
__global__ void adm_cm_reduce_line_kernel_##val_per_thread ( \
int h, int w, int scale, int buffer_h, \
int buffer_stride, const int32_t *buffer, int64_t *accum) \
{ \
adm_cm_reduce_line_kernel<val_per_thread>( \
h, w, scale, buffer_h, buffer_stride, buffer, accum); \
}

#define ADM_CM_LINE(rows_per_thread) \
__global__ void adm_cm_line_kernel_##rows_per_thread ( \
AdmBufferCuda buf, int h, int w, int top, \
int bottom, int left, int right, int start_row, int end_row, int start_col, \
int end_col, int src_stride, int csf_a_stride, int buffer_h, \
int buffer_stride, int32_t *accum_per_block, AdmFixedParametersCuda params, \
int scale, int64_t* accum_global, WarpShift ws, \
const uint32_t shift_inner_accum, const uint32_t add_shift_inner_accum) \
int buffer_stride, int32_t *accum_per_block, AdmFixedParametersCuda params) \
{ \
adm_cm_line_kernel<rows_per_thread>( \
buf, h, w, top, bottom, left, right, start_row, end_row, start_col, \
end_col, src_stride, csf_a_stride, buffer_h, buffer_stride, \
accum_per_block, params,scale, accum_global, \
ws, shift_inner_accum, add_shift_inner_accum); \
accum_per_block, params); \
}



extern "C" {
// 128 = warps_per_thread * val_per_thread = 32 * 4 -- assuming 32 threads per warp, this might change in the future
ADM_CM_REDUCE_LINE(4); // adm_cm_reduce_line_kernel_4
ADM_CM_LINE(8); // adm_cm_line_kernel_8
}
59 changes: 17 additions & 42 deletions libvmaf/src/feature/cuda/integer_adm_cuda.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,6 @@

#define RES_BUFFER_SIZE 4 * 3 * 2

typedef struct WarpShift {
uint32_t shift_cub[3];
uint32_t add_shift_cub[3];
uint32_t shift_sq[3];
uint32_t add_shift_sq[3];
} WarpShift;

typedef struct AdmStateCuda {
size_t integer_stride;
AdmBufferCuda buf;
Expand Down Expand Up @@ -75,7 +68,7 @@ typedef struct AdmStateCuda {
func_adm_csf_den_scale_line_kernel,
func_adm_csf_den_s123_line_kernel,
// adm_cm kernel
func_adm_cm_reduce_line_kernel_4,
func_adm_cm_reduce_line_kernel,
func_adm_cm_line_kernel_8,
func_i4_adm_cm_line_kernel;

Expand Down Expand Up @@ -440,15 +433,14 @@ void i4_adm_cm_device(AdmStateCuda *s, AdmBufferCuda *buf, int w, int h, int src
0, c_stream, args, NULL));
}
{
const int val_per_thread = 4;
const int warps_per_cta = 4;
const int BLOCKX = VMAF_CUDA_THREADS_PER_WARP * warps_per_cta;

void* args[] = {
&h, &w, &scale, &buffer_h, &buffer_stride,
&buf->tmp_accum->data, &buf->adm_cm[scale]};
CHECK_CUDA(cu_f, cuLaunchKernel(s->func_adm_cm_reduce_line_kernel_4,
DIV_ROUND_UP(buffer_stride, BLOCKX * val_per_thread), buffer_h, 3,
CHECK_CUDA(cu_f, cuLaunchKernel(s->func_adm_cm_reduce_line_kernel,
1, buffer_h, 3,
BLOCKX, 1, 1,
0, c_stream, args, NULL));
}
Expand All @@ -472,33 +464,6 @@ void adm_cm_device(AdmStateCuda *s, AdmBufferCuda *buf, int w, int h, int src_st
int buffer_stride = end_col - start_col;
int buffer_h = end_row - start_row;

// precompute warp shift per band
//const int32_t shift_sub[3] = {10, 10, 12};
const int fixed_shift[3] = {4, 4, 3};

// accumulation
const int32_t shift_xsq[3] = {29, 29, 30};
const int32_t add_shift_xsq[3] = {268435456, 268435456, 536870912};

const int NUM_BANDS = 3;
WarpShift ws;
for (int band = 0;band < NUM_BANDS;++band) {
ws.shift_cub[band] = (uint32_t)(ceil(log2f(w)));
if (scale == 0) {
ws.shift_cub[band] -= fixed_shift[band];
ws.shift_sq[band] = shift_xsq[band];
ws.add_shift_sq[band] = add_shift_xsq[band];
} else {
ws.shift_sq[band] = 30;
ws.add_shift_sq[band] = (1 << (ws.shift_sq[band]-1));
}
ws.add_shift_cub[band] = 1 << (ws.shift_cub[band] - 1);
}

// precompute global shift
uint32_t shift_inner_accum = (uint32_t)(ceil(log2f(h)));
uint32_t add_shift_inner_accum = 1 << (shift_inner_accum - 1);

// fused
{
const int rows_per_thread = 8;
Expand All @@ -507,16 +472,26 @@ void adm_cm_device(AdmStateCuda *s, AdmBufferCuda *buf, int w, int h, int src_st
void* args[] = {
&*buf, &h, &w, &top, &bottom, &left, &right, &start_row, &end_row, &start_col,
&end_col, &src_stride, &csf_a_stride, &buffer_h, &buffer_stride,
&buf->tmp_accum->data, &*p,
&scale, &buf->adm_cm[scale], &ws,
&shift_inner_accum, &add_shift_inner_accum
&buf->tmp_accum->data, &*p
};

CHECK_CUDA(cu_f, cuLaunchKernel(s->func_adm_cm_line_kernel_8,
DIV_ROUND_UP(buffer_stride, BLOCKX), DIV_ROUND_UP(buffer_h, BLOCKY * rows_per_thread), 3,
BLOCKX, BLOCKY, 1,
0, c_stream, args, NULL));
}
{
const int warps_per_cta = 4;
const int BLOCKX = VMAF_CUDA_THREADS_PER_WARP * warps_per_cta;

void* args[] = {
&h, &w, &scale, &buffer_h, &buffer_stride,
&buf->tmp_accum->data, &buf->adm_cm[scale]};
CHECK_CUDA(cu_f, cuLaunchKernel(s->func_adm_cm_reduce_line_kernel,
1, buffer_h, 3,
BLOCKX, 1, 1,
0, c_stream, args, NULL));
}
}

static void conclude_adm_cm(int64_t *accum, int h,
Expand Down Expand Up @@ -1038,7 +1013,7 @@ static int init_fex_cuda(VmafFeatureExtractor *fex, enum VmafPixelFormat pix_fmt
CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_adm_csf_den_scale_line_kernel, adm_csf_den_module, "adm_csf_den_scale_line_kernel_8_128"));
CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_adm_csf_den_s123_line_kernel, adm_csf_den_module, "adm_csf_den_s123_line_kernel_8_128"));

CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_adm_cm_reduce_line_kernel_4, adm_cm_module, "adm_cm_reduce_line_kernel_4"));
CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_adm_cm_reduce_line_kernel, adm_cm_module, "adm_cm_reduce_line_kernel"));
CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_adm_cm_line_kernel_8, adm_cm_module, "adm_cm_line_kernel_8"));
CHECK_CUDA(cu_f, cuModuleGetFunction(&s->func_i4_adm_cm_line_kernel, adm_cm_module, "i4_adm_cm_line_kernel"));

Expand Down
18 changes: 18 additions & 0 deletions libvmaf/test/meson.build
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,8 +204,26 @@ test_cuda_pic_preallocation = executable('test_cuda_pic_preallocation',
c_args: ['-DHAVE_CUDA=1']
)

test_cuda_adm_cm_row_rounding = executable('test_cuda_adm_cm_row_rounding',
['test.c', 'test_cuda_adm_cm_row_rounding.c'],
include_directories : [libvmaf_inc, test_inc],
link_with : get_option('default_library') == 'both' ? libvmaf.get_static_lib() : libvmaf,
dependencies: [math_lib, cuda_dependency],
c_args: ['-DHAVE_CUDA=1']
)

test_cuda_adm_cm_scale0_border = executable('test_cuda_adm_cm_scale0_border',
['test.c', 'test_cuda_adm_cm_scale0_border.c'],
include_directories : [libvmaf_inc, test_inc],
link_with : get_option('default_library') == 'both' ? libvmaf.get_static_lib() : libvmaf,
dependencies: [math_lib, cuda_dependency],
c_args: ['-DHAVE_CUDA=1']
)

test('test_ring_buffer', test_ring_buffer)
test('test_cuda_pic_preallocation', test_cuda_pic_preallocation)
test('test_cuda_adm_cm_row_rounding', test_cuda_adm_cm_row_rounding, timeout : 300)
test('test_cuda_adm_cm_scale0_border', test_cuda_adm_cm_scale0_border, timeout : 300)
endif

test_pic_preallocation = executable('test_pic_preallocation',
Expand Down
Loading