[SYCL][sycl-post-link] Add post-struct padding in descs and fix blob size calculations (#5376)

These changes fix two related issues:
1. The runtime does not correctly handle the offset in specialization
   constant descriptors as it does not take into account the size of the
   previous element in composite types, effectively always adding
   padding between each element equal to at least the size of the
   previous element. This is fixed by changing it to keep track of the
   local offset into the current composite type and subtracting that
   from the offset of the descriptor to find the required padding.
2. Composite types may have padding at the end. sycl-post-link does not
   currently generate enough information for the runtime to take the
   additional padding into account when calculating the size of the
   specialization constant default value blob. This is fixed by
   inserting an additional descriptor after the last element of
   a composite type with end-padding. This "padding-descriptor" will
   have the offset right after the last element of the composite type
   and a size corresponding to the size of the required padding.
   padding-descriptors all have the same ID, namely the maximum value
   of the (unsigned) ID.

Signed-off-by: Steffen Larsen <steffen.larsen@intel.com>

Author: steffenlarsen

llvm/test/tools/sycl-post-link/spec-constants/composite-padding-desc.ll

@@ -0,0 +1,73 @@
+; RUN: sycl-post-link -spec-const=rt %s -o %t.files.table
+; RUN: FileCheck %s -input-file=%t.files_0.prop
+;
+; This test checks that composite specialization constants with implicit padding
+; at the end of the composite type will have an additional padding descriptor at
+; the end of the descriptor list. 
+
+; ModuleID = 'test.bc'
+source_filename = "llvm-link"
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64"
+target triple = "spir64-unknown-unknown"
+
+%"class.cl::sycl::specialization_id" = type { %struct.TestStruct }
+%struct.TestStruct = type { i32, i8 }
+
+$_ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E10KernelName = comdat any
+
+@__usid_str = private unnamed_addr constant [33 x i8] c"fb86570d411366d1____ZL9SpecConst\00", align 1
+@_ZL9SpecConst = internal addrspace(1) constant %"class.cl::sycl::specialization_id" zeroinitializer, align 4
+
+; Function Attrs: convergent norecurse
+define weak_odr dso_local spir_kernel void @_ZTSZZ4mainENKUlRN2cl4sycl7handlerEE_clES2_E10KernelName() local_unnamed_addr #0 comdat !kernel_arg_buffer_location !5 !sycl_kernel_omit_args !6 {
+entry:
+  %tmp.i = alloca %struct.TestStruct, align 4
+  %tmp.ascast.i = addrspacecast %struct.TestStruct* %tmp.i to %struct.TestStruct addrspace(4)*
+  %0 = bitcast %struct.TestStruct* %tmp.i to i8*
+  call void @llvm.lifetime.start.p0i8(i64 8, i8* nonnull %0) #3
+  call spir_func void @_Z40__sycl_getComposite2020SpecConstantValueI10TestStructET_PKcPKvS5_(%struct.TestStruct addrspace(4)* sret(%struct.TestStruct) align 4 %tmp.ascast.i, i8 addrspace(4)* addrspacecast (i8* getelementptr inbounds ([33 x i8], [33 x i8]* @__usid_str, i64 0, i64 0) to i8 addrspace(4)*), i8 addrspace(4)* addrspacecast (i8 addrspace(1)* bitcast (%"class.cl::sycl::specialization_id" addrspace(1)* @_ZL9SpecConst to i8 addrspace(1)*) to i8 addrspace(4)*), i8 addrspace(4)* null) #4
+  call void @llvm.lifetime.end.p0i8(i64 8, i8* nonnull %0) #3
+  ret void
+}
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #1
+
+; Function Attrs: convergent
+declare dso_local spir_func void @_Z40__sycl_getComposite2020SpecConstantValueI10TestStructET_PKcPKvS5_(%struct.TestStruct addrspace(4)* sret(%struct.TestStruct) align 4, i8 addrspace(4)*, i8 addrspace(4)*, i8 addrspace(4)*) local_unnamed_addr #2
+
+; Function Attrs: argmemonly nofree nosync nounwind willreturn
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #1
+
+attributes #0 = { convergent norecurse "frame-pointer"="all" "min-legal-vector-width"="0" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "sycl-module-id"="test.cpp" "uniform-work-group-size"="true" }
+attributes #1 = { argmemonly nofree nosync nounwind willreturn }
+attributes #2 = { convergent "frame-pointer"="all" "no-trapping-math"="true" "stack-protector-buffer-size"="8" }
+attributes #3 = { nounwind }
+attributes #4 = { convergent }
+
+!opencl.spir.version = !{!0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0, !0}
+!spirv.Source = !{!1, !1, !1, !1, !1, !1, !1, !1, !1, !1, !1, !1}
+!llvm.ident = !{!2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2, !2}
+!llvm.module.flags = !{!3, !4}
+
+!0 = !{i32 1, i32 2}
+!1 = !{i32 4, i32 100000}
+!2 = !{!"clang version 14.0.0"}
+!3 = !{i32 1, !"wchar_size", i32 4}
+!4 = !{i32 7, !"frame-pointer", i32 2}
+!5 = !{i32 -1}
+!6 = !{i1 true}
+
+; We expect the following in the descriptor list for SpecConst:
+; First 8 bytes are the size of the list.
+; Each 12 bytes after the size comprise a descriptor consisting of 3 32-bit
+; unsigned integers. For SpecConst these are:
+;            ID             |     Composite offset    |          Size
+; 0x00 0x00 0x00 0x00 (0)   | 0x00 0x00 0x00 0x00 (0) | 0x04 0x00 0x00 0x00 (4)
+; 0x01 0x00 0x00 0x00 (1)   | 0x04 0x00 0x00 0x00 (4) | 0x01 0x00 0x00 0x00 (1)
+; 0xff 0xff 0xff 0xff (max) | 0x05 0x00 0x00 0x00 (5) | 0x03 0x00 0x00 0x00 (3)
+; Most important for this test is the last descriptor which represents 3-bytes
+; implicit padding at the end of the composite type of the spec constant.
+;
+; CHECK: [SYCL/specialization constants]
+; CHECK-NEXT: fb86570d411366d1____ZL9SpecConst=2|gEAAAAAAAAAAAAAAAAAAAQAAAAQAAAAAEAAAAEAAAAw/////FAAAAMAAAAA

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

