Skip to content

Commit

Permalink
move tests to cuda_test.cc
Browse files Browse the repository at this point in the history
  • Loading branch information
@jorisvandenbossche @pitrou
jorisvandenbossche authored and pitrou committed Feb 28, 2024
1 parent 57f3e8d commit a38e318
Show file tree
Hide file tree
Showing 3 changed files with 117 additions and 132 deletions.
18 changes: 1 addition & 17 deletions cpp/src/arrow/c/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,23 +15,7 @@
# specific language governing permissions and limitations
# under the License.

if(ARROW_CUDA)
if(ARROW_BUILD_SHARED)
set(ARROW_CUDA_LIBRARY arrow_cuda_shared)
else()
set(ARROW_CUDA_LIBRARY arrow_cuda_static)
endif()

set(ARROW_C_TEST_LINK_LIBS ${ARROW_CUDA_LIBRARY} ${ARROW_TEST_LINK_LIBS})
else()
set(ARROW_C_TEST_LINK_LIBS ${ARROW_TEST_LINK_LIBS})
endif()

add_arrow_test(bridge_test
PREFIX
"arrow-c"
STATIC_LINK_LIBS
${ARROW_C_TEST_LINK_LIBS})
add_arrow_test(bridge_test PREFIX "arrow-c")
add_arrow_test(dlpack_test)

add_arrow_benchmark(bridge_benchmark)
Expand Down
115 changes: 0 additions & 115 deletions cpp/src/arrow/c/bridge_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,10 +51,6 @@
#include "arrow/compute/api_vector.h"
#endif

#ifdef ARROW_CUDA
#include "arrow/gpu/cuda_api.h"
#endif

namespace arrow {

using internal::ArrayExportGuard;
Expand Down Expand Up @@ -4334,117 +4330,6 @@ TEST_F(TestDeviceArrayRoundtrip, Struct) {
TestWithJSON(mm, type, R"([[4, null], null, [5, "foo"]])");
}

#ifdef ARROW_CUDA

class TestCUDADeviceArrayRoundtrip : public ::testing::Test {
public:
using ArrayFactory = std::function<Result<std::shared_ptr<Array>>()>;

static Result<std::shared_ptr<MemoryManager>> DeviceMapper(ArrowDeviceType type,
int64_t id) {
if (type != ARROW_DEVICE_CUDA) {
return Status::NotImplemented("should only be CUDA device");
}

ARROW_ASSIGN_OR_RAISE(auto manager, cuda::CudaDeviceManager::Instance());
ARROW_ASSIGN_OR_RAISE(auto device, manager->GetDevice(id));
return device->default_memory_manager();
}

static ArrayFactory JSONArrayFactory(std::shared_ptr<DataType> type, const char* json) {
return [=]() { return ArrayFromJSON(type, json); };
}

template <typename ArrayFactory>
void TestWithArrayFactory(ArrayFactory&& factory) {
TestWithArrayFactory(factory, factory);
}

template <typename ArrayFactory, typename ExpectedArrayFactory>
void TestWithArrayFactory(ArrayFactory&& factory,
ExpectedArrayFactory&& factory_expected) {
ASSERT_OK_AND_ASSIGN(auto manager, cuda::CudaDeviceManager::Instance());
ASSERT_OK_AND_ASSIGN(auto device, manager->GetDevice(0));
auto mm = device->default_memory_manager();

std::shared_ptr<Array> array;
std::shared_ptr<Array> device_array;
ASSERT_OK_AND_ASSIGN(array, factory());
ASSERT_OK_AND_ASSIGN(device_array, array->CopyTo(mm));

struct ArrowDeviceArray c_array {};
struct ArrowSchema c_schema {};
ArrayExportGuard array_guard(&c_array.array);
SchemaExportGuard schema_guard(&c_schema);

ASSERT_OK(ExportType(*device_array->type(), &c_schema));
std::shared_ptr<Device::SyncEvent> sync{nullptr};
ASSERT_OK(ExportDeviceArray(*device_array, sync, &c_array));

std::shared_ptr<Array> device_array_roundtripped;
ASSERT_OK_AND_ASSIGN(device_array_roundtripped,
ImportDeviceArray(&c_array, &c_schema, DeviceMapper));
ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema));
ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array));

