ncnn GLSL extension

rationale

Different GPUs support different features, some support fp16 as buffer storage type, some support fp16 as operand variable, some old GPUs only support fp32

When the GPU supports the VK_KHR_16bit_storage extension, in order to minimize the memory bandwidth consumption of the GPU, we will give priority to using fp16 as the storage type. Otherwise, we use packHalf2x16 and unpackHalf2x16 in GLSL 4.2 to compress 2 fp32 to uint, reducing read and write bandwidth.

Similarly, when the gpu supports the VK_KHR_shader_float16_int8 extension, in order to speed up the calculation efficiency, we will give priority to using fp16 as the operation operand, which usually doubles the speed. Otherwise, we use fp32.

To ensure the widest compatibility, the following code for declaring descriptor binding and loading data will be written

#if NCNN_fp16_storage // gpu supports 16bit storage
layout (binding = 0) buffer blob { f16vec4 blob_data[]; };
#elif NCNN_fp16_packed // gpu supports GLSL 4.2
layout (binding = 0) buffer blob { uvec2 blob_data[]; };
#else // gpu only supports fp32
layout (binding = 0) buffer blob { vec4 blob_data[]; };
#endif

void main()
{
    const int i = int(gl_GlobalInvocationID.x);

#if NCNN_fp16_storage && NCNN_fp16_arithmetic // gpu supports 16bit storage and shader float16
    f16vec4 x = blob_data[i];
#elif NCNN_fp16_storage // gpu supports 16bit storage but no shader float16
    vec4 x = vec4(blob_data[i]);
#elif NCNN_fp16_packed && NCNN_fp16_arithmetic // gpu supports GLSL 4.2 and shader float16
    f16vec4 x = f16vec4(unpackFloat2x16(blob_data[i].x), unpackFloat2x16(blob_data[i].y));
#elif NCNN_fp16_packed // gpu supports GLSL 4.2
    vec4 x = vec4(unpackHalf2x16(blob_data[i].x), unpackHalf2x16(blob_data[i].y));
#else // gpu only supports fp32
    vec4 x = blob_data[i];
#endif
}

As you can see, just declaring the buffer type and reading a value consumes a lot of lines of code, which is a maintenance nightmare. Therefore, ncnn adds more flexible data types and auxiliary functions to reduce the size of the code and improve readability, and will automatically expand to the most efficient implementation according to the feature level supported by the GPU.

The above code, by using the ncnn glsl extension, can be simplified to

layout (binding = 0) buffer blob { sfpvec4 blob_data[]; };

void main()
{
    const int i = int(gl_GlobalInvocationID.x);

    afpvec4 x = buffer_ld4(blob_data, i);
}

The ncnn glsl extension provides the necessary data types for storage, computation, shared memory, and load, store, conversion functions for buffers and images. We also provide some buffer and image copy functions to prevent loss of precision when using fp16 as the intermediate data type, and to avoid unnecessary unpackHalf2x16 and packHalf2x16 pair.

entrypoint for compiling GLSL

The gpu.h header in the ncnn library exposes 3 APIs for compiling glsl code into spir-v binary, they support ncnn glsl extension, these 3 functions accept opt switch to control the expansion form of ncnn glsl extension. The first two accept raw glsl code strings, and the last one is used to create ncnn's built-in shader.

namespace ncnn {

// online spirv compilation
NCNN_EXPORT int compile_spirv_module(const char* comp_string, const Option& opt, std::vector<uint32_t>& spirv);
NCNN_EXPORT int compile_spirv_module(const char* comp_data, int comp_data_size, const Option& opt, std::vector<uint32_t>& spirv);
NCNN_EXPORT int compile_spirv_module(int shader_type_index, const Option& opt, std::vector<uint32_t>& spirv);

} // namespace ncnn

compile ncnn extended GLSL code directly

You can write shader code with ncnn glsl extension, compiled to spir-v using ncnn functions. The compiled product is a standard-compliant spir-v binary, which can be directly used to create a pipeline object in the vulkan api

static const char my_glsl_data[] = R"(
#version 450

layout (binding = 0) readonly buffer a_blob { sfpvec4 a_blob_data[]; };
layout (binding = 1) writeonly buffer b_blob { sfpvec4 b_blob_data[]; };

void main()
{
    const int i = int(gl_GlobalInvocationID.x);

    afpvec4 v = buffer_ld4(a_blob_data, i);

    v = v + 123;

    buffer_st4(b_blob_data, i, v);
}
)";

