mirror of
https://github.com/aykhans/AzSuicideDataVisualization.git
synced 2025-04-22 10:28:02 +00:00
108 lines
4.7 KiB
Cython
108 lines
4.7 KiB
Cython
# 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://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
# 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)
|
|
|
|
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)
|