JointMatrix canonicalized GEP patching for opaque pointers

This commit revisits functionality introduced by b74e645.
Previous solution didn't work on llvm16 + opaque pointers.
Additionally it assumed specific bitcast -> GEP pattern which wasn't
guaranteed to work in the future if SYCL codegen changed.
New algorithm does recursive search on users of AllocaInst
to find GEP and lifetime instructions that need fixing.

Author: fda0

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

@@ -2764,36 +2764,6 @@ void JointMatrixFuncsResolutionPass::visitCallInst(CallInst &CI) {
     return;
   }
 
-  // Update size of allocated element in llvm.lifetime.start/end intrincics
-  if (auto II = dyn_cast<IntrinsicInst>(&CI)) {
-    if (II->getIntrinsicID() == Intrinsic::lifetime_start || II->getIntrinsicID() == Intrinsic::lifetime_end) {
-
-      // track pointer operand to alloca instr
-      auto &DL = CI.getModule()->getDataLayout();
-      Value *obj = IGCLLVM::getUnderlyingObject(II->getOperand(1), DL);
-
-      if (AllocaInst *AI = dyn_cast_or_null<AllocaInst>(obj)) {
-        // if alloca requires resolving, resolve alloca, otherwise do not touch intrinsic
-        // as it is not related to Joint Matrix type
-        if (!isOrContainsMatrixType(AI->getAllocatedType()))
-          return;
-
-        ResolveSIMDSize(CI.getParent()->getParent());
-        AllocaInst *NAI = cast<AllocaInst>(Resolve(AI));
-        auto allocaSizeInBits = IGCLLVM::makeOptional(NAI->getAllocationSizeInBits(DL));
-        if (!allocaSizeInBits.has_value())
-          return;
-        uint64_t newSize = (uint64_t)(allocaSizeInBits.value() / 8);
-
-        // update first argument, if it is constant int
-        if (auto *ConstInt = dyn_cast<ConstantInt>(CI.getOperand(0))) {
-          CI.setOperand(0, ConstantInt::get(ConstInt->getType(), newSize));
-          LLVM_DEBUG(dbgs() << "   -- UPDATED CALL: " << CI << "\n");
-        }
-      }
-    }
-  }
-
   if (isAnyFunctionArgMatrixType(func))
     ResolveCallFuncWithMatrixArgs(ResolvedFuncSignatures[func], &CI);
 }
@@ -2879,6 +2849,50 @@ DIType *getOrCreateType(Type *T, Module *M) {
   return diType;
 }
 
