[Backport to llvm_release_220] Update OCL nan builtin translation to work with bfloat type (#3638)

Backport of PR #3558 into `llvm_release_220`.

All commits applied cleanly.

Co-authored-by: Viktoria Maximova <[email protected]>

Author: github-actions[bot]

lib/SPIRV/SPIRVUtil.cpp

@@ -2660,7 +2660,7 @@ class OpenCLStdToSPIRVFriendlyIRMangleInfo : public BuiltinFuncMangleInfo {
 public:
   OpenCLStdToSPIRVFriendlyIRMangleInfo(OCLExtOpKind ExtOpId,
                                        ArrayRef<Type *> ArgTys, Type *RetTy)
-      : ExtOpId(ExtOpId), ArgTys(ArgTys) {
+      : ExtOpId(ExtOpId), ArgTys(ArgTys), RetTy(RetTy) {
 
     std::string Postfix = "";
     if (needRetTypePostfix())
@@ -2676,6 +2676,11 @@ class OpenCLStdToSPIRVFriendlyIRMangleInfo : public BuiltinFuncMangleInfo {
     case OpenCLLIB::Vloada_halfn:
     case OpenCLLIB::Vloadn:
       return true;
+    case OpenCLLIB::Nan:
+      // Only add return type mangling for bfloat16 to disambiguate from half
+      // (both are represented as i16 in LLVM). Float and half use traditional
+      // naming for backward compatibility.
+      return RetTy->getScalarType()->isBFloatTy();
     default:
       return false;
     }
@@ -2741,6 +2746,7 @@ class OpenCLStdToSPIRVFriendlyIRMangleInfo : public BuiltinFuncMangleInfo {
 private:
   OCLExtOpKind ExtOpId;
   ArrayRef<Type *> ArgTys;
+  Type *RetTy;
 };
 } // namespace
 

test/transcoding/OpenCL/nan_bfloat.ll

@@ -0,0 +1,35 @@
+; RUN: llvm-spirv %s --spirv-ext=+SPV_KHR_bfloat16 -o %t.spv
+; RUN: llvm-spirv %t.spv -to-text -o %t.spt
+; RUN: FileCheck < %t.spt %s --check-prefix=CHECK-SPIRV
+
+; RUN: llvm-spirv -r %t.spv --spirv-target-env=SPV-IR -o - | llvm-dis -o %t.rev.ll
+; RUN: FileCheck < %t.rev.ll %s --check-prefixes=CHECK-SPV-IR
+
+; Check OpenCL built-in nan translation.
+; Verify it's possible to distinguish between bfloat and half versions.
+
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
+target triple = "spir64"
+
+; CHECK-SPIRV: TypeFloat [[#BFLOAT:]] 16 0 {{$}}
+; CHECK-SPIRV: TypeFloat [[#HALF:]] 16 {{$}}
+; CHECK-SPIRV: ExtInst [[#BFLOAT]] [[#]] [[#]] nan
+; CHECK-SPIRV: ExtInst [[#HALF]] [[#]] [[#]] nan
+
+; CHECK-SPV-IR: call spir_func bfloat @_Z22__spirv_ocl_nan_RDF16bt(
+; CHECK-SPV-IR: call spir_func half @_Z15__spirv_ocl_nant(
+
+define dso_local spir_kernel void @test_bfloat(ptr addrspace(1) align 2 %a, i16 %b) {
+entry:
+  %call = tail call spir_func bfloat @_Z23__spirv_ocl_nan__RDF16bt(i16 %b)
+  %call2 = tail call spir_func half @_Z22__spirv_ocl_nan__Rhalft(i16 %b)
+  ret void
+}
+
+declare spir_func bfloat @_Z23__spirv_ocl_nan__RDF16bt(i16)
+declare spir_func half @_Z22__spirv_ocl_nan__Rhalft(i16)
+
+
+!opencl.ocl.version = !{!0}
+
+!0 = !{i32 3, i32 0}