ceph/src/arrow/c_glib/test/test-cuda.rb

   1 # Licensed to the Apache Software Foundation (ASF) under one
   2 # or more contributor license agreements.  See the NOTICE file
   3 # distributed with this work for additional information
   4 # regarding copyright ownership.  The ASF licenses this file
   5 # to you under the Apache License, Version 2.0 (the
   6 # "License"); you may not use this file except in compliance
   7 # with the License.  You may obtain a copy of the License at
   8 #
   9 #   http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing,
  12 # software distributed under the License is distributed on an
  13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 # KIND, either express or implied.  See the License for the
  15 # specific language governing permissions and limitations
  16 # under the License.
  17
  18 class TestCUDA < Test::Unit::TestCase
  19   include Helper::Buildable
  20   include Helper::Omittable
  21
  22   def setup
  23     omit("Arrow CUDA is required") unless defined?(::ArrowCUDA)
  24     @manager = ArrowCUDA::DeviceManager.new
  25     omit("At least one GPU is required") if @manager.n_devices.zero?
  26     @context = @manager.get_context(0)
  27   end
  28
  29   sub_test_case("Context") do
  30     def test_allocated_size
  31       allocated_size_before = @context.allocated_size
  32       size = 128
  33       buffer = ArrowCUDA::Buffer.new(@context, size)
  34       assert_equal(size,
  35                    @context.allocated_size - allocated_size_before)
  36     end
  37   end
  38
  39   sub_test_case("Buffer") do
  40     def setup
  41       super
  42       @buffer = ArrowCUDA::Buffer.new(@context, 128)
  43     end
  44
  45     def test_copy
  46       @buffer.copy_from_host("Hello World")
  47       assert_equal("llo W", @buffer.copy_to_host(2, 5).to_s)
  48     end
  49
  50     def test_export
  51       require_gi_bindings(3, 3, 9)
  52       @buffer.copy_from_host("Hello World")
  53       handle = @buffer.export
  54       serialized_handle = handle.serialize.data
  55       Tempfile.open("arrow-cuda-export") do |output|
  56         pid = spawn(RbConfig.ruby, "-e", <<-SCRIPT)
  57 require "gi"
  58
  59 Gio = GI.load("Gio")
  60 Arrow = GI.load("Arrow")
  61 ArrowCUDA = GI.load("ArrowCUDA")
  62
  63 manager = ArrowCUDA::DeviceManager.new
  64 context = manager.get_context(0)
  65 serialized_handle = #{serialized_handle.to_s.dump}
  66 handle = ArrowCUDA::IPCMemoryHandle.new(serialized_handle)
  67 buffer = ArrowCUDA::Buffer.new(context, handle)
  68 File.open(#{output.path.dump}, "w") do |output|
  69   output.print(buffer.copy_to_host(0, 6).to_s)
  70 end
  71         SCRIPT
  72         Process.waitpid(pid)
  73         assert_equal("Hello ", output.read)
  74       end
  75     end
  76
  77     def test_context
  78       assert_equal(@context.allocated_size,
  79                    @buffer.context.allocated_size)
  80     end
  81
  82     sub_test_case("#read_record_batch") do
  83       def setup
  84         super
  85         @field = Arrow::Field.new("enabled", Arrow::BooleanDataType.new)
  86         @schema = Arrow::Schema.new([@field])
  87         @columns = [
  88           build_boolean_array([true]),
  89         ]
  90         @cpu_record_batch = Arrow::RecordBatch.new(@schema, 1, @columns)
  91
  92         @buffer = ArrowCUDA::Buffer.new(@context, @cpu_record_batch)
  93       end
  94
  95       def test_default
  96         gpu_record_batch = @buffer.read_record_batch(@schema)
  97         assert_equal(@cpu_record_batch.n_rows,
  98                      gpu_record_batch.n_rows)
  99       end
 100
 101       def test_options
 102         options = Arrow::ReadOptions.new
 103         gpu_record_batch = @buffer.read_record_batch(@schema, options)
 104         assert_equal(@cpu_record_batch.n_rows,
 105                      gpu_record_batch.n_rows)
 106       end
 107     end
 108   end
 109
 110   sub_test_case("HostBuffer") do
 111     def test_new
 112       buffer = ArrowCUDA::HostBuffer.new(0, 128)
 113       assert_equal(128, buffer.size)
 114     end
 115   end
 116
 117   sub_test_case("BufferInputStream") do
 118     def test_new
 119       buffer = ArrowCUDA::Buffer.new(@context, 128)
 120       buffer.copy_from_host("Hello World")
 121       stream = ArrowCUDA::BufferInputStream.new(buffer)
 122       begin
 123         assert_equal("Hello Worl", stream.read(5).copy_to_host(0, 10).to_s)
 124       ensure
 125         stream.close
 126       end
 127     end
 128   end
 129
 130   sub_test_case("BufferOutputStream") do
 131     def setup
 132       super
 133       @buffer = ArrowCUDA::Buffer.new(@context, 128)
 134       @buffer.copy_from_host("\x00" * @buffer.size)
 135       @stream = ArrowCUDA::BufferOutputStream.new(@buffer)
 136     end
 137
 138     def cleanup
 139       super
 140       @stream.close
 141     end
 142
 143     def test_new
 144       @stream.write("Hello World")
 145       assert_equal("Hello World", @buffer.copy_to_host(0, 11).to_s)
 146     end
 147
 148     def test_buffer
 149       assert_equal(0, @stream.buffer_size)
 150       @stream.buffer_size = 5
 151       assert_equal(5, @stream.buffer_size)
 152       @stream.write("Hell")
 153       assert_equal(4, @stream.buffered_size)
 154       assert_equal("\x00" * 5, @buffer.copy_to_host(0, 5).to_s)
 155       @stream.write("o")
 156       assert_equal("Hello", @buffer.copy_to_host(0, 5).to_s)
 157     end
 158   end
 159 end