Repository: KhronosGroup/MoltenVK

Issue #231: Shared/threadgroup buffers break specialization constants
Created: 2018-08-29T22:20:58Z
Attempting to create a compute shader with a shared buffer generates incorrect MSL if the buffer is sized based on a specialization constant.
This GLSL
#version 450

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 1) const uint blockSize = 1;

shared float bucket[blockSize];

void main() {
    const uint tid = gl_LocalInvocationIndex;

    bucket[tid] = float(tid);

    memoryBarrierShared();
    barrier();
}

generates this SPIR-V
           OpCapability Shader
      %1 = OpExtInstImport "GLSL.std.450"
           OpMemoryModel Logical GLSL450
           OpEntryPoint GLCompute %4 "main" %10
           OpExecutionMode %4 LocalSize 1 1 1
           OpSource GLSL 450
           OpName %4 "main"
           OpName %8 "tid"
           OpName %10 "gl_LocalInvocationIndex"
           OpName %13 "blockSize"
           OpName %16 "bucket"
           OpDecorate %10 BuiltIn LocalInvocationIndex
           OpDecorate %13 SpecId 1
           OpDecorate %25 SpecId 0
           OpDecorate %27 BuiltIn WorkgroupSize
      %2 = OpTypeVoid
      %3 = OpTypeFunction %2
      %6 = OpTypeInt 32 0
      %7 = OpTypePointer Function %6
      %9 = OpTypePointer Input %6
     %10 = OpVariable %9 Input
     %12 = OpTypeFloat 32
     %13 = OpSpecConstant %6 1
     %14 = OpTypeArray %12 %13
     %15 = OpTypePointer Workgroup %14
     %16 = OpVariable %15 Workgroup
     %20 = OpTypePointer Workgroup %12
     %22 = OpConstant %6 1
     %23 = OpConstant %6 264
     %24 = OpConstant %6 2
     %25 = OpSpecConstant %6 1
     %26 = OpTypeVector %6 3
     %27 = OpSpecConstantComposite %26 %25 %22 %22
      %4 = OpFunction %2 None %3
      %5 = OpLabel
      %8 = OpVariable %7 Function
     %11 = OpLoad %6 %10
           OpStore %8 %11
     %17 = OpLoad %6 %8
     %18 = OpLoad %6 %8
     %19 = OpConvertUToF %12 %18
     %21 = OpAccessChain %20 %16 %17
           OpStore %21 %19
           OpMemoryBarrier %22 %23
           OpControlBarrier %24 %24 %23
           OpReturn
           OpFunctionEnd

which in turn generates incorrect MSL
#include <metal_stdlib>
#include <simd/simd.h>

using namespace metal;

constant uint _25_tmp [[function_constant(0)]];
constant uint _25 = is_function_constant_defined(_25_tmp) ? _25_tmp : 1u;
constant uint3 gl_WorkGroupSize = uint3(_25, 1u, 1u);

kernel void main0(uint gl_LocalInvocationIndex [[thread_index_in_threadgroup]])
{
    threadgroup float bucket[1u];
    uint tid = gl_LocalInvocationIndex;
    bucket[tid] = float(tid);
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

If instead I hardcode the buffer size
#version 450

layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 1) const uint blockSize = 1;

shared float bucket[256];

void main() {
    const uint tid = gl_LocalInvocationIndex;

    bucket[tid] = float(tid);

    memoryBarrierShared();
    barrier();
}

it generates this SPIR-V
           OpCapability Shader
      %1 = OpExtInstImport "GLSL.std.450"
           OpMemoryModel Logical GLSL450
           OpEntryPoint GLCompute %4 "main" %10
           OpExecutionMode %4 LocalSize 1 1 1
           OpSource GLSL 450
           OpName %4 "main"
           OpName %8 "tid"
           OpName %10 "gl_LocalInvocationIndex"
           OpName %16 "bucket"
           OpName %28 "blockSize"
           OpDecorate %10 BuiltIn LocalInvocationIndex
           OpDecorate %25 SpecId 0
           OpDecorate %27 BuiltIn WorkgroupSize
           OpDecorate %28 SpecId 1
      %2 = OpTypeVoid
      %3 = OpTypeFunction %2
      %6 = OpTypeInt 32 0
      %7 = OpTypePointer Function %6
      %9 = OpTypePointer Input %6
     %10 = OpVariable %9 Input
     %12 = OpTypeFloat 32
     %13 = OpConstant %6 256
     %14 = OpTypeArray %12 %13
     %15 = OpTypePointer Workgroup %14
     %16 = OpVariable %15 Workgroup
     %20 = OpTypePointer Workgroup %12
     %22 = OpConstant %6 1
     %23 = OpConstant %6 264
     %24 = OpConstant %6 2
     %25 = OpSpecConstant %6 1
     %26 = OpTypeVector %6 3
     %27 = OpSpecConstantComposite %26 %25 %22 %22
     %28 = OpSpecConstant %6 1
      %4 = OpFunction %2 None %3
      %5 = OpLabel
      %8 = OpVariable %7 Function
     %11 = OpLoad %6 %10
           OpStore %8 %11
     %17 = OpLoad %6 %8
     %18 = OpLoad %6 %8
     %19 = OpConvertUToF %12 %18
     %21 = OpAccessChain %20 %16 %17
           OpStore %21 %19
           OpMemoryBarrier %22 %23
           OpControlBarrier %24 %24 %23
           OpReturn
           OpFunctionEnd

which generates the expected correct MSL
#include <metal_stdlib>
#include <simd/simd.h>

using namespace metal;

constant uint _25_tmp [[function_constant(0)]];
constant uint _25 = is_function_constant_defined(_25_tmp) ? _25_tmp : 1u;
constant uint blockSize_tmp [[function_constant(1)]];
constant uint blockSize = is_function_constant_defined(blockSize_tmp) ? blockSize_tmp : 1u;
constant uint3 gl_WorkGroupSize = uint3(_25, 1u, 1u);

kernel void main0(uint gl_LocalInvocationIndex [[thread_index_in_threadgroup]])
{
    threadgroup float bucket[256];
    uint tid = gl_LocalInvocationIndex;
    bucket[tid] = float(tid);
    threadgroup_barrier(mem_flags::mem_threadgroup);
}

It seems the addition of a shared buffer that depends on the specialization constant causes it to be optimized out and use the default value.
For brevity, the pipeline is compiled with
uint32_t lims[] = {256, 256};
static const vk::SpecializationMapEntry mapEntries[] = {
    {0, 0*sizeof(uint32_t), sizeof(uint32_t)},
    {1, 1*sizeof(uint32_t), sizeof(uint32_t)}
};
const vk::SpecializationInfo specializationInfo(2, mapEntries, 2*sizeof(uint32_t), reinterpret_cast<void*>(lims));
pipeline = device.createComputePipeline({}, {{}, {{}, vk::ShaderStageFlagBits::eCompute, shader, "main", &specializationInfo}, pipelineLayout});

Duplicate issue #1423: Specialize shader functions by regenerating MSL from SPIR-V when appropriate
Created: 2021-08-14T15:04:06Z
Some Vulkan shader specializations are not supported in Metal:

Array length

Such specializations could be handled by dynamically regenerating the MSL code from SPIR-V to hardcode the specialized value in the MSL source code at pipeline creation time. This would be slower than using Metal's native specialization, so we'd only want to do it where Vulkan needs to perform a particular specialization in the list above.

Closing comment (2024-02-22T18:43:21Z): Looks like a duplicate of #231, no ?