Add 1xtfloat capability to pairwise_matrix distance computations

cpp/bench/prims/distance/tune_pairwise/bench_cutlass.cu

@@ -0,0 +1,149 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.
+ */
+
+// Tuning benchmarks.
+//
+// Goals:
+//
+// 1. Fast compile times to maintain iteration speed.
+// 2. Create benchmarks that can inform the design of the kernels.
+//
+// Non-goals:
+//
+// 1. Measure every distance operation. Instead measures just one distance
+// operation at the same time.
+// 2. Be useful for finding performance regressions. This is handled by the
+// normal benchmarks.
+//
+// So far, both goals are partly achieved.
+//
+// RE (1), COMPILE TIMES: kernel.cu is fast to compile. This file is not.
+// When the internals of a pairwise distance kernel is changed, this file is not
+// recompiled.
+//
+// RE 2, benchmarks with intent: this file contains a benchmark to check the
+// maximal throughput of a kernel. Measuring other things, like performance on
+// skinny or wide matrices is not yet implemented.
+
+#include "kernel_cutlass.cuh" // launch_kernel
+#include <algorithm> // std::min
+#include <common/benchmark.hpp> // RAFT_BENCH_REGISTER
+#include <raft/distance/detail/pairwise_matrix/params.cuh> // pairwise_matrix_params
+#include <rmm/device_uvector.hpp> // rmm::device_uvector
+#include <vector> // std::vector
+
+namespace raft::bench::distance::tune_cutlass {
+
+// Max throughput benchmark.
+//
+// Goal: Measure the maximum distances/sec that can be computed.
+//
+// To achieve this, we make sure that:
+//
+// - Input data size is a multiple of the block tile size.
+//
+// - Perfect distribution of work between SMs, i.e. the number of block tiles is
+// a large multiple (num_waves) of the number of blocks (#SMs * occupancy).
+//
+// - Multiple iterations over Kblk are executed (num_k_iters).
+struct throughput_param {
+ int m, n, k;
+ bool use_1x_tfloat;
+};
+
+const std::vector<throughput_param> throughput_params{
+ {1024, 1024, 1024, true},
+ {1024, 1024, 1 << 11, true},
+ {1024, 1024, 1 << 12, true},
+ {1024, 1024, 1 << 13, true},
+ {1024, 1 << 14, 1024, true},
+ {1024, 1 << 14, 1 << 11, true},
+ {1024, 1 << 14, 1 << 12, true},
+ {1024, 1 << 14, 1 << 13, true},
+
+ {1024, 1024, 1024, false},
+ {1024, 1024, 1 << 11, false},
+ {1024, 1024, 1 << 12, false},
+ {1024, 1024, 1 << 13, false},
+ {1024, 1 << 14, 1024, false},
+ {1024, 1 << 14, 1 << 11, false},
+ {1024, 1 << 14, 1 << 12, false},
+ {1024, 1 << 14, 1 << 13, false},
+};
+
+struct throughput_cutlass : public fixture {
+ const throughput_param p;
+
+ throughput_cutlass(const throughput_param& p_) : p(p_) {}
+
+ void run_benchmark(::benchmark::State& state) override
+ {
+ size_t m = p.m;
+ size_t n = p.n;
+ size_t k = p.k;
+
+ // DataT, OutT, IdxT, etc, are defined in tuned_kernel.cuh
+ rmm::device_uvector<DataT> x_vec(m * k, stream);
+ rmm::device_uvector<DataT> y_vec(n * k, stream);
+ rmm::device_uvector<DataT> x_norm_vec(m, stream);
+ rmm::device_uvector<DataT> y_norm_vec(n, stream);
+ rmm::device_uvector<OutT> out_vec(m * n, stream);
+
+ auto x = x_vec.data();
+ auto y = y_vec.data();
+ auto x_norm = x_norm_vec.data();
+ auto y_norm = y_norm_vec.data();
+ auto out = out_vec.data();
+ FinOpT fin_op{};
+
+ // Create kernel parameter struct. Flip x and y if column major.
+ IdxT ldx = row_major ? k : m;
+ IdxT ldy = row_major ? k : n;
+ IdxT ld_out = row_major ? n : m;
+
+ // Template parameters of pairwise_matrix_params are defined in kernel.cuh
+ pairwise_matrix_params kparams{
+ IdxT(m), IdxT(n), IdxT(k), ldx, ldy, ld_out, x, y, x_norm, y_norm, out, fin_op, row_major};
+
+ // Run benchmark
+ loop_on_state(state, [&]() { launch_kernel(kparams, p.use_1x_tfloat, stream); });
+
+ // Report metrics. We don't report flop/s because we do not know for each
+ // distance operation how many flops it costs. For L2_unexp and l1, we can
+ // double this number to get the flop/s. For l2 expanded, core_ops/s should
+ // equal flop/s (modulo the sqrt and subtracting from the norm).
+ size_t num_core_ops = m * n * k;
+ size_t read_elts = n * k + m * k;
+ size_t write_elts = m * n;
+
+ state.counters["m"] = benchmark::Counter(m);
+ state.counters["n"] = benchmark::Counter(n);
+ state.counters["k"] = benchmark::Counter(k);
+ state.counters["1xtfloat"] = benchmark::Counter(p.use_1x_tfloat);
+
+ state.counters["core_ops/s"] = benchmark::Counter(num_core_ops,
+ benchmark::Counter::kIsIterationInvariantRate,
+ benchmark::Counter::OneK::kIs1000);
+
+ state.counters["BW"] = benchmark::Counter(write_elts * sizeof(OutT) + read_elts * sizeof(DataT),
+ benchmark::Counter::kIsIterationInvariantRate,
+ benchmark::Counter::OneK::kIs1000);
+ }
+};
+
+RAFT_BENCH_REGISTER(throughput_cutlass, "", throughput_params);
+
+} // namespace raft::bench::distance::tune_cutlass

