adding split k reduction to conv2d profiler

masahi · Feb 4, 2022 · be89334 · be89334
1 parent ae09b0f
commit be89334
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 1 deletion.
diff --git a/python/tvm/contrib/cutlass/conv2d_profiler.py b/python/tvm/contrib/cutlass/conv2d_profiler.py
@@ -149,6 +149,34 @@ def __init__(self):
   for (int iteration = 0; iteration < 100; ++iteration) {
     auto status = implicit_gemm_op();
     CUTLASS_CHECK(status);
+    if (split_k_slices > 1) {
+      ReductionDevice reduction_op;
+      auto& status = result.status;
+      cutlass::conv::Operator kConvolutionalOperator = ImplicitGemm::kConvolutionalOperator;
+      auto stride_idx = ImplicitGemm::ImplicitGemmKernel::kTensorCStrideIdx;
+      typename ReductionDevice::Arguments reduction_args(
+        cutlass::conv::implicit_gemm_problem_size(kConvolutionalOperator, problem_size).mn(),
+        problem_size.split_k_slices,
+        cutlass::conv::implicit_gemm_tensor_c_size(kConvolutionalOperator, problem_size),
+        {
+          reinterpret_cast<ElementAccumulator*> (workspace.get()),
+          ReductionStrideIndex(tensor_c.stride()[stride_idx])
+        },
+        {
+          tensor_d.device_data(),
+          ReductionStrideIndex(tensor_d.stride()[stride_idx])
+        },
+        {
+          tensor_c.device_data(),
+          ReductionStrideIndex(tensor_c.stride()[stride_idx])
+        },
+        {ElementComputeEpilogue(1), ElementComputeEpilogue(0)},
+      );
+
+      status = reduction_op.initialize(reduction_args, nullptr);
+      status = reduction_op();
+    }
+
   }
 
   cudaEventRecord(events[1]);

diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py
@@ -151,7 +151,9 @@ def enumerate_conv2d_operators(
                 ret.append(
                     {
                         "src": profiler_emitter.emit(
-                            kernel_emitter.emit(op), op.procedural_name(), split_k_slice
+                            kernel_emitter.emit(op, emit_reduction=split_k_slices > 1),
+                            op.procedural_name(),
+                            split_k_slice,
                         ),
                         "name": op.procedural_name(),
                         "tile_description": tile,