types.cuh

/*
 * Copyright (c) 2025, NVIDIA CORPORATION.
 *
 * 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.
 */
 
#pragma once
 
#include <cudf/types.hpp>
#include <cudf/utilities/bit.hpp>
#include <cudf/utilities/traits.hpp>
 
#include <cuda/std/type_traits>
 
namespace cudf {
namespace jit {
namespace detail {
 
class alignas(16) column_device_view_base {
 public:
  CUDF_HOST_DEVICE column_device_view_base(data_type type,
                                           size_type size,
                                           void const* data,
                                           bitmask_type const* null_mask,
                                           size_type offset)
    : _type{type}, _size{size}, _data{data}, _null_mask{null_mask}, _offset{offset}
  {
  }
 
  template <typename T = void,
            CUDF_ENABLE_IF(cuda::std::is_same_v<T, void> or is_rep_layout_compatible<T>())>
  [[nodiscard]] CUDF_HOST_DEVICE T const* head() const noexcept
  {
    return static_cast<T const*>(_data);
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] CUDF_HOST_DEVICE T const* data() const noexcept
  {
    return head<T>() + _offset;
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type size() const noexcept { return _size; }
 
  [[nodiscard]] CUDF_HOST_DEVICE data_type type() const noexcept { return _type; }
 
  [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const noexcept { return nullptr != _null_mask; }
 
  [[nodiscard]] CUDF_HOST_DEVICE bitmask_type const* null_mask() const noexcept
  {
    return _null_mask;
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type offset() const noexcept { return _offset; }
 
  [[nodiscard]] __device__ bool is_valid(size_type element_index) const noexcept
  {
    return not nullable() or is_valid_nocheck(element_index);
  }
 
  [[nodiscard]] __device__ bool is_valid_nocheck(size_type element_index) const noexcept
  {
    return bit_is_set(_null_mask, offset() + element_index);
  }
 
  [[nodiscard]] __device__ bool is_null(size_type element_index) const noexcept
  {
    return not is_valid(element_index);
  }
 
  [[nodiscard]] __device__ bool is_null_nocheck(size_type element_index) const noexcept
  {
    return not is_valid_nocheck(element_index);
  }
 
  [[nodiscard]] __device__ bitmask_type get_mask_word(size_type word_index) const noexcept
  {
    return null_mask()[word_index];
  }
 
 protected:
  data_type _type;
 
  size_type _size;
 
  void const* _data;
 
  bitmask_type const* _null_mask;
 
  size_type _offset;
};
 
}  // namespace detail
 
class alignas(16) column_device_view : public detail::column_device_view_base {
 public:
  CUDF_HOST_DEVICE column_device_view(data_type type,
                                      size_type size,
                                      void const* data,
                                      bitmask_type const* null_mask,
                                      size_type offset,
                                      column_device_view* children,
                                      size_type num_children)
    : column_device_view_base(type, size, data, null_mask, offset),
      d_children(children),
      _num_children(num_children)
  {
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    return data<T>()[element_index];
  }
 
  template <typename T, CUDF_ENABLE_IF(is_fixed_point<T>())>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    using namespace numeric;
    using rep        = typename T::rep;
    auto const scale = scale_type{_type.scale()};
    return T{scaled_integer<rep>{data<rep>()[element_index], scale}};
  }
 
  [[nodiscard]] __device__ column_device_view child(size_type child_index) const noexcept
  {
    return d_children[child_index];
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type num_child_columns() const noexcept
  {
    return _num_children;
  }
 
 protected:
  column_device_view* d_children = nullptr;
 
  size_type _num_children = 0;
};
 
class alignas(16) mutable_column_device_view : public detail::column_device_view_base {
 public:
  mutable_column_device_view(data_type type,
                             size_type size,
                             void* data,
                             bitmask_type const* null_mask,
                             size_type offset,
                             mutable_column_device_view* children,
                             size_type num_children)
    : column_device_view_base(type, size, data, null_mask, offset),
      d_children(children),
      _num_children(num_children)
  {
  }
 
  template <typename T = void,
            CUDF_ENABLE_IF(cuda::std::is_same_v<T, void> or is_rep_layout_compatible<T>())>
  CUDF_HOST_DEVICE T* head() const noexcept
  {
    return const_cast<T*>(detail::column_device_view_base::head<T>());
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  CUDF_HOST_DEVICE T* data() const noexcept
  {
    return const_cast<T*>(detail::column_device_view_base::data<T>());
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  [[nodiscard]] __device__ T& element(size_type element_index) const noexcept
  {
    return data<T>()[element_index];
  }
 
  template <typename T, CUDF_ENABLE_IF(is_fixed_point<T>())>
  [[nodiscard]] __device__ T element(size_type element_index) const noexcept
  {
    using namespace numeric;
    using rep        = typename T::rep;
    auto const scale = scale_type{_type.scale()};
    return T{scaled_integer<rep>{data<rep>()[element_index], scale}};
  }
 
  template <typename T, CUDF_ENABLE_IF(is_rep_layout_compatible<T>())>
  __device__ void assign(size_type element_index, T value) const noexcept
  {
    data<T>()[element_index] = value;
  }
 
  template <typename T, CUDF_ENABLE_IF(is_fixed_point<T>())>
  __device__ void assign(size_type element_index, T value) const noexcept
  {
    // consider asserting that the scale matches
    using namespace numeric;
    using rep                  = typename T::rep;
    data<rep>()[element_index] = value.value();
  }
 
  [[nodiscard]] __device__ mutable_column_device_view child(size_type child_index) const noexcept
  {
    return d_children[child_index];
  }
 
  [[nodiscard]] CUDF_HOST_DEVICE size_type num_child_columns() const noexcept
  {
    return _num_children;
  }
 
 private:
  mutable_column_device_view* d_children;  
  size_type _num_children;                 
};
 
}  // namespace jit
}  // namespace cudf