@@ -236,29 +236,44 @@ 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, unsigned &Index, unsigned &Offset,
+    const Module &M, Type *Ty, const unsigned *&IDIter, unsigned &Offset,
     std::vector<SpecConstantDescriptor> &Result) {
   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
       // just make a single recursive call here and use it to populate Result
       // in a loop.
-      collectCompositeElementsInfoRecursive(M, ArrTy->getElementType(), Index,
+      collectCompositeElementsInfoRecursive(M, ArrTy->getElementType(), IDIter,
                                             Offset, Result);
     }
   } else if (auto *StructTy = dyn_cast<StructType>(Ty)) {
     const StructLayout *SL = M.getDataLayout().getStructLayout(StructTy);
+    const unsigned BaseOffset = Offset;
+    unsigned LocalOffset = Offset;
     for (size_t I = 0, E = StructTy->getNumElements(); I < E; ++I) {
       auto *ElTy = StructTy->getElementType(I);
       // When handling elements of a structure, we do not use manually
       // calculated offsets (which are sum of sizes of all previously
       // encountered elements), but instead rely on data provided for us by
       // DataLayout, because the structure can be unpacked, i.e. padded in
       // order to ensure particular alignment of its elements.
-      unsigned LocalOffset = Offset + SL->getElementOffset(I);
-      collectCompositeElementsInfoRecursive(M, ElTy, Index, LocalOffset,
+      LocalOffset = Offset + SL->getElementOffset(I);
+      collectCompositeElementsInfoRecursive(M, ElTy, IDIter, LocalOffset,
                                             Result);
     }
+
+    // Add a special descriptor if the struct has padding at the end.
+    const unsigned PostStructPadding =
+        BaseOffset + SL->getSizeInBytes() - LocalOffset;
+    if (PostStructPadding > 0) {
+      SpecConstantDescriptor Desc;
+      // ID of padding descriptors is the max value possible.
+      Desc.ID = std::numeric_limits<unsigned>::max();
+      Desc.Offset = LocalOffset;
+      Desc.Size = PostStructPadding;
+      Result.push_back(Desc);
+    }
+
     // Update "global" offset according to the total size of a handled struct
     // type.
     Offset += SL->getSizeInBytes();
@@ -267,15 +282,18 @@ void collectCompositeElementsInfoRecursive(
       // 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(), Index,
+      collectCompositeElementsInfoRecursive(M, VecTy->getElementType(), IDIter,
                                             Offset, Result);
     }
   } else { // Assume that we encountered some scalar element
     SpecConstantDescriptor Desc;
-    Desc.ID = 0; // To be filled later
+    Desc.ID = *IDIter;
     Desc.Offset = Offset;
     Desc.Size = M.getDataLayout().getTypeStoreSize(Ty);
-    Result[Index++] = Desc;
+    Result.push_back(Desc);
+
+    // Move current ID and offset
+    ++IDIter;
     Offset += Desc.Size;
   }
 }
