vulkan: add col2im_1d op (#24425)
* vulkan: add GGML_OP_COL2IM_1D, follow-up to the CPU op * vulkan: col2im_1d bounded gather loop instead of full-K scan with modulo * vulkan: col2im_1d address review from @jeffbolznv * vulkan: col2im_1d return nullptr for unsupported types, address review from @0cc4m
This commit is contained in:
@@ -0,0 +1,61 @@
|
||||
#version 450
|
||||
|
||||
#include "types.glsl"
|
||||
|
||||
layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; // columns: [K_OC, T_in]
|
||||
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];}; // output: [T_out, OC]
|
||||
|
||||
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
|
||||
|
||||
layout (push_constant) uniform parameter {
|
||||
uint32_t T_out;
|
||||
uint32_t OC;
|
||||
uint32_t K_OC;
|
||||
uint32_t T_in;
|
||||
uint32_t K;
|
||||
int32_t stride;
|
||||
int32_t p0;
|
||||
} p;
|
||||
|
||||
// Load A_TYPE to float
|
||||
float load_col(uint32_t idx) {
|
||||
#if defined(DATA_A_BF16)
|
||||
return bf16_to_fp32(uint32_t(data_a[idx]));
|
||||
#else
|
||||
return float(data_a[idx]);
|
||||
#endif
|
||||
}
|
||||
|
||||
// Store float as D_TYPE
|
||||
void store_dst(uint32_t idx, float v) {
|
||||
#if defined(DATA_A_BF16)
|
||||
data_d[idx] = D_TYPE(fp32_to_bf16(v));
|
||||
#else
|
||||
data_d[idx] = D_TYPE(v);
|
||||
#endif
|
||||
}
|
||||
|
||||
void main() {
|
||||
const uint32_t t_out = gl_GlobalInvocationID.x;
|
||||
const uint32_t oc = gl_GlobalInvocationID.y;
|
||||
if (t_out >= p.T_out || oc >= p.OC) return;
|
||||
|
||||
const int32_t t_abs = int32_t(t_out) + p.p0; // absolute position in uncropped signal
|
||||
|
||||
// Gather: only the ceil(K/stride) columns that scatter into t_abs, no modulo
|
||||
int32_t t_in_min = (t_abs - int32_t(p.K) + p.stride) / p.stride;
|
||||
if (t_in_min < 0) t_in_min = 0;
|
||||
int32_t t_in_max = t_abs / p.stride;
|
||||
if (t_in_max >= int32_t(p.T_in)) t_in_max = int32_t(p.T_in) - 1;
|
||||
|
||||
float val = 0.0;
|
||||
for (int32_t t_in = t_in_min; t_in <= t_in_max; t_in++) {
|
||||
int32_t k = t_abs - t_in * p.stride;
|
||||
// col layout: [K_OC, T_in], column index = oc * K + k
|
||||
uint32_t col_idx = (oc * p.K + uint32_t(k)) + uint32_t(t_in) * p.K_OC;
|
||||
val += load_col(col_idx);
|
||||
}
|
||||
|
||||
// dst layout: [T_out, OC], element (t_out, oc) = t_out + oc * T_out
|
||||
store_dst(t_out + oc * p.T_out, val);
|
||||
}
|
||||
@@ -1003,6 +1003,9 @@ void process_shaders() {
|
||||
string_to_spv("timestep_embedding_f32", "timestep_embedding.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}}));
|
||||
|
||||
string_to_spv("conv_transpose_1d_f32", "conv_transpose_1d.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}});
|
||||
string_to_spv("col2im_1d_f32", "col2im_1d.comp", {{"DATA_A_F32", "1"}, {"A_TYPE", "float"}, {"D_TYPE", "float"}});
|
||||
string_to_spv("col2im_1d_f16", "col2im_1d.comp", {{"DATA_A_F16", "1"}, {"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
|
||||
string_to_spv("col2im_1d_bf16", "col2im_1d.comp", {{"DATA_A_BF16", "1"}, {"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}});
|
||||
|
||||
string_to_spv("snake_f32", "snake.comp", {{"DATA_A_F32", "1"}, {"A_TYPE", "float"}, {"D_TYPE", "float"}});
|
||||
string_to_spv("snake_f16", "snake.comp", {{"DATA_A_F16", "1"}, {"A_TYPE", "float16_t"}, {"D_TYPE", "float16_t"}});
|
||||
|
||||
Reference in New Issue
Block a user