Enable GAS resolution for arguments

When we can prove that pointer coming from kernel function
arguments is not changed inside the kernel, we assume that
it points to global address space

Author: ekochetk

IGC/Compiler/CISACodeGen/ResolveGAS.cpp

@@ -35,6 +35,7 @@ SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 #include "llvm/Support/Debug.h"
 #include "llvmWrapper/IR/Constant.h"
 #include <llvm/ADT/DenseSet.h>
+#include <llvm/ADT/SmallVector.h>
 #include <llvm/ADT/PostOrderIterator.h>
 #include <llvm/Analysis/LoopInfo.h>
 #include <llvm/IR/IRBuilder.h>
@@ -79,6 +80,8 @@ namespace {
         void getAnalysisUsage(AnalysisUsage& AU) const override {
             AU.setPreservesCFG();
             AU.addRequired<LoopInfoWrapperPass>();
+            AU.addRequired<AAResultsWrapperPass>();
+            AU.addRequired<MetaDataUtilsWrapper>();
         }
 
         bool isResolvableLoopPHI(PHINode* PN) const {
@@ -89,11 +92,17 @@ namespace {
         bool resolveOnFunction(Function*) const;
         bool resolveOnBasicBlock(BasicBlock*) const;
 
+        bool resolveMemoryFromHost(Function&) const;
+
         void populateResolvableLoopPHIs();
         void populateResolvableLoopPHIsForLoop(const Loop*);
 
         bool isAddrSpaceResolvable(PHINode* PN, const Loop* L,
             BasicBlock* BackEdge) const;
+
+        bool checkGenericArguments(Function& F) const;
+        void convertLoadToGlobal(LoadInst* LI) const;
+        bool isLoadGlobalCandidate(LoadInst* LI) const;
     };
 
     class GASPropagator : public InstVisitor<GASPropagator, bool> {
@@ -148,6 +157,8 @@ namespace IGC {
     IGC_INITIALIZE_PASS_BEGIN(GASResolving, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY,
         PASS_ANALYSIS)
         IGC_INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)
+        IGC_INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
+        IGC_INITIALIZE_PASS_DEPENDENCY(MetaDataUtilsWrapper)
         IGC_INITIALIZE_PASS_END(GASResolving, PASS_FLAG, PASS_DESC, PASS_CFG_ONLY,
             PASS_ANALYSIS)
 }
@@ -158,6 +169,8 @@ bool GASResolving::runOnFunction(Function& F) {
     IRB = &TheBuilder;
     Propagator = &ThePropagator;
 
+    resolveMemoryFromHost(F);
+
     populateResolvableLoopPHIs();
 
     bool Changed = false;
@@ -676,3 +689,173 @@ bool GASPropagator::visitCallInst(CallInst& I) {
 
     return false;
 }
+
+bool GASResolving::resolveMemoryFromHost(Function& F) const {
+    MetaDataUtils* pMdUtils = getAnalysis<MetaDataUtilsWrapper>().getMetaDataUtils();
+
+    // skip all non-entry functions
+    if (!isEntryFunc(pMdUtils, &F))
+        return false;
+
+    // early check in order not to iterate whole function
+    if (!checkGenericArguments(F))
+        return false;
+
+    SmallVector<StoreInst*, 32> Stores;
+    SmallVector<LoadInst*, 32> Loads;
+    AliasAnalysis* AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
+
+    // collect load candidates and in parallel check for unsafe instructions
+    // visitor may be a more beautiful way to do this
+    bool HasASCast = false; // if there exists addrspace cast from non global/generic space
+    bool HasPtoi = false; // if there exists ptrtoint with global/generic space
+    for (BasicBlock& B : F) {
+        for (Instruction& I : B) {
+            if (auto LI = dyn_cast<LoadInst>(&I)) {
+                if (isLoadGlobalCandidate(LI)) {
+                    Loads.push_back(LI);
+                }
+            }
+            else if (auto CI = dyn_cast<CallInst>(&I)) {
+                if (CI->onlyReadsMemory())
+                    continue;
+
+                // currently recognize only these ones
+                // in fact intrinsics should be marked as read-only
+                if (auto II = dyn_cast<IntrinsicInst>(CI)) {
+                    if (II->getIntrinsicID() == Intrinsic::lifetime_start ||
+                        II->getIntrinsicID() == Intrinsic::lifetime_end)
+                        continue;
+                }
+
+                // if we have an unsafe call in the kernel, abort
+                // to improve we can collect arguments of writing calls as memlocations for alias analysis
+                return false;
+            }
+            else if (auto PI = dyn_cast<PtrToIntInst>(&I)) {
+                // if we have a ptrtoint we need to check data flow which we don't want to
+                if (PI->getPointerAddressSpace() != ADDRESS_SPACE_GLOBAL &&
+                    PI->getPointerAddressSpace() != ADDRESS_SPACE_GENERIC)
+                    return false;
+                else {
+                    HasPtoi = true;
+                }
+
+                return false;
+            }
+            else if (auto AI = dyn_cast<AddrSpaceCastInst>(&I)) {
+                if (AI->getSrcAddressSpace() != ADDRESS_SPACE_GLOBAL &&
+                    AI->getSrcAddressSpace() != ADDRESS_SPACE_GENERIC) {
+                    HasASCast = true;
+                }
+            }
+            else if (auto SI = dyn_cast<StoreInst>(&I)) {
+                Value* V = SI->getValueOperand();
+                if (isa<PointerType>(V->getType())) {
+                    // this store can potentially write non-global pointer to memory
+                    Stores.push_back(SI);
+                }
+            }
+            else if (I.mayWriteToMemory()) {
+                // unsupported instruction poisoning memory
+                return false;
+            }
+        }
+    }
+    if (HasASCast && HasPtoi)
+        return false;
+
+    if (Loads.empty())
+        return false;
+
+    bool Changed = false;
+    while (!Loads.empty())
+    {
+        LoadInst* LI = Loads.pop_back_val();
+
+        // check that we don't have aliasing stores for this load
+        // we expect to have basic and addrspace AA available at the moment
+        // on optimization phase
+        bool aliases = false;
+        for (auto SI : Stores) {
+            if (AA->alias(MemoryLocation::get(SI), MemoryLocation::get(LI))) {
+                aliases = true;
+                break;
+            }
+        }
+        if (aliases)
+            continue;
+
+        convertLoadToGlobal(LI);
+        Changed = true;
+    }
+    return Changed;
+}
+
+bool GASResolving::isLoadGlobalCandidate(LoadInst* LI) const {
+    // first check that loaded address has generic address space
+    // otherwise it is not our candidate
+    PointerType* PtrTy = dyn_cast<PointerType>(LI->getType());
+    if (!PtrTy || PtrTy->getAddressSpace() != ADDRESS_SPACE_GENERIC)
+        return false;
+
+    // next check that it is a load from function argument + offset
+    // which is necessary to prove that this address has global addrspace
+    Value* LoadBase = LI->getPointerOperand()->stripInBoundsOffsets();
+    if (!isa<Argument>(LoadBase))
+        return false;
+
+    // don't want to process cases when argument is from local address space
+    auto LoadTy = cast<PointerType>(LoadBase->getType());
+    if (LoadTy->getAddressSpace() != ADDRESS_SPACE_GLOBAL)
+        return false;
+
+    // TODO: skip cases that have been fixed on previous traversals
+
+    return true;
+}
+
+void GASResolving::convertLoadToGlobal(LoadInst* LI) const {
+    // create two addressspace casts: generic -> global -> generic
+    // the next scalar phase of this pass will propagate global to all uses of the load
+
+    PointerType* PtrTy = cast<PointerType>(LI->getType());
+    IRB->SetInsertPoint(LI->getNextNode());
+    PointerType* GlobalPtrTy = PointerType::get(PtrTy->getElementType(), ADDRESS_SPACE_GLOBAL);
+    Value* GlobalAddr = IRB->CreateAddrSpaceCast(LI, GlobalPtrTy);
+    Value* GenericCopyAddr = IRB->CreateAddrSpaceCast(GlobalAddr, PtrTy);
+
+    for (auto UI = LI->use_begin(), UE = LI->use_end(); UI != UE; /*EMPTY*/) {
+        Use& U = *UI++;
+        if (U.getUser() == GlobalAddr)
+            continue;
+        U.set(GenericCopyAddr);
+    }
+}
+
+bool GASResolving::checkGenericArguments(Function& F) const {
+    // check that we have a pointer to pointer or pointer to struct that has pointer elements
+    // and main pointer type is global while underlying pointer type is generic
+
+    auto* FT = F.getFunctionType();
+    for (unsigned p = 0; p < FT->getNumParams(); ++p) {
+        if (auto Ty = dyn_cast<PointerType>(FT->getParamType(p))) {
+            if (Ty->getAddressSpace() != ADDRESS_SPACE_GLOBAL)
+                continue;
+            auto PteeTy = Ty->getElementType();
+            if (auto PTy = dyn_cast<PointerType>(PteeTy)) {
+                if (PTy->getAddressSpace() == ADDRESS_SPACE_GENERIC)
+                    return true;
+            }
+            if (auto STy = dyn_cast<StructType>(PteeTy)) {
+                for (unsigned e = 0; e < STy->getNumElements(); ++e) {
+                    if (auto ETy = dyn_cast<PointerType>(STy->getElementType(e))) {
+                        if (ETy->getAddressSpace() == ADDRESS_SPACE_GENERIC)
+                            return true;
+                    }
+                }
+            }
+        }
+    }
+    return false;
+}