[sycl-post-link] Add undefined value handling

Author: Pavel Samolysov

llvm/tools/sycl-post-link/SpecConstants.cpp

@@ -51,6 +51,14 @@ constexpr char SPEC_CONST_MD_STRING[] = "sycl.specialization-constants";
 constexpr char SPEC_CONST_DEFAULT_VAL_MD_STRING[] =
     "sycl.specialization-constants-default-values";
 
+/// Spec. Constant ID is a pair of Id and a flag whether this Id belongs to an
+/// undefined value. Undefined values ('undef' in the IR) are used to get the
+/// required alignment and should be handled in a special manner as padding.
+struct ID {
+  unsigned ID;
+  bool Undef;
+};
+
 StringRef getStringLiteralArg(const CallInst *CI, unsigned ArgNo,
                               SmallVectorImpl<Instruction *> &DelInsts) {
   Value *V = CI->getArgOperand(ArgNo)->stripPointerCasts();
@@ -236,8 +244,13 @@ MDNode *generateSpecConstDefaultValueMetadata(StringRef SymID, Value *Default) {
 /// Recursively iterates over a composite type in order to collect information
 /// about its scalar elements.
 void collectCompositeElementsInfoRecursive(
-    const Module &M, Type *Ty, const unsigned *&IDIter, unsigned &Offset,
+    const Module &M, Type *Ty, const ID *&IDIter, unsigned &Offset,
     std::vector<SpecConstantDescriptor> &Result) {
+  if (IDIter->Undef) {
+    // We can just skip undefined values because every such value is just a
+    // padding and will be handled in a different manner.
+    return;
+  }
   if (auto *ArrTy = dyn_cast<ArrayType>(Ty)) {
     for (size_t I = 0; I < ArrTy->getNumElements(); ++I) {
       // TODO: this is a spot for potential optimization: for arrays we could
@@ -246,7 +259,9 @@ void collectCompositeElementsInfoRecursive(
       collectCompositeElementsInfoRecursive(M, ArrTy->getElementType(), IDIter,
                                             Offset, Result);
     }
-  } else if (auto *StructTy = dyn_cast<StructType>(Ty)) {
+    return;
+  }
+  if (auto *StructTy = dyn_cast<StructType>(Ty)) {
     const StructLayout *SL = M.getDataLayout().getStructLayout(StructTy);
     const unsigned BaseOffset = Offset;
     unsigned LocalOffset = Offset;
@@ -267,7 +282,12 @@ void collectCompositeElementsInfoRecursive(
         BaseOffset + SL->getSizeInBytes() - LocalOffset;
     if (PostStructPadding > 0) {
       SpecConstantDescriptor Desc;
-      // ID of padding descriptors is the max value possible.
+      // ID of padding descriptors is the max value possible. This value is a
+      // magic value for the runtime and will just be skipped. Even if there
+      // are many specialization constants and every constant has padding of
+      // a different length, everything will work regardless rewriting
+      // the descriptions with Desc.ID equals to the max value: they will just
+      // be ignored at all.
       Desc.ID = std::numeric_limits<unsigned>::max();
       Desc.Offset = LocalOffset;
       Desc.Size = PostStructPadding;
@@ -277,25 +297,29 @@ void collectCompositeElementsInfoRecursive(
     // Update "global" offset according to the total size of a handled struct
     // type.
     Offset += SL->getSizeInBytes();
-  } else if (auto *VecTy = dyn_cast<FixedVectorType>(Ty)) {
+    return;
+  }
+  if (auto *VecTy = dyn_cast<FixedVectorType>(Ty)) {
     for (size_t I = 0; I < VecTy->getNumElements(); ++I) {
       // TODO: this is a spot for potential optimization: for vectors we could
       // just make a single recursive call here and use it to populate Result
       // in a loop.
       collectCompositeElementsInfoRecursive(M, VecTy->getElementType(), IDIter,
                                             Offset, Result);
     }
-  } else { // Assume that we encountered some scalar element
-    SpecConstantDescriptor Desc;
-    Desc.ID = *IDIter;
-    Desc.Offset = Offset;
-    Desc.Size = M.getDataLayout().getTypeStoreSize(Ty);
-    Result.push_back(Desc);
-
-    // Move current ID and offset
-    ++IDIter;
-    Offset += Desc.Size;
+    return;
   }
+
+  // Assume that we encountered some scalar element
+  SpecConstantDescriptor Desc;
+  Desc.ID = IDIter->ID;
+  Desc.Offset = Offset;
+  Desc.Size = M.getDataLayout().getTypeStoreSize(Ty);
+  Result.push_back(Desc);
+
+  // Move current ID and offset
+  ++IDIter;
+  Offset += Desc.Size;
 }
 
 /// Recursively iterates over a composite type in order to collect information
@@ -400,7 +424,7 @@ void collectCompositeElementsDefaultValuesRecursive(
 }
 
 MDNode *generateSpecConstantMetadata(const Module &M, StringRef SymbolicID,
-                                     Type *SCTy, ArrayRef<unsigned> IDs,
+                                     Type *SCTy, ArrayRef<ID> IDs,
                                      bool IsNativeSpecConstant) {
   SmallVector<Metadata *, 16> MDOps;
   LLVMContext &Ctx = M.getContext();
@@ -413,7 +437,7 @@ MDNode *generateSpecConstantMetadata(const Module &M, StringRef SymbolicID,
     std::vector<SpecConstantDescriptor> Result;
     Result.reserve(IDs.size());
     unsigned Offset = 0;
-    const unsigned *IDPtr = IDs.data();
+    const ID *IDPtr = IDs.data();
     collectCompositeElementsInfoRecursive(M, SCTy, IDPtr, Offset, Result);
 
     // We may have padding elements so size should be at least the same size as
@@ -432,7 +456,7 @@ MDNode *generateSpecConstantMetadata(const Module &M, StringRef SymbolicID,
     assert(IDs.size() == 1 &&
            "There must be a single ID for emulated spec constant");
     MDOps.push_back(ConstantAsMetadata::get(
-        Constant::getIntegerValue(Int32Ty, APInt(32, IDs[0]))));
+        Constant::getIntegerValue(Int32Ty, APInt(32, IDs[0].ID))));
     // Second element is always zero here
     MDOps.push_back(ConstantAsMetadata::get(
         Constant::getIntegerValue(Int32Ty, APInt(32, 0))));
@@ -548,25 +572,35 @@ Instruction *emitSpecConstantComposite(Type *Ty, ArrayRef<Value *> Elements,
 /// composite (plus for the top-level composite). Also enumerates all
 /// encountered scalars and assigns them IDs (or re-uses existing ones).
 Instruction *emitSpecConstantRecursiveImpl(Type *Ty, Instruction *InsertBefore,
-                                           SmallVectorImpl<unsigned> &IDs,
+                                           SmallVectorImpl<ID> &IDs,
                                            unsigned &Index,
                                            Constant *DefaultValue) {
   if (!Ty->isArrayTy() && !Ty->isStructTy() && !Ty->isVectorTy()) { // Scalar
     if (Index >= IDs.size()) {
       // If it is a new specialization constant, we need to generate IDs for
       // scalar elements, starting with the second one.
-      IDs.push_back(IDs.back() + 1);
+      assert(!isa_and_nonnull<UndefValue>(DefaultValue) &&
+             "All scalar values should be defined");
+      IDs.push_back({IDs.back().ID + 1, false});
     }
-    return emitSpecConstant(IDs[Index++], Ty, InsertBefore, DefaultValue);
+    return emitSpecConstant(IDs[Index++].ID, Ty, InsertBefore, DefaultValue);
   }
 
   SmallVector<Value *, 8> Elements;
+  auto HandleUndef = [&](Constant *Def) {
+    if (Index >= IDs.size()) {
+      // If it is a new specialization constant, we need to generate IDs for
+      // the whole undef value.
+      IDs.push_back({IDs.back().ID + 1, true});
+    }
+    Elements.push_back(Def);
+  };
   auto LoopIteration = [&](Type *Ty, unsigned LocalIndex) {
     // Select corresponding element of the default value if it was provided
     Constant *Def =
         DefaultValue ? DefaultValue->getAggregateElement(LocalIndex) : nullptr;
     if (isa_and_nonnull<UndefValue>(Def))
-      Elements.push_back(Def);
+      HandleUndef(Def);
     else
       Elements.push_back(
           emitSpecConstantRecursiveImpl(Ty, InsertBefore, IDs, Index, Def));
@@ -576,7 +610,7 @@ Instruction *emitSpecConstantRecursiveImpl(Type *Ty, Instruction *InsertBefore,
     // If the default value is a composite and has the value 'undef', we should
     // not generate a bunch of __spirv_SpecConstant for its elements but
     // pass it into __spirv_SpecConstantComposite as is.
-    Elements.push_back(DefaultValue);
+    HandleUndef(DefaultValue);
   } else if (auto *ArrTy = dyn_cast<ArrayType>(Ty)) {
     for (size_t I = 0; I < ArrTy->getNumElements(); ++I) {
       LoopIteration(ArrTy->getElementType(), I);
@@ -599,7 +633,7 @@ Instruction *emitSpecConstantRecursiveImpl(Type *Ty, Instruction *InsertBefore,
 
 /// Wrapper intended to hide IsFirstElement argument from the caller
 Instruction *emitSpecConstantRecursive(Type *Ty, Instruction *InsertBefore,
-                                       SmallVectorImpl<unsigned> &IDs,
+                                       SmallVectorImpl<ID> &IDs,
                                        Constant *DefaultValue) {
   unsigned Index = 0;
   return emitSpecConstantRecursiveImpl(Ty, InsertBefore, IDs, Index,
@@ -610,9 +644,9 @@ Instruction *emitSpecConstantRecursive(Type *Ty, Instruction *InsertBefore,
 
 PreservedAnalyses SpecConstantsPass::run(Module &M,
                                          ModuleAnalysisManager &MAM) {
-  unsigned NextID = 0;
+  ID NextID = {0, false};
   unsigned NextOffset = 0;
-  StringMap<SmallVector<unsigned, 1>> IDMap;
+  StringMap<SmallVector<ID, 1>> IDMap;
   StringMap<unsigned> OffsetMap;
   MapVector<StringRef, MDNode *> SCMetadata;
   SmallVector<MDNode *, 4> DefaultsMetadata;
@@ -693,9 +727,8 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
       if (SetValAtRT) {
         // 2. Spec constant value will be set at run time - then add the literal
         // to a "spec const string literal ID" -> "vector of integer IDs" map,
-        // uniquing the integer IDs if this is a new literal
-        auto Ins =
-            IDMap.insert(std::make_pair(SymID, SmallVector<unsigned, 1>{}));
+        // making the integer IDs unique if this is a new literal
+        auto Ins = IDMap.insert(std::make_pair(SymID, SmallVector<ID, 1>{}));
         IsNewSpecConstant = Ins.second;
         auto &IDs = Ins.first->second;
         if (IsNewSpecConstant) {
@@ -711,7 +744,7 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
           // emitSpecConstantRecursive might emit more than one spec constant
           // (because of composite types) and therefore, we need to adjust
           // NextID according to the actual amount of emitted spec constants.
-          NextID += IDs.size();
+          NextID.ID += IDs.size();
 
           // Generate necessary metadata which later will be pulled by
           // sycl-post-link and transformed into device image properties
@@ -743,7 +776,7 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
             SCMetadata[SymID] = generateSpecConstantMetadata(
                 M, SymID, SCTy, NextID, /* is native spec constant */ false);
 
-            ++NextID;
+            ++NextID.ID;
             NextOffset += Size;
           }