Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Migrate binary operations to pylibcudf #14821

Merged
merged 10 commits into from
Jan 23, 2024
Merged

Migrate binary operations to pylibcudf #14821

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
ed9871e
Switch to using an enum class for binary operations
vyasr Jan 17, 2024
47ffa4a
Add an empty pyx file to generate the Python enum
vyasr Jan 17, 2024
b2825a7
Add binary op functionality to pylibcudf
vyasr Jan 17, 2024
d048040
Ignore warning for now
vyasr Jan 17, 2024
332c765
Switch to using pylibcudf for column binops
vyasr Jan 17, 2024
65a2810
Remove now unused jit binaryop code path
vyasr Jan 17, 2024
a948a55
Add docs for binaryop
vyasr Jan 19, 2024
bcb97f1
Remove jit binary op bindings
vyasr Jan 22, 2024
e675cd1
Remove pytest.ini change
vyasr Jan 22, 2024
b7209b3
Remove unnecessary upper
vyasr Jan 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions docs/cudf/source/user_guide/api_docs/pylibcudf/binaryop.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
========
binaryop
========

.. automodule:: cudf._lib.pylibcudf.binaryop
:members:
1 change: 1 addition & 0 deletions docs/cudf/source/user_guide/api_docs/pylibcudf/index.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ This page provides API documentation for pylibcudf.
:maxdepth: 1
:caption: API Documentation

binaryop
column
copying
gpumemoryview
Expand Down
5 changes: 0 additions & 5 deletions python/cudf/cudf/_lib/binaryop.pxd
View file
Open in desktop

This file was deleted.

261 changes: 41 additions & 220 deletions python/cudf/cudf/_lib/binaryop.pyx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,160 +1,30 @@
# Copyright (c) 2020-2022, NVIDIA CORPORATION.
# Copyright (c) 2020-2024, NVIDIA CORPORATION.

from enum import IntEnum

from libcpp.memory cimport unique_ptr
from libcpp.string cimport string
from libcpp.utility cimport move

from cudf._lib.binaryop cimport underlying_type_t_binary_operator
from cudf._lib.column cimport Column

from cudf._lib.scalar import as_device_scalar

from cudf._lib.scalar cimport DeviceScalar
from cudf._lib.types cimport dtype_to_pylibcudf_type

from cudf._lib.types import SUPPORTED_NUMPY_TO_LIBCUDF_TYPES

from cudf._lib.cpp.column.column cimport column
from cudf._lib.cpp.column.column_view cimport column_view
from cudf._lib.cpp.scalar.scalar cimport scalar
from cudf._lib.cpp.types cimport data_type, type_id
from cudf._lib.types cimport dtype_to_data_type, underlying_type_t_type_id

from cudf.api.types import is_scalar
from cudf._lib import pylibcudf
from cudf._lib.scalar import as_device_scalar
from cudf.core.buffer import acquire_spill_lock

cimport cudf._lib.cpp.binaryop as cpp_binaryop
from cudf._lib.cpp.binaryop cimport binary_operator
import cudf


class BinaryOperation(IntEnum):
ADD = (
<underlying_type_t_binary_operator> binary_operator.ADD
)
SUB = (
<underlying_type_t_binary_operator> binary_operator.SUB
)
MUL = (
<underlying_type_t_binary_operator> binary_operator.MUL
)
DIV = (
<underlying_type_t_binary_operator> binary_operator.DIV
)
TRUEDIV = (
<underlying_type_t_binary_operator> binary_operator.TRUE_DIV
)
FLOORDIV = (
<underlying_type_t_binary_operator> binary_operator.FLOOR_DIV
)
MOD = (
<underlying_type_t_binary_operator> binary_operator.PYMOD
)
POW = (
<underlying_type_t_binary_operator> binary_operator.POW
)
INT_POW = (
<underlying_type_t_binary_operator> binary_operator.INT_POW
)
EQ = (
<underlying_type_t_binary_operator> binary_operator.EQUAL
)
NE = (
<underlying_type_t_binary_operator> binary_operator.NOT_EQUAL
)
LT = (
<underlying_type_t_binary_operator> binary_operator.LESS
)
GT = (
<underlying_type_t_binary_operator> binary_operator.GREATER
)
LE = (
<underlying_type_t_binary_operator> binary_operator.LESS_EQUAL
)
GE = (
<underlying_type_t_binary_operator> binary_operator.GREATER_EQUAL
)
AND = (
<underlying_type_t_binary_operator> binary_operator.BITWISE_AND
)
OR = (
<underlying_type_t_binary_operator> binary_operator.BITWISE_OR
)
XOR = (
<underlying_type_t_binary_operator> binary_operator.BITWISE_XOR
)
L_AND = (
<underlying_type_t_binary_operator> binary_operator.LOGICAL_AND
)
L_OR = (
<underlying_type_t_binary_operator> binary_operator.LOGICAL_OR
)
GENERIC_BINARY = (
<underlying_type_t_binary_operator> binary_operator.GENERIC_BINARY
)
NULL_EQUALS = (
<underlying_type_t_binary_operator> binary_operator.NULL_EQUALS
)


cdef binaryop_v_v(Column lhs, Column rhs,
binary_operator c_op, data_type c_dtype):
cdef column_view c_lhs = lhs.view()
cdef column_view c_rhs = rhs.view()

