[Bug][Vulkan] SPIR-V backend generates sequentially consistant OpControlBarriers and OpMemberDecorate DecorationOffset on non-interface blocks

[kistenklaus]

Expected behavior

I was trying to get TVM to work, just to benchmark it against other frameworks. Specifically i am interessted in DNN compilers, which compile down networks down to the instruction level.
That's why i was interessted in TVM, but it didn't work out the box as expected, even for very simple networks, that only contain a convolution and a relu layer.

I also think that i already found the bug, later more (last section)

Actual behavior

Network didn't compile, see error log later.

Environment

System:

OS: Manjaro Linux x86_64 
Host: B650I AORUS ULTRA -CF 
Kernel: 6.6.126-1-MANJARO 
Uptime: 51 mins 
Packages: 1948 (pacman) 
Shell: bash 5.3.9 
Resolution: 1920x1080 
WM: i3 
Theme: Breeze [GTK2/3] 
Icons: breeze [GTK2/3] 
Terminal: Neovim Terminal 
CPU: AMD Ryzen 7 7800X3D (16) @ 5.050GHz 
GPU: NVIDIA GeForce RTX 4070 
GPU: AMD ATI 10:00.0 Raphael 
Memory: 16369MiB / 31200MiB               

TVM directly from main (c9fb8cd)

Steps to reproduce

build with USE_VULKAN=ON and set(USE_KHRONOS_SPIRV /usr).
Then run this script:

import tvm
from tvm import relax
from tvm.relax.frontend import nn
from tvm.s_tir import dlight as dl


