Unify representation of scalar and composite spec constants

Both scalar and composite spec constants are not communicated between
the device compiler and runtime in a single propery set using unified
format (the same as previously used for composite spec constants).

This change was suggested on code review here:
intel#2779 (comment)

Author: AlexeySachkov

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

@@ -198,23 +198,28 @@ void setSpecConstSymIDMetadata(Instruction *I, StringRef SymID,
   I->setMetadata(SPEC_CONST_SYM_ID_MD_STRING, Entry);
 }
 
-std::pair<StringRef, unsigned>
+std::pair<StringRef, std::vector<SpecConstantDescriptor>>
 getScalarSpecConstMetadata(const Instruction *I) {
   const MDNode *N = I->getMetadata(SPEC_CONST_SYM_ID_MD_STRING);
   if (!N)
-    return std::make_pair("", 0);
+    return std::make_pair("",
+                          std::vector<SpecConstantDescriptor>{});
   const auto *MDSym = cast<MDString>(N->getOperand(0));
   const auto *MDInt = cast<ConstantAsMetadata>(N->getOperand(1));
   unsigned ID = static_cast<unsigned>(
       cast<ConstantInt>(MDInt->getValue())->getValue().getZExtValue());
-  return std::make_pair(MDSym->getString(), ID);
+  std::vector<SpecConstantDescriptor> Res(1);
+  Res[0].ID = ID;
+  Res[0].Size = I->getType()->getPrimitiveSizeInBits() / /* bits in byte */8;
+  Res[0].Offset = 0;
+  return std::make_pair(MDSym->getString(), Res);
 }
 
 /// 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,
-    std::vector<CompositeSpecConstElementDescriptor> &Result) {
+    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
@@ -248,7 +253,7 @@ void collectCompositeElementsInfoRecursive(
                                             Offset, Result);
     }
   } else { // Assume that we encountered some scalar element
-    CompositeSpecConstElementDescriptor Desc;
+    SpecConstantDescriptor Desc;
     Desc.ID = 0; // To be filled later
     Desc.Offset = Offset;
     Desc.Size = Ty->getPrimitiveSizeInBits() / 8;
@@ -257,15 +262,15 @@ void collectCompositeElementsInfoRecursive(
   }
 }
 
-std::pair<StringRef, std::vector<CompositeSpecConstElementDescriptor>>
+std::pair<StringRef, std::vector<SpecConstantDescriptor>>
 getCompositeSpecConstMetadata(const Instruction *I) {
   const MDNode *N = I->getMetadata(SPEC_CONST_SYM_ID_MD_STRING);
   if (!N)
     return std::make_pair("",
-                          std::vector<CompositeSpecConstElementDescriptor>{});
+                          std::vector<SpecConstantDescriptor>{});
   const auto *MDSym = cast<MDString>(N->getOperand(0));
 
