Support to let users create Array from raw device pointers

The example show cases the usage using CUDA API, but it should be
similar to OpenCL or CPU. In the case of OpenCL backend, the pointer
would be cl_mem

Author: 9prady9

Cargo.toml

@@ -15,6 +15,15 @@ exclude = [
     "arrayfire/*",
 ]
 
+[workspace]
+members = [
+    "cuda-interop",
+]
+
+[lib]
+name = "arrayfire"
+path = "src/lib.rs"
+
 [package.metadata.docs.rs]
 rustdoc-args = [ "--html-in-header", "./scripts/mathjax.script", ]
 
@@ -53,10 +62,6 @@ serde_derive = "1.0"
 serde = "1.0"
 rustc_version = "0.2"
 
-[lib]
-name = "arrayfire"
-path = "src/lib.rs"
-
 [[example]]
 name = "helloworld"
 path = "examples/helloworld.rs"

cuda-interop/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "af-cuda-interop"
+version = "0.1.0"
+authors = ["Pradeep Garigipati <pradeep@arrayfire.com>"]
+edition = "2018"
+
+[dependencies]
+libc = "0.2"
+arrayfire = { path = "../" }
+cuda-runtime-sys = "0.3.0-alpha.1"
+
+[dev-dependencies]
+rustacuda = "0.1"
+rustacuda_core = "0.1"
+
+[[example]]
+name = "custom_kernel"
+path = "examples/custom_kernel.rs"

cuda-interop/examples/cuda_af_app.rs

@@ -0,0 +1,61 @@
+use arrayfire::{af_print, dim4, info, set_device, Array};
+use rustacuda::prelude::*;
+use rustacuda::*;
+
+fn main() {
+    // MAKE SURE to do all rustacuda initilization before arrayfire API's
+    // first call. It seems like some CUDA context state is getting messed up
+    // if we mix CUDA context init(device, context, module, stream) with ArrayFire API
+    match rustacuda::init(CudaFlags::empty()) {
+        Ok(()) => {}
+        Err(e) => panic!("rustacuda init failure: {:?}", e),
+    }
+    let device = match Device::get_device(0) {
+        Ok(d) => d,
+        Err(e) => panic!("Failed to get device: {:?}", e),
+    };
+    let _context =
+        match Context::create_and_push(ContextFlags::MAP_HOST | ContextFlags::SCHED_AUTO, device) {
+            Ok(c) => c,
+            Err(e) => panic!("Failed to create context: {:?}", e),
+        };
+    let stream = match Stream::new(StreamFlags::NON_BLOCKING, None) {
+        Ok(s) => s,
+        Err(e) => panic!("Failed to create stream: {:?}", e),
+    };
+
+    let mut in_x = DeviceBuffer::from_slice(&[1.0f32; 10]).unwrap();
+    let mut in_y = DeviceBuffer::from_slice(&[2.0f32; 10]).unwrap();
+
+    // wait for any prior kernels to finish before passing
+    // the device pointers to ArrayFire
+    match stream.synchronize() {
+        Ok(()) => {}
+        Err(e) => panic!("Stream sync failure: {:?}", e),
+    };
+
+    set_device(0);
+    info();
+
+    let x = Array::new_from_device_ptr(in_x.as_device_ptr().as_raw_mut(), dim4!(10));
+    let y = Array::new_from_device_ptr(in_y.as_device_ptr().as_raw_mut(), dim4!(10));
+
+    // Lock so that ArrayFire doesn't free pointers from RustaCUDA
+    // But we have to make sure these pointers stay in valid scope
+    // as long as the associated ArrayFire Array objects are valid
+    x.lock();
+    y.lock();
+
+    af_print!("x", x);
+    af_print!("y", y);
+
+    let o = x + y;
+    af_print!("out", o);
+
+    let _o_dptr = unsafe { o.device_ptr() }; // Calls an implicit lock
+
+    // User has to call unlock if they want to relenquish control to ArrayFire
+
+    // Once the non-arrayfire operations are done, call unlock.
+    o.unlock(); // After this, there is no guarantee that value of o_dptr is valid
+}

cuda-interop/examples/custom_kernel.rs

