NVIDIA / NVIDIA/cuvs

Using a wrapper `byte_array` to support both `uint8_t` and `int8_t`

Open
#1,418 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Cuda
Stars
854
Forks
236
Avg merge
3d 3h
Merged PRs (30d)
62

Description

Currently, libcuvs instantiates both uint8_t and int8_t which causes binary size bloat. Given the relative similarity of both types and the ease of normalization of values that are encompassed by them, I propose that we combine the two types using a wrapper cuvs::detail::byte_array.

The public API remains the same:

namespace cuvs {

void foo(raft::device_matrix_view<uint8_t, ...>(...);

void foo(raft::device_matrix_view<int8_t, ...>(...);

}

The definition of cuvs::detail::byte_array, which is a struct that would always normalize int8_t values to uint8_t when returning and de-normalize when assigning:

namespace cuvs::detail {

struct byte_array {
  void* data = nullptr;
  bool is_signed = false;

  byte_array(void* ptr, bool signed_flag)
      : data(ptr), is_signed(signed_flag) {}

  // Proxy that references an element in the array
  struct byte {
    byte_array* parent = nullptr;
    int64_t idx = -1;
    uint8_t value = 0;   // used for detached proxies

    // Constructor for live proxy
    byte(byte_array& p, int64_t i) : parent(&p), idx(i) {}

    // Copy constructor: detached copy stores the current value
    byte(const byte& other)
        : parent(nullptr), idx(-1), value(static_cast<uint8_t>(other)) {}

    // Copy assignment: detached copy stores value
    byte& operator=(const byte& other) {
        parent = nullptr;
        idx = -1;
        value = static_cast<uint8_t>(other);
        return *this;
    }

    // Deleted move operations
    byte(byte&& other) = delete;
    byte& operator=(byte&& other) = delete;

    // Conversion to uint8_t
    operator uint8_t() const {
        if (parent) {
            if (parent->is_signed) {
                int8_t val = reinterpret_cast<int8_t*>(parent->data)[idx];
                return static_cast<uint8_t>(static_cast<int16_t>(val) + 128);
            } else {
                return reinterpret_cast<uint8_t*>(parent->data)[idx];
            }
        } else {
            return value;  // return local value if detached
        }
    }

    // Assignment from uint8_t
    byte& operator=(uint8_t normalized_value) {
        if (parent) {
            if (parent->is_signed) {
                reinterpret_cast<int8_t*>(parent->data)[idx] =
                    static_cast<int8_t>(static_cast<int16_t>(normalized_value) - 128);
            } else {
                reinterpret_cast<uint8_t*>(parent->data)[idx] = normalized_value;
            }
        } else {
            value = normalized_value;  // store in local value if detached
        }
        return *this;
    }
  };

  // Non-const index access: returns live proxy
  byte operator[](int64_t idx) { return byte(*this, idx); }

  // Const index access: returns immediate value
  uint8_t operator[](int64_t idx) const {
    if (is_signed) {
      int8_t val = reinterpret_cast<int8_t*>(data)[idx];
      return static_cast<uint8_t>(static_cast<int16_t>(val) + 128);
    } else {
      return reinterpret_cast<uint8_t*>(data)[idx];
    }
  }

  // Dereference (like *ptr)
  uint8_t operator*() const { return (*this)[0]; }
  byte operator*() { return byte(*this, 0); }

  // Pointer arithmetic
  byte_array operator+(int64_t offset) const {
    if (is_signed)
      return byte_array(static_cast<int8_t*>(data) + offset, true);
    else
      return byte_array(static_cast<uint8_t*>(data) + offset, false);
  }

  bool operator==(const byte_array& other) const { return data == other.data; }
  bool operator!=(const byte_array& other) const { return !(*this == other); }
};

} // namespace cuvs::detail

This would change the kernel definition to:

template <typename T, typename DataT = std::conditional_t<std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>, cuvs::detail::byte_array, T*>
__global__ void kernel(DataT data, ...) {
  ...
} 

And kernel instantiation would now look like:

// T is still uint8_t or int8_t at this point
template <typename T>
void launch_kernel(raft::device_matrix_view<T, ...> data, ...) {
   if constexpr (std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>) {
       cuvs::detail::byte_array data_b{data.data_handle(), std::is_same_v<T, int8_t>);
       kernel<uint8_t><<<...>>>(data_b, ...)
   }
   else {
        // normal instantiation, DataT is initialized to T*
        kernel<<<...>>>(data.data_handle(), ...)
    }
}

This approach will allow us to instantiate kernels and do all arithmetic in uint8_t without having to modify user data. At the time of search, we still do arithmetic in uint8_t and can write back data in int8_t if that is what the user requests.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by locating the kernel definitions and launch_kernel entry points that currently instantiate uint8_t and int8_t. Trace their data access and instantiation paths, then verify that the wrapper removes duplicate kernel instantiations, preserves signed and unsigned normalization, and reduces binary size without changing the public API.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.