import quincy 17.2.0

[ceph.git] / ceph / src / arrow / c_glib / test / test-cuda.rb

diff --git a/ceph/src/arrow/c_glib/test/test-cuda.rb b/ceph/src/arrow/c_glib/test/test-cuda.rb
new file mode 100644 (file)
index 0000000..c9b3492
--- /dev/null
+++ b/ceph/src/arrow/c_glib/test/test-cuda.rb
@@ -0,0 +1,159 @@
+# 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.
+
+class TestCUDA < Test::Unit::TestCase
+  include Helper::Buildable
+  include Helper::Omittable
+
+  def setup
+    omit("Arrow CUDA is required") unless defined?(::ArrowCUDA)
+    @manager = ArrowCUDA::DeviceManager.new
+    omit("At least one GPU is required") if @manager.n_devices.zero?
+    @context = @manager.get_context(0)
+  end
+
+  sub_test_case("Context") do
+    def test_allocated_size
+      allocated_size_before = @context.allocated_size
+      size = 128
+      buffer = ArrowCUDA::Buffer.new(@context, size)
+      assert_equal(size,
+                   @context.allocated_size - allocated_size_before)
+    end
+  end
+
+  sub_test_case("Buffer") do
+    def setup
+      super
+      @buffer = ArrowCUDA::Buffer.new(@context, 128)
+    end
+
+    def test_copy
+      @buffer.copy_from_host("Hello World")
+      assert_equal("llo W", @buffer.copy_to_host(2, 5).to_s)
+    end
+
+    def test_export
+      require_gi_bindings(3, 3, 9)
+      @buffer.copy_from_host("Hello World")
+      handle = @buffer.export
+      serialized_handle = handle.serialize.data
+      Tempfile.open("arrow-cuda-export") do |output|
+        pid = spawn(RbConfig.ruby, "-e", <<-SCRIPT)
+require "gi"
+
+Gio = GI.load("Gio")
+Arrow = GI.load("Arrow")
+ArrowCUDA = GI.load("ArrowCUDA")
+
+manager = ArrowCUDA::DeviceManager.new
+context = manager.get_context(0)
+serialized_handle = #{serialized_handle.to_s.dump}
+handle = ArrowCUDA::IPCMemoryHandle.new(serialized_handle)
+buffer = ArrowCUDA::Buffer.new(context, handle)
+File.open(#{output.path.dump}, "w") do |output|
+  output.print(buffer.copy_to_host(0, 6).to_s)
+end
+        SCRIPT
+        Process.waitpid(pid)
+        assert_equal("Hello ", output.read)
+      end
+    end
+
+    def test_context
+      assert_equal(@context.allocated_size,
+                   @buffer.context.allocated_size)
+    end
+
+    sub_test_case("#read_record_batch") do
+      def setup
+        super
+        @field = Arrow::Field.new("enabled", Arrow::BooleanDataType.new)
+        @schema = Arrow::Schema.new([@field])
+        @columns = [
+          build_boolean_array([true]),
+        ]
+        @cpu_record_batch = Arrow::RecordBatch.new(@schema, 1, @columns)
+
+        @buffer = ArrowCUDA::Buffer.new(@context, @cpu_record_batch)
+      end
+
+      def test_default
+        gpu_record_batch = @buffer.read_record_batch(@schema)
+        assert_equal(@cpu_record_batch.n_rows,
+                     gpu_record_batch.n_rows)
+      end
+
+      def test_options
+        options = Arrow::ReadOptions.new
+        gpu_record_batch = @buffer.read_record_batch(@schema, options)
+        assert_equal(@cpu_record_batch.n_rows,
+                     gpu_record_batch.n_rows)
+      end
+    end
+  end
+
+  sub_test_case("HostBuffer") do
+    def test_new
+      buffer = ArrowCUDA::HostBuffer.new(0, 128)
+      assert_equal(128, buffer.size)
+    end
+  end
+
+  sub_test_case("BufferInputStream") do
+    def test_new
+      buffer = ArrowCUDA::Buffer.new(@context, 128)
+      buffer.copy_from_host("Hello World")
+      stream = ArrowCUDA::BufferInputStream.new(buffer)
+      begin
+        assert_equal("Hello Worl", stream.read(5).copy_to_host(0, 10).to_s)
+      ensure
+        stream.close
+      end
+    end
+  end
+
+  sub_test_case("BufferOutputStream") do
+    def setup
+      super
+      @buffer = ArrowCUDA::Buffer.new(@context, 128)
+      @buffer.copy_from_host("\x00" * @buffer.size)
+      @stream = ArrowCUDA::BufferOutputStream.new(@buffer)
+    end
+
+    def cleanup
+      super
+      @stream.close
+    end
+
+    def test_new
+      @stream.write("Hello World")
+      assert_equal("Hello World", @buffer.copy_to_host(0, 11).to_s)
+    end
+
+    def test_buffer
+      assert_equal(0, @stream.buffer_size)
+      @stream.buffer_size = 5
+      assert_equal(5, @stream.buffer_size)
+      @stream.write("Hell")
+      assert_equal(4, @stream.buffered_size)
+      assert_equal("\x00" * 5, @buffer.copy_to_host(0, 5).to_s)
+      @stream.write("o")
+      assert_equal("Hello", @buffer.copy_to_host(0, 5).to_s)
+    end
+  end
+end