@@ -0,0 +1,83 @@
+use arrayfire as af;
+use rustacuda::prelude::*;
+use rustacuda::*;
+
+use std::ffi::CString;
+
+fn main() {
+    // MAKE SURE to do all rustacuda initilization before arrayfire API's
+    // first call. It seems like some CUDA context state is getting messed up
+    // if we mix CUDA context init(device, context, module, stream) with ArrayFire API
+    match rustacuda::init(CudaFlags::empty()) {
+        Ok(()) => {}
+        Err(e) => panic!("rustacuda init failure: {:?}", e),
+    }
+    let device = match Device::get_device(0) {
+        Ok(d) => d,
+        Err(e) => panic!("Failed to get device: {:?}", e),
+    };
+    let _context =
+        match Context::create_and_push(ContextFlags::MAP_HOST | ContextFlags::SCHED_AUTO, device) {
+            Ok(c) => c,
+            Err(e) => panic!("Failed to create context: {:?}", e),
+        };
+    let ptx = CString::new(include_str!("./resources/add.ptx")).unwrap();
+    let module = match Module::load_from_string(&ptx) {
+        Ok(m) => m,
+        Err(e) => panic!("Failed to load module from string: {:?}", e),
+    };
+    let stream = match Stream::new(StreamFlags::NON_BLOCKING, None) {
+        Ok(s) => s,
+        Err(e) => panic!("Failed to create stream: {:?}", e),
+    };
+
+    af::set_device(0);
+    af::info();
+
+    let num: i32 = 10;
+    let x = af::constant(1f32, af::dim4!(10));
+    let y = af::constant(2f32, af::dim4!(10));
+    let out = af::constant(0f32, af::dim4!(10));
+
+    af::af_print!("x", x);
+    af::af_print!("y", y);
+    af::af_print!("out(init)", out);
+
+    //TODO Figure out how to use Stream returned by ArrayFire with Rustacuda
+    // let af_id = get_device();
+    // let cuda_id = get_device_native_id(af_id);
+    // let af_cuda_stream = get_stream(cuda_id);
+
+    //TODO Figure out how to use Stream returned by ArrayFire with Rustacuda
+    // let stream = Stream {inner: mem::transmute(af_cuda_stream)};
+
+    // Run a custom CUDA kernel in the ArrayFire CUDA stream
+    unsafe {
+        // Obtain device pointers from ArrayFire using Array::device() method
+        let d_x: *mut f32 = x.device_ptr() as *mut f32;
+        let d_y: *mut f32 = y.device_ptr() as *mut f32;
+        let d_o: *mut f32 = out.device_ptr() as *mut f32;
+
+        match launch!(module.sum<<<1, 1, 0, stream>>>(
+        memory::DevicePointer::wrap(d_x),
+        memory::DevicePointer::wrap(d_y),
+        memory::DevicePointer::wrap(d_o),
+        num
+        )) {
+            Ok(()) => {}
+            Err(e) => panic!("Kernel Launch failure: {:?}", e),
+        }
+
+        // wait for the kernel to finish as it is async call
+        match stream.synchronize() {
+            Ok(()) => {}
+            Err(e) => panic!("Stream sync failure: {:?}", e),
+        };
+
+        // Return control of Array memory to ArrayFire using unlock
+        x.unlock();
+        y.unlock();
+        out.unlock();
+    }
+    af::af_print!("sum after kernel launch", out);
+}

cuda-interop/examples/resources/add.cu

@@ -0,0 +1,7 @@
+extern "C" __constant__ int my_constant = 314;
+
+extern "C" __global__ void sum(const float* x, const float* y, float* out, int count) {
+    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x) {
+        out[i] = x[i] + y[i];
+    }
+}

cuda-interop/examples/resources/add.ptx

@@ -0,0 +1,61 @@
+//
+// Generated by NVIDIA NVVM Compiler
+//
+// Compiler Build ID: CL-24817639
+// Cuda compilation tools, release 10.0, V10.0.130
+// Based on LLVM 3.4svn
+//
+
+.version 3.2
+.target sm_20
+.address_size 64
+
+        // .globl       sum
+.const .align 4 .u32 my_constant = 314;
+
+.visible .entry sum(
+        .param .u64 sum_param_0,
+        .param .u64 sum_param_1,
+        .param .u64 sum_param_2,
+        .param .u32 sum_param_3
+)
+{
+        .reg .pred      %p<3>;
+        .reg .f32       %f<4>;
+        .reg .b32       %r<11>;
+        .reg .b64       %rd<11>;
+
+
+        ld.param.u64    %rd4, [sum_param_0];
+        ld.param.u64    %rd5, [sum_param_1];
+        ld.param.u64    %rd6, [sum_param_2];
+        ld.param.u32    %r6, [sum_param_3];
+        mov.u32         %r1, %ntid.x;
+        mov.u32         %r7, %ctaid.x;
+        mov.u32         %r8, %tid.x;
+        mad.lo.s32      %r10, %r1, %r7, %r8;
+        setp.ge.s32     %p1, %r10, %r6;
+        @%p1 bra        BB0_3;
+
+        cvta.to.global.u64      %rd1, %rd6;
+        cvta.to.global.u64      %rd2, %rd5;
+        cvta.to.global.u64      %rd3, %rd4;
+        mov.u32         %r9, %nctaid.x;
+        mul.lo.s32      %r3, %r9, %r1;
+
+BB0_2:
+        mul.wide.s32    %rd7, %r10, 4;
+        add.s64         %rd8, %rd3, %rd7;
+        add.s64         %rd9, %rd2, %rd7;
+        ld.global.f32   %f1, [%rd9];
+        ld.global.f32   %f2, [%rd8];
+        add.f32         %f3, %f2, %f1;
+        add.s64         %rd10, %rd1, %rd7;
+        st.global.f32   [%rd10], %f3;
+        add.s32         %r10, %r3, %r10;
+        setp.lt.s32     %p2, %r10, %r6;
+        @%p2 bra        BB0_2;
+
+BB0_3:
+        ret;
+}