@@ -392,13 +410,19 @@ MDNode *generateSpecConstantMetadata(const Module &M, StringRef SymbolicID,
   MDOps.push_back(MDString::get(Ctx, SymbolicID));
 
   if (IsNativeSpecConstant) {
-    std::vector<SpecConstantDescriptor> Result(IDs.size());
-    unsigned Index = 0, Offset = 0;
-    collectCompositeElementsInfoRecursive(M, SCTy, Index, Offset, Result);
+    std::vector<SpecConstantDescriptor> Result;
+    Result.reserve(IDs.size());
+    unsigned Offset = 0;
+    const unsigned *IDPtr = IDs.data();
+    collectCompositeElementsInfoRecursive(M, SCTy, IDPtr, Offset, Result);
+
+    // We may have padding elements so size should be at least the same size as
+    // the ID vector.
+    assert(Result.size() >= IDs.size());
 
     for (unsigned I = 0; I < Result.size(); ++I) {
       MDOps.push_back(ConstantAsMetadata::get(
-          Constant::getIntegerValue(Int32Ty, APInt(32, IDs[I]))));
+          Constant::getIntegerValue(Int32Ty, APInt(32, Result[I].ID))));
       MDOps.push_back(ConstantAsMetadata::get(
           Constant::getIntegerValue(Int32Ty, APInt(32, Result[I].Offset))));
       MDOps.push_back(ConstantAsMetadata::get(

sycl/source/detail/device_image_impl.hpp

@@ -267,16 +267,22 @@ class device_image_impl {
         auto *It = reinterpret_cast<const std::uint32_t *>(&Descriptors[8]);
         auto *End = reinterpret_cast<const std::uint32_t *>(&Descriptors[0] +
                                                             Descriptors.size());
-        unsigned PrevOffset = 0;
+        unsigned LocalOffset = 0;
         while (It != End) {
           // Make sure that alignment is correct in blob.
-          BlobOffset += /*Offset*/ It[1] - PrevOffset;
-          PrevOffset = It[1];
-          // The map is not locked here because updateSpecConstSymMap() is only
-          // supposed to be called from c'tor.
-          MSpecConstSymMap[std::string{SCName}].push_back(
-              SpecConstDescT{/*ID*/ It[0], /*CompositeOffset*/ It[1],
-                             /*Size*/ It[2], BlobOffset});
+          const unsigned OffsetFromLast = /*Offset*/ It[1] - LocalOffset;
+          BlobOffset += OffsetFromLast;
+          // Composites may have a special padding element at the end which
+          // should not have a descriptor. These padding elements all have max
+          // ID value.
+          if (It[0] != std::numeric_limits<std::uint32_t>::max()) {
+            // The map is not locked here because updateSpecConstSymMap() is
+            // only supposed to be called from c'tor.
+            MSpecConstSymMap[std::string{SCName}].push_back(
+                SpecConstDescT{/*ID*/ It[0], /*CompositeOffset*/ It[1],
+                               /*Size*/ It[2], BlobOffset});
+          }
+          LocalOffset += OffsetFromLast + /*Size*/ It[2];
           BlobOffset += /*Size*/ It[2];
           It += NumElements;
         }
@@ -288,6 +294,9 @@ class device_image_impl {
       if (HasDefaultValues) {
         pi::ByteArray DefValDescriptors =
             pi::DeviceBinaryProperty(*SCDefValRange.begin()).asByteArray();
+        assert(DefValDescriptors.size() - 8 == MSpecConstsBlob.size() &&
+               "Specialization constant default value blob do not have the "
+               "expected size.");
         std::uninitialized_copy(&DefValDescriptors[8],
                                 &DefValDescriptors[8] + MSpecConstsBlob.size(),
                                 MSpecConstsBlob.data());