|
|
# Licensed to the Apache Software Foundation (ASF) under one |
|
|
# or more contributor license agreements. See the NOTICE file |
|
|
# distributed with this work for additional information |
|
|
# regarding copyright ownership. The ASF licenses this file |
|
|
# to you under the Apache License, Version 2.0 (the |
|
|
# "License"); you may not use this file except in compliance |
|
|
# with the License. You may obtain a copy of the License at |
|
|
# |
|
|
# http: |
|
|
# |
|
|
# Unless required by applicable law or agreed to in writing, |
|
|
# software distributed under the License is distributed on an |
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
|
|
# KIND, either express or implied. See the License for the |
|
|
# specific language governing permissions and limitations |
|
|
# under the License. |
|
|
|
|
|
# distutils: language = c++ |
|
|
|
|
|
from pyarrow.includes.libarrow cimport * |
|
|
|
|
|
cdef extern from "arrow/gpu/cuda_api.h" namespace "arrow::cuda" nogil: |
|
|
|
|
|
cdef cppclass CCudaDeviceManager" arrow::cuda::CudaDeviceManager": |
|
|
@staticmethod |
|
|
CResult[CCudaDeviceManager*] Instance() |
|
|
CResult[shared_ptr[CCudaContext]] GetContext(int gpu_number) |
|
|
CResult[shared_ptr[CCudaContext]] GetSharedContext(int gpu_number, |
|
|
void* handle) |
|
|
CStatus AllocateHost(int device_number, int64_t nbytes, |
|
|
shared_ptr[CCudaHostBuffer]* buffer) |
|
|
int num_devices() const |
|
|
|
|
|
cdef cppclass CCudaContext" arrow::cuda::CudaContext": |
|
|
CResult[shared_ptr[CCudaBuffer]] Allocate(int64_t nbytes) |
|
|
CResult[shared_ptr[CCudaBuffer]] View(uint8_t* data, int64_t nbytes) |
|
|
CResult[shared_ptr[CCudaBuffer]] OpenIpcBuffer( |
|
|
const CCudaIpcMemHandle& ipc_handle) |
|
|
CStatus Synchronize() |
|
|
int64_t bytes_allocated() const |
|
|
const void* handle() const |
|
|
int device_number() const |
|
|
CResult[uintptr_t] GetDeviceAddress(uintptr_t addr) |
|
|
shared_ptr[CDevice] device() const |
|
|
shared_ptr[CMemoryManager] memory_manager() const |
|
|
|
|
|
cdef cppclass CCudaIpcMemHandle" arrow::cuda::CudaIpcMemHandle": |
|
|
@staticmethod |
|
|
CResult[shared_ptr[CCudaIpcMemHandle]] FromBuffer( |
|
|
const void* opaque_handle) |
|
|
CResult[shared_ptr[CBuffer]] Serialize(CMemoryPool* pool) const |
|
|
|
|
|
cdef cppclass CCudaBuffer" arrow::cuda::CudaBuffer"(CBuffer): |
|
|
CCudaBuffer(uint8_t* data, int64_t size, |
|
|
const shared_ptr[CCudaContext]& context, |
|
|
c_bool own_data=false, c_bool is_ipc=false) |
|
|
CCudaBuffer(const shared_ptr[CCudaBuffer]& parent, |
|
|
const int64_t offset, const int64_t size) |
|
|
|
|
|
@staticmethod |
|
|
CResult[shared_ptr[CCudaBuffer]] FromBuffer(shared_ptr[CBuffer] buf) |
|
|
|
|
|
CStatus CopyToHost(const int64_t position, const int64_t nbytes, |
|
|
void* out) const |
|
|
CStatus CopyFromHost(const int64_t position, const void* data, |
|
|
int64_t nbytes) |
|
|
CStatus CopyFromDevice(const int64_t position, const void* data, |
|
|
int64_t nbytes) |
|
|
CStatus CopyFromAnotherDevice(const shared_ptr[CCudaContext]& src_ctx, |
|
|
const int64_t position, const void* data, |
|
|
int64_t nbytes) |
|
|
CResult[shared_ptr[CCudaIpcMemHandle]] ExportForIpc() |
|
|
shared_ptr[CCudaContext] context() const |
|
|
|
|
|
cdef cppclass \ |
|
|
CCudaHostBuffer" arrow::cuda::CudaHostBuffer"(CMutableBuffer): |
|
|
pass |
|
|
|
|
|
cdef cppclass \ |
|
|
CCudaBufferReader" arrow::cuda::CudaBufferReader"(CBufferReader): |
|
|
CCudaBufferReader(const shared_ptr[CBuffer]& buffer) |
|
|
CResult[int64_t] Read(int64_t nbytes, void* buffer) |
|
|
CResult[shared_ptr[CBuffer]] Read(int64_t nbytes) |
|
|
|
|
|
cdef cppclass \ |
|
|
CCudaBufferWriter" arrow::cuda::CudaBufferWriter"(WritableFile): |
|
|
CCudaBufferWriter(const shared_ptr[CCudaBuffer]& buffer) |
|
|
CStatus Close() |
|
|
CStatus Write(const void* data, int64_t nbytes) |
|
|
CStatus WriteAt(int64_t position, const void* data, int64_t nbytes) |
|
|
CStatus SetBufferSize(const int64_t buffer_size) |
|
|
int64_t buffer_size() |
|
|
int64_t num_bytes_buffered() const |
|
|
|
|
|
CResult[shared_ptr[CCudaHostBuffer]] AllocateCudaHostBuffer( |
|
|
int device_number, const int64_t size) |
|
|
|
|
|
# Cuda prefix is added to avoid picking up arrow::cuda functions |
|
|
# from arrow namespace. |
|
|
CResult[shared_ptr[CCudaBuffer]] \ |
|
|
CudaSerializeRecordBatch" arrow::cuda::SerializeRecordBatch"\ |
|
|
(const CRecordBatch& batch, |
|
|
CCudaContext* ctx) |
|
|
CResult[shared_ptr[CRecordBatch]] \ |
|
|
CudaReadRecordBatch" arrow::cuda::ReadRecordBatch"\ |
|
|
(const shared_ptr[CSchema]& schema, |
|
|
CDictionaryMemo* dictionary_memo, |
|
|
const shared_ptr[CCudaBuffer]& buffer, |
|
|
CMemoryPool* pool) |
|
|
|