cdef unique_ptr[column] c_result

with nogil:
c_result = move(
cpp_binaryop.binary_operation(
c_lhs,
c_rhs,
c_op,
c_dtype
)
)

return Column.from_unique_ptr(move(c_result))


cdef binaryop_v_s(Column lhs, DeviceScalar rhs,
binary_operator c_op, data_type c_dtype):
cdef column_view c_lhs = lhs.view()
cdef const scalar* c_rhs = rhs.get_raw_ptr()

cdef unique_ptr[column] c_result

with nogil:
c_result = move(
cpp_binaryop.binary_operation(
c_lhs,
c_rhs[0],
c_op,
c_dtype
)
)

return Column.from_unique_ptr(move(c_result))

cdef binaryop_s_v(DeviceScalar lhs, Column rhs,
binary_operator c_op, data_type c_dtype):
cdef const scalar* c_lhs = lhs.get_raw_ptr()
cdef column_view c_rhs = rhs.view()

cdef unique_ptr[column] c_result

with nogil:
c_result = move(
cpp_binaryop.binary_operation(
c_lhs[0],
c_rhs,
c_op,
c_dtype
)
)

return Column.from_unique_ptr(move(c_result))
# Map pandas operation names to pylibcudf operation names.
_op_map = {
"TRUEDIV": "TRUE_DIV",
"FLOORDIV": "FLOOR_DIV",
"MOD": "PYMOD",
"EQ": "EQUAL",
"NE": "NOT_EQUAL",
"LT": "LESS",
"GT": "GREATER",
"LE": "LESS_EQUAL",
"GE": "GREATER_EQUAL",
"AND": "BITWISE_AND",
"OR": "BITWISE_OR",
"XOR": "BITWISE_XOR",
"L_AND": "LOGICAL_AND",
"L_OR": "LOGICAL_OR",
}


@acquire_spill_lock()
Expand All @@ -166,74 +36,25 @@ def binaryop(lhs, rhs, op, dtype):
# pipeline for libcudf binops that don't map to Python binops.
if op not in {"INT_POW", "NULL_EQUALS"}:
op = op[2:-2]

op = BinaryOperation[op.upper()]
cdef binary_operator c_op = <binary_operator> (
<underlying_type_t_binary_operator> op
)

cdef data_type c_dtype = dtype_to_data_type(dtype)

if is_scalar(lhs) or lhs is None:
s_lhs = as_device_scalar(lhs, dtype=rhs.dtype if lhs is None else None)
result = binaryop_s_v(
s_lhs,
rhs,
c_op,
c_dtype
)

elif is_scalar(rhs) or rhs is None:
s_rhs = as_device_scalar(rhs, dtype=lhs.dtype if rhs is None else None)
result = binaryop_v_s(
lhs,
s_rhs,
c_op,
c_dtype
)

else:
result = binaryop_v_v(
lhs,
rhs,
c_op,
c_dtype
)
return result


@acquire_spill_lock()
def binaryop_udf(Column lhs, Column rhs, udf_ptx, dtype):
"""
Apply a user-defined binary operator (a UDF) defined in `udf_ptx` on
the two input columns `lhs` and `rhs`. The output type of the UDF
has to be specified in `dtype`, a numpy data type.
Currently ONLY int32, int64, float32 and float64 are supported.
"""
cdef column_view c_lhs = lhs.view()
cdef column_view c_rhs = rhs.view()

cdef type_id tid = (
<type_id> (
<underlying_type_t_type_id> (
SUPPORTED_NUMPY_TO_LIBCUDF_TYPES[cudf.dtype(dtype)]
)
op = op.upper()
op = _op_map.get(op, op)

return Column.from_pylibcudf(
# Check if the dtype args are desirable here.
pylibcudf.binaryop.binary_operation(
lhs.to_pylibcudf(mode="read") if isinstance(lhs, Column)
else (
<DeviceScalar> as_device_scalar(
lhs, dtype=rhs.dtype if lhs is None else None
)
).c_value,
rhs.to_pylibcudf(mode="read") if isinstance(rhs, Column)
else (
<DeviceScalar> as_device_scalar(
rhs, dtype=lhs.dtype if rhs is None else None
)
).c_value,
pylibcudf.binaryop.BinaryOperator[op],
dtype_to_pylibcudf_type(dtype),
)
)
cdef data_type c_dtype = data_type(tid)

cdef string cpp_str = udf_ptx.encode("UTF-8")

cdef unique_ptr[column] c_result

with nogil:
c_result = move(
cpp_binaryop.binary_operation(
c_lhs,
c_rhs,
cpp_str,
c_dtype
)
)

return Column.from_unique_ptr(move(c_result))
4 changes: 2 additions & 2 deletions python/cudf/cudf/_lib/cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# =============================================================================
# Copyright (c) 2022-2023, NVIDIA CORPORATION.
# Copyright (c) 2022-2024, NVIDIA CORPORATION.
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License. You may obtain a copy of the License at
Expand All @@ -12,7 +12,7 @@
# the License.
# =============================================================================

set(cython_sources copying.pyx types.pyx)
set(cython_sources binaryop.pyx copying.pyx types.pyx)

set(linked_libraries cudf::cudf)

Expand Down
Loading
Loading