[mlir][amdgpu] Improve Chipset version utility (#106169)

* Fix an OOB access
* Add comparison operators
* Add documentation
* Add unit tests

Author: kuhar

mlir/include/mlir/Dialect/AMDGPU/Utils/Chipset.h

@@ -9,19 +9,41 @@
 #define MLIR_DIALECT_AMDGPU_UTILS_CHIPSET_H_
 
 #include "mlir/Support/LLVM.h"
+#include <utility>
 
-namespace mlir {
-namespace amdgpu {
+namespace mlir::amdgpu {
+
+/// Represents the amdgpu gfx chipset version, e.g., gfx90a, gfx942, gfx1103.
+/// Note that the leading digits form a decimal number, while the last two
+/// digits for a hexadecimal number. For example:
+///   gfx942  --> major = 9, minor = 0x42
+///   gfx90a  --> major = 9, minor = 0xa
+///   gfx1103 --> major = 10, minor = 0x3
 struct Chipset {
   Chipset() = default;
   Chipset(unsigned majorVersion, unsigned minorVersion)
       : majorVersion(majorVersion), minorVersion(minorVersion){};
+
+  /// Parses the chipset version string and returns the chipset on success, and
+  /// failure otherwise.
   static FailureOr<Chipset> parse(StringRef name);
 
-  unsigned majorVersion = 0;
-  unsigned minorVersion = 0;
+  friend bool operator==(const Chipset &lhs, const Chipset &rhs) {
+    return lhs.majorVersion == rhs.majorVersion &&
+           lhs.minorVersion == rhs.minorVersion;
+  }
+  friend bool operator!=(const Chipset &lhs, const Chipset &rhs) {
+    return !(lhs == rhs);
+  }
+  friend bool operator<(const Chipset &lhs, const Chipset &rhs) {
+    return std::make_pair(lhs.majorVersion, lhs.minorVersion) <
+           std::make_pair(rhs.majorVersion, rhs.minorVersion);
+  }
+
+  unsigned majorVersion = 0; // The major version (decimal).
+  unsigned minorVersion = 0; // The minor version (hexadecimal).
 };
-} // end namespace amdgpu
-} // end namespace mlir
+
+} // namespace mlir::amdgpu
 
 #endif

mlir/lib/Dialect/AMDGPU/Utils/Chipset.cpp

@@ -1,4 +1,4 @@
-//===- Chipset.cpp - AMDGPU Chipset version struct parsing -----------===//
+//===- Chipset.cpp - AMDGPU Chipset version struct parsing ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -7,22 +7,26 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
-#include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 
-using namespace mlir;
-using namespace mlir::amdgpu;
+namespace mlir::amdgpu {
 
 FailureOr<Chipset> Chipset::parse(StringRef name) {
-  if (!name.starts_with("gfx"))
+  if (!name.consume_front("gfx"))
     return failure();
+  if (name.size() < 3)
+    return failure();
+
   unsigned major = 0;
   unsigned minor = 0;
-  StringRef majorRef = name.drop_front(3).drop_back(2);
+
+  StringRef majorRef = name.drop_back(2);
   StringRef minorRef = name.take_back(2);
   if (majorRef.getAsInteger(10, major))
     return failure();
   if (minorRef.getAsInteger(16, minor))
     return failure();
   return Chipset(major, minor);
 }
+
+} // namespace mlir::amdgpu

mlir/unittests/Dialect/AMDGPU/AMDGPUUtilsTest.cpp

@@ -0,0 +1,57 @@
+//===- AMDGPUUtilsTest.cpp - Unit tests for AMDGPU dialect utils ----------===//
+//
+// 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/AMDGPU/Utils/Chipset.h"
+#include "gtest/gtest.h"
+
+namespace mlir::amdgpu {
+namespace {
+
+TEST(ChipsetTest, Parsing) {
+  FailureOr<Chipset> chipset = Chipset::parse("gfx90a");
+  ASSERT_TRUE(succeeded(chipset));
+  EXPECT_EQ(chipset->majorVersion, 9u);
+  EXPECT_EQ(chipset->minorVersion, 0x0au);
+
+  chipset = Chipset::parse("gfx940");
+  ASSERT_TRUE(succeeded(chipset));
+  EXPECT_EQ(chipset->majorVersion, 9u);
+  EXPECT_EQ(chipset->minorVersion, 0x40u);
+
+  chipset = Chipset::parse("gfx1103");
+  ASSERT_TRUE(succeeded(chipset));
+  EXPECT_EQ(chipset->majorVersion, 11u);
+  EXPECT_EQ(chipset->minorVersion, 0x03u);
+}
+
+TEST(ChipsetTest, ParsingInvalid) {
+  EXPECT_TRUE(failed(Chipset::parse("navi33")));
+  EXPECT_TRUE(failed(Chipset::parse("rdna2")));
+  EXPECT_TRUE(failed(Chipset::parse("sm_80")));
+  EXPECT_TRUE(failed(Chipset::parse("GFX940")));
+  EXPECT_TRUE(failed(Chipset::parse("Gfx940")));
+  EXPECT_TRUE(failed(Chipset::parse("gfx9")));
+  EXPECT_TRUE(failed(Chipset::parse("gfx_940")));
+  EXPECT_TRUE(failed(Chipset::parse("gfx940_")));
+  EXPECT_TRUE(failed(Chipset::parse("gfxmeow")));
+  EXPECT_TRUE(failed(Chipset::parse("gfx1fff")));
+}
+
+TEST(ChipsetTest, Comparison) {
+  EXPECT_EQ(Chipset(9, 0x40), Chipset(9, 0x40));
+  EXPECT_NE(Chipset(9, 0x40), Chipset(9, 0x42));
+  EXPECT_NE(Chipset(9, 0x00), Chipset(10, 0x00));
+
+  EXPECT_LT(Chipset(9, 0x00), Chipset(10, 0x00));
+  EXPECT_LT(Chipset(9, 0x0a), Chipset(9, 0x42));
+  EXPECT_FALSE(Chipset(9, 0x42) < Chipset(9, 0x42));
+  EXPECT_FALSE(Chipset(9, 0x42) < Chipset(9, 0x40));
+}
+
+} // namespace
+} // namespace mlir::amdgpu

mlir/unittests/Dialect/AMDGPU/CMakeLists.txt

@@ -0,0 +1,7 @@
+add_mlir_unittest(MLIRAMDGPUTests
+  AMDGPUUtilsTest.cpp
+)
+target_link_libraries(MLIRAMDGPUTests
+  PRIVATE
+  MLIRAMDGPUUtils
+)

mlir/unittests/Dialect/CMakeLists.txt

@@ -6,6 +6,7 @@ target_link_libraries(MLIRDialectTests
   MLIRIR
   MLIRDialect)
 
+add_subdirectory(AMDGPU)
 add_subdirectory(ArmSME)
 add_subdirectory(Index)
 add_subdirectory(LLVMIR)

utils/bazel/llvm-project-overlay/mlir/unittests/BUILD.bazel

@@ -137,6 +137,19 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "amdgpu_tests",
+    size = "small",
+    srcs = glob([
+        "Dialect/AMDGPU/*.cpp",
+    ]),
+    deps = [
+        "//mlir:AMDGPUUtils",
+        "//third-party/unittest:gtest",
+        "//third-party/unittest:gtest_main",
+    ],
+)
+
 cc_test(
     name = "memref_tests",
     size = "small",