Skip to content

Commit 45581d5

Browse files
author
ssjia
committed
Update (base update)
[ghstack-poisoned]
1 parent b1385db commit 45581d5

11 files changed

Lines changed: 1669 additions & 0 deletions

File tree

‎backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.glsl‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
/*
10+
* conv2d_gemm: GEMM step of im2col-backed conv2d.
11+
*
12+
* Reads the im2col'd input produced by conv2d_im2col.glsl as a 2D matrix
13+
* of shape [M, K_total] (M = H_out * W_out, K_total = Kh*Kw*Cin_padded)
14+
* and writes the conv2d output as texture3D channels-packed
15+
* logical shape [1, C_out, H_out, W_out].
16+
*
17+
* The im2col input can be any of:
18+
* - texture2d, width-packed: texel at (k4, m) holds 4 K values for row m.
19+
* IN_STORAGE=texture2d codegen.
20+
* - texture3d, channels-packed: texel at (ow, oh, k4) holds 4 K values
21+
* for output spatial position (oh, ow). Used when M would exceed
22+
* max_texture2d_dim. IN_STORAGE=texture3d codegen.
23+
* - buffer: vec4 at offset m*K4 + k4, same K packing.
24+
* IN_STORAGE=buffer codegen.
25+
*
26+
* The matmul interpretation is:
27+
* out[m, n] = sum_k im2col[m, k] * weight[n, k] + bias[n]
28+
* with M = H_out * W_out, K = K_total, N = C_out.
29+
*/
30+
31+
#version 450 core
32+
33+
#define PRECISION ${PRECISION}
34+
35+
$if IN_STORAGE == "buffer" and DTYPE == "half":
36+
${define_explicit_type_extensions(DTYPE)}
37+
38+
// VEC4_T is the input storage's natural texel type, which is also the tile type
39+
// (the linear_fp_*_tile headers default the tile vec4 type to VEC4_T). For the
40+
// buffer/half path this resolves to f16vec4, so the GEMM inner loop accumulates
41+
// in true FP16 — the fma emits mad.f16 and the accumulators live in half-width
42+
// registers. Texture-sampled half always returns vec4, so FP16 accumulation is
43+
// naturally confined to the buffer (Mali) path; the texture variants (Adreno),
44+
// where FP16 accumulation regresses, stay vec4 / FP32 with no extra gating.
45+
#define VEC4_T ${texel_load_type(DTYPE, IN_STORAGE)}
46+
47+
// OUT_VEC4_T is the output surface type. t_out is always texture3d, whose
48+
// imageStore ABI takes vec4 (fp32) regardless of DTYPE, so the accumulator tile
49+
// is cast from VEC4_T to OUT_VEC4_T at store time.
50+
#define OUT_VEC4_T ${texel_load_type(DTYPE, "texture3d")}
51+
52+
#define TILE_M4 ${TILE_M4}
53+
#define TILE_K4 ${TILE_K4}
54+
#define TILE_N4 ${TILE_N4}
55+
56+
#define TILE_M ${TILE_M}
57+
#define TILE_K ${TILE_K4 * 4}
58+
#define TILE_N ${TILE_N4 * 4}
59+
60+
$if IN_STORAGE == "buffer":
61+
#define INPUT_BUFFER
62+
$elif IN_STORAGE == "texture3d":
63+
#define INPUT_TEXTURE3D
64+
65+
${define_required_extensions("texture3d", DTYPE)}
66+
$if IN_STORAGE == "buffer":
67+
${define_required_extensions("buffer", DTYPE)}
68+
69+
layout(std430) buffer;
70+
71+
#include "common.glslh"
72+
73+
${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")}
74+
$if IN_STORAGE == "buffer":
75+
${layout_declare_tensor(B, "r", "t_in", DTYPE, "buffer", is_scalar_array=False)}
76+
$else:
77+
${layout_declare_tensor(B, "r", "t_in", DTYPE, IN_STORAGE)}
78+
${layout_declare_tensor(B, "r", "t_weight_packed", DTYPE, "texture2d")}
79+
${layout_declare_tensor(B, "r", "t_bias", DTYPE, "texture2d")}
80+
81+
${layout_declare_ubo(B, "ivec4", "out_sizes")}
82+
83+
// Push constants are uploaded in 16-byte chunks (one ivec4 each).
84+
layout(push_constant) uniform restrict Block {
85+
ivec4 gemm_dims; // (K_total, K4_total, M, _unused)
86+
vec4 clamp_vals; // (out_min, out_max, _unused, _unused)
87+
};
88+
89+
#define K_TOTAL gemm_dims.x
90+
#define K4_TOTAL gemm_dims.y
91+
#define M_TOTAL gemm_dims.z
92+
#define OUT_MIN clamp_vals.x
93+
#define OUT_MAX clamp_vals.y
94+
95+
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
96+
97+
${layout_declare_spec_const(C, "int", "activation_type", "0")}
98+
99+
#include "linear_fp_input_tile.glslh"
100+
#include "linear_fp_packed_weight_tile_load.glslh"
101+
#include "linear_fp_output_tile_fp_compute.glslh"
102+
103+
/*
104+
* Load TILE_M rows × TILE_K4 K-tiles of the im2col'd input.
105+
* The im2col output is a contiguous (M, K_total/4) matrix of vec4s, so the
106+
* load is a plain 2D fetch — no spatial decomposition.
107+
*/
108+
void load_input_tile_with_checks(
109+
out FPInputTile tile,
110+
const int k4_start,
111+
const int m_start,
112+
const int K4,
113+
const int M,
114+
const int W_out) {
115+
// W_out is only consumed by the texture3d variant below.
116+
[[unroll]] for (int m = 0; m < TILE_M; ++m) {
117+
[[unroll]] for (int k4 = 0; k4 < TILE_K4; ++k4) {
118+
if (k4_start + k4 < K4 && m_start + m < M) {
119+
const int row = m_start + m;
120+
const int col = k4_start + k4;
121+
#if defined(INPUT_BUFFER)
122+
// Cast SSBO texel into the input tile type (f16vec4 for half, vec4 for
123+
// float).
124+
tile.data[m][k4] = LINEAR_FP_INPUT_TILE_VEC4_T(t_in[row * K4 + col]);
125+
#elif defined(INPUT_TEXTURE3D)
126+
// texture3d layout: row (the flat M index) decomposes into (ow, oh)
127+
// and K4 is along the Z axis. texelFetch returns vec4 (fp32); cast to
128+
// the input tile type.
129+
tile.data[m][k4] = LINEAR_FP_INPUT_TILE_VEC4_T(
130+
texelFetch(t_in, ivec3(row % W_out, row / W_out, col), 0));
131+
#else
132+
tile.data[m][k4] =
133+
LINEAR_FP_INPUT_TILE_VEC4_T(texelFetch(t_in, ivec2(col, row), 0));
134+
#endif
135+
} else {
136+
tile.data[m][k4] = LINEAR_FP_INPUT_TILE_VEC4_T(0.0);
137+
}
138+
}
139+
}
140+
}
141+
142+
void store_output_tile_with_checks(
143+
const FPOutTile out_tile,
144+
const int n4_start,
145+
const int m_start,
146+
const int N4,
147+
const int M,
148+
const int W_out) {
149+
[[unroll]] for (int m = 0; m < TILE_M; ++m) {
150+
[[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) {
151+
if (m_start + m < M && n4_start + n4 < N4) {
152+
const int spatial = m_start + m;
153+
// Cast the accumulator (f16vec4 for the buffer/half path) to the
154+
// texture3d output surface type for the activation clamp and store.
155+
OUT_VEC4_T texel = OUT_VEC4_T(out_tile.data[m][n4]);
156+
if (activation_type == 1) {
157+
texel = max(texel, OUT_VEC4_T(0.0));
158+
} else if (activation_type == 2) {
159+
texel = clamp(texel, OUT_VEC4_T(OUT_MIN), OUT_VEC4_T(OUT_MAX));
160+
}
161+
imageStore(
162+
t_out, ivec3(spatial % W_out, spatial / W_out, n4_start + n4), texel);
163+
}
164+
}
165+
}
166+
}
167+
168+
void main() {
169+
const int tile_idx_n = int(gl_GlobalInvocationID.x);
170+
const int tile_idx_m = int(gl_GlobalInvocationID.y);
171+
172+
const int n4_start = tile_idx_n * TILE_N4;
173+
const int m_start = tile_idx_m * TILE_M;
174+
175+
const int W_out = out_sizes.x;
176+
const int H_out = out_sizes.y;
177+
const int M = M_TOTAL;
178+
const int K4 = K4_TOTAL;
179+
const int N = out_sizes.z;
180+
const int N4 = div_up_4(N);
181+
182+
if (n4_start >= N4 || m_start >= M) {
183+
return;
184+
}
185+
186+
FPOutTile out_tile;
187+
initialize(out_tile);
188+
189+
FPInputTile in_tile;
190+
FPWeightTile w_tile;
191+
192+
for (int k4 = 0; k4 < K4; k4++) {
193+
load_input_tile_with_checks(in_tile, k4, m_start, K4, M, W_out);
194+
load_packed_weight_tile_with_checks(w_tile, n4_start, k4, 0, N4, K4);
195+
fp_accumulate_with_fp_weight(out_tile, in_tile, w_tile);
196+
}
197+
198+
// Apply bias
199+
[[unroll]] for (int m = 0; m < TILE_M; ++m) {
200+
[[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) {
201+
if (n4_start + n4 < N4) {
202+
// t_bias is an fp32 texture2d; cast its texel to the accumulator type.
203+
out_tile.data[m][n4] += LINEAR_FP_OUTPUT_TILE_VEC4_T(
204+
texelFetch(t_bias, ivec2(n4_start + n4, 0), 0));
205+
}
206+
}
207+
}
208+
209+
store_output_tile_with_checks(out_tile, n4_start, m_start, N4, M, W_out);
210+
}

‎backends/vulkan/runtime/graph/ops/glsl/conv2d_gemm.yaml‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
conv2d_gemm:
8+
parameter_names_with_default_values:
9+
DTYPE: float
10+
IN_STORAGE: texture2d
11+
TILE_M4: 1
12+
TILE_K4: 1
13+
TILE_N4: 1
14+
TILE_M: 4
15+
generate_variant_forall:
16+
combination:
17+
parameter_names: [IN_STORAGE, DTYPE]
18+
combos:
19+
- parameter_values: [texture2d, float]
20+
- parameter_values: [texture2d, half]
21+
- parameter_values: [texture3d, float]
22+
- parameter_values: [texture3d, half]
23+
- parameter_values: [buffer, float]
24+
- parameter_values: [buffer, half]
25+
shader_variants:
26+
- NAME: conv2d_gemm

‎backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.glsl‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
/*
10+
* Im2col transformation for FP32 / FP16 conv2d.
11+
*
12+
* The output is a 2D matrix of shape [M, K_total] where
13+
* M = H_out * W_out (number of output spatial positions)
14+
* K_total = Kh * Kw * align_up_4(C_in) (flattened receptive field)
15+
*
16+
* K layout (so a 4-tile in K — one vec4 — holds the same kernel position):
17+
* K = (ki * Kw + kj) * Cin_padded + ci
18+
*
19+
* Three codegen'd storage variants of the output tensor:
20+
* - texture2d, width-packed: texel at (k4, m) holds 4 K values for spatial
21+
* position m. Extents = (K_total/4, M).
22+
* - texture3d, channels-packed: texel at (ow, oh, k4) holds 4 K values
23+
* for output spatial position (oh, ow). Extents = (W_out, H_out, K4).
24+
* Used as a fallback when M would exceed max_texture2d_dim.
25+
* - buffer: vec4 at offset (m * K4 + k4), same K packing.
26+
*
27+
* The caller picks storage per device (Mali → buffer; others → texture2d
28+
* when its 2D extents fit, texture3d when its 3D extents fit, else buffer).
29+
*/
30+
31+
#version 450 core
32+
33+
#define PRECISION ${PRECISION}
34+
35+
#define VEC4_T ${texel_load_type(DTYPE, "texture3d")}
36+
37+
$if OUT_STORAGE == "buffer":
38+
#define OUTPUT_BUFFER
39+
#define VEC4_BUF_T ${texel_load_type(DTYPE, "buffer")}
40+
$elif OUT_STORAGE == "texture3d":
41+
#define OUTPUT_TEXTURE3D
42+
43+
${define_required_extensions("texture3d", DTYPE)}
44+
$if OUT_STORAGE == "buffer":
45+
${define_required_extensions("buffer", DTYPE)}
46+
47+
layout(std430) buffer;
48+
49+
$if OUT_STORAGE == "buffer":
50+
${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer", is_scalar_array=False)}
51+
$else:
52+
${layout_declare_tensor(B, "w", "t_out", DTYPE, OUT_STORAGE)}
53+
${layout_declare_tensor(B, "r", "t_in", DTYPE, "texture3d")}
54+
55+
${layout_declare_ubo(B, "ivec4", "in_sizes")}
56+
57+
// Push constants are uploaded in 16-byte chunks (one ivec4 each) to comply
58+
// with the per-entry size limit.
59+
layout(push_constant) uniform restrict Block {
60+
ivec4 kernel_stride; // (Kh, Kw, Sh, Sw)
61+
ivec4 padding_dil; // (Ph, Pw, Dh, Dw)
62+
ivec4 dims; // (Cin_padded, W_out, H_out, K4_total)
63+
};
64+
65+
#define KERNEL_H kernel_stride.x
66+
#define KERNEL_W kernel_stride.y
67+
#define STRIDE_H kernel_stride.z
68+
#define STRIDE_W kernel_stride.w
69+
#define PADDING_H padding_dil.x
70+
#define PADDING_W padding_dil.y
71+
#define DILATION_H padding_dil.z
72+
#define DILATION_W padding_dil.w
73+
#define CIN_PADDED dims.x
74+
#define W_OUT dims.y
75+
#define H_OUT dims.z
76+
#define K4_TOTAL dims.w
77+
78+
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
79+
80+
void main() {
81+
const int k4 = int(gl_GlobalInvocationID.x);
82+
const int m = int(gl_GlobalInvocationID.y);
83+
const int M = H_OUT * W_OUT;
84+
85+
if (k4 >= K4_TOTAL || m >= M) {
86+
return;
87+
}
88+
89+
const int k_start = k4 * 4;
90+
91+
// K = (ki * Kw + kj) * Cin_padded + ci ; since Cin_padded % 4 == 0, all 4
92+
// K values in this texel share the same (ki, kj) and span 4 consecutive
93+
// ci values starting at ci_start.
94+
const int krow_idx = k_start / CIN_PADDED; // ki * Kw + kj
95+
const int ci_start = k_start % CIN_PADDED;
96+
const int kj = krow_idx % KERNEL_W;
97+
const int ki = krow_idx / KERNEL_W;
98+
const int ci_blk = ci_start >> 2; // ci_start / 4
99+
100+
// Decompose flat output position m back into (oh, ow).
101+
const int ow = m % W_OUT;
102+
const int oh = m / W_OUT;
103+
104+
// Compute the input spatial position for this (oh, ow, ki, kj).
105+
const int ih = oh * STRIDE_H - PADDING_H + ki * DILATION_H;
106+
const int iw = ow * STRIDE_W - PADDING_W + kj * DILATION_W;
107+
108+
VEC4_T out_texel = VEC4_T(0);
109+
if (ih >= 0 && ih < in_sizes.y && iw >= 0 && iw < in_sizes.x) {
110+
out_texel = texelFetch(t_in, ivec3(iw, ih, ci_blk), 0);
111+
}
112+
113+
#if defined(OUTPUT_BUFFER)
114+
t_out[m * K4_TOTAL + k4] = VEC4_BUF_T(out_texel);
115+
#elif defined(OUTPUT_TEXTURE3D)
116+
imageStore(t_out, ivec3(ow, oh, k4), out_texel);
117+
#else
118+
imageStore(t_out, ivec2(k4, m), out_texel);
119+
#endif
120+
}

‎backends/vulkan/runtime/graph/ops/glsl/conv2d_im2col.yaml‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
#
4+
# This source code is licensed under the BSD-style license found in the
5+
# LICENSE file in the root directory of this source tree.
6+
7+
conv2d_im2col:
8+
parameter_names_with_default_values:
9+
DTYPE: float
10+
OUT_STORAGE: texture2d
11+
generate_variant_forall:
12+
combination:
13+
parameter_names: [OUT_STORAGE, DTYPE]
14+
combos:
15+
- parameter_values: [texture2d, float]
16+
- parameter_values: [texture2d, half]
17+
- parameter_values: [texture3d, float]
18+
- parameter_values: [texture3d, half]
19+
- parameter_values: [buffer, float]
20+
- parameter_values: [buffer, half]
21+
shader_variants:
22+
- NAME: conv2d_im2col

0 commit comments

Comments
 (0)