cuda-interop/src/lib.rs

@@ -0,0 +1,60 @@
+//! af-cuda-interop package is to used only when the application intends to mix
+//! arrayfire code with raw CUDA code.
+
+use arrayfire::{handle_error_general, AfError};
+use cuda_runtime_sys::cudaStream_t;
+use libc::c_int;
+
+extern "C" {
+    fn afcu_get_native_id(native_id: *mut c_int, id: c_int) -> c_int;
+    fn afcu_set_native_id(native_id: c_int) -> c_int;
+    fn afcu_get_stream(out: *mut cudaStream_t, id: c_int) -> c_int;
+}
+
+/// Get active device's id in CUDA context
+///
+/// # Parameters
+///
+/// - `id` is the integer identifier of concerned CUDA device as per ArrayFire context
+///
+/// # Return Values
+///
+/// Integer identifier of device in CUDA context
+pub fn get_device_native_id(id: i32) -> i32 {
+    unsafe {
+        let mut temp: i32 = 0;
+        let err_val = afcu_get_native_id(&mut temp as *mut c_int, id);
+        handle_error_general(AfError::from(err_val));
+        temp
+    }
+}
+
+/// Set active device using CUDA context's id
+///
+/// # Parameters
+///
+/// - `id` is the identifier of GPU in CUDA context
+pub fn set_device_native_id(native_id: i32) {
+    unsafe {
+        let err_val = afcu_set_native_id(native_id);
+        handle_error_general(AfError::from(err_val));
+    }
+}
+
+/// Get CUDA stream of active CUDA device
+///
+/// # Parameters
+///
+/// - `id` is the identifier of device in ArrayFire context
+///
+/// # Return Values
+///
+/// [cudaStream_t](https://docs.rs/cuda-runtime-sys/0.3.0-alpha.1/cuda_runtime_sys/type.cudaStream_t.html) handle.
+pub fn get_stream(native_id: i32) -> cudaStream_t {
+    unsafe {
+        let mut ret_val: cudaStream_t = std::ptr::null_mut();
+        let err_val = afcu_get_stream(&mut ret_val as *mut cudaStream_t, native_id);
+        handle_error_general(AfError::from(err_val));
+        ret_val
+    }
+}

scripts/generate_documentation.sh

@@ -3,5 +3,6 @@
 # this script meant to be run from the root of arrayfire-rust
 
 cargo rustdoc -p arrayfire -- --html-in-header ./scripts/mathjax.script
+cargo rustdoc -p af-cuda-interop -- --html-in-header ./scripts/mathjax.script
 
-mdbook build tutorials-book && cp -r tutorials-book/book ./target/doc/arrayfire/
+mdbook build tutorials-book && cp -r tutorials-book/book ./target/doc/

src/blas/mod.rs

@@ -1,4 +1,6 @@
-use super::core::{af_array, AfError, Array, FloatingPoint, HasAfEnum, MatProp, HANDLE_ERROR};
+use super::core::{
+    af_array, AfError, Array, CublasMathMode, FloatingPoint, HasAfEnum, MatProp, HANDLE_ERROR,
+};
 
 use libc::{c_int, c_uint, c_void};
 use std::vec::Vec;
@@ -32,6 +34,8 @@ extern "C" {
 
     fn af_transpose(out: *mut af_array, arr: af_array, conjugate: bool) -> c_int;
     fn af_transpose_inplace(arr: af_array, conjugate: bool) -> c_int;
+
+    fn afcu_cublasSetMathMode(mode: c_int) -> c_int;
 }
 
 /// BLAS general matrix multiply (GEMM) of two Array objects
@@ -237,3 +241,20 @@ pub fn transpose_inplace<T: HasAfEnum>(arr: &mut Array<T>, conjugate: bool) {
         HANDLE_ERROR(AfError::from(err_val));
     }
 }
+
+/// Sets the cuBLAS math mode for the internal handle.
+///
+/// See the cuBLAS documentation for additional details
+///
+/// # Parameters
+///
+/// - `mode` takes a value of [CublasMathMode](./enum.CublasMathMode.html) enum
+pub fn set_cublas_mode(mode: CublasMathMode) {
+    unsafe {
+        afcu_cublasSetMathMode(mode as c_int);
+        //let err_val = afcu_cublasSetMathMode(mode as c_int);
+        // FIXME(wonder if this something to throw off,
+        // the program state is not invalid or anything
+        // HANDLE_ERROR(AfError::from(err_val));
+    }
+}