Fix handling of multiple usage of composite spec constant

Also updated enumeration algorithm, so IDs are not reserved for __spirv_SpecConstantComposite entries anymore.
intel · Nov 26, 2020 · 9a82d53 · 9a82d53
1 parent 4c9ce32
commit 9a82d53
Show file tree

Hide file tree

Showing 3 changed files with 176 additions and 30 deletions.
diff --git a/llvm/test/tools/sycl-post-link/composite-spec-constant.ll b/llvm/test/tools/sycl-post-link/composite-spec-constant.ll
@@ -9,21 +9,21 @@
 ; CHECK: %[[#NS1:]] = call float @_Z20__spirv_SpecConstantif(i32 [[#ID + 1]], float
 ; CHECK: %[[#NA0:]] = call %struct._ZTS1A.A @_Z29__spirv_SpecConstantCompositeif(i32 %[[#NS0]], float %[[#NS1]])
 ;
-; CHECK: %[[#NS2:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 3]], i32
-; CHECK: %[[#NS3:]] = call float @_Z20__spirv_SpecConstantif(i32 [[#ID + 4]], float
+; CHECK: %[[#NS2:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 2]], i32
+; CHECK: %[[#NS3:]] = call float @_Z20__spirv_SpecConstantif(i32 [[#ID + 3]], float
 ; CHECK: %[[#NA1:]] = call %struct._ZTS1A.A @_Z29__spirv_SpecConstantCompositeif(i32 %[[#NS2]], float %[[#NS3]])
 ;
 ; CHECK: %[[#NA:]] = call [2 x %struct._ZTS1A.A] @_Z29__spirv_SpecConstantCompositestruct._ZTS1A.Astruct._ZTS1A.A(%struct._ZTS1A.A %[[#NA0]], %struct._ZTS1A.A %[[#NA1]])
 ;
-; CHECK: %[[#B0:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 7]], i32{{.*}})
-; CHECK: %[[#B1:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 8]], i32{{.*}})
+; CHECK: %[[#B0:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 4]], i32{{.*}})
+; CHECK: %[[#B1:]] = call i32 @_Z20__spirv_SpecConstantii(i32 [[#ID + 5]], i32{{.*}})
 ; CHECK: %[[#BV:]] = call <2 x i32> @_Z29__spirv_SpecConstantCompositeii(i32 %[[#B0]], i32 %[[#B1]])
 ; CHECK: %[[#B:]] = call %"class._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec" @_Z29__spirv_SpecConstantCompositeDv2_i(<2 x i32> %[[#BV]])
 ;
 ; CHECK: %[[#POD:]] = call %struct._ZTS3POD.POD @"_Z29__spirv_SpecConstantCompositeAstruct._ZTS1A.Aclass._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec"([2 x %struct._ZTS1A.A] %[[#NA]], %"class._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec" %[[#B]]), !SYCL_SPEC_CONST_SYM_ID ![[#MD:]]
 ; CHECK: store %struct._ZTS3POD.POD %[[#POD]]
 ;
-; CHECK: ![[#MD]] = !{!"_ZTS3POD", i32 [[#ID]], i32 [[#ID + 1]], i32 [[#ID + 3]], i32 [[#ID + 4]], i32 [[#ID + 7]], i32 [[#ID + 8]]}
+; CHECK: ![[#MD]] = !{!"_ZTS3POD", i32 [[#ID]], i32 [[#ID + 1]], i32 [[#ID + 2]], i32 [[#ID + 3]], i32 [[#ID + 4]], i32 [[#ID + 5]]}
 
 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-sycldevice"

diff --git a/llvm/test/tools/sycl-post-link/multiple-composite-spec-const-usages.ll b/llvm/test/tools/sycl-post-link/multiple-composite-spec-const-usages.ll
@@ -0,0 +1,98 @@
+; RUN: sycl-post-link -spec-const=rt --ir-output-only %s -S -o - \
+; RUN: | FileCheck %s --implicit-check-not __sycl_getCompositeSpecConstantValue
+;
+; This test is intended to check that sycl-post-link tool is capable of handling
+; situations when the same composite specialization constants is used more than
+; once
+;
+; CHECK-LABEL: @foo1
+; CHECK: call %struct._ZTS3POD.POD @"_Z29__spirv_SpecConstantCompositeAstruct._ZTS1A.Aclass._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec"({{.*}}), !SYCL_SPEC_CONST_SYM_ID ![[#MD0:]]
+; CHECK-LABEL: @_ZTS4Test
+; CHECK: call %struct._ZTS3POD.POD @"_Z29__spirv_SpecConstantCompositeAstruct._ZTS1A.Aclass._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec"({{.*}}), !SYCL_SPEC_CONST_SYM_ID ![[#MD1:]]
+; CHECK-LABEL: @foo2
+; CHECK: call %struct._ZTS3POD.POD @"_Z29__spirv_SpecConstantCompositeAstruct._ZTS1A.Aclass._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec"({{.*}}), !SYCL_SPEC_CONST_SYM_ID ![[#MD0:]]
+;
+; CHECK-DAG: ![[#MD0]] = !{!"_ZTS3PO2", i32 [[#ID:]],
+; CHECK-SAME: i32 [[#ID + 1]], i32 [[#ID + 2]], i32 [[#ID + 3]], i32 [[#ID + 4]], i32 [[#ID + 5]]}
+; CHECK-DAG: ![[#MD1]] = !{!"_ZTS3POD", i32 [[#ID1:]],
+; CHECK-SAME: i32 [[#ID1 + 1]], i32 [[#ID1 + 2]], i32 [[#ID1 + 3]], i32 [[#ID1 + 4]], i32 [[#ID1 + 5]]}
+
+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-sycldevice"
+
+%struct._ZTS3POD.POD = type { [2 x %struct._ZTS1A.A], %"class._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec" }
+%struct._ZTS1A.A = type { i32, float }
+%"class._ZTSN2cl4sycl3vecIiLi2EEE.cl::sycl::vec" = type { <2 x i32> }
+%"class._ZTSN2cl4sycl5rangeILi1EEE.cl::sycl::range" = type { %"class._ZTSN2cl4sycl6detail5arrayILi1EEE.cl::sycl::detail::array" }
+%"class._ZTSN2cl4sycl6detail5arrayILi1EEE.cl::sycl::detail::array" = type { [1 x i64] }
+%"class._ZTSN2cl4sycl2idILi1EEE.cl::sycl::id" = type { %"class._ZTSN2cl4sycl6detail5arrayILi1EEE.cl::sycl::detail::array" }
+
+$_ZTS4Test = comdat any
+
+@__builtin_unique_stable_name._ZNK2cl4sycl6ONEAPI12experimental13spec_constantI3PODS4_E3getIS4_EENSt9enable_ifIXsr3std6is_podIT_EE5valueES8_E4typeEv = private unnamed_addr addrspace(1) constant [9 x i8] c"_ZTS3POD\00", align 1
+@__builtin_unique_stable_name.2 = private unnamed_addr addrspace(1) constant [9 x i8] c"_ZTS3PO2\00", align 1
+
+define spir_func void @foo1() {
+  %ref.tmp.i = alloca %struct._ZTS3POD.POD, align 8
+  %1 = addrspacecast %struct._ZTS3POD.POD* %ref.tmp.i to %struct._ZTS3POD.POD addrspace(4)*
+  call spir_func void @_Z36__sycl_getCompositeSpecConstantValueI3PODET_PKc(%struct._ZTS3POD.POD addrspace(4)* sret align 8 %1, i8 addrspace(4)* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([9 x i8], [9 x i8] addrspace(1)* @__builtin_unique_stable_name.2, i64 0, i64 0) to i8 addrspace(4)*)) #4
+  ret void
+}
+
+; Function Attrs: convergent norecurse uwtable
+define weak_odr dso_local spir_kernel void @_ZTS4Test(%struct._ZTS3POD.POD addrspace(1)* %_arg_, %"class._ZTSN2cl4sycl5rangeILi1EEE.cl::sycl::range"* byval(%"class._ZTSN2cl4sycl5rangeILi1EEE.cl::sycl::range") align 8 %_arg_1, %"class._ZTSN2cl4sycl5rangeILi1EEE.cl::sycl::range"* byval(%"class._ZTSN2cl4sycl5rangeILi1EEE.cl::sycl::range") align 8 %_arg_2, %"class._ZTSN2cl4sycl2idILi1EEE.cl::sycl::id"* byval(%"class._ZTSN2cl4sycl2idILi1EEE.cl::sycl::id") align 8 %_arg_3) local_unnamed_addr #0 comdat !kernel_arg_buffer_location !4 {
+entry:
+  %ref.tmp.i = alloca %struct._ZTS3POD.POD, align 8
+  %0 = getelementptr inbounds %"class._ZTSN2cl4sycl2idILi1EEE.cl::sycl::id", %"class._ZTSN2cl4sycl2idILi1EEE.cl::sycl::id"* %_arg_3, i64 0, i32 0, i32 0, i64 0
+  %1 = load i64, i64* %0, align 8
+  %add.ptr.i = getelementptr inbounds %struct._ZTS3POD.POD, %struct._ZTS3POD.POD addrspace(1)* %_arg_, i64 %1
+  %2 = bitcast %struct._ZTS3POD.POD* %ref.tmp.i to i8*
+  call void @llvm.lifetime.start.p0i8(i64 24, i8* nonnull %2) #3
+  %3 = addrspacecast %struct._ZTS3POD.POD* %ref.tmp.i to %struct._ZTS3POD.POD addrspace(4)*
+  call spir_func void @_Z36__sycl_getCompositeSpecConstantValueI3PODET_PKc(%struct._ZTS3POD.POD addrspace(4)* sret align 8 %3, i8 addrspace(4)* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([9 x i8], [9 x i8] addrspace(1)* @__builtin_unique_stable_name._ZNK2cl4sycl6ONEAPI12experimental13spec_constantI3PODS4_E3getIS4_EENSt9enable_ifIXsr3std6is_podIT_EE5valueES8_E4typeEv, i64 0, i64 0) to i8 addrspace(4)*)) #4
+  %4 = bitcast %struct._ZTS3POD.POD addrspace(1)* %add.ptr.i to i8 addrspace(1)*
+  %5 = addrspacecast i8 addrspace(1)* %4 to i8 addrspace(4)*
+  call void @llvm.memcpy.p4i8.p0i8.i64(i8 addrspace(4)* align 8 dereferenceable(24) %5, i8* nonnull align 8 dereferenceable(24) %2, i64 24, i1 false), !tbaa.struct !5
+  call void @llvm.lifetime.end.p0i8(i64 24, i8* nonnull %2) #3
+  ret void
+}
+
+define spir_func void @foo2() {
+  %ref.tmp.i = alloca %struct._ZTS3POD.POD, align 8
+  %1 = addrspacecast %struct._ZTS3POD.POD* %ref.tmp.i to %struct._ZTS3POD.POD addrspace(4)*
+  call spir_func void @_Z36__sycl_getCompositeSpecConstantValueI3PODET_PKc(%struct._ZTS3POD.POD addrspace(4)* sret align 8 %1, i8 addrspace(4)* addrspacecast (i8 addrspace(1)* getelementptr inbounds ([9 x i8], [9 x i8] addrspace(1)* @__builtin_unique_stable_name.2, i64 0, i64 0) to i8 addrspace(4)*)) #4
+  ret void
+}
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.lifetime.start.p0i8(i64 immarg, i8* nocapture) #1
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.lifetime.end.p0i8(i64 immarg, i8* nocapture) #1
+
+; Function Attrs: argmemonly nounwind willreturn
+declare void @llvm.memcpy.p4i8.p0i8.i64(i8 addrspace(4)* noalias nocapture writeonly, i8* noalias nocapture readonly, i64, i1 immarg) #1
+
+; Function Attrs: convergent
+declare dso_local spir_func void @_Z36__sycl_getCompositeSpecConstantValueI3PODET_PKc(%struct._ZTS3POD.POD addrspace(4)* sret align 8, i8 addrspace(4)*) local_unnamed_addr #2
+
+attributes #0 = { convergent norecurse uwtable "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "min-legal-vector-width"="0" "no-infs-fp-math"="false" "no-jump-tables"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "sycl-module-id"="../sycl/test/spec_const/composite.cpp" "tune-cpu"="generic" "uniform-work-group-size"="true" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #1 = { argmemonly nounwind willreturn }
+attributes #2 = { convergent "disable-tail-calls"="false" "frame-pointer"="all" "less-precise-fpmad"="false" "no-infs-fp-math"="false" "no-nans-fp-math"="false" "no-signed-zeros-fp-math"="false" "no-trapping-math"="true" "stack-protector-buffer-size"="8" "tune-cpu"="generic" "unsafe-fp-math"="false" "use-soft-float"="false" }
+attributes #3 = { nounwind }
+attributes #4 = { convergent }
+
+!llvm.module.flags = !{!0}
+!opencl.spir.version = !{!1}
+!spirv.Source = !{!2}
+!llvm.ident = !{!3}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{i32 1, i32 2}
+!2 = !{i32 4, i32 100000}
+!3 = !{!"clang version 12.0.0 "}
+!4 = !{i32 -1, i32 -1, i32 -1, i32 -1}
+!5 = !{i64 0, i64 16, !6, i64 16, i64 8, !6}
+!6 = !{!7, !7, i64 0}
+!7 = !{!"omnipotent char", !8, i64 0}
+!8 = !{!"Simple C++ TBAA"}
diff --git a/llvm/tools/sycl-post-link/SpecConstants.cpp b/llvm/tools/sycl-post-link/SpecConstants.cpp
@@ -210,6 +210,8 @@ getScalarSpecConstMetadata(const Instruction *I) {
   return std::make_pair(MDSym->getString(), ID);
 }
 
+/// Recursively iterates over a composite type in order to collect information
+/// about its scalar elements.
 void collectCompositeElementsInfoRecursive(
     const Type *Ty, unsigned &Index, unsigned &Offset,
     std::vector<CompositeSpecConstElementDescriptor> &Result) {
@@ -303,21 +305,52 @@ Instruction *emitSpecConstantComposite(Type *Ty,
   return emitCall(Ty, SPIRV_GET_SPEC_CONST_COMPOSITE, Args, InsertBefore);
 }
 
-Instruction *
-emitSpecConstantRecursive(unsigned &NextID, Type *Ty,
-                          SmallVectorImpl<unsigned> &GeneratedScalarIDs,
-                          Instruction *InsertBefore) {
-  if (!Ty->isArrayTy() && !Ty->isStructTy() && !Ty->isVectorTy()) {
-    // assume that this is a scalar
-    GeneratedScalarIDs.push_back(NextID);
-    return emitSpecConstant(NextID, Ty, InsertBefore);
+/// For specified specialization constant type emits LLVM IR which is required
+/// in order to correctly handle it later during LLVM IR -> SPIR-V translation.
+///
+/// @param Ty [in] Specialization constant type to handle.
+/// @param InsertBefore [in] Location in the module where new instructions
+/// should be inserted.
+/// @param IDs [in,out] List of IDs which are assigned for scalar specialization
+/// constants. If \c IsNewSpecConstant is true, this vector is expected to
+/// contain a single element with ID of the first spec constant - the rest of
+/// generated spec constants will have their IDs generated by incrementing that
+/// first ID. If  \c IsNewSpecConstant is false, this vector is expected to
+/// contain enough elements to assign ID to each scalar element encountered in
+/// the specified composite type.
+/// @param IsNewSpecConstant [in] Flag to specify whether \c IDs vector should
+/// be filled with new IDs or it should be used as-is to replicate an existing
+/// spec constant
+/// @param [in,out] IsFirstElement Flag indicating whether this function is
+/// handling the first scalar element encountered in the specified composite
+/// type \c Ty or not.
+///
+/// @returns Instruction* representing specialization constant in LLVM IR, which
+/// is in SPIR-V friendly LLVM IR form.
+/// For scalar types it results in a single __spirv_SpecConstant call.
+/// For composite types it results in a number of __spirv_SpecConstant calls
+/// for each scalar member of the composite plus in a number of
+/// __spirvSpecConstantComposite calls for each composite member of the
+/// 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,
+                                           bool IsNewSpecConstant,
+                                           bool &IsFirstElement) {
+  if (!Ty->isArrayTy() && !Ty->isStructTy() && !Ty->isVectorTy()) { // Scalar
+    if (IsNewSpecConstant && !IsFirstElement) {
+      // 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);
+    }
+    IsFirstElement = false;
+    return emitSpecConstant(IDs.back(), Ty, InsertBefore);
   }
 
   SmallVector<Instruction *, 8> Elements;
   auto LoopIteration = [&](Type *Ty) {
-    ++NextID; // The first NextID is reserved for SpecConstantComposite below
-    Elements.push_back(emitSpecConstantRecursive(NextID, Ty, GeneratedScalarIDs,
-                                                 InsertBefore));
+    Elements.push_back(emitSpecConstantRecursiveImpl(
+        Ty, InsertBefore, IDs, IsNewSpecConstant, IsFirstElement));
   };
 
   if (auto *ArrTy = dyn_cast<ArrayType>(Ty)) {
@@ -339,12 +372,21 @@ emitSpecConstantRecursive(unsigned &NextID, Type *Ty,
   return emitSpecConstantComposite(Ty, Elements, InsertBefore);
 }
 
+/// Wrapper intended to hide IsFirstElement argument from the caller
+Instruction *emitSpecConstantRecursive(Type *Ty, Instruction *InsertBefore,
+                                       SmallVectorImpl<unsigned> &IDs,
+                                       bool IsNewSpecConstant) {
+  bool IsFirstElement = true;
+  return emitSpecConstantRecursiveImpl(Ty, InsertBefore, IDs, IsNewSpecConstant,
+                                       IsFirstElement);
+}
+
 } // namespace
 
 PreservedAnalyses SpecConstantsPass::run(Module &M,
                                          ModuleAnalysisManager &MAM) {
   unsigned NextID = 0;
-  StringMap<unsigned> IDMap;
+  StringMap<SmallVector<unsigned, 1>> IDMap;
 
   // Iterate through all declarations of instances of function template
   // template <typename T> T __sycl_getSpecConstantValue(const char *ID)
@@ -380,7 +422,7 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
       DelInsts.push_back(CI);
       Type *SCTy = CI->getType();
       unsigned NameArgNo = 0;
-      if (IsComposite) { // structs are returned via sret arguments
+      if (IsComposite) { // structs are returned via sret arguments.
         NameArgNo = 1;
         auto *PtrTy = cast<PointerType>(CI->getArgOperand(0)->getType());
         SCTy = PtrTy->getElementType();
@@ -389,22 +431,28 @@ 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" -> "integer ID" map, uniquing
-        // the integer ID if this is new literal
-        auto Ins = IDMap.insert(std::make_pair(SymID, 0));
-        if (Ins.second)
-          Ins.first->second = NextID;
-        unsigned ID = Ins.first->second;
+        // to a "spec const string literal ID" -> "integer ID" map or
+        // "composite spec const string literal ID" -> "vector of integer IDs"
+        // map, uniquing the integer IDs if this is new literal
+        auto Ins =
+            IDMap.insert(std::make_pair(SymID, SmallVector<unsigned, 1>{}));
+        bool IsNewSpecConstant = Ins.second;
+        auto &IDs = Ins.first->second;
+        if (IsNewSpecConstant) {
+          // For any spec constant type there will be always at least one ID
+          // generatedA.
+          IDs.push_back(NextID);
+        }
 
         //  3. Transform to spirv intrinsic _Z*__spirv_SpecConstant* or
         //  _Z*__spirv_SpecConstantComposite
-        SmallVector<unsigned, 4> GeneratedIDs;
-        auto *SPIRVCall = emitSpecConstantRecursive(ID, SCTy, GeneratedIDs, CI);
-        if (Ins.second) {
+        auto *SPIRVCall =
+            emitSpecConstantRecursive(SCTy, CI, IDs, IsNewSpecConstant);
+        if (IsNewSpecConstant) {
           // emitSpecConstantRecursive might emit more than one spec constant
           // (because of composite types) and therefore, we need to ajudst
-          // NextID according to the actual amount of emitted spec constants
-          NextID += GeneratedIDs.size();
+          // NextID according to the actual amount of emitted spec constants.
+          NextID += IDs.size();
         }
 
         if (IsComposite) {
@@ -418,7 +466,7 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
 
         // Mark the instruction with <symbolic_id, int_ids...> list for later
         // recollection by collectSpecConstantMetadata method.
-        setSpecConstSymIDMetadata(SPIRVCall, SymID, GeneratedIDs);
+        setSpecConstSymIDMetadata(SPIRVCall, SymID, IDs);
         // Example of the emitted call when spec constant is integer:
         // %6 = call i32 @_Z20__spirv_SpecConstantii(i32 0, i32 0), \
         //                                          !SYCL_SPEC_CONST_SYM_ID !22