cpp/bench/prims/distance/tune_pairwise/kernel_cutlass.cu

@@ -0,0 +1,48 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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 "kernel_cutlass.cuh"
+#include <raft/distance/detail/distance_ops/all_ops.cuh> // distance_op
+#include <raft/distance/detail/pairwise_matrix/dispatch_sm80.cuh>
+#include <raft/distance/detail/pairwise_matrix/params.cuh>
+#include <raft/distance/distance_types.hpp> // Compute_options
+#include <raft/util/arch.cuh> // raft::util::arch::SM_compute_arch
+
+namespace raft::bench::distance::tune_cutlass {
+
+// Distance op
+using OpT = raft::distance::detail::ops::l2_exp_distance_op<DataT, AccT, IdxT>;
+
+constexpr bool perform_sqrt = false;
+OpT distance_op{perform_sqrt};
+
+// Architecture
+namespace arch = raft::util::arch;
+constexpr auto sm_compat_range = arch::SM_range(arch::SM_80(), arch::SM_future());
+
+void launch_kernel(pairwise_matrix_params params, bool use_1x_tfloat, cudaStream_t stream)
+{
+ raft::distance::detail::pairwise_matrix_sm80_dispatch(
+ distance_op,
+ use_1x_tfloat ? raft::distance::Compute_options::Fast_Reduced_Precision
+ : raft::distance::Compute_options::Fast_Similar_Precision,
+ params,
+ sm_compat_range,
+ stream);
+ RAFT_CUDA_TRY(cudaGetLastError());
+}
+
+} // namespace raft::bench::distance::tune_cutlass

cpp/bench/prims/distance/tune_pairwise/kernel_cutlass.cuh

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include <raft/core/operators.hpp> // raft::identity_op
+#include <raft/distance/detail/pairwise_matrix/params.cuh> // pairwise_matrix_params
+
+namespace raft::bench::distance::tune_cutlass {
+
+// Launch one specific kernel with the following template parameters
+constexpr bool row_major = true;
+using DataT = float;
+using AccT = float;
+using OutT = DataT;
+using IdxT = int;
+
+using FinOpT = raft::identity_op;
+
+using pairwise_matrix_params =
+ raft::distance::detail::pairwise_matrix_params<IdxT, DataT, OutT, FinOpT>;
+
+void launch_kernel(pairwise_matrix_params params, bool use_1x_tfloat, cudaStream_t stream);
+
+} // namespace raft::bench::distance::tune_cutlass