add implementation of Eval(ctx, data, size, lambda) (k2-fsa#164)

Author: qindazhu

.clang-format

@@ -3,7 +3,7 @@ BasedOnStyle: Google
 ---
 Language:               Cpp
 Cpp11BracedListStyle:   true
-Standard:               c++11
+Standard:               Cpp11
 DerivePointerAlignment: false
 PointerAlignment:       Right
 ---

k2/csrc/array.h

@@ -61,19 +61,12 @@ class Array1 {
   // with CUDA) and also on the CPU.  We'll do src(i) to evaluate element i.
   // NOTE: we assume this thread is already set to use the device associated
   // with the context in 'ctx', if it's a CUDA context.
-  // TODO(haowen): require Callable to be a function, the compiler may confuse
-  // with Array1(ctx, size, elem)
-  /*
   template <typename Callable>
   Array1(ContextPtr ctx, int32_t size, Callable &&callable) {
     Init(ctx, size);
-    K2_LOG(FATAL) << "Not Implemented";
-
-    // TODO(haowen): there's no such definition
-    // `Eval(ContextPtr, T*, int32_t, Callable&)` now
-    // Eval(ctx, Data(), size, std::forward<Callable>(callable));
+    T *data = Data();
+    Eval(ctx, data, size, std::forward<Callable>(callable));
   }
-  */
 
   Array1(ContextPtr ctx, int32_t size) { Init(ctx, size); }
 

k2/csrc/array_test.cu

@@ -62,14 +62,8 @@ void TestArray1() {
     MemoryCopy(static_cast<void *>(array_data),
                static_cast<void *>(data.data()),
                array.Dim() * array.ElementSize(), kind);
-    // copy data from CPU/GPU to CPU
-    kind = GetMemoryCopyKind(*array.Context(), *cpu);
-    std::vector<T> cpu_data(array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(array_data),
-               array.Dim() * array.ElementSize(), kind);
     for (int32_t i = 0; i < array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], i);
+      EXPECT_EQ(array[i], i);
     }
   }
 
@@ -79,32 +73,30 @@ void TestArray1() {
     ASSERT_EQ(array.Dim(), 5);
     // operator=(T t)
     array = 2;
-    // copy data from CPU/GPU to CPU
-    const T *array_data = array.Data();
-    auto kind = GetMemoryCopyKind(*array.Context(), *cpu);
-    std::vector<T> cpu_data(array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(array_data),
-               array.Dim() * array.ElementSize(), kind);
     for (int32_t i = 0; i < array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], 2);
       EXPECT_EQ(array[i], 2);
     }
   }
 
   {
     // created with Array1(ContextPtr, int32_t size, T elem)
-    Array1<T> array(context, 5, 2);
+    Array1<T> array(context, 5, T(2));
     ASSERT_EQ(array.Dim(), 5);
     // copy data from CPU/GPU to CPU
-    const T *array_data = array.Data();
-    auto kind = GetMemoryCopyKind(*array.Context(), *cpu);
-    std::vector<T> cpu_data(array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(array_data),
-               array.Dim() * array.ElementSize(), kind);
     for (int32_t i = 0; i < array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], 2);
+      EXPECT_EQ(array[i], 2);
+    }
+  }
+
+  {
+    // created with Array1(ContextPtr, int32_t size, Callable &&callable)
+    auto lambda_set_values = [] __host__ __device__(int32_t i) -> T {
+      return i * i;
+    };
+    Array1<T> array(context, 5, lambda_set_values);
+    ASSERT_EQ(array.Dim(), 5);
+    for (int32_t i = 0; i < array.Dim(); ++i) {
+      EXPECT_EQ(array[i], i * i);
     }
   }
 
@@ -114,15 +106,8 @@ void TestArray1() {
     std::iota(data.begin(), data.end(), 0);
     Array1<T> array(context, data);
     ASSERT_EQ(array.Dim(), 5);
-    // copy data from CPU/GPU to CPU
-    const T *array_data = array.Data();
-    auto kind = GetMemoryCopyKind(*array.Context(), *cpu);
-    std::vector<T> cpu_data(array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(array_data),
-               array.Dim() * array.ElementSize(), kind);
     for (int32_t i = 0; i < array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], data[i]);
+      EXPECT_EQ(array[i], data[i]);
     }
   }
 
@@ -135,15 +120,8 @@ void TestArray1() {
     int32_t size = 6;
     Array1<T> sub_array = array.Range(start, size);
     ASSERT_EQ(sub_array.Dim(), size);
-    // copy data from CPU/GPU to CPU
-    const T *sub_array_data = sub_array.Data();
-    auto kind = GetMemoryCopyKind(*sub_array.Context(), *cpu);
-    std::vector<T> cpu_data(sub_array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(sub_array_data),
-               sub_array.Dim() * sub_array.ElementSize(), kind);
     for (int32_t i = 0; i < sub_array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], data[i + start]);
