print cuda kernel configs for convolution, dot and avg_pool

xysmlx · Sep 29, 2020 · 12b5a5e · 12b5a5e
1 parent 0493fa9
commit 12b5a5e
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/src/nnfusion/core/kernels/cuda_gpu/kernels/avg_pool.cpp b/src/nnfusion/core/kernels/cuda_gpu/kernels/avg_pool.cpp
@@ -278,6 +278,16 @@ cuda::AvgPoolmD::AvgPoolmD(shared_ptr<KernelContext> ctx)
         << join(window_stride, "_") << "_pb" << join(padding_below, "_") << "_pb"
         << join(padding_above, "_");
     custom_tag = tag.str();
+
+    // for kernel codegen
+    std::ofstream fout;
+    std::stringstream codegen_tag;
+    codegen_tag << "cudnn_avgpool_dtype_" << output_type << "|i+" << join(input_shape, "_") << "|o+"
+                << join(output_shape, "_") << "|ws+" << join(window_shape, "_") << "|wst+"
+                << join(window_stride, "_") << "|pb+" << join(padding_below, "_");
+    fout.open("/tmp/nnfusion_rt_avgpool_kernels.txt", std::ios_base::app);
+    fout << codegen_tag.str() << std::endl;
+    fout.close();
 }
 
 LanguageUnit_p cuda::AvgPoolmD::emit_function_body()

diff --git a/src/nnfusion/core/kernels/cuda_gpu/kernels/convolution.cpp b/src/nnfusion/core/kernels/cuda_gpu/kernels/convolution.cpp
@@ -28,6 +28,17 @@ cuda::ConvolutionCudnn::ConvolutionCudnn(shared_ptr<KernelContext> ctx)
         << join(window_movement_strides, "_") << "_wd" << join(window_dilation_strides, "_") << "_p"
         << join(padding_below_diff, "_");
     custom_tag = tag.str();
+
+    // for kernel codegen
+    std::ofstream fout;
+    std::stringstream codegen_tag;
+    codegen_tag << "cudnn_convolution_op_" << dtype << "|i+" << join(input_shape, "_") << "|w+"
+                << join(filter_shape, "_") << "|o+" << join(output_shape, "_") << "|ws+"
+                << join(window_movement_strides, "_") << "|wd+"
+                << join(window_dilation_strides, "_") << "|p+" << join(padding_below_diff, "_");
+    fout.open("/tmp/nnfusion_rt_conv_kernels.txt", std::ios_base::app);
+    fout << codegen_tag.str() << std::endl;
+    fout.close();
 }
 
 LanguageUnit_p cuda::ConvolutionCudnn::emit_function_body()

diff --git a/src/nnfusion/core/kernels/cuda_gpu/kernels/dot.cpp b/src/nnfusion/core/kernels/cuda_gpu/kernels/dot.cpp
@@ -22,6 +22,16 @@ cuda::Dot::Dot(shared_ptr<KernelContext> ctx)
         << "_dtype_" << dtype.c_type_string() << "_reduction_axes_count_" << reduction_axes << "_i_"
         << join(arg0_shape, "_") << "_i_" << join(arg1_shape, "_");
     custom_tag = tag.str();
+
+    // for kernel codegen
+    std::ofstream fout;
+    std::stringstream codegen_tag;
+    codegen_tag << "cublas_dot_op_" << dtype << "|_reduction_axes_count_+" << reduction_axes
+                << "|arg0_shape+" << join(arg0_shape, "_") << "|arg1_shape+"
+                << join(arg1_shape, "_") << "|out_shape+" << join(out_shape, "_");
+    fout.open("/tmp/nnfusion_rt_dot_kernels.txt", std::ios_base::app);
+    fout << codegen_tag.str() << std::endl;
+    fout.close();
 }
 
 LanguageUnit_p cuda::Dot::emit_function_body()