[SPIR-V] Add atomic_flag builtin implementation

[mpaszkowski]

Summary

This change:

• Adds implementation details for the atomic_flag_test_and_set, atomic_flag_test_and_set_explicit, atomic_flag_clear, atomic_flag_clear_explicit builtins
• Adds an atomic_flag.ll test from the SPIR-V translator with additional checks

The atomic_flag.ll passes after this change.

[iliya-diyachkov]

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

634–647

It seems that this fragment is very similar to MemSemanticsReg evaluation in buildAtomicRMWInst(), so we could separate it into a static functionit and avoid the code duplication.

649–652

Move the if condition to the assertion.

654–664

This fragment also looks similar to getting ScopeRegister in buildAtomicRMWInst().

[arsenm]

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

649–652

Agree

llvm/test/CodeGen/SPIRV/transcoding/atomic_flag.ll

20

Should use opaque pointers in new tests

23

Use named values in tests

[mpaszkowski]

Thank you @iliya-diyachkov and @arsenm for your review! I apologize that this took so long, but I have now addressed all the comments in the new diff.

[arsenm]

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

438–440

Why is this something that would be materialized in a register instead of an immediate argument?

445

Don't need the 0

[mpaszkowski]

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

438–440

This is the way it is done by the Khronos LLVM SPIR-V Translator and we want the SPIR-V backend to be compatible.

Translator:

4 Constant 2 14 42 
...
7 AtomicUMin 2 30 4 12 13 14

Backend:

%10 = OpConstant %4 42 
...	
%30 = OpAtomicUMin %4 %18 %12 %13 %10

[iliya-diyachkov]

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

438–440

@arsenm, according to the SPIRV specification (see the link), the last operand of OpAtomicUMin instruction should be a register (<id> in the specification notation), not an immediate.

[iliya-diyachkov]

The patch looks good to me. Thank you for the implementation.

llvm/lib/Target/SPIRV/SPIRVBuiltins.cpp

 /// Helper function for building an atomic load instruction.
 static bool buildAtomicLoadInst(const SPIRV::IncomingCall *Call,
                                 MachineIRBuilder &MIRBuilder,
                                 SPIRVGlobalRegistry *GR) {
   Register PtrRegister = Call->Arguments[0];
   // TODO: if true insert call to __translate_ocl_memory_sccope before
   // OpAtomicLoad and the function implementation. We can use Translator's
   // output for transcoding/atomic_explicit_arguments.cl as an example.
   Register ScopeRegister;
   if (Call->Arguments.size() > 1)
     ScopeRegister = Call->Arguments[1];
   else
     ScopeRegister = buildConstantIntReg(SPIRV::Scope::Device, MIRBuilder, GR);
 
   Register MemSemanticsReg;
   if (Call->Arguments.size() > 2) {
     // TODO: Insert call to __translate_ocl_memory_order before OpAtomicLoad.
     MemSemanticsReg = Call->Arguments[2];
   } else {
     int Semantics =
         SPIRV::MemorySemantics::SequentiallyConsistent |
         getMemSemanticsForStorageClass(GR->getPointerStorageClass(PtrRegister));
     MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
   }
MIRBuilder.buildInstr(Opcode)
       .addUse(GR->getSPIRVTypeID(Call->ReturnType))
       .addUse(PtrRegister)
       .addUse(ScopeRegister)
       .addUse(MemSemanticsReg)
       .addUse(Call->Arguments[1]);
   return true;
 }
 
