From fde2e408de13e014103a4f836725be8b8ecc6e57 Mon Sep 17 00:00:00 2001
From: Yunsong Wang <yunsongw@nvidia.com>
Date: Mon, 7 Feb 2022 13:09:35 -0500
Subject: [PATCH] Add stream benchmark example

---
 README.md               |  1 +
 examples/CMakeLists.txt |  1 +
 examples/stream.cu      | 62 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 64 insertions(+)
 create mode 100644 examples/stream.cu

diff --git a/README.md b/README.md
index a7240ec..28c5568 100644
--- a/README.md
+++ b/README.md
@@ -60,6 +60,7 @@ various NVBench features and usecases:
 - [Enums and compile-time-constant-integral parameter axes](examples/enums.cu)
 - [Reporting item/sec and byte/sec throughput statistics](examples/throughput.cu)
 - [Skipping benchmark configurations](examples/skip.cu)
+- [Benchmarking on a specific stream](examples/stream.cu)
 - [Benchmarks that sync CUDA devices: `nvbench::exec_tag::sync`](examples/exec_tag_sync.cu)
 - [Manual timing: `nvbench::exec_tag::timer`](examples/exec_tag_timer.cu)
 
diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt
index a6adc80..4532bfe 100644
--- a/examples/CMakeLists.txt
+++ b/examples/CMakeLists.txt
@@ -4,6 +4,7 @@ set(example_srcs
   exec_tag_sync.cu
   exec_tag_timer.cu
   skip.cu
+  stream.cu
   throughput.cu
   auto_throughput.cu
 )
diff --git a/examples/stream.cu b/examples/stream.cu
new file mode 100644
index 0000000..d0ca0c8
--- /dev/null
+++ b/examples/stream.cu
@@ -0,0 +1,62 @@
+/*
+ *  Copyright 2022 NVIDIA Corporation
+ *
+ *  Licensed under the Apache License, Version 2.0 with the LLVM exception
+ *  (the "License"); you may not use this file except in compliance with
+ *  the License.
+ *
+ *  You may obtain a copy of the License at
+ *
+ *      http://llvm.org/foundation/relicensing/LICENSE.txt
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include <nvbench/nvbench.cuh>
+
+// Grab some testing kernels from NVBench:
+#include <nvbench/test_kernels.cuh>
+
+// Thrust vectors simplify memory management:
+#include <thrust/device_vector.h>
+
+// A function to benchmark but does not expose an explicit stream argument.
+void copy(int32_t *input, int32_t *output, std::size_t const num_values)
+{
+  nvbench::copy_kernel<<<256, 256>>>(input, output, num_values);
+}
+
+// `stream_bench` copies a 64 MiB buffer of int32_t on a CUDA stream specified
+// by the user.
+//
+// By default, NVBench creates and provides an explicit stream via
+// `launch::get_stream()` to pass to every stream-ordered operation. Sometimes
+// it is inconvenient or impossible to specify an explicit CUDA stream to every
+// stream-ordered operation. In this case, users may specify a target stream via
+// `state::set_cuda_stream`. It is assumed that all work of interest executes on
+// or synchronizes with this stream.
+void stream_bench(nvbench::state &state)
+{
+  // Allocate input data:
+  const std::size_t num_values = 64 * 1024 * 1024 / sizeof(nvbench::int32_t);
+  thrust::device_vector<nvbench::int32_t> input(num_values);
+  thrust::device_vector<nvbench::int32_t> output(num_values);
+
+  // Set CUDA default stream as the target stream. Note the default stream
+  // is non-owning.
+  cudaStream_t default_stream = 0;
+  state.set_cuda_stream(
+    nvbench::cuda_stream{default_stream, false /*owning = false*/});
+
+  state.exec([&input, &output, num_values](nvbench::launch &) {
+    copy(thrust::raw_pointer_cast(input.data()),
+         thrust::raw_pointer_cast(output.data()),
+         num_values);
+  });
+}
+
+NVBENCH_BENCH(stream_bench);