Add handling for GEPs when trying to find TargetExtensionType of Opaque Pointer

This patch adds missing case of GetElementPtrInst when trying to figure out the TargetExtensionType of opaque pointer

Author: bokrzesi

IGC/Compiler/Optimizer/OpenCLPasses/JointMatrixFuncsResolutionPass/JointMatrixFuncsResolutionPass.cpp

@@ -1312,12 +1312,12 @@ Type *JointMatrixFuncsResolutionPass::TryFindTargetExtensionTypeOfOpaquePtr(Valu
       auto aiTy = ai->getAllocatedType();
       if (IGCLLVM::isTargetExtTy(aiTy))
         return aiTy;
-    } else if (auto *ci = dyn_cast<CallInst>(use)) {
-      auto funcReturnType = ci->getFunction()->getReturnType();
-      if (IGCLLVM::isTargetExtTy(funcReturnType))
-        return funcReturnType;
-    } else if (auto *spaceCast = dyn_cast<AddrSpaceCastInst>(use)) {
-      return TryFindTargetExtensionTypeOfOpaquePtr(spaceCast->getPointerOperand());
+    } else if (auto *cast = dyn_cast<CastInst>(use)) {
+      return TryFindTargetExtensionTypeOfOpaquePtr(cast->getOperand(0));
+    } else if (auto *gep = dyn_cast<GetElementPtrInst>(use)) {
+      auto gepTy = gep->getResultElementType();
+      if (IGCLLVM::isTargetExtTy(gepTy))
+        return gepTy;
     }
   }
 
