Fix crashing when finding type of ocl_image* types

Author: arsenm

llvm/lib/Target/AMDGPU/AMDGPULibFunc.cpp

@@ -620,17 +620,17 @@ bool ItaniumParamParser::parseItaniumParam(StringRef& param,
   // parse type
   char const TC = param.front();
   if (isDigit(TC)) {
-    res.ArgType = StringSwitch<AMDGPULibFunc::EType>
-      (eatLengthPrefixedName(param))
-      .Case("ocl_image1darray" , AMDGPULibFunc::IMG1DA)
-      .Case("ocl_image1dbuffer", AMDGPULibFunc::IMG1DB)
-      .Case("ocl_image2darray" , AMDGPULibFunc::IMG2DA)
-      .Case("ocl_image1d"      , AMDGPULibFunc::IMG1D)
-      .Case("ocl_image2d"      , AMDGPULibFunc::IMG2D)
-      .Case("ocl_image3d"      , AMDGPULibFunc::IMG3D)
-      .Case("ocl_event"        , AMDGPULibFunc::DUMMY)
-      .Case("ocl_sampler"      , AMDGPULibFunc::DUMMY)
-      .Default(AMDGPULibFunc::DUMMY);
+    res.ArgType =
+        StringSwitch<AMDGPULibFunc::EType>(eatLengthPrefixedName(param))
+            .StartsWith("ocl_image1d_array", AMDGPULibFunc::IMG1DA)
+            .StartsWith("ocl_image1d_buffer", AMDGPULibFunc::IMG1DB)
+            .StartsWith("ocl_image2d_array", AMDGPULibFunc::IMG2DA)
+            .StartsWith("ocl_image1d", AMDGPULibFunc::IMG1D)
+            .StartsWith("ocl_image2d", AMDGPULibFunc::IMG2D)
+            .StartsWith("ocl_image3d", AMDGPULibFunc::IMG3D)
+            .Case("ocl_event", AMDGPULibFunc::DUMMY)
+            .Case("ocl_sampler", AMDGPULibFunc::DUMMY)
+            .Default(AMDGPULibFunc::DUMMY);
   } else {
     drop_front(param);
     switch (TC) {

llvm/test/CodeGen/AMDGPU/amdgpu-simplify-libcall-image-function-signatures.ll

@@ -0,0 +1,65 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=amdgpu-simplifylib %s | FileCheck %s
+
+; Make sure we can produce a valid FunctionType for the expected
+; signature of image functions.
+
+declare i32 @_Z16get_image_height20ocl_image2d_depth_rw(ptr addrspace(4))
+
+define i32 @call_ocl_image2d_depth(ptr addrspace(4) %img) {
+; CHECK-LABEL: define i32 @call_ocl_image2d_depth(
+; CHECK-SAME: ptr addrspace(4) [[IMG:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @_Z16get_image_height20ocl_image2d_depth_rw(ptr addrspace(4) [[IMG]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %result = call i32 @_Z16get_image_height20ocl_image2d_depth_rw(ptr addrspace(4) %img)
+  ret i32 %result
+}
+
+declare i32 @_Z15get_image_width14ocl_image3d_ro(ptr addrspace(4))
+
+define i32 @call_ocl_image3d_depth(ptr addrspace(4) %img) {
+; CHECK-LABEL: define i32 @call_ocl_image3d_depth(
+; CHECK-SAME: ptr addrspace(4) [[IMG:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @_Z15get_image_width14ocl_image3d_ro(ptr addrspace(4) [[IMG]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %result = call i32 @_Z15get_image_width14ocl_image3d_ro(ptr addrspace(4) %img)
+  ret i32 %result
+}
+
+declare i32 @_Z15get_image_width14ocl_image1d_ro(ptr addrspace(4))
+
+define i32 @call_get_image_width14ocl_image1d_ro(ptr addrspace(4) %img) {
+; CHECK-LABEL: define i32 @call_get_image_width14ocl_image1d_ro(
+; CHECK-SAME: ptr addrspace(4) [[IMG:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @_Z15get_image_width14ocl_image1d_ro(ptr addrspace(4) [[IMG]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %result = call i32 @_Z15get_image_width14ocl_image1d_ro(ptr addrspace(4) %img)
+  ret i32 %result
+}
+
+declare <2 x i32> @_Z13get_image_dim20ocl_image2d_array_ro(ptr addrspace(4))
+
+define <2 x i32> @call_Z13get_image_dim20ocl_image2d_array_ro(ptr addrspace(4) %img) {
+; CHECK-LABEL: define <2 x i32> @call_Z13get_image_dim20ocl_image2d_array_ro(
+; CHECK-SAME: ptr addrspace(4) [[IMG:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call <2 x i32> @_Z13get_image_dim20ocl_image2d_array_ro(ptr addrspace(4) [[IMG]])
+; CHECK-NEXT:    ret <2 x i32> [[RESULT]]
+;
+  %result = call <2 x i32> @_Z13get_image_dim20ocl_image2d_array_ro(ptr addrspace(4) %img)
+  ret <2 x i32> %result
+}
+
+declare i32 @_Z15get_image_width20ocl_image1d_array_ro(ptr addrspace(4))
+
+define i32 @call_Z15get_image_width20ocl_image1d_array_ro(ptr addrspace(4) %img) {
+; CHECK-LABEL: define i32 @call_Z15get_image_width20ocl_image1d_array_ro(
+; CHECK-SAME: ptr addrspace(4) [[IMG:%.*]]) {
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @_Z15get_image_width20ocl_image1d_array_ro(ptr addrspace(4) [[IMG]])
+; CHECK-NEXT:    ret i32 [[RESULT]]
+;
+  %result = call i32 @_Z15get_image_width20ocl_image1d_array_ro(ptr addrspace(4) %img)
+  ret i32 %result
+}