diff --git a/mlir/unittests/Dialect/Rock/CMakeLists.txt b/mlir/unittests/Dialect/Rock/CMakeLists.txt index 5f2f03d432e7..fb5c2648b75d 100644 --- a/mlir/unittests/Dialect/Rock/CMakeLists.txt +++ b/mlir/unittests/Dialect/Rock/CMakeLists.txt @@ -3,6 +3,7 @@ set(ROCK_UNITTEST_SOURCES TosaUtilsTests.cpp loweringUtilsTests.cpp transformMapUtilsTests.cpp + TuningParamSpaceTests.cpp GeneralGemmParamsAttrTests.cpp AccelGemmParamsAttrTests.cpp GemmGemmParamsAttrTests.cpp diff --git a/mlir/unittests/Dialect/Rock/TuningParamSpaceTests.cpp b/mlir/unittests/Dialect/Rock/TuningParamSpaceTests.cpp new file mode 100644 index 000000000000..aa60a5c11bc1 --- /dev/null +++ b/mlir/unittests/Dialect/Rock/TuningParamSpaceTests.cpp @@ -0,0 +1,112 @@ +//===- TuningParamSpaceTests.cpp - Tests for Rock tuning space generation -===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Rock/IR/Rock.h" +#include "mlir/Dialect/Rock/Tuning/RockTuning.h" +#include "mlir/IR/DialectRegistry.h" +#include "mlir/Parser/Parser.h" +#include "gtest/gtest.h" + +using namespace mlir; +using namespace mlir::rock; + +namespace { + +static constexpr llvm::StringLiteral kNonAccelGemmModule = R"mlir( +module { + func.func @non_accel_gemm(%a: memref<1x384x3072xf16>, %b: memref<1x3072x768xf16>, %c: memref<1x384x768xf16>) attributes {kernel, mhal.arch = "amdgcn-amd-amdhsa:gfx900"} { + rock.gemm %c = %a * %b features = none storeMethod = set {arch = "amdgcn-amd-amdhsa:gfx900", perf_config = "v3:32,64,8,32,16,8,1,1,2,1,1"} : memref<1x384x768xf16> = memref<1x384x3072xf16> * memref<1x3072x768xf16> + return + } +} +)mlir"; + +static constexpr llvm::StringLiteral kAccelGemmModule = R"mlir( +module { + func.func @accel_gemm(%a: memref<1x384x3072xf16>, %b: memref<1x3072x768xf16>, %c: memref<1x384x768xf16>) attributes {kernel, mhal.arch = "amdgcn-amd-amdhsa:gfx90a:sramecc+:xnack-"} { + rock.gemm %c = %a * %b features = mfma|dot|atomic_add|atomic_add_f16 storeMethod = set {arch = "amdgcn-amd-amdhsa:gfx90a:sramecc+:xnack-", perf_config = "v4:64,64,8,32,32,16,4,4,1,2,0,0,1,1"} : memref<1x384x768xf16> = memref<1x384x3072xf16> * memref<1x3072x768xf16> + return + } +} +)mlir"; + +class TuningParamSpaceTest : public ::testing::Test { +public: + TuningParamSpaceTest() { + DialectRegistry registry; + registry.insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + } + +protected: + MLIRContext context; +}; + +TEST_F(TuningParamSpaceTest, GreedyFallsBackToExhaustiveForNonAccelGemm) { + OwningOpRef module = + parseSourceString(kNonAccelGemmModule, &context); + ASSERT_TRUE(module) << "Failed to parse non-accel test module"; + + TuningParamSpaceSettings settings{0, ""}; + std::unique_ptr tuningSpace( + createTunableParamSpace(*module, TuningParamSetKind::Greedy, settings)); + ASSERT_TRUE(tuningSpace); + EXPECT_FALSE(tuningSpace->tuningRange.empty()); + EXPECT_EQ(tuningSpace->effectiveKind, TuningParamSetKind::Exhaustive); + EXPECT_EQ(getNumberOfIterations(tuningSpace->effectiveKind), 1u); + EXPECT_FALSE(needToUpdateBest(tuningSpace->effectiveKind)); +} + +TEST_F(TuningParamSpaceTest, GreedyRemainsGreedyForAccelGemm) { + OwningOpRef module = + parseSourceString(kAccelGemmModule, &context); + ASSERT_TRUE(module) << "Failed to parse accel test module"; + + TuningParamSpaceSettings settings{0, ""}; + std::unique_ptr tuningSpace( + createTunableParamSpace(*module, TuningParamSetKind::Greedy, settings)); + ASSERT_TRUE(tuningSpace); + EXPECT_FALSE(tuningSpace->tuningRange.empty()); + EXPECT_EQ(tuningSpace->effectiveKind, TuningParamSetKind::Greedy); + EXPECT_EQ(getNumberOfIterations(tuningSpace->effectiveKind), 3u); + EXPECT_TRUE(needToUpdateBest(tuningSpace->effectiveKind)); +} + +TEST_F(TuningParamSpaceTest, + GreedyFallbackForNonAccelGemmIgnoresGreedyIterationSettings) { + OwningOpRef module = + parseSourceString(kNonAccelGemmModule, &context); + ASSERT_TRUE(module) << "Failed to parse non-accel test module"; + + TuningParamSpaceSettings settings{/*iteration=*/2, + /*winningConfig=*/"invalid-config"}; + std::unique_ptr tuningSpace( + createTunableParamSpace(*module, TuningParamSetKind::Greedy, settings)); + ASSERT_TRUE(tuningSpace); + EXPECT_FALSE(tuningSpace->tuningRange.empty()); + EXPECT_EQ(tuningSpace->effectiveKind, TuningParamSetKind::Exhaustive); + EXPECT_EQ(getNumberOfIterations(tuningSpace->effectiveKind), 1u); + EXPECT_FALSE(needToUpdateBest(tuningSpace->effectiveKind)); +} + +TEST(TuningParamSpaceHelpersTest, IterationHelpersMatchTuningKindContract) { + EXPECT_EQ(getNumberOfIterations(TuningParamSetKind::Quick), 1u); + EXPECT_EQ(getNumberOfIterations(TuningParamSetKind::Full), 1u); + EXPECT_EQ(getNumberOfIterations(TuningParamSetKind::Exhaustive), 1u); + EXPECT_EQ(getNumberOfIterations(TuningParamSetKind::Greedy), 3u); + + EXPECT_FALSE(needToUpdateBest(TuningParamSetKind::Quick)); + EXPECT_FALSE(needToUpdateBest(TuningParamSetKind::Full)); + EXPECT_FALSE(needToUpdateBest(TuningParamSetKind::Exhaustive)); + EXPECT_TRUE(needToUpdateBest(TuningParamSetKind::Greedy)); +} + +} // namespace