+/// Helper function for building atomic flag instructions (e.g.
+/// OpAtomicFlagTestAndSet).
+static bool buildAtomicFlagInst(const SPIRV::IncomingCall *Call,
+                                unsigned Opcode, MachineIRBuilder &MIRBuilder,
+                                SPIRVGlobalRegistry *GR) {
+  const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
+
+  Register PtrRegister = Call->Arguments[0];
+
+  Register MemSemanticsReg;
+  unsigned Semantics = SPIRV::MemorySemantics::SequentiallyConsistent;
+  if (Call->Arguments.size() >= 2) {
+    std::memory_order Order =
+        static_cast<std::memory_order>(getIConstVal(Call->Arguments[1], MRI));
+
+    Semantics =
+        getSPIRVMemSemantics(Order) |
+        getMemSemanticsForStorageClass(GR->getPointerStorageClass(PtrRegister));
+    if (Order == Semantics)
+      MemSemanticsReg = Call->Arguments[1];
+  }
+  if (!MemSemanticsReg.isValid())
+    MemSemanticsReg = buildConstantIntReg(Semantics, MIRBuilder, GR);
+
+  if (Opcode == SPIRV::OpAtomicFlagClear)
+    assert((Semantics != SPIRV::MemorySemantics::Acquire &&
+            Semantics != SPIRV::MemorySemantics::AcquireRelease) &&
+           "Invalid memory order argument!");
+
+  Register ScopeRegister;
+  SPIRV::Scope::Scope Scope = SPIRV::Scope::Device;
+  if (Call->Arguments.size() >= 3) {
+    auto CLScope = static_cast<SPIRV::CLMemoryScope>(
+        getIConstVal(Call->Arguments[2], MRI));
+    Scope = getSPIRVScope(CLScope);
+    if (CLScope == static_cast<unsigned>(Scope))
+      ScopeRegister = Call->Arguments[2];
+  }
+  if (!ScopeRegister.isValid())
+    ScopeRegister = buildConstantIntReg(Scope, MIRBuilder, GR);
+
+  auto MIB = MIRBuilder.buildInstr(Opcode);
+  if (Opcode == SPIRV::OpAtomicFlagTestAndSet)
+    MIB.addDef(Call->ReturnRegister)
+        .addUse(GR->getSPIRVTypeID(Call->ReturnType));
+
+  MIB.addUse(PtrRegister).addUse(ScopeRegister).addUse(MemSemanticsReg);
+  return true;
+}
+
 /// Helper function for building barriers, i.e., memory/control ordering
 /// operations.
 static bool buildBarrierInst(const SPIRV::IncomingCall *Call, unsigned Opcode,
                              MachineIRBuilder &MIRBuilder,
                              SPIRVGlobalRegistry *GR) {
   const MachineRegisterInfo *MRI = MIRBuilder.getMRI();
   unsigned MemFlags = getIConstVal(Call->Arguments[0], MRI);
   unsigned MemSemantics = SPIRV::MemorySemantics::None;
static bool generateAtomicInst(const SPIRV::IncomingCall *Call,
   case SPIRV::OpAtomicISub:
   case SPIRV::OpAtomicOr:
   case SPIRV::OpAtomicXor:
   case SPIRV::OpAtomicAnd:
   case SPIRV::OpAtomicExchange:
     return buildAtomicRMWInst(Call, Opcode, MIRBuilder, GR);
   case SPIRV::OpMemoryBarrier:
     return buildBarrierInst(Call, SPIRV::OpMemoryBarrier, MIRBuilder, GR);
+  case SPIRV::OpAtomicFlagTestAndSet:
+  case SPIRV::OpAtomicFlagClear:
+    return buildAtomicFlagInst(Call, Opcode, MIRBuilder, GR);
   default:
     return false;
   }
 }
 
 static bool generateBarrierInst(const SPIRV::IncomingCall *Call,
                                 MachineIRBuilder &MIRBuilder,
                                 SPIRVGlobalRegistry *GR) {

llvm/lib/Target/SPIRV/SPIRVBuiltins.td

 defm : DemangledNativeBuiltin<"atomic_fetch_or", OpenCL_std, Atomic, 2, 4, OpAtomicOr>;
 defm : DemangledNativeBuiltin<"atomic_fetch_xor", OpenCL_std, Atomic, 2, 4, OpAtomicXor>;
 defm : DemangledNativeBuiltin<"atomic_fetch_and", OpenCL_std, Atomic, 2, 4, OpAtomicAnd>;
 defm : DemangledNativeBuiltin<"atomic_fetch_add_explicit", OpenCL_std, Atomic, 4, 6, OpAtomicIAdd>;
 defm : DemangledNativeBuiltin<"atomic_fetch_sub_explicit", OpenCL_std, Atomic, 4, 6, OpAtomicISub>;
 defm : DemangledNativeBuiltin<"atomic_fetch_or_explicit", OpenCL_std, Atomic, 4, 6, OpAtomicOr>;
 defm : DemangledNativeBuiltin<"atomic_fetch_xor_explicit", OpenCL_std, Atomic, 4, 6, OpAtomicXor>;
 defm : DemangledNativeBuiltin<"atomic_fetch_and_explicit", OpenCL_std, Atomic, 4, 6, OpAtomicAnd>;
+defm : DemangledNativeBuiltin<"atomic_flag_test_and_set", OpenCL_std, Atomic, 1, 1, OpAtomicFlagTestAndSet>;
+defm : DemangledNativeBuiltin<"atomic_flag_test_and_set_explicit", OpenCL_std, Atomic, 2, 3, OpAtomicFlagTestAndSet>;
+defm : DemangledNativeBuiltin<"atomic_flag_clear", OpenCL_std, Atomic, 1, 1, OpAtomicFlagClear>;
+defm : DemangledNativeBuiltin<"atomic_flag_clear_explicit", OpenCL_std, Atomic, 2, 3, OpAtomicFlagClear>;
 
 // Barrier builtin records:
 defm : DemangledNativeBuiltin<"barrier", OpenCL_std, Barrier, 1, 3, OpControlBarrier>;
 defm : DemangledNativeBuiltin<"work_group_barrier", OpenCL_std, Barrier, 1, 3, OpControlBarrier>;
 
 // Kernel enqueue builtin records:
 defm : DemangledNativeBuiltin<"retain_event", OpenCL_std, Enqueue, 1, 1, OpRetainEvent>;
 defm : DemangledNativeBuiltin<"release_event", OpenCL_std, Enqueue, 1, 1, OpReleaseEvent>;

llvm/test/CodeGen/SPIRV/transcoding/atomic_flag.ll

+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+
+;; Types:
+; CHECK-DAG:  %[[#INT:]] = OpTypeInt 32
+; CHECK-DAG:  %[[#BOOL:]] = OpTypeBool
+; CHECK-DAG:  %[[#PTR:]] = OpTypePointer Generic %[[#]]
+;; Constants:
+; CHECK-DAG:  %[[#SEQ_CONSI_SEMAN:]] = OpConstant %[[#INT]] 16
+; CHECK-DAG:  %[[#DEVICE_SCOPE:]] = OpConstant %[[#INT]] 1
+; CHECK-DAG:  %[[#WORKGROUP_SCOPE:]] = OpConstant %[[#INT]] 2
+;; Instructions:
+; CHECK-DAG:  %[[#CAST:]] = OpPtrCastToGeneric %[[#PTR]] %[[#]]
+; CHECK:      %[[#]] = OpAtomicFlagTestAndSet %[[#BOOL]] %[[#CAST]] %[[#DEVICE_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+; CHECK:      %[[#]] = OpAtomicFlagTestAndSet %[[#BOOL]] %[[#CAST]] %[[#DEVICE_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+; CHECK:      %[[#]] = OpAtomicFlagTestAndSet %[[#BOOL]] %[[#CAST]] %[[#WORKGROUP_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+; CHECK:      OpAtomicFlagClear %[[#CAST]] %[[#DEVICE_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+; CHECK:      OpAtomicFlagClear %[[#CAST]] %[[#DEVICE_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+; CHECK:      OpAtomicFlagClear %[[#CAST]] %[[#WORKGROUP_SCOPE]] %[[#SEQ_CONSI_SEMAN]]
+
+define dso_local spir_kernel void @testAtomicFlag(i32 addrspace(1)* nocapture noundef %res) {
+entry:
+  %f = alloca i32, align 4
+  %0 = bitcast i32* %f to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* nonnull %0)
+  %f.ascast = addrspacecast i32* %f to i32 addrspace(4)*
+  %call = call spir_func zeroext i1 @_Z24atomic_flag_test_and_setPU3AS4VU7_Atomici(i32 addrspace(4)* noundef %f.ascast)
+  %conv = zext i1 %call to i32
+  store i32 %conv, i32 addrspace(1)* %res, align 4
+  %call2 = call spir_func zeroext i1 @_Z33atomic_flag_test_and_set_explicitPU3AS4VU7_Atomici12memory_order(i32 addrspace(4)* noundef %f.ascast, i32 noundef 5)
+  %conv3 = zext i1 %call2 to i32
+  %1 = load i32, i32 addrspace(1)* %res, align 4
+  %add = add nsw i32 %1, %conv3
+  store i32 %add, i32 addrspace(1)* %res, align 4
+  %call5 = call spir_func zeroext i1 @_Z33atomic_flag_test_and_set_explicitPU3AS4VU7_Atomici12memory_order12memory_scope(i32 addrspace(4)* noundef %f.ascast, i32 noundef 5, i32 noundef 1)
+  %conv6 = zext i1 %call5 to i32
+  %2 = load i32, i32 addrspace(1)* %res, align 4
+  %add7 = add nsw i32 %2, %conv6
+  store i32 %add7, i32 addrspace(1)* %res, align 4
+  call spir_func void @_Z17atomic_flag_clearPU3AS4VU7_Atomici(i32 addrspace(4)* noundef %f.ascast)
+  call spir_func void @_Z26atomic_flag_clear_explicitPU3AS4VU7_Atomici12memory_order(i32 addrspace(4)* noundef %f.ascast, i32 noundef 5)
+  call spir_func void @_Z26atomic_flag_clear_explicitPU3AS4VU7_Atomici12memory_order12memory_scope(i32 addrspace(4)* noundef %f.ascast, i32 noundef 5, i32 noundef 1)
+  call void @llvm.lifetime.end.p0i8(i64 4, i8* nonnull %0)
+  ret void
+}
+
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture)
+
+declare spir_func zeroext i1 @_Z24atomic_flag_test_and_setPU3AS4VU7_Atomici(i32 addrspace(4)* noundef)
+
+declare spir_func zeroext i1 @_Z33atomic_flag_test_and_set_explicitPU3AS4VU7_Atomici12memory_order(i32 addrspace(4)* noundef, i32 noundef)
+
+declare spir_func zeroext i1 @_Z33atomic_flag_test_and_set_explicitPU3AS4VU7_Atomici12memory_order12memory_scope(i32 addrspace(4)* noundef, i32 noundef, i32 noundef)
+
+declare spir_func void @_Z17atomic_flag_clearPU3AS4VU7_Atomici(i32 addrspace(4)* noundef)
+
+declare spir_func void @_Z26atomic_flag_clear_explicitPU3AS4VU7_Atomici12memory_order(i32 addrspace(4)* noundef, i32 noundef)
+
+declare spir_func void @_Z26atomic_flag_clear_explicitPU3AS4VU7_Atomici12memory_order12memory_scope(i32 addrspace(4)* noundef, i32 noundef, i32 noundef)
+
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture)