2424namespace {
2525constexpr int PBT_ITERATIONS = 100 ;
2626
27+ // Tensor Core 默认 fallback(使用 bank-conflict-free)
28+ auto defaultTensorCoreFallback () {
29+ return [](const float * A, const float * B, float * C, int M, int K, int N,
30+ cudaStream_t stream) {
31+ launch_bank_conflict_free_sgemm<32 >(A, B, C, M, K, N, stream);
32+ };
33+ }
34+ } // namespace
35+
2736std::vector<std::tuple<int , int , int >> getStandardDimensions () {
2837 return {
2938 {1 , 1 , 1 }, {2 , 2 , 2 }, {4 , 4 , 4 }, {16 , 16 , 16 },
@@ -246,7 +255,10 @@ TEST_P(TensorCoreSGEMMTest, FastPathCorrectnessProperty) {
246255 ASSERT_TRUE (tensorCoreDimensionsSupported (M_ , K_ , N_ ));
247256
248257 VerifyResult result = runKernelAndCompare (
249- [&] { launch_tensor_core_sgemm (d_A_->get (), d_B_->get (), d_C_->get (), M_ , K_ , N_ ); },
258+ [&] {
259+ launch_tensor_core_sgemm_with_fallback (d_A_->get (), d_B_->get (), d_C_->get (),
260+ M_ , K_ , N_ , defaultTensorCoreFallback ());
261+ },
250262 kTensorCoreVerifyTolerance );
251263
252264 EXPECT_TRUE (result.passed ) << " TensorCore SGEMM fast path failed for dimensions " << M_ << " x"
@@ -261,7 +273,10 @@ class TensorCoreFallbackTest : public SGEMMKernelTest {};
261273
262274TEST_P (TensorCoreFallbackTest, NonAlignedInputsFallbackSafely) {
263275 VerifyResult result = runKernelAndCompare (
264- [&] { launch_tensor_core_sgemm (d_A_->get (), d_B_->get (), d_C_->get (), M_ , K_ , N_ ); },
276+ [&] {
277+ launch_tensor_core_sgemm_with_fallback (d_A_->get (), d_B_->get (), d_C_->get (),
278+ M_ , K_ , N_ , defaultTensorCoreFallback ());
279+ },
265280 kStandardVerifyTolerance );
266281
267282 EXPECT_TRUE (result.passed ) << " TensorCore fallback failed for dimensions " << M_ << " x" << K_
@@ -272,9 +287,10 @@ INSTANTIATE_TEST_SUITE_P(TensorCoreFallbackDimensions, TensorCoreFallbackTest,
272287 ::testing::ValuesIn (getTensorCoreFallbackDimensions()));
273288
274289TEST (TensorCoreWrapperTest, ZeroSizeInputsReturnSafely) {
275- EXPECT_NO_THROW (launch_tensor_core_sgemm (nullptr , nullptr , nullptr , 0 , 16 , 16 ));
276- EXPECT_NO_THROW (launch_tensor_core_sgemm (nullptr , nullptr , nullptr , 16 , 0 , 16 ));
277- EXPECT_NO_THROW (launch_tensor_core_sgemm (nullptr , nullptr , nullptr , 16 , 16 , 0 ));
290+ auto fallback = [](const float *, const float *, float *, int , int , int , cudaStream_t) {};
291+ EXPECT_NO_THROW (launch_tensor_core_sgemm_with_fallback (nullptr , nullptr , nullptr , 0 , 16 , 16 , fallback));
292+ EXPECT_NO_THROW (launch_tensor_core_sgemm_with_fallback (nullptr , nullptr , nullptr , 16 , 0 , 16 , fallback));
293+ EXPECT_NO_THROW (launch_tensor_core_sgemm_with_fallback (nullptr , nullptr , nullptr , 16 , 16 , 0 , fallback));
278294}
279295
280296class DimensionInvarianceTest : public ::testing::Test {
0 commit comments