// Check value of imported array (copy to CPU to assert equality)
std::shared_ptr<Array> array_roundtripped;
ASSERT_OK_AND_ASSIGN(array_roundtripped,
device_array_roundtripped->CopyTo(default_cpu_memory_manager()));
{
std::shared_ptr<Array> expected;
ASSERT_OK_AND_ASSIGN(expected, factory_expected());
AssertTypeEqual(*expected->type(), *array_roundtripped->type());
AssertArraysEqual(*expected, *array_roundtripped, true);
}

// Re-export and re-import, now both at once
ASSERT_OK(ExportDeviceArray(*device_array, sync, &c_array, &c_schema));
device_array_roundtripped.reset();
ASSERT_OK_AND_ASSIGN(device_array_roundtripped,
ImportDeviceArray(&c_array, &c_schema, DeviceMapper));
ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema));
ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array));

// Check value of imported array (copy to CPU to assert equality)
array_roundtripped.reset();
ASSERT_OK_AND_ASSIGN(array_roundtripped,
device_array_roundtripped->CopyTo(default_cpu_memory_manager()));
{
std::shared_ptr<Array> expected;
ASSERT_OK_AND_ASSIGN(expected, factory_expected());
AssertTypeEqual(*expected->type(), *array_roundtripped->type());
AssertArraysEqual(*expected, *array_roundtripped, true);
}
}

void TestWithJSON(std::shared_ptr<DataType> type, const char* json) {
TestWithArrayFactory(JSONArrayFactory(type, json));
}
};

TEST_F(TestCUDADeviceArrayRoundtrip, Primitive) { TestWithJSON(int32(), "[4, 5, null]"); }

TEST_F(TestCUDADeviceArrayRoundtrip, Struct) {
auto type = struct_({field("ints", int16()), field("strs", utf8())});

TestWithJSON(type, "[]");
TestWithJSON(type, R"([[4, "foo"], [5, "bar"]])");
TestWithJSON(type, R"([[4, null], null, [5, "foo"]])");
}

TEST_F(TestCUDADeviceArrayRoundtrip, Dictionary) {
auto factory = []() {
auto values = ArrayFromJSON(utf8(), R"(["foo", "bar", "quux"])");
auto indices = ArrayFromJSON(uint16(), "[0, 2, 1, null, 1]");
return DictionaryArray::FromArrays(dictionary(indices->type(), values->type()),
indices, values);
};
TestWithArrayFactory(factory);
}

#endif

////////////////////////////////////////////////////////////////////////////
// Array stream export tests

Expand Down
116 changes: 116 additions & 0 deletions cpp/src/arrow/gpu/cuda_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,8 @@

#include "gtest/gtest.h"

#include "arrow/c/bridge.h"
#include "arrow/c/util_internal.h"
#include "arrow/io/memory.h"
#include "arrow/ipc/api.h"
#include "arrow/ipc/dictionary.h"
Expand All @@ -38,8 +40,11 @@