Option opt;
 // you can control the extension behavior
 // even if the gpu supports 16bit storage
opt.use_fp16_storage = false;

std::vector<uint32_t> spirv;
ncnn::compile_spirv_module(my_glsl_data, sizeof(my_glsl_data) - 1, opt, spirv);

// To create pipeline object later
// ncnn::Pipeline pipeline(vkdev);
// pipeline.set_local_size_xyz(64, 1, 1);
// pipeline.create(spirv.data(), spirv.size() * 4, specializations);

ncnn built-in shader

The shader index inside ncnn is exposed in the layer_shader_type.h header and can be used if needed

#include "layer_shader_type.h"

int shader_type_index = LayerShaderType::convert_ycbcr;

Option opt;

std::vector<uint32_t> spirv;
int retc = compile_spirv_module(shader_type_index, opt, spirv);

data types

storage type

declare buffer data layout in descriptor binding

layout (binding = 0) buffer top_blob { sfpvec4 top_blob_data[]; };

arithmetic type

declare local variable in glsl code

void main()
{
    afpvec4 v = a * b;
}

local type

declare variable in shared local memory

shared lfp tmp_a[8][4][2];

integer type

declare int8/int16 buffer data layout and local variables in glsl code

layout (binding = 0) readonly buffer bottom_blob { sint8vec4 bottom_blob_data[]; };
layout (binding = 1) readonly buffer weight_blob { sint16 weight_blob_data[]; };

sint8vec4 uses one int to hold four signed int8 lanes in all int8 storage modes. This keeps pack4 data in packed form for integer dot-product and shared-memory paths. Use i8buffer_ld4 to unpack it to ivec4, and use i8buffer_sm4 to load the raw packed int.

sint16 uses one int to hold two signed int16 lanes when opt.use_int16_packed is enabled, and uses native int16_t when opt.use_int16_storage is enabled. sint16vec4 stores four logical int16 lanes as two packed int values in int16p mode and as native i16vec4 in int16s mode. lint16 and lint16vec4 are the shared/local-memory counterparts.

buffer functions

• load typed value from src[offset]

afp buffer_ld1(sfp src, int offset);
afpvec2 buffer_ld2(sfpvec2 src, int offset);
afpvec4 buffer_ld4(sfpvec4 src, int offset);

• store typed value to dst[offset]

void buffer_st1(sfp dst, int offset, afp v);
void buffer_st2(sfpvec2 dst, int offset, afpvec2 v);
void buffer_st4(sfpvec4 dst, int offset, afpvec4 v);

• copy typed value from src[src_offset] to dst[dst_offset]

void buffer_cp1(sfp dst, int dst_offset, sfp src, int src_offset);
void buffer_cp2(sfpvec2 dst, int dst_offset, sfpvec2 src, int src_offset);
void buffer_cp4(sfpvec4 dst, int dst_offset, sfpvec4 src, int src_offset);

• copy and pack value from src[src_offsets[0],src_offsets[1],...] to dst[dst_offset]

void buffer_cp1to4(sfpvec4 dst, int dst_offset, sfp src, ivec4 src_offsets);

• copy and unpack value from src[src_offset] to dst[dst_offsets[0],dst_offsets[1],...]

void buffer_cp4to1(sfp dst, ivec4 dst_offsets, sfpvec4 src, int src_offset);

integer buffer functions

• load integer typed value from src[offset]

aint8 i8buffer_ld1(sint8 src, int offset);
aint8vec4 i8buffer_ld4(sint8vec4 src, int offset);
int i8buffer_sm4(sint8vec4 src, int offset);
int i16buffer_ld1(sint16 src, int offset);
ivec2 i16buffer_ld2(sint16 src, int offset);
sint16vec4 i16buffer_sm4(sint16vec4 src, int offset);
aint16vec4 i16buffer_ld4(sint16vec4 src, int offset);
aint16 lint162aint16(lint16 v);
aint16vec4 lint162aint16vec4(lint16vec4 v);

i8buffer_sm4 loads the raw packed int representation of four int8 lanes. It is useful for shared-memory staging and dotPacked4x8EXT paths where unpacking to ivec4 would be wasteful.

i16buffer_ld1 and i16buffer_ld2 load signed int16 lanes as int and ivec2. Without native int16 storage, offset is still the logical int16 lane offset, and packed storage groups two adjacent lanes in one int.

