nvfuser-cu121-torch25 0.2.25.dev20250201__cp312-cp312-manylinux_2_28_x86_64.whl

nvfuser/include/nvfuser/flatbuffers/grpc.h

@@ -0,0 +1,300 @@
+/*
+ * Copyright 2014 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FLATBUFFERS_GRPC_H_
+#define FLATBUFFERS_GRPC_H_
+
+// Helper functionality to glue FlatBuffers and GRPC.
+
+#include "flatbuffers/flatbuffers.h"
+#include "grpc/byte_buffer_reader.h"
+#include "grpcpp/support/byte_buffer.h"
+#include "grpcpp/support/slice.h"
+
+namespace flatbuffers {
+namespace grpc {
+
+// Message is a typed wrapper around a buffer that manages the underlying
+// `grpc_slice` and also provides flatbuffers-specific helpers such as `Verify`
+// and `GetRoot`. Since it is backed by a `grpc_slice`, the underlying buffer
+// is refcounted and ownership is be managed automatically.
+template<class T> class Message {
+ public:
+  Message() {}
+
+  Message(::grpc::Slice slice) : slice_(slice) {}
+
+  Message &operator=(const Message &other) = delete;
+
+  Message(Message &&other) = default;
+
+  Message(const Message &other) = delete;
+
+  Message &operator=(Message &&other) = default;
+
+  const uint8_t *mutable_data() const { return slice_.begin(); }
+
+  const uint8_t *data() const { return slice_.begin(); }
+
+  size_t size() const { return slice_.size(); }
+
+  bool Verify() const {
+    Verifier verifier(data(), size());
+    return verifier.VerifyBuffer<T>(nullptr);
+  }
+
+  T *GetMutableRoot() { return flatbuffers::GetMutableRoot<T>(mutable_data()); }
+
+  const T *GetRoot() const { return flatbuffers::GetRoot<T>(data()); }
+
+  // This is only intended for serializer use, or if you know what you're doing
+  const ::grpc::Slice &BorrowSlice() const { return slice_; }
+
+ private:
+  ::grpc::Slice slice_;
+};
+
+class MessageBuilder;
+
+// SliceAllocator is a gRPC-specific allocator that uses the `grpc_slice`
+// refcounted slices to manage memory ownership. This makes it easy and
+// efficient to transfer buffers to gRPC.
+class SliceAllocator : public Allocator {
+ public:
+  SliceAllocator() {}
+
+  SliceAllocator(const SliceAllocator &other) = delete;
+  SliceAllocator &operator=(const SliceAllocator &other) = delete;
+
+  SliceAllocator(SliceAllocator &&other) {
+    // default-construct and swap idiom
+    swap(other);
+  }
+
+  SliceAllocator &operator=(SliceAllocator &&other) {
+    // move-construct and swap idiom
+    SliceAllocator temp(std::move(other));
+    swap(temp);
+    return *this;
+  }
+
+  void swap(SliceAllocator &other) {
+    using std::swap;
+    swap(slice_, other.slice_);
+  }
+
+  virtual ~SliceAllocator() {}
+
+  virtual uint8_t *allocate(size_t size) override {
+    FLATBUFFERS_ASSERT(slice_.size() == 0);
+    slice_ = ::grpc::Slice(size);
+    return const_cast<uint8_t *>(slice_.begin());
+  }
+
+  virtual void deallocate(uint8_t *p, size_t size) override {
+    FLATBUFFERS_ASSERT(p == slice_.begin());
+    FLATBUFFERS_ASSERT(size == slice_.size());
+    slice_ = ::grpc::Slice();
+  }
+
+  virtual uint8_t *reallocate_downward(uint8_t *old_p, size_t old_size,
+                                       size_t new_size, size_t in_use_back,
+                                       size_t in_use_front) override {
+    FLATBUFFERS_ASSERT(old_p == slice_.begin());
+    FLATBUFFERS_ASSERT(old_size == slice_.size());
+    FLATBUFFERS_ASSERT(new_size > old_size);
+    ::grpc::Slice old_slice = slice_;
+    ::grpc::Slice new_slice = ::grpc::Slice(new_size);
+    uint8_t *new_p = const_cast<uint8_t *>(new_slice.begin());
+    memcpy_downward(old_p, old_size, new_p, new_size, in_use_back,
+                    in_use_front);
+    slice_ = new_slice;
+    return new_p;
+  }
+
+ private:
+  ::grpc::Slice &get_slice(uint8_t *p, size_t size) {
+    FLATBUFFERS_ASSERT(p == slice_.begin());
+    FLATBUFFERS_ASSERT(size == slice_.size());
+    return slice_;
+  }
+
+  ::grpc::Slice slice_;
+
+  friend class MessageBuilder;
+};
+
+// SliceAllocatorMember is a hack to ensure that the MessageBuilder's
+// slice_allocator_ member is constructed before the FlatBufferBuilder, since
+// the allocator is used in the FlatBufferBuilder ctor.
+namespace detail {
+struct SliceAllocatorMember {
+  SliceAllocator slice_allocator_;
+};
+}  // namespace detail
+
+// MessageBuilder is a gRPC-specific FlatBufferBuilder that uses SliceAllocator
+// to allocate gRPC buffers.
+class MessageBuilder : private detail::SliceAllocatorMember,
+                       public FlatBufferBuilder {
+ public:
+  explicit MessageBuilder(uoffset_t initial_size = 1024)
+      : FlatBufferBuilder(initial_size, &slice_allocator_, false) {}
+
+  MessageBuilder(const MessageBuilder &other) = delete;
+  MessageBuilder &operator=(const MessageBuilder &other) = delete;
+
+  MessageBuilder(MessageBuilder &&other)
+      : FlatBufferBuilder(1024, &slice_allocator_, false) {
+    // Default construct and swap idiom.
+    Swap(other);
+  }
+
+  /// Create a MessageBuilder from a FlatBufferBuilder.
+  explicit MessageBuilder(FlatBufferBuilder &&src,
+                          void (*dealloc)(void *,
+                                          size_t) = &DefaultAllocator::dealloc)
+      : FlatBufferBuilder(1024, &slice_allocator_, false) {
+    src.Swap(*this);
+    src.SwapBufAllocator(*this);
+    if (buf_.capacity()) {
+      uint8_t *buf = buf_.scratch_data();  // pointer to memory
+      size_t capacity = buf_.capacity();   // size of memory
+      slice_allocator_.slice_ = ::grpc::Slice(buf, capacity, dealloc);
+    } else {
+      slice_allocator_.slice_ = ::grpc::Slice();
+    }
+  }
+
+  /// Move-assign a FlatBufferBuilder to a MessageBuilder.
+  /// Only FlatBufferBuilder with default allocator (basically, nullptr) is
+  /// supported.
+  MessageBuilder &operator=(FlatBufferBuilder &&src) {
+    // Move construct a temporary and swap
+    MessageBuilder temp(std::move(src));
+    Swap(temp);
+    return *this;
+  }
+
+  MessageBuilder &operator=(MessageBuilder &&other) {
+    // Move construct a temporary and swap
+    MessageBuilder temp(std::move(other));
+    Swap(temp);
+    return *this;
+  }
+
+  void Swap(MessageBuilder &other) {
+    slice_allocator_.swap(other.slice_allocator_);
+    FlatBufferBuilder::Swap(other);
+    // After swapping the FlatBufferBuilder, we swap back the allocator, which
+    // restores the original allocator back in place. This is necessary because
+    // MessageBuilder's allocator is its own member (SliceAllocatorMember). The
+    // allocator passed to FlatBufferBuilder::vector_downward must point to this
+    // member.
+    buf_.swap_allocator(other.buf_);
+  }
+
+  // Releases the ownership of the buffer pointer.
+  // Returns the size, offset, and the original grpc_slice that
+  // allocated the buffer. Also see grpc_slice_unref().
+  uint8_t *ReleaseRaw(size_t &size, size_t &offset, ::grpc::Slice &slice) {
+    uint8_t *buf = FlatBufferBuilder::ReleaseRaw(size, offset);
+    slice = slice_allocator_.slice_;
+    slice_allocator_.slice_ = ::grpc::Slice();
+    return buf;
+  }
+
+  ~MessageBuilder() {}
+
+  // GetMessage extracts the subslice of the buffer corresponding to the
+  // flatbuffers-encoded region and wraps it in a `Message<T>` to handle buffer
+  // ownership.
+  template<class T> Message<T> GetMessage() {
+    auto buf_data = buf_.scratch_data();  // pointer to memory
+    auto buf_size = buf_.capacity();      // size of memory
+    auto msg_data = buf_.data();          // pointer to msg
+    auto msg_size = buf_.size();          // size of msg
+    // Do some sanity checks on data/size
+    FLATBUFFERS_ASSERT(msg_data);
+    FLATBUFFERS_ASSERT(msg_size);
+    FLATBUFFERS_ASSERT(msg_data >= buf_data);
+    FLATBUFFERS_ASSERT(msg_data + msg_size <= buf_data + buf_size);
+    // Calculate offsets from the buffer start
+    auto begin = msg_data - buf_data;
+    auto end = begin + msg_size;
+    // Get the slice we are working with (no refcount change)
+    ::grpc::Slice slice = slice_allocator_.get_slice(buf_data, buf_size);
+    // Extract a subslice of the existing slice (increment refcount)
+    ::grpc::Slice subslice = slice.sub(begin, end);
+    // Wrap the subslice in a `Message<T>`, but don't increment refcount
+    Message<T> msg(subslice);
+    return msg;
+  }
+
+  template<class T> Message<T> ReleaseMessage() {
+    Message<T> msg = GetMessage<T>();
+    Reset();
+    return msg;
+  }
+
+ private:
+  // SliceAllocator slice_allocator_;  // part of SliceAllocatorMember
+};
+
+}  // namespace grpc
+}  // namespace flatbuffers
+
+namespace grpc {
+
+template<class T> class SerializationTraits<flatbuffers::grpc::Message<T>> {
+ public:
+  static grpc::Status Serialize(const flatbuffers::grpc::Message<T> &msg,
+                                ByteBuffer *buffer, bool *own_buffer) {
+    // Package the single slice into a `ByteBuffer`,
+    // incrementing the refcount in the process.
+    *buffer = ByteBuffer(&msg.BorrowSlice(), 1);
+    *own_buffer = true;
+    return grpc::Status::OK;
+  }
+
+  // Deserialize by pulling the
+  static grpc::Status Deserialize(ByteBuffer *buf,
+                                  flatbuffers::grpc::Message<T> *msg) {
+    Slice slice;
+    if (!buf->TrySingleSlice(&slice).ok()) {
+      if (!buf->DumpToSingleSlice(&slice).ok()) {
+        buf->Clear();
+        return ::grpc::Status(::grpc::StatusCode::INTERNAL, "No payload");
+      }
+    }
+    *msg = flatbuffers::grpc::Message<T>(slice);
+    buf->Clear();
+#if FLATBUFFERS_GRPC_DISABLE_AUTO_VERIFICATION
+    return ::grpc::Status::OK;
+#else
+    if (msg->Verify()) {
+      return ::grpc::Status::OK;
+    } else {
+      return ::grpc::Status(::grpc::StatusCode::INTERNAL,
+                            "Message verification failed");
+    }
+#endif
+  }
+};
+
+}  // namespace grpc
+
+#endif  // FLATBUFFERS_GRPC_H_

nvfuser/include/nvfuser/flatbuffers/hash.h

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2015 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef FLATBUFFERS_HASH_H_
+#define FLATBUFFERS_HASH_H_
+
+#include <cstdint>
+#include <cstring>
+
+#include "flatbuffers/flatbuffers.h"
+
+namespace flatbuffers {
+
+template<typename T> struct FnvTraits {
+  static const T kFnvPrime;
+  static const T kOffsetBasis;
+};
+
+template<> struct FnvTraits<uint32_t> {
+  static const uint32_t kFnvPrime = 0x01000193;
+  static const uint32_t kOffsetBasis = 0x811C9DC5;
+};
+
+template<> struct FnvTraits<uint64_t> {
+  static const uint64_t kFnvPrime = 0x00000100000001b3ULL;
+  static const uint64_t kOffsetBasis = 0xcbf29ce484222645ULL;
+};
+
+template<typename T> T HashFnv1(const char *input) {
+  T hash = FnvTraits<T>::kOffsetBasis;
+  for (const char *c = input; *c; ++c) {
+    hash *= FnvTraits<T>::kFnvPrime;
+    hash ^= static_cast<unsigned char>(*c);
+  }
+  return hash;
+}
+
+template<typename T> T HashFnv1a(const char *input) {
+  T hash = FnvTraits<T>::kOffsetBasis;
+  for (const char *c = input; *c; ++c) {
+    hash ^= static_cast<unsigned char>(*c);
+    hash *= FnvTraits<T>::kFnvPrime;
+  }
+  return hash;
+}
+
+template<> inline uint16_t HashFnv1<uint16_t>(const char *input) {
+  uint32_t hash = HashFnv1<uint32_t>(input);
+  return (hash >> 16) ^ (hash & 0xffff);
+}
+
+template<> inline uint16_t HashFnv1a<uint16_t>(const char *input) {
+  uint32_t hash = HashFnv1a<uint32_t>(input);
+  return (hash >> 16) ^ (hash & 0xffff);
+}
+
+template<typename T> struct NamedHashFunction {
+  const char *name;
+
+  typedef T (*HashFunction)(const char *);
+  HashFunction function;
+};
+
+const NamedHashFunction<uint16_t> kHashFunctions16[] = {
+  { "fnv1_16", HashFnv1<uint16_t> },
+  { "fnv1a_16", HashFnv1a<uint16_t> },
+};
+
+const NamedHashFunction<uint32_t> kHashFunctions32[] = {
+  { "fnv1_32", HashFnv1<uint32_t> },
+  { "fnv1a_32", HashFnv1a<uint32_t> },
+};
+
+const NamedHashFunction<uint64_t> kHashFunctions64[] = {
+  { "fnv1_64", HashFnv1<uint64_t> },
+  { "fnv1a_64", HashFnv1a<uint64_t> },
+};
+
+inline NamedHashFunction<uint16_t>::HashFunction FindHashFunction16(
+    const char *name) {
+  std::size_t size = sizeof(kHashFunctions16) / sizeof(kHashFunctions16[0]);
+  for (std::size_t i = 0; i < size; ++i) {
+    if (std::strcmp(name, kHashFunctions16[i].name) == 0) {
+      return kHashFunctions16[i].function;
+    }
+  }
+  return nullptr;
+}
+
+inline NamedHashFunction<uint32_t>::HashFunction FindHashFunction32(
+    const char *name) {
+  std::size_t size = sizeof(kHashFunctions32) / sizeof(kHashFunctions32[0]);
+  for (std::size_t i = 0; i < size; ++i) {
+    if (std::strcmp(name, kHashFunctions32[i].name) == 0) {
+      return kHashFunctions32[i].function;
+    }
+  }
+  return nullptr;
+}
+
+inline NamedHashFunction<uint64_t>::HashFunction FindHashFunction64(
+    const char *name) {
+  std::size_t size = sizeof(kHashFunctions64) / sizeof(kHashFunctions64[0]);
+  for (std::size_t i = 0; i < size; ++i) {
+    if (std::strcmp(name, kHashFunctions64[i].name) == 0) {
+      return kHashFunctions64[i].function;
+    }
+  }
+  return nullptr;
+}
+
+}  // namespace flatbuffers
+
+#endif  // FLATBUFFERS_HASH_H_