Fix bugs, add more lit tests for predicated memory optimizations

Adds more lit tests to verify the functionality of the predicated memory optimizations.
Fixes a few bugs found during the testing.

Author: YuriPlyakhin

IGC/Compiler/CISACodeGen/MemOpt.cpp

@@ -88,6 +88,8 @@ namespace {
         // A list of memory references (within a BB) with the distance to the begining of the BB.
         typedef std::vector<std::pair<Instruction*, unsigned> > MemRefListTy;
         typedef std::vector<Instruction*> TrivialMemRefListTy;
+        // ALoadInst, Offset, MemRefListTy::iterator, LeadingLoad's int2PtrOffset
+        typedef SmallVector<std::tuple<Instruction *, int64_t, MemRefListTy::iterator>, 8> MergeVector;
 
     public:
         static char ID;
@@ -134,6 +136,10 @@ namespace {
         Value* getShuffle(Value* ShflId, Instruction* BlockReadToOptimize,
             Value* SgId, llvm::IRBuilder<>& Builder, unsigned& ToOptSize);
 
+        unsigned getNumElements(Type* Ty) {
+            return Ty->isVectorTy() ? (unsigned)cast<IGCLLVM::FixedVectorType>(Ty)->getNumElements() : 1;
+        }
+
         Type* getVectorElementType(Type* Ty) const {
             return isa<VectorType>(Ty) ? cast<VectorType>(Ty)->getElementType() : Ty;
         }
@@ -145,6 +151,95 @@ namespace {
             return DL->getTypeStoreSize(A) == DL->getTypeStoreSize(B);
         }
 
+        Value* createBitOrPointerCast(Value* V, Type* DestTy,
+            IGCIRBuilder<>& Builder) {
+            if (V->getType() == DestTy)
+                return V;
+
+            if (V->getType()->isPointerTy() && DestTy->isPointerTy()) {
+                PointerType* SrcPtrTy = cast<PointerType>(V->getType());
+                PointerType* DstPtrTy = cast<PointerType>(DestTy);
+                if (SrcPtrTy->getPointerAddressSpace() !=
+                    DstPtrTy->getPointerAddressSpace())
+                    return Builder.CreateAddrSpaceCast(V, DestTy);
+            }
+
+            if (V->getType()->isPointerTy()) {
+                if (DestTy->isIntegerTy()) {
+                    return Builder.CreatePtrToInt(V, DestTy);
+                }
+                else if (DestTy->isFloatingPointTy()) {
+                    uint32_t Size = (uint32_t)DestTy->getPrimitiveSizeInBits();
+                    Value* Cast = Builder.CreatePtrToInt(
+                        V, Builder.getIntNTy(Size));
+                    return Builder.CreateBitCast(Cast, DestTy);
+                }
+            }
+
+            if (DestTy->isPointerTy()) {
+                if (V->getType()->isIntegerTy()) {
+                    return Builder.CreateIntToPtr(V, DestTy);
+                }
+                else if (V->getType()->isFloatingPointTy()) {
+                    uint32_t Size = (uint32_t)V->getType()->getPrimitiveSizeInBits();
+                    Value* Cast = Builder.CreateBitCast(
+                        V, Builder.getIntNTy(Size));
+                    return Builder.CreateIntToPtr(Cast, DestTy);
+                }
+            }
+
+            return Builder.CreateBitCast(V, DestTy);
+        }
+
+        /**
+         * @brief Creates a new merge value for merged load from a set of predicated loads' merge values.
+         *
+         * This function constructs a new combined merge value by merging the merge values of multiple predicated load intrinsics.
+         * Merge value from each input predicated load is inserted into the appropriate position in the resulting merge vector value,
+         * based on its offset and the scalar size. The function handles both scalar and vector merge input values.
+         *
+         * @param MergeValTy The type of the merged value to be created.
+         * @param LoadsToMerge A vector of tuples, each containing a load instruction and its associated offset.
+         * @param LdScalarSize The size (in bytes) of the scalar element being loaded in the combined load.
+         * @param NumElts Number of elements in the merged value vector.
+         * @return Value* The newly created merged value, or nullptr if we are merging generic loads, not predicated.
+         */
+        Value* CreateNewMergeValue(IGCIRBuilder<>& Builder, Type* MergeValTy,
+                                  const MergeVector& LoadsToMerge, unsigned LdScalarSize,
+                                  unsigned& NumElts) {
+            Value* NewMergeValue = UndefValue::get(MergeValTy);
+            unsigned Pos = 0;
+            int64_t FirstOffset = std::get<1>(LoadsToMerge.front());
+
+            for (auto& I : LoadsToMerge) {
+                PredicatedLoadIntrinsic* PLI = ALoadInst::get(std::get<0>(I))->getPredicatedLoadIntrinsic();
+                if (!PLI)
+                    return nullptr;
+
+                Value* MergeValue = PLI->getMergeValue();
+                unsigned MergeValNumElements = getNumElements(MergeValue->getType());
+                Type* MergeValScalarTy = MergeValTy->getScalarType();
+                Pos = unsigned((std::get<1>(I) - FirstOffset) / LdScalarSize);
+
+                if (MergeValNumElements == 1) {
+                    IGC_ASSERT_MESSAGE(Pos < NumElts, "Index is larger than the number of elements, we cannot update merge value.");
+                    MergeValue = createBitOrPointerCast(MergeValue, MergeValScalarTy, Builder);
+                    NewMergeValue = Builder.CreateInsertElement(NewMergeValue, MergeValue, Builder.getInt32(Pos));
+                    continue;
+                }
+
+                IGC_ASSERT_MESSAGE(Pos + MergeValNumElements <= NumElts,
+                    "Index is larger than the number of elements, we cannot update merge value.");
+
+                for (unsigned i = 0; i < MergeValNumElements; ++i) {
+                    Value* ExtractValue = Builder.CreateExtractElement(MergeValue, Builder.getInt32(i));
+                    ExtractValue = createBitOrPointerCast(ExtractValue, MergeValScalarTy, Builder);
+                    NewMergeValue = Builder.CreateInsertElement(NewMergeValue, ExtractValue, Builder.getInt32(Pos + i));
+                }
+            }
+            return NewMergeValue;
+        }
+
         bool isSafeToMergeLoad(const ALoadInst& Ld,
             const SmallVectorImpl<Instruction*>& checkList) const;
         bool isSafeToMergeStores(
@@ -644,76 +739,6 @@ bool MemOpt::removeRedBlockRead(GenIntrinsicInst* LeadingBlockRead,
     return true;
 }
 
-namespace {
-    unsigned getNumElements(Type* Ty) {
-        return Ty->isVectorTy() ? (unsigned)cast<IGCLLVM::FixedVectorType>(Ty)->getNumElements() : 1;
-    }
-
-    template <typename T>
-    Value* createBitOrPointerCast(Value* V, Type* DestTy,
-        IGCIRBuilder<T>& Builder) {
-        if (V->getType() == DestTy)
-            return V;
-        if (V->getType()->isPointerTy() && DestTy->isPointerTy()) {
-            PointerType* SrcPtrTy = cast<PointerType>(V->getType());
-            PointerType* DstPtrTy = cast<PointerType>(DestTy);
-            if (SrcPtrTy->getPointerAddressSpace() !=
-                DstPtrTy->getPointerAddressSpace())
-                return Builder.CreateAddrSpaceCast(V, DestTy);
-        }
-        if (V->getType()->isPointerTy()) {
-            if (DestTy->isIntegerTy()) {
-                return Builder.CreatePtrToInt(V, DestTy);
-            }
-            else if (DestTy->isFloatingPointTy()) {
-                uint32_t Size = (uint32_t)DestTy->getPrimitiveSizeInBits();
-                Value* Cast = Builder.CreatePtrToInt(
-                    V, Builder.getIntNTy(Size));
-                return Builder.CreateBitCast(Cast, DestTy);
-            }
-        }
-        if (DestTy->isPointerTy()) {
-            if (V->getType()->isIntegerTy()) {
-                return Builder.CreateIntToPtr(V, DestTy);
-            }
-            else if (V->getType()->isFloatingPointTy()) {
-                uint32_t Size = (uint32_t)V->getType()->getPrimitiveSizeInBits();
-                Value* Cast = Builder.CreateBitCast(
-                    V, Builder.getIntNTy(Size));
-                return Builder.CreateIntToPtr(Cast, DestTy);
-            }
-        }
-        return Builder.CreateBitCast(V, DestTy);
-    }
-
-    template <typename T>
-    Value* CreateNewMergeValue(IGCIRBuilder<T>& Builder, Type* Ty,
-                               const SmallVector<Instruction *>& LoadsToMerge) {
-        Value* NewMergeValue = UndefValue::get(Ty);
-        unsigned idx = 0;
-        for (auto* Inst : LoadsToMerge) {
-            PredicatedLoadIntrinsic* PLI = ALoadInst::get(Inst)->getPredicatedLoadIntrinsic();
-            if (!PLI)
-                return nullptr;
-            Value* MergeValue = PLI->getMergeValue();
-            unsigned NumElements = getNumElements(MergeValue->getType());
-            Type* ScalarTy = Ty->getScalarType();
-            if (NumElements == 1) {
-                MergeValue = createBitOrPointerCast(MergeValue, ScalarTy, Builder);
-                NewMergeValue = Builder.CreateInsertElement(NewMergeValue, MergeValue, Builder.getInt32(idx++));
-                continue;
-            }
-
-            for (unsigned i = 0; i < NumElements; ++i) {
-                Value* ExtractValue = Builder.CreateExtractElement(MergeValue, Builder.getInt32(i));
-                ExtractValue = createBitOrPointerCast(ExtractValue, ScalarTy, Builder);
-                NewMergeValue = Builder.CreateInsertElement(NewMergeValue, ExtractValue, Builder.getInt32(idx++));
-            }
-        }
-        return NewMergeValue;
-    }
-}
-
 //Removes redundant blockread if both blockreads are scalar.
 void MemOpt::removeScalarBlockRead(Instruction* BlockReadToOptimize,
     Instruction* BlockReadToRemove, Value* SgId,
@@ -1151,8 +1176,7 @@ bool MemOpt::mergeLoad(ALoadInst& LeadingLoad,
     }
 
     // ALoadInst, Offset, MemRefListTy::iterator, LeadingLoad's int2PtrOffset
-    SmallVector<std::tuple<Instruction *, int64_t, MemRefListTy::iterator>, 8>
-        LoadsToMerge;
+    MergeVector LoadsToMerge;
     LoadsToMerge.push_back(std::make_tuple(LeadingLoad.inst(), 0, MI));
 
     // Loads to be merged is scanned in the program order and will be merged into
@@ -1375,10 +1399,8 @@ bool MemOpt::mergeLoad(ALoadInst& LeadingLoad,
     Value* NewPointer = Builder.CreateBitCast(Ptr, NewPointerType);
 
     // Prepare Merge Value if needed:
-    SmallVector<Instruction*> LoadsToMergeInsts;
-    for (auto& I : LoadsToMerge)
-        LoadsToMergeInsts.push_back(std::get<0>(I));
-    Value* NewMergeValue = CreateNewMergeValue(Builder, NewLoadType, LoadsToMergeInsts);
+    Value* NewMergeValue = CreateNewMergeValue(Builder, NewLoadType, LoadsToMerge,
+                                               LdScalarSize, NumElts);
 
     Instruction* NewLoad =
         FirstLoad.CreateAlignedLoad(Builder, NewLoadType, NewPointer, NewMergeValue);
@@ -3486,7 +3508,7 @@ void LdStCombine::combineLoads()
                     ++numInsts;
 
                 // cannot merge beyond fence or window limit
-                if ((I->isFenceLike() && !isa<PredicatedLoadIntrinsic>(I)) || numInsts > LDWINDOWSIZE) {
+                if ((I->isFenceLike() && !isa<PredicatedLoadIntrinsic>(I) && !isa<PredicatedStoreIntrinsic>(I)) || numInsts > LDWINDOWSIZE) {
                     LLVM_DEBUG(dbgs() << "- - Stop at fence or window limit\n");
                     break;
                 }
@@ -3772,6 +3794,7 @@ void LdStCombine::createBundles(BasicBlock* BB, InstAndOffsetPairs& LoadStores)
         {
             // 1. The first one is the leading store.
             const LdStInfo* leadLSI = &LoadStores[i];
+            LLVM_DEBUG(llvm::dbgs() << "Try leading LdSt: " << *leadLSI->getInst() << "\n");
             if (isBundled(leadLSI, m_combinedInsts) ||
                 (i+1) == SZ) /* skip for last one */ {
                 ++i;
@@ -3780,7 +3803,7 @@ void LdStCombine::createBundles(BasicBlock* BB, InstAndOffsetPairs& LoadStores)
 
             if (m_WI && isUniform &&
                 !m_WI->isUniform(leadLSI->getValueOperand())) {
-                // no combining for *uniform-ptr = non-uniform value
+                LLVM_DEBUG(llvm::dbgs() << "No combining for *uniform-ptr = non-uniform value\n");
                 ++i;
                 continue;
             }
@@ -3816,6 +3839,7 @@ void LdStCombine::createBundles(BasicBlock* BB, InstAndOffsetPairs& LoadStores)
             uint32_t vecSize = -1;
             for (int j = i + 1; j < SZ; ++j) {
                 const LdStInfo* LSI = &LoadStores[j];
+                LLVM_DEBUG(llvm::dbgs() << "Try to make bundle with: " << *LSI->getInst() << "\n");
                 if (isBundled(LSI, m_combinedInsts) ||
                     (leadLSI->getByteOffset() + totalBytes) != LSI->getByteOffset())
                 {
@@ -3824,7 +3848,7 @@ void LdStCombine::createBundles(BasicBlock* BB, InstAndOffsetPairs& LoadStores)
                 }
                 if (m_WI && isUniform &&
                     !m_WI->isUniform(LSI->getValueOperand())) {
-                    // no combining for *uniform-ptr = non-uniform value
+                    LLVM_DEBUG(llvm::dbgs() << "No combining for *uniform-ptr = non-uniform value\n");
                     break;
                 }
 

IGC/Compiler/tests/LdStCombine/checkmetadata-predicated.ll

@@ -0,0 +1,106 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+
+; REQUIRES: regkeys
+; RUN: igc_opt --typed-pointers %s -S -inputocl -igc-ldstcombine -regkey=EnableLdStCombine=5 -platformbmg -instcombine | FileCheck %s
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f16:16:16-f32:32:32-f64:64:64-f80:128:128-v16:16:16-v24:32:32-v32:32:32-v48:64:64-v64:64:64-v96:128:128-v128:128:128-v192:256:256-v256:256:256-v512:512:512-v1024:1024:1024-a:64:64-f80:128:128-n8:16:32:64"
+
+define void @f0(i32* %dst, i32* %src) {
+entry:
+  %0 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %src, i64 4, i1 true, i32 42), !nontemporal !5
+  %arrayidx1 = getelementptr inbounds i32, i32* %src, i64 1
+  %1 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx1, i64 4, i1 true, i32 43), !nontemporal !5
+  %arrayidx2 = getelementptr inbounds i32, i32* %src, i64 2
+  %2 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx2, i64 4, i1 true, i32 44), !nontemporal !5
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %dst, i32 %0, i64 4, i1 true), !nontemporal !5
+  %arrayidx4 = getelementptr inbounds i32, i32* %dst, i64 1
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx4, i32 %1, i64 4, i1 true), !nontemporal !5
+  %arrayidx5 = getelementptr inbounds i32, i32* %dst, i64 2
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx5, i32 %2, i64 4, i1 true), !nontemporal !5
+  ret void
+}
+
+; CHECK-LABEL: @f0(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[SRC:%.*]] to <3 x i32>*
+; CHECK-NEXT:    [[TMP1:%.*]] = call <3 x i32> @llvm.genx.GenISA.PredicatedLoad.v3i32.p0v3i32.v3i32(<3 x i32>* [[TMP0]], i64 4, i1 true, <3 x i32> <i32 42, i32 43, i32 44>)
+; CHECK-SAME:    !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i32* [[DST:%.*]] to <3 x i32>*
+; CHECK-NEXT:    call void @llvm.genx.GenISA.PredicatedStore.p0v3i32.v3i32(<3 x i32>* [[TMP2]], <3 x i32> [[TMP1]], i64 4, i1 true)
+; CHECK-SAME:    !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    ret void
+
+
+define void @f1(i32* %dst, i32* %src) {
+entry:
+  %0 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %src, i64 4, i1 true, i32 42)
+  %arrayidx1 = getelementptr inbounds i32, i32* %src, i64 1
+  %1 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx1, i64 4, i1 true, i32 43), !nontemporal !5
+  %arrayidx2 = getelementptr inbounds i32, i32* %src, i64 2
+  %2 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx2, i64 4, i1 true, i32 44)
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %dst, i32 %0, i64 4, i1 true)
+  %arrayidx4 = getelementptr inbounds i32, i32* %dst, i64 1
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx4, i32 %1, i64 4, i1 true)
+  %arrayidx5 = getelementptr inbounds i32, i32* %dst, i64 2
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx5, i32 %2, i64 4, i1 true), !nontemporal !5
+  ret void
+}
+
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[SRC:%.*]] to <3 x i32>*
+; CHECK-NEXT:    [[TMP1:%.*]] = call <3 x i32> @llvm.genx.GenISA.PredicatedLoad.v3i32.p0v3i32.v3i32(<3 x i32>* [[TMP0]], i64 4, i1 true, <3 x i32> <i32 42, i32 43, i32 44>)
+; CHECK-SAME:    !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i32* [[DST:%.*]] to <3 x i32>*
+; CHECK-NEXT:    call void @llvm.genx.GenISA.PredicatedStore.p0v3i32.v3i32(<3 x i32>* [[TMP2]], <3 x i32> [[TMP1]], i64 4, i1 true)
+; CHECK-SAME:    !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    ret void
+
+define void @f2(i32* %dst, i32* %src) {
+entry:
+  %0 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %src, i64 4, i1 true, i32 42)
+  %arrayidx1 = getelementptr inbounds i32, i32* %src, i64 1
+  %1 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx1, i64 4, i1 true, i32 43)
+  %arrayidx2 = getelementptr inbounds i32, i32* %src, i64 2
+  %2 = call i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32* %arrayidx2, i64 4, i1 true, i32 44)
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %dst, i32 %0, i64 4, i1 true)
+  %arrayidx4 = getelementptr inbounds i32, i32* %dst, i64 1
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx4, i32 %1, i64 4, i1 true)
+  %arrayidx5 = getelementptr inbounds i32, i32* %dst, i64 2
+  call void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32* %arrayidx5, i32 %2, i64 4, i1 true)
+  ret void
+}
+
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = bitcast i32* [[SRC:%.*]] to <3 x i32>*
+; CHECK-NEXT:    [[TMP1:%.*]] = call <3 x i32> @llvm.genx.GenISA.PredicatedLoad.v3i32.p0v3i32.v3i32(<3 x i32>* [[TMP0]], i64 4, i1 true, <3 x i32> <i32 42, i32 43, i32 44>)
+; CHECK-NOT:     !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast i32* [[DST:%.*]] to <3 x i32>*
+; CHECK-NEXT:    call void @llvm.genx.GenISA.PredicatedStore.p0v3i32.v3i32(<3 x i32>* [[TMP2]], <3 x i32> [[TMP1]], i64 4, i1 true)
+; CHECK-NOT:     !nontemporal ![[NONTEMPORAL:[0-9]+]]
+; CHECK-NEXT:    ret void
+
+; CHECK: ![[NONTEMPORAL]] = !{i32 1}
+
+; Function Attrs: nounwind readonly
+declare i32 @llvm.genx.GenISA.PredicatedLoad.i32.p0i32.i32(i32*, i64, i1, i32) #0
+
+declare void @llvm.genx.GenISA.PredicatedStore.p0i32.i32(i32*, i32, i64, i1)
+
+attributes #0 = { nounwind readonly }
+
+!igc.functions = !{!0, !3, !4}
+
+!0 = !{void (i32*, i32*)* @f0, !1}
+!1 = !{!2}
+!2 = !{!"function_type", i32 0}
+!3 = !{void (i32*, i32*)* @f1, !1}
+!4 = !{void (i32*, i32*)* @f2, !1}
+!5 = !{i32 1}