@@ -1334,11 +1334,8 @@ Type *JointMatrixFuncsResolutionPass::TryFindTypeOfOpaquePtr(Value *V) {
     if (auto *ai = dyn_cast<AllocaInst>(use)) {
       auto aiTy = ai->getAllocatedType();
       return aiTy;
-    } else if (auto *ci = dyn_cast<CallInst>(use)) {
-      auto funcReturnType = ci->getFunction()->getReturnType();
-      return funcReturnType;
-    } else if (auto *spaceCast = dyn_cast<AddrSpaceCastInst>(use)) {
-      return TryFindTypeOfOpaquePtr(spaceCast->getPointerOperand());
+    } else if (auto *cast = dyn_cast<CastInst>(use)) {
+      return TryFindTypeOfOpaquePtr(cast->getOperand(0));
     } else if (auto *gep = dyn_cast<GetElementPtrInst>(use)) {
       return gep->getResultElementType();
     } else if (auto *bitcast = dyn_cast<BitCastInst>(use)) {
@@ -1470,7 +1467,10 @@ Instruction *JointMatrixFuncsResolutionPass::ResolvePrefetch(CallInst *CI) {
     ptrElemType = IGCLLVM::getNonOpaquePtrEltTy(ptrType);
   }
 
-  IGC_ASSERT_MESSAGE(ptrElemType, "Pointer type not found");
+  if (!ptrElemType) {
+    m_Ctx->EmitError("Pointer type not found when resolving prefetch", ptrVal);
+    return nullptr;
+  }
 
   if (StructType *structTy = dyn_cast<StructType>(ptrElemType)) {
     if (structTy->getNumElements() == 1) {
@@ -2288,8 +2288,11 @@ bool JointMatrixFuncsResolutionPass::preprocessAccessChain(Function *F) {
 #if LLVM_VERSION_MAJOR >= 16
     if (IGCLLVM::isOpaquePointerTy(operand0->getType())) {
       chainBaseTy = TryFindTargetExtensionTypeOfOpaquePtr(operand0);
-      IGC_ASSERT_MESSAGE(chainBaseTy, "__spirv_AccessChain call 1st argument must be "
-                                      "pointer to target extension type.");
+
+      if (!chainBaseTy) {
+        m_Ctx->EmitError("__spirv_AccessChain call 1st argument must be pointer to target extension type", operand0);
+        continue;
+      }
     } else
 #endif
     {

IGC/Compiler/tests/JointMatrixFuncsResolutionPass/access-chain-traverse-for-targetextensiontype.ll

@@ -0,0 +1,39 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; REQUIRES: llvm-16-plus
+
+; RUN: igc_opt --opaque-pointers -platformpvc -igc-joint-matrix-resolution -S 2>&1 < %s | FileCheck %s
+; ------------------------------------------------
+; JointMatrixFuncsResolutionPass
+; ------------------------------------------------
+; Walk thru uses in order to figure out the TET type
+
+%"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix.3" = type { target("spirv.CooperativeMatrixKHR", i16, 3, 8, 16, 0) }
+
+; CHECK:       define spir_kernel void @test()
+; CHECK:       %addressCast = addrspacecast ptr %gep3 to ptr addrspace(4)
+; CHECK:       %bitcast2 = bitcast ptr addrspace(4) %addressCast to ptr addrspace(4)
+; CHECK-NEXT:  %0 = load <8 x i16>, ptr addrspace(4) %bitcast2, align 8
+
+; Resolution flow: Call operand0 -> Bitcast -> Address Space Cast -> GEP -> TET
+
+; Function Attrs: nounwind
+define spir_kernel void @test() {
+entry:
+    %alloca = alloca target("spirv.CooperativeMatrixKHR", i16, 3, 8, 16, 0), align 8
+    %gep = getelementptr inbounds %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix.3", ptr %alloca, i64 0, i32 0
+    %addressCast = addrspacecast ptr %gep to ptr addrspace(4)
+    %bitcast2 = bitcast ptr addrspace(4) %addressCast to ptr addrspace(4)
+    %call = call spir_func ptr addrspace(4) @_Z19__spirv_AccessChainPU3AS4PU3AS144__spirv_CooperativeMatrixKHR__short_3_8_16_0l(ptr addrspace(4) %bitcast2, i64 0)
+
+    store i16 5, ptr addrspace(4) %call, align 2
+    ret void
+}
+
+; Function Attrs: nounwind
+declare spir_func ptr addrspace(4) @_Z19__spirv_AccessChainPU3AS4PU3AS144__spirv_CooperativeMatrixKHR__short_3_8_16_0l(ptr addrspace(4), i64) #0

IGC/Compiler/tests/JointMatrixFuncsResolutionPass/prefetch.ll

@@ -6,15 +6,14 @@
 ;
 ;============================ end_copyright_notice =============================
 ; REQUIRES: llvm-16-plus
-; RUN: igc_opt --opaque-pointers -igc-joint-matrix-resolution --platformpvc -S 2>&1 < %s | FileCheck %s
+; RUN: igc_opt --opaque-pointers -igc-joint-matrix-resolution --platformpvc -S 2>&1 < %s | FileCheck %s --implicit-check-not error:
 ; ------------------------------------------------
 ; Written based on IR generated from this test: llvm/sycl/test-e2e/Matrix/joint_matrix_prefetch.cpp
 ; The purpose of this test is to check whether we figure out
 ; the type of pointer (operand 0) during ResolvePrefetch() correctly
 ; ------------------------------------------------
 
 ; CHECK:     call void @__builtin_spriv_OpJointMatrixPrefetchINTEL_SG16_8x16_i16(ptr addrspace(4) %add.ptr
-; CHECK-NOT: error
 
 %"class.sycl::_V1::ext::oneapi::bfloat16" = type { i16 }
 ; Function Attrs: nounwind
@@ -27,5 +26,18 @@ define spir_kernel void @test(ptr addrspace(1) align 2 %A, i64 %mul1) {
     ret void
 }
 
+; CHECK:     define spir_kernel void @test2(ptr addrspace(4) align 2 %A) {
+; CHECK:     %bitcast2 = bitcast ptr addrspace(4) %gep to ptr addrspace(4)
+; CHECK:     call void @__builtin_spriv_OpJointMatrixPrefetchINTEL_SG16_8x32_i8(ptr addrspace(4) %bitcast2, i64 256, i32 4)
+
+; Comes from 'test-e2e/Matrix/Output/joint_matrix_bf16_fill_k_cache_prefetch.cpp'
+define spir_kernel void @test2(ptr addrspace(4) align 2 %A) {
+  entry:
+    %gep = getelementptr inbounds i8, ptr addrspace(4) %A, i64 1
+    %bitcast2 = bitcast ptr addrspace(4) %gep to ptr addrspace(4)
+    call spir_func void @"_Z38__spirv_CooperativeMatrixPrefetchINTELPU3AS438class.sycl::_V1::ext::oneapi::bfloat16iiiil"(ptr addrspace(4) %bitcast2, i32 8, i32 32, i32 0, i32 0, i64 256) #0
+    ret void
+}
+
 ; Function Attrs: nounwind
 declare spir_func void @"_Z38__spirv_CooperativeMatrixPrefetchINTELPU3AS438class.sycl::_V1::ext::oneapi::bfloat16iiiil"(ptr addrspace(4) %0, i32 %1, i32 %2, i32 %3, i32 %4, i64 %5) #0