class TinyConv(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2D(
            in_channels=3,
            out_channels=32,
            kernel_size=3,
            stride=1,
            padding=1,
            bias=True,
            dtype="float16",
            data_layout="NCHW",
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.conv(x))


def main():
    input_shape = (1, 3, 1088, 1920)
    input_dtype = "float16"

    model = TinyConv()

    mod, _ = model.export_tvm(
        {"forward": {"x": nn.spec.Tensor(input_shape, input_dtype)}}
    )

    dev = tvm.vulkan(0)
    print("Vulkan exists:", tvm.vulkan().exist)
    print("Device:", dev)
    if not tvm.vulkan().exist:
        raise RuntimeError("TVM does not see a Vulkan device")

    target = tvm.target.Target.from_device(dev)

    with target:
        mod = tvm.ir.transform.Sequential(
            [
                relax.get_pipeline("zero"),
                dl.ApplyDefaultSchedule(
                    dl.gpu.Matmul(),
                    dl.gpu.GEMV(),
                    dl.gpu.Reduction(),
                    dl.gpu.GeneralReduction(),
                    dl.gpu.Fallback(),
                ),
            ]
        )(mod)

    print("Starting compile...")
    ex = relax.build(mod, target)
    print("Compile OK")
    print(type(ex))


if __name__ == "__main__":
    main()

This fails with:

[kistenklaus@kiste tvm-fork]$ python bench.py 
Vulkan exists: True
Device: vulkan:0
Starting compile...
Traceback (most recent call last):
  File "/home/kistenklaus/Documents/tvm-fork/bench.py", line 65, in <module>
    main()
  File "/home/kistenklaus/Documents/tvm-fork/bench.py", line 59, in main
    ex = relax.build(mod, target)
         ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kistenklaus/Documents/tvm-fork/.venv/lib/python3.11/site-packages/
tvm/relax/vm_build.py", line 262, in build
    return _vmlink(
           ^^^^^^^^
  File "/home/kistenklaus/Documents/tvm-fork/.venv/lib/python3.11/site-packages/tvm/relax/vm_build.py", line 157, in _vmlink
    lib = tvm.tir.build(tir_mod, target=target, pipeline=tir_pipeline)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kistenklaus/Documents/tvm-fork/.venv/lib/python3.11/site-packages/tvm/tir/build.py", line 238, in build
    return tir_to_runtime(host_mod, device_mod_dict, target_host)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kistenklaus/Documents/tvm-fork/.venv/lib/python3.11/site-packages/tvm/tir/build.py", line 146, in tir_to_runtime
    device_modules.append(codegen_build(device_mod, target))
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kistenklaus/Documents/tvm-fork/.venv/lib/python3.11/site-packages/tvm/tir/build.py", line 130, in codegen_build
    return bf(mod, target)
           ^^^^^^^^^^^^^^^
  File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
  File "<unknown>", line 0, in tvm::codegen::BuildSPIRV(tvm::IRModule, tvm::Target)
  File "<unknown>", line 0, in tvm::codegen::LowerToSPIRV[abi:cxx11](tvm::IRModule, tvm::Target)
  File "/home/kistenklaus/Documents/tvm-fork/src/target/spirv/spirv_utils.cc", line 101, in void tvm::codegen::SPIRVTools::ValidateShader(const std::vecto
r<unsigned int>&)
    TVM_FFI_ICHECK_EQ(res, SPV_SUCCESS)

tvm.error.InternalError: Check failed: res == SPV_SUCCESS (-14 vs. 0) :  index=270 error:[VUID-StandaloneSpirv-MemorySemantics-10866] ControlBarrier: Memo
ry Semantics with SequentiallyConsistent memory order must not be used in the Vulkan API
  OpControlBarrier %int_2 %int_2 %int_272

Possible Layman Fix

I grepped through the error messages a bit and the only line where spv::MemorySemanticsSequentiallyConsistentMask is used is in src/target/spirv/codegen_spirv.cc.
In the following function:

spirv::Value CodeGenSPIRV::CreateStorageSync(const CallNode* op) {
  const std::string& sync = op->args[0].as<StringImmNode>()->value;
  spirv::Value value;

  uint32_t vulkan_api_version = spirv_support_.vulkan_api_version;

  int64_t sync_scope;
  int64_t memory_semantics = spv::MemorySemanticsSequentiallyConsistentMask;
  if ((sync == "warp") && (vulkan_api_version >= VK_API_VERSION_1_1)) {
    // Synchronize control at the Subgroup level, but memory at the
    // Workgroup level.  This is because different invocations in a
    // subgroup may have each modified memory that exists at the
    // workgroup scope.  This should be changed if/when tir exposes
    // more information as to which memory access needs to be
    // synchronized.
    sync_scope = spv::ScopeSubgroup;
    memory_semantics |=
        spv::MemorySemanticsSubgroupMemoryMask | spv::MemorySemanticsWorkgroupMemoryMask;

  } else if ((sync == "shared") || (sync == "warp")) {
    sync_scope = spv::ScopeWorkgroup;
    memory_semantics |= spv::MemorySemanticsWorkgroupMemoryMask;
  } else {
    TVM_FFI_THROW(InternalError) << "Do not support sync " << sync;
  }

  auto type_int = builder_->GetSType(DataType::Int(32));
  builder_->MakeInst(spv::OpControlBarrier, builder_->IntImm(type_int, sync_scope),
                     builder_->IntImm(type_int, sync_scope),
                     builder_->IntImm(type_int, memory_semantics));

  return value;
}

Now iam not very familiar with the rest of the TVM codebase, so i might be wrong and this might break correctness somewhere, but for vulkan a normal workgroup or subgroup barrier has acquire release semantics not sequentially consistant.
My suspicion would be that this came from some code related to OpenCL, because OpenCL requires OpControlBarriers which are sequentially consistant (no idea why, doesn't make any sense to me).

I tried to verify this in the SPIR-V spec, but they are broken yet again =^(.
Regarless iam quite confident that sequential consistancy within a vulkan context for OpControlBarrier is a bug.

Maybe to further motivate this claim https://github.com/KhronosGroup/GLSL/blob/main/extensions/khr/GL_KHR_memory_scope_semantics.txt also defaults to acquire release semantics for barrier();

This doesn't make it run out of the box, but i think i found another bug, but i will split this up into another issue.

Once i get a reply confirming that this is a bug, i can create a quick PR fixing it.

Triage

• needs-triage
• actionable
• backend:vulkan

[kistenklaus]

Iam to lazy to create another issue for this, after changing sequential consistant to acquire release semantics, i get another error, which is:

tvm.error.InternalError: Check failed: res == SPV_SUCCESS (-10 vs. 0) :  index=78 error:[VUID-StandaloneSpirv-None-10684] Invalid explicit layout decorati
ons on type for operand '36[%_ptr_Workgroup__struct_35]', the Workgroup storage class has a explicit layout from the Offset decoration.
  %pad_temp_reindex_pad_shared = OpVariable %_ptr_Workgroup__struct_35 Workgroup

It took me a while to parse this, but i think what's happening is that in src/target/spirv/ir_builder.cc in the function:

SType IRBuilder::GetStructArrayType(const SType& value_type, uint32_t num_elems,
                                    bool interface_block) {
  auto key = std::make_tuple(value_type.id, num_elems, interface_block);
  auto it = struct_array_type_tbl_.find(key);
  if (it != struct_array_type_tbl_.end()) {
    return it->second;
  }

  SType arr_type;
  arr_type.id = id_counter_++;
  arr_type.type = DataType::Handle();
  arr_type.element_type_id = value_type.id;

  if (num_elems != 0) {
    Value length = UIntImm(GetSType(DataType::UInt(32)), num_elems);
    ib_.Begin(spv::OpTypeArray).AddSeq(arr_type, value_type, length).Commit(&global_);
  } else {
    ib_.Begin(spv::OpTypeRuntimeArray).AddSeq(arr_type, value_type).Commit(&global_);
  }
  int nbits = value_type.type.bits() * value_type.type.lanes();
  TVM_FFI_ICHECK_EQ(nbits % 8, 0);
  uint32_t nbytes = static_cast<uint32_t>(nbits) / 8;
  // decorate the array type.
  this->Decorate(spv::OpDecorate, arr_type, spv::DecorationArrayStride, nbytes);
  // declare struct of array
  SType struct_type;
  struct_type.id = id_counter_++;
  struct_type.type = DataType::Handle();
  struct_type.element_type_id = value_type.id;
  ib_.Begin(spv::OpTypeStruct).AddSeq(struct_type, arr_type).Commit(&global_);
  // decorate the array type.
  ib_.Begin(spv::OpMemberDecorate)
      .AddSeq(struct_type, 0, spv::DecorationOffset, 0)
      .Commit(&decorate_);

  if (interface_block) {
    // Runtime array are always decorated as Block or BufferBlock
    // (shader storage buffer)
    if (spirv_support_.supports_storage_buffer_storage_class) {
      // If SPIRV 1.3+, or with extension
      // SPV_KHR_storage_buffer_storage_class, BufferBlock is
      // deprecated.
      extensions_used_.insert("SPV_KHR_storage_buffer_storage_class");
      this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBlock);
    } else {
      if (num_elems == 0) {
        this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBufferBlock);
      }
    }
  }
  struct_array_type_tbl_[key] = struct_type;
  return struct_type;
}

We add OpMemberDecorate DecorateOffset to non interface blocks, which is not valid, moving it inside the if statement fixes it for me.

Now if somebody has a working spir-v specs running locally it's probably worth verifying, but iam to lazy to build the spec right now.

Again once i get a reply, i will create a PR.

[kistenklaus]

#18914 seems to fix the second bug. Could not find a PR for the sequentially consistant bug.