i16buffer_sm4 loads the raw sint16vec4 representation of four logical int16 lanes from buffer storage. i16buffer_ld4 loads four logical int16 lanes from buffer storage as aint16vec4. lint162aint16 and lint162aint16vec4 convert shared/local int16 values to arithmetic int16 values.

• store integer typed value to dst[offset]

void i8buffer_st1(sint8 dst, int offset, aint8 v);
void i8buffer_st4(sint8vec4 dst, int offset, aint8vec4 v);
void i16buffer_st1(sint16 dst, int offset, int v);
void i16buffer_st2(sint16 dst, int offset, ivec2 v);
void i16buffer_st4(sint16vec4 dst, int offset, ivec4 v);
void i16buffer_st4(lint16vec4 dst, int offset, ivec4 v);

Without native int8 storage, i8buffer_st1 updates one byte lane inside a packed int and may use an atomic compare-and-swap loop.

Without native int16 storage, i16buffer_st1 updates one int16 lane inside a packed int and may use an atomic compare-and-swap loop. i16buffer_st2 stores complete packed words directly when offset is aligned. i16buffer_st4 writes four logical int16 lanes to sint16vec4 storage or lint16vec4 shared/local memory.

• copy int8 typed value from src[src_offset] to dst[dst_offset]

void i8buffer_cp1(sint8 dst, int dst_offset, sint8 src, int src_offset);
void i8buffer_cp4(sint8vec4 dst, int dst_offset, sint8vec4 src, int src_offset);

• copy and pack int8 typed values from src[src_offsets[0],src_offsets[1],...] to dst[dst_offset]

void i8buffer_cp1to4(sint8vec4 dst, int dst_offset, sint8 src, ivec4 src_offsets);

• copy and unpack int8 typed values from src[src_offset] to dst[dst_offsets[0],dst_offsets[1],...]

void i8buffer_cp4to1(sint8 dst, ivec4 dst_offsets, sint8vec4 src, int src_offset);

• pack and unpack signed integer lanes

ivec4 unpackInt4x8(int v);
int packInt4x8(ivec4 v);
ivec2 unpackInt2x16(int v);
int packInt2x16(ivec2 v);
int float2int8(float v);
ivec4 float2int8vec4(vec4 v);

packInt4x8 stores .r/.g/.b/.a in the low-to-high bytes of one int. packInt2x16 stores .r/.g in the low-to-high 16-bit lanes of one int. float2int8 and float2int8vec4 round half away from zero and saturate to [-127, 127] for deterministic int8 quantization.

local data conversion functions

• storage buffer to local memory

lfp buffer_sm1(sfp src, int offset);
lfpvec4 buffer_sm4(sfpvec4 src, int offset);

• local memory to local variable

afp lfp2afp(lfp v);
afpvec4 lfp2afpvec4(lfpvec4 v);

• local variable to local memory

lfp afp2lfp(afp v);
lfpvec4 afp2lfpvec4(afpvec4 v);

Note: The common usage of local memory is to read from global memory first, store it in local memory, and then read local variables from local memory for subsequent use. Therefore, only storage type to local type and local type to arithmetic type conversion functions are provided here.

misc functions

• prefer specialization constant over push constant

T psc(T x)

Declare the same variable in specialization constant AND push constant section, then psc(x) will become a compile-time constant when specialization constant given non-zero or be dynamic via push constant otherwise. This is often used for tensor shape specialization. We can usually resolve all shape information and make them be compile-time constants for more aggressive shader optimization.

layout (constant_id = 0) const int size = 0;

layout (push_constant) uniform parameter
{
    int size;
} p;

void main()
{
    const int s = psc(size);
}

platform macros

judge if the current platform is moltenvk, for enabling some platform-specific workaround

#if NCNN_moltenvk
// enable workaround for moltenvk
#endif

ncnn adds additional macro definitions in the new version, which may conflict or confuse the existing glsl code. In order to obtain cross-version compatibility of ncnn, you can switch between the old and new codes according to the ncnn_glsl_version macro version.

#if ncnn_glsl_version >= 1
// use device macros introduced since version 1
#endif

ncnn additionally defines most of the vulkan device-related features as macros, which we can use to distinguish different platforms, device extensions, features, and properties

extension macros

When the device supports an extension, ncnn_<extension_name> is defined as the extension version

void main()
{
#if ncnn_VK_KHR_16bit_storage
    // here is the code for any device that supports VK_KHR_16bit_storage
#endif

#if ncnn_VK_KHR_sampler_ycbcr_conversion >= 10
    // here is the code for any device that supports VK_KHR_sampler_ycbcr_conversion and version >= 10
#endif
}

