Fix IVDep translation for accesses to kernel closure fields

When IVDep is applied to a regular kernel-scope array/pointer variable,
the GEP instructions that access this variable are marked with index
group metadata.
When IVDep is applied to a captured argument of a kernel, however,
the pointer operand of index group-marked instructions becomes a GEP
itself. The latter GEP accesses the closure of a kernel lambda with
an offset that signifies the number of the captured parameter in the
closure layout.

To treat the second case:
1. In forward translation, we ensure that for each memory block, all
pointer accesses marked into IVDep index groups are recognized and
translated into LoopControl parameters. Previously, only the latest
pointer access would get listed in the SPIR-V representation of
attribute parameters.
2. In backward translation, we also differentiate between GEP
instructions that access an array variable directly and the GEPs
accessing the "closure field accesses". An additional map stores
1 index group metadata node per each "pointer to closure + integer
offset" pair. This ensures that all accesses to the same kernel
parameter are marked into the same index group on a given loop nest
level.

Signed-off-by: Artem Gindinson <artem.gindinson@intel.com>

Author: Artem Gindinson

llvm-spirv/lib/SPIRV/SPIRVReader.cpp

@@ -53,6 +53,7 @@
 #include "VectorComputeUtil.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/IR/Constants.h"
@@ -803,18 +804,18 @@ void SPIRVToLLVM::setLLVMLoopMetadata(const LoopInstType *LM,
            "Missing loop control parameter!");
   }
   if (LC & LoopControlDependencyArrayINTELMask) {
-    // Collect array variable <-> safelen information
-    std::map<Value *, unsigned> ArraySflnMap;
+    // Collect pointer variable <-> safelen information
+    std::map<Value *, unsigned> PointerSflnMap;
     unsigned NumOperandPairs = LoopControlParameters[NumParam];
     unsigned OperandsEndIndex = NumParam + NumOperandPairs * 2;
     assert(OperandsEndIndex <= LoopControlParameters.size() &&
            "Missing loop control parameter!");
     SPIRVModule *M = LM->getModule();
     while (NumParam < OperandsEndIndex) {
       SPIRVId ArraySPIRVId = LoopControlParameters[++NumParam];
-      Value *ArrayVar = ValueMap[M->getValue(ArraySPIRVId)];
+      Value *PointerVar = ValueMap[M->getValue(ArraySPIRVId)];
       unsigned Safelen = LoopControlParameters[++NumParam];
-      ArraySflnMap.emplace(ArrayVar, Safelen);
+      PointerSflnMap.emplace(PointerVar, Safelen);
     }
 
     // A single run over the loop to retrieve all GetElementPtr instructions
@@ -826,24 +827,54 @@ void SPIRVToLLVM::setLLVMLoopMetadata(const LoopInstType *LM,
         if (!GEP)
           continue;
 
-        Value *AccessedArray = GEP->getPointerOperand();
-        auto ArraySflnIt = ArraySflnMap.find(AccessedArray);
-        if (ArraySflnIt != ArraySflnMap.end())
-          ArrayGEPMap[AccessedArray].push_back(GEP);
+        Value *AccessedPointer = GEP->getPointerOperand();
+        auto PointerSflnIt = PointerSflnMap.find(AccessedPointer);
+        if (PointerSflnIt != PointerSflnMap.end()) {
+          ArrayGEPMap[AccessedPointer].push_back(GEP);
+        }
       }
     }
 
-    // Create index group metadata nodes - one per each array
+    // Create index group metadata nodes - one per each of the array
     // variables. Mark each GEP accessing a particular array variable
     // into a corresponding index group
