[llvm] [SPIR-V] Use mangled names for deducing builtin funcs arg types (PR #94255)

[Michal Paszkowski via llvm-commits]

https://github.com/michalpaszkowski created https://github.com/llvm/llvm-project/pull/94255

Before this change, the method for deducing function argument types assumed that any argument of untyped pointer type must be either: 
1) A pointer of an LLVM IR element type, passed byval/byref. 
2) An OpenCL/SPIR-V builtin type if there is spv_assign_type intrinsic assigning a TargetExtType.
3) Just a pointer (with default size).

This does not take into consideration builtin functions which might also have arguments of OpenCL/SPIR-V builtin type. Since builtins have just their prototypes inside a module (no body), no spv_assign_type intrinsics are generared for their arguments. Hence, a fourth option:

4) An OpenCL/SPIR-V builtin type if the mangled function name contains type information.

A test mimicking SPIR-V Translator behavior was added.

>From 7b36f94a55f11b71e46d8bc99ad50143308d7d30 Mon Sep 17 00:00:00 2001
From: Michal Paszkowski <michal at paszkowski.org>
Date: Mon, 3 Jun 2024 09:49:09 -0700
Subject: [PATCH] [SPIR-V] Use mangled names for deducing builtin funcs arg
 types

Before this change, the method for deducing function argument types
assumed that any argument of untyped pointer type must be either:
1) A pointer of an LLVM IR element type, passed byval/byref.
2) An OpenCL/SPIR-V builtin type if there is spv_assign_type intrinsic
   assigning a TargetExtType.
3) Just a pointer (with default size)

This does not take into consideration builtin functions which might
also have arguments of OpenCL/SPIR-V builtin type. Since builtins have
just their prototypes inside a module (no body), no spv_assign_type
intrinsics are generared for their arguments. Hence, a fourth option:
4) An OpenCL/SPIR-V builtin type if the mangled function name contains
   type information.

A test mimicking SPIR-V Translator behavior was added.
---
 llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp      | 16 ++++++++++++++--
 .../builtin-function-ptr-arg-of-builtin-type.ll  | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/SPIRV/pointers/builtin-function-ptr-arg-of-builtin-type.ll

diff --git a/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp b/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp
index f4daab7d06eb5..51b0bce46691e 100644
--- a/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp
@@ -211,12 +211,15 @@ static SPIRVType *getArgSPIRVType(const Function &F, unsigned ArgIdx,
         addressSpaceToStorageClass(getPointerAddressSpace(ArgType), ST));
   }
 
-  // In case OriginalArgType is of untyped pointer type, there are three
+  // In case OriginalArgType is of untyped pointer type, there are four
   // possibilities:
   // 1) This is a pointer of an LLVM IR element type, passed byval/byref.
   // 2) This is an OpenCL/SPIR-V builtin type if there is spv_assign_type
   //    intrinsic assigning a TargetExtType.
-  // 3) This is a pointer, try to retrieve pointer element type from a
+  // 3) This is an OpenCL/SPIR-V builtin type if the mangled function name
+  //    contains type information (the Arg's function is a builtin, has no
+  //    body).
+  // 4) This is a pointer, try to retrieve pointer element type from a
   // spv_assign_ptr_type intrinsic or otherwise use default pointer element
   // type.
   if (hasPointeeTypeAttr(Arg)) {
@@ -255,6 +258,15 @@ static SPIRVType *getArgSPIRVType(const Function &F, unsigned ArgIdx,
             cast<ConstantInt>(II->getOperand(2))->getZExtValue(), ST));
   }
 
+  std::string DemangledFuncName =
+      getOclOrSpirvBuiltinDemangledName(F.getName());
+  if (!DemangledFuncName.empty()) {
+    Type *BuiltinType = SPIRV::parseBuiltinCallArgumentBaseType(
+        DemangledFuncName, ArgIdx, F.getContext());
+    if (BuiltinType && BuiltinType->isTargetExtTy())
+      return GR->getOrCreateSPIRVType(BuiltinType, MIRBuilder, ArgAccessQual);
+  }
+
   // Replace PointerType with TypedPointerType to be able to map SPIR-V types to
   // LLVM types in a consistent manner
   if (isUntypedPointerTy(OriginalArgType)) {
diff --git a/llvm/test/CodeGen/SPIRV/pointers/builtin-function-ptr-arg-of-builtin-type.ll b/llvm/test/CodeGen/SPIRV/pointers/builtin-function-ptr-arg-of-builtin-type.ll
new file mode 100644
index 0000000000000..ff386affdc4c5
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/pointers/builtin-function-ptr-arg-of-builtin-type.ll
@@ -0,0 +1,15 @@
+; RUN: llc -O0 -mtriple=spirv64-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv64-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; CHECK-DAG: %[[#INT8:]] = OpTypeInt 8 0
+; CHECK-DAG: %[[#PTR_INT8:]] = OpTypePointer Function %[[#INT8]]
+; CHECK-DAG: %[[#EVENT:]] = OpTypeEvent
+; CHECK-DAG: %[[#FUNC_TY:]] = OpTypeFunction %[[#]] %[[#PTR_INT8]] %[[#PTR_INT8]] %[[#]] %[[#]] %[[#EVENT]]
+; CHECK-DAG: %[[#]] = OpFunction %[[#]] None %[[#FUNC_TY]]
+
+define spir_kernel void @foo(ptr %a, ptr %b) {
+  %call = call spir_func ptr @_Z29async_work_group_strided_copyPU3AS3hPU3AS1Khmm9ocl_event(ptr %a, ptr %b, i64 1, i64 1, ptr null)
+  ret void
+}
+
+declare spir_func ptr @_Z29async_work_group_strided_copyPU3AS3hPU3AS1Khmm9ocl_event(ptr, ptr, i64, i64, ptr)