device feature and property macros

ncnn will query device features and properties and then define them as macros.

The macro name is ncnn_<feature_name> or ncnn_<property_name>

The GL_EXT_shader_explicit_arithmetic_types_int64 extension will be automatically enabled without explicit code indication when the device supports shaderInt64

The GL_EXT_shader_explicit_arithmetic_types_int16 extension will be automatically enabled without explicit code indication when the device supports shaderInt16

void main()
{
#if ncnn_robustBufferAccess
    // here is the code for any device that supports robustBufferAccess feature
#endif

#if ncnn_vendorID == 4318
    // here is the vendor specific code, 4318 is nvidia graphics
#endif

#if ncnn_subgroupSize == 32
    // here is the code path optimized for subgroup_size == 32
#endif

#if ncnn_VK_KHR_shader_integer_dot_product && ncnn_shaderIntegerDotProduct && ncnn_integerDotProduct4x8BitPackedSignedAccelerated
    // here is the packed int8 dot-product path
#endif

#if ncnn_VK_KHR_cooperative_matrix
    // here is the KHR cooperative matrix path
#elif ncnn_VK_NV_cooperative_matrix
    // here is the NV cooperative matrix path
#endif

    // use macro definitions
    uint size; // dynamic value from some previous routines
    if (size < ncnn_subgroupSize)
    {
#if ncnn_supportedOperations & 4
        // subgroup support arithmetic
#endif

#if ncnn_subgroup_arithmetic
        // shorthand style for checking subgroup arithmetic :P
#endif
    }
}

Cooperative matrix shape and component-type combinations are selected on the host side. Use GpuInfo::support_cooperative_matrix(), GpuInfo::support_int8_cooperative_matrix(), GpuInfo::support_bf16_cooperative_matrix(), and GpuInfo::get_optimal_cooperative_matrix_mnk() before creating a cooperative matrix pipeline.

For signed int8 cooperative matrix kernels, ncnn requires signed int8 A/B and signed int32 accumulator/result cooperative matrix support at subgroup scope. The shader still uses the normal ncnn_VK_KHR_cooperative_matrix / ncnn_VK_NV_cooperative_matrix extension macros to select GLSL syntax, while the host selects this path with support_int8_cooperative_matrix().

In int8 cooperative matrix and integer dot-product shaders, prefer keeping pack4 data in the packed sint8vec4 representation and use i8buffer_sm4 for shared-memory staging when the layout is already packed. Use i8buffer_ld4 only when arithmetic needs unpacked ivec4 lanes.

validation layer macros

ncnn will define some additional convenient macros when the vulkan validation layer enabled

• ncnn_enable_validation_layer
• NCNN_LOGE

currently, you have to modify the ENABLE_VALIDATION_LAYER definition at the beginning of src/gpu.cpp to 1 to enable these macros.

The GL_EXT_debug_printf extension will be enabled automatically without explicitly specifying it in your code.

void main()
{
    int gx = int(gl_GlobalInvocationID.x);

#if ncnn_enable_validation_layer
    NCNN_LOGE("gx = %d\n", gx);
#endif
}

At runtime, NCNN_LOGE will print out the value of gx

option macros

enable glsl extension only if user enable some options

The GL_EXT_shader_16bit_storage extension will be automatically enabled without explicit code indication when the device supports 16-bit storage and the user turns on opt.use_fp16_storage, opt.use_bf16_storage, or opt.use_int16_storage

The GL_EXT_shader_explicit_arithmetic_types_float16 extension will be automatically enabled without explicit code indication when the device supports 16-bit arithmetic and the user turns on opt.use_fp16_arithmetic

The GL_EXT_shader_8bit_storage extension will be automatically enabled without explicit code indication when the device supports 8-bit storage and the user turns on opt.use_int8_storage

The GL_EXT_shader_explicit_arithmetic_types_int8 extension will be automatically enabled without explicit code indication when the device supports 8-bit arithmetic and the user turns on opt.use_int8_arithmetic

The GL_EXT_bfloat16 extension will be automatically enabled without explicit code indication when the device supports bfloat16 storage and the user turns on opt.use_bf16_storage

void main()
{
#if NCNN_fp16_storage
    // the user enable fp16 storage option and the device has fp16 storage support
#endif

#if NCNN_fp16_arithmetic
    // the user enable fp16 arithmetic option and the device has fp16 arithmetic support
#endif
}

Source: docs/developer-guide/glsl-extension.md