namespace arrow {

using internal::ArrayExportGuard;
using internal::ArrayStreamExportGuard;
using internal::checked_cast;
using internal::checked_pointer_cast;
using internal::SchemaExportGuard;

namespace cuda {

Expand Down Expand Up @@ -703,5 +708,116 @@ TEST_F(TestCudaArrowIpc, DictionaryWriteRead) {
CompareBatch(*batch, *cpu_batch);
}

// ------------------------------------------------------------------------
// Test C Device Interface export/import with CUDA
// (equivalent tests for non-CUDA live in bridge_test.cc)

class TestCudaDeviceArrayRoundtrip : public ::testing::Test {
public:
using ArrayFactory = std::function<Result<std::shared_ptr<Array>>()>;

static Result<std::shared_ptr<MemoryManager>> DeviceMapper(ArrowDeviceType type,
int64_t id) {
if (type != ARROW_DEVICE_CUDA) {
return Status::NotImplemented("should only be CUDA device");
}

ARROW_ASSIGN_OR_RAISE(auto manager, cuda::CudaDeviceManager::Instance());
ARROW_ASSIGN_OR_RAISE(auto device, manager->GetDevice(id));
return device->default_memory_manager();
}

static ArrayFactory JSONArrayFactory(std::shared_ptr<DataType> type, const char* json) {
return [=]() { return ArrayFromJSON(type, json); };
}

template <typename ArrayFactory>
void TestWithArrayFactory(ArrayFactory&& factory) {
TestWithArrayFactory(factory, factory);
}

template <typename ArrayFactory, typename ExpectedArrayFactory>
void TestWithArrayFactory(ArrayFactory&& factory,
ExpectedArrayFactory&& factory_expected) {
ASSERT_OK_AND_ASSIGN(auto manager, cuda::CudaDeviceManager::Instance());
ASSERT_OK_AND_ASSIGN(auto device, manager->GetDevice(0));
auto mm = device->default_memory_manager();

std::shared_ptr<Array> array;
std::shared_ptr<Array> device_array;
ASSERT_OK_AND_ASSIGN(array, factory());
ASSERT_OK_AND_ASSIGN(device_array, array->CopyTo(mm));

struct ArrowDeviceArray c_array {};
struct ArrowSchema c_schema {};
ArrayExportGuard array_guard(&c_array.array);
SchemaExportGuard schema_guard(&c_schema);

ASSERT_OK(ExportType(*device_array->type(), &c_schema));
std::shared_ptr<Device::SyncEvent> sync{nullptr};
ASSERT_OK(ExportDeviceArray(*device_array, sync, &c_array));

std::shared_ptr<Array> device_array_roundtripped;
ASSERT_OK_AND_ASSIGN(device_array_roundtripped,
ImportDeviceArray(&c_array, &c_schema, DeviceMapper));
ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema));
ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array));

// Check value of imported array (copy to CPU to assert equality)
std::shared_ptr<Array> array_roundtripped;
ASSERT_OK_AND_ASSIGN(array_roundtripped,
device_array_roundtripped->CopyTo(default_cpu_memory_manager()));
{
std::shared_ptr<Array> expected;
ASSERT_OK_AND_ASSIGN(expected, factory_expected());
AssertTypeEqual(*expected->type(), *array_roundtripped->type());
AssertArraysEqual(*expected, *array_roundtripped, true);
}

// Re-export and re-import, now both at once
ASSERT_OK(ExportDeviceArray(*device_array, sync, &c_array, &c_schema));
device_array_roundtripped.reset();
ASSERT_OK_AND_ASSIGN(device_array_roundtripped,
ImportDeviceArray(&c_array, &c_schema, DeviceMapper));
ASSERT_TRUE(ArrowSchemaIsReleased(&c_schema));
ASSERT_TRUE(ArrowArrayIsReleased(&c_array.array));

// Check value of imported array (copy to CPU to assert equality)
array_roundtripped.reset();
ASSERT_OK_AND_ASSIGN(array_roundtripped,
device_array_roundtripped->CopyTo(default_cpu_memory_manager()));
{
std::shared_ptr<Array> expected;
ASSERT_OK_AND_ASSIGN(expected, factory_expected());
AssertTypeEqual(*expected->type(), *array_roundtripped->type());
AssertArraysEqual(*expected, *array_roundtripped, true);
}
}

void TestWithJSON(std::shared_ptr<DataType> type, const char* json) {
TestWithArrayFactory(JSONArrayFactory(type, json));
}
};

TEST_F(TestCudaDeviceArrayRoundtrip, Primitive) { TestWithJSON(int32(), "[4, 5, null]"); }

TEST_F(TestCudaDeviceArrayRoundtrip, Struct) {
auto type = struct_({field("ints", int16()), field("strs", utf8())});

TestWithJSON(type, "[]");
TestWithJSON(type, R"([[4, "foo"], [5, "bar"]])");
TestWithJSON(type, R"([[4, null], null, [5, "foo"]])");
}

TEST_F(TestCudaDeviceArrayRoundtrip, Dictionary) {
auto factory = []() {
auto values = ArrayFromJSON(utf8(), R"(["foo", "bar", "quux"])");
auto indices = ArrayFromJSON(uint16(), "[0, 2, 1, null, 1]");
return DictionaryArray::FromArrays(dictionary(indices->type(), values->type()),
indices, values);
};
TestWithArrayFactory(factory);
}

} // namespace cuda
} // namespace arrow

0 comments on commit a38e318

Please sign in to comment.