diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5a0b2f95e830..7a043d91e698 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -337,6 +337,7 @@ add_library( src/aggregation/result_cache.cpp src/ast/expression_parser.cpp src/ast/expressions.cpp + src/ast/jit_expressions.cpp src/ast/operators.cpp src/binaryop/binaryop.cpp src/binaryop/compiled/ATan2.cu diff --git a/cpp/include/cudf/ast/detail/operator_functor.cuh b/cpp/include/cudf/ast/detail/operator_functor.cuh index 043b2c47746d..5abe32887054 100644 --- a/cpp/include/cudf/ast/detail/operator_functor.cuh +++ b/cpp/include/cudf/ast/detail/operator_functor.cuh @@ -778,12 +778,5 @@ struct operator_functor { } }; -constexpr bool flatten_predicate(possibly_null_value_t value) { return value; } - -constexpr bool flatten_predicate(possibly_null_value_t value) -{ - return value.has_value() && *value; -} - } // namespace ast::detail } // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/ast/expressions.hpp b/cpp/include/cudf/ast/expressions.hpp index 76fdf2d10120..22ace86028e0 100644 --- a/cpp/include/cudf/ast/expressions.hpp +++ b/cpp/include/cudf/ast/expressions.hpp @@ -59,7 +59,7 @@ class expression_transformer; * This class is a part of a "visitor" pattern with the `expression_parser` class. * Expressions inheriting from this class can accept parsers as visitors. */ -struct expression { +struct [[nodiscard]] expression { /** * @brief Accepts a visitor class. * @@ -514,17 +514,17 @@ class operation : public expression { namespace detail { -/// @brief An expression that represents a filter predicate. +/// @brief An expression that represents a predicate. /// /// This is an internal expression used in filter operations. It is not intended to be used by /// external code and is not a part of the public API. -class filter_predicate : public expression { +class predicate : public expression { public: /** * @brief Construct a new filter predicate object * @param source The source expression from which the predicate value is taken */ - filter_predicate(expression const& source) : source_{source} {} + predicate(expression const& source) : source_{source} {} /** * @copydoc expression::accept diff --git a/cpp/include/cudf/ast/jit_expressions.hpp b/cpp/include/cudf/ast/jit_expressions.hpp new file mode 100644 index 000000000000..37ea00912537 --- /dev/null +++ b/cpp/include/cudf/ast/jit_expressions.hpp @@ -0,0 +1,442 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +namespace CUDF_EXPORT cudf { +namespace ast { + +/** + * @addtogroup expressions + * @{ + * @file + */ + +namespace jit { +namespace detail { + +struct operation : public ast::expression { + /** + * @brief Construct a new operation object. + * @param op The opcode for this operation + * @param args The arguments for this operation + */ + operation(cudf::detail::row_ir::opcode op, + std::vector> args) + : op_{op}, args_{std::move(args)} + { + } + + /** + * @brief Construct a new operation object with a target scale (for rescale and precision check + * operations). + * @param op The opcode for this operation + * @param args The arguments for this operation + * @param target_scale The target scale for this operation (only applicable for rescale and + * precision check operations) + */ + operation(cudf::detail::row_ir::opcode op, + std::vector> args, + int32_t target_scale) + : op_{op}, args_{std::move(args)}, target_scale_{target_scale} + { + } + + operation(operation const&) = default; //< Copy constructor + operation(operation&&) = default; //< Move constructor + operation& operator=(operation const&) = default; //< Copy assignment + operation& operator=(operation&&) = default; //< Move assignment + ~operation() override = default; //< Destructor + + /** + * @brief Get the opcode. + * + * @return The opcode + */ + [[nodiscard]] cudf::detail::row_ir::opcode get_opcode() const { return op_; } + + /** + * @brief Get the operands. + * + * @return Vector of operands + */ + [[nodiscard]] std::span const> get_arguments() const + { + return args_; + } + + /** + * @brief Get the target scale for rescale and precision check operations. + * + * @return The target scale if applicable, std::nullopt otherwise + */ + [[nodiscard]] std::optional get_target_scale() const { return target_scale_; } + + /** + * @copydoc expression::accept + */ + cudf::size_type accept(cudf::ast::detail::expression_parser& visitor) const override; + + /** + * @copydoc expression::accept + */ + std::reference_wrapper accept( + cudf::ast::detail::expression_transformer& visitor) const override; + + [[nodiscard]] bool may_evaluate_null(table_view const& left, + table_view const& right, + rmm::cuda_stream_view stream) const override; + + /** + * @copydoc expression::accept + */ + [[nodiscard]] std::unique_ptr accept( + cudf::detail::row_ir::ast_converter& visitor) const override; + + private: + cudf::detail::row_ir::opcode op_; + std::vector> args_; + std::optional target_scale_ = std::nullopt; +}; + +} // namespace detail + +/** + * @brief Creates an expression that evaluates to `NULL` if the condition is true, and the value of + * `a` otherwise. + * @param tree The expression tree to which this expression will be added + * @param a The expression to nullify if the condition is true + * @param condition The condition under which to nullify the value + * @return An expression representing the nullified value + */ +expression const& nullify_if(ast::tree& tree, expression const& a, expression const& condition); + +/** + * @brief Creates an expression that evaluates to the first non-null value among its arguments. + * @param tree The expression tree to which this expression will be added + * @param a The first expression to coalesce + * @param b The second expression to coalesce + * @return An expression representing the coalesced value + */ +expression const& coalesce(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that evaluates to `true` if the condition is true and not null, and + * `false` otherwise. This is used to implement predicates in the JIT. + * @param tree The expression tree to which this expression will be added + * @param condition The condition to evaluate as a predicate + * @return An expression representing the result of the predicate + */ +expression const& predicate(ast::tree& tree, expression const& condition); + +/** + * @brief Creates an expression that performs ANSI-compliant addition of `a` and `b`, which throws + * an error on overflow. + * @param tree The expression tree to which this expression will be added + * @param a The first addend + * @param b The second addend + * @return An expression representing the result of the addition + */ +expression const& ansi_add(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant subtraction of `a` and `b`, which + * throws an error on overflow. + * @param tree The expression tree to which this expression will be added + * @param a The minuend + * @param b The subtrahend + * @return An expression representing the result of the subtraction + */ +expression const& ansi_sub(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant multiplication of `a` and `b`, which + * throws an error on overflow. + * @param tree The expression tree to which this expression will be added + * @param a The first factor + * @param b The second factor + * @return An expression representing the result of the multiplication + */ +expression const& ansi_mul(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant division of `a` by `b`, which throws an + * error on division by zero. + * @param tree The expression tree to which this expression will be added + * @param a The dividend + * @param b The divisor + * @return An expression representing the result of the division + */ +expression const& ansi_div(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant modulus of `a` by `b`, which throws an + * error on division by zero. + * @param tree The expression tree to which this expression will be added + * @param a The value to be divided + * @param b The divisor + * @return An expression representing the result of the modulus operation + */ +expression const& ansi_mod(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant absolute value of `a`, which throws an + * error on overflow. + * @param tree The expression tree to which this expression will be added + * @param a The value for which to compute the absolute value + * @return An expression representing the absolute value + */ +expression const& ansi_abs(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that performs ANSI-compliant negation of `a`, which throws an error + * on overflow. + * @param tree The expression tree to which this expression will be added + * @param a The value to negate + * @return An expression representing the negated value + */ +expression const& ansi_neg(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that performs an ANSI-compliant precision check on `a` with the + * given precision, which throws an error if the value of `a` exceeds the specified precision. + * @param tree The expression tree to which this expression will be added + * @param a The value for which to perform the precision check + * @param precision The precision to check against + * @return An expression representing the result of the precision check + */ +expression const& ansi_precision_check(ast::tree& tree, + expression const& a, + expression const& precision); + +/** + * @brief Creates an expression that performs ANSI-compliant addition of `a` and `b`, which returns + * `NULL` on overflow instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The first addend + * @param b The second addend + * @return An expression representing the result of the addition, or `NULL` if overflow occurs + */ +expression const& ansi_try_add(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant subtraction of `a` and `b`, which + * returns `NULL` on overflow instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The minuend + * @param b The subtrahend + * @return An expression representing the result of the subtraction, or `NULL` if overflow occurs + */ +expression const& ansi_try_sub(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant multiplication of `a` and `b`, which + * returns `NULL` on overflow instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The first factor + * @param b The second factor + * @return An expression representing the result of the multiplication, or `NULL` if overflow occurs + */ +expression const& ansi_try_mul(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant division of `a` by `b`, which returns + * `NULL` on division by zero instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The dividend + * @param b The divisor + * @return An expression representing the result of the division, or `NULL` if division by zero + * occurs + */ +expression const& ansi_try_div(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant modulus of `a` by `b`, which returns + * `NULL` on division by zero instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The value to be divided + * @param b The divisor + * @return An expression representing the result of the modulus operation, or `NULL` if division by + * zero occurs + */ +expression const& ansi_try_mod(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs ANSI-compliant absolute value of `a`, which returns + * `NULL` on overflow instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The value for which to compute the absolute value + * @return An expression representing the absolute value, or `NULL` if overflow occurs + */ +expression const& ansi_try_abs(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that performs ANSI-compliant negation of `a`, which returns `NULL` + * on overflow instead of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The value to negate + * @return An expression representing the negated value, or `NULL` if overflow occurs + */ +expression const& ansi_try_neg(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that performs an ANSI-compliant precision check on `a` with the + * given precision, which returns `NULL` if the value of `a` exceeds the specified precision instead + * of throwing an error. + * @param tree The expression tree to which this expression will be added + * @param a The value for which to perform the precision check + * @param precision The precision to check against + * @return An expression representing the result of the precision check, or `NULL` if the value + * exceeds the specified precision + */ +expression const& ansi_try_precision_check(ast::tree& tree, + expression const& a, + expression const& precision); + +/** + * @brief Creates an expression that performs a bitwise left shift of `a` by `b`. + * @param tree The expression tree to which this expression will be added + * @param a The value to shift + * @param b The number of bits by which to shift + * @return An expression representing the result of the bitwise left shift + */ +expression const& bit_shift_left(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that performs a bitwise right shift of `a` by `b`. + * @param tree The expression tree to which this expression will be added + * @param a The value to shift + * @param b The number of bits by which to shift + * @return An expression representing the result of the bitwise right shift + */ +expression const& bit_shift_right(ast::tree& tree, expression const& a, expression const& b); + +/** + * @brief Creates an expression that casts `a` to a boolean type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_b8(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 8-bit signed integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_i8(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 16-bit signed integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_i16(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 32-bit signed integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_i32(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 64-bit signed integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_i64(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to an 8-bit unsigned integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_u8(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 16-bit unsigned integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_u16(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 32-bit unsigned integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_u32(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 64-bit unsigned integer type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_u64(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 32-bit floating point type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_f32(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 64-bit floating point type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_f64(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 32-bit decimal type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_dec32(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 64-bit decimal type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_dec64(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that casts `a` to a 128-bit decimal type. + * @param tree The expression tree to which this expression will be added + * @param a The value to cast + * @return An expression representing the result of the cast + */ +expression const& cast_to_dec128(ast::tree& tree, expression const& a); + +/** + * @brief Creates an expression that rescales a decimal expression `a` to a new scale `new_scale`. + * @param tree The expression tree to which this expression will be added + * @param a The decimal expression to rescale + * @param new_scale The new scale to which to rescale the decimal expression + * @return An expression representing the rescaled decimal value + */ +expression const& rescale(ast::tree& tree, expression const& a, int32_t new_scale); + +} // namespace jit + +} // namespace ast +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/ansi_arithmetic.cuh b/cpp/include/cudf/operators/ansi_arithmetic.cuh new file mode 100644 index 000000000000..ab7c1889d5f4 --- /dev/null +++ b/cpp/include/cudf/operators/ansi_arithmetic.cuh @@ -0,0 +1,627 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +namespace CUDF_EXPORT cudf { +namespace ops { +namespace detail { + +template +struct promoted_t; + +template <> +struct promoted_t { + using type = int16_t; +}; + +template <> +struct promoted_t { + using type = uint16_t; +}; + +template <> +struct promoted_t { + using type = int32_t; +}; + +template <> +struct promoted_t { + using type = uint32_t; +}; + +template <> +struct promoted_t { + using type = int64_t; +}; + +template <> +struct promoted_t { + using type = uint64_t; +}; + +template <> +struct promoted_t { + using type = __int128; +}; + +template <> +struct promoted_t { + using type = unsigned __int128; +}; + +template +using promote = typename promoted_t::type; + +} // namespace detail + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_add(T* out, T const* a, T const* b) +{ + using P = detail::promote; + auto r = static_cast

(*a) + static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max())) { return errc::OVERFLOW; } + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_add(T* out, T const* a, T const* b) +{ + using P = detail::promote; + auto r = static_cast

(*a) + static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max()) || + r < static_cast