+      EXPECT_EQ(sub_array[i], data[i + start]);
     }
   }
 
@@ -239,16 +217,9 @@ void TestArray1() {
     Array1<int32_t> indexes_array(context, indexes);
     std::vector<T> expected_data = {1, 2, 3, 6, 2, 7, 9, 10, 3, 5, 7, 4};
     Array1<T> ans_array = array[indexes_array];
-    // copy data from CPU/GPU to CPU
     ASSERT_EQ(ans_array.Dim(), expected_data.size());
-    const T *ans_array_data = ans_array.Data();
-    auto kind = GetMemoryCopyKind(*ans_array.Context(), *cpu);
-    std::vector<T> cpu_data(ans_array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(ans_array_data),
-               ans_array.Dim() * ans_array.ElementSize(), kind);
     for (int32_t i = 0; i < ans_array.Dim(); ++i) {
-      EXPECT_EQ(cpu_data[i], expected_data[i]);
+      EXPECT_EQ(ans_array[i], expected_data[i]);
     }
   }
 
@@ -272,14 +243,9 @@ void TestArray1() {
     EXPECT_EQ(array.Dim(), new_size);
     // copy data from CPU/GPU to CPU
     const T *array_data = array.Data();
-    auto kind = GetMemoryCopyKind(*array.Context(), *cpu);
-    std::vector<T> cpu_data(array.Dim());
-    MemoryCopy(static_cast<void *>(cpu_data.data()),
-               static_cast<const void *>(array_data),
-               array.Dim() * array.ElementSize(), kind);
     // data.size() == 5, array.Dim() == 8, there are 3 uninitialized elements.
     for (int32_t i = 0; i < data.size(); ++i) {
-      EXPECT_EQ(cpu_data[i], data[i]);
+      EXPECT_EQ(array[i], data[i]);
     }
   }
 }

k2/csrc/context.h

@@ -337,6 +337,14 @@ __global__ void eval_lambda(int32_t n, LambdaT lambda) {
   }
 }
 
+template <typename T, typename LambdaT>
+__global__ void eval_lambda(T *data, int32_t n, LambdaT lambda) {
+  int32_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  if (i < n) {
+    data[i] = lambda(i);
+  }
+}
+
 template <typename LambdaT>
 __global__ void eval_lambda2(int32_t m, int32_t n, LambdaT lambda) {
   // actually threadIdx.y will always be 1 for now so we could drop that part of
@@ -380,6 +388,34 @@ void Eval(ContextPtrType c, int32_t n, LambdaT &lambda) {
   Eval(c->GetCudaStream(), n, lambda);
 }
 
+/* Eval() will do `data[i] = lambda(i)` for 0 <= i < n, on the appropriate
+   device (CPU or GPU) */
+template <typename T, typename LambdaT>
+void Eval(cudaStream_t stream, T *data, int32_t n, LambdaT &lambda) {
+  if (n <= 0) return;  // actually it would be an error if n < 0.
+  if (stream == kCudaStreamInvalid) {
+    // TODO: if n is very large, we'll eventually support running this with
+    // multiple threads.
+    for (int32_t i = 0; i < n; ++i) {
+      data[i] = lambda(i);
+    }
+  } else {
+    int32_t block_size = 256;
+    int32_t grid_size = NumBlocks(n, block_size);
+    eval_lambda<T, LambdaT>
+        <<<grid_size, block_size, 0, stream>>>(data, n, lambda);
+    auto err = cudaGetLastError();
+    K2_DCHECK_CUDA_ERROR(err);
+  }
+}
+
+template <typename ContextPtrType,  // Context*  or ContextPtr ==
+                                    // std::shared_ptr<Context>
+          typename T, typename LambdaT>
+void Eval(ContextPtrType c, T *data, int32_t n, LambdaT &lambda) {
+  Eval(c->GetCudaStream(), data, n, lambda);
+}
+
 /*
   This is a form of Eval() where the lambda takes  two arguments.
 

k2/csrc/fsa.cu

@@ -76,7 +76,7 @@ int32_t GetFsaVecBasicProperties(FsaVec &fsa_vec) {
   // the final-state of its FSA (i.e. last-numbered) or has at least one arc
   // leaving it, not counting self-loops. Again, it's a looser condition than
   // being 'co-accessible' in FSA terminology.
-  Array1<char> reachable(c, num_states * 2 + 1, 0);
+  Array1<char> reachable(c, num_states * 2 + 1, static_cast<char>(0));
   Array1<char> flag = reachable.Range(num_states * 1, 1);
   Array1<char> co_reachable = reachable.Range(num_states, num_states);
   reachable = reachable.Range(0, num_states);

k2/csrc/tensor.cu

@@ -133,6 +133,7 @@ void Tensor::Init(ContextPtr c) {
 Tensor ToContiguous(const Tensor &tensor) {
   // TODO(haowen): implement
   K2_LOG(FATAL) << "Not implemented";
+  return tensor;
 }
 
 }  // namespace k2