+void JointMatrixFuncsResolutionPass::RecursiveSearchAndFixCanonicalizdGEPandLifetime(
+    std::unordered_set<llvm::Value *> &visited, const DataLayout &DL, Value *root, uint64_t matrixTypeAllocSize,
+    uint64_t totalAllocationSize) {
+  auto insertToVisited = visited.insert(root);
+  if (!insertToVisited.second) // root was already visited
+    return;
+
+  // Depth first recursive traversal of root users
+  for (auto U : root->users()) {
+    // Only traverse children nodes if current node is a cast or a PHI node.
+    if (isa<CastInst>(U) || isa<PHINode>(U)) {
+      LLVM_DEBUG(dbgs() << "DFS: visiting users of " << *U << '\n');
+      RecursiveSearchAndFixCanonicalizdGEPandLifetime(visited, DL, U, matrixTypeAllocSize, totalAllocationSize);
+      continue;
+    }
+
+    if (auto GEP = dyn_cast<GetElementPtrInst>(U)) {
+      // Update canonicalized i8 GEP:
+      //   getelementptr i8, ptr %x, i64 <const>
+      if (GEP->getSourceElementType()->isIntegerTy(8) && GEP->hasAllConstantIndices() && GEP->getNumIndices() == 1) {
+        LLVM_DEBUG(dbgs() << "Found canonicalized i8 GEP: " << *GEP << "\n");
+        auto offset = cast<ConstantInt>(GEP->getOperand(1));
+        uint64_t pointerSize = DL.getPointerSizeInBits(GEP->getPointerAddressSpace()) / 8;
+        uint64_t offsetInElements = offset->getZExtValue() / pointerSize;
+        uint64_t correctOffset = offsetInElements * matrixTypeAllocSize;
+        ConstantInt *newOffsetConstant = ConstantInt::get(offset->getType(), correctOffset);
+        GEP->setOperand(1, newOffsetConstant);
+        LLVM_DEBUG(dbgs().indent(2) << "Fixed index: " << *GEP << "\n");
+      }
+    } else if (auto II = dyn_cast<IntrinsicInst>(U)) {
+      // Update size for lifetime intrinsics
+      if (II->getIntrinsicID() == Intrinsic::lifetime_start || II->getIntrinsicID() == Intrinsic::lifetime_end) {
+        if (auto constInt = dyn_cast<ConstantInt>(II->getOperand(0))) {
+          LLVM_DEBUG(dbgs() << "Found lifetime intrinsic for joint matrix array allocation" << *II << '\n');
+          II->setOperand(0, ConstantInt::get(constInt->getType(), totalAllocationSize));
+          LLVM_DEBUG(dbgs().indent(2) << "Fixed size: " << *II << "\n");
+        }
+      }
+    } else {
+      LLVM_DEBUG(dbgs() << "Skipping joint matrix array alloca user: " << *U << '\n');
+    }
+  }
+}
+
 void JointMatrixFuncsResolutionPass::visitAllocaInst(AllocaInst &I) {
   LLVM_DEBUG(dbgs() << " - VISIT: " << I << "\n");
 
@@ -2890,9 +2904,10 @@ void JointMatrixFuncsResolutionPass::visitAllocaInst(AllocaInst &I) {
 
   ResolveSIMDSize(I.getParent()->getParent());
 
-  Value *newInst = ResolveGeneric(&I);
+  AllocaInst *newInst = cast<AllocaInst>(ResolveGeneric(&I));
 
-  if (newInst) {
+  // update debug info
+  {
     TinyPtrVector<DbgDeclareInst *> DDIs;
     for (DbgVariableIntrinsic *DVI : FindDbgAddrUses(&I))
       if (auto *DDI = dyn_cast<DbgDeclareInst>(DVI))
@@ -2913,6 +2928,20 @@ void JointMatrixFuncsResolutionPass::visitAllocaInst(AllocaInst &I) {
       ddi->eraseFromParent();
     }
   }
+
+  // update GEPs and lifetime intrinsics
+  {
+    Type *unresolvedMatrixType = getContainedMatrixType(I.getAllocatedType());
+    Type *resolvedMatrixType = ResolveTypes(unresolvedMatrixType);
+
+    const DataLayout &DL = newInst->getModule()->getDataLayout();
+    uint64_t matrixTypeSize = DL.getTypeAllocSize(resolvedMatrixType);
+    uint64_t totalAllocSize =
+        IGCLLVM::makeOptional(newInst->getAllocationSizeInBits(DL)).value_or(TypeSize(0, false)) / 8;
+    // We have to use old alloca instruction I because its uses weren't replaced by newInst yet.
+    std::unordered_set<Value *> visited;
+    RecursiveSearchAndFixCanonicalizdGEPandLifetime(visited, DL, &I, matrixTypeSize, totalAllocSize);
+  }
 }
 
 void JointMatrixFuncsResolutionPass::visitAddrSpaceCastInst(AddrSpaceCastInst &I) {
@@ -2969,46 +2998,6 @@ void JointMatrixFuncsResolutionPass::visitGetElementPtrInst(GetElementPtrInst &G
     return;
 
   Type *GEPEltType = GEP.getSourceElementType();
-
-  // After constant GEPs are canonicalized to i8 types, we may get patterns like below:
-  //
-  // %8 = bitcast [4 x [4 x %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix"]]* %tC.i to i8*
-  // %arrayctor.end.i = getelementptr inbounds i8, i8* %8, i64 128
-  //
-  // It is not correct, because
-  // original offset was 16 elements of matrix type. Matrix type before resolution is represented as pointer
-  // Pointer is typically 8 bytes, hence offset of 128 bytes is calculated as 16 x 8 = 128
-  // The real offset would be 16 matrix types after resolution, not pointer types.
-  // So to fix the offset, we need calculate the offset in matrix type, taking into account pointer type size
-  // Then we need calculate real matrix type size after resolution in bytes
-  // Then real offset in bytes will be multiplicaiton of offset in matrix types and size of matrix type in bytes
-  //
-  // For example, if matrix type was resolved like that:
-  // %"struct.sycl::_V1::ext::oneapi::experimental::matrix::joint_matrix.resolved" = type { <8 x float> }
-  // offset will be 16 * (8 * 4) = 512:
-  //
-  // %arrayctor.end.i = getelementptr inbounds i8, i8* %8, i64 512
-  BitCastInst *BC = dyn_cast<BitCastInst>(GEP.getOperand(0));
-
-  if (GEPEltType->isIntegerTy(8) && BC && (GEP.getNumIndices() == 1) && GEP.hasAllConstantIndices()) {
-    if (Type *BCSrcTy = BC->getSrcTy(); BCSrcTy->isPointerTy()) {
-      if (Type *unresolvedMatTy = getContainedMatrixType(BCSrcTy)) {
-
-        // Calculate offset based on matrix type
-        ConstantInt *index = cast<ConstantInt>(GEP.getOperand(1));
-        auto &DL = GEP.getModule()->getDataLayout();
-        uint64_t pointerSizeInBytes = DL.getPointerSizeInBits(GEP.getPointerAddressSpace()) / 8;
-        uint64_t offsetInElements = index->getZExtValue() / pointerSizeInBytes;
-
-        // Calculate correct offset in bytes and update GEP
-        uint64_t elementSize = (uint64_t)DL.getTypeAllocSize(ResolveTypes(unresolvedMatTy));
-        uint64_t correctOffset = offsetInElements * elementSize;
-        GEP.idx_begin()->set(ConstantInt::get(index->getType(), correctOffset));
-        return;
-      }
-    }
-  }
-
   if (!isOrContainsMatrixType(GEPEltType))
     return;
 

IGC/Compiler/Optimizer/OpenCLPasses/JointMatrixFuncsResolutionPass/JointMatrixFuncsResolutionPass.h

@@ -84,6 +84,9 @@ class JointMatrixFuncsResolutionPass final : public llvm::ModulePass,
                                          IGC::JointMatrixTypeDescription *outDescription);
 #endif
 
+  void RecursiveSearchAndFixCanonicalizdGEPandLifetime(std::unordered_set<llvm::Value *> &visited,
+                                                       const llvm::DataLayout &DL, llvm::Value *rootValue,
+                                                       uint64_t matrixTypeAllocSize, uint64_t totalAllocationSize);
   llvm::StringRef GetMatrixTypeName(llvm::Type *opaqueType);
   bool SetLayoutFromUse(IGC::JointMatrixTypeDescription *outDescription);
   unsigned GetUseFromLegacyLayout(unsigned int legacyLayout);

IGC/Compiler/tests/JointMatrixFuncsResolutionPass/offset_correction-typed-pointers-targetextensiontype.ll

@@ -30,8 +30,9 @@ target triple = "spir64-unknown-unknown"
 ; CHECK-LABEL: define spir_kernel void @test(
 ; CHECK-SAME: i64 [[OFFSET:%.*]], i8 addrspace(1)* [[PTR:%.*]], i8* [[PTR1:%.*]]) {
 define spir_kernel void @test(i64 %offset, i8 addrspace(1)* %ptr, i8* %ptr1) {
+entry:
 
-; CHECK-NEXT:    [[TC_I1:%.*]] = alloca [4 x [4 x %"struct::joint_matrix::C.resolved"]], align 8
+; CHECK:         [[TC_I1:%.*]] = alloca [4 x [4 x %"struct::joint_matrix::C.resolved"]], align 8
 ; CHECK-NEXT:    [[TA_I3:%.*]] = alloca [4 x [2 x %"struct::joint_matrix::A.resolved"]], align 8
 ; CHECK-NEXT:    [[TB_I5:%.*]] = alloca [4 x [2 x %"struct::joint_matrix::B.resolved"]], align 8
 ; CHECK-NEXT:    [[TI_I:%.*]] = alloca [4 x [4 x %"struct::almost_joint_matrix"]], align 8
@@ -49,74 +50,89 @@ define spir_kernel void @test(i64 %offset, i8 addrspace(1)* %ptr, i8* %ptr1) {
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 128, i8* [[TMP2]])
 ; CHECK-NEXT:    [[TMP3:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::B.resolved"]]* [[TB_I5]] to i8*
 ; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 256, i8* [[TMP3]])
-  %1 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
-  call void @llvm.lifetime.start.p0i8(i64 128, i8* %1)
-  %2 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
-  call void @llvm.lifetime.start.p0i8(i64 64, i8* %2)
-  %3 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
-  call void @llvm.lifetime.start.p0i8(i64 64, i8* %3)
-
-; Update GEP offsets
+  %b1 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
+  call void @llvm.lifetime.start.p0i8(i64 128, i8* %b1)
+  %b2 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
+  call void @llvm.lifetime.start.p0i8(i64 64, i8* %b2)
+  %b3 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
+  call void @llvm.lifetime.start.p0i8(i64 64, i8* %b3)
+
+; Update GEP offsets coming from alloca directly
 ; CHECK-NEXT:    [[TMP4:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* [[TC_I1]] to i8*
 ; CHECK-NEXT:    [[I1:%.*]] = getelementptr inbounds i8, i8* [[TMP4]], i64 512
 ; CHECK-NEXT:    [[TMP5:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::A.resolved"]]* [[TA_I3]] to i8*
 ; CHECK-NEXT:    [[I2:%.*]] = getelementptr inbounds i8, i8* [[TMP5]], i64 128
 ; CHECK-NEXT:    [[TMP6:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::B.resolved"]]* [[TB_I5]] to i8*
 ; CHECK-NEXT:    [[I3:%.*]] = getelementptr inbounds i8, i8* [[TMP6]], i64 256
-  %4 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
-  %i1 = getelementptr inbounds i8, i8* %4, i64 128
-  %5 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
-  %i2 = getelementptr inbounds i8, i8* %5, i64 64
-  %6 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
-  %i3 = getelementptr inbounds i8, i8* %6, i64 64
+  %b4 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
+  %i1 = getelementptr inbounds i8, i8* %b4, i64 128
+  %b5 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
+  %i2 = getelementptr inbounds i8, i8* %b5, i64 64
+  %b6 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
+  %i3 = getelementptr inbounds i8, i8* %b6, i64 64
 
 ; Do not touch if offest is not a constant
 ; CHECK-NEXT:    [[TMP7:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* [[TC_I1]] to i8*
 ; CHECK-NEXT:    [[I4:%.*]] = getelementptr inbounds i8, i8* [[TMP7]], i64 [[OFFSET]]
-  %7 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
-  %i4 = getelementptr inbounds i8, i8* %7, i64 %offset
+  %b7 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
+  %i4 = getelementptr inbounds i8, i8* %b7, i64 %offset
 
-; no change: GEP operand is not a result of bitcast
+; no change: GEP operand is comming from kernel argument - not from matrix type
 ; CHECK-NEXT:    [[I5:%.*]] = getelementptr inbounds i8, i8* [[PTR1]], i64 128
   %i5 = getelementptr inbounds i8, i8* %ptr1, i64 128
 
-; Do not touch if bitcast is not for matrix type
+; no change: GEP comes from bitcast that doesn't come from matrix type
 ; CHECK-NEXT:    [[TMP8:%.*]] = bitcast [4 x [4 x %"struct::almost_joint_matrix"]]* [[TI_I]] to i8*
 ; CHECK-NEXT:    [[I6:%.*]] = getelementptr inbounds i8, i8* [[TMP8]], i64 128
-  %8 = bitcast [4 x [4 x %"struct::almost_joint_matrix"]]* %tI.i to i8*
-  %i6 = getelementptr inbounds i8, i8* %8, i64 128
+  %b8 = bitcast [4 x [4 x %"struct::almost_joint_matrix"]]* %tI.i to i8*
+  %i6 = getelementptr inbounds i8, i8* %b8, i64 128
 
 ; no change - GEP is not based on i8
 ; CHECK-NEXT:    [[TMP9:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* [[TC_I1]] to i16*
 ; CHECK-NEXT:    [[ARRAYCTOR_END_I:%.*]] = getelementptr inbounds i16, i16* [[TMP9]], i64 128
-  %9 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i16*
-  %arrayctor.end.i = getelementptr inbounds i16, i16* %9, i64 128
+  %b9 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i16*
+  %arrayctor.end.i = getelementptr inbounds i16, i16* %b9, i64 128
+  br label %loop_header
 
+; Test going through phi value + check that we aren't doing ininite recursion
+loop_header:
+  %loop_cond = phi i1 [1, %entry], [0, %loop_header]
+  %loop_matrix = phi [4 x [4 x %"struct::joint_matrix::C"]]* [%tC.i, %entry], [%loop_matrix, %loop_header]
+
+; Update GEP offsets coming from phi value
+; CHECK:         [[TMP10:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* {{.*}} to i8*
+; CHECK-NEXT:    {{.*}} = getelementptr inbounds i8, i8* [[TMP10]], i64 128
+  %b10 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
+  %i7 = getelementptr inbounds i8, i8* %b10, i64 32
+
+  br i1 %loop_cond, label %loop_header, label %after_loop
+
+after_loop:
 ; Life time end size update
-; CHECK-NEXT:    [[TMP10:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::B.resolved"]]* [[TB_I5]] to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 256, i8* [[TMP10]])
-; CHECK-NEXT:    [[TMP11:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::A.resolved"]]* [[TA_I3]] to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 128, i8* [[TMP11]])
-; CHECK-NEXT:    [[TMP12:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* [[TC_I1]] to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 512, i8* [[TMP12]])
-  %10 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
-  call void @llvm.lifetime.end.p0i8(i64 64, i8* %10)
-  %11 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
-  call void @llvm.lifetime.end.p0i8(i64 64, i8* %11)
-  %12 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
-  call void @llvm.lifetime.end.p0i8(i64 128, i8* %12)
+; CHECK:         [[TMP11:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::B.resolved"]]* [[TB_I5]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 256, i8* [[TMP11]])
+; CHECK-NEXT:    [[TMP12:%.*]] = bitcast [4 x [2 x %"struct::joint_matrix::A.resolved"]]* [[TA_I3]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 128, i8* [[TMP12]])
+; CHECK-NEXT:    [[TMP13:%.*]] = bitcast [4 x [4 x %"struct::joint_matrix::C.resolved"]]* [[TC_I1]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 512, i8* [[TMP13]])
+  %b11 = bitcast [4 x [2 x %"struct::joint_matrix::B"]]* %tB.i to i8*
+  call void @llvm.lifetime.end.p0i8(i64 64, i8* %b11)
+  %b12 = bitcast [4 x [2 x %"struct::joint_matrix::A"]]* %tA.i to i8*
+  call void @llvm.lifetime.end.p0i8(i64 64, i8* %b12)
+  %b13 = bitcast [4 x [4 x %"struct::joint_matrix::C"]]* %tC.i to i8*
+  call void @llvm.lifetime.end.p0i8(i64 128, i8* %b13)
 
 ; do not touch life time intrinsics if not for Joint Matrix types
 ; CHECK-NEXT:    [[GROUPID_ASCAST:%.*]] = addrspacecast [3 x i64]* [[GROUPID]] to [3 x i64] addrspace(4)*
-; CHECK-NEXT:    [[TMP13:%.*]] = bitcast [3 x i64]* [[GROUPID]] to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 24, i8* [[TMP13]])
 ; CHECK-NEXT:    [[TMP14:%.*]] = bitcast [3 x i64]* [[GROUPID]] to i8*
-; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 24, i8* [[TMP14]])
+; CHECK-NEXT:    call void @llvm.lifetime.start.p0i8(i64 24, i8* [[TMP14]])
+; CHECK-NEXT:    [[TMP15:%.*]] = bitcast [3 x i64]* [[GROUPID]] to i8*
+; CHECK-NEXT:    call void @llvm.lifetime.end.p0i8(i64 24, i8* [[TMP15]])
   %GroupID.ascast = addrspacecast [3 x i64]* %GroupID to [3 x i64] addrspace(4)*
-  %13 = bitcast [3 x i64]* %GroupID to i8*
-  call void @llvm.lifetime.start.p0i8(i64 24, i8* %13)
-  %14 = bitcast [3 x i64]* %GroupID to i8*
-  call void @llvm.lifetime.end.p0i8(i64 24, i8* %14)
+  %b14 = bitcast [3 x i64]* %GroupID to i8*
+  call void @llvm.lifetime.start.p0i8(i64 24, i8* %b14)
+  %b15 = bitcast [3 x i64]* %GroupID to i8*
+  call void @llvm.lifetime.end.p0i8(i64 24, i8* %b15)
 
 ; CHECK-NEXT:    ret void
   ret void