-    std::map<unsigned, std::vector<MDNode *>> SafelenIdxGroupMap;
+    std::map<unsigned, SmallSet<MDNode *, 4>> SafelenIdxGroupMap;
+    // Whenever a kernel closure field access is pointed to instead of
+    // an array/pointer variable, ensure that all GEPs to that memory
+    // share the same index group by hashing the newly added index groups.
+    // "Memory offset info" represents a handle to the whole closure block
+    // + an integer offset to a particular captured parameter.
+    using MemoryOffsetInfo = std::pair<Value *, unsigned>;
+    std::map<MemoryOffsetInfo, MDNode *> OffsetIdxGroupMap;
+
     for (auto &ArrayGEPIt : ArrayGEPMap) {
-      // Emit a distinct index group that will be referenced from
-      // llvm.loop.parallel_access_indices metadata
-      auto *CurrentDepthIdxGroup = llvm::MDNode::getDistinct(*Context, None);
-      unsigned Safelen = ArraySflnMap.find(ArrayGEPIt.first)->second;
-      SafelenIdxGroupMap[Safelen].push_back(CurrentDepthIdxGroup);
+      MDNode *CurrentDepthIdxGroup = nullptr;
+      if (auto *PrecedingGEP = dyn_cast<GetElementPtrInst>(ArrayGEPIt.first)) {
+        Value *ClosureFieldPointer = PrecedingGEP->getPointerOperand();
+        unsigned Offset =
+            cast<ConstantInt>(PrecedingGEP->getOperand(2))->getZExtValue();
+        MemoryOffsetInfo Info{ClosureFieldPointer, Offset};
+        auto OffsetIdxGroupIt = OffsetIdxGroupMap.find(Info);
+        if (OffsetIdxGroupIt == OffsetIdxGroupMap.end()) {
+          // This is the first GEP encountered for this closure field.
+          // Emit a distinct index group that will be referenced from
+          // llvm.loop.parallel_access_indices metadata; hash the new
+          // MDNode for future accesses to the same memory.
+          CurrentDepthIdxGroup = llvm::MDNode::getDistinct(*Context, None);
+          OffsetIdxGroupMap.emplace(Info, CurrentDepthIdxGroup);
+        } else {
+          // Previous accesses to that field have already been indexed,
+          // just use the already-existing metadata.
+          CurrentDepthIdxGroup = OffsetIdxGroupIt->second;
+        }
+      } else /* Regular kernel-scope array/pointer variable */ {
+        // Emit a distinct index group that will be referenced from
+        // llvm.loop.parallel_access_indices metadata
+        CurrentDepthIdxGroup = llvm::MDNode::getDistinct(*Context, None);
+      }
 
+      unsigned Safelen = PointerSflnMap.find(ArrayGEPIt.first)->second;
+      SafelenIdxGroupMap[Safelen].insert(CurrentDepthIdxGroup);
       for (auto *GEP : ArrayGEPIt.second) {
         StringRef IdxGroupMDName("llvm.index.group");
         llvm::MDNode *PreviousIdxGroup = GEP->getMetadata(IdxGroupMDName);

llvm-spirv/lib/SPIRV/SPIRVWriter.cpp

@@ -951,7 +951,8 @@ class LLVMParallelAccessIndices {
       // survive the translation. This check should be replaced with an
       // assertion once all known cases are handled.
       if (IdxGroupArrayPairIt != IndexGroupArrayMap.end())
-        ArrayVariablesVec.push_back(IdxGroupArrayPairIt->second);
+        for (SPIRVId ArrayAccessId : IdxGroupArrayPairIt->second)
+          ArrayVariablesVec.push_back(ArrayAccessId);
     }
   }
 
@@ -1476,13 +1477,13 @@ SPIRVValue *LLVMToSPIRV::transValueWithoutDecoration(Value *V,
       // 1) The metadata node has no operands. It will be directly referenced
       //    from within the optimization hint metadata.
       if (NumOperands == 0)
-        IndexGroupArrayMap[IndexGroup] = AccessedArrayId;
+        IndexGroupArrayMap[IndexGroup].insert(AccessedArrayId);
       // 2) The metadata node has several operands. It serves to link an index
       //    group specific to some embedded loop with other index groups that
       //    mark the same array variable for the outer loop(s).
       for (unsigned I = 0; I < NumOperands; ++I) {
         auto *ContainedIndexGroup = getMDOperandAsMDNode(IndexGroup, I);
-        IndexGroupArrayMap[ContainedIndexGroup] = AccessedArrayId;
+        IndexGroupArrayMap[ContainedIndexGroup].insert(AccessedArrayId);
       }
     }
 

llvm-spirv/lib/SPIRV/SPIRVWriter.h

@@ -53,6 +53,7 @@
 #include "SPIRVType.h"
 #include "SPIRVValue.h"
 
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/IR/IntrinsicInst.h"
 
@@ -138,7 +139,7 @@ class LLVMToSPIRV : public ModulePass {
 
   typedef DenseMap<Type *, SPIRVType *> LLVMToSPIRVTypeMap;
   typedef DenseMap<Value *, SPIRVValue *> LLVMToSPIRVValueMap;
-  typedef DenseMap<MDNode *, SPIRVId> LLVMToSPIRVMetadataMap;
+  typedef DenseMap<MDNode *, SmallSet<SPIRVId, 2>> LLVMToSPIRVMetadataMap;
 
 private:
   Module *M;