(cuda::std::numeric_limits::min())) { + return errc::OVERFLOW; + } + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v) +__device__ inline errc ansi_add(T* out, T const* a, T const* b) +{ + *out = *a + *b; + return errc::OK; +} + +template +__device__ inline errc ansi_add(decimal* out, decimal const* a, decimal const* b) +{ + auto scale = cuda::std::min(a->scale(), b->scale()); + + if (numeric::addition_overflow(a->rescaled(scale).value(), b->rescaled(scale).value())) { + return errc::OVERFLOW; + } + + *out = decimal{numeric::scaled_integer{ + a->rescaled(scale).value() + b->rescaled(scale).value(), numeric::scale_type{scale}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_add(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_add(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_sub(T* out, T const* a, T const* b) +{ + if (*a < *b) { return errc::OVERFLOW; } + auto r = *a - *b; + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_sub(T* out, T const* a, T const* b) +{ + using P = detail::promote; + auto r = static_cast

(*a) - static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max()) || + r < static_cast

(cuda::std::numeric_limits::min())) { + return errc::OVERFLOW; + } + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v) +__device__ inline errc ansi_sub(T* out, T const* a, T const* b) +{ + *out = *a - *b; + return errc::OK; +} + +template +__device__ inline errc ansi_sub(decimal* out, decimal const* a, decimal const* b) +{ + auto scale = cuda::std::min(a->scale(), b->scale()); + + if (numeric::subtraction_overflow(a->rescaled(scale).value(), b->rescaled(scale).value())) { + return errc::OVERFLOW; + } + + *out = decimal{numeric::scaled_integer{ + a->rescaled(scale).value() - b->rescaled(scale).value(), numeric::scale_type{scale}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_sub(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_sub(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_mul(T* out, T const* a, T const* b) +{ + using P = detail::promote; + auto r = static_cast

(*a) * static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max())) { return errc::OVERFLOW; } + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_mul(T* out, T const* a, T const* b) +{ + using P = detail::promote; + auto r = static_cast

(*a) * static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max()) || + r < static_cast

(cuda::std::numeric_limits::min())) { + return errc::OVERFLOW; + } + *out = static_cast(r); + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v) +__device__ inline errc ansi_mul(T* out, T const* a, T const* b) +{ + *out = *a * *b; + return errc::OK; +} + +template +__device__ inline errc ansi_mul(decimal* out, decimal const* a, decimal const* b) +{ + if (numeric::multiplication_overflow(a->value(), b->value())) { return errc::OVERFLOW; } + + *out = decimal{numeric::scaled_integer{a->value() * b->value(), + numeric::scale_type{a->scale() + b->scale()}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_mul(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_mul(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_div(T* out, T const* a, T const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + *out = static_cast(static_cast(*a) / static_cast(*b)); + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_div(T* out, T const* a, T const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + if (*a == cuda::std::numeric_limits::min() && *b == -1) { return errc::OVERFLOW; } + *out = static_cast(static_cast(*a) / static_cast(*b)); + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v) +__device__ inline errc ansi_div(T* out, T const* a, T const* b) +{ + using P = detail::promote; + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + auto r = static_cast

(*a) / static_cast

(*b); + if (r > static_cast

(cuda::std::numeric_limits::max()) || + r < static_cast

(cuda::std::numeric_limits::lowest())) { + return errc::OVERFLOW; + } + *out = static_cast(r); + return errc::OK; +} + +template +__device__ inline errc ansi_div(decimal* out, decimal const* a, decimal const* b) +{ + if (numeric::division_overflow(a->value(), b->value()) || b->value() == 0) { + return errc::OVERFLOW; + } + + *out = decimal{numeric::scaled_integer{a->value() / b->value(), + numeric::scale_type{a->scale() - b->scale()}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_div(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_div(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_mod(T* out, T const* a, T const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + T r = *a % *b; + if (r != 0 && ((r > 0) != (*b > 0))) { r += *b; } + *out = r; + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_mod(T* out, T const* a, T const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + *out = *a % *b; + return errc::OK; +} + +__device__ inline errc ansi_mod(float* out, float const* a, float const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + *out = (*a) - (*b) * ::floorf((*a) / (*b)); + return errc::OK; +} + +__device__ inline errc ansi_mod(double* out, double const* a, double const* b) +{ + if (*b == 0) { return errc::DIVISION_BY_ZERO; } + *out = (*a) - (*b) * ::floor((*a) / (*b)); + return errc::OK; +} + +template +__device__ inline errc ansi_mod(decimal* out, decimal const* a, decimal const* b) +{ + if (b->value() == 0) { return errc::DIVISION_BY_ZERO; } + + decimal div; + + if (errc e = ansi_div(&div, a, b); e != errc::OK) { return e; } + + decimal quotient; + floor("ient, &div); + *out = *a - *b * quotient; + return errc::OK; +} + +template +__device__ inline errc ansi_mod(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_mod(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_signed_v) +__device__ inline errc ansi_abs(T* out, T const* a) +{ + if (*a == cuda::std::numeric_limits::min()) { return errc::OVERFLOW; } + *out = (*a < 0) ? -(*a) : *a; + return errc::OK; +} + +template + requires(cuda::std::is_integral_v && cuda::std::is_unsigned_v) +__device__ inline errc ansi_abs(T* out, T const* a) +{ + *out = *a; + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v) +__device__ inline errc ansi_abs(T* out, T const* a) +{ + *out = (*a < 0) ? -(*a) : *a; + return errc::OK; +} + +template +__device__ inline errc ansi_abs(decimal* out, decimal const* a) +{ + if (a->value() == cuda::std::numeric_limits::min()) { return errc::OVERFLOW; } + auto rep = a->value() < 0 ? -a->value() : a->value(); + *out = decimal{numeric::scaled_integer{rep, numeric::scale_type{a->scale()}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_abs(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + if (errc e = ansi_abs(&r, &a->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_signed_v) +__device__ inline errc ansi_neg(T* out, T const* a) +{ + if (*a == cuda::std::numeric_limits::min()) { return errc::OVERFLOW; } + *out = -(*a); + return errc::OK; +} + +template +__device__ inline errc ansi_neg(decimal* out, decimal const* a) +{ + if (a->value() == cuda::std::numeric_limits::min()) { return errc::OVERFLOW; } + auto rep = -a->value(); + *out = decimal{numeric::scaled_integer{rep, numeric::scale_type{a->scale()}}}; + return errc::OK; +} + +template +__device__ inline errc ansi_neg(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + if (errc e = ansi_neg(&r, &a->value()); e != errc::OK) { + *out = nullopt; + return e; + } + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_precision_check(decimal* out, + decimal const* a, + int32_t const* precision) +{ + if (*precision <= 0) { return errc::OVERFLOW; } + + auto value = a->value(); + if (value == cuda::std::numeric_limits::min()) { return errc::OVERFLOW; } + + auto abs_value = value < 0 ? -value : value; + + if (abs_value >= detail::ipow10(static_cast(*precision))) { return errc::OVERFLOW; } + + *out = *a; + return errc::OK; +} + +template +__device__ inline errc ansi_precision_check(optional* out, + optional const* a, + optional const* precision) +{ + if (a->has_value()) { + T r; + if (errc e = ansi_precision_check(&r, &a->value(), &precision->value()); e != errc::OK) { + *out = nullopt; + return e; + } else { + *out = r; + return errc::OK; + } + } else { + *out = nullopt; + return errc::OK; + } +} + +template +__device__ inline errc ansi_try_add(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_add(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_sub(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_sub(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_mul(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_mul(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_div(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_div(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_mod(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + if (errc e = ansi_mod(&r, &a->value(), &b->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_abs(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + if (errc e = ansi_abs(&r, &a->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_neg(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + if (errc e = ansi_neg(&r, &a->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc ansi_try_precision_check(optional>* out, + optional> const* a, + optional const* precision) +{ + if (a->has_value() && precision->has_value()) { + decimal r; + if (errc e = ansi_precision_check(&r, &a->value(), &precision->value()); e != errc::OK) { + *out = nullopt; + } else { + *out = r; + } + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/arithmetic.cuh b/cpp/include/cudf/operators/arithmetic.cuh new file mode 100644 index 000000000000..4163328385b4 --- /dev/null +++ b/cpp/include/cudf/operators/arithmetic.cuh @@ -0,0 +1,278 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template + requires(cuda::std::is_signed_v || cuda::std::is_floating_point_v) +__device__ inline errc abs(T* out, T const* a) +{ + *out = (*a < 0) ? -*a : *a; + return errc::OK; +} + +template + requires(cuda::std::is_unsigned_v) +__device__ inline errc abs(T* out, T const* a) +{ + *out = *a; + return errc::OK; +} + +template +__device__ inline errc abs(decimal* out, decimal const* a) +{ + auto rep = a->value() < 0 ? -a->value() : a->value(); + *out = decimal{numeric::scaled_integer{rep, a->scale()}}; + return errc::OK; +} + +template +__device__ inline errc abs(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + abs(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc add(T* out, T const* a, T const* b) +{ + *out = (*a + *b); + return errc::OK; +} + +template +__device__ inline errc add(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + add(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc div(T* out, T const* a, T const* b) +{ + *out = (*a / *b); + return errc::OK; +} + +template +__device__ inline errc div(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + div(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_integral_v) +__device__ inline errc floor_div(T* out, T const* a, T const* b) +{ + *out = cudf::detail::integral_floor_div(*a, *b); + return errc::OK; +} + +__device__ inline errc floor_div(float* out, float const* a, float const* b) +{ + *out = ::floorf(*a / *b); + return errc::OK; +} + +__device__ inline errc floor_div(double* out, double const* a, double const* b) +{ + *out = ::floor(*a / *b); + return errc::OK; +} + +template +__device__ inline errc floor_div(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + floor_div(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc mod(T* out, T const* a, T const* b) +{ + *out = (*a % *b); + return errc::OK; +} + +__device__ inline errc mod(float* out, float const* a, float const* b) +{ + *out = ::fmodf(*a, *b); + return errc::OK; +} + +__device__ inline errc mod(double* out, double const* a, double const* b) +{ + *out = ::fmod(*a, *b); + return errc::OK; +} + +template +__device__ inline errc mod(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + mod(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc pymod(T* out, T const* a, T const* b) +{ + *out = (*a % *b + *b) % *b; + return errc::OK; +} + +__device__ inline errc pymod(float* out, float const* a, float const* b) +{ + *out = ::fmodf(::fmodf(*a, *b) + *b, *b); + return errc::OK; +} + +__device__ inline errc pymod(double* out, double const* a, double const* b) +{ + *out = ::fmod(::fmod(*a, *b) + *b, *b); + return errc::OK; +} + +template +__device__ inline errc pymod(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + pymod(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc mul(T* out, T const* a, T const* b) +{ + *out = (*a * *b); + return errc::OK; +} + +template +__device__ inline errc mul(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + mul(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_signed_v) +__device__ inline errc neg(T* out, T const* a) +{ + *out = -(*a); + return errc::OK; +} + +template +__device__ inline errc neg(decimal* out, decimal const* a) +{ + auto rep = -a->value(); + *out = decimal{numeric::scaled_integer{rep, a->scale()}}; + return errc::OK; +} + +template +__device__ inline errc neg(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + neg(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc sub(T* out, T const* a, T const* b) +{ + *out = *a - *b; + return errc::OK; +} + +template +__device__ inline errc sub(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + sub(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(cuda::std::is_floating_point_v || cuda::std::is_integral_v) +__device__ inline errc true_div(double* out, T const* a, T const* b) +{ + *out = static_cast(*a) / static_cast(*b); + return errc::OK; +} + +template +__device__ inline errc true_div(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + double r; + true_div(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/bitwise.cuh b/cpp/include/cudf/operators/bitwise.cuh new file mode 100644 index 000000000000..18943bbd56ab --- /dev/null +++ b/cpp/include/cudf/operators/bitwise.cuh @@ -0,0 +1,133 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +__device__ inline errc bit_and(T* out, T const* a, T const* b) +{ + *out = (*a & *b); + return errc::OK; +} + +template +__device__ inline errc bit_and(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + bit_and(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc bit_invert(T* out, T const* a) +{ + *out = ~(*a); + return errc::OK; +} + +template +__device__ inline errc bit_invert(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + bit_invert(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc bit_or(T* out, T const* a, T const* b) +{ + *out = (*a | *b); + return errc::OK; +} + +template +__device__ inline errc bit_or(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + bit_or(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc bit_xor(T* out, T const* a, T const* b) +{ + *out = (*a ^ *b); + return errc::OK; +} + +template +__device__ inline errc bit_xor(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + bit_xor(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc bit_shift_left(T* out, T const* a, T const* b) +{ + *out = (*a << *b); + return errc::OK; +} + +template +__device__ inline errc bit_shift_left(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + bit_shift_left(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc bit_shift_right(T* out, T const* a, T const* b) +{ + *out = (*a >> *b); + return errc::OK; +} + +template +__device__ inline errc bit_shift_right(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + bit_shift_right(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/casts.cuh b/cpp/include/cudf/operators/casts.cuh new file mode 100644 index 000000000000..4b0fd81619a3 --- /dev/null +++ b/cpp/include/cudf/operators/casts.cuh @@ -0,0 +1,345 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +__device__ inline errc cast_to_b8(bool* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_b8(optional* out, optional const* a) +{ + if (a->has_value()) { + bool r; + cast_to_b8(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_i8(int8_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_i8(optional* out, optional const* a) +{ + if (a->has_value()) { + int8_t r; + cast_to_i8(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_i16(int16_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_i16(optional* out, optional const* a) +{ + if (a->has_value()) { + int16_t r; + cast_to_i16(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_i32(int32_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_i32(optional* out, optional const* a) +{ + if (a->has_value()) { + int32_t r; + cast_to_i32(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_i64(int64_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_i64(optional* out, optional const* a) +{ + if (a->has_value()) { + int64_t r; + cast_to_i64(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_u8(uint8_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_u8(optional* out, optional const* a) +{ + if (a->has_value()) { + uint8_t r; + cast_to_u8(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_u16(uint16_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_u16(optional* out, optional const* a) +{ + if (a->has_value()) { + uint16_t r; + cast_to_u16(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_u32(uint32_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_u32(optional* out, optional const* a) +{ + if (a->has_value()) { + uint32_t r; + cast_to_u32(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_u64(uint64_t* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_u64(optional* out, optional const* a) +{ + if (a->has_value()) { + uint64_t r; + cast_to_u64(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} +template + requires(std::is_integral_v || std::is_floating_point_v) +__device__ inline errc cast_to_f32(float* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_f32(float* out, decimal const* a) +{ + *out = convert_fixed_to_floating(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_f32(optional* out, optional const* a) +{ + if (a->has_value()) { + float r; + cast_to_f32(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template + requires(std::is_integral_v || std::is_floating_point_v) +__device__ inline errc cast_to_f64(double* out, T const* a) +{ + *out = static_cast(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_f64(double* out, decimal const* a) +{ + *out = convert_fixed_to_floating(*a); + return errc::OK; +} + +template +__device__ inline errc cast_to_f64(optional* out, optional const* a) +{ + if (a->has_value()) { + double r; + cast_to_f64(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +namespace detail { + +template +__device__ inline errc decimal_cast(decimal* out, decimal const* a) +{ + auto rep = static_cast(a->value()); + *out = decimal{numeric::scaled_integer{rep, a->scale()}}; + return errc::OK; +} + +} // namespace detail + +// TODO: CAST_TO_DEC32 for int & float + +template +__device__ inline errc cast_to_dec32(numeric::decimal32* out, decimal const* a) +{ + return detail::decimal_cast(out, a); +} + +template +__device__ inline errc cast_to_dec32(optional* out, + optional> const* a) +{ + if (a->has_value()) { + numeric::decimal32 r; + cast_to_dec32(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_dec64(numeric::decimal64* out, decimal const* a) +{ + return detail::decimal_cast(out, a); +} + +template +__device__ inline errc cast_to_dec64(optional* out, + optional> const* a) +{ + if (a->has_value()) { + numeric::decimal64 r; + cast_to_dec64(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc cast_to_dec128(numeric::decimal128* out, decimal const* a) +{ + return detail::decimal_cast(out, a); +} + +template +__device__ inline errc cast_to_dec128(optional* out, + optional> const* a) +{ + if (a->has_value()) { + numeric::decimal128 r; + cast_to_dec128(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc rescale(decimal* out, decimal const* a, int32_t const* new_scale) +{ + *out = a->rescaled(numeric::scale_type{*new_scale}); + return errc::OK; +} + +template +__device__ inline errc rescale(optional>* out, + optional> const* a, + optional const* new_scale) +{ + if (a->has_value() && new_scale->has_value()) { + decimal r; + rescale(&r, &a->value(), new_scale->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/comparison.cuh b/cpp/include/cudf/operators/comparison.cuh new file mode 100644 index 000000000000..39ed1145eaf8 --- /dev/null +++ b/cpp/include/cudf/operators/comparison.cuh @@ -0,0 +1,159 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +__device__ inline errc equal(bool* out, T const* a, T const* b) +{ + *out = (*a == *b); + return errc::OK; +} + +template +__device__ inline errc equal(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + equal(&r, &a->value(), &b->value()); + *out = r; + } else if (!a->has_value() && !b->has_value()) { + *out = true; + } else { + *out = false; + } + return errc::OK; +} + +template +__device__ inline errc not_equal(bool* out, T const* a, T const* b) +{ + *out = (*a != *b); + return errc::OK; +} + +template +__device__ inline errc not_equal(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + not_equal(&r, &a->value(), &b->value()); + *out = r; + } else if (!a->has_value() && !b->has_value()) { + *out = false; + } else { + *out = true; + } + return errc::OK; +} + +template +__device__ inline errc greater(bool* out, T const* a, T const* b) +{ + *out = (*a > *b); + return errc::OK; +} + +template +__device__ inline errc greater(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + greater(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = false; + } + return errc::OK; +} + +template +__device__ inline errc greater_equal(bool* out, T const* a, T const* b) +{ + *out = (*a >= *b); + return errc::OK; +} + +template +__device__ inline errc greater_equal(optional* out, + optional const* a, + optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + greater_equal(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = false; + } + return errc::OK; +} + +template +__device__ inline errc less(bool* out, T const* a, T const* b) +{ + *out = (*a < *b); + return errc::OK; +} + +template +__device__ inline errc less(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + less(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = false; + } + return errc::OK; +} + +template +__device__ inline errc less_equal(bool* out, T const* a, T const* b) +{ + *out = (*a <= *b); + return errc::OK; +} + +template +__device__ inline errc less_equal(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + less_equal(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = false; + } + return errc::OK; +} + +template +__device__ inline errc null_equal(bool* out, T const* a, T const* b) +{ + *out = (*a == *b); + return errc::OK; +} + +template +__device__ inline errc null_equal(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + *out = (*(*a) == *(*b)); + } else if (!a->has_value() && !b->has_value()) { + *out = true; + } else { + *out = false; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/error.hpp b/cpp/include/cudf/operators/error.hpp new file mode 100644 index 000000000000..52e6465a1e47 --- /dev/null +++ b/cpp/include/cudf/operators/error.hpp @@ -0,0 +1,28 @@ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +enum errc : int { OK = 0, OVERFLOW = 1, DIVISION_BY_ZERO = 2 }; + +inline char const* to_string(errc error_code) +{ + switch (error_code) { + case errc::OK: return "cudf::ops::errc::OK"; + case errc::OVERFLOW: return "cudf::ops::errc::OVERFLOW"; + case errc::DIVISION_BY_ZERO: return "cudf::ops::errc::DIVISION_BY_ZERO"; + default: return "UNKNOWN_ERROR"; + } +} + +enum class error_mode : char { IGNORE = 0, ANY_ROW = 1 }; + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/logic.cuh b/cpp/include/cudf/operators/logic.cuh new file mode 100644 index 000000000000..fbbc8f8148e5 --- /dev/null +++ b/cpp/include/cudf/operators/logic.cuh @@ -0,0 +1,150 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +__device__ inline errc null_logical_and(bool* out, T const* a, T const* b) +{ + *out = (*a && *b); + return errc::OK; +} + +template +__device__ inline errc null_logical_and(optional* out, + optional const* a, + optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + null_logical_and(&r, &a->value(), &b->value()); + *out = r; + } else if (!a->has_value() && !b->has_value()) { + *out = nullopt; + } else { + if (a->has_value() ? *(*a) : *(*b)) { + *out = nullopt; + } else { + *out = false; + } + } + return errc::OK; +} + +template +__device__ inline errc null_logical_or(bool* out, T const* a, T const* b) +{ + *out = (*a || *b); + return errc::OK; +} + +template +__device__ inline errc null_logical_or(optional* out, + optional const* a, + optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + null_logical_or(&r, &a->value(), &b->value()); + *out = r; + } else if (!a->has_value() && !b->has_value()) { + *out = nullopt; + } else { + if (a->has_value() ? *(*a) : *(*b)) { + *out = true; + } else { + *out = nullopt; + } + } + return errc::OK; +} + +template +__device__ inline errc logical_and(bool* out, T const* a, T const* b) +{ + *out = (*a && *b); + return errc::OK; +} + +template +__device__ inline errc logical_and(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + logical_and(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc logical_or(bool* out, T const* a, T const* b) +{ + *out = (*a || *b); + return errc::OK; +} + +template +__device__ inline errc logical_or(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + bool r; + logical_or(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc logical_not(bool* out, T const* a) +{ + *out = !(*a); + return errc::OK; +} + +template +__device__ inline errc logical_not(optional* out, optional const* a) +{ + if (a->has_value()) { + bool r; + logical_not(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc if_else(T* out, T const* true_value, T const* false_value, bool const* pred) +{ + *out = *pred ? *true_value : *false_value; + return errc::OK; +} + +template +__device__ inline errc if_else(optional* out, + optional const* true_value, + optional const* false_value, + optional const* pred) +{ + if (pred->has_value() && true_value->has_value() && false_value->has_value()) { + if_else(&out->value(), &pred->value(), &true_value->value(), &false_value->value()); + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/math.cuh b/cpp/include/cudf/operators/math.cuh new file mode 100644 index 000000000000..30f759b3d012 --- /dev/null +++ b/cpp/include/cudf/operators/math.cuh @@ -0,0 +1,243 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +__device__ inline errc cbrt(float* out, float const* a) +{ + *out = ::cbrtf(*a); + return errc::OK; +} + +__device__ inline errc cbrt(double* out, double const* a) +{ + *out = ::cbrt(*a); + return errc::OK; +} + +template +__device__ inline errc cbrt(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + cbrt(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc ceil(float* out, float const* a) +{ + *out = ::ceilf(*a); + return errc::OK; +} + +__device__ inline errc ceil(double* out, double const* a) +{ + *out = ::ceil(*a); + return errc::OK; +} + +template +__device__ inline errc ceil(decimal* out, decimal const* a) +{ + auto factor = detail::ipow10(static_cast(a->scale())); + auto div = a->value() / factor; + auto rem = a->value() % factor; + if (rem == 0) { + *out = *a; + } else { + auto val = a->value() > 0 ? (div + 1) : div; + *out = decimal{numeric::scaled_integer{val, a->scale()}}; + } + return errc::OK; +} + +template +__device__ inline errc ceil(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + ceil(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc exp(float* out, float const* a) +{ + *out = ::expf(*a); + return errc::OK; +} + +__device__ inline errc exp(double* out, double const* a) +{ + *out = ::exp(*a); + return errc::OK; +} + +template +__device__ inline errc exp(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + exp(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc floor(float* out, float const* a) +{ + *out = ::floorf(*a); + return errc::OK; +} + +__device__ inline errc floor(double* out, double const* a) +{ + *out = ::floor(*a); + return errc::OK; +} + +template +__device__ inline errc floor(decimal* out, decimal const* a) +{ + auto factor = detail::ipow10(static_cast(a->scale())); + auto div = a->value() / factor; + auto rem = a->value() % factor; + if (rem == 0) { + *out = *a; + } else { + auto val = a->value() > 0 ? div : (div - 1); + *out = decimal{numeric::scaled_integer{val, a->scale()}}; + } + return errc::OK; +} + +template +__device__ inline errc floor(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + floor(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc log(float* out, float const* a) +{ + *out = ::logf(*a); + return errc::OK; +} + +__device__ inline errc log(double* out, double const* a) +{ + *out = ::log(*a); + return errc::OK; +} + +template +__device__ inline errc log(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + log(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc pow(float* out, float const* a, float const* b) +{ + *out = ::powf(*a, *b); + return errc::OK; +} + +__device__ inline errc pow(double* out, double const* a, double const* b) +{ + *out = ::pow(*a, *b); + return errc::OK; +} + +template +__device__ inline errc pow(optional* out, optional const* a, optional const* b) +{ + if (a->has_value() && b->has_value()) { + T r; + pow(&r, &a->value(), &b->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc rint(float* out, float const* a) +{ + *out = ::rintf(*a); + return errc::OK; +} + +__device__ inline errc rint(double* out, double const* a) +{ + *out = ::rint(*a); + return errc::OK; +} + +template +__device__ inline errc rint(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + rint(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc sqrt(float* out, float const* a) +{ + *out = ::sqrtf(*a); + return errc::OK; +} + +__device__ inline errc sqrt(double* out, double const* a) +{ + *out = ::sqrt(*a); + return errc::OK; +} + +template +__device__ inline errc sqrt(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + sqrt(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/null_handling.cuh b/cpp/include/cudf/operators/null_handling.cuh new file mode 100644 index 000000000000..c107e9283d13 --- /dev/null +++ b/cpp/include/cudf/operators/null_handling.cuh @@ -0,0 +1,80 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +__device__ inline errc is_null(bool* out, T const* a) +{ + *out = false; + return errc::OK; +} + +template +__device__ inline errc is_null(optional* out, optional const* a) +{ + *out = !a->has_value(); + return errc::OK; +} + +template +__device__ inline errc nullify_if(optional* out, + optional const* a, + optional const* condition) +{ + if (condition->has_value() && a->has_value()) { + if (condition->value()) { + *out = nullopt; + } else { + *out = a->value(); + } + } else { + *out = nullopt; + } + return errc::OK; +} + +template +__device__ inline errc coalesce(T* out, T const* a, T const* b) +{ + *out = *a; + return errc::OK; +} + +template +__device__ inline errc coalesce(optional* out, optional const* a, optional const* b) +{ + if (a->has_value()) { + *out = a->value(); + } else if (b->has_value()) { + *out = b->value(); + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc predicate(bool* out, bool const* a) +{ + *out = *a; + return errc::OK; +} + +__device__ inline errc predicate(optional* out, optional const* a) +{ + if (a->has_value()) { + *out = a->value(); + } else { + *out = false; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/op_traits.hpp b/cpp/include/cudf/operators/op_traits.hpp new file mode 100644 index 000000000000..8d3cdfaa8471 --- /dev/null +++ b/cpp/include/cudf/operators/op_traits.hpp @@ -0,0 +1,428 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include + +#include +#include + +namespace cudf::detail::row_ir { + +enum [[nodiscard]] type : uint64_t { + NONE = 0x0, + BOOL8 = 0x1, + INT8 = 0x2, + INT16 = 0x4, + INT32 = 0x8, + INT64 = 0x10, + UINT8 = 0x20, + UINT16 = 0x40, + UINT32 = 0x80, + UINT64 = 0x100, + FLOAT32 = 0x200, + FLOAT64 = 0x400, + DECIMAL32 = 0x800, + DECIMAL64 = 0x1000, + DECIMAL128 = 0x2000, + TIMESTAMP_DAYS = 0x4000, + TIMESTAMP_SECONDS = 0x8000, + TIMESTAMP_MILLISECONDS = 0x10000, + TIMESTAMP_MICROSECONDS = 0x20000, + TIMESTAMP_NANOSECONDS = 0x40000, + DURATION_DAYS = 0x80000, + DURATION_SECONDS = 0x100000, + DURATION_MILLISECONDS = 0x200000, + DURATION_MICROSECONDS = 0x400000, + DURATION_NANOSECONDS = 0x800000, + STRING = 0x1000000, + SIGNED_INTEGERS = INT8 | INT16 | INT32 | INT64, + UNSIGNED_INTEGERS = UINT8 | UINT16 | UINT32 | UINT64, + INTEGERS = SIGNED_INTEGERS | UNSIGNED_INTEGERS, + FLOATS = FLOAT32 | FLOAT64, + DECIMALS = DECIMAL32 | DECIMAL64 | DECIMAL128, + ARITHMETIC = INTEGERS | FLOATS | DECIMALS, + SIGNED_ARITHMETIC = SIGNED_INTEGERS | FLOATS | DECIMALS, + ALL = 0x0FFFFFFF, + ARG_MASK = 0x10000000, + ARG0 = 0x10000000, + ARG1 = 0x10000001, + ARG2 = 0x10000002, + ARG3 = 0x10000003, + INPUT = 0x20000000, +}; + +struct [[nodiscard]] op_type { + type output = type::NONE; + std::vector args = {}; +}; + +/** + * @brief Indicates how an operator propagates null values + */ +enum class [[nodiscard]] null_output : uint8_t { + PROPAGATE = 0, + ALWAYS_VALID = 1, + ALWAYS_NULLABLE = 2, +}; + +[[nodiscard]] inline std::string_view get_op_name(opcode op) +{ + switch (op) { + case opcode::GET_INPUT: return "get_input"; + case opcode::SET_OUTPUT: return "set_output"; + case opcode::IDENTITY: return "identity"; + case opcode::IS_NULL: return "is_null"; + case opcode::NULLIFY_IF: return "nullify_if"; + case opcode::COALESCE: return "coalesce"; + case opcode::PREDICATE: return "predicate"; + case opcode::ABS: return "abs"; + case opcode::ADD: return "add"; + case opcode::DIV: return "div"; + case opcode::TRUE_DIV: return "true_div"; + case opcode::FLOOR_DIV: return "floor_div"; + case opcode::MOD: return "mod"; + case opcode::PYMOD: return "pymod"; + case opcode::MUL: return "mul"; + case opcode::NEG: return "neg"; + case opcode::SUB: return "sub"; + case opcode::ANSI_ADD: return "ansi_add"; + case opcode::ANSI_SUB: return "ansi_sub"; + case opcode::ANSI_MUL: return "ansi_mul"; + case opcode::ANSI_DIV: return "ansi_div"; + case opcode::ANSI_MOD: return "ansi_mod"; + case opcode::ANSI_ABS: return "ansi_abs"; + case opcode::ANSI_NEG: return "ansi_neg"; + case opcode::ANSI_PRECISION_CHECK: return "ansi_precision_check"; + case opcode::ANSI_TRY_ADD: return "ansi_try_add"; + case opcode::ANSI_TRY_SUB: return "ansi_try_sub"; + case opcode::ANSI_TRY_MUL: return "ansi_try_mul"; + case opcode::ANSI_TRY_DIV: return "ansi_try_div"; + case opcode::ANSI_TRY_MOD: return "ansi_try_mod"; + case opcode::ANSI_TRY_ABS: return "ansi_try_abs"; + case opcode::ANSI_TRY_NEG: return "ansi_try_neg"; + case opcode::ANSI_TRY_PRECISION_CHECK: return "ansi_try_precision_check"; + case opcode::BIT_AND: return "bit_and"; + case opcode::BIT_INVERT: return "bit_invert"; + case opcode::BIT_OR: return "bit_or"; + case opcode::BIT_XOR: return "bit_xor"; + case opcode::BIT_SHIFT_LEFT: return "bit_shift_left"; + case opcode::BIT_SHIFT_RIGHT: return "bit_shift_right"; + case opcode::CAST_TO_B8: return "cast_to_b8"; + case opcode::CAST_TO_I8: return "cast_to_i8"; + case opcode::CAST_TO_I16: return "cast_to_i16"; + case opcode::CAST_TO_I32: return "cast_to_i32"; + case opcode::CAST_TO_I64: return "cast_to_i64"; + case opcode::CAST_TO_U8: return "cast_to_u8"; + case opcode::CAST_TO_U16: return "cast_to_u16"; + case opcode::CAST_TO_U32: return "cast_to_u32"; + case opcode::CAST_TO_U64: return "cast_to_u64"; + case opcode::CAST_TO_F32: return "cast_to_f32"; + case opcode::CAST_TO_F64: return "cast_to_f64"; + case opcode::CAST_TO_DEC32: return "cast_to_dec32"; + case opcode::CAST_TO_DEC64: return "cast_to_dec64"; + case opcode::CAST_TO_DEC128: return "cast_to_dec128"; + case opcode::RESCALE: return "rescale"; + case opcode::EQUAL: return "equal"; + case opcode::NOT_EQUAL: return "not_equal"; + case opcode::GREATER: return "greater"; + case opcode::GREATER_EQUAL: return "greater_equal"; + case opcode::LESS: return "less"; + case opcode::LESS_EQUAL: return "less_equal"; + case opcode::NULL_EQUAL: return "null_equal"; + case opcode::NULL_LOGICAL_AND: return "null_logical_and"; + case opcode::NULL_LOGICAL_OR: return "null_logical_or"; + case opcode::LOGICAL_AND: return "logical_and"; + case opcode::LOGICAL_OR: return "logical_or"; + case opcode::LOGICAL_NOT: return "logical_not"; + case opcode::IF_ELSE: return "if_else"; + case opcode::CBRT: return "cbrt"; + case opcode::CEIL: return "ceil"; + case opcode::FLOOR: return "floor"; + case opcode::RINT: return "rint"; + case opcode::SQRT: return "sqrt"; + case opcode::POW: return "pow"; + case opcode::EXP: return "exp"; + case opcode::LOG: return "log"; + case opcode::ARCCOS: return "arccos"; + case opcode::ARCCOSH: return "arccosh"; + case opcode::ARCSIN: return "arcsin"; + case opcode::ARCSINH: return "arcsinh"; + case opcode::ARCTAN: return "arctan"; + case opcode::ARCTANH: return "arctanh"; + case opcode::COS: return "cos"; + case opcode::COSH: return "cosh"; + case opcode::SIN: return "sin"; + case opcode::SINH: return "sinh"; + case opcode::TAN: return "tan"; + case opcode::TANH: return "tanh"; + default: CUDF_UNREACHABLE("Invalid opcode"); + } +} + +[[nodiscard]] inline null_output get_op_null_output(opcode op) +{ + switch (op) { + case opcode::IS_NULL: + case opcode::NULL_EQUAL: + case opcode::PREDICATE: return null_output::ALWAYS_VALID; + + case opcode::NULLIFY_IF: + case opcode::COALESCE: + case opcode::ANSI_TRY_ADD: + case opcode::ANSI_TRY_SUB: + case opcode::ANSI_TRY_MUL: + case opcode::ANSI_TRY_DIV: + case opcode::ANSI_TRY_MOD: + case opcode::ANSI_TRY_ABS: + case opcode::ANSI_TRY_NEG: + case opcode::ANSI_TRY_PRECISION_CHECK: + case opcode::NULL_LOGICAL_AND: + case opcode::NULL_LOGICAL_OR: return null_output::ALWAYS_NULLABLE; + + default: return null_output::PROPAGATE; + } +} + +/** + * @brief Indicates whether the output of the operator will be different when it is called with or + * without the null-ness of a value. + */ +[[nodiscard]] inline bool get_op_requires_nulls(opcode op) +{ + switch (op) { + case opcode::COALESCE: + case opcode::IS_NULL: + case opcode::NULL_EQUAL: + case opcode::NULL_LOGICAL_AND: + case opcode::NULL_LOGICAL_OR: + case opcode::PREDICATE: return true; + + default: return false; + } +} + +[[nodiscard]] inline bool get_op_is_fallible(opcode op) +{ + switch (op) { + case opcode::ANSI_ADD: + case opcode::ANSI_SUB: + case opcode::ANSI_MUL: + case opcode::ANSI_DIV: + case opcode::ANSI_MOD: + case opcode::ANSI_ABS: + case opcode::ANSI_NEG: + case opcode::ANSI_PRECISION_CHECK: return true; + + default: return false; + } +} + +[[nodiscard]] inline int32_t op_rescale(opcode op, + std::span arg_scales, + std::optional target_scale) +{ + switch (op) { + case opcode::GET_INPUT: return 0; + case opcode::SET_OUTPUT: + case opcode::IDENTITY: + case opcode::COALESCE: + case opcode::PREDICATE: + case opcode::IS_NULL: + case opcode::ABS: + case opcode::NEG: + case opcode::ANSI_ABS: + case opcode::ANSI_NEG: + case opcode::ANSI_TRY_NEG: + case opcode::ANSI_TRY_ABS: + case opcode::CAST_TO_B8: + case opcode::CAST_TO_I8: + case opcode::CAST_TO_I16: + case opcode::CAST_TO_I32: + case opcode::CAST_TO_I64: + case opcode::CAST_TO_U8: + case opcode::CAST_TO_U16: + case opcode::CAST_TO_U32: + case opcode::CAST_TO_U64: + case opcode::CAST_TO_F32: + case opcode::CAST_TO_F64: + case opcode::CAST_TO_DEC32: + case opcode::CAST_TO_DEC64: + case opcode::CAST_TO_DEC128: + case opcode::NULLIFY_IF: + case opcode::EQUAL: + case opcode::GREATER: + case opcode::GREATER_EQUAL: + case opcode::LESS: + case opcode::LESS_EQUAL: + case opcode::NOT_EQUAL: + case opcode::NULL_EQUAL: + case opcode::NULL_LOGICAL_AND: + case opcode::NULL_LOGICAL_OR: + case opcode::LOGICAL_AND: + case opcode::LOGICAL_OR: + case opcode::LOGICAL_NOT: + case opcode::IF_ELSE: + case opcode::ARCCOS: + case opcode::ARCCOSH: + case opcode::ARCSIN: + case opcode::ARCSINH: + case opcode::ARCTAN: + case opcode::ARCTANH: + case opcode::COS: + case opcode::COSH: + case opcode::SIN: + case opcode::SINH: + case opcode::TAN: + case opcode::TANH: + case opcode::ANSI_PRECISION_CHECK: + case opcode::ANSI_TRY_PRECISION_CHECK: + case opcode::BIT_AND: + case opcode::BIT_INVERT: + case opcode::BIT_OR: + case opcode::BIT_XOR: + case opcode::BIT_SHIFT_LEFT: + case opcode::BIT_SHIFT_RIGHT: + case opcode::CBRT: + case opcode::CEIL: + case opcode::FLOOR: + case opcode::RINT: + case opcode::SQRT: + case opcode::POW: + case opcode::EXP: + case opcode::TRUE_DIV: + case opcode::LOG: return arg_scales[0]; + case opcode::FLOOR_DIV: + case opcode::ANSI_DIV: + case opcode::DIV: + case opcode::ANSI_TRY_DIV: return arg_scales[0] - arg_scales[1]; + case opcode::ADD: + case opcode::SUB: + case opcode::ANSI_ADD: + case opcode::ANSI_SUB: + case opcode::ANSI_TRY_ADD: + case opcode::ANSI_TRY_SUB: + case opcode::MOD: + case opcode::ANSI_MOD: + case opcode::ANSI_TRY_MOD: + case opcode::PYMOD: return std::min(arg_scales[0], arg_scales[1]); + case opcode::MUL: + case opcode::ANSI_MUL: + case opcode::ANSI_TRY_MUL: return arg_scales[0] + arg_scales[1]; + case opcode::RESCALE: return target_scale.value(); + default: CUDF_UNREACHABLE("Invalid opcode"); + } +} + +/** + * @brief Get the typing information for a given operator + * This function returns the expected input and output types for a given operator. The typing + * information can be used for type checking and inference when constructing expression trees. + * @param op The operator for which to get the typing information + * @return An `op_typing` struct containing the expected output type and input types for the + * operator + */ +[[nodiscard]] inline op_type get_op_typing(opcode op) +{ + switch (op) { + case opcode::GET_INPUT: return {type::INPUT, {}}; + case opcode::SET_OUTPUT: return {type::NONE, {type::ALL}}; + case opcode::IDENTITY: return {type::ARG0, {type::ALL}}; + case opcode::IS_NULL: return {type::BOOL8, {type::ALL}}; + case opcode::NULLIFY_IF: return {type::ARG0, {type::ALL, type::BOOL8}}; + case opcode::COALESCE: return {type::ARG0, {type::ALL, type::ARG0}}; + case opcode::PREDICATE: return {type::ARG0, {type::BOOL8}}; + case opcode::ABS: + case opcode::NEG: + case opcode::ANSI_ABS: + case opcode::ANSI_NEG: + case opcode::ANSI_TRY_NEG: + case opcode::ANSI_TRY_ABS: return {type::ARG0, {type::ARITHMETIC}}; + case opcode::FLOOR_DIV: return {type::ARG0, {type{type::FLOATS | type::INTEGERS}, type::ARG0}}; + case opcode::TRUE_DIV: + return {type::FLOAT64, {type{type::FLOATS | type::INTEGERS}, type::ARG0}}; + case opcode::ADD: + case opcode::DIV: + case opcode::MOD: + case opcode::PYMOD: + case opcode::MUL: + case opcode::SUB: + case opcode::ANSI_ADD: + case opcode::ANSI_SUB: + case opcode::ANSI_MUL: + case opcode::ANSI_DIV: + case opcode::ANSI_MOD: + case opcode::ANSI_TRY_ADD: + case opcode::ANSI_TRY_SUB: + case opcode::ANSI_TRY_MUL: + case opcode::ANSI_TRY_DIV: + case opcode::ANSI_TRY_MOD: return {type::ARG0, {type::ARITHMETIC, type::ARG0}}; + case opcode::ANSI_PRECISION_CHECK: + case opcode::ANSI_TRY_PRECISION_CHECK: return {type::ARG0, {type::DECIMALS, type::INT32}}; + case opcode::BIT_AND: + case opcode::BIT_INVERT: + case opcode::BIT_OR: + case opcode::BIT_XOR: + case opcode::BIT_SHIFT_LEFT: + case opcode::BIT_SHIFT_RIGHT: return {type::ARG0, {type::INTEGERS, type::ARG0}}; + case opcode::CAST_TO_B8: return {type::BOOL8, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_I8: return {type::INT8, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_I16: return {type::INT16, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_I32: return {type::INT32, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_I64: return {type::INT64, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_U8: return {type::UINT8, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_U16: return {type::UINT16, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_U32: return {type::UINT32, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_U64: return {type::UINT64, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_F32: return {type::FLOAT32, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_F64: return {type::FLOAT64, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::CAST_TO_DEC32: return {type::DECIMAL32, {type::DECIMALS}}; + case opcode::CAST_TO_DEC64: return {type::DECIMAL64, {type::DECIMALS}}; + case opcode::CAST_TO_DEC128: return {type::DECIMAL128, {type::DECIMALS}}; + case opcode::RESCALE: return {type::ARG0, {type::DECIMALS, type::INT32}}; + case opcode::EQUAL: + case opcode::GREATER: + case opcode::GREATER_EQUAL: + case opcode::LESS: + case opcode::LESS_EQUAL: + case opcode::NOT_EQUAL: + case opcode::NULL_EQUAL: return {type::BOOL8, {type::ALL, type::ARG0}}; + case opcode::NULL_LOGICAL_AND: + case opcode::NULL_LOGICAL_OR: + case opcode::LOGICAL_AND: + case opcode::LOGICAL_OR: + return {type::BOOL8, {type{type::ARITHMETIC | type::BOOL8}, type::ARG0}}; + case opcode::LOGICAL_NOT: return {type::BOOL8, {type{type::ARITHMETIC | type::BOOL8}}}; + case opcode::IF_ELSE: return {type::ARG0, {type::ALL, type::ARG0, type::BOOL8}}; + case opcode::CBRT: + case opcode::CEIL: + case opcode::FLOOR: + case opcode::RINT: + case opcode::SQRT: + case opcode::POW: + case opcode::EXP: + case opcode::LOG: + case opcode::ARCCOS: + case opcode::ARCCOSH: + case opcode::ARCSIN: + case opcode::ARCSINH: + case opcode::ARCTAN: + case opcode::ARCTANH: + case opcode::COS: + case opcode::COSH: + case opcode::SIN: + case opcode::SINH: + case opcode::TAN: + case opcode::TANH: return {type::ARG0, {type::FLOATS}}; + default: CUDF_UNREACHABLE("Invalid opcode"); + } +} + +[[nodiscard]] inline int32_t get_op_arity(opcode op) +{ + return static_cast(get_op_typing(op).args.size()); +} + +} // namespace cudf::detail::row_ir diff --git a/cpp/include/cudf/operators/opcodes.hpp b/cpp/include/cudf/operators/opcodes.hpp new file mode 100644 index 000000000000..b0e9c9cfcf3e --- /dev/null +++ b/cpp/include/cudf/operators/opcodes.hpp @@ -0,0 +1,127 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once +#include +#include + +namespace cudf { +namespace detail { +namespace row_ir { + +enum class [[nodiscard]] opcode : int32_t { + GET_INPUT, + SET_OUTPUT, + + // Identity operators + IDENTITY, + + // Null handling operators + IS_NULL, + NULLIFY_IF, + + COALESCE, + PREDICATE, + + /// Arithmetic operators + ABS, + ADD, + DIV, + TRUE_DIV, + FLOOR_DIV, + MOD, + PYMOD, + MUL, + NEG, + SUB, + + /// ANSI Arithmetic functions. raise errors on overflow, division by zero, etc. + ANSI_ADD, + ANSI_SUB, + ANSI_MUL, + ANSI_DIV, + ANSI_MOD, + ANSI_ABS, + ANSI_NEG, + ANSI_PRECISION_CHECK, + + /// ANSI TRY arithmetic functions. return NULL instead of raising errors + ANSI_TRY_ADD, + ANSI_TRY_SUB, + ANSI_TRY_MUL, + ANSI_TRY_DIV, + ANSI_TRY_MOD, + ANSI_TRY_ABS, + ANSI_TRY_NEG, + ANSI_TRY_PRECISION_CHECK, + + /// Bitwise operators + BIT_AND, + BIT_INVERT, + BIT_OR, + BIT_XOR, + BIT_SHIFT_LEFT, + BIT_SHIFT_RIGHT, + + /// Type conversion operators + CAST_TO_B8, + CAST_TO_I8, + CAST_TO_I16, + CAST_TO_I32, + CAST_TO_I64, + CAST_TO_U8, + CAST_TO_U16, + CAST_TO_U32, + CAST_TO_U64, + CAST_TO_F32, + CAST_TO_F64, + CAST_TO_DEC32, + CAST_TO_DEC64, + CAST_TO_DEC128, + RESCALE, + + /// Comparison & Logic operators + EQUAL, + NOT_EQUAL, + GREATER, + GREATER_EQUAL, + LESS, + LESS_EQUAL, + NULL_EQUAL, + NULL_LOGICAL_AND, + NULL_LOGICAL_OR, + LOGICAL_AND, + LOGICAL_OR, + LOGICAL_NOT, + IF_ELSE, + + /// Mathematical operators + CBRT, + CEIL, + FLOOR, + RINT, + SQRT, + POW, + EXP, + LOG, + + /// Trigonometric operators + ARCCOS, + ARCCOSH, + ARCSIN, + ARCSINH, + ARCTAN, + ARCTANH, + COS, + COSH, + SIN, + SINH, + TAN, + TANH, +}; + +} // namespace row_ir +} // namespace detail +} // namespace cudf diff --git a/cpp/include/cudf/operators/trigonometric.cuh b/cpp/include/cudf/operators/trigonometric.cuh new file mode 100644 index 000000000000..1156c3eba2e3 --- /dev/null +++ b/cpp/include/cudf/operators/trigonometric.cuh @@ -0,0 +1,313 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +__device__ inline errc arccos(float* out, float const* a) +{ + *out = ::acosf(*a); + return errc::OK; +} + +__device__ inline errc arccos(double* out, double const* a) +{ + *out = ::acos(*a); + return errc::OK; +} + +template +__device__ inline errc arccos(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arccos(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc arccosh(float* out, float const* a) +{ + *out = ::acoshf(*a); + return errc::OK; +} + +__device__ inline errc arccosh(double* out, double const* a) +{ + *out = ::acosh(*a); + return errc::OK; +} + +template +__device__ inline errc arccosh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arccosh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc arcsin(float* out, float const* a) +{ + *out = ::asinf(*a); + return errc::OK; +} + +__device__ inline errc arcsin(double* out, double const* a) +{ + *out = ::asin(*a); + return errc::OK; +} + +template +__device__ inline errc arcsin(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arcsin(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc arcsinh(float* out, float const* a) +{ + *out = ::asinhf(*a); + return errc::OK; +} + +__device__ inline errc arcsinh(double* out, double const* a) +{ + *out = ::asinh(*a); + return errc::OK; +} + +template +__device__ inline errc arcsinh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arcsinh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc arctan(float* out, float const* a) +{ + *out = ::atanf(*a); + return errc::OK; +} + +__device__ inline errc arctan(double* out, double const* a) +{ + *out = ::atan(*a); + return errc::OK; +} + +template +__device__ inline errc arctan(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arctan(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc arctanh(float* out, float const* a) +{ + *out = ::atanhf(*a); + return errc::OK; +} + +__device__ inline errc arctanh(double* out, double const* a) +{ + *out = ::atanh(*a); + return errc::OK; +} + +template +__device__ inline errc arctanh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + arctanh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc cos(float* out, float const* a) +{ + *out = ::cosf(*a); + return errc::OK; +} + +__device__ inline errc cos(double* out, double const* a) +{ + *out = ::cos(*a); + return errc::OK; +} + +template +__device__ inline errc cos(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + cos(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc cosh(float* out, float const* a) +{ + *out = ::coshf(*a); + return errc::OK; +} + +__device__ inline errc cosh(double* out, double const* a) +{ + *out = ::cosh(*a); + return errc::OK; +} + +template +__device__ inline errc cosh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + cosh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc sin(float* out, float const* a) +{ + *out = ::sinf(*a); + return errc::OK; +} + +__device__ inline errc sin(double* out, double const* a) +{ + *out = ::sin(*a); + return errc::OK; +} + +template +__device__ inline errc sin(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + sin(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc sinh(float* out, float const* a) +{ + *out = ::sinhf(*a); + return errc::OK; +} + +__device__ inline errc sinh(double* out, double const* a) +{ + *out = ::sinh(*a); + return errc::OK; +} + +template +__device__ inline errc sinh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + sinh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc tan(float* out, float const* a) +{ + *out = ::tanf(*a); + return errc::OK; +} + +__device__ inline errc tan(double* out, double const* a) +{ + *out = ::tan(*a); + return errc::OK; +} + +template +__device__ inline errc tan(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + tan(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +__device__ inline errc tanh(float* out, float const* a) +{ + *out = ::tanhf(*a); + return errc::OK; +} + +__device__ inline errc tanh(double* out, double const* a) +{ + *out = ::tanh(*a); + return errc::OK; +} + +template +__device__ inline errc tanh(optional* out, optional const* a) +{ + if (a->has_value()) { + T r; + tanh(&r, &a->value()); + *out = r; + } else { + *out = nullopt; + } + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/operators/types.cuh b/cpp/include/cudf/operators/types.cuh new file mode 100644 index 000000000000..673a6b883016 --- /dev/null +++ b/cpp/include/cudf/operators/types.cuh @@ -0,0 +1,69 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace CUDF_EXPORT cudf { +namespace ops { + +template +using optional = cuda::std::optional; + +inline constexpr auto nullopt = cuda::std::nullopt; + +template +using decimal = numeric::fixed_point; + +template +using duration = cuda::std::chrono::duration; + +namespace detail { + +template +__device__ constexpr T ipow10(T exponent) +{ + if (exponent == 0) { return 1; } + + T extra = 1; + T square = 10; + T n = exponent; + + while (n > 1) { + if ((n & 1) == 1) { extra *= square; } + n >>= 1; + square *= square; + } + + return square * extra; +} + +} // namespace detail + +template +__device__ inline errc identity(T* out, T const* a) +{ + *out = *a; + return errc::OK; +} + +template +__device__ inline errc identity(optional* out, optional const* a) +{ + *out = *a; + return errc::OK; +} + +} // namespace ops +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/transform.hpp b/cpp/include/cudf/transform.hpp index a2a99ed0c23b..1936c5284b2e 100644 --- a/cpp/include/cudf/transform.hpp +++ b/cpp/include/cudf/transform.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -163,6 +164,7 @@ std::unique_ptr transform_extended( * input columns. * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned column's device memory + * @param error_mode The error handling mode to use during the transform * @return A table containing the columns resulting from applying the transform * function to every element of the input according to the output specifications * @@ -176,6 +178,7 @@ std::unique_ptr multi_transform( std::span outputs, std::vector>&& string_offsets, std::optional row_size, + ops::error_mode error_mode = ops::error_mode::IGNORE, rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); diff --git a/cpp/src/ast/expressions.cpp b/cpp/src/ast/expressions.cpp index d51d3f323498..9e6b8cf3ce8c 100644 --- a/cpp/src/ast/expressions.cpp +++ b/cpp/src/ast/expressions.cpp @@ -80,25 +80,24 @@ bool operation::may_evaluate_null(table_view const& left, }); }; -cudf::size_type detail::filter_predicate::accept(detail::expression_parser& visitor) const +cudf::size_type detail::predicate::accept(detail::expression_parser& visitor) const { - CUDF_FAIL( - "filter_predicate is an internal expression and should not be visited by expression_parser", - std::invalid_argument); + CUDF_FAIL("predicate is an internal expression and should not be visited by expression_parser", + std::invalid_argument); } -std::reference_wrapper detail::filter_predicate::accept( +std::reference_wrapper detail::predicate::accept( detail::expression_transformer& visitor) const { CUDF_FAIL( - "filter_predicate is an internal expression and should not be visited by " + "predicate is an internal expression and should not be visited by " "expression_transformer", std::invalid_argument); } -bool detail::filter_predicate::may_evaluate_null(table_view const& left, - table_view const& right, - rmm::cuda_stream_view stream) const +bool detail::predicate::may_evaluate_null(table_view const& left, + table_view const& right, + rmm::cuda_stream_view stream) const { return false; } @@ -135,7 +134,7 @@ std::unique_ptr column_name_reference::accept( std::invalid_argument); } -std::unique_ptr detail::filter_predicate::accept( +std::unique_ptr detail::predicate::accept( cudf::detail::row_ir::ast_converter& converter) const { return converter.add_ir_node(*this); diff --git a/cpp/src/ast/jit_expressions.cpp b/cpp/src/ast/jit_expressions.cpp new file mode 100644 index 000000000000..3911b1ee6c8a --- /dev/null +++ b/cpp/src/ast/jit_expressions.cpp @@ -0,0 +1,232 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#include "cudf/ast/jit_expressions.hpp" + +#include "jit/row_ir.hpp" + +namespace cudf { +namespace ast { + +namespace jit::detail { + +cudf::size_type operation::accept(cudf::ast::detail::expression_parser& visitor) const +{ + CUDF_FAIL("predicate is an internal expression and should not be visited by expression_parser", + std::invalid_argument); +} + +std::reference_wrapper operation::accept( + cudf::ast::detail::expression_transformer& visitor) const +{ + CUDF_FAIL( + "predicate is an internal expression and should not be visited by " + "expression_transformer", + std::invalid_argument); +} + +bool operation::may_evaluate_null(table_view const& left, + table_view const& right, + rmm::cuda_stream_view stream) const +{ + CUDF_FAIL("predicate is an internal expression and should not be evaluated directly", + std::invalid_argument); +} + +std::unique_ptr operation::accept( + cudf::detail::row_ir::ast_converter& converter) const +{ + return converter.add_ir_node(*this); +} + +} // namespace jit::detail + +expression const& jit::nullify_if(ast::tree& tree, expression const& a, expression const& condition) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::NULLIFY_IF, {a, condition})); +} + +expression const& jit::coalesce(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::COALESCE, {a, b})); +} + +expression const& jit::predicate(ast::tree& tree, expression const& condition) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::PREDICATE, {condition})); +} + +expression const& jit::ansi_add(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_ADD, {a, b})); +} + +expression const& jit::ansi_sub(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_SUB, {a, b})); +} + +expression const& jit::ansi_mul(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_MUL, {a, b})); +} + +expression const& jit::ansi_div(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_DIV, {a, b})); +} + +expression const& jit::ansi_mod(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_MOD, {a, b})); +} + +expression const& jit::ansi_abs(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_ABS, {a})); +} + +expression const& jit::ansi_neg(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_NEG, {a})); +} + +expression const& jit::ansi_precision_check(ast::tree& tree, + expression const& a, + expression const& precision) +{ + return tree.push( + detail::operation(cudf::detail::row_ir::opcode::ANSI_PRECISION_CHECK, {a, precision})); +} + +expression const& jit::ansi_try_add(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_ADD, {a, b})); +} + +expression const& jit::ansi_try_sub(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_SUB, {a, b})); +} + +expression const& jit::ansi_try_mul(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_MUL, {a, b})); +} + +expression const& jit::ansi_try_div(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_DIV, {a, b})); +} + +expression const& jit::ansi_try_mod(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_MOD, {a, b})); +} + +expression const& jit::ansi_try_abs(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_ABS, {a})); +} + +expression const& jit::ansi_try_neg(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_NEG, {a})); +} + +expression const& jit::ansi_try_precision_check(ast::tree& tree, + expression const& a, + expression const& precision) +{ + return tree.push( + detail::operation(cudf::detail::row_ir::opcode::ANSI_TRY_PRECISION_CHECK, {a, precision})); +} + +expression const& jit::bit_shift_left(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::BIT_SHIFT_LEFT, {a, b})); +} + +expression const& jit::bit_shift_right(ast::tree& tree, expression const& a, expression const& b) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::BIT_SHIFT_RIGHT, {a, b})); +} + +expression const& jit::cast_to_b8(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_B8, {a})); +} + +expression const& jit::cast_to_i8(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_I8, {a})); +} + +expression const& jit::cast_to_i16(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_I16, {a})); +} + +expression const& jit::cast_to_i32(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_I32, {a})); +} + +expression const& jit::cast_to_i64(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_I64, {a})); +} + +expression const& jit::cast_to_u8(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_U8, {a})); +} + +expression const& jit::cast_to_u16(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_U16, {a})); +} + +expression const& jit::cast_to_u32(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_U32, {a})); +} + +expression const& jit::cast_to_u64(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_U64, {a})); +} + +expression const& jit::cast_to_f32(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_F32, {a})); +} + +expression const& jit::cast_to_f64(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_F64, {a})); +} + +expression const& jit::cast_to_dec32(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_DEC32, {a})); +} + +expression const& jit::cast_to_dec64(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_DEC64, {a})); +} + +expression const& jit::cast_to_dec128(ast::tree& tree, expression const& a) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::CAST_TO_DEC128, {a})); +} + +expression const& jit::rescale(ast::tree& tree, expression const& a, int32_t target_scale) +{ + return tree.push(detail::operation(cudf::detail::row_ir::opcode::RESCALE, {a}, target_scale)); +} + +} // namespace ast +} // namespace cudf diff --git a/cpp/src/jit/column_accessor.cuh b/cpp/src/jit/column_accessor.cuh index 303e8cb029f3..ef1536cb409d 100644 --- a/cpp/src/jit/column_accessor.cuh +++ b/cpp/src/jit/column_accessor.cuh @@ -16,12 +16,13 @@ namespace cudf { namespace jit { -template +template struct column_accessor { - static constexpr int32_t index = Index; - using column_type = Column; - using element_type = Element; - using optional_element_type = cuda::std::optional; + static constexpr int32_t index = Index; + static constexpr int32_t table_index = TableIndex; + using column_type = Column; + using element_type = Element; + using optional_element_type = cuda::std::optional; static constexpr bool as_scalar = AsScalar; diff --git a/cpp/src/jit/error_sink.cuh b/cpp/src/jit/error_sink.cuh new file mode 100644 index 000000000000..9063663d9615 --- /dev/null +++ b/cpp/src/jit/error_sink.cuh @@ -0,0 +1,38 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include + +namespace cudf { +namespace jit { + +struct error_sink { + ops::errc any_error_ = ops::errc::OK; + + constexpr error_sink() = default; + + template + __device__ void report(ops::errc error) + { + if constexpr (mode == ops::error_mode::IGNORE) { + return; + } else { + if (error != ops::errc::OK) [[unlikely]] { + cuda::std::atomic_ref any_error_ref{any_error_}; + any_error_ref.store(error, cuda::std::memory_order_relaxed); + } + } + } + + [[nodiscard]] __host__ __device__ constexpr ops::errc any_error() const { return any_error_; } +}; + +} // namespace jit +} // namespace cudf diff --git a/cpp/src/jit/join_column_accessor.cuh b/cpp/src/jit/join_column_accessor.cuh deleted file mode 100644 index e78a9ec1991c..000000000000 --- a/cpp/src/jit/join_column_accessor.cuh +++ /dev/null @@ -1,129 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include - -#include -#include - -namespace cudf { -namespace jit { - -// Join-specific accessor for indexed table access. -// Receives both left and right table pointers plus both row indices, -// and selects the appropriate table based on the Side template parameter. -enum class join_side : bool { LEFT, RIGHT }; - -template -struct join_column_accessor { - using type = T; - static constexpr int32_t index = Index; - - static __device__ T element(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - cudf::size_type left_row_idx, - cudf::size_type right_row_idx, - cudf::size_type /* thread_idx */) - { - if constexpr (Side == join_side::LEFT) { - return left_tables[index].template element(left_row_idx); - } else { - return right_tables[index].template element(right_row_idx); - } - } - - static __device__ bool is_null(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - cudf::size_type left_row_idx, - cudf::size_type right_row_idx, - cudf::size_type /* thread_idx */) - { - if constexpr (Side == join_side::LEFT) { - return left_tables[index].is_null(left_row_idx); - } else { - return right_tables[index].is_null(right_row_idx); - } - } - - static __device__ bool is_valid(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - cudf::size_type left_row_idx, - cudf::size_type right_row_idx, - cudf::size_type /* thread_idx */) - { - if constexpr (Side == join_side::LEFT) { - return left_tables[index].is_valid(left_row_idx); - } else { - return right_tables[index].is_valid(right_row_idx); - } - } - - static __device__ cuda::std::optional nullable_element( - cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - cudf::size_type left_row_idx, - cudf::size_type right_row_idx, - cudf::size_type thread_idx) - { - if (is_null(left_tables, right_tables, left_row_idx, right_row_idx, thread_idx)) { - return cuda::std::nullopt; - } - return element(left_tables, right_tables, left_row_idx, right_row_idx, thread_idx); - } -}; - -// Join-specific accessor for scalar (literal) values. -// Scalar columns are appended to the left table's device views. -// Always reads at row 0 since scalar columns have size 1. -template -struct join_scalar_accessor { - using type = T; - static constexpr int32_t index = Index; - - static __device__ T element(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const*, - cudf::size_type, - cudf::size_type, - cudf::size_type) - { - return left_tables[index].template element(0); - } - - static __device__ bool is_null(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const*, - cudf::size_type, - cudf::size_type, - cudf::size_type) - { - return left_tables[index].is_null(0); - } - - static __device__ bool is_valid(cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const*, - cudf::size_type, - cudf::size_type, - cudf::size_type) - { - return left_tables[index].is_valid(0); - } - - static __device__ cuda::std::optional nullable_element( - cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - cudf::size_type left_row_idx, - cudf::size_type right_row_idx, - cudf::size_type thread_idx) - { - if (is_null(left_tables, right_tables, left_row_idx, right_row_idx, thread_idx)) { - return cuda::std::nullopt; - } - return element(left_tables, right_tables, left_row_idx, right_row_idx, thread_idx); - } -}; - -} // namespace jit -} // namespace cudf diff --git a/cpp/src/jit/row_ir.cpp b/cpp/src/jit/row_ir.cpp index 82cd308c2f5d..cf331c54d0f1 100644 --- a/cpp/src/jit/row_ir.cpp +++ b/cpp/src/jit/row_ir.cpp @@ -12,22 +12,36 @@ #include #include #include -#include #include #include #include #include -namespace cudf { +namespace cudf::detail::row_ir { -namespace detail { - -namespace row_ir { +int32_t instance_context::add_output() +{ + auto id = static_cast(output_vars_.size()); + auto id_str = std::format("out_{}", id); + output_vars_.emplace_back(std::move(id_str)); + return id; +} -std::string cuda_type(cudf::data_type type, bool nullable) +int32_t instance_context::add_input(input in) { - auto name = type_to_name(type); - return nullable ? std::format("cuda::std::optional<{}>", name) : name; + auto id = static_cast(inputs_.size()); + auto id_str = std::format("in_{}", id); + + data_type type{type_id::EMPTY, 0}; + if (auto* col = std::get_if(&in)) { + type = col->column.type(); + } else { + auto& scalar = std::get(in); + type = scalar.scalar_column->type(); + } + inputs_.emplace_back(std::move(in)); + input_vars_.emplace_back(std::move(id_str), type); + return id; } std::string instance_context::make_tmp_id() @@ -39,641 +53,637 @@ bool instance_context::has_nulls() const { return has_nulls_; } void instance_context::set_has_nulls(bool has_nulls) { has_nulls_ = has_nulls; } -get_input::get_input(int32_t input) : id_(), input_(input), type_() {} +std::span instance_context::get_inputs() const { return inputs_; } -std::string_view get_input::get_id() { return id_; } +std::span instance_context::get_input_vars() const { return input_vars_; } -data_type get_input::get_type() { return type_; } +std::span instance_context::get_output_vars() const { return output_vars_; } -bool get_input::is_null_aware() { return false; } - -bool get_input::is_always_valid() { return false; } - -void get_input::instantiate(instance_context& ctx, instance_info const& info) +node::node(opcode op, std::optional target_scale, std::vector> args) + : op_{op}, target_scale_{target_scale}, args_{std::move(args)} { - id_ = ctx.make_tmp_id(); - auto const& input = info.inputs[input_]; - type_ = input.type; + CUDF_EXPECTS(op != opcode::GET_INPUT && op != opcode::SET_OUTPUT, + std::format("Invalid opcode `{}` for operation node.", get_op_name(op))); + if (op_ != opcode::RESCALE) { + CUDF_EXPECTS(args_.size() == static_cast(get_op_arity(op)), + std::format("Invalid number of arguments for operator `{}`. Expected {}, Got {}.", + get_op_name(op), + get_op_arity(op), + args_.size())); + } else { + CUDF_EXPECTS(args_.size() == 1, + std::format("RESCALE operator expects exactly 1 argument. Got {}.", args_.size())); + CUDF_EXPECTS( + target_scale_.has_value(), + std::format("Target scale must be provided for RESCALE operator and must be nullopt " + "for other operators.")); + } } -std::string get_input::generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) +node::node(input_reference input) : reference_{input}, op_{opcode::GET_INPUT} {} + +node::node(output_reference reference, std::unique_ptr arg) + : reference_{reference}, op_{opcode::SET_OUTPUT} { - switch (info.id) { - case target::CUDA: { - return std::format( - "{} {} = {};", cuda_type(type_, ctx.has_nulls()), id_, instance.inputs[input_].id); - } - default: - CUDF_FAIL("Unsupported target: " + std::to_string(static_cast(info.id)), - std::invalid_argument); - } + args_.emplace_back(std::move(arg)); } -set_output::set_output(int32_t output, std::unique_ptr source) - : id_(), output_(output), source_(std::move(source)), type_(), output_id_() +node::node(output_reference reference, node arg) + : node{reference, std::make_unique(std::move(arg))} { } -std::string_view set_output::get_id() { return id_; } +std::string_view node::get_id() const { return id_; } -data_type set_output::get_type() { return type_; } +data_type node::get_type() const { return type_; } -bool set_output::is_null_aware() { return source_->is_null_aware(); } +std::optional node::get_target_scale() const { return target_scale_; } -bool set_output::is_always_valid() { return source_->is_always_valid(); } +opcode node::get_opcode() const { return op_; } -node& set_output::get_source() { return *source_; } +std::span const> node::get_args() const { return args_; } -void set_output::instantiate(instance_context& ctx, instance_info const& info) +bool node::is_null_aware() const { - source_->instantiate(ctx, info); - id_ = ctx.make_tmp_id(); - auto source_type = source_->get_type(); - type_ = source_type; - output_id_ = info.outputs[output_].id; -} + if (op_ == opcode::GET_INPUT) { return false; } -std::string set_output::generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) -{ - switch (info.id) { - case target::CUDA: { - auto source_code = source_->generate_code(ctx, info, instance); - return std::format( - "{}\n" - "{} {} = {};\n" - "*{} = {};", - source_code, - cuda_type(type_, ctx.has_nulls()), - id_, - source_->get_id(), - output_id_, - id_); - } - default: - CUDF_FAIL("Unsupported target: " + std::to_string(static_cast(info.id)), - std::invalid_argument); - } + // to emit nulls for always-nullable operators, we need to mark them as null-aware + if (get_op_null_output(op_) == null_output::ALWAYS_NULLABLE) { return true; } + + if (get_op_requires_nulls(op_)) { return true; } + + CUDF_EXPECTS(!args_.empty(), + "Unexpectedly found an operator node with no arguments. All operator nodes should " + "have at least one argument."); + + return std::any_of(args_.begin(), args_.end(), [](auto& a) { return a->is_null_aware(); }); } -operation::operation(opcode op, std::unique_ptr* move_begin, std::unique_ptr* move_end) - : id_(), op_(op), operands_(), type_() +bool node::is_always_valid() const { - operands_.insert( - operands_.begin(), std::make_move_iterator(move_begin), std::make_move_iterator(move_end)); - CUDF_EXPECTS(static_cast(operands_.size()) == ast::detail::ast_operator_arity(op), - "Invalid number of arguments for operator.", - std::invalid_argument); - CUDF_EXPECTS( - operands_.size() > 0, "Operator must have at least one operand", std::invalid_argument); + if (op_ == opcode::GET_INPUT) { return false; } + + if (get_op_null_output(op_) == null_output::ALWAYS_VALID) { return true; } + + CUDF_EXPECTS(!args_.empty(), + "Unexpectedly found an operator node with no arguments. All operator nodes should " + "have at least one argument."); + + return std::all_of(args_.begin(), args_.end(), [](auto& a) { return a->is_always_valid(); }); } -operation::operation(opcode op, std::vector> operands) - : operation(op, operands.data(), operands.data() + operands.size()) +bool node::is_fallible() const { -} + if (op_ == opcode::GET_INPUT) { return false; } + + if (get_op_is_fallible(op_)) { return true; } -std::string_view operation::get_id() { return id_; } + CUDF_EXPECTS(!args_.empty(), + "Unexpectedly found an operator node with no arguments. All operator nodes should " + "have at least one argument."); -data_type operation::get_type() { return type_; } + return std::any_of(args_.begin(), args_.end(), [](auto& a) { return a->is_fallible(); }); +} -inline bool is_operator_null_aware(opcode op) +row_ir::type as_typing(data_type type) { - switch (op) { - case ast::ast_operator::IS_NULL: - case ast::ast_operator::NULL_EQUAL: - case ast::ast_operator::NULL_LOGICAL_AND: - case ast::ast_operator::NULL_LOGICAL_OR: return true; - - case ast::ast_operator::ADD: - case ast::ast_operator::SUB: - case ast::ast_operator::MUL: - case ast::ast_operator::DIV: - case ast::ast_operator::TRUE_DIV: - case ast::ast_operator::FLOOR_DIV: - case ast::ast_operator::MOD: - case ast::ast_operator::PYMOD: - case ast::ast_operator::POW: - case ast::ast_operator::NOT_EQUAL: - case ast::ast_operator::EQUAL: - case ast::ast_operator::LESS: - case ast::ast_operator::GREATER: - case ast::ast_operator::LESS_EQUAL: - case ast::ast_operator::GREATER_EQUAL: - case ast::ast_operator::BITWISE_AND: - case ast::ast_operator::BITWISE_OR: - case ast::ast_operator::BITWISE_XOR: - case ast::ast_operator::LOGICAL_AND: - case ast::ast_operator::LOGICAL_OR: - case ast::ast_operator::IDENTITY: - case ast::ast_operator::SIN: - case ast::ast_operator::COS: - case ast::ast_operator::TAN: - case ast::ast_operator::ARCSIN: - case ast::ast_operator::ARCCOS: - case ast::ast_operator::ARCTAN: - case ast::ast_operator::SINH: - case ast::ast_operator::COSH: - case ast::ast_operator::TANH: - case ast::ast_operator::ARCSINH: - case ast::ast_operator::ARCCOSH: - case ast::ast_operator::ARCTANH: - case ast::ast_operator::EXP: - case ast::ast_operator::LOG: - case ast::ast_operator::SQRT: - case ast::ast_operator::CBRT: - case ast::ast_operator::CEIL: - case ast::ast_operator::FLOOR: - case ast::ast_operator::ABS: - case ast::ast_operator::RINT: - case ast::ast_operator::BIT_INVERT: - case ast::ast_operator::NOT: - case ast::ast_operator::CAST_TO_INT64: - case ast::ast_operator::CAST_TO_UINT64: - case ast::ast_operator::CAST_TO_FLOAT64: return false; - - default: CUDF_UNREACHABLE("Unrecognized operator type."); + switch (type.id()) { + case type_id::BOOL8: return type::BOOL8; + case type_id::INT8: return type::INT8; + case type_id::INT16: return type::INT16; + case type_id::INT32: return type::INT32; + case type_id::INT64: return type::INT64; + case type_id::UINT8: return type::UINT8; + case type_id::UINT16: return type::UINT16; + case type_id::UINT32: return type::UINT32; + case type_id::UINT64: return type::UINT64; + case type_id::FLOAT32: return type::FLOAT32; + case type_id::FLOAT64: return type::FLOAT64; + case type_id::DECIMAL32: return type::DECIMAL32; + case type_id::DECIMAL64: return type::DECIMAL64; + case type_id::DECIMAL128: return type::DECIMAL128; + case type_id::TIMESTAMP_DAYS: return type::TIMESTAMP_DAYS; + case type_id::TIMESTAMP_SECONDS: return type::TIMESTAMP_SECONDS; + case type_id::TIMESTAMP_MILLISECONDS: return type::TIMESTAMP_MILLISECONDS; + case type_id::TIMESTAMP_MICROSECONDS: return type::TIMESTAMP_MICROSECONDS; + case type_id::TIMESTAMP_NANOSECONDS: return type::TIMESTAMP_NANOSECONDS; + case type_id::DURATION_DAYS: return type::DURATION_DAYS; + case type_id::DURATION_SECONDS: return type::DURATION_SECONDS; + case type_id::DURATION_MILLISECONDS: return type::DURATION_MILLISECONDS; + case type_id::DURATION_MICROSECONDS: return type::DURATION_MICROSECONDS; + case type_id::DURATION_NANOSECONDS: return type::DURATION_NANOSECONDS; + case type_id::STRING: return type::STRING; + default: + CUDF_FAIL(std::format("Unsupported data type for Row IR: {}", type_to_name(type)), + std::invalid_argument); } } -bool operation::is_null_aware() -{ - return is_operator_null_aware(op_) || - std::any_of( - operands_.begin(), operands_.end(), [](auto& op) { return op->is_null_aware(); }); +type_id as_type_id(type type) +{ + switch (type) { + case type::BOOL8: return type_id::BOOL8; + case type::INT8: return type_id::INT8; + case type::INT16: return type_id::INT16; + case type::INT32: return type_id::INT32; + case type::INT64: return type_id::INT64; + case type::UINT8: return type_id::UINT8; + case type::UINT16: return type_id::UINT16; + case type::UINT32: return type_id::UINT32; + case type::UINT64: return type_id::UINT64; + case type::FLOAT32: return type_id::FLOAT32; + case type::FLOAT64: return type_id::FLOAT64; + case type::DECIMAL32: return type_id::DECIMAL32; + case type::DECIMAL64: return type_id::DECIMAL64; + case type::DECIMAL128: return type_id::DECIMAL128; + case type::TIMESTAMP_DAYS: return type_id::TIMESTAMP_DAYS; + case type::TIMESTAMP_SECONDS: return type_id::TIMESTAMP_SECONDS; + case type::TIMESTAMP_MILLISECONDS: return type_id::TIMESTAMP_MILLISECONDS; + case type::TIMESTAMP_MICROSECONDS: return type_id::TIMESTAMP_MICROSECONDS; + case type::TIMESTAMP_NANOSECONDS: return type_id::TIMESTAMP_NANOSECONDS; + case type::DURATION_DAYS: return type_id::DURATION_DAYS; + case type::DURATION_SECONDS: return type_id::DURATION_SECONDS; + case type::DURATION_MILLISECONDS: return type_id::DURATION_MILLISECONDS; + case type::DURATION_MICROSECONDS: return type_id::DURATION_MICROSECONDS; + case type::DURATION_NANOSECONDS: return type_id::DURATION_NANOSECONDS; + case type::STRING: return type_id::STRING; + default: + CUDF_FAIL(std::format("Invalid typing for {}: {}", __FUNCTION__, static_cast(type)), + std::invalid_argument); + } } -inline bool is_operator_always_valid(opcode op) +opcode as_opcode(ast::ast_operator op) { switch (op) { - case ast::ast_operator::IS_NULL: - case ast::ast_operator::NULL_EQUAL: return true; - - case ast::ast_operator::NULL_LOGICAL_AND: - case ast::ast_operator::NULL_LOGICAL_OR: - case ast::ast_operator::ADD: - case ast::ast_operator::SUB: - case ast::ast_operator::MUL: - case ast::ast_operator::DIV: - case ast::ast_operator::TRUE_DIV: - case ast::ast_operator::FLOOR_DIV: - case ast::ast_operator::MOD: - case ast::ast_operator::PYMOD: - case ast::ast_operator::POW: - case ast::ast_operator::NOT_EQUAL: - case ast::ast_operator::EQUAL: - case ast::ast_operator::LESS: - case ast::ast_operator::GREATER: - case ast::ast_operator::LESS_EQUAL: - case ast::ast_operator::GREATER_EQUAL: - case ast::ast_operator::BITWISE_AND: - case ast::ast_operator::BITWISE_OR: - case ast::ast_operator::BITWISE_XOR: - case ast::ast_operator::LOGICAL_AND: - case ast::ast_operator::LOGICAL_OR: - case ast::ast_operator::IDENTITY: - case ast::ast_operator::SIN: - case ast::ast_operator::COS: - case ast::ast_operator::TAN: - case ast::ast_operator::ARCSIN: - case ast::ast_operator::ARCCOS: - case ast::ast_operator::ARCTAN: - case ast::ast_operator::SINH: - case ast::ast_operator::COSH: - case ast::ast_operator::TANH: - case ast::ast_operator::ARCSINH: - case ast::ast_operator::ARCCOSH: - case ast::ast_operator::ARCTANH: - case ast::ast_operator::EXP: - case ast::ast_operator::LOG: - case ast::ast_operator::SQRT: - case ast::ast_operator::CBRT: - case ast::ast_operator::CEIL: - case ast::ast_operator::FLOOR: - case ast::ast_operator::ABS: - case ast::ast_operator::RINT: - case ast::ast_operator::BIT_INVERT: - case ast::ast_operator::NOT: - case ast::ast_operator::CAST_TO_INT64: - case ast::ast_operator::CAST_TO_UINT64: - case ast::ast_operator::CAST_TO_FLOAT64: return false; - - default: CUDF_UNREACHABLE("Unrecognized operator type."); + case ast::ast_operator::ADD: return opcode::ADD; + case ast::ast_operator::SUB: return opcode::SUB; + case ast::ast_operator::MUL: return opcode::MUL; + case ast::ast_operator::DIV: return opcode::DIV; + case ast::ast_operator::TRUE_DIV: return opcode::TRUE_DIV; + case ast::ast_operator::FLOOR_DIV: return opcode::FLOOR_DIV; + case ast::ast_operator::MOD: return opcode::MOD; + case ast::ast_operator::PYMOD: return opcode::PYMOD; + case ast::ast_operator::POW: return opcode::POW; + case ast::ast_operator::EQUAL: return opcode::EQUAL; + case ast::ast_operator::NULL_EQUAL: return opcode::NULL_EQUAL; + case ast::ast_operator::NOT_EQUAL: return opcode::NOT_EQUAL; + case ast::ast_operator::LESS: return opcode::LESS; + case ast::ast_operator::GREATER: return opcode::GREATER; + case ast::ast_operator::LESS_EQUAL: return opcode::LESS_EQUAL; + case ast::ast_operator::GREATER_EQUAL: return opcode::GREATER_EQUAL; + case ast::ast_operator::BITWISE_AND: return opcode::BIT_AND; + case ast::ast_operator::BITWISE_OR: return opcode::BIT_OR; + case ast::ast_operator::BITWISE_XOR: return opcode::BIT_XOR; + case ast::ast_operator::LOGICAL_AND: return opcode::LOGICAL_AND; + case ast::ast_operator::NULL_LOGICAL_AND: return opcode::NULL_LOGICAL_AND; + case ast::ast_operator::LOGICAL_OR: return opcode::LOGICAL_OR; + case ast::ast_operator::NULL_LOGICAL_OR: return opcode::NULL_LOGICAL_OR; + case ast::ast_operator::IDENTITY: return opcode::IDENTITY; + case ast::ast_operator::IS_NULL: return opcode::IS_NULL; + case ast::ast_operator::SIN: return opcode::SIN; + case ast::ast_operator::COS: return opcode::COS; + case ast::ast_operator::TAN: return opcode::TAN; + case ast::ast_operator::ARCSIN: return opcode::ARCSIN; + case ast::ast_operator::ARCCOS: return opcode::ARCCOS; + case ast::ast_operator::ARCTAN: return opcode::ARCTAN; + case ast::ast_operator::SINH: return opcode::SINH; + case ast::ast_operator::COSH: return opcode::COSH; + case ast::ast_operator::TANH: return opcode::TANH; + case ast::ast_operator::ARCSINH: return opcode::ARCSINH; + case ast::ast_operator::ARCCOSH: return opcode::ARCCOSH; + case ast::ast_operator::ARCTANH: return opcode::ARCTANH; + case ast::ast_operator::EXP: return opcode::EXP; + case ast::ast_operator::LOG: return opcode::LOG; + case ast::ast_operator::SQRT: return opcode::SQRT; + case ast::ast_operator::CBRT: return opcode::CBRT; + case ast::ast_operator::CEIL: return opcode::CEIL; + case ast::ast_operator::FLOOR: return opcode::FLOOR; + case ast::ast_operator::ABS: return opcode::ABS; + case ast::ast_operator::RINT: return opcode::RINT; + case ast::ast_operator::BIT_INVERT: return opcode::BIT_INVERT; + case ast::ast_operator::NOT: return opcode::LOGICAL_NOT; + case ast::ast_operator::CAST_TO_INT64: return opcode::CAST_TO_I64; + case ast::ast_operator::CAST_TO_UINT64: return opcode::CAST_TO_U64; + case ast::ast_operator::CAST_TO_FLOAT64: return opcode::CAST_TO_F64; + default: CUDF_UNREACHABLE("Invalid opcode"); } } -bool operation::is_always_valid() +std::string to_cuda_type(cudf::data_type type, bool nullable) { - return is_operator_always_valid(op_) || - std::all_of( - operands_.begin(), operands_.end(), [](auto& op) { return op->is_always_valid(); }); + auto name = type_to_name(type); + return nullable ? std::format("cuda::std::optional<{}>", name) : name; } -opcode operation::get_opcode() const { return op_; } - -std::span const> operation::get_operands() const { return operands_; } - -void operation::instantiate(instance_context& ctx, instance_info const& info) +data_type get_return_type(opcode op, + std::span args, + std::optional target_scale) { - for (auto& arg : operands_) { - arg->instantiate(ctx, info); - } + std::vector arg_types; + std::vector arg_scales; - id_ = ctx.make_tmp_id(); - std::vector operand_types; + for (auto& type : args) { + arg_types.emplace_back(as_typing(type)); + arg_scales.emplace_back(type.scale()); + } - for (auto& arg : operands_) { - operand_types.emplace_back(arg->get_type()); + auto op_type_match = get_op_typing(op); + auto rescaled = op_rescale(op, arg_scales, target_scale); + + for (size_t i = 0; i < args.size(); ++i) { + auto required_type = op_type_match.args[i]; + auto arg_type = arg_types[i]; + + if ((required_type & type::ARG_MASK) != type::NONE) { + auto src_index = static_cast(required_type & ~type::ARG_MASK); + CUDF_EXPECTS( + src_index < i, + std::format( + "Invalid type match rule for operator `{}` at argument #{}", get_op_name(op), i), + std::runtime_error); + CUDF_EXPECTS(args[i].id() == args[src_index].id(), + std::format("Argument #{} of operator `{}` does not match type of argument " + "#{}. Got `{}`, expected `{}`", + i, + get_op_name(op), + src_index, + type_to_name(args[i]), + type_to_name(args[src_index]))); + } else { + CUDF_EXPECTS( + (arg_type & required_type) != 0, + std::format("Argument #{} of operator `{}` does not match expected types. Got {}", + i, + get_op_name(op), + type_to_name(args[i]))); + } } - type_ = ast::detail::ast_operator_return_type(op_, operand_types); + if ((op_type_match.output & type::ARG_MASK) != type::NONE) { + auto arg_index = static_cast(op_type_match.output & ~type::ARG_MASK); + auto type = args[arg_index].id(); + auto scale = numeric::scale_type{is_fixed_point(data_type{type}) ? rescaled : 0}; + return data_type{type, scale}; + } else { + CUDF_EXPECTS( + op_type_match.output != type::NONE, + std::format("Invalid type match rule for operator `{}` return type", get_op_name(op)), + std::runtime_error); + auto type = as_type_id(op_type_match.output); + auto scale = numeric::scale_type{is_fixed_point(data_type{type}) ? rescaled : 0}; + return data_type{type, scale}; + } } -std::string operation::generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) +void node::instantiate(instance_context& ctx) { - std::string operands_code; + id_ = ctx.make_tmp_id(); - for (auto& arg : operands_) { - operands_code = - std::format("{}{}{}", operands_code, arg->generate_code(ctx, info, instance), "\n"); + for (auto& arg : args_) { + arg->instantiate(ctx); } - auto operation_code = [&]() { - switch (info.id) { - case target::CUDA: { - auto first_operand = operands_[0]->get_id(); - auto operands_str = (operands_.size() == 1) - ? std::string{first_operand} - : std::accumulate(operands_.begin() + 1, - operands_.end(), - std::string{first_operand}, - [](auto const& a, auto& node) { - return std::format("{}, {}", a, node->get_id()); - }); - - auto cuda = std::format( - "{} {} = cudf::ast::detail::operator_functor{{}}({});", - cuda_type(type_, ctx.has_nulls()), - id_, - ast::detail::ast_operator_string(op_), - ctx.has_nulls(), - operands_str); - return cuda; + switch (op_) { + case opcode::GET_INPUT: { + type_ = ctx.get_input_vars()[std::get(reference_).index].type; + } break; + case opcode::SET_OUTPUT: { + type_ = args_[0]->get_type(); + } break; + default: { + std::vector arg_types; + for (auto& arg : args_) { + arg_types.emplace_back(arg->get_type()); } - default: - CUDF_FAIL("Unsupported target: " + std::to_string(static_cast(info.id)), - std::invalid_argument); - } - }(); - return operands_code + operation_code; -} + if (op_ == opcode::RESCALE) { + scale_reference_ = + input_reference{ctx.add_input(cudf::numeric_scalar{target_scale_.value_or(0)})}; + } -filter_predicate::filter_predicate(std::unique_ptr source) : id_(), source_(std::move(source)) -{ + type_ = get_return_type(op_, arg_types, target_scale_); + } break; + } } -std::string_view filter_predicate::get_id() { return id_; } - -data_type filter_predicate::get_type() { return data_type{type_id::BOOL8}; } +void node::emit_code(instance_context& instance, target_info const& info, code_sink& sink) const +{ + for (auto& arg : args_) { + arg->emit_code(instance, info, sink); + } -bool filter_predicate::is_null_aware() { return source_->is_null_aware(); } + switch (info.id) { + case target::CUDA: { + auto type = to_cuda_type(type_, instance.has_nulls()); -bool filter_predicate::is_always_valid() { return true; } + switch (op_) { + case opcode::GET_INPUT: { + sink.emit( + std::format(R"***({} {} = {}; +)***", + type, + id_, + instance.get_input_vars()[std::get(reference_).index].id)); + } break; + + case opcode::SET_OUTPUT: { + sink.emit(std::format( + R"***({} {} = {}; +*{} = {}; +)***", + type, + id_, + args_[0]->get_id(), + instance.get_output_vars()[std::get(reference_).index].id, + id_)); + } break; + + default: { + auto first_arg = std::format("&{}", args_[0]->get_id()); + auto args_str = (args_.size() == 1) + ? std::string{first_arg} + : std::accumulate(args_.begin() + 1, + args_.end(), + std::string{first_arg}, + [](auto const& a, auto& node) { + return std::format("{}, &{}", a, node->get_id()); + }); + + if (op_ == opcode::RESCALE) { + args_str = std::format( + "{}, &{}", args_str, instance.get_input_vars()[scale_reference_.index].id); + } -node& filter_predicate::get_source() { return *source_; } + bool fallible = get_op_is_fallible(op_); + auto op_name = get_op_name(op_); -void filter_predicate::instantiate(instance_context& ctx, instance_info const& info) -{ - source_->instantiate(ctx, info); - CUDF_EXPECTS(source_->get_type().id() == type_id::BOOL8, - "Filter predicate source must be boolean.", - std::invalid_argument); - id_ = ctx.make_tmp_id(); -} + if (!fallible) { + sink.emit(std::format( + R"***({} {}; +cudf::ops::{}(&{}, {}); +)***", + type, + id_, + op_name, + id_, + args_str)); + } else { + sink.emit(std::format( + R"***({} {}; +if(cudf::ops::errc e = cudf::ops::{}(&{}, {}); e != cudf::ops::errc::OK) {{ +return e; +}} +)***", + type, + id_, + op_name, + id_, + args_str)); + } + } break; + } + } break; -[[nodiscard]] std::string filter_predicate::generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) -{ - switch (info.id) { - case target::CUDA: { - auto source_code = source_->generate_code(ctx, info, instance); - return std::format( - "{}\n" - "bool {} = cudf::ast::detail::flatten_predicate({});\n", - source_code, - id_, - source_->get_id()); - } default: - CUDF_FAIL("Unsupported target: " + std::to_string(static_cast(info.id)), + CUDF_FAIL(std::format("Unsupported target: {}", static_cast(info.id)), std::invalid_argument); } } -std::span ast_converter::get_input_specs() const { return input_specs_; } - -int32_t ast_converter::add_ast_input(ast_input_spec in) -{ - auto id = static_cast(input_specs_.size()); - input_specs_.push_back(std::move(in)); - return id; -} - std::unique_ptr ast_converter::add_ir_node(ast::literal const& expr) { - auto index = add_ast_input( - ast_scalar_input_spec{expr.get_scalar(), - expr.get_value(), - make_column_from_scalar(expr.get_scalar(), 1, stream_, mr_)}); - return std::make_unique(index); + auto id = instance_.add_input(expr.get_scalar()); + return std::make_unique(input_reference{id}); } std::unique_ptr ast_converter::add_ir_node(ast::column_reference const& expr) { - auto index = - add_ast_input(ast_column_input_spec{expr.get_table_source(), expr.get_column_index()}); - return std::make_unique(index); -} + // resolve the table for a column input spec, preferring left_table/right_table for join cases, + // falling back to args.table for the single-table case. + auto resolve = [&](ast::table_reference ref) { + CUDF_EXPECTS(ref == ast::table_reference::LEFT || ref == ast::table_reference::RIGHT, + "Invalid table reference in column expression"); + return ref == ast::table_reference::LEFT ? left_table_ : right_table_; + }; -std::unique_ptr ast_converter::add_ir_node(ast::operation const& expr) -{ - std::vector> operands; - for (auto const& operand : expr.get_operands()) { - operands.push_back(operand.get().accept(*this)); - } - return std::make_unique(expr.get_operator(), std::move(operands)); + auto table = resolve(expr.get_table_source()); + auto id = instance_.add_input( + column_input{.column = table.column(expr.get_column_index()), + .table_source = (expr.get_table_source() == ast::table_reference::LEFT ? 0 : 1), + .column_index = static_cast(expr.get_column_index())}); + return std::make_unique(input_reference{id}); } -std::unique_ptr ast_converter::add_ir_node(ast::detail::filter_predicate const& expr) -{ - auto operand = expr.get_operand().accept(*this); - return std::make_unique(std::move(operand)); -} - -// Resolve the table for a column input spec, preferring left_table/right_table for join cases, -// falling back to args.table for the single-table case. -table_view const& resolve_table(ast_column_input_spec const& in, ast_args const& args) +std::unique_ptr ast_converter::add_ir_node(ast::operation const& expr) { - if (in.table == ast::table_reference::LEFT) { - return args.left_table.num_columns() > 0 ? args.left_table : args.table; + std::vector> args; + for (auto& operand : expr.get_operands()) { + args.emplace_back(operand.get().accept(*this)); } - return args.right_table; -} - -void ast_converter::add_input_var(ast_column_input_spec const& in, ast_args const& args) -{ - // TODO(lamarrr): consider mangling column name to make debugging easier - auto id = std::format("in_{}", input_vars_.size()); - auto type = resolve_table(in, args).column(in.column).type(); - input_vars_.emplace_back(std::move(id), type); + return std::make_unique( + as_opcode(expr.get_operator()), std::nullopt, std::move(args)); } -void ast_converter::add_input_var(ast_scalar_input_spec const& in, - [[maybe_unused]] ast_args const& args) +std::unique_ptr ast_converter::add_ir_node(ast::detail::predicate const& expr) { - auto id = std::format("in_{}", input_vars_.size()); - auto type = in.ref.get().type(); - input_vars_.emplace_back(std::move(id), type); + return std::make_unique( + row_ir::opcode::PREDICATE, std::nullopt, expr.get_operand().accept(*this)); } -void ast_converter::add_output_var() +std::unique_ptr ast_converter::add_ir_node(ast::jit::detail::operation const& expr) { - auto id = std::format("out_{}", output_vars_.size()); - output_vars_.emplace_back(std::move(id)); -} - -template -decltype(auto) dispatch_input_spec(ast_input_spec const& in, Fn&& fn, Args&&... args) -{ - if (std::holds_alternative(in)) { - return fn(std::get(in), std::forward(args)...); - } else if (std::holds_alternative(in)) { - return fn(std::get(in), std::forward(args)...); - } else { - CUDF_FAIL("Unsupported input type"); + std::vector> args; + for (auto& arg : expr.get_arguments()) { + args.emplace_back(arg.get().accept(*this)); } + return std::make_unique( + expr.get_opcode(), expr.get_target_scale(), std::move(args)); } -std::variant get_column_view(ast_column_input_spec const& spec, - ast_args const& args) -{ - return resolve_table(spec, args).column(spec.column); -} +bool is_nullable(scalar_input const& in) { return in.scalar_column->view().nullable(); } -std::variant get_column_view(ast_scalar_input_spec const& spec, - ast_args const& args) -{ - return scalar_column_view{spec.broadcast_column->view()}; -} +bool is_nullable(column_input const& in) { return in.column.nullable(); } -std::tuple ast_converter::generate_code(target target_id, - ast::expression const& expr, - ast_args const& args) +std::tuple ast_converter::generate_code( + target target_id, ast::expression const& expr, std::string_view function_name) { - auto output_expr_ir = expr.accept(*this); - output_irs_.emplace_back(std::make_unique(0, std::move(output_expr_ir))); + // add 1 auto-deduced output variable + [[maybe_unused]] auto output_id = instance_.add_output(); - // resolve the flattened input references into IR input variables - for (auto const& input : input_specs_) { - dispatch_input_spec(input, [this](auto&... args) { add_input_var(args...); }, args); - } + output_irs_.emplace_back(std::make_unique(output_reference{0}, expr.accept(*this))); bool has_nullable_inputs = - std::any_of(input_specs_.begin(), input_specs_.end(), [&](auto const& input) { - return dispatch_input_spec( - input, - [](auto&... args) { - auto col = get_column_view(args...); - return std::visit([](auto& view) { return view.nullable(); }, col); - }, - args); + std::any_of(instance_.inputs_.begin(), instance_.inputs_.end(), [&](auto& in) { + return std::visit([](auto& c) { return is_nullable(c); }, in); }); - // add 1 auto-deduced output variable - add_output_var(); - - instance_context instance_ctx; - instance_info instance{input_vars_, output_vars_}; - - auto is_null_aware = - std::any_of( - output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_null_aware(); }) - ? null_aware::YES - : null_aware::NO; + bool is_null_aware = std::any_of( + output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_null_aware(); }); bool output_is_always_valid = std::all_of( output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_always_valid(); }); - bool may_evaluate_null = !output_is_always_valid && has_nullable_inputs; + bool may_evaluate_null = !output_is_always_valid || has_nullable_inputs; + auto null_policy = may_evaluate_null ? output_nullability::PRESERVE : output_nullability::ALL_VALID; - instance_ctx.set_has_nulls(is_null_aware == null_aware::YES); + auto is_fallible = std::any_of( + output_irs_.cbegin(), output_irs_.cend(), [](auto& ir) { return ir->is_fallible(); }); + + instance_.set_has_nulls(is_null_aware); // instantiate the IR nodes for (auto& ir : output_irs_) { - ir->instantiate(instance_ctx, instance); + ir->instantiate(instance_); } target_info target{target_id}; - std::string body; + CUDF_EXPECTS( + target.id == target::CUDA, "Unsupported target for code generation", std::invalid_argument); - for (auto& ir : output_irs_) { - body = std::format("{}{}{}", body, ir->generate_code(instance_ctx, target, instance), "\n"); + auto output_decl = [&](auto i) { + auto& var = instance_.output_vars_[i]; + auto& ir = output_irs_[i]; + return std::format("{}* {}", to_cuda_type(ir->get_type(), instance_.has_nulls()), var.id); + }; + + auto input_decl = [&](auto i) { + auto& var = instance_.input_vars_[i]; + return std::format("{} {}", to_cuda_type(var.type, instance_.has_nulls()), var.id); + }; + + std::vector arg_decls; + + for (size_t i = 0; i < instance_.output_vars_.size(); ++i) { + arg_decls.emplace_back(output_decl(i)); } - switch (target.id) { - case target::CUDA: { - { - auto output_decl = [&](size_t i) { - auto const& var = output_vars_[i]; - auto const& ir = output_irs_[i]; - auto output_type = ir->get_type(); - return std::format("{}* {}", cuda_type(output_type, instance_ctx.has_nulls()), var.id); - }; - - auto input_decl = [&](size_t i) { - auto const& var = input_vars_[i]; - return std::format("{} {}", cuda_type(var.type, instance_ctx.has_nulls()), var.id); - }; - - std::vector params_decls; - - for (size_t i = 0; i < output_vars_.size(); ++i) { - params_decls.push_back(output_decl(i)); - } - - for (size_t i = 0; i < input_vars_.size(); ++i) { - params_decls.push_back(input_decl(i)); - } - - auto params_decl = [&] { - if (params_decls.empty()) { - return std::string{}; - } else if (params_decls.size() == 1) { - return params_decls[0]; - } else { - return std::accumulate( - params_decls.begin() + 1, - params_decls.end(), - params_decls[0], - [](auto const& a, auto const& b) { return std::format("{}, {}", a, b); }); - } - }(); - - code_ = std::format( - R"***( -__device__ void expression({}) -{{ -{} -return; -}} -)***", - params_decl, - body); + for (size_t i = 0; i < instance_.input_vars_.size(); ++i) { + arg_decls.emplace_back(input_decl(i)); + } - return {is_null_aware, null_policy}; - } - break; + auto args_decl = [&] { + if (arg_decls.empty()) { + return std::string{}; + } else if (arg_decls.size() == 1) { + return arg_decls[0]; + } else { + return std::accumulate( + arg_decls.begin() + 1, arg_decls.end(), arg_decls[0], [](auto const& a, auto const& b) { + return std::format("{}, {}", a, b); + }); } - default: - CUDF_FAIL("Unsupported target: " + std::to_string(static_cast(target.id)), - std::invalid_argument); + }(); + + code_sink sink; + sink.emit(std::format("__device__ cudf::ops::errc {}(", function_name)); + sink.emit(args_decl); + sink.emit(")\n{\n"); + for (auto& ir : output_irs_) { + ir->emit_code(instance_, target, sink); } + sink.emit("return cudf::ops::errc::OK;\n}"); + return { + sink.get_code(), is_null_aware ? null_aware::YES : null_aware::NO, null_policy, is_fallible}; +} + +std::variant get_column_view(scalar_input const& in) +{ + return scalar_column_view{in.scalar_column->view()}; +} + +std::variant get_column_view(column_input const& in) +{ + return column_view{in.column}; } // Due to the AST expression tree structure, we can't generate the IR without the target // tables transform_args ast_converter::compute_column(target target_id, ast::expression const& expr, - ast_args const& args, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - ast_converter converter{stream, mr}; + ast_converter converter{stream, mr, left_table, right_table}; // TODO(lamarrr): consider deduplicating ast expression's input column references. See // TransformTest/1.DeeplyNestedArithmeticLogicalExpression for reference - auto [is_null_aware, output_nullability] = converter.generate_code(target_id, expr, args); - + auto [code, is_null_aware, output_nullability, is_fallible] = + converter.generate_code(target_id, expr, function_name); std::vector> inputs; std::vector> scalar_columns; + std::vector> table_sources; + std::vector> column_indices; + + for (auto& input : converter.instance_.inputs_) { + if (std::holds_alternative(input)) { + auto& col = std::get(input); + table_sources.emplace_back(col.table_source); + column_indices.emplace_back(col.column_index); + } else { + table_sources.emplace_back(std::nullopt); + column_indices.emplace_back(std::nullopt); + } - for (auto& input : converter.input_specs_) { - auto column_view = - dispatch_input_spec(input, [](auto&... args) { return get_column_view(args...); }, args); - inputs.emplace_back(column_view); + auto view = std::visit([](auto& in) { return get_column_view(in); }, input); + inputs.emplace_back(view); - if (std::holds_alternative(input)) { - auto& scalar_input = std::get(input); - scalar_columns.push_back(std::move(scalar_input.broadcast_column)); + if (std::holds_alternative(input)) { + auto& scalar = std::get(input); + scalar_columns.emplace_back(std::move(scalar.scalar_column)); } } auto& out = converter.output_irs_[0]; auto output_column_type = out->get_type(); - - auto result = transform_args{.scalar_columns = std::move(scalar_columns), - .inputs = inputs, - .udf = std::move(converter.code_), - .output_type = output_column_type, - .source_type = cudf::udf_source_type::CUDA, - .user_data = std::nullopt, - .is_null_aware = is_null_aware, - .null_policy = output_nullability, - .row_size = args.table.num_rows(), - .input_specs = std::move(converter.input_specs_)}; - + auto output = transform_output{.type = output_column_type, .nullability = output_nullability}; + auto row_size = std::max({left_table.num_rows(), right_table.num_rows()}); + auto result = + transform_args{.scalar_columns = std::move(scalar_columns), + .input_table_sources = std::move(table_sources), + .input_column_indices = std::move(column_indices), + .udf = std::move(code), + .source_type = cudf::udf_source_type::CUDA, + .is_null_aware = is_null_aware, + .user_data = std::nullopt, + .inputs = inputs, + .outputs{output}, + .string_offsets{}, + .row_size = row_size, + .error_mode = is_fallible ? ops::error_mode::ANY_ROW : ops::error_mode::IGNORE}; if (get_context().dump_codegen()) { - std::cout << "Generated code for transform: " << result.udf << std::endl; + std::cout << "Generated code for transform: \n" << result.udf << std::endl; } return result; } -filter_args ast_converter::filter(target target_id, - ast::expression const& expr, - ast_args const& args, - table_view const& filter_table, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +transform_args ast_converter::filter(target target_id, + ast::expression const& expr, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { - auto filter = ast::detail::filter_predicate{expr}; - auto transform = compute_column(target_id, filter, args, stream, mr); + auto filter = ast::detail::predicate{expr}; + auto transform = + compute_column(target_id, filter, left_table, right_table, function_name, stream, mr); - CUDF_EXPECTS(transform.output_type.id() == type_id::BOOL8, + CUDF_EXPECTS(transform.outputs.size() == 1, + "Filter expression must have exactly one output column."); + CUDF_EXPECTS(transform.outputs[0].type.id() == type_id::BOOL8, "Filter expression must return a boolean type.", std::invalid_argument); - std::vector filter_columns; - std::transform(filter_table.begin(), - filter_table.end(), - std::back_inserter(filter_columns), - [](auto const& col) { return col; }); - - auto result = filter_args{.scalar_columns = std::move(transform.scalar_columns), - .inputs = std::move(transform.inputs), - .filter_columns = std::move(filter_columns), - .udf = std::move(transform.udf), - .source_type = transform.source_type, - .user_data = transform.user_data, - .is_null_aware = transform.is_null_aware, - .predicate_nullability = transform.null_policy, - .input_specs = std::move(transform.input_specs)}; - - return result; + return transform; } -} // namespace row_ir -} // namespace detail -} // namespace cudf +} // namespace cudf::detail::row_ir diff --git a/cpp/src/jit/row_ir.hpp b/cpp/src/jit/row_ir.hpp index 5e3cd2633e68..1686faed997b 100644 --- a/cpp/src/jit/row_ir.hpp +++ b/cpp/src/jit/row_ir.hpp @@ -6,7 +6,12 @@ #pragma once #include #include +#include +#include #include +#include +#include +#include #include #include #include @@ -15,10 +20,7 @@ #include #include -#include -#include #include -#include #include #include #include @@ -60,18 +62,41 @@ struct untyped_var_info { }; /** - * @brief The information needed to instantiate the IR nodes + * @brief The information about the target for which the IR is generated. */ -struct instance_info { - std::span inputs; ///< The input variables - std::span outputs; ///< The output variables +struct target_info { + target id = target::CUDA; ///< The target identifier +}; + +struct scalar_input { + std::unique_ptr scalar_column = + nullptr; ///< The scalar value represented as a column with a single element +}; + +struct column_input { + column_view column = {}; ///< The column input + std::optional table_source = std::nullopt; + std::optional column_index = std::nullopt; }; +using input = std::variant; + /** - * @brief The information about the target for which the IR is generated. + * @brief The arguments needed to invoke a `cudf::transform` */ -struct target_info { - target id = target::CUDA; ///< The target identifier +struct [[nodiscard]] transform_args { + std::vector> scalar_columns = {}; + std::vector> input_table_sources = {}; + std::vector> input_column_indices = {}; + std::string udf = {}; + udf_source_type source_type = cudf::udf_source_type::CUDA; + null_aware is_null_aware = null_aware::NO; + std::optional user_data = std::nullopt; + std::vector inputs = {}; + std::vector outputs = {}; + std::vector> string_offsets = {}; + std::optional row_size = std::nullopt; + ops::error_mode error_mode = ops::error_mode::IGNORE; }; /** @@ -81,12 +106,25 @@ struct target_info { */ struct [[nodiscard]] instance_context { private: - int32_t num_tmp_vars_ = 0; ///< The number of temporary variables generated - std::string tmp_prefix_ = "tmp_"; ///< The prefix for temporary variable identifiers - bool has_nulls_ = false; ///< If expressions involve null values + int32_t num_tmp_vars_ = 0; ///< The number of temporary variables generated + std::string tmp_prefix_ = "tmp_"; ///< The prefix for temporary variable identifiers + bool has_nulls_ = false; ///< If expressions involve null values + std::vector inputs_; ///< The inputs for the IR + std::vector input_vars_; ///< The input variables for the IR + std::vector output_vars_; ///< The output variables for the IR + rmm::cuda_stream_view + stream_; ///< The CUDA stream for any device operations during IR generation + rmm::device_async_resource_ref + mr_; ///< The device memory resource for any device memory allocation during IR generation public: - instance_context() = default; ///< Default constructor + friend struct ast_converter; + friend struct node; + + instance_context(rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) + : stream_(stream), mr_(mr) + { + } instance_context(instance_context const&) = delete; @@ -98,6 +136,21 @@ struct [[nodiscard]] instance_context { ~instance_context() = default; ///< Destructor + [[nodiscard]] int32_t add_output(); + + [[nodiscard]] int32_t add_input(input in); + + [[nodiscard]] int32_t add_input(scalar const& scalar) + { + return add_input( + scalar_input{.scalar_column = make_column_from_scalar(scalar, 1, stream_, mr_)}); + } + + [[nodiscard]] int32_t add_input(column_view const& column) + { + return add_input(column_input{.column = column}); + } + /** * @brief Generate a globally unique temporary variable identifier * @return A unique temporary variable identifier @@ -114,406 +167,206 @@ struct [[nodiscard]] instance_context { * @param has_nulls True if expressions involve null values */ void set_has_nulls(bool has_nulls); -}; - -struct [[nodiscard]] node { - /** - * @brief Get the identifier of the IR node - * @return The identifier of the IR node - */ - virtual std::string_view get_id() = 0; - - /** - * @brief Get the type info of the IR node - * @return The type information of the IR node - */ - [[nodiscard]] virtual data_type get_type() = 0; - - /** - * @brief Returns `false` if this node forwards nulls from its inputs to its output. - * e.g., `ADD` operator is not null-aware because if any of its inputs is null, the output is - * null. but `NULL_EQUAL` operator is null-aware because it can produce a non-null output even if - * its inputs are null. - */ - [[nodiscard]] virtual bool is_null_aware() = 0; /** - * @brief Returns `true` if this node always produces a valid output even if its inputs are - * nullable, e.g., `IS_NULL` operator produces a valid boolean output regardless of the - * nullability of its input. + * @brief Get the input values for the IR + * @return A span of input values for the IR */ - [[nodiscard]] virtual bool is_always_valid() = 0; + [[nodiscard]] std::span get_inputs() const; /** - * @brief Instantiate the IR node with the given context and instance information, setting up any - * necessary state and preprocessing needed for code generation. - * @param ctx The context within which the IR is instantiated - * @param info The instance information + * @brief Get the input variables for the IR + * @return A span of input variable information */ - virtual void instantiate(instance_context& ctx, instance_info const& info) = 0; + [[nodiscard]] std::span get_input_vars() const; /** - * @brief Generate the code for the IR node based on the instance context and target information. - * @param ctx The context within which the IR is instantiated - * @param info The target information - * @param instance The instance information - * @return The generated code for the IR node + * @brief Get the output variables for the IR + * @return A span of output variable information */ - [[nodiscard]] virtual std::string generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) = 0; - - virtual ~node() = default; + [[nodiscard]] std::span get_output_vars() const; }; -/** - * @brief The operation code used in the IR nodes. - */ -using opcode = ast::ast_operator; - -/** - * @brief An IR node that retrieves an input variable by its index. - * This node is used to access input variables in the IR. - */ -struct [[nodiscard]] get_input final : node { +struct [[nodiscard]] code_sink { private: - std::string id_; ///< The identifier of the IR node - int32_t input_; ///< The index of the input variable - data_type type_; ///< The type information of the IR node + std::string code_; public: - /** - * @brief Construct a new get_input IR node - * @param input The index of the input variable - */ - get_input(int32_t input); - - get_input(get_input const&) = delete; - - get_input& operator=(get_input const&) = delete; + void emit(std::string_view code) { code_ += code; } - get_input(get_input&&) = default; ///< Move constructor - - get_input& operator=(get_input&&) = default; ///< Move assignment operator - - ~get_input() override = default; ///< Destructor - - /** - * @copydoc node::get_id - */ - [[nodiscard]] std::string_view get_id() override; - - /** - * @copydoc node::get_type - */ - [[nodiscard]] data_type get_type() override; - - /** - * @copydoc node::is_null_aware - */ - [[nodiscard]] bool is_null_aware() override; - - /** - * @copydoc node::is_always_valid - */ - [[nodiscard]] bool is_always_valid() override; + [[nodiscard]] std::string const& get_code() const { return code_; } +}; - /** - * @copydoc node::instantiate - */ - void instantiate(instance_context& ctx, instance_info const& info) override; +struct [[nodiscard]] input_reference { + int32_t index = 0; ///< The index of the input variable +}; - /** - * @copydoc node::generate_code - */ - [[nodiscard]] std::string generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) override; +struct [[nodiscard]] output_reference { + int32_t index = 0; ///< The index of the output variable }; -/** - * @brief An IR node that sets the output variable to the value of a source IR node. - */ -struct [[nodiscard]] set_output final : node { +struct [[nodiscard]] node { private: - std::string id_; ///< The identifier of the IR node - int32_t output_; ///< The index of the output variable - std::unique_ptr source_; ///< The source IR node from which the value is taken - data_type type_; ///< The type information of the IR node - std::string output_id_; ///< The identifier of the output variable - - public: - /** - * @brief Construct a new set_output IR node - * @param output The index of the output variable - * @param source The source IR node from which the value is taken - */ - set_output(int32_t output, std::unique_ptr source); - - set_output(set_output const&) = delete; - - set_output& operator=(set_output const&) = delete; - - set_output(set_output&&) = default; ///< Move constructor - - set_output& operator=(set_output&&) = default; ///< Move assignment operator - - ~set_output() override = default; ///< Destructor - - /** - * @copydoc node::get_id - */ - [[nodiscard]] std::string_view get_id() override; - - /** - * @copydoc node::get_type - */ - [[nodiscard]] data_type get_type() override; + std::variant reference_ = + std::monostate{}; ///< The index of the input/output variable + opcode op_ = opcode::GET_INPUT; ///< The operation code + std::optional target_scale_ = std::nullopt; ///< The target scale for decimal + std::vector> args_ = {}; ///< The arguments of the operation - /** - * @copydoc node::is_null_aware - */ - [[nodiscard]] bool is_null_aware() override; + data_type type_ = {}; ///< The resolved type information of the IR node - /** - * @copydoc node::is_always_valid - */ - [[nodiscard]] bool is_always_valid() override; - - /** - * @brief Get the source IR node from which the value is taken - */ - [[nodiscard]] node& get_source(); + std::string id_ = {}; ///< The identifier of the IR node + input_reference + scale_reference_; ///< The index of the scale variable for decimal rescaling if applicable /** - * @copydoc node::instantiate + * @brief Create a set of argument IR nodes */ - void instantiate(instance_context& ctx, instance_info const& info) override; + template + requires(std::is_same_v && ...) + static std::vector> arguments(T... args) + { + std::vector> result; + (result.emplace_back(std::make_unique(std::move(args))), ...); + return result; + } /** - * @copydoc node::generate_code + * @brief Create a set of argument IR nodes */ - [[nodiscard]] std::string generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) override; -}; - -/** - * @brief An IR node that represents an operation with zero or more operands. - */ -struct [[nodiscard]] operation final : node { - private: - std::string id_; ///< The identifier of the IR node - opcode op_; ///< The operation code - std::vector> operands_; ///< The operands of the operation - data_type type_; ///< The type information of the IR node - - operation(opcode op, std::unique_ptr* move_begin, std::unique_ptr* move_end); + template + requires(std::is_same_v, T> && ...) + static std::vector> arguments(T... args) + { + std::vector> result; + (result.emplace_back(std::move(args)), ...); + return result; + } public: /** - * @brief Create a set of operand IR nodes + * @brief Construct a new operation IR node + * @param op The operation code + * @param args The arguments of the operation */ - template - requires(std::is_base_of_v && ...) - static std::array, sizeof...(T)> operands(T&&... args) - { - return {std::make_unique(std::forward(args))...}; - } + node(opcode op, std::optional target_scale, std::vector> args); /** - * @brief Create a set of operand IR nodes from existing unique pointers + * @brief Construct a new operation IR node + * @param op The operation code + * @param args The arguments of the operation */ template - requires(std::is_base_of_v && ...) - static std::array, sizeof...(T)> operands(std::unique_ptr&&... args) + requires(std::is_same_v && ...) + node(opcode op, std::optional target_scale, T... args) + : node(op, target_scale, arguments(std::move(args)...)) { - return {std::move(args)...}; } /** * @brief Construct a new operation IR node * @param op The operation code - * @param operands The operands of the operation + * @param args The arguments of the operation */ - operation(opcode op, std::vector> operands); - - template - operation(opcode op, std::array, N> operands) - : operation{op, operands.data(), operands.data() + N} + template + requires(std::is_same_v && ...) + node(opcode op, std::optional target_scale, std::unique_ptr... args) + : node(op, target_scale, arguments(std::move(args)...)) { } - operation(operation const&) = delete; - - operation& operator=(operation const&) = delete; - - operation(operation&&) = default; ///< Move constructor - - operation& operator=(operation&&) = default; ///< Move assignment operator - - ~operation() override = default; ///< Destructor - /** - * @copydoc node::get_id + * @brief Construct a new input reference IR node + * @param input The index of the input variable */ - [[nodiscard]] std::string_view get_id() override; + node(input_reference input); /** - * @copydoc node::get_type + * @brief Construct a new output reference IR node + * @param output The index of the output variable + * @param arg The argument node that produces the value to be set to the output variable */ - [[nodiscard]] data_type get_type() override; + node(output_reference reference, std::unique_ptr arg); /** - * @copydoc node::is_null_aware + * @brief Construct a new output reference IR node + * @param output The index of the output variable + * @param arg The argument node that produces the value to be set to the output variable */ - [[nodiscard]] bool is_null_aware() override; + node(output_reference reference, node arg); - /** - * @copydoc node::is_always_valid - */ - [[nodiscard]] bool is_always_valid() override; + node(node const& other) = delete; + node(node&& other) = default; ///< Move constructor + node& operator=(node const& other) = delete; + node& operator=(node&& other) = default; ///< Move assignment operator + ~node() = default; ///< Destructor /** - * @brief Get the operation code of the operation - * @return The operation code of the operation - */ - [[nodiscard]] opcode get_opcode() const; - - /** @brief Get the operands of the operation - * @return A span of unique pointers to the operands of the operation + * @brief Get the identifier of the IR node + * @return The identifier of the IR node */ - [[nodiscard]] std::span const> get_operands() const; + [[nodiscard]] std::string_view get_id() const; /** - * @copydoc node::instantiate + * @brief Get the type info of the IR node + * @return The type information of the IR node */ - void instantiate(instance_context& ctx, instance_info const& info) override; + [[nodiscard]] data_type get_type() const; /** - * @copydoc node::generate_code + * @brief Get the target scale for decimal rescaling if applicable + * @return The target scale for decimal rescaling if applicable, std::nullopt otherwise */ - [[nodiscard]] std::string generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) override; -}; - -/** - * @brief An IR node that flattens a boolean predicate to be used in a filter operation. - * This node replaces null values with false. - */ -struct [[nodiscard]] filter_predicate final : node { - private: - std::string id_; ///< The identifier of the IR node - std::unique_ptr source_; ///< The source IR node from which the predicate value is taken - - public: - filter_predicate(std::unique_ptr source); + [[nodiscard]] std::optional get_target_scale() const; /** - * @copydoc node::get_id + * @brief Get the operation code of the operation + * @return The operation code of the operation */ - [[nodiscard]] std::string_view get_id() override; + opcode get_opcode() const; - /** - * @copydoc node::get_type + /** @brief Get the arguments of the operation + * @return A span of unique pointers to the arguments of the operation */ - [[nodiscard]] data_type get_type() override; + [[nodiscard]] std::span const> get_args() const; /** - * @copydoc node::is_null_aware + * @brief Returns `false` if this node forwards nulls from its inputs to its output. + * e.g., `ADD` operator is not null-aware because if any of its inputs is null, the output is + * null. but `NULL_EQUAL` operator is null-aware because it can produce a non-null output even if + * its inputs are null. */ - [[nodiscard]] bool is_null_aware() override; + [[nodiscard]] bool is_null_aware() const; /** - * @copydoc node::is_always_valid + * @brief Returns `true` if this node always produces a valid output even if its inputs are + * nullable, e.g., `IS_NULL` operator produces a valid boolean output regardless of the + * nullability of its input. */ - [[nodiscard]] bool is_always_valid() override; + [[nodiscard]] bool is_always_valid() const; /** - * @brief Get the source IR node from which the value is taken + * @brief Get if the IR node can raise an error during evaluation. + * @return `true` if the IR node can raise an error during evaluation, `false` otherwise */ - [[nodiscard]] node& get_source(); + [[nodiscard]] bool is_fallible() const; /** - * @copydoc node::instantiate + * @brief Instantiate the IR node with the given context and instance information, setting up any + * necessary state and preprocessing needed for code generation. + * @param ctx The context within which the IR is instantiated + * @param info The instance information */ - void instantiate(instance_context& ctx, instance_info const& info) override; + void instantiate(instance_context& ctx); /** - * @copydoc node::generate_code + * @brief Generate the code for the IR node based on the instance context and target information. + * @param ctx The context within which the IR is instantiated + * @param info The target information + * @param instance The instance information + * @param sink The code sink to which the generated code is emitted */ - [[nodiscard]] std::string generate_code(instance_context& ctx, - target_info const& info, - instance_info const& instance) override; -}; - -/** - * @brief A specification of an input column to the AST - */ -struct ast_column_input_spec { - ast::table_reference table = {}; ///< The table reference (LEFT or RIGHT) - int32_t column = 0; ///< The column index in the referenced table -}; - -/** - * @brief A specification of an input scalar to the AST - */ -struct ast_scalar_input_spec { - std::reference_wrapper ref; ///< The scalar value - ast::generic_scalar_device_view view; ///< The device view of the scalar value - std::unique_ptr broadcast_column = - nullptr; ///< The broadcasted column, a column of size 1 -}; - -/** - * @brief An input specification for the AST - */ -using ast_input_spec = std::variant; - -/** - * @brief The arguments needed to invoke a `cudf::transform` - */ -struct [[nodiscard]] transform_args { - std::vector> scalar_columns = - {}; ///< The scalar columns created during the expression conversion - std::vector> inputs = - {}; ///< The input columns to the transform UDF - std::string udf = {}; ///< The user-defined function to apply - data_type output_type = data_type{type_id::EMPTY}; ///< The output type of the transform - cudf::udf_source_type source_type = cudf::udf_source_type::CUDA; ///< The source type of the UDF - std::optional user_data = std::nullopt; ///< User data to pass to the transform - null_aware is_null_aware = null_aware::NO; ///< Whether the transform is null-aware - output_nullability null_policy = output_nullability::PRESERVE; ///< Null-transformation policy - std::optional row_size = std::nullopt; ///< The row size of the transform operation - std::vector input_specs = {}; ///< The input specs (table ref + column index) -}; - -/** - * @brief The arguments needed to invoke a `cudf::filter` - */ -struct [[nodiscard]] filter_args { - std::vector> scalar_columns = - {}; ///< The scalar columns created during the expression conversion - std::vector> inputs = - {}; ///< The input columns to the transform UDF - std::vector filter_columns = {}; ///< The input columns to the filter - std::string udf = {}; ///< The user-defined function to apply as a predicate - cudf::udf_source_type source_type = cudf::udf_source_type::CUDA; ///< The source type of the UDF - std::optional user_data = std::nullopt; ///< User data to pass to the filter - null_aware is_null_aware = null_aware::NO; ///< Whether the filter is null-aware - output_nullability predicate_nullability = - output_nullability::PRESERVE; ///< Null-transformation policy for the predicate output - std::vector input_specs = {}; ///< The input specs (table ref + column index) -}; - -/** - * @brief The AST input column arguments used to resolve the column expressions - */ -struct ast_args { - table_view table = {}; ///< The table view containing the columns (single-table case) - table_view left_table = {}; ///< The left table for join predicates - table_view right_table = {}; ///< The right table for join predicates + void emit_code(instance_context& ctx, target_info const& info, code_sink& sink) const; }; /** @@ -521,15 +374,14 @@ struct ast_args { */ struct [[nodiscard]] ast_converter { private: - std::vector input_specs_; ///< The input specs for the AST - std::vector input_vars_; ///< The input variables for the IR - std::vector output_vars_; ///< The output variables for the IR - std::vector> output_irs_; ///< The output IR nodes - std::string code_; ///< The generated code for the IR + std::vector> output_irs_; ///< The output IR nodes rmm::cuda_stream_view stream_; ///< CUDA stream used for device memory operations and kernel launches. rmm::device_async_resource_ref mr_; ///< Device memory resource used to allocate the returned table's device memory + instance_context instance_; ///< The instance context used during the IR generation + table_view left_table_; ///< The left input table for the expression + table_view right_table_; ///< The right input table for the expression public: /** @@ -537,8 +389,15 @@ struct [[nodiscard]] ast_converter { * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory */ - ast_converter(rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) - : stream_(std::move(stream)), mr_(std::move(mr)) + ast_converter(rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr, + table_view left_table, + table_view right_table) + : stream_(std::move(stream)), + mr_(std::move(mr)), + instance_(stream_, mr_), + left_table_(std::move(left_table)), + right_table_(std::move(right_table)) { } @@ -551,51 +410,37 @@ struct [[nodiscard]] ast_converter { ~ast_converter() = default; ///< Destructor - private: - friend class ast::literal; - friend class ast::column_reference; - friend class ast::operation; - friend class ast::column_name_reference; - friend class ast::detail::filter_predicate; - + public: [[nodiscard]] std::unique_ptr add_ir_node(ast::literal const& expr); [[nodiscard]] std::unique_ptr add_ir_node(ast::column_reference const& expr); [[nodiscard]] std::unique_ptr add_ir_node(ast::operation const& expr); - [[nodiscard]] std::unique_ptr add_ir_node( - ast::detail::filter_predicate const& expr); + [[nodiscard]] std::unique_ptr add_ir_node(ast::detail::predicate const& expr); - [[nodiscard]] std::span get_input_specs() const; + [[nodiscard]] std::unique_ptr add_ir_node(ast::jit::detail::operation const& expr); - /** - * @brief add an AST input/input_reference and return its reference index - */ - [[nodiscard]] int32_t add_ast_input(ast_input_spec in); - - void add_input_var(ast_column_input_spec const& in, ast_args const& args); - - void add_input_var(ast_scalar_input_spec const& in, ast_args const& args); + [[nodiscard]] std::tuple generate_code( + target target, ast::expression const& expr, std::string_view function_name); - void add_output_var(); - - [[nodiscard]] std::tuple generate_code( - target target, ast::expression const& expr, ast_args const& args); - - public: /** * @brief Convert an AST `compute_column` expression to a `cudf::transform` * @param target The target for which the IR is generated * @param expr The AST expression to convert - * @param args The arguments needed to resolve the AST expression + * @param left_table The left input table for the expression + * @param right_table The right input table for the expression + * @param table The input table for the expression + * @param function_name The name of the generated function * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory * @return The result of the conversion, containing the transform arguments and scalar columns */ static transform_args compute_column(target target, ast::expression const& expr, - ast_args const& args, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); @@ -603,18 +448,21 @@ struct [[nodiscard]] ast_converter { * @brief Convert an AST `filter` expression to a `cudf::filter` * @param target The target for which the IR is generated * @param expr The AST expression to convert - * @param args The arguments needed to resolve the AST expression - * @param filter_table The table to be filtered + * @param left_table The left input table for the expression + * @param right_table The right input table for the expression + * @param table The input table for the expression + * @param function_name The name of the generated function * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table's device memory * @return The result of the conversion, containing the filter arguments and scalar columns */ - static filter_args filter(target target, - ast::expression const& expr, - ast_args const& args, - table_view const& filter_table, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr); + static transform_args filter(target target, + ast::expression const& expr, + table_view const& left_table, + table_view const& right_table, + std::string_view function_name, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); }; } // namespace row_ir diff --git a/cpp/src/join/filter_join_indices_jit.cu b/cpp/src/join/filter_join_indices_jit.cu index d7a2565df341..d841f46fb5bd 100644 --- a/cpp/src/join/filter_join_indices_jit.cu +++ b/cpp/src/join/filter_join_indices_jit.cu @@ -48,33 +48,35 @@ namespace detail { namespace { -// Build template parameters for JIT kernel -jitify2::StringVec build_join_filter_template_params(std::vector const& left_columns, - std::vector const& right_columns, - bool has_user_data, - null_aware is_null_aware) +jitify2::StringVec build_join_filter_template_params( + std::span inputs, + std::span> table_sources, + null_aware is_null_aware) { jitify2::StringVec template_params; - - template_params.emplace_back(jitify2::reflection::reflect(has_user_data)); + template_params.emplace_back(jitify2::reflection::reflect(false)); // has_user_data = false template_params.emplace_back(jitify2::reflection::reflect(is_null_aware)); - // Add left column accessors - for (std::size_t i = 0; i < left_columns.size(); ++i) { - auto const& col = left_columns[i]; - std::string type_name = cudf::type_to_name(col.type()); - template_params.emplace_back( - jitify2::reflection::Template("cudf::jit::join_column_accessor") - .instantiate(type_name, std::to_string(i), "cudf::jit::join_side::LEFT")); - } - - // Add right column accessors - for (std::size_t i = 0; i < right_columns.size(); ++i) { - auto const& col = right_columns[i]; - std::string type_name = cudf::type_to_name(col.type()); - template_params.emplace_back( - jitify2::reflection::Template("cudf::jit::join_column_accessor") - .instantiate(type_name, std::to_string(i), "cudf::jit::join_side::RIGHT")); + for (size_t i = 0; i < inputs.size(); ++i) { + auto const& input = inputs[i]; + if (auto* col = std::get_if(&input)) { + auto element = cudf::type_to_name(col->type()); + template_params.emplace_back( + jitify2::reflection::Template("cudf::jit::column_accessor") + .instantiate( + i, "cudf::column_device_view_core", element, false, table_sources[i].value())); + } else { + auto& scalar = std::get(input); + auto element = cudf::type_to_name(scalar.as_column_view().type()); + template_params.emplace_back( + jitify2::reflection::Template("cudf::jit::column_accessor") + .instantiate(i, + "cudf::column_device_view_core", + element, + true, + 0 // scalars dont belong to a table, so just use 0 as placeholder + )); + } } return template_params; @@ -82,8 +84,8 @@ jitify2::StringVec build_join_filter_template_params(std::vector co // Build the JIT kernel for join filtering jitify2::ConfiguredKernel build_join_filter_kernel(std::string const& predicate_code, - std::vector const& left_columns, - std::vector const& right_columns, + std::span inputs, + std::span> table_sources, bool is_ptx, bool has_user_data, null_aware is_null_aware, @@ -92,24 +94,28 @@ jitify2::ConfiguredKernel build_join_filter_kernel(std::string const& predicate_ { CUDF_FUNC_RANGE(); + std::vector ptx_output_types{"bool"}; + std::vector ptx_input_types; + + for (auto const& input : inputs) { + if (auto* col = std::get_if(&input)) { + ptx_input_types.push_back(cudf::type_to_name(col->type())); + } else { + auto& scalar = std::get(input); + ptx_input_types.push_back(cudf::type_to_name(scalar.type())); + } + } + // Parse predicate code auto const cuda_source = is_ptx ? cudf::jit::parse_single_function_ptx( predicate_code, "GENERIC_JOIN_FILTER_OP", - [&] { - std::vector left_types, right_types; - for (auto const& col : left_columns) - left_types.push_back(cudf::type_to_name(col.type())); - for (auto const& col : right_columns) - right_types.push_back(cudf::type_to_name(col.type())); - return cudf::jit::build_ptx_params(left_types, right_types, has_user_data); - }()) + cudf::jit::build_ptx_params(ptx_output_types, ptx_input_types, has_user_data)) : cudf::jit::parse_single_function_cuda(predicate_code, "GENERIC_JOIN_FILTER_OP"); // Build template parameters and kernel name - auto template_args = - build_join_filter_template_params(left_columns, right_columns, has_user_data, is_null_aware); + auto template_args = build_join_filter_template_params(inputs, table_sources, is_null_aware); auto kernel_name = jitify2::reflection::Template("cudf::join::jit::filter_join_kernel").instantiate(template_args); @@ -122,44 +128,45 @@ jitify2::ConfiguredKernel build_join_filter_kernel(std::string const& predicate_ // Launch the JIT kernel for join filtering void launch_join_filter_kernel(jitify2::ConfiguredKernel& kernel, - cudf::table_view const& left, - cudf::table_view const& right, cudf::device_span left_indices, cudf::device_span right_indices, + std::span inputs, bool* predicate_results, std::optional user_data, - std::vector const& extra_left_cols, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); // Create device views of tables - std::vector left_cols(left.begin(), left.end()); - left_cols.insert(left_cols.end(), extra_left_cols.begin(), extra_left_cols.end()); - std::vector right_cols(right.begin(), right.end()); + std::vector column_views; + for (auto const& input : inputs) { + if (auto* col = std::get_if(&input)) { + column_views.push_back(*col); + } else { + auto& scalar = std::get(input); + column_views.push_back(scalar.as_column_view()); + } + } - auto [left_handles, left_device_views] = - cudf::jit::column_views_to_device(left_cols, stream, mr); - auto [right_handles, right_device_views] = - cudf::jit::column_views_to_device(right_cols, stream, mr); + auto [handles, device_views] = + cudf::jit::column_views_to_device(column_views, stream, mr); // Set up kernel parameters - use JIT-compatible span type - cudf::jit::device_span left_span{left_indices.data(), left_indices.size()}; - cudf::jit::device_span right_span{right_indices.data(), - right_indices.size()}; - cudf::column_device_view_core const* left_tables_ptr = left_device_views.data(); - cudf::column_device_view_core const* right_tables_ptr = right_device_views.data(); - void* user_data_ptr = user_data.value_or(nullptr); - - std::array args{&left_span, - &right_span, - &left_tables_ptr, - &right_tables_ptr, - &predicate_results, - &user_data_ptr}; - - kernel->launch_raw(args.data()); + cudf::size_type num_rows = left_indices.size(); + cudf::size_type const* left_indices_ptr = left_indices.data(); + cudf::size_type const* right_indices_ptr = right_indices.data(); + cudf::column_device_view_core const* columns_ptr = device_views.data(); + void* user_data_ptr = user_data.value_or(nullptr); + + void* args[]{&num_rows, + &left_indices_ptr, + &right_indices_ptr, + &columns_ptr, + &predicate_results, + &user_data_ptr}; + + kernel->launch_raw(args); } // Same join semantics handling as the AST version @@ -348,43 +355,6 @@ apply_join_semantics(cudf::table_view const& left, } } -// Build template parameters from AST input specs (preserves expression input order) -jitify2::StringVec build_join_filter_template_params_from_specs( - std::vector const& input_specs, - cudf::table_view const& left, - cudf::table_view const& right, - null_aware is_null_aware) -{ - jitify2::StringVec template_params; - template_params.emplace_back(jitify2::reflection::reflect(false)); // has_user_data = false - template_params.emplace_back(jitify2::reflection::reflect(is_null_aware)); - - // Scalar columns are appended to the left table's device views, - // starting at index left.num_columns(). - auto scalar_index = left.num_columns(); - - for (auto const& spec : input_specs) { - if (std::holds_alternative(spec)) { - auto const& col_spec = std::get(spec); - auto const& table = col_spec.table == ast::table_reference::LEFT ? left : right; - auto const side_str = col_spec.table == ast::table_reference::LEFT - ? "cudf::jit::join_side::LEFT" - : "cudf::jit::join_side::RIGHT"; - auto type_name = cudf::type_to_name(table.column(col_spec.column).type()); - template_params.emplace_back( - jitify2::reflection::Template("cudf::jit::join_column_accessor") - .instantiate(type_name, std::to_string(col_spec.column), side_str)); - } else if (std::holds_alternative(spec)) { - auto const& scalar_spec = std::get(spec); - auto type_name = cudf::type_to_name(scalar_spec.ref.get().type()); - template_params.emplace_back(jitify2::reflection::Template("cudf::jit::join_scalar_accessor") - .instantiate(type_name, std::to_string(scalar_index++))); - } - } - - return template_params; -} - void validate_column_types(cudf::table_view const& table, char const* side) { for (auto const& col : table) { @@ -433,12 +403,20 @@ filter_join_indices_jit(cudf::table_view const& left, if (left_indices.empty()) { return make_empty_result(); } // Compile JIT kernel - std::vector left_cols(left.begin(), left.end()); - std::vector right_cols(right.begin(), right.end()); + std::vector inputs; + std::vector> table_sources; + for (auto const& col : left) { + inputs.emplace_back(col); + table_sources.emplace_back(0); + } + for (auto const& col : right) { + inputs.emplace_back(col); + table_sources.emplace_back(1); + } auto kernel = build_join_filter_kernel(predicate_code, - left_cols, - right_cols, + inputs, + table_sources, is_ptx, false, // has_user_data = false for now null_aware::NO, @@ -450,13 +428,11 @@ filter_join_indices_jit(cudf::table_view const& left, // Launch kernel launch_join_filter_kernel(kernel, - left, - right, left_indices, right_indices, + inputs, predicate_results.data(), std::nullopt, // no user data for now - {}, stream, mr); @@ -496,13 +472,11 @@ filter_join_indices_jit(cudf::table_view const& left, } // Convert AST predicate to JIT code - row_ir::ast_args ast_args{.left_table = left, .right_table = right}; auto filter_result = row_ir::ast_converter::filter( - row_ir::target::CUDA, predicate, ast_args, table_view{}, stream, mr); + row_ir::target::CUDA, predicate, left, right, "filter_operation", stream, mr); - // Build template params matching the AST input order - auto template_args = build_join_filter_template_params_from_specs( - filter_result.input_specs, left, right, filter_result.is_null_aware); + auto template_args = build_join_filter_template_params( + filter_result.inputs, filter_result.input_table_sources, filter_result.is_null_aware); auto const cuda_source = cudf::jit::parse_single_function_cuda(filter_result.udf, "GENERIC_JOIN_FILTER_OP"); @@ -513,23 +487,14 @@ filter_join_indices_jit(cudf::table_view const& left, cudf::jit::get_udf_kernel(*join_jit_filter_join_kernel_cu_jit, kernel_name, cuda_source); auto configured_kernel = kernel->configure_1d_max_occupancy(0, 0, nullptr, stream.value()); - // Collect scalar columns to append to left device views so join_scalar_accessor - // can read them at indices >= left.num_columns(). - std::vector scalar_cols; - for (auto const& col : filter_result.scalar_columns) { - scalar_cols.push_back(col->view()); - } - // Allocate and compute predicate results auto predicate_results = rmm::device_uvector(left_indices.size(), stream); launch_join_filter_kernel(configured_kernel, - left, - right, left_indices, right_indices, + filter_result.inputs, predicate_results.data(), std::nullopt, - scalar_cols, stream, mr); diff --git a/cpp/src/join/jit/filter_join_kernel.cu b/cpp/src/join/jit/filter_join_kernel.cu index c8b07bc7a625..61f26d261067 100644 --- a/cpp/src/join/jit/filter_join_kernel.cu +++ b/cpp/src/join/jit/filter_join_kernel.cu @@ -10,8 +10,9 @@ #include #include +#include -#include +#include #include #include @@ -31,62 +32,64 @@ namespace cudf::join::jit { // This must match the definition in cudf/join/join.hpp constexpr cudf::size_type JoinNoMatch = cuda::std::numeric_limits::min(); -template -CUDF_KERNEL void filter_join_kernel(cudf::jit::device_span left_indices, - cudf::jit::device_span right_indices, - cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - bool* predicate_results, - void* user_data) +template +__device__ void execute_predicate_op(void* user_data, + size_type row_index, + cuda::std::tuple args) +{ + if constexpr (has_user_data) { + cuda::std::apply([&](auto&&... args) { GENERIC_JOIN_FILTER_OP(user_data, row_index, args...); }, + args); + } else { + cuda::std::apply([&](auto&&... args) { GENERIC_JOIN_FILTER_OP(args...); }, args); + } +} + +template +CUDF_KERNEL void filter_join_kernel(cudf::size_type num_rows, + cudf::size_type const* __restrict__ left_indices, + cudf::size_type const* __restrict__ right_indices, + cudf::column_device_view_core const* __restrict__ columns, + bool* __restrict__ predicate_results, + void* __restrict__ user_data) { auto const start = cudf::detail::grid_1d::global_thread_id(); auto const stride = cudf::detail::grid_1d::grid_stride(); - auto const size = left_indices.size(); - - for (auto i = start; i < size; i += stride) { - auto const left_idx = left_indices[i]; - auto const right_idx = right_indices[i]; + for (auto i = start; i < num_rows; i += stride) { // Skip if either index is JoinNoMatch - if (left_idx == JoinNoMatch || right_idx == JoinNoMatch) { + if (left_indices[i] == JoinNoMatch || right_indices[i] == JoinNoMatch) { predicate_results[i] = false; continue; } + cudf::size_type const* indices[] = {left_indices, right_indices}; + // Each accessor receives both tables and both indices, and internally selects // the appropriate table based on whether it's a left or right accessor. if constexpr (is_null_aware == null_aware::YES) { // Null-aware path: pass optional inputs, get optional result cuda::std::optional result{false}; - if constexpr (has_user_data) { - GENERIC_JOIN_FILTER_OP( - user_data, - i, - &result, - InputAccessors::nullable_element(left_tables, right_tables, left_idx, right_idx, i)...); - } else { - GENERIC_JOIN_FILTER_OP( - &result, - InputAccessors::nullable_element(left_tables, right_tables, left_idx, right_idx, i)...); - } + auto inputs = Accessors::map([&]() { + return cuda::std::tuple{A::nullable_element(columns, indices[A::table_index][i])...}; + }); + execute_predicate_op( + user_data, i, cuda::std::tuple_cat(cuda::std::tuple{&result}, inputs)); predicate_results[i] = result.has_value() && result.value(); } else { // Non-null-aware path: if any input is null, predicate is false - if ((InputAccessors::is_null(left_tables, right_tables, left_idx, right_idx, i) || ...)) { + auto any_null = Accessors::map( + [&]() { return (A::is_null(columns, indices[A::table_index][i]) || ...); }); + if (any_null) { predicate_results[i] = false; continue; } bool result = false; - if constexpr (has_user_data) { - GENERIC_JOIN_FILTER_OP( - user_data, - i, - &result, - InputAccessors::element(left_tables, right_tables, left_idx, right_idx, i)...); - } else { - GENERIC_JOIN_FILTER_OP( - &result, InputAccessors::element(left_tables, right_tables, left_idx, right_idx, i)...); - } + auto inputs = Accessors::map([&]() { + return cuda::std::tuple{A::element(columns, indices[A::table_index][i])...}; + }); + execute_predicate_op( + user_data, i, cuda::std::tuple_cat(cuda::std::tuple{&result}, inputs)); predicate_results[i] = result; } } diff --git a/cpp/src/join/jit/filter_join_kernel.cuh b/cpp/src/join/jit/filter_join_kernel.cuh index 5a9215e810a8..f5d086d215cc 100644 --- a/cpp/src/join/jit/filter_join_kernel.cuh +++ b/cpp/src/join/jit/filter_join_kernel.cuh @@ -17,20 +17,21 @@ namespace cudf::join::jit { * * @tparam has_user_data Whether the predicate function requires user data * @tparam is_null_aware Whether the expression needs input validity as part of its computation - * @tparam InputAccessors Variadic template for input column accessors + * @tparam Accessors type list of accessors for columns used in the predicate * @param left_indices Device span of left table indices * @param right_indices Device span of right table indices - * @param left_tables Device view of left table columns - * @param right_tables Device view of right table columns + * @param left_table Device view of left table columns + * @param right_table Device view of right table columns + * @param scalars Device view of scalar values used in the predicate * @param predicate_results Output array for predicate evaluation results * @param user_data Optional user data for predicate function */ -template -CUDF_KERNEL void filter_join_kernel(cudf::jit::device_span left_indices, - cudf::jit::device_span right_indices, - cudf::column_device_view_core const* left_tables, - cudf::column_device_view_core const* right_tables, - bool* predicate_results, - void* user_data); +template +CUDF_KERNEL void filter_join_kernel(cudf::size_type num_rows, + cudf::size_type const* __restrict__ left_indices, + cudf::size_type const* __restrict__ right_indices, + cudf::column_device_view_core const* __restrict__ columns, + bool* __restrict__ predicate_results, + void* __restrict__ user_data); } // namespace cudf::join::jit diff --git a/cpp/src/stream_compaction/filter/filter.cu b/cpp/src/stream_compaction/filter/filter.cu index 407fe0c4d3b7..f7eea83b7af9 100644 --- a/cpp/src/stream_compaction/filter/filter.cu +++ b/cpp/src/stream_compaction/filter/filter.cu @@ -22,41 +22,42 @@ namespace cudf { namespace detail { -std::vector> filter( - std::span const> predicate_inputs, - std::string const& predicate_udf, - std::vector const& filter_columns, - cudf::udf_source_type source_type, - std::optional user_data, - null_aware is_null_aware, - output_nullability predicate_nullability, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +std::unique_ptr
filter(std::string const& predicate_udf, + cudf::udf_source_type source_type, + null_aware is_null_aware, + std::optional user_data, + std::span predicate_inputs, + table_view const& filter_table, + ops::error_mode error_mode, + output_nullability predicate_nullability, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { - CUDF_EXPECTS(!filter_columns.empty(), + CUDF_EXPECTS(filter_table.num_columns() > 0, "At least one column must be provided to filter.", std::invalid_argument); - auto row_size = filter_columns[0].size(); - CUDF_EXPECTS(std::all_of(filter_columns.begin(), - filter_columns.end(), + auto row_size = filter_table.num_rows(); + CUDF_EXPECTS(std::all_of(filter_table.begin(), + filter_table.end(), [&](auto const& col) { return col.size() == row_size; }), "All columns to filter must have the same number of rows.", std::invalid_argument); - auto predicate = cudf::transform_extended(predicate_inputs, - predicate_udf, - data_type{type_id::BOOL8}, - source_type, - user_data, - is_null_aware, - row_size, - predicate_nullability, - stream, - mr); - - return apply_mask( - cudf::table_view{filter_columns}, predicate->view(), mask_type::RETENTION, stream, mr) - ->release(); + transform_output outputs[] = {transform_output{data_type{type_id::BOOL8}, predicate_nullability}}; + + auto result = cudf::multi_transform(predicate_udf, + source_type, + is_null_aware, + user_data, + predicate_inputs, + outputs, + {}, + filter_table.num_rows(), + error_mode, + stream, + mr); + + return apply_mask(filter_table, result->get_column(0), mask_type::RETENTION, stream, mr); } } // namespace detail @@ -67,19 +68,24 @@ std::unique_ptr
filter(table_view const& predicate_table, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - cudf::detail::row_ir::ast_args ast_args{.table = predicate_table}; - auto args = cudf::detail::row_ir::ast_converter::filter( - cudf::detail::row_ir::target::CUDA, predicate_expr, ast_args, filter_table, stream, mr); - - return std::make_unique
(cudf::detail::filter(args.inputs, - args.udf, - args.filter_columns, - args.source_type, - args.user_data, - args.is_null_aware, - args.predicate_nullability, - stream, - mr)); + auto args = cudf::detail::row_ir::ast_converter::filter(cudf::detail::row_ir::target::CUDA, + predicate_expr, + predicate_table, + {}, + "filter_operation", + stream, + mr); + + return detail::filter(args.udf, + args.source_type, + args.is_null_aware, + args.user_data, + args.inputs, + filter_table, + args.error_mode, + args.outputs[0].nullability, + stream, + mr); } std::vector> filter_extended( @@ -94,15 +100,17 @@ std::vector> filter_extended( rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); - return detail::filter(predicate_inputs, - predicate_udf, - filter_columns, - source_type, - user_data, - is_null_aware, - predicate_nullability, - stream, - mr); + auto table = detail::filter(predicate_udf, + source_type, + is_null_aware, + user_data, + predicate_inputs, + table_view{filter_columns}, + ops::error_mode::IGNORE, + predicate_nullability, + stream, + mr); + return table->release(); } std::vector> filter(std::vector const& predicate_columns, @@ -130,15 +138,17 @@ std::vector> filter(std::vector const& pred } } - return detail::filter(inputs, - predicate_udf, - filter_columns, - is_ptx ? cudf::udf_source_type::PTX : cudf::udf_source_type::CUDA, - user_data, - is_null_aware, - predicate_nullability, - stream, - mr); + auto table = detail::filter(predicate_udf, + is_ptx ? cudf::udf_source_type::PTX : cudf::udf_source_type::CUDA, + is_null_aware, + user_data, + inputs, + table_view{filter_columns}, + ops::error_mode::IGNORE, + predicate_nullability, + stream, + mr); + return table->release(); } } // namespace cudf diff --git a/cpp/src/transform/jit/kernel.cu b/cpp/src/transform/jit/kernel.cu index 2456433d0bc6..f95bb0cfcda3 100644 --- a/cpp/src/transform/jit/kernel.cu +++ b/cpp/src/transform/jit/kernel.cu @@ -6,6 +6,16 @@ #include #include #include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include #include #include @@ -18,6 +28,7 @@ #include #include +#include #include #include @@ -34,20 +45,39 @@ namespace cudf { namespace jit { -template -__device__ void execute_transform_op(void* user_data, size_type element_idx, Args args) +template +__device__ void execute_transform_op(error_sink* __restrict__ error_sink, + void* user_data, + size_type element_idx, + Args args) { // TODO: static assert invocable if constexpr (has_user_data) { - cuda::std::apply([&](auto... a) { GENERIC_TRANSFORM_OP(a...); }, - cuda::std::tuple_cat(cuda::std::tuple{user_data, element_idx}, args)); + cuda::std::apply( + [&](auto... a) { + if constexpr (mode == ops::error_mode::IGNORE) { + GENERIC_TRANSFORM_OP(a...); + } else { + error_sink->report(GENERIC_TRANSFORM_OP(a...)); + } + }, + cuda::std::tuple_cat(cuda::std::tuple{user_data, element_idx}, args)); } else { - cuda::std::apply([&](auto... a) { GENERIC_TRANSFORM_OP(a...); }, args); + cuda::std::apply( + [&](auto... a) { + if constexpr (mode == ops::error_mode::IGNORE) { + GENERIC_TRANSFORM_OP(a...); + } else { + error_sink->report(GENERIC_TRANSFORM_OP(a...)); + } + }, + args); } } /// @brief The generic transform kernel. Supports all types and nullability combinations. -template @@ -55,7 +85,8 @@ CUDF_KERNEL void transform_kernel(size_type row_size, bitmask_type const* __restrict__ stencil, void* __restrict__ user_data, column_device_view_core const* __restrict__ input_cols, - mutable_column_device_view_core const* __restrict__ output_cols) + mutable_column_device_view_core const* __restrict__ output_cols, + error_sink* __restrict__ error_sink) { // TODO: ensure block size is a multiple of warp size for correct warp-synchronous behavior auto start = detail::grid_1d::global_thread_id(); @@ -75,8 +106,8 @@ CUDF_KERNEL void transform_kernel(size_type row_size, auto out_ptrs = cuda::std::apply([&](auto&... args) { return cuda::std::tuple{&args...}; }, outs); - execute_transform_op( - user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins)); + execute_transform_op( + error_sink, user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins)); OutputAccessors::map([&]() { (A::assign(output_cols, element_idx, cuda::std::get(outs)), ...); @@ -96,8 +127,8 @@ CUDF_KERNEL void transform_kernel(size_type row_size, auto out_ptrs = cuda::std::apply([&](auto&... args) { return cuda::std::tuple{&args...}; }, outs); - execute_transform_op( - user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins)); + execute_transform_op( + error_sink, user_data, element_idx, cuda::std::tuple_cat(out_ptrs, ins)); OutputAccessors::map([&]() { (A::assign(output_cols, element_idx, *cuda::std::get(outs)), ...); diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index 436316241432..a8886eb8273e 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -24,6 +24,7 @@ #include #include +#include #include #include #include @@ -158,7 +159,8 @@ using handle = std::variant< namespace jit_transform { -jitify2::Kernel instantiate(null_aware is_null_aware, +jitify2::Kernel instantiate(ops::error_mode error_handling_mode, + null_aware is_null_aware, bool has_user_data, std::string const& ins, std::string const& outs, @@ -176,7 +178,7 @@ jitify2::Kernel instantiate(null_aware is_null_aware, : jit::parse_single_function_cuda(udf, "GENERIC_TRANSFORM_OP"); auto kernel = jitify2::reflection::Template("cudf::jit::transform_kernel") - .instantiate(is_null_aware, has_user_data, ins, outs); + .instantiate(error_handling_mode, is_null_aware, has_user_data, ins, outs); return jit::get_udf_kernel( *transform_jit_kernel_cu_jit, kernel, cuda_source, {"-restrict", "--dopt=on"}); @@ -188,10 +190,11 @@ void launch(jitify2::Kernel const& kernel, void* user_data, column_device_view_core const* input_cols, mutable_column_device_view_core const* output_cols, + jit::error_sink* error_sink, rmm::cuda_stream_view stream) { CUDF_FUNC_RANGE(); - void* args[] = {&row_size, &stencil, &user_data, &input_cols, &output_cols}; + void* args[] = {&row_size, &stencil, &user_data, &input_cols, &output_cols, &error_sink}; kernel->configure_1d_max_occupancy(0, 0, nullptr, stream.value())->launch_raw(args); } @@ -245,7 +248,7 @@ auto reflect(udf_source_type source_type, auto element = std::visit([](auto& c) { return reflect_input_element(c); }, in); bool as_scalar = std::holds_alternative(in); auto accessor = jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate(i, column, element, as_scalar); + .instantiate(i, column, element, as_scalar, 0); in_types.push_back(accessor); } @@ -257,7 +260,7 @@ auto reflect(udf_source_type source_type, auto element = std::visit([](auto& c) { return reflect_output_element(c); }, out); bool as_scalar = false; // never scalar auto accessor = jitify2::reflection::Template("cudf::jit::column_accessor") - .instantiate(i, column, element, as_scalar); + .instantiate(i, column, element, as_scalar, 0); out_types.push_back(accessor); } @@ -319,20 +322,23 @@ auto to_args(std::span inputs, return std::make_tuple(std::move(d_args), std::move(handles)); } -void run(null_aware is_null_aware, +void run(ops::error_mode error_handling_mode, + null_aware is_null_aware, bool has_user_data, size_type row_size, bitmask_type const* d_stencil, void* user_data, std::span inputs, std::span outputs, + jit::error_sink* d_error_sink, std::string const& udf, udf_source_type source_type, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { auto [in_types, out_types, ptx_in_types, ptx_out_types] = reflect(source_type, inputs, outputs); - auto kernel = instantiate(is_null_aware, + auto kernel = instantiate(error_handling_mode, + is_null_aware, has_user_data, in_types, out_types, @@ -344,7 +350,8 @@ void run(null_aware is_null_aware, auto* input_cols = reinterpret_cast(cols.data()); auto* output_cols = reinterpret_cast(input_cols + inputs.size()); - return launch(kernel, row_size, d_stencil, user_data, input_cols, output_cols, stream); + return launch( + kernel, row_size, d_stencil, user_data, input_cols, output_cols, d_error_sink, stream); } } // namespace jit_transform @@ -797,6 +804,7 @@ auto finalize_outputs(null_aware is_null_aware, std::unique_ptr
execute_transform(std::string const& udf, udf_source_type source_type, + ops::error_mode error_handling_mode, null_aware is_null_aware, std::optional in_row_size, std::optional user_data, @@ -819,19 +827,47 @@ std::unique_ptr
execute_transform(std::string const& udf, auto stencil_arg = stencil.has_value() ? stencil->first : nullptr; auto stencil_has_nulls = stencil.has_value() ? (stencil->second > 0) : false; - jit_transform::run(is_null_aware, + + std::optional> d_error_sink = std::nullopt; + + switch (error_handling_mode) { + case ops::error_mode::IGNORE: break; + case ops::error_mode::ANY_ROW: + d_error_sink = rmm::device_scalar(jit::error_sink{}, stream); + break; + } + + jit_transform::run(error_handling_mode, + is_null_aware, user_data.has_value(), row_size, stencil_has_nulls ? stencil_arg : nullptr, user_data.value_or(nullptr), inputs, output_columns, + d_error_sink.has_value() ? d_error_sink->data() : nullptr, udf, source_type, stream, mr); auto finalized = finalize_outputs(is_null_aware, row_size, std::move(output_columns), stream, mr); + + switch (error_handling_mode) { + case ops::error_mode::IGNORE: { + } break; + case ops::error_mode::ANY_ROW: { + auto error = d_error_sink->value(stream).any_error(); + switch (error) { + case ops::errc::OK: break; + case ops::errc::OVERFLOW: CUDF_FAIL("Overflow error in transform UDF", std::overflow_error); + case ops::errc::DIVISION_BY_ZERO: + CUDF_FAIL("Division by zero error in transform UDF", std::overflow_error); + default: CUDF_FAIL("Unknown error in transform UDF", std::runtime_error); + } + } break; + } + return std::make_unique
(std::move(finalized)); } @@ -845,6 +881,7 @@ std::unique_ptr
multi_transform(std::string const& udf, std::span outputs, std::vector>&& string_offsets, std::optional row_size, + ops::error_mode error_handling_mode, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { @@ -852,6 +889,7 @@ std::unique_ptr
multi_transform(std::string const& udf, perform_checks(source_type, is_null_aware, row_size, inputs, outputs, string_offsets); return execute_transform(udf, source_type, + error_handling_mode, is_null_aware, row_size, user_data, @@ -874,9 +912,18 @@ std::unique_ptr transform_extended(std::span inpu rmm::device_async_resource_ref mr) { transform_output outputs[] = {{.type = output_type, .nullability = null_policy}}; - auto table = multi_transform( - udf, source_type, is_null_aware, user_data, inputs, outputs, {}, row_size, stream, mr); - auto cols = table->release(); + auto table = multi_transform(udf, + source_type, + is_null_aware, + user_data, + inputs, + outputs, + {}, + row_size, + ops::error_mode::IGNORE, + stream, + mr); + auto cols = table->release(); return std::move(cols[0]); } @@ -922,19 +969,21 @@ std::unique_ptr compute_column_jit(table_view const& table, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - detail::row_ir::ast_args ast_args{.table = table}; auto args = detail::row_ir::ast_converter::compute_column( - detail::row_ir::target::CUDA, expr, ast_args, stream, mr); - return transform_extended(args.inputs, - args.udf, - args.output_type, - args.source_type, - args.user_data, - args.is_null_aware, - args.row_size, - args.null_policy, - stream, - mr); + detail::row_ir::target::CUDA, expr, table, {}, "compute_operation", stream, mr); + auto result = multi_transform(args.udf, + args.source_type, + args.is_null_aware, + args.user_data, + args.inputs, + args.outputs, + std::move(args.string_offsets), + args.row_size, + args.error_mode, + stream, + mr); + auto cols = result->release(); + return std::move(cols[0]); } } // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 68cde65c57bb..9945490e0120 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -681,6 +681,10 @@ ConfigureTest(ENCODE_TEST encode/encode_tests.cpp) # * ast tests ------------------------------------------------------------------------------------- ConfigureTest(AST_TEST ast/transform_tests.cpp ast/ast_tree_tests.cpp) +# ################################################################################################## +# * jit-ast tests ------------------------------------------------------------------------------------- +ConfigureTest(JIT_AST_TEST ast/jit_ast_tests.cpp) + # ################################################################################################## # * lists tests ---------------------------------------------------------------------------------- ConfigureTest( diff --git a/cpp/tests/ast/jit_ast_tests.cpp b/cpp/tests/ast/jit_ast_tests.cpp new file mode 100644 index 000000000000..cb3ddbde7718 --- /dev/null +++ b/cpp/tests/ast/jit_ast_tests.cpp @@ -0,0 +1,707 @@ + +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include + +constexpr cudf::test::debug_output_level VERBOSITY{cudf::test::debug_output_level::ALL_ERRORS}; + +template +using column_wrapper = cudf::test::fixed_width_column_wrapper; + +template +using decimal_column_wrapper = cudf::test::fixed_point_column_wrapper; + +struct JITExpressionTest : public cudf::test::BaseFixture {}; + +template +struct JITIntegerArithmeticTest : public cudf::test::BaseFixture { + static constexpr T MAX = std::numeric_limits::max(); + static constexpr T MIN = std::numeric_limits::min(); +}; + +template +struct JITSignedIntegerArithmeticTest : public JITIntegerArithmeticTest {}; + +template +struct JITDecimalArithmeticTest : public JITIntegerArithmeticTest {}; + +using SignedIntegralTypesNotBool = cudf::test::Types; + +TYPED_TEST_SUITE(JITIntegerArithmeticTest, cudf::test::IntegralTypesNotBool); +TYPED_TEST_SUITE(JITSignedIntegerArithmeticTest, SignedIntegralTypesNotBool); +TYPED_TEST_SUITE(JITDecimalArithmeticTest, cudf::test::FixedPointTypes); + +TEST_F(JITExpressionTest, NullifyIf) +{ + auto a = column_wrapper{{3, 20, 1, 50, 0, 20}}; + auto condition = column_wrapper{{false, true, false, true, false, true}}; + auto expected = column_wrapper{{3, 0, 1, 0, 0, 0}, {1, 0, 1, 0, 1, 0}}; + auto table = cudf::table_view{{a, condition}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto condition_ref = cudf::ast::column_reference(1); + auto& nullify_if = cudf::ast::jit::nullify_if(tree, a_ref, condition_ref); + auto result = cudf::compute_column_jit(table, nullify_if); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +TEST_F(JITExpressionTest, Coalesce) +{ + auto a = column_wrapper{{1, 3, 5, 7, 9, 11}, {1, 0, 0, 1, 0, 0}}; + auto b = column_wrapper{{2, 4, 6, 8, 10, 12}, {1, 1, 1, 0, 1, 0}}; + auto expected = column_wrapper{{1, 4, 6, 7, 10, 0}, {1, 1, 1, 1, 1, 0}}; + auto table = cudf::table_view{{a, b}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto& coalesce = cudf::ast::jit::coalesce(tree, a_ref, b_ref); + auto result = cudf::compute_column_jit(table, coalesce); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +TYPED_TEST(JITIntegerArithmeticTest, AnsiAdd) +{ + using T = TypeParam; + auto a = column_wrapper{{3, 20, 1, 50}}; + auto b = column_wrapper{{10, 7, 20, 0}}; + auto b_fail = column_wrapper{{T{10}, this->MAX, T{20}, T{0}}}; + auto expected = column_wrapper{{13, 27, 21, 50}}; + auto expected_fail = column_wrapper{{13, 0, 21, 50}, {1, 0, 1, 1}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto& add = cudf::ast::jit::ansi_add(tree, a_ref, b_ref); + auto& add_fail = cudf::ast::jit::ansi_add(tree, a_ref, b_fail_ref); + auto& try_add_fail = cudf::ast::jit::ansi_try_add(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, add); + auto result_fail = cudf::compute_column_jit(table, try_add_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, add_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiAdd) +{ + using T = TypeParam; + using R = typename T::rep; + auto a = decimal_column_wrapper{{3, 20, 1, 50}, numeric::scale_type{0}}; + auto b = decimal_column_wrapper{{10, 7, 20, 0}, numeric::scale_type{0}}; + auto b_fail = decimal_column_wrapper{{R{10}, this->MAX, R{20}, R{0}}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{13, 27, 21, 50}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{13, 0, 21, 50}, {1, 0, 1, 1}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto& add = cudf::ast::jit::ansi_add(tree, a_ref, b_ref); + auto& add_fail = cudf::ast::jit::ansi_add(tree, a_ref, b_fail_ref); + auto& try_add_fail = cudf::ast::jit::ansi_try_add(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, add); + auto result_fail = cudf::compute_column_jit(table, try_add_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, add_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITSignedIntegerArithmeticTest, AnsiSub) +{ + using T = TypeParam; + auto a = column_wrapper{{3, 20, 1, 50}}; + auto b = column_wrapper{{10, 7, 20, 0}}; + auto b_fail = column_wrapper{{T{10}, T{this->MIN}, T{20}, T{0}}}; + auto expected = column_wrapper{{-7, 13, -19, 50}}; + auto expected_fail = column_wrapper{{-7, 0, -19, 50}, {1, 0, 1, 1}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto& sub = cudf::ast::jit::ansi_sub(tree, a_ref, b_ref); + auto& sub_fail = cudf::ast::jit::ansi_sub(tree, a_ref, b_fail_ref); + auto& try_sub_fail = cudf::ast::jit::ansi_try_sub(tree, a_ref, b_fail_ref); + + auto result = cudf::compute_column_jit(table, sub); + auto result_fail = cudf::compute_column_jit(table, try_sub_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, sub_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiSub) +{ + using T = TypeParam; + using R = typename T::rep; + auto a = decimal_column_wrapper{{3, 20, 1, 50}, numeric::scale_type{0}}; + auto b = decimal_column_wrapper{{10, 7, 20, 0}, numeric::scale_type{0}}; + auto b_fail = + decimal_column_wrapper{{R{10}, R{this->MIN}, R{20}, R{0}}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{-7, 13, -19, 50}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{-7, 0, -19, 50}, {1, 0, 1, 1}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto& sub = cudf::ast::jit::ansi_sub(tree, a_ref, b_ref); + auto& sub_fail = cudf::ast::jit::ansi_sub(tree, a_ref, b_fail_ref); + auto& try_sub_fail = cudf::ast::jit::ansi_try_sub(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, sub); + auto result_fail = cudf::compute_column_jit(table, try_sub_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, sub_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITIntegerArithmeticTest, AnsiMul) +{ + using T = TypeParam; + auto a = column_wrapper{{3, 20, 2, 50}}; + auto b = column_wrapper{{10, 2, 1, 0}}; + auto b_fail = column_wrapper{{T{10}, T{this->MAX}, T{1}, T{0}}}; + auto expected = column_wrapper{{30, 40, 2, 0}}; + auto expected_fail = column_wrapper{{30, 0, 2, 0}, {1, 0, 1, 1}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& mul = cudf::ast::jit::ansi_mul(tree, a_ref, b_ref); + auto& mul_fail = cudf::ast::jit::ansi_mul(tree, a_ref, b_fail_ref); + auto& try_mul_fail = cudf::ast::jit::ansi_try_mul(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, mul); + auto result_fail = cudf::compute_column_jit(table, try_mul_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, mul_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiMul) +{ + using T = TypeParam; + using R = typename T::rep; + auto a = decimal_column_wrapper{{3, 20, 2, 50}, numeric::scale_type{0}}; + auto b = decimal_column_wrapper{{10, 7, 1, 0}, numeric::scale_type{0}}; + auto b_fail = + decimal_column_wrapper{{R{10}, R{this->MAX}, R{1}, R{0}}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{30, 140, 2, 0}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{30, 0, 2, 0}, {1, 0, 1, 1}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& mul = cudf::ast::jit::ansi_mul(tree, a_ref, b_ref); + auto& mul_fail = cudf::ast::jit::ansi_mul(tree, a_ref, b_fail_ref); + auto& try_mul_fail = cudf::ast::jit::ansi_try_mul(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, mul); + auto result_fail = cudf::compute_column_jit(table, try_mul_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, mul_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITIntegerArithmeticTest, AnsiDiv) +{ + using T = TypeParam; + auto a = column_wrapper{{3, 20, 1, 50}}; + auto b = column_wrapper{{10, 7, 2, 1}}; + auto b_fail = column_wrapper{{10, 1, 20, 0}}; + auto expected = column_wrapper{{0, 2, 0, 50}}; + auto expected_fail = column_wrapper{{0, 20, 0, 50}, {1, 1, 1, 0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& div = cudf::ast::jit::ansi_div(tree, a_ref, b_ref); + auto& div_fail = cudf::ast::jit::ansi_div(tree, a_ref, b_fail_ref); + auto& try_div_fail = cudf::ast::jit::ansi_try_div(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, div); + auto result_fail = cudf::compute_column_jit(table, try_div_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, div_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiDiv) +{ + using T = TypeParam; + auto a = decimal_column_wrapper{{3, 20, 1, 50}, numeric::scale_type{0}}; + auto b = decimal_column_wrapper{{10, 7, 2, 1}, numeric::scale_type{0}}; + auto b_fail = decimal_column_wrapper{{10, 1, 20, 0}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{0, 2, 0, 50}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{0, 20, 0, 50}, {1, 1, 1, 0}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& div = cudf::ast::jit::ansi_div(tree, a_ref, b_ref); + auto& div_fail = cudf::ast::jit::ansi_div(tree, a_ref, b_fail_ref); + auto& try_div_fail = cudf::ast::jit::ansi_try_div(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, div); + auto result_fail = cudf::compute_column_jit(table, try_div_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, div_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITIntegerArithmeticTest, AnsiMod) +{ + using T = TypeParam; + auto a = column_wrapper{{3, 20, 1, 50}}; + auto b = column_wrapper{{10, 7, 2, 1}}; + auto b_fail = column_wrapper{{10, 1, 20, 0}}; + auto expected = column_wrapper{{3, 6, 1, 0}}; + auto expected_fail = column_wrapper{{3, 0, 1, 0}, {1, 1, 1, 0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& mod = cudf::ast::jit::ansi_mod(tree, a_ref, b_ref); + auto& mod_fail = cudf::ast::jit::ansi_mod(tree, a_ref, b_fail_ref); + auto& try_mod_fail = cudf::ast::jit::ansi_try_mod(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, mod); + auto result_fail = cudf::compute_column_jit(table, try_mod_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, mod_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiMod) +{ + using T = TypeParam; + auto a = decimal_column_wrapper{{3, 20, 1, 50}, numeric::scale_type{0}}; + auto b = decimal_column_wrapper{{10, 7, 2, 1}, numeric::scale_type{0}}; + auto b_fail = decimal_column_wrapper{{10, 1, 20, 0}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{3, 6, 1, 0}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{3, 0, 1, 0}, {1, 1, 1, 0}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, b, b_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto b_fail_ref = cudf::ast::column_reference(2); + auto tree = cudf::ast::tree{}; + auto& mod = cudf::ast::jit::ansi_mod(tree, a_ref, b_ref); + auto& mod_fail = cudf::ast::jit::ansi_mod(tree, a_ref, b_fail_ref); + auto& try_mod_fail = cudf::ast::jit::ansi_try_mod(tree, a_ref, b_fail_ref); + auto result = cudf::compute_column_jit(table, mod); + auto result_fail = cudf::compute_column_jit(table, try_mod_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, mod_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITSignedIntegerArithmeticTest, AnsiAbs) +{ + using T = TypeParam; + auto a = column_wrapper{{T{3}, T{-20}, T{1}, T{-50}, this->MAX, T{this->MIN + 1}, T{0}}}; + auto a_fail = column_wrapper{{T{3}, T{-20}, T{1}, T{-50}, this->MIN, T{1}, T{0}}}; + auto expected = + column_wrapper{{T{3}, T{20}, T{1}, T{50}, this->MAX, T{std::abs(this->MIN + 1)}, T{0}}}; + auto expected_fail = column_wrapper{{3, 20, 1, 50, 0, 1, 0}, {1, 1, 1, 1, 0, 1, 1}}; + auto table = cudf::table_view{{a, a_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto a_fail_ref = cudf::ast::column_reference(1); + auto tree = cudf::ast::tree{}; + auto& abs = cudf::ast::jit::ansi_abs(tree, a_ref); + auto& abs_fail = cudf::ast::jit::ansi_abs(tree, a_fail_ref); + auto& try_abs_fail = cudf::ast::jit::ansi_try_abs(tree, a_fail_ref); + auto result = cudf::compute_column_jit(table, abs); + auto result_fail = cudf::compute_column_jit(table, try_abs_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, abs_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiAbs) +{ + using T = TypeParam; + using R = typename T::rep; + auto a = decimal_column_wrapper{ + {R{3}, R{-20}, R{1}, R{-50}, this->MAX, R{this->MIN + 1}, R{0}}, numeric::scale_type{0}}; + auto a_fail = decimal_column_wrapper{{R{3}, R{-20}, R{1}, R{-50}, this->MIN, R{1}, R{0}}, + numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{ + {R{3}, R{20}, R{1}, R{50}, this->MAX, R{std::abs(this->MIN + 1)}, R{0}}, + numeric::scale_type{0}}; + auto expected_fail = decimal_column_wrapper{ + {3, 20, 1, 50, 0, 1, 0}, {1, 1, 1, 1, 0, 1, 1}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, a_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto a_fail_ref = cudf::ast::column_reference(1); + auto tree = cudf::ast::tree{}; + auto& abs = cudf::ast::jit::ansi_abs(tree, a_ref); + auto& abs_fail = cudf::ast::jit::ansi_abs(tree, a_fail_ref); + auto& try_abs_fail = cudf::ast::jit::ansi_try_abs(tree, a_fail_ref); + auto result = cudf::compute_column_jit(table, abs); + auto result_fail = cudf::compute_column_jit(table, try_abs_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, abs_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITSignedIntegerArithmeticTest, AnsiNeg) +{ + using T = TypeParam; + auto a = column_wrapper{{T{3}, T{-20}, T{1}, T{-50}, this->MAX, T{-this->MAX}, T{0}}}; + auto a_fail = column_wrapper{{T{3}, T{-20}, T{1}, T{-50}, this->MIN, T{1}, T{0}}}; + auto expected = column_wrapper{{T{-3}, T{20}, T{-1}, T{50}, T{-this->MAX}, this->MAX, T{0}}}; + auto expected_fail = column_wrapper{{-3, 20, -1, 50, 0, -1, 0}, {1, 1, 1, 1, 0, 1, 1}}; + auto table = cudf::table_view{{a, a_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto a_fail_ref = cudf::ast::column_reference(1); + auto tree = cudf::ast::tree{}; + auto& neg = cudf::ast::jit::ansi_neg(tree, a_ref); + auto& neg_fail = cudf::ast::jit::ansi_neg(tree, a_fail_ref); + auto& try_neg_fail = cudf::ast::jit::ansi_try_neg(tree, a_fail_ref); + auto result = cudf::compute_column_jit(table, neg); + auto result_fail = cudf::compute_column_jit(table, try_neg_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, neg_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiNeg) +{ + using T = TypeParam; + using R = typename T::rep; + auto a = decimal_column_wrapper{{R{3}, R{-20}, R{1}, R{-50}, this->MAX, R{-this->MAX}, R{0}}, + numeric::scale_type{0}}; + auto a_fail = decimal_column_wrapper{{R{3}, R{-20}, R{1}, R{-50}, this->MIN, R{1}, R{0}}, + numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{ + {R{-3}, R{20}, R{-1}, R{50}, R{-this->MAX}, this->MAX, R{0}}, numeric::scale_type{0}}; + auto expected_fail = decimal_column_wrapper{ + {-3, 20, -1, 50, 0, -1, 0}, {1, 1, 1, 1, 0, 1, 1}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a, a_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto a_fail_ref = cudf::ast::column_reference(1); + auto tree = cudf::ast::tree{}; + auto& neg = cudf::ast::jit::ansi_neg(tree, a_ref); + auto& neg_fail = cudf::ast::jit::ansi_neg(tree, a_fail_ref); + auto& try_neg_fail = cudf::ast::jit::ansi_try_neg(tree, a_fail_ref); + auto result = cudf::compute_column_jit(table, neg); + auto result_fail = cudf::compute_column_jit(table, try_neg_fail); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, neg_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, AnsiPrecisionCheck) +{ + using T = TypeParam; + auto a = decimal_column_wrapper{{3, 200, 250, 200}, numeric::scale_type{0}}; + auto a_fail = decimal_column_wrapper{{3, 200, 250, 20000}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{3, 200, 250, 200}, numeric::scale_type{0}}; + auto expected_fail = + decimal_column_wrapper{{3, 200, 250, 200}, {1, 1, 1, 0}, numeric::scale_type{0}}; + auto max_precision = cudf::numeric_scalar(3); + auto table = cudf::table_view{{a, a_fail}}; + auto a_ref = cudf::ast::column_reference(0); + auto a_fail_ref = cudf::ast::column_reference(1); + auto tree = cudf::ast::tree{}; + auto precision = cudf::ast::literal(max_precision); + auto& precision_check = cudf::ast::jit::ansi_precision_check(tree, a_ref, precision); + auto& precision_check_fail = cudf::ast::jit::ansi_precision_check(tree, a_fail_ref, precision); + auto& try_precision_check = cudf::ast::jit::ansi_try_precision_check(tree, a_fail_ref, precision); + auto result = cudf::compute_column_jit(table, precision_check); + auto result_fail = cudf::compute_column_jit(table, try_precision_check); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); + + EXPECT_THROW(result = cudf::compute_column_jit(table, precision_check_fail), std::overflow_error); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_fail, result_fail->view(), VERBOSITY); +} + +TEST_F(JITExpressionTest, BitShiftLeft) +{ + auto a = column_wrapper{0b111111, 0b111110, 0b101111, 0b1100}; + auto expected = column_wrapper{0b11111100, 0b11111000, 0b10111100, 0b110000}; + auto shift = cudf::numeric_scalar(2); + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + auto shift_literal = cudf::ast::literal(shift); + auto& shift_left = cudf::ast::jit::bit_shift_left(tree, a_ref, shift_literal); + auto result = cudf::compute_column_jit(table, shift_left); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +TEST_F(JITExpressionTest, BitShiftRight) +{ + auto a = column_wrapper{0b1111, 0b10111, 0b11100, 0b11110011}; + auto expected = column_wrapper{0b11, 0b101, 0b111, 0b111100}; + auto shift = cudf::numeric_scalar(2); + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + auto shift_literal = cudf::ast::literal(shift); + auto& shift_right = cudf::ast::jit::bit_shift_right(tree, a_ref, shift_literal); + auto result = cudf::compute_column_jit(table, shift_right); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +template +void test_cast() +{ + auto a = column_wrapper{{0, 1, 2, 3, 4, 5}}; + auto expected = column_wrapper{{0, 1, 2, 3, 4, 5}}; + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + + cudf::ast::expression const* cast = nullptr; + + if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_b8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i16(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i32(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i64(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u16(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u32(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u64(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_f32(tree, a_ref); + } else { + static_assert(std::is_same_v); + cast = &cudf::ast::jit::cast_to_f64(tree, a_ref); + } + + auto result = cudf::compute_column_jit(table, *cast); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +template +void test_from_decimal_cast() +{ + auto a = decimal_column_wrapper{{0, 1, 2, 3, 4, 5}, numeric::scale_type{0}}; + auto expected = column_wrapper{0, 1, 2, 3, 4, 5}; + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + + cudf::ast::expression const* cast = nullptr; + + if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_b8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i16(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i32(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_i64(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u8(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u16(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u32(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_u64(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_f32(tree, a_ref); + } else { + static_assert(std::is_same_v); + cast = &cudf::ast::jit::cast_to_f64(tree, a_ref); + } + + auto result = cudf::compute_column_jit(table, *cast); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +template +void test_cast_to() +{ + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_cast(); + test_from_decimal_cast(); + test_from_decimal_cast(); + test_from_decimal_cast(); +} + +TEST_F(JITExpressionTest, Cast) +{ + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); + test_cast_to(); +} + +template +void test_decimal_cast() +{ + auto a = decimal_column_wrapper{{0, 1, 2, 3, 4, 5}, numeric::scale_type{0}}; + auto expected = decimal_column_wrapper{{0, 1, 2, 3, 4, 5}, numeric::scale_type{0}}; + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + + cudf::ast::expression const* cast = nullptr; + + if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_dec32(tree, a_ref); + } else if constexpr (std::is_same_v) { + cast = &cudf::ast::jit::cast_to_dec64(tree, a_ref); + } else if constexpr (std::is_same_v) { + static_assert(std::is_same_v); + cast = &cudf::ast::jit::cast_to_dec128(tree, a_ref); + } + + auto result = cudf::compute_column_jit(table, *cast); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +TYPED_TEST(JITDecimalArithmeticTest, CastTo) +{ + using T = TypeParam; + test_decimal_cast(); + test_decimal_cast(); + test_decimal_cast(); +} + +TEST_F(JITExpressionTest, Rescale) +{ + auto a = cudf::test::fixed_point_column_wrapper{{123, 1234, 12345, 123456, 1234567}, + numeric::scale_type{0}}; + auto expected = cudf::test::fixed_point_column_wrapper{ + {12300, 123400, 1234500, 12345600, 123456700}, numeric::scale_type{-2}}; + auto table = cudf::table_view{{a}}; + auto a_ref = cudf::ast::column_reference(0); + auto tree = cudf::ast::tree{}; + auto& rescaled = cudf::ast::jit::rescale(tree, a_ref, -2); + auto result = cudf::compute_column_jit(table, rescaled); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +TEST_F(JITExpressionTest, AnsiFused) +{ + constexpr auto I32_MAX = std::numeric_limits::max(); + auto a = column_wrapper{{1, 3, 20, 1, 50, 10}}; + auto b = column_wrapper{{1, 10, 7, 20, I32_MAX, 2}}; + auto c = column_wrapper{{1, 5, 4, I32_MAX, 2, 5}}; + auto d = column_wrapper{{0, 1, 0, 0, 1, 5}}; + auto expected = column_wrapper{{0, 65, 0, 0, 0, 12}, {0, 1, 0, 0, 0, 1}}; + auto table = cudf::table_view{{a, b, c, d}}; + auto tree = cudf::ast::tree{}; + auto a_ref = cudf::ast::column_reference(0); + auto b_ref = cudf::ast::column_reference(1); + auto c_ref = cudf::ast::column_reference(2); + auto d_ref = cudf::ast::column_reference(3); + auto& add = cudf::ast::jit::ansi_try_add(tree, a_ref, b_ref); + auto& mul = cudf::ast::jit::ansi_try_mul(tree, add, c_ref); + auto& div = cudf::ast::jit::ansi_try_div(tree, mul, d_ref); + auto result = cudf::compute_column_jit(table, div); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view(), VERBOSITY); +} + +CUDF_TEST_PROGRAM_MAIN() diff --git a/cpp/tests/jit/row_ir.cpp b/cpp/tests/jit/row_ir.cpp index 32022af56ddf..16eaddec081b 100644 --- a/cpp/tests/jit/row_ir.cpp +++ b/cpp/tests/jit/row_ir.cpp @@ -20,37 +20,51 @@ namespace row_ir = cudf::detail::row_ir; -struct RowIRCudaCodeGenTest : public ::testing::Test {}; +struct RowIRCudaCodeGenTest : public ::testing::Test { + std::unique_ptr f32 = + cudf::test::fixed_width_column_wrapper({1.0f, 2.0f, 3.0f}).release(); + std::unique_ptr f64 = + cudf::test::fixed_width_column_wrapper({1.0, 2.0, 3.0}).release(); + std::unique_ptr d32 = + cudf::test::fixed_point_column_wrapper({1, 2, 3}, numeric::scale_type{2}).release(); + std::unique_ptr i32 = + cudf::test::fixed_width_column_wrapper({1, 2, 3}).release(); + std::unique_ptr b8 = + cudf::test::fixed_width_column_wrapper({true, false, true}).release(); + cudf::table_view table = cudf::table_view({*f32, *f64, *d32, *i32}); +}; TEST_F(RowIRCudaCodeGenTest, GetInput) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = {{"in_0", {cudf::data_type{cudf::type_id::INT32}}}, - {"in_1", {cudf::data_type{cudf::type_id::FLOAT32}}}}; - - row_ir::instance_info info{inputs, {}}; - { - row_ir::instance_context ctx{}; - row_ir::get_input get_input_0{0}; - get_input_0.instantiate(ctx, info); - auto code = get_input_0.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + row_ir::code_sink sink; + row_ir::node get_input_0{row_ir::input_reference{0}}; + get_input_0.instantiate(ctx); + get_input_0.emit_code(ctx, target_info, sink); - auto expected_code = "int32_t tmp_0 = in_0;"; + auto expected_code = "int32_t tmp_0 = in_0;\n"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } { - row_ir::instance_context ctx{}; - row_ir::get_input get_input_1{1}; - get_input_1.instantiate(ctx, info); - auto null_code = get_input_1.generate_code(ctx, target_info, info); - - auto expected_null_code = "float tmp_0 = in_1;"; - - EXPECT_EQ(null_code, expected_null_code); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + row_ir::code_sink sink; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*f32); + row_ir::node get_input_1{row_ir::input_reference{1}}; + get_input_1.instantiate(ctx); + get_input_1.emit_code(ctx, target_info, sink); + + auto expected_null_code = "float tmp_0 = in_1;\n"; + + EXPECT_EQ(sink.get_code(), expected_null_code); } } @@ -58,39 +72,49 @@ TEST_F(RowIRCudaCodeGenTest, SetOutput) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = {{"in_0", {cudf::data_type{cudf::type_id::INT32}}}, - {"in_1", {cudf::data_type{cudf::type_id::FLOAT32}}}}; - - row_ir::untyped_var_info outputs[] = {{"out_0"}, {"out_1"}}; - - row_ir::instance_info info{inputs, outputs}; - { - row_ir::instance_context ctx{}; - row_ir::set_output set_output_0{0, std::make_unique(0)}; - set_output_0.instantiate(ctx, info); - auto code = set_output_0.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*f32); + [[maybe_unused]] auto out0 = ctx.add_output(); + [[maybe_unused]] auto out1 = ctx.add_output(); + row_ir::code_sink sink; + row_ir::node set_output_0{row_ir::output_reference{0}, + row_ir::node{row_ir::input_reference{0}}}; + set_output_0.instantiate(ctx); + set_output_0.emit_code(ctx, target_info, sink); auto expected_code = - R"***(int32_t tmp_0 = in_0; -int32_t tmp_1 = tmp_0; -*out_0 = tmp_1;)***"; + R"***(int32_t tmp_1 = in_0; +int32_t tmp_0 = tmp_1; +*out_0 = tmp_0; +)***"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } { - row_ir::instance_context ctx{}; - row_ir::set_output set_output_1{1, std::make_unique(1)}; - set_output_1.instantiate(ctx, info); - auto code = set_output_1.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + row_ir::code_sink sink; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*f32); + [[maybe_unused]] auto out0 = ctx.add_output(); + [[maybe_unused]] auto out1 = ctx.add_output(); + row_ir::node set_output_1{row_ir::output_reference{1}, + row_ir::node{row_ir::input_reference{1}}}; + set_output_1.instantiate(ctx); + set_output_1.emit_code(ctx, target_info, sink); auto expected_code = - R"***(float tmp_0 = in_1; -float tmp_1 = tmp_0; -*out_1 = tmp_1;)***"; + R"***(float tmp_1 = in_1; +float tmp_0 = tmp_1; +*out_1 = tmp_0; +)***"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } } @@ -98,39 +122,45 @@ TEST_F(RowIRCudaCodeGenTest, UnaryOperation) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = {{"in_0", {cudf::data_type{cudf::type_id::INT32}}}, - {"in_1", {cudf::data_type{cudf::type_id::DECIMAL32}}}}; - - row_ir::untyped_var_info outputs[] = {{"out_0"}, {"out_1"}}; - - row_ir::instance_info info{inputs, outputs}; - { - row_ir::instance_context ctx{}; - row_ir::operation op{row_ir::opcode::IDENTITY, - row_ir::operation::operands(row_ir::get_input(0))}; - op.instantiate(ctx, info); - auto code = op.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*f32); + + row_ir::code_sink sink; + row_ir::node op{ + row_ir::opcode::IDENTITY, std::nullopt, row_ir::node{row_ir::input_reference{0}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); auto expected_code = - R"***(int32_t tmp_0 = in_0; -int32_t tmp_1 = cudf::ast::detail::operator_functor{}(tmp_0);)***"; + R"***(int32_t tmp_1 = in_0; +int32_t tmp_0; +cudf::ops::identity(&tmp_0, &tmp_1); +)***"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } { - row_ir::instance_context ctx{}; - row_ir::operation op{row_ir::opcode::IDENTITY, - row_ir::operation::operands(row_ir::get_input(1))}; - op.instantiate(ctx, info); - auto null_code = op.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*d32); + row_ir::code_sink sink; + row_ir::node op{ + row_ir::opcode::IDENTITY, std::nullopt, row_ir::node{row_ir::input_reference{1}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); auto expected_null_code = - R"***(numeric::decimal32 tmp_0 = in_1; -numeric::decimal32 tmp_1 = cudf::ast::detail::operator_functor{}(tmp_0);)***"; + R"***(numeric::decimal32 tmp_1 = in_1; +numeric::decimal32 tmp_0; +cudf::ops::identity(&tmp_0, &tmp_1); +)***"; - EXPECT_EQ(null_code, expected_null_code); + EXPECT_EQ(sink.get_code(), expected_null_code); } } @@ -138,41 +168,50 @@ TEST_F(RowIRCudaCodeGenTest, BinaryOperation) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = {{"in_0", {cudf::data_type{cudf::type_id::INT32}}}, - {"in_1", {cudf::data_type{cudf::type_id::DECIMAL32}}}}; - - row_ir::untyped_var_info outputs[] = {{"out_0"}, {"out_1"}}; - - row_ir::instance_info info{inputs, outputs}; - { - row_ir::instance_context ctx{}; - row_ir::operation op{row_ir::opcode::ADD, - row_ir::operation::operands(row_ir::get_input(0), row_ir::get_input(0))}; - op.instantiate(ctx, info); - auto code = op.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*d32); + row_ir::code_sink sink; + row_ir::node op{row_ir::opcode::ADD, + std::nullopt, + row_ir::node{row_ir::input_reference{0}}, + row_ir::node{row_ir::input_reference{0}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); auto expected_code = - R"***(int32_t tmp_0 = in_0; -int32_t tmp_1 = in_0; -int32_t tmp_2 = cudf::ast::detail::operator_functor{}(tmp_0, tmp_1);)***"; + R"***(int32_t tmp_1 = in_0; +int32_t tmp_2 = in_0; +int32_t tmp_0; +cudf::ops::add(&tmp_0, &tmp_1, &tmp_2); +)***"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } { - row_ir::instance_context ctx{}; - row_ir::operation op{row_ir::opcode::ADD, - row_ir::operation::operands(row_ir::get_input(1), row_ir::get_input(1))}; - op.instantiate(ctx, info); - auto null_code = op.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*d32); + row_ir::code_sink sink; + row_ir::node op{row_ir::opcode::ADD, + std::nullopt, + row_ir::node{row_ir::input_reference{1}}, + row_ir::node{row_ir::input_reference{1}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); auto expected_null_code = - R"***(numeric::decimal32 tmp_0 = in_1; -numeric::decimal32 tmp_1 = in_1; -numeric::decimal32 tmp_2 = cudf::ast::detail::operator_functor{}(tmp_0, tmp_1);)***"; + R"***(numeric::decimal32 tmp_1 = in_1; +numeric::decimal32 tmp_2 = in_1; +numeric::decimal32 tmp_0; +cudf::ops::add(&tmp_0, &tmp_1, &tmp_2); +)***"; - EXPECT_EQ(null_code, expected_null_code); + EXPECT_EQ(sink.get_code(), expected_null_code); } } @@ -180,59 +219,57 @@ TEST_F(RowIRCudaCodeGenTest, VectorLengthOperation) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = { - {"in_0", {cudf::data_type{cudf::type_id::FLOAT64}}}, - {"in_1", {cudf::data_type{cudf::type_id::FLOAT64}}}, - {"in_2", {cudf::data_type{cudf::type_id::FLOAT64}}}, - {"in_3", {cudf::data_type{cudf::type_id::FLOAT64}}}, - }; - - row_ir::untyped_var_info outputs[] = {{"out_0"}, {"out_1"}}; - - row_ir::instance_info info{inputs, outputs}; - auto length_operation = [&](int32_t input0, int32_t input1, int32_t output) { // This function generates the IR for the vector length operation: // length(v) = sqrt(x^2 + y^2) // where v = (x, y) and v is a 2D vector. - auto x2 = std::make_unique( - row_ir::opcode::MUL, - row_ir::operation::operands(row_ir::get_input(input0), row_ir::get_input(input0))); + auto x2 = row_ir::node(row_ir::opcode::MUL, + std::nullopt, + row_ir::node{row_ir::input_reference{input0}}, + row_ir::node{row_ir::input_reference{input0}}); - auto y2 = std::make_unique( - row_ir::opcode::MUL, - row_ir::operation::operands(row_ir::get_input(input1), row_ir::get_input(input1))); + auto y2 = row_ir::node(row_ir::opcode::MUL, + std::nullopt, + row_ir::node{row_ir::input_reference{input1}}, + row_ir::node{row_ir::input_reference{input1}}); - auto sum = std::make_unique( - row_ir::opcode::ADD, row_ir::operation::operands(std::move(x2), std::move(y2))); + auto sum = row_ir::node(row_ir::opcode::ADD, std::nullopt, std::move(x2), std::move(y2)); - auto length = std::make_unique(row_ir::opcode::SQRT, - row_ir::operation::operands(std::move(sum))); + auto length = row_ir::node(row_ir::opcode::SQRT, std::nullopt, std::move(sum)); - return std::make_unique(output, std::move(length)); + return row_ir::node(row_ir::output_reference{0}, std::move(length)); }; { - row_ir::instance_context ctx{}; + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*f64); + [[maybe_unused]] auto in1 = ctx.add_input(*f64); + [[maybe_unused]] auto out0 = ctx.add_output(); + row_ir::code_sink sink; auto expr_ir = length_operation(0, 1, 0); - expr_ir->instantiate(ctx, info); - - auto code = expr_ir->generate_code(ctx, target_info, info); + expr_ir.instantiate(ctx); + expr_ir.emit_code(ctx, target_info, sink); auto expected_code = - R"***(double tmp_0 = in_0; -double tmp_1 = in_0; -double tmp_2 = cudf::ast::detail::operator_functor{}(tmp_0, tmp_1); -double tmp_3 = in_1; -double tmp_4 = in_1; -double tmp_5 = cudf::ast::detail::operator_functor{}(tmp_3, tmp_4); -double tmp_6 = cudf::ast::detail::operator_functor{}(tmp_2, tmp_5); -double tmp_7 = cudf::ast::detail::operator_functor{}(tmp_6); -double tmp_8 = tmp_7; -*out_0 = tmp_8;)***"; - - EXPECT_EQ(code, expected_code); + R"***(double tmp_4 = in_0; +double tmp_5 = in_0; +double tmp_3; +cudf::ops::mul(&tmp_3, &tmp_4, &tmp_5); +double tmp_7 = in_1; +double tmp_8 = in_1; +double tmp_6; +cudf::ops::mul(&tmp_6, &tmp_7, &tmp_8); +double tmp_2; +cudf::ops::add(&tmp_2, &tmp_3, &tmp_6); +double tmp_1; +cudf::ops::sqrt(&tmp_1, &tmp_2); +double tmp_0 = tmp_1; +*out_0 = tmp_0; +)***"; + + EXPECT_EQ(sink.get_code(), expected_code); } } @@ -252,12 +289,12 @@ TEST_F(RowIRCudaCodeGenTest, AstConversionBasic) auto expected = cudf::test::fixed_width_column_wrapper(expected_iter, expected_iter + column->size()); - row_ir::ast_args args{.table = cudf::table_view{{column->view()}}}; - auto transform_args = row_ir::ast_converter::compute_column(row_ir::target::CUDA, add_op, - args, + cudf::table_view{{*column}}, + cudf::table_view{}, + "expression", cudf::get_default_stream(), cudf::get_current_device_resource_ref()); @@ -265,8 +302,9 @@ TEST_F(RowIRCudaCodeGenTest, AstConversionBasic) ASSERT_EQ(transform_args.scalar_columns[0]->view().size(), 1); EXPECT_EQ(transform_args.source_type, cudf::udf_source_type::CUDA); EXPECT_EQ(transform_args.is_null_aware, cudf::null_aware::NO); - EXPECT_EQ(transform_args.null_policy, cudf::output_nullability::ALL_VALID); - EXPECT_EQ(transform_args.output_type, cudf::data_type{cudf::type_id::INT32}); + EXPECT_EQ(transform_args.outputs.size(), 1); + EXPECT_EQ(transform_args.outputs[0].nullability, cudf::output_nullability::PRESERVE); + EXPECT_EQ(transform_args.outputs[0].type, cudf::data_type{cudf::type_id::INT32}); ASSERT_EQ(transform_args.inputs.size(), 2); /// The first input should be a scalar value of 42 @@ -282,66 +320,127 @@ TEST_F(RowIRCudaCodeGenTest, AstConversionBasic) EXPECT_EQ(std::get(transform_args.inputs[1]).null_count(), column->null_count()); - auto expected_udf = R"***( -__device__ void expression(int32_t* out_0, int32_t in_0, int32_t in_1) + auto expected_udf = + R"***(__device__ cudf::ops::errc expression(int32_t* out_0, int32_t in_0, int32_t in_1) { -int32_t tmp_0 = in_0; -int32_t tmp_1 = in_1; -int32_t tmp_2 = cudf::ast::detail::operator_functor{}(tmp_0, tmp_1); -int32_t tmp_3 = tmp_2; -*out_0 = tmp_3; - -return; -} -)***"; +int32_t tmp_2 = in_0; +int32_t tmp_3 = in_1; +int32_t tmp_1; +cudf::ops::add(&tmp_1, &tmp_2, &tmp_3); +int32_t tmp_0 = tmp_1; +*out_0 = tmp_0; +return cudf::ops::errc::OK; +})***"; EXPECT_EQ(transform_args.udf, expected_udf); - auto result = cudf::transform_extended(transform_args.inputs, - transform_args.udf, - transform_args.output_type, - transform_args.source_type, - transform_args.user_data, - transform_args.is_null_aware, - transform_args.row_size, - transform_args.null_policy); - - CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); + auto result = cudf::multi_transform(transform_args.udf, + transform_args.source_type, + transform_args.is_null_aware, + transform_args.user_data, + transform_args.inputs, + transform_args.outputs, + std::move(transform_args.string_offsets), + transform_args.row_size, + transform_args.error_mode); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->get_column(0).view()); } TEST_F(RowIRCudaCodeGenTest, FilterPredicate) { row_ir::target_info target_info{row_ir::target::CUDA}; - row_ir::var_info inputs[] = {{"in_0", {cudf::data_type{cudf::type_id::BOOL8}}}}; + { + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*b8); + row_ir::code_sink sink; + row_ir::node filter_predicate( + row_ir::opcode::PREDICATE, std::nullopt, row_ir::node{row_ir::input_reference{0}}); + filter_predicate.instantiate(ctx); + filter_predicate.emit_code(ctx, target_info, sink); + + auto expected_code = R"***(bool tmp_1 = in_0; +bool tmp_0; +cudf::ops::predicate(&tmp_0, &tmp_1); +)***"; - row_ir::instance_info info{inputs, {}}; + EXPECT_EQ(sink.get_code(), expected_code); + } { - row_ir::instance_context ctx{}; - row_ir::filter_predicate filter_predicate(std::make_unique(0)); - filter_predicate.instantiate(ctx, info); - auto code = filter_predicate.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*b8); + row_ir::code_sink sink; + row_ir::node filter_predicate( + row_ir::opcode::PREDICATE, std::nullopt, row_ir::node{row_ir::input_reference{0}}); + ctx.set_has_nulls(true); + filter_predicate.instantiate(ctx); + filter_predicate.emit_code(ctx, target_info, sink); - auto expected_code = R"***(bool tmp_0 = in_0; -bool tmp_1 = cudf::ast::detail::flatten_predicate(tmp_0); + auto expected_code = R"***(cuda::std::optional tmp_1 = in_0; +cuda::std::optional tmp_0; +cudf::ops::predicate(&tmp_0, &tmp_1); )***"; - EXPECT_EQ(code, expected_code); + EXPECT_EQ(sink.get_code(), expected_code); } +} + +TEST_F(RowIRCudaCodeGenTest, FallibleBinaryOperation) +{ + row_ir::target_info target_info{row_ir::target::CUDA}; { - row_ir::instance_context ctx{}; - row_ir::filter_predicate filter_predicate(std::make_unique(0)); - ctx.set_has_nulls(true); - filter_predicate.instantiate(ctx, info); - auto null_code = filter_predicate.generate_code(ctx, target_info, info); + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*d32); + row_ir::code_sink sink; + row_ir::node op{row_ir::opcode::ANSI_ADD, + std::nullopt, + row_ir::node{row_ir::input_reference{0}}, + row_ir::node{row_ir::input_reference{0}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); + + auto expected_code = + R"***(int32_t tmp_1 = in_0; +int32_t tmp_2 = in_0; +int32_t tmp_0; +if(cudf::ops::errc e = cudf::ops::ansi_add(&tmp_0, &tmp_1, &tmp_2); e != cudf::ops::errc::OK) { +return e; +} +)***"; - auto expected_code = R"***(cuda::std::optional tmp_0 = in_0; -bool tmp_1 = cudf::ast::detail::flatten_predicate(tmp_0); + EXPECT_EQ(sink.get_code(), expected_code); + } + + { + row_ir::instance_context ctx{cudf::get_default_stream(), + cudf::get_current_device_resource_ref()}; + [[maybe_unused]] auto in0 = ctx.add_input(*i32); + [[maybe_unused]] auto in1 = ctx.add_input(*d32); + row_ir::code_sink sink; + row_ir::node op{row_ir::opcode::ANSI_ADD, + std::nullopt, + row_ir::node{row_ir::input_reference{1}}, + row_ir::node{row_ir::input_reference{1}}}; + op.instantiate(ctx); + op.emit_code(ctx, target_info, sink); + + auto expected_null_code = + R"***(numeric::decimal32 tmp_1 = in_1; +numeric::decimal32 tmp_2 = in_1; +numeric::decimal32 tmp_0; +if(cudf::ops::errc e = cudf::ops::ansi_add(&tmp_0, &tmp_1, &tmp_2); e != cudf::ops::errc::OK) { +return e; +} )***"; - EXPECT_EQ(null_code, expected_code); + EXPECT_EQ(sink.get_code(), expected_null_code); } } diff --git a/cpp/tests/transform/integration/unary_transform_test.cpp b/cpp/tests/transform/integration/unary_transform_test.cpp index 9c342a8f06c7..cdbc50aa8dac 100644 --- a/cpp/tests/transform/integration/unary_transform_test.cpp +++ b/cpp/tests/transform/integration/unary_transform_test.cpp @@ -200,6 +200,62 @@ __device__ inline void fdsf ( test_udf(ptx, op, data_init, 0, cudf::udf_source_type::PTX); } +TEST_F(UnaryOperationIntegrationTest, Transform_ErrorHandling) +{ + // c = a*a*a*a + std::string const cuda = + R"***( +__device__ cudf::ops::errc expression ( + int* C, + int a, + int b +) +{ + return cudf::ops::ansi_div(C, &a, &b); +} +)***"; + + cudf::test::fixed_width_column_wrapper a{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + cudf::test::fixed_width_column_wrapper b_fail{1, 2, 3, 4, 5, 6, 0, 8, 9, 10}; + cudf::test::fixed_width_column_wrapper b{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}; + cudf::test::fixed_width_column_wrapper expected{1, 1, 1, 1, 1, 1, 1, 1, 1, 1}; + + { + cudf::transform_input inputs[] = {a, b}; + cudf::transform_output outputs[] = { + {cudf::data_type(cudf::type_id::INT32), cudf::output_nullability::ALL_VALID}}; + std::unique_ptr result; + EXPECT_NO_THROW(result = cudf::multi_transform(cuda, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::ops::error_mode::ANY_ROW)); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected); + } + + { + cudf::transform_input inputs[] = {a, b_fail}; + cudf::transform_output outputs[] = { + {cudf::data_type(cudf::type_id::INT32), cudf::output_nullability::ALL_VALID}}; + std::unique_ptr result; + EXPECT_THROW(result = cudf::multi_transform(cuda, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt, + cudf::ops::error_mode::ANY_ROW), + std::overflow_error); + } +} + TEST_F(UnaryOperationIntegrationTest, Transform_INT32_INT32) { // c = a * a - a