type inference for atomic instructions

llvm · Jul 8, 2024 · 9612af4 · 9612af4
1 parent b6537c8
commit 9612af4
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 7 deletions.
diff --git a/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp b/llvm/lib/Target/SPIRV/SPIRVEmitIntrinsics.cpp
@@ -46,6 +46,10 @@
 using namespace llvm;
 
 namespace llvm {
+namespace SPIRV {
+#define GET_BuiltinGroup_DECL
+#include "SPIRVGenTables.inc"
+} // namespace SPIRV
 void initializeSPIRVEmitIntrinsicsPass(PassRegistry &);
 } // namespace llvm
 
@@ -661,6 +665,32 @@ void SPIRVEmitIntrinsics::deduceOperandElementType(Instruction *I) {
               KnownElemTy = ElemTy; // src will rewrite dest if both are defined
             Ops.push_back(std::make_pair(Op, i));
           }
+        } else if (Grp == SPIRV::Atomic || Grp == SPIRV::AtomicFloating) {
+          if (CI->arg_size() < 2)
+            return;
+          Value *Op = CI->getArgOperand(0);
+          if (!isPointerTy(Op->getType()))
+            return;
+          switch (Opcode) {
+          case SPIRV::OpAtomicLoad:
+          case SPIRV::OpAtomicCompareExchangeWeak:
+          case SPIRV::OpAtomicCompareExchange:
+          case SPIRV::OpAtomicExchange:
+          case SPIRV::OpAtomicIAdd:
+          case SPIRV::OpAtomicISub:
+          case SPIRV::OpAtomicOr:
+          case SPIRV::OpAtomicXor:
+          case SPIRV::OpAtomicAnd:
+          case SPIRV::OpAtomicUMin:
+          case SPIRV::OpAtomicUMax:
+          case SPIRV::OpAtomicSMin:
+          case SPIRV::OpAtomicSMax: {
+            KnownElemTy = getAtomicElemTy(GR, I, Op);
+            if (!KnownElemTy)
+              return;
+            Ops.push_back(std::make_pair(Op, 0));
+          } break;
+          }
         }
       }
     }

diff --git a/llvm/test/CodeGen/SPIRV/instructions/atomic.ll b/llvm/test/CodeGen/SPIRV/instructions/atomic.ll
@@ -17,6 +17,7 @@
 ; CHECK-DAG: [[I32Ty:%.*]] = OpTypeInt 32 0
 ; CHECK-DAG: [[PtrI32Ty:%.*]] = OpTypePointer Function [[I32Ty]]
 ; CHECK-DAG: [[I64Ty:%.*]] = OpTypeInt 64 0
+; CHECK-DAG: [[PtrI64Ty:%.*]] = OpTypePointer Generic [[I64Ty]]
 ;; Device scope is encoded with constant 1
 ; CHECK-DAG: [[SCOPE:%.*]] = OpConstant [[I32Ty]] 1
 ;; "monotonic" maps to the relaxed memory semantics, encoded with constant 0
@@ -131,13 +132,15 @@ define i32 @test_xor(i32* %ptr, i32 %val) {
 }
 
 ; CHECK: OpFunction
-; CHECK: [[Arg1:%.*]] = OpFunctionParameter
-; CHECK: [[Arg2:%.*]] = OpFunctionParameter
-; CHECK: OpAtomicSMin [[I64Ty]] %[[#]] [[SCOPE]] [[RELAXED]] [[Arg2]]
-; CHECK: OpAtomicSMax [[I64Ty]] %[[#]] [[SCOPE]] [[RELAXED]] [[Arg2]]
-; CHECK: OpAtomicUMin [[I64Ty]] %[[#]] [[SCOPE]] [[RELAXED]] [[Arg2]]
-; CHECK: OpAtomicUMax [[I64Ty]] %[[#]] [[SCOPE]] [[RELAXED]] [[Arg2]]
-; CHECK: OpFunctionEnd
+; CHECK-NEXT: [[Arg1:%.*]] = OpFunctionParameter [[PtrI64Ty]]
+; CHECK-NEXT: [[Arg2:%.*]] = OpFunctionParameter [[I64Ty]]
+; CHECK-NEXT: OpLabel
+; CHECK-NEXT: OpAtomicSMin [[I64Ty]] [[Arg1]] [[SCOPE]] [[RELAXED]] [[Arg2]]
+; CHECK-NEXT: OpAtomicSMax [[I64Ty]] [[Arg1]] [[SCOPE]] [[RELAXED]] [[Arg2]]
+; CHECK-NEXT: OpAtomicUMin [[I64Ty]] [[Arg1]] [[SCOPE]] [[RELAXED]] [[Arg2]]
+; CHECK-NEXT: OpAtomicUMax [[I64Ty]] [[Arg1]] [[SCOPE]] [[RELAXED]] [[Arg2]]
+; CHECK-NEXT: OpReturn
+; CHECK-NEXT: OpFunctionEnd
 define dso_local spir_kernel void @test_wrappers(ptr addrspace(4) %arg, i64 %val) {
   %r1 = call spir_func i64 @__spirv_AtomicSMin(ptr addrspace(4) %arg, i32 1, i32 0, i64 %val)
   %r2 = call spir_func i64 @__spirv_AtomicSMax(ptr addrspace(4) %arg, i32 1, i32 0, i64 %val)