[SYCL] Add property set types and JSON representation (#147321)

This PR adds the `PropertySet` type, along with a pair of functions used
to serialize and deserialize into a JSON representation. A property set
is a key-value map, with values being one of 2 types - uint32 or byte
array. A property set registry is a collection of property sets, indexed
by a "category" name.

In SYCL offloading, property sets will be used to communicate metadata
about device images needed by the SYCL runtime. For example, there is a
property set which has a byte array containing the numeric ID, offset,
and size of each SYCL2020 spec constant. Another example is a property
set describing the optional kernel features used in the module: does it
use fp64? fp16? atomic64?

This metadata will be computed by `clang-sycl-linker` and the JSON
representation will be inserted in the string table of each
output `OffloadBinary`. This JSON will be consumed the SYCL offload
wrapper and will be lowered to the binary form SYCL runtime expects.

For example, consider this SYCL program that calls a kernel that uses
fp64:

```c++
#include <sycl/sycl.hpp>

using namespace sycl;
class MyKernel;

int main() {
  queue q;
  auto *p = malloc_shared<double>(1, q);
  *p = .1;
  q.single_task<MyKernel>([=]{ *p *= 2; }).wait();
  std::cout << *p << "\n";
  free(p, q);
}
```

The device code for this program would have the kernel marked with
`!sycl_used_aspects`:

```
define spir_kernel void @_ZTS8MyKernel([...]) !sycl_used_aspects !n { [...] }
!n = {i32 6}
```

`clang-sycl-linker` would recognize this metadata and then would output
the following JSON in the `OffloadBinary`'s key-value map:

```
{
  "SYCL/device requirements": {
    // aspects contains a list of sycl::aspect values used
    // by the module; in this case just the value 6 encoded 
    // as a 4-byte little-endian integer 
    "aspects": "BjAwMA=="
  }
}
```

The SYCL offload wrapper would lower those property sets to something
like this:

```c++
struct _sycl_device_binary_property_set_struct {
  char *CategoryName;
  _sycl_device_binary_property *PropertiesBegin;
  _sycl_device_binary_property *PropertiesEnd;
};

struct _sycl_device_binary_property_struct {
  char *PropertyName;  
  void *ValAddr;     
  uint64_t ValSize; 
};

//  
_sycl_device_binary_property_struct device_requirements[] = {
  /* PropertyName */ "aspects",
  /* ValAddr */ [pointer to the bytes 0x06 0x00 0x00 0x00],
  /* ValSize */ 4,
};

_sycl_device_binary_property_set_struct properties[] = {
  /* CategoryName */ "SYCL/device requirements",
  /* PropertiesBegin */ device_requirements,
  /* PropertiesEnd */ std::end(device_requirments),
}
```

---------

Co-authored-by: Arvind Sudarsanam <[email protected]>

Author: jzc

llvm/include/llvm/Frontend/Offloading/PropertySet.h

@@ -0,0 +1,33 @@
+///===- llvm/Frontend/Offloading/PropertySet.h ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+///===---------------------------------------------------------------------===//
+/// \file This file defines PropertySetRegistry and PropertyValue types and
+/// provides helper functions to translate PropertySetRegistry from/to JSON.
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Error.h"
+
+#include <map>
+#include <variant>
+
+namespace llvm {
+class raw_ostream;
+class MemoryBufferRef;
+
+namespace offloading {
+
+using ByteArray = SmallVector<unsigned char, 0>;
+using PropertyValue = std::variant<uint32_t, ByteArray>;
+using PropertySet = std::map<std::string, PropertyValue>;
+using PropertySetRegistry = std::map<std::string, PropertySet>;
+
+void writePropertiesToJSON(const PropertySetRegistry &P, raw_ostream &O);
+Expected<PropertySetRegistry> readPropertiesFromJSON(MemoryBufferRef Buf);
+
+} // namespace offloading
+} // namespace llvm

llvm/lib/Frontend/Offloading/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_llvm_component_library(LLVMFrontendOffloading
   Utility.cpp
   OffloadWrapper.cpp
+  PropertySet.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Frontend

llvm/lib/Frontend/Offloading/PropertySet.cpp

@@ -0,0 +1,100 @@
+///===- llvm/Frontend/Offloading/PropertySet.cpp --------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Frontend/Offloading/PropertySet.h"
+#include "llvm/Support/Base64.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBufferRef.h"
+
+using namespace llvm;
+using namespace llvm::offloading;
+
+void llvm::offloading::writePropertiesToJSON(
+    const PropertySetRegistry &PSRegistry, raw_ostream &Out) {
+  json::OStream J(Out);
+  J.object([&] {
+    for (const auto &[CategoryName, PropSet] : PSRegistry) {
+      J.attributeObject(CategoryName, [&] {
+        for (const auto &[PropName, PropVal] : PropSet) {
+          switch (PropVal.index()) {
+          case 0:
+            J.attribute(PropName, std::get<uint32_t>(PropVal));
+            break;
+          case 1:
+            J.attribute(PropName, encodeBase64(std::get<ByteArray>(PropVal)));
+            break;
+          default:
+            llvm_unreachable("unsupported property type");
+          }
+        }
+      });
+    }
+  });
+}
+
+// note: createStringError has an overload that takes a format string,
+// but it uses llvm::format instead of llvm::formatv, which does
+// not work with json::Value. This is a helper function to use
+// llvm::formatv with createStringError.
+template <typename... Ts> auto createStringErrorV(Ts &&...Args) {
+  return createStringError(formatv(std::forward<Ts>(Args)...));
+}
+
+Expected<PropertyValue>
+readPropertyValueFromJSON(const json::Value &PropValueVal) {
+  if (std::optional<uint64_t> Val = PropValueVal.getAsUINT64())
+    return PropertyValue(static_cast<uint32_t>(*Val));
+
+  if (std::optional<StringRef> Val = PropValueVal.getAsString()) {
+    std::vector<char> Decoded;
+    if (Error E = decodeBase64(*Val, Decoded))
+      return createStringErrorV("unable to base64 decode the string {0}: {1}",
+                                Val, toString(std::move(E)));
+    return PropertyValue(ByteArray(Decoded.begin(), Decoded.end()));
+  }
+
+  return createStringErrorV("expected a uint64 or a string, got {0}",
+                            PropValueVal);
+}
+
+Expected<PropertySetRegistry>
+llvm::offloading::readPropertiesFromJSON(MemoryBufferRef Buf) {
+  PropertySetRegistry Res;
+  Expected<json::Value> V = json::parse(Buf.getBuffer());
+  if (Error E = V.takeError())
+    return E;
+
+  const json::Object *O = V->getAsObject();
+  if (!O)
+    return createStringErrorV(
+        "error while deserializing property set registry: "
+        "expected JSON object, got {0}",
+        *V);
+
+  for (const auto &[CategoryName, Value] : *O) {
+    const json::Object *PropSetVal = Value.getAsObject();
+    if (!PropSetVal)
+      return createStringErrorV("error while deserializing property set {0}: "
+                                "expected JSON array, got {1}",
+                                CategoryName.str(), Value);
+
+    PropertySet &PropSet = Res[CategoryName.str()];
+    for (const auto &[PropName, PropValueVal] : *PropSetVal) {
+      Expected<PropertyValue> Prop = readPropertyValueFromJSON(PropValueVal);
+      if (Error E = Prop.takeError())
+        return createStringErrorV(
+            "error while deserializing property {0} in property set {1}: {2}",
+            PropName.str(), CategoryName.str(), toString(std::move(E)));
+
+      auto [It, Inserted] =
+          PropSet.try_emplace(PropName.str(), std::move(*Prop));
+      assert(Inserted && "Property already exists in PropertySet");
+    }
+  }
+  return Res;
+}

llvm/unittests/Frontend/CMakeLists.txt

@@ -22,6 +22,7 @@ add_llvm_unittest(LLVMFrontendTests
   OpenMPDecompositionTest.cpp
   OpenMPDirectiveNameTest.cpp
   OpenMPDirectiveNameParserTest.cpp
+  PropertySetRegistryTest.cpp
 
   DEPENDS
   acc_gen

llvm/unittests/Frontend/PropertySetRegistryTest.cpp

@@ -0,0 +1,76 @@
+//===- llvm/unittest/Frontend/PropertySetRegistry.cpp ---------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Frontend/Offloading/PropertySet.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "gtest/gtest.h"
+
+using namespace llvm::offloading;
+using namespace llvm;
+
+void checkSerialization(const PropertySetRegistry &PSR) {
+  SmallString<0> Serialized;
+  raw_svector_ostream OS(Serialized);
+  writePropertiesToJSON(PSR, OS);
+  auto PSR2 = readPropertiesFromJSON({Serialized, ""});
+  ASSERT_EQ("", toString(PSR2.takeError()));
+  EXPECT_EQ(PSR, *PSR2);
+}
+
+TEST(PropertySetRegistryTest, PropertySetRegistry) {
+  PropertySetRegistry PSR;
+  checkSerialization(PSR);
+
+  PSR["Category1"]["Prop1"] = 42U;
+  PSR["Category1"]["Prop2"] = ByteArray(StringRef("Hello").bytes());
+  PSR["Category2"]["A"] = ByteArray{0, 4, 16, 32, 255};
+  checkSerialization(PSR);
+
+  PSR = PropertySetRegistry();
+  PSR["ABC"]["empty_array"] = ByteArray();
+  PSR["ABC"]["max_val"] = std::numeric_limits<uint32_t>::max();
+  checkSerialization(PSR);
+}
+
+TEST(PropertySetRegistryTest, IllFormedJSON) {
+  SmallString<0> Input;
+
+  // Invalid json
+  Input = "{ invalid }";
+  auto Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  Input = "";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  // Not a JSON object
+  Input = "[1, 2, 3]";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  // Property set not an object
+  Input = R"({ "Category": 42 })";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  // Property value has non string/non-integer type
+  Input = R"({ "Category": { "Prop": [1, 2, 3] } })";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  // Property value is an invalid base64 string
+  Input = R"({ "Category": { "Prop": ";" } })";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+
+  Input = R"({ "Category": { "Prop": "!@#$" } })";
+  Res = readPropertiesFromJSON({Input, ""});
+  EXPECT_NE("", toString(Res.takeError()));
+}