-  std::vector<CompositeSpecConstElementDescriptor> Result(N->getNumOperands() -
+  std::vector<SpecConstantDescriptor> Result(N->getNumOperands() -
                                                           1);
   unsigned Index = 0, Offset = 0;
   collectCompositeElementsInfoRecursive(I->getModule(), I->getType(), Index,
@@ -505,9 +510,8 @@ PreservedAnalyses SpecConstantsPass::run(Module &M,
   return IRModified ? PreservedAnalyses::none() : PreservedAnalyses::all();
 }
 
-bool SpecConstantsPass::collectSpecConstantMetadata(
-    Module &M, ScalarSpecIDMapTy &ScalarIDMap,
-    CompositeSpecIDMapTy &CompositeIDMap) {
+bool SpecConstantsPass::collectSpecConstantMetadata(Module &M,
+                                                    SpecIDMapTy &IDMap) {
   bool Met = false;
 
   for (Function &F : M) {
@@ -521,18 +525,16 @@ bool SpecConstantsPass::collectSpecConstantMetadata(
       if (!CI || CI->isIndirectCall() || !(Callee = CI->getCalledFunction()))
         continue;
 
+      std::pair<StringRef, std::vector<SpecConstantDescriptor>> Res;
       if (Callee->getName().contains(SPIRV_GET_SPEC_CONST_COMPOSITE)) {
-        auto Res = getCompositeSpecConstMetadata(CI);
-        if (!Res.first.empty()) {
-          CompositeIDMap[Res.first] = Res.second;
-          Met = true;
-        }
+        Res = getCompositeSpecConstMetadata(CI);
       } else if (Callee->getName().contains(SPIRV_GET_SPEC_CONST_VAL)) {
-        auto Res = getScalarSpecConstMetadata(CI);
-        if (!Res.first.empty()) {
-          ScalarIDMap[Res.first] = Res.second;
-          Met = true;
-        }
+        Res = getScalarSpecConstMetadata(CI);
+      }
+
+      if (!Res.first.empty()) {
+        IDMap[Res.first] = Res.second;
+        Met = true;
       }
     }
   }

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

@@ -24,24 +24,27 @@
 
 using namespace llvm;
 
-using ScalarSpecIDMapTy = std::map<StringRef, unsigned>;
-// Represents an element of a composite speciailization constant - at SYCL RT
-// level composite specialization constants are being represented as a single
-// byte-array, while at SPIR-V level they are represented by a number of scalar
-// specialization constants.
-struct CompositeSpecConstElementDescriptor {
+// Represents either an element of a composite speciailization constant or a
+// single scalar specialization constant - at SYCL RT level composite
+// specialization constants are being represented as a single byte-array, while
+// at SPIR-V level they are represented by a number of scalar specialization
+// constants.
+// The same representation is re-used for scalar specialization constants in
+// order to unify they processing with composite ones.
+struct SpecConstantDescriptor {
   // Encodes ID of a scalar specialization constants which is a leaf of some
   // composite specialization constant.
   unsigned ID;
   // Encodes offset from the beginning of composite, where scalar resides, i.e.
   // location of the scalar value within a byte-array containing the whole
-  // composite specialization constant.
+  // composite specialization constant. If descriptor is used to represent a
+  // whole scalar specialization constant instead of an element of a composite,
+  // this field should be contain zero.
   unsigned Offset;
   // Encodes size of scalar specialization constant.
   unsigned Size;
 };
-using CompositeSpecIDMapTy =
-    std::map<StringRef, std::vector<CompositeSpecConstElementDescriptor>>;
+using SpecIDMapTy = std::map<StringRef, std::vector<SpecConstantDescriptor>>;
 
 class SpecConstantsPass : public PassInfoMixin<SpecConstantsPass> {
 public:
@@ -57,9 +60,7 @@ class SpecConstantsPass : public PassInfoMixin<SpecConstantsPass> {
   // constants and
   // "spec constant name" -> vector<"spec constant int ID"> map for composite
   // spec constants
-  static bool collectSpecConstantMetadata(Module &M,
-                                          ScalarSpecIDMapTy &ScalarIDMap,
-                                          CompositeSpecIDMapTy &CompositeIDMap);
+  static bool collectSpecConstantMetadata(Module &M, SpecIDMapTy &IDMap);
 
 private:
   bool SetValAtRT;

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

@@ -545,16 +545,12 @@ static string_vector saveDeviceImageProperty(
     if (ImgPSInfo.DoSpecConst && ImgPSInfo.SetSpecConstAtRT) {
       if (ImgPSInfo.SpecConstsMet) {
         // extract spec constant maps per each module
-        ScalarSpecIDMapTy TmpScalarSpecIDMap;
-        CompositeSpecIDMapTy TmpCompositeSpecIDMap;
-        SpecConstantsPass::collectSpecConstantMetadata(
-            *ResultModules[I].get(), TmpScalarSpecIDMap, TmpCompositeSpecIDMap);
+        SpecIDMapTy TmpSpecIDMap;
+        SpecConstantsPass::collectSpecConstantMetadata(*ResultModules[I].get(),
+                                                       TmpSpecIDMap);
         PropSet.add(
             llvm::util::PropertySetRegistry::SYCL_SPECIALIZATION_CONSTANTS,
-            TmpScalarSpecIDMap);
-        PropSet.add(llvm::util::PropertySetRegistry::
-                        SYCL_COMPOSITE_SPECIALIZATION_CONSTANTS,
-                    TmpCompositeSpecIDMap);
+            TmpSpecIDMap);
       }
     }
     if (ImgPSInfo.EmitKernelParamInfo) {

sycl/include/CL/sycl/detail/pi.h

@@ -660,12 +660,8 @@ static const uint8_t PI_DEVICE_BINARY_OFFLOAD_KIND_SYCL = 4;
 /// Name must be consistent with
 /// PropertySetRegistry::SYCL_SPECIALIZATION_CONSTANTS defined in
 /// PropertySetIO.h
-#define __SYCL_PI_PROPERTY_SET_SCALAR_SPEC_CONST_MAP                           \
+#define __SYCL_PI_PROPERTY_SET_SPEC_CONST_MAP                           \
   "SYCL/specialization constants"
-/// PropertySetRegistry::SYCL_COMPOSITE_SPECIALIZATION_CONSTANTS defined in
-/// PropertySetIO.h
-#define __SYCL_PI_PROPERTY_SET_COMPOSITE_SPEC_CONST_MAP                        \
-  "SYCL/composite specialization constants"
 /// PropertySetRegistry::SYCL_DEVICELIB_REQ_MASK defined in PropertySetIO.h
 #define __SYCL_PI_PROPERTY_SET_DEVICELIB_REQ_MASK "SYCL/devicelib req mask"
 /// PropertySetRegistry::SYCL_KERNEL_PARAM_OPT_INFO defined in PropertySetIO.h

sycl/include/CL/sycl/detail/pi.hpp

@@ -283,20 +283,18 @@ class DeviceBinaryImage {
     return Format;
   }
 
-  /// Gets the iterator range over scalar specialization constants in this
-  /// binary image. For each property pointed to by an iterator within the
+  /// Gets the iterator range over specialization constants in this binary
+  /// image. For each property pointed to by an iterator within the
   /// range, the name of the property is the specialization constant symbolic ID
-  /// and the value is 32-bit unsigned integer ID.
-  const PropertyRange &getScalarSpecConstants() const {
-    return ScalarSpecConstIDMap;
-  }
-  /// Gets the iterator range over composite specialization constants in this
-  /// binary image. For each property pointed to by an iterator within the
-  /// range, the name of the property is the specialization constant symbolic ID
-  /// and the value is a list of tuples of 32-bit unsigned integer values, which
-  /// encode scalar specialization constants, that form the composite one.
+  /// and the value is a list of 3-element tuples of 32-bit unsigned integers,
+  /// describing the specialization constant.
+  /// This is done in order to unify representation of both scalar and composite
+  /// specialization constants: composite specialization constant is represented
+  /// by its leaf elements, so for scalars the list contains only a single
+  /// tuple, while for composite there might be more of them.
   /// Each tuple consists of ID of scalar specialization constant, its location
-  /// within a composite (offset in bytes from the beginning) and its size.
+  /// within a composite (offset in bytes from the beginning or 0 if it is not
+  /// an element of a composite specialization constant) and its size.
   /// For example, for the following structure:
   /// struct A { int a; float b; };
   /// struct POD { A a[2]; int b; };
@@ -306,8 +304,11 @@ class DeviceBinaryImage {
   /// { ID2, 8, 4 },  // .a[1].a
   /// { ID3, 12, 4 }, // .a[1].b
   /// { ID4, 16, 4 }, // .b
-  const PropertyRange &getCompositeSpecConstants() const {
-    return CompositeSpecConstIDMap;
+  /// And for an interger specialization constant, the list of tuples will look
+  /// like:
+  /// { ID5, 0, 4 }
+  const PropertyRange &getSpecConstants() const {
+    return SpecConstIDMap;
   }
   const PropertyRange &getDeviceLibReqMask() const { return DeviceLibReqMask; }
   const PropertyRange &getKernelParamOptInfo() const {
@@ -321,8 +322,7 @@ class DeviceBinaryImage {
 
   pi_device_binary Bin;
   pi::PiDeviceBinaryType Format = PI_DEVICE_BINARY_TYPE_NONE;
-  DeviceBinaryImage::PropertyRange ScalarSpecConstIDMap;
-  DeviceBinaryImage::PropertyRange CompositeSpecConstIDMap;
+  DeviceBinaryImage::PropertyRange SpecConstIDMap;
   DeviceBinaryImage::PropertyRange DeviceLibReqMask;
   DeviceBinaryImage::PropertyRange KernelParamOptInfo;
 };

sycl/source/detail/pi.cpp

@@ -593,9 +593,7 @@ void DeviceBinaryImage::init(pi_device_binary Bin) {
     // try to determine the format; may remain "NONE"
     Format = getBinaryImageFormat(Bin->BinaryStart, getSize());
 
-  ScalarSpecConstIDMap.init(Bin, __SYCL_PI_PROPERTY_SET_SCALAR_SPEC_CONST_MAP);
-  CompositeSpecConstIDMap.init(Bin,
-                               __SYCL_PI_PROPERTY_SET_COMPOSITE_SPEC_CONST_MAP);
+  SpecConstIDMap.init(Bin, __SYCL_PI_PROPERTY_SET_SPEC_CONST_MAP);
   DeviceLibReqMask.init(Bin, __SYCL_PI_PROPERTY_SET_DEVICELIB_REQ_MASK);
   KernelParamOptInfo.init(Bin, __SYCL_PI_PROPERTY_SET_KERNEL_PARAM_OPT_INFO);
 }

sycl/source/detail/program_impl.cpp

@@ -523,29 +523,14 @@ void program_impl::flush_spec_constants(const RTDeviceBinaryImage &Img,
                                         RT::PiProgram NativePrg) const {
   // iterate via all specialization constants the program's image depends on,
   // and set each to current runtime value (if any)
-  const pi::DeviceBinaryImage::PropertyRange &ScalarSCRange =
-      Img.getScalarSpecConstants();
-  const pi::DeviceBinaryImage::PropertyRange &CompositeSCRange =
-      Img.getCompositeSpecConstants();
+  const pi::DeviceBinaryImage::PropertyRange &SCRange = Img.getSpecConstants();
   ContextImplPtr Ctx = getSyclObjImpl(get_context());
   using SCItTy = pi::DeviceBinaryImage::PropertyRange::ConstIterator;
 
   auto LockGuard = Ctx->getKernelProgramCache().acquireCachedPrograms();
   NativePrg = NativePrg ? NativePrg : getHandleRef();
 
-  for (SCItTy SCIt : ScalarSCRange) {
-    auto SCEntry = SpecConstRegistry.find((*SCIt)->Name);
-    if (SCEntry == SpecConstRegistry.end())
-      // spec constant has not been set in user code - SPIR-V will use default
-      continue;
-    const spec_constant_impl &SC = SCEntry->second;
-    assert(SC.isSet() && "uninitialized spec constant");
-    pi_uint32 ID = pi::DeviceBinaryProperty(*SCIt).asUint32();
-    Ctx->getPlugin().call<PiApiKind::piextProgramSetSpecializationConstant>(
-        NativePrg, ID, SC.getSize(), SC.getValuePtr());
-  }
-
-  for (SCItTy SCIt : CompositeSCRange) {
+  for (SCItTy SCIt : SCRange) {
     auto SCEntry = SpecConstRegistry.find((*SCIt)->Name);
     if (SCEntry == SpecConstRegistry.end())
       // spec constant has not been set in user code - SPIR-V will use default
@@ -557,9 +542,9 @@ void program_impl::flush_spec_constants(const RTDeviceBinaryImage &Img,
     assert(Descriptors.size() > 8 && "Unexpected property size");
     // Expected layout is vector of 3-component tuples (flattened into a vector
     // of scalars), where each tuple consists of: ID of a scalar spec constant,
-    // which is a member of the composite; offset, which is used to calculate
-    // location of scalar member within the composite; size of a scalar member
-    // of the composite.
+    // (which might be a member of the composite); offset, which is used to
+    // calculate location of scalar member within the composite or zero for
+    // scalar spec constants; size of a spec constant
     assert(((Descriptors.size() - 8) / sizeof(std::uint32_t)) % 3 == 0 &&
            "unexpected layout of composite spec const descriptors");
     auto *It = reinterpret_cast<const std::uint32_t *>(&Descriptors[8]);

sycl/test/basic_tests/spec_const_types.cpp

@@ -44,35 +44,35 @@ int main() {
   // Create specialization constants.
   cl::sycl::ONEAPI::experimental::spec_constant<bool, MyBoolConst> i1 =
       program.set_spec_constant<MyBoolConst>((bool)get_value());
-  // CHECK-DAG: _ZTS11MyBoolConst=1|0
+  // CHECK-DAG: _ZTS11MyBoolConst=2|
 
   cl::sycl::ONEAPI::experimental::spec_constant<int8_t, MyInt8Const> i8 =
       program.set_spec_constant<MyInt8Const>((int8_t)get_value());
-  // CHECK-DAG: _ZTS11MyInt8Const=1|1
+  // CHECK-DAG: _ZTS11MyInt8Const=2|
   cl::sycl::ONEAPI::experimental::spec_constant<uint8_t, MyUInt8Const> ui8 =
       program.set_spec_constant<MyUInt8Const>((uint8_t)get_value());
-  // CHECK-DAG: _ZTS12MyUInt8Const=1|2
+  // CHECK-DAG: _ZTS12MyUInt8Const=2|
 
   cl::sycl::ONEAPI::experimental::spec_constant<int16_t, MyInt16Const> i16 =
       program.set_spec_constant<MyInt16Const>((int16_t)get_value());
-  // CHECK-DAG: _ZTS12MyInt16Const=1|3
+  // CHECK-DAG: _ZTS12MyInt16Const=2|
   cl::sycl::ONEAPI::experimental::spec_constant<uint16_t, MyUInt16Const> ui16 =
       program.set_spec_constant<MyUInt16Const>((uint16_t)get_value());
-  // CHECK-DAG: _ZTS13MyUInt16Const=1|4
+  // CHECK-DAG: _ZTS13MyUInt16Const=2|
 
   cl::sycl::ONEAPI::experimental::spec_constant<int32_t, MyInt32Const> i32 =
       program.set_spec_constant<MyInt32Const>((int32_t)get_value());
-  // CHECK-DAG: _ZTS12MyInt32Const=1|5
+  // CHECK-DAG: _ZTS12MyInt32Const=2|
   cl::sycl::ONEAPI::experimental::spec_constant<uint32_t, MyUInt32Const> ui32 =
       program.set_spec_constant<MyUInt32Const>((uint32_t)get_value());
-  // CHECK-DAG: _ZTS13MyUInt32Const=1|6
+  // CHECK-DAG: _ZTS13MyUInt32Const=2|
 
   cl::sycl::ONEAPI::experimental::spec_constant<int64_t, MyInt64Const> i64 =
       program.set_spec_constant<MyInt64Const>((int64_t)get_value());
-  // CHECK-DAG: _ZTS12MyInt64Const=1|7
+  // CHECK-DAG: _ZTS12MyInt64Const=2|
   cl::sycl::ONEAPI::experimental::spec_constant<uint64_t, MyUInt64Const> ui64 =
       program.set_spec_constant<MyUInt64Const>((uint64_t)get_value());
-  // CHECK-DAG: _ZTS13MyUInt64Const=1|8
+  // CHECK-DAG: _ZTS13MyUInt64Const=2|
 
 #define HALF 0 // TODO not yet supported
 #if HALF
@@ -82,11 +82,11 @@ int main() {
 
   cl::sycl::ONEAPI::experimental::spec_constant<float, MyFloatConst> f32 =
       program.set_spec_constant<MyFloatConst>((float)get_value());
-  // CHECK-DAG: _ZTS12MyFloatConst=1|9
+  // CHECK-DAG: _ZTS12MyFloatConst=2|
 
   cl::sycl::ONEAPI::experimental::spec_constant<double, MyDoubleConst> f64 =
       program.set_spec_constant<MyDoubleConst>((double)get_value());
-  // CHECK-DAG: _ZTS13MyDoubleConst=1|10
+  // CHECK-DAG: _ZTS13MyDoubleConst=2|
 
   program.build_with_kernel_type<SpecializedKernel>();