gpu slice optimization (#1172)

* add gpu_timer to external function

* compiled version

* working version

* using block_begin and block_end

* add the missing '
;'

* move slice to cuda emiter

* change size_t to uint32_t in kernel

* working version

* change block size from 1 to 64

* fix bugs

* nthreads need to be size_t in broadcast op

* add rank to kernel name hash

* update slice in convolution

* resolve index conflict

* change align to align_to_blocksize, add overflow check

* add gird size check and fix pool merge bug

* code style, change names

src/ngraph/runtime/gpu/cuda_emitter.hpp

@@ -47,7 +47,7 @@ namespace ngraph
                                  GPUShape pad_interior,
                                  const std::string& pad_value = "");
 
-                size_t build_pad_dynamic(const runtime::gpu::GPURuntimeContext* ctx,
+                size_t build_pad_dynamic(const GPURuntimeContext* ctx,
                                          const std::array<std::string, 2>& dtypes,
                                          GPUShape input_shape,
                                          GPUShape output_shape,
@@ -70,6 +70,13 @@ namespace ngraph
                                       GPUShape padding_below,
                                       bool include_pad = false);
 
+                size_t build_slice(const GPURuntimeContext* ctx,
+                                   const std::array<std::string, 2>& dtypes,
+                                   GPUShape input_shape,
+                                   GPUShape lower_bounds,
+                                   GPUShape slice_strides,
+                                   GPUShape output_shape);
+
                 size_t build_reduce_window(const GPURuntimeContext* ctx,
                                            const OpName op_name,
                                            const std::vector<std::string>& dtypes,
@@ -144,6 +151,7 @@ namespace ngraph
 
             private:
                 CUDAEmitter(GPUPrimitiveEmitter* emitter);
+                uint32_t align_to_block_size(uint32_t grid_size, uint32_t block_size);
                 void print_tensor_from_gpu(codegen::CodeWriter& writer,
                                            const std::string& tensor_name,
                                            GPUShape shape);

src/ngraph/runtime/gpu/gpu_cuda_kernel_builder.cpp

@@ -363,30 +363,36 @@ void runtime::gpu::CudaKernelBuilder::get_reverse_sequence_op(
 }
 void runtime::gpu::CudaKernelBuilder::get_slice_op(codegen::CodeWriter& writer,
                                                    const std::string& name,
-                                                   const std::array<std::string, 2>& data_types)
+                                                   const std::array<std::string, 2>& data_types,
+                                                   size_t rank)
 {
     writer << "extern \"C\" __global__ void cuda_" << name << "(" << data_types[0] << "* in, "
-           << data_types[1] << "* out, size_t* input_strides, size_t* lower_bounds, size_t* "
-                               "slice_strides, size_t* output_strides, size_t rank, size_t n)\n";
+           << data_types[1] << "* out, uint32_t* input_strides, uint32_t* lower_bounds, uint32_t* "
+                               "slice_strides, uint32_t* output_strides, uint32_t n)\n";
     writer.block_begin();
     {
-        writer << "size_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n";
+        writer << "uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;\n";
         writer << "if (tid < n)\n";
         writer.block_begin();
         {
-            writer << "size_t input_idx = 0;\n";
-            writer << "size_t output_idx = tid;\n";
-
-            writer << "for(size_t i = 0; i < rank; i++)\n";
-            writer.block_begin();
+            writer << "uint32_t input_idx = 0;\n";
+            writer << "uint32_t output_idx = tid;\n";
+            size_t i = 0;
+            for (; i < rank - 1; i++)
             {
-                writer << "input_idx += (((output_idx / output_strides[i]) * slice_strides[i]) + "
-                          "lower_bounds[i]) * input_strides[i];\n";
-                writer << "output_idx %= output_strides[i];\n";
+                writer << "input_idx += (((output_idx / output_strides[" << i
+                       << "]) * slice_strides[" << i << "]) + "
+                                                        "lower_bounds["
+                       << i << "]) * input_strides[" << i << "];\n";
+                writer << "output_idx %= output_strides[" << i << "];\n";
             }
-            writer.block_end();
+            writer << "input_idx += (((output_idx / output_strides[" << i << "]) * slice_strides["
+                   << i << "]) + "
+                           "lower_bounds["
+                   << i << "]) * input_strides[" << i << "];\n";
             writer << "out[tid] = in[input_idx];\n";
         }
+
         writer.block_end();
     }
     writer.block_end();

src/ngraph/runtime/gpu/gpu_cuda_kernel_builder.hpp

@@ -59,7 +59,8 @@ namespace ngraph
 
                 static void get_slice_op(codegen::CodeWriter& writer,
                                          const std::string& name,
-                                         const std::array<std::string, 2>& data_types);
+                                         const std::array<std::string, 2>& data_types,
+                                         size_t rank);
 
                 static void get_reverse_op(codegen::CodeWriter& writer,
                                            const std::string& name,

src/ngraph/runtime/gpu/gpu_cuda_kernel_emitters.cpp

@@ -97,46 +97,6 @@ void runtime::gpu::emit_reshape(const std::string& name,
     CUDA_SAFE_CALL(cuCtxSynchronize()); // Retrieve and print output.
 }
 
-void runtime::gpu::emit_slice(const std::string& name,
-                              CUdeviceptr in,
-                              CUdeviceptr out,
-                              const std::array<std::string, 2>& data_types,
-                              GPURuntimeContext* ctx,
-                              CUdeviceptr input_strides,
-                              CUdeviceptr lower_bounds,
-                              CUdeviceptr slice_strides,
-                              CUdeviceptr output_strides,
-                              size_t rank,
-                              size_t count)
-{
-    std::string name_signature = name + "_" + data_types[0] + "_" + data_types[1];
-    std::replace(name_signature.begin(), name_signature.end(), ' ', '_');
-    auto compiled_kernel = ctx->compiled_kernel_pool->get(name_signature);
-    if (compiled_kernel == nullptr)
-    {
-        codegen::CodeWriter writer;
-        CudaKernelBuilder::add_pod_typedefs(writer);
-        CudaKernelBuilder::get_slice_op(writer, name_signature, data_types);
-        std::string kernel = writer.get_code();
-        compiled_kernel = ctx->compiled_kernel_pool->set(name_signature, kernel);
-    }
-
-    void* args_list[] = {
-        &in, &out, &input_strides, &lower_bounds, &slice_strides, &output_strides, &rank, &count};
-    CUDA_SAFE_CALL(cuLaunchKernel(*compiled_kernel.get(),
-                                  static_cast<uint32_t>(count),
-                                  1,
-                                  1, // grid dim
-                                  1,
-                                  1,
-                                  1, // block dim
-                                  0,
-                                  NULL, // shared mem and stream
-                                  args_list,
-                                  0));  // arguments
-    CUDA_SAFE_CALL(cuCtxSynchronize()); // Retrieve and print output.
-}
-
 void runtime::gpu::emit_reverse(const std::string& name,
                                 CUdeviceptr in,
                                 CUdeviceptr out,

src/ngraph/runtime/gpu/gpu_cuda_kernel_emitters.hpp

@@ -55,18 +55,6 @@ namespace ngraph
                               size_t rank,
                               size_t count);
 
-            void emit_slice(const std::string& name,
-                            CUdeviceptr in,
-                            CUdeviceptr out,
-                            const std::array<std::string, 2>& data_types,
-                            GPURuntimeContext* ctx,
-                            CUdeviceptr input_strides,
-                            CUdeviceptr lower_bounds,
-                            CUdeviceptr slice_strides,
-                            CUdeviceptr output_strides,
-                            size_t rank,
-                            size_t count);
-
             void emit_reverse(const std::string& name,
                               CUdeviceptr in,
                               CUdeviceptr out,

src/ngraph/runtime/gpu/gpu_emitter.cpp

@@ -466,45 +466,20 @@ namespace ngraph
                 // since we padded output with temp buffer, we need to copy back to real ouput
                 if (pad_required || is_deconvolution)
                 {
-                    const auto arg_rank = output_shape.size();
-                    const auto input_strides = row_major_strides(output_shape_padded);
-                    const auto output_strides = row_major_strides(output_shape);
-                    GPUAllocator allocator =
-                        external_function->get_primitive_emitter()->get_memory_allocator();
-                    size_t idx_input_strides = allocator.reserve_argspace(
-                        input_strides.data(), input_strides.size() * sizeof(size_t));
-                    size_t idx_output_strides = allocator.reserve_argspace(
-                        output_strides.data(), output_strides.size() * sizeof(size_t));
-                    size_t idx_lower_bounds = allocator.reserve_argspace(
-                        padding_below_back.data(), padding_below_back.size() * sizeof(size_t));
-                    size_t idx_slice_strides =
-                        allocator.reserve_argspace(padding_interior_back.data(),
-                                                   padding_interior_back.size() * sizeof(size_t));
-
-                    writer << "size_t rank = " << arg_rank << ";\n";
-                    writer << "void* input_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_input_strides
-                           << ");\n";
-                    writer << "void* output_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_output_strides
-                           << ");\n";
-                    writer << "void* slice_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_slice_strides
-                           << ");\n";
-                    writer << "void* lower_bounds_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_lower_bounds
-                           << ");\n";
+                    auto& cuda_emitter =
+                        external_function->get_primitive_emitter()->get_cuda_emitter();
+                    auto slice_index =
+                        cuda_emitter->build_slice(external_function->ctx().get(),
+                                                  {{args[0].get_type(), out[0].get_type()}},
+                                                  output_shape_padded,
+                                                  padding_below_back,
+                                                  padding_interior_back,
+                                                  output_shape);
 
-                    writer << "runtime::gpu::emit_slice(\"" << node->description()
-                           << "\", CUdeviceptr(pad_buffer), CUdeviceptr(" << out[0].get_name()
-                           << ")"
-                           << ", {\"" << args[0].get_type() << "\", \"" << out[0].get_type()
-                           << "\"}"
-                           << ", "
-                           << "ctx, "
-                           << "CUdeviceptr(input_strides_d), CUdeviceptr(lower_bounds_d), "
-                              "CUdeviceptr(slice_strides_d), CUdeviceptr(output_strides_d)"
-                           << ", " << arg_rank << ", " << out[0].get_size() << ");\n";
+                    writer << "gpu::invoke_primitive(ctx, " << slice_index << ", ";
+                    writer << "std::vector<void*>{pad_buffer}.data(), ";
+                    writer << "std::vector<void*>{" << out[0].get_name() << "}.data()";
+                    writer << ");\n";
                 }
                 writer.block_end();
             }
@@ -1075,12 +1050,9 @@ namespace ngraph
                 auto slice = static_cast<const op::Slice*>(node);
 
                 const auto arg_shape = args[0].get_shape();
-                const auto arg_rank = arg_shape.size();
                 const auto result_shape = out[0].get_shape();
                 const Coordinate& lower_bounds = slice->get_lower_bounds();
                 const Strides slice_strides = slice->get_strides();
-                const auto input_strides = row_major_strides(arg_shape);
-                const auto output_strides = row_major_strides(result_shape);
 
                 writer.block_begin();
                 if (args[0].get_size() == out[0].get_size())
@@ -1089,41 +1061,20 @@ namespace ngraph
                 }
                 else
                 {
-                    GPUAllocator allocator =
-                        external_function->get_primitive_emitter()->get_memory_allocator();
-                    size_t idx_input_strides = allocator.reserve_argspace(
-                        input_strides.data(), input_strides.size() * sizeof(size_t));
-                    size_t idx_output_strides = allocator.reserve_argspace(
-                        output_strides.data(), output_strides.size() * sizeof(size_t));
-                    size_t idx_lower_bounds = allocator.reserve_argspace(
-                        lower_bounds.data(), lower_bounds.size() * sizeof(size_t));
-                    size_t idx_slice_strides = allocator.reserve_argspace(
-                        slice_strides.data(), slice_strides.size() * sizeof(size_t));
-
-                    writer << "size_t rank = " << arg_rank << ";\n";
-                    writer << "void* input_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_input_strides
-                           << ");\n";
-                    writer << "void* output_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_output_strides
-                           << ");\n";
-                    writer << "void* slice_strides_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_slice_strides
-                           << ");\n";
-                    writer << "void* lower_bounds_d = "
-                           << " runtime::gpu::invoke_memory_primitive(ctx, " << idx_lower_bounds
-                           << ");\n";
+                    auto& cuda_emitter =
+                        external_function->get_primitive_emitter()->get_cuda_emitter();
+                    auto index =
+                        cuda_emitter->build_slice(external_function->ctx().get(),
+                                                  {{args[0].get_type(), out[0].get_type()}},
+                                                  arg_shape,
+                                                  lower_bounds,
+                                                  slice_strides,
+                                                  result_shape);
 
-                    writer << "runtime::gpu::emit_slice(\"" << node->description()
-                           << "\", CUdeviceptr(" << args[0].get_name() << "), CUdeviceptr("
-                           << out[0].get_name() << ")"
-                           << ", {\"" << args[0].get_type() << "\", \"" << out[0].get_type()
-                           << "\"}"
-                           << ", "
-                           << "ctx, "
-                           << "CUdeviceptr(input_strides_d), CUdeviceptr(lower_bounds_d), "
-                              "CUdeviceptr(slice_strides_d), CUdeviceptr(output_strides_d)"
-                           << ", " << arg_rank << ", " << out[0].get_size() << ");\n";
+                    writer << "gpu::invoke_primitive(ctx, " << index << ", ";
+                    writer << "std::vector<void*>{" << args[0].get_name() << "}.data(), ";
+                    writer << "std::vector<void*>{" << out[0].get_name() << "}.data()";
+                    writer << ");\n";
                 }
                 writer.block_end();
             }