paleolimbot commented on code in PR #205: URL: https://github.com/apache/arrow-nanoarrow/pull/205#discussion_r1235468628
########## extensions/nanoarrow_device/src/nanoarrow/nanoarrow_device_metal.cc: ########## @@ -0,0 +1,331 @@ + +// 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. + +#include <errno.h> +#include <string.h> +#include <unistd.h> + +#define NS_PRIVATE_IMPLEMENTATION +#define MTL_PRIVATE_IMPLEMENTATION +#include <Metal/Metal.hpp> + +#include "nanoarrow_device.hpp" + +#include "nanoarrow_device_metal.h" + +// If non-null, caller must ->release() the return value. This doesn't +// release the underlying memory (which must be managed separately). +static MTL::Buffer* ArrowDeviceMetalWrapBufferNonOwning(MTL::Device* mtl_device, + const void* arbitrary_addr, + int64_t size_bytes) { + // We can wrap any zero-size buffer + if (size_bytes == 0) { + return mtl_device->newBuffer(0, MTL::ResourceStorageModeShared); + } + + // Cache the page size from the system call + static int pagesize = 0; + if (pagesize == 0) { + pagesize = getpagesize(); + } + + int64_t allocation_size; + if (size_bytes % pagesize == 0) { + allocation_size = size_bytes; + } else { + allocation_size = (size_bytes / pagesize) + 1 * pagesize; + } + + // Will return nullptr if the memory is improperly aligned + return mtl_device->newBuffer(arbitrary_addr, allocation_size, + MTL::ResourceStorageModeShared, nullptr); +} + +static uint8_t* ArrowDeviceMetalAllocatorReallocate( + struct ArrowBufferAllocator* allocator, uint8_t* ptr, int64_t old_size, + int64_t new_size) { + // Cache the page size from the system call + static int pagesize = 0; + if (pagesize == 0) { + pagesize = getpagesize(); + } + + int64_t allocation_size; + if (new_size % pagesize == 0) { + allocation_size = new_size; + } else { + allocation_size = (new_size / pagesize) + 1 * pagesize; + } + + // If growing an existing buffer but the allocation size is still big enough, + // return the same pointer and do nothing. + if (ptr != nullptr && new_size >= old_size && new_size <= allocation_size) { + return ptr; + } + + int64_t copy_size; + if (new_size > old_size) { + copy_size = old_size; + } else { + copy_size = new_size; + } + + void* new_ptr = nullptr; + posix_memalign(&new_ptr, pagesize, allocation_size); + if (new_ptr != nullptr && ptr != nullptr) { + memcpy(new_ptr, ptr, copy_size); + } + + if (ptr != nullptr) { + free(ptr); + } + + return reinterpret_cast<uint8_t*>(new_ptr); +} + +static void ArrowDeviceMetalAllocatorFree(struct ArrowBufferAllocator* allocator, + uint8_t* ptr, int64_t old_size) { + free(ptr); +} + +void ArrowDeviceMetalInitBuffer(struct ArrowBuffer* buffer) { + buffer->allocator.reallocate = &ArrowDeviceMetalAllocatorReallocate; + buffer->allocator.free = &ArrowDeviceMetalAllocatorFree; + buffer->allocator.private_data = nullptr; + buffer->data = nullptr; Review Comment: For CUDA we definitely need this...for Metal as implemented here this is just `posix_memalign()` and `free()` regardless of the device identifier. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
