Support Predicated Memory intrinsics in codegen memory optimization passes

Disabled PromoteToPredicatedMemory for illegal int types.
Implemented support for Predicated Memory intrinsics
in Memory Optimization passes: MemOpt, MemOpt2,
CombineLdSt, AdvMemOpt, PrepareLoadsStoresPass.
Disabled PromoteToPredicatedMemory for
instructions/calls which may create unexpected behavior.

Author: YuriPlyakhin

IGC/Compiler/CISACodeGen/AdvMemOpt.cpp

@@ -13,6 +13,7 @@ SPDX-License-Identifier: MIT
 #include <llvm/Pass.h>
 #include <llvm/Transforms/Utils/Local.h>
 #include <llvm/ADT/Optional.h>
+#include "llvmWrapper/Analysis/TargetLibraryInfo.h"
 #include "llvmWrapper/Transforms/Utils/LoopUtils.h"
 #include "common/LLVMWarningsPop.hpp"
 #include "Compiler/CISACodeGen/ShaderCodeGen.hpp"
@@ -29,13 +30,16 @@ using namespace llvm::PatternMatch;
 using namespace IGC;
 using namespace IGC::IGCMD;
 
+#define DEBUG_TYPE "AdvMemOpt"
+
 namespace {
 
     class AdvMemOpt : public FunctionPass {
         DominatorTree* DT = nullptr;
         LoopInfo* LI = nullptr;
         PostDominatorTree* PDT = nullptr;
         WIAnalysis* WI = nullptr;
+        TargetLibraryInfo* TLI = nullptr;
 
     public:
         static char ID;
@@ -58,6 +62,7 @@ namespace {
             AU.addRequired<DominatorTreeWrapperPass>();
             AU.addRequired<LoopInfoWrapperPass>();
             AU.addRequired<PostDominatorTreeWrapperPass>();
+            AU.addRequired<TargetLibraryInfoWrapperPass>();
         }
 
         bool collectOperandInst(SmallPtrSetImpl<Instruction*>&,
@@ -96,10 +101,13 @@ namespace IGC {
         IGC_INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
         IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
         IGC_INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass);
+        IGC_INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
     IGC_INITIALIZE_PASS_END(AdvMemOpt, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY, PASS_ANALYSIS)
 } // End namespace IGC
 
 bool AdvMemOpt::runOnFunction(Function& F) {
+    bool Changed = false;
+
     // Skip non-kernel function.
     MetaDataUtils* MDU = nullptr;
     MDU = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
@@ -111,6 +119,7 @@ bool AdvMemOpt::runOnFunction(Function& F) {
     PDT = &getAnalysis<PostDominatorTreeWrapperPass>().getPostDomTree();
     LI = &getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
     WI = &getAnalysis<WIAnalysis>();
+    TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 
     SmallVector<Loop*, 8> InnermostLoops;
     for (auto I = LI->begin(), E = LI->end(); I != E; ++I)
@@ -138,7 +147,7 @@ bool AdvMemOpt::runOnFunction(Function& F) {
                 }
             }
         }
-        hoistUniformLoad(Line);
+        Changed |= hoistUniformLoad(Line);
     }
 
     auto* Ctx = getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
@@ -151,8 +160,8 @@ bool AdvMemOpt::runOnFunction(Function& F) {
         // because, once ballot-loop is added, vISA finalizer cannot schedule
         // those sample operations.
         auto& DL = F.getParent()->getDataLayout();
-        IRBuilder<> IRB(F.getContext());
-        Cluster.init(Ctx, &DL, nullptr/*AA*/, 32);
+        IGCIRBuilder<> IRB(F.getContext());
+        Cluster.init(Ctx, &DL, nullptr/*AA*/, TLI, 32);
         for (Function::iterator I = F.begin(), E = F.end(); I != E;
              ++I) {
             BasicBlock *BB = &*I;
@@ -182,24 +191,26 @@ bool AdvMemOpt::runOnFunction(Function& F) {
                     if (!WI->isUniform(LI->getResourceValue())) {
                         NumResourceVarying++;
                     }
-                } else if (auto *SI = dyn_cast<StoreInst>(I)) {
-                    if (!WI->isUniform(SI))
+                } else if (auto SI = AStoreInst::get(I); SI.has_value()) {
+                    if (!WI->isUniform(SI->inst()))
                       continue;
 
                     unsigned AS = SI->getPointerAddressSpace();
                     if (AS != ADDRESS_SPACE_PRIVATE &&
                         AS != ADDRESS_SPACE_GLOBAL)
                       continue;
 
-                    IRB.SetInsertPoint(SI);
+                    IRB.SetInsertPoint(SI->inst());
 
-                    if (auto NewSI = expand64BitStore(IRB, DL, SI)) {
+                    if (auto NewSI = expand64BitStore(IRB, DL, SI.value())) {
+                      auto NewASI = AStoreInst::get(NewSI);
                       WI->incUpdateDepend(NewSI, WIAnalysis::UNIFORM_THREAD);
-                      WI->incUpdateDepend(NewSI->getValueOperand(),
+                      WI->incUpdateDepend(NewASI->getValueOperand(),
                                           WIAnalysis::UNIFORM_THREAD);
-                      WI->incUpdateDepend(NewSI->getPointerOperand(),
+                      WI->incUpdateDepend(NewASI->getPointerOperand(),
                                           WIAnalysis::UNIFORM_THREAD);
-                      SI->eraseFromParent();
+                      SI->inst()->eraseFromParent();
+                      Changed = true;
                     }
                 }
             }
@@ -209,24 +220,29 @@ bool AdvMemOpt::runOnFunction(Function& F) {
                     NumResourceVarying;
                 // clustering method cannot handle memory dependence
                 if (!HasStore)
-                    Cluster.runForGFX(BB);
+                    Changed |= Cluster.runForGFX(BB);
             }
         }
     }
-    return false;
+    return Changed;
 }
 
 bool AdvMemOpt::isLeadCandidate(BasicBlock* BB) const {
-    // A candidate lead should have at least one uniform loads. In addition,
+    // A candidate lead should have at least one uniform load. In addition,
     // there's no instruction might to write memory from the last uniform loads
     // to the end.
+    LLVM_DEBUG(dbgs() << "Check lead candidate: " << BB->getName() << "\n");
     for (auto II = BB->rbegin(), IE = BB->rend(); II != IE; ++II) {
-        if (II->mayWriteToMemory())
+        if (II->mayWriteToMemory()) {
+            LLVM_DEBUG(dbgs() <<" - May write to memory. Bail out: " << *II << "\n");
             return false;
-        LoadInst* LD = dyn_cast<LoadInst>(&*II);
-        if (!LD || !WI->isUniform(LD))
+        }
+        std::optional<ALoadInst> LD = ALoadInst::get(&*II);
+        if (!LD.has_value() || !WI->isUniform(LD->inst())) {
+            LLVM_DEBUG(dbgs() << " - Not uniform load. Skip: " << *II << "\n");
             continue;
-        // Found uniform loads.
+        }
+        LLVM_DEBUG(dbgs() << "Found uniform loads.\n");
         return true;
     }
     return false;
@@ -351,42 +367,66 @@ bool AdvMemOpt::hoistInst(Instruction* LD, BasicBlock* BB) const {
 
 bool AdvMemOpt::hoistUniformLoad(ArrayRef<BasicBlock*> Line) const {
     bool Changed = false;
-    // Find the lead BB where to hoist uniform load.
+    LLVM_DEBUG(dbgs() << "Find the lead BB where to hoist uniform load.\n");
+
     auto BI = Line.begin();
     auto BE = Line.end();
+
     while (BI != BE) {
         if (!isLeadCandidate(*BI)) {
             ++BI;
             continue;
         }
+
         // Found lead.
         BasicBlock* Lead = *BI++;
-        BasicBlock* Prev = Lead;
+        LLVM_DEBUG(dbgs() << "Found lead to hoist to: " << Lead->getName() << "\n");
+
         for (; BI != BE; ++BI) {
             BasicBlock* Curr = *BI;
+            LLVM_DEBUG(dbgs() << " - Try to hoist from: " << Curr->getName() << "\n");
             // Check whether it's safe to hoist uniform loads from Curr to Lead by
             // checking all blocks between Prev and Curr.
-            if (hasMemoryWrite(Prev, Curr))
+            if (hasMemoryWrite(Lead, Curr)) {
+                LLVM_DEBUG(dbgs() << "- Memory write between Lead and Curr. Bail out.\n");
                 break;
+            }
+
             // Hoist uniform loads from Curr into Lead.
             for (auto II = Curr->getFirstNonPHI()->getIterator(),
                 IE = Curr->end(); II != IE; /*EMPTY*/) {
-                if (II->mayWriteToMemory())
+                LLVM_DEBUG(dbgs() << " - - Try hoisting: " << *II << "\n");
+
+                if (II->mayWriteToMemory()) {
+                    LLVM_DEBUG(dbgs() << " - - May write to memory. Bail out.\n");
                     break;
-                LoadInst* LD = dyn_cast<LoadInst>(&*II++);
-                if (!LD || !WI->isUniform(LD))
+                }
+
+                std::optional<ALoadInst> LD = ALoadInst::get(&*II++);
+                if (!LD.has_value() || !WI->isUniform(LD->inst())) {
+                    LLVM_DEBUG(dbgs() << " - - Not uniform load. Skip.\n");
                     continue;
-                if (!hoistInst(LD, Lead))
-                    break; // Bail out if any uniform load could not be hoisted safely.
-                  // Reset iterator
-                II = Curr->getFirstNonPHI()->getIterator();
+                }
+
+                if (!hoistInst(LD->inst(), Lead)) {
+                    LLVM_DEBUG(dbgs() << " - - Uniform load could not be hoisted safely. Bail out.\n");
+                    break;
+                }
                 Changed = true;
+                LLVM_DEBUG(dbgs() << " - - Hoisted!\n");
+
+                // Reset iterator
+                II = Curr->getFirstNonPHI()->getIterator();
             }
+
             // After hoisting uniform loads safely, if Curr has memory write, stop
             // hoisting further.
-            if (hasMemoryWrite(Curr))
+            if (hasMemoryWrite(Curr)) {
+                LLVM_DEBUG(dbgs() << "- Curr has memory write. Bail out.\n");
                 break;
+            }
         }
     }
+
     return Changed;
 }

IGC/Compiler/CISACodeGen/MemOpt2.cpp

@@ -11,6 +11,7 @@ SPDX-License-Identifier: MIT
 #include "common/LLVMWarningsPush.hpp"
 #include <llvm/Pass.h>
 #include <llvm/Transforms/Utils/Local.h>
+#include <llvmWrapper/Analysis/TargetLibraryInfo.h>
 #include <llvmWrapper/IR/DerivedTypes.h>
 #include "common/LLVMWarningsPop.hpp"
 
@@ -117,9 +118,9 @@ bool MemInstCluster::runForOCL(Function& F) {
 bool MemInstCluster::isSafeToMoveTo(Instruction* I, Instruction* Pos, const SmallVectorImpl<Instruction*>* CheckList) const {
     // TODO: So far, we simply don't allow rescheduling load/atomic operations.
     // Add alias analysis to allow memory operations to be rescheduled.
-    if (auto LD = dyn_cast<LoadInst>(I)) {
+    if (auto LD = ALoadInst::get(I); LD.has_value()) {
         if (CheckList)
-            return isSafeToScheduleLoad(LD, CheckList);
+            return isSafeToScheduleLoad(LD.value(), CheckList);
         return false;
     }
     if (GenIntrinsicInst * GII = dyn_cast<GenIntrinsicInst>(I)) {
@@ -292,26 +293,26 @@ bool MemInstCluster::clusterLoad(BasicBlock* BB) {
     unsigned MaxLiveOutByte = getMaxLiveOutThreshold() * 4;
     unsigned CountByte = 0;
     for (auto BI = BB->begin(), BE = BB->end(); BI != BE; ++BI) {
-        LoadInst* Lead = dyn_cast<LoadInst>(BI);
-        if (!Lead || !Lead->isSimple())
+        std::optional<ALoadInst> Lead = ALoadInst::get(&(*BI));
+        if (!Lead.has_value() || !Lead->isSimple())
             continue;
         if (Lead->getPointerAddressSpace() != ADDRESS_SPACE_LOCAL &&
             Lead->getPointerAddressSpace() != ADDRESS_SPACE_GLOBAL)
             continue;
 
-        CountByte = getNumLiveOutBytes(Lead);
+        CountByte = getNumLiveOutBytes(Lead->inst());
         if (CountByte > MaxLiveOutByte)
             continue;
 
         SmallVector<Instruction*, 8> CheckList;
-        InsertPos = Lead;
+        InsertPos = Lead->inst();
         // Find candidate the cluster them.
         BasicBlock::iterator I = BasicBlock::iterator(InsertPos), E;
         for (I = std::next(I), E = BB->end(); I != E; ++I) {
             if (I->mayWriteToMemory())
                 CheckList.push_back(&(*I));
-            LoadInst* Next = dyn_cast<LoadInst>(I);
-            if (!Next || !Next->isSimple())
+            std::optional<ALoadInst> Next = ALoadInst::get(&(*I));
+            if (!Next.has_value() || !Next->isSimple())
                 continue;
             // Skip memory accesses on different memory address space.
             // FIXME: GetUnderlyingObject() cannot track through `inttoptr`
@@ -320,21 +321,21 @@ bool MemInstCluster::clusterLoad(BasicBlock* BB) {
             // same buffer.
             if (Next->getPointerAddressSpace() != Lead->getPointerAddressSpace())
                 continue;
-            CountByte += getNumLiveOutBytes(Next);
+            CountByte += getNumLiveOutBytes(Next->inst());
             if (CountByte > MaxLiveOutByte) {
-                BasicBlock::iterator I = BasicBlock::iterator(Next);
+                BasicBlock::iterator I = BasicBlock::iterator(Next->inst());
                 BI = std::prev(I);
                 break;
             }
-            Changed |= schedule(BB, Next, InsertPos, &CheckList);
+            Changed |= schedule(BB, Next->inst(), InsertPos, &CheckList);
         }
     }
     return Changed;
 }
 
-bool MemInstCluster::isSafeToScheduleLoad(const LoadInst* LD,
+bool MemInstCluster::isSafeToScheduleLoad(const ALoadInst& LD,
     const SmallVectorImpl<Instruction*>* CheckList) const {
-    MemoryLocation A = MemoryLocation::get(LD);
+    MemoryLocation A = getLocation(LD.inst(), TLI);
 
     for (auto* I : *CheckList) {
         // Skip instructions never writing to memory.
@@ -343,7 +344,7 @@ bool MemInstCluster::isSafeToScheduleLoad(const LoadInst* LD,
         if (!AA)
             return false;
         // Unsafe if there's alias.
-        MemoryLocation B = getLocation(I);
+        MemoryLocation B = getLocation(I, TLI);
         if (!A.Ptr || !B.Ptr || AA->alias(A, B))
             return false;
     }
@@ -419,6 +420,7 @@ class MemOpt2 : public FunctionPass {
         AU.addRequired<AAResultsWrapperPass>();
         AU.addRequired<CodeGenContextWrapper>();
         AU.addRequired<MetaDataUtilsWrapper>();
+        AU.addRequired<TargetLibraryInfoWrapperPass>();
     }
     MemInstCluster Cluster;
     unsigned MaxLiveOutThreshold = 16;
@@ -438,6 +440,7 @@ IGC_INITIALIZE_PASS_BEGIN(MemOpt2, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY,
                           PASS_ANALYSIS)
 IGC_INITIALIZE_PASS_DEPENDENCY(CodeGenContextWrapper)
 IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
+IGC_INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 IGC_INITIALIZE_PASS_END(MemOpt2, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY,
                         PASS_ANALYSIS)
 
@@ -456,7 +459,8 @@ bool MemOpt2::runOnFunction(Function &F) {
         getAnalysis<CodeGenContextWrapper>().getCodeGenContext();
     auto DL = &F.getParent()->getDataLayout();
     auto AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
+    auto TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 
-    Cluster.init(cgCtx, DL, AA, MaxLiveOutThreshold);
+    Cluster.init(cgCtx, DL, AA, TLI, MaxLiveOutThreshold);
     return Cluster.runForOCL(F);
 }

IGC/Compiler/CISACodeGen/MemOpt2.h

@@ -23,10 +23,13 @@ SPDX-License-Identifier: MIT
 #include "Compiler/IGCPassSupport.h"
 #include "Compiler/MetaDataUtilsWrapper.h"
 
+#include "MemOptUtils.h"
+
 class MemInstCluster {
   IGC::CodeGenContext *CTX = nullptr;
   const DataLayout *DL = nullptr;
   AliasAnalysis *AA = nullptr;
+  TargetLibraryInfo* TLI = nullptr;
   unsigned MaxLiveOutThreshold = 0;
   llvm::DenseSet<Instruction *> Scheduled;
 
@@ -35,14 +38,15 @@ class MemInstCluster {
   ~MemInstCluster() {}
 
   MemInstCluster(IGC::CodeGenContext *pCTX, const DataLayout *pDL,
-                 AliasAnalysis *pAA, unsigned MLT) {
-    init(pCTX, pDL, pAA, MLT);
+                 AliasAnalysis *pAA, TargetLibraryInfo* pTLI, unsigned MLT) {
+    init(pCTX, pDL, pAA, pTLI, MLT);
   }
   void init(IGC::CodeGenContext *pCTX, const DataLayout *pDL,
-            AliasAnalysis *pAA, unsigned MLT) {
+            AliasAnalysis *pAA, TargetLibraryInfo* pTLI, unsigned MLT) {
     CTX = pCTX;
     DL = pDL;
     AA = pAA;
+    TLI = pTLI;
     MaxLiveOutThreshold = MLT;
   }
   ///  Called by MemOpt2 to cluster GPGPU kernels
@@ -61,19 +65,11 @@ class MemInstCluster {
   bool clusterLoad(BasicBlock *BB);
   bool isDefinedBefore(BasicBlock *BB, Instruction *I, Instruction *Pos) const;
   bool
-  isSafeToScheduleLoad(const LoadInst *LD,
+  isSafeToScheduleLoad(const IGC::ALoadInst& LD,
                        const SmallVectorImpl<Instruction *> *CheckList) const;
   bool schedule(BasicBlock *BB, Value *V, Instruction *&InsertPos,
                 const SmallVectorImpl<Instruction *> *CheckList = nullptr);
 
-  MemoryLocation getLocation(Instruction *I) const {
-    if (LoadInst *LD = dyn_cast<LoadInst>(I))
-      return MemoryLocation::get(LD);
-    if (StoreInst *ST = dyn_cast<StoreInst>(I))
-      return MemoryLocation::get(ST);
-    return MemoryLocation();
-  }
-
   unsigned getNumLiveOuts(Instruction *I) const;
 
   unsigned getNumLiveOutBytes(Instruction *I) const;