Merge pull request #27 from arrayfire/feature/index

Indexing

Author: 9prady9

arrayfire

@@ -7,7 +7,12 @@ fn main() {
     set_device(0);
     info();
 
-    let dims = Dim4::new(&[5, 3, 1, 1]);
+    let num_rows: u64 = 5;
+    let num_cols: u64 = 3;
+    let values: &[f32] = &[1.0, 2.0, 3.0];
+    let indices = Array::new(Dim4::new(&[3, 1, 1, 1]), values, Aftype::F32).unwrap();
+
+    let dims = Dim4::new(&[num_rows, num_cols, 1, 1]);
 
     println!("Create a 5-by-3 matrix of random floats on the GPU");
     let a = match randu(dims, Aftype::F32) {
@@ -36,16 +41,35 @@ fn main() {
     let test = &a + &b;
     println!("a + b"); print(&test);
 
+    // Index array using sequences
+    let seqs = &[Seq::new(1.0, 3.0, 1.0), Seq::default()];
+    let sub = index(&a, seqs).unwrap();
+    println!("a(seq(1,3,1), span)"); print(&sub);
+
+    //Index array using array and sequence
+    let seq4gen = Seq::new(0.0, 2.0, 1.0);
+
+    let mut idxrs = match Indexer::new() {
+        Ok(v) => v,
+        Err(e) => panic!("{}",e),
+    };
+    idxrs.set_index(&indices, 0, None);
+    idxrs.set_index(&seq4gen, 1, Some(false));
+
+    let sub2 = index_gen(&a, idxrs).unwrap();
+    println!("a(indices, seq(0, 2, 1))"); print(&sub2);
+
     // printf("Negate the first three elements of second column\n");
     // B(seq(0, 2), 1) = B(seq(0, 2), 1) * -1;
     // af_print(B);
 
     println!("Fourier transform the result");
     fft(&b, 1.0, 0).map(|x| print(&x));
 
-    // printf("Grab last row\n");
-    // array c = C.row(end);
-    // af_print(c);
+    println!("Grab last row & col of the random matrix");
+    print(&a);
+    print(&row(&a, num_rows - 1).unwrap());
+    print(&col(&a, num_cols - 1).unwrap());
 
     println!("Create 2-by-3 matrix from host data");
     let d_dims = Dim4::new(&[2, 3, 1, 1]);

examples/helloworld.rs

@@ -153,7 +153,7 @@ impl Array {
         self.handle
     }
 
-    pub fn host(&self, data:&mut [f64]) -> Result<(), AfError> {
+    pub fn host<T>(&self, data: &mut [T]) -> Result<(), AfError> {
         unsafe {
             let ret_val = af_get_data_ptr(data.as_mut_ptr() as *mut c_void, self.handle as AfArray);
             match ret_val {

src/array.rs

@@ -67,7 +67,7 @@ impl Window {
             let mut temp: u64 = 0;
             let err_val = af_create_window(&mut temp as MutWndHandle,
                                            width as c_int, height as c_int,
-                                           title.as_bytes().as_ptr() as *const u8);
+                                           title.clone().as_bytes().as_ptr() as *const u8);
             match err_val {
                 0 => Ok(Window::from(temp)),
                 _ => Err(AfError::from(err_val)),
@@ -88,7 +88,7 @@ impl Window {
     pub fn set_title(&self, title: String) -> Result<(), AfError> {
         unsafe {
             let err_val = af_set_title(self.handle as WndHandle,
-                                       title.as_bytes().as_ptr() as *const u8);
+                                       title.clone().as_bytes().as_ptr() as *const u8);
             match err_val {
                 0 => Ok(()),
                 _ => Err(AfError::from(err_val)),

src/graphics.rs

@@ -93,7 +93,7 @@ pub fn load_image(filename: String, is_color: bool) -> Result<Array, AfError> {
     unsafe {
         let mut temp: i64 = 0;
         let err_val = af_load_image(&mut temp as MutAfArray,
-                                    filename.as_bytes().as_ptr() as *const u8,
+                                    filename.clone().as_bytes().as_ptr() as *const u8,
                                     is_color as c_int);
         match err_val {
             0 => Ok(Array::from(temp)),
@@ -105,7 +105,7 @@ pub fn load_image(filename: String, is_color: bool) -> Result<Array, AfError> {
 #[allow(unused_mut)]
 pub fn save_image(filename: String, input: &Array) -> Result<(), AfError> {
     unsafe {
-        let err_val = af_save_image(filename.as_bytes().as_ptr() as *const u8,
+        let err_val = af_save_image(filename.clone().as_bytes().as_ptr() as *const u8,
                                     input.get() as AfArray);
         match err_val {
             0 => Ok(()),

src/image/mod.rs

@@ -0,0 +1,187 @@
+extern crate libc;
+
+use array::Array;
+use defines::AfError;
+use seq::Seq;
+use self::libc::{c_int, c_uint, c_longlong};
+
+type MutAfIndex    = *mut self::libc::c_longlong;
+type MutAfArray    = *mut self::libc::c_longlong;
+type AfArray       = self::libc::c_longlong;
+type DimT          = self::libc::c_longlong;
+type IndexT        = self::libc::c_longlong;
+
+#[allow(dead_code)]
+extern {
+    fn af_create_indexers(indexers: MutAfIndex) -> c_int;
+    fn af_set_array_indexer(indexer: MutAfIndex, idx: AfArray, dim: DimT) -> c_int;
+    fn af_set_seq_indexer(indexer: MutAfIndex, idx: *const Seq, dim: DimT, is_batch: c_int) -> c_int;
+    fn af_release_indexers(indexers: MutAfIndex) -> c_int;
+
+    fn af_index(out: MutAfArray, input: AfArray, ndims: c_uint, index: *const Seq) -> c_int;
+    fn af_lookup(out: MutAfArray, arr: AfArray, indices: AfArray, dim: c_uint) -> c_int;
+    fn af_assign_seq(out: MutAfArray, lhs: AfArray, ndims: c_uint, indices: *const Seq, rhs: AfArray) -> c_int;
+    fn af_index_gen(out: MutAfArray, input: AfArray, ndims: DimT, indices: *const IndexT) -> c_int;
+    fn af_assign_gen(out: MutAfArray, lhs: AfArray, ndims: DimT, indices: *const IndexT, rhs: AfArray) -> c_int;
+}
+
+pub struct Indexer {
+    handle: i64,
+    count: u32,
+}
+
+pub trait Indexable {
+    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError>;
+}
+
+impl Indexable for Array {
+    #[allow(unused_variables)]
+    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
+        unsafe {
+            let err_val = af_set_array_indexer(idxr.clone().get() as MutAfIndex,
+                                             self.get() as AfArray,
+                                             dim as DimT);
+            match err_val {
+                0 => Ok(()),
+                _ => Err(AfError::from(err_val)),
+            }
+        }
+    }
+}
+
+impl Indexable for Seq {
+    fn set(&self, idxr: &Indexer, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
+        unsafe {
+            let err_val = af_set_seq_indexer(idxr.clone().get() as MutAfIndex, self as *const Seq,
+                                           dim as DimT, is_batch.unwrap() as c_int);
+            match err_val {
+                0 => Ok(()),
+                _ => Err(AfError::from(err_val)),
+            }
+        }
+    }
+}
+
+impl Indexer {
+    #[allow(unused_mut)]
+    pub fn new() -> Result<Indexer, AfError> {
+        unsafe {
+            let mut temp: i64 = 0;
+            let err_val = af_create_indexers(&mut temp as MutAfIndex);
+            match err_val {
+                0 => Ok(Indexer{handle: temp, count: 0}),
+                _ => Err(AfError::from(err_val)),
+            }
+        }
+    }
+
+    pub fn set_index<T: Indexable>(&mut self, idx: &T, dim: u32, is_batch: Option<bool>) -> Result<(), AfError> {
+        self.count = self.count + 1;
+        idx.set(self, dim, is_batch)
+    }
+
+    pub fn get(&self) -> i64 {
+        self.handle
+    }
+
+    pub fn len(&self) -> u32 {
+        self.count
+    }
+}
+
+impl Drop for Indexer {
+    fn drop(&mut self) {
+        unsafe {
+            let ret_val = af_release_indexers(self.handle as MutAfIndex);
+            match ret_val {
+                0 => (),
+                _ => panic!("Failed to release indexers resource: {}", ret_val),
+            }
+        }
+    }
+}
+
+pub fn index(input: &Array, seqs: &[Seq]) -> Result<Array, AfError> {
+    unsafe {
+        let mut temp: i64 = 0;
+        let err_val = af_index(&mut temp as MutAfArray
+                               , input.get() as AfArray, seqs.len() as u32
+                               , seqs.as_ptr() as *const Seq);
+        match err_val {
+            0 => Ok(Array::from(temp)),
+            _ => Err(AfError::from(err_val)),
+        }
+    }
+}
+
+pub fn row(input: &Array, row_num: u64) -> Result<Array, AfError> {
+    let dims_err = input.dims();
+    let dims = match dims_err {
+        Ok(dim) =>  dim.clone(),
+        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
+    };
+
+    index(input, &[Seq::new(row_num as f64, row_num as f64, 1.0)
+                    ,Seq::new(0.0, dims[1] as f64 - 1.0, 1.0)])
+}
+
+pub fn col(input: &Array, col_num: u64) -> Result<Array, AfError> {
+    let dims_err = input.dims();
+    let dims = match dims_err {
+        Ok(dim) =>  dim.clone(),
+        Err(e)  =>  panic!("Error unwrapping dims in row(): {}", e),
+    };
+
+    index(input, &[Seq::new(0.0, dims[0] as f64 - 1.0, 1.0)
+                    ,Seq::new(col_num as f64, col_num as f64, 1.0)])
+}
+
+pub fn lookup(input: &Array, indices: &Array, seq_dim: i32) -> Result<Array, AfError> {
+    unsafe {
+        let mut temp: i64 = 0;
+        let err_val = af_lookup(&mut temp as MutAfArray, input.get() as AfArray,
+                                indices.get() as AfArray, seq_dim as c_uint);
+        match err_val {
+            0 => Ok(Array::from(temp)),
+            _ => Err(AfError::from(err_val)),
+        }
+    }
+}
+
+pub fn assign_seq(lhs: &Array, ndims: usize, seqs: &[Seq], rhs: &Array) -> Result<Array, AfError> {
+    unsafe{
+        let mut temp: i64 = 0;
+        let err_val = af_assign_seq(&mut temp as MutAfArray, lhs.get() as AfArray,
+                                    ndims as c_uint, seqs.as_ptr() as *const Seq,
+                                    rhs.get() as AfArray);
+        match err_val {
+            0 => Ok(Array::from(temp)),
+            _ => Err(AfError::from(err_val)),
+        }
+    }
+}
+
+pub fn index_gen(input: &Array, indices: Indexer) -> Result<Array, AfError> {
+    unsafe{
+        let mut temp: i64 = 0;
+        let err_val = af_index_gen(&mut temp as MutAfArray, input.get() as AfArray,
+                                   indices.len() as DimT, indices.get() as *const IndexT);
+        match err_val {
+            0 => Ok(Array::from(temp)),
+            _ => Err(AfError::from(err_val)),
+        }
+    }
+}
+
+pub fn assign_gen(lhs: &Array, indices: &Indexer, rhs: &Array) -> Result<Array, AfError> {
+    unsafe{
+        let mut temp: i64 = 0;
+        let err_val = af_assign_gen(&mut temp as MutAfArray, lhs.get() as AfArray,
+                                    indices.len() as DimT, indices.get() as *const IndexT,
+                                    rhs.get() as AfArray);
+        match err_val {
+            0 => Ok(Array::from(temp)),
+            _ => Err(AfError::from(err_val)),
+        }
+    }
+}

src/index.rs

@@ -44,6 +44,12 @@ mod defines;
 pub use dim4::Dim4;
 mod dim4;
 
+pub use index::{Indexer, index, row, col, lookup, assign_seq, index_gen, assign_gen};
+mod index;
+
+pub use seq::Seq;
+mod seq;
+
 pub use graphics::Window;
 mod graphics;
 

src/lib.rs

@@ -0,0 +1,43 @@
+extern crate libc;
+
+use std::fmt;
+use std::default::Default;
+use self::libc::{c_double};
+
+#[derive(Copy, Clone)]
+#[repr(C)]
+pub struct Seq {
+    begin: c_double,
+    end: c_double,
+    step: c_double,
+}
+
+impl Default for Seq {
+    fn default() -> Seq {
+        Seq { begin: 1.0, end: 1.0, step: 0.0, }
+    }
+}
+
+impl fmt::Display for Seq {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "[begin: {}, end: {}, step: {}]", self.begin, self.end, self.step)
+    }
+}
+
+impl Seq {
+    pub fn new(begin: f64, end: f64, step: f64) -> Seq {
+        Seq { begin: begin, end: end, step: step, }
+    }
+
+    pub fn begin(&self) -> f64 {
+        self.begin as f64
+    }
+
+    pub fn end(&self) -> f64 {
+        self.end as f64
+    }
+
+    pub fn step(&self) -> f64 {
+        self.step as f64
+    }
+}