[Mlir-commits] [mlir] [mlir][NVVM] Reject malformed target chip names (PR #220873)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/220873

>From 4fcf451312e524bdd559da4571afa1acc8303eb8 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 2 Sep 2026 14:58:42 -0700
Subject: [PATCH] [mlir][NVVM] Reject malformed target chip names

Parse the complete sm_<version>[a|f] grammar before using the SM version.
This prevents a failed integer conversion from leaving the version
uninitialized and also rejects overflowing encodings.

Valgrind did not observe the uninitialized read in the optimized pre-fix
binary; the arbitrary value reached the existing diagnostic:

  error: unexpected error: Minimum NVVM target SM version is sm_20
  ==534292== ERROR SUMMARY: 0 errors from 0 contexts (suppressed: 0 from 0)

Found by Coverity.

Assisted-by: Codex
---
 .../Dialect/LLVMIR/NVVMRequiresSMTraits.h     | 32 ++++++++++++++-----
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp    | 11 +++++--
 .../Dialect/LLVMIR/nvvm-target-invalid.mlir   |  8 +++++
 3 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
index fa7eb4194dfac..d60b7c8bb6867 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
@@ -15,7 +15,10 @@
 
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StorageUniquerSupport.h"
-#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <limits>
+#include <optional>
 
 namespace mlir {
 
@@ -23,6 +26,7 @@ namespace NVVM {
 
 // Struct to store and check compatibility of SM versions.
 struct NVVMCheckSMVersion {
+  static constexpr llvm::StringLiteral kStreamingMultiprocessorPrefix = "sm_";
   static constexpr char kArchAcceleratedSuffix = 'a';
   static constexpr char kFamilySpecificSuffix = 'f';
 
@@ -62,16 +66,28 @@ struct NVVMCheckSMVersion {
 
   // Parses an SM version string and returns an equivalent full SM version
   // integer.
-  static unsigned getTargetFullSmVersionFromStr(StringRef smVersionString) {
-    bool isAA = smVersionString.back() == kArchAcceleratedSuffix;
-    bool isFS = smVersionString.back() == kFamilySpecificSuffix;
+  static std::optional<unsigned>
+  getTargetFullSmVersionFromStr(StringRef smVersionString) {
+    if (!smVersionString.consume_front(kStreamingMultiprocessorPrefix))
+      return std::nullopt;
+
+    unsigned suffix = 0;
+    if (!smVersionString.empty()) {
+      if (smVersionString.back() == kArchAcceleratedSuffix)
+        suffix = 3;
+      else if (smVersionString.back() == kFamilySpecificSuffix)
+        suffix = 2;
+      if (suffix)
+        smVersionString = smVersionString.drop_back();
+    }
 
     unsigned smVersion;
-    smVersionString.drop_front(3)
-        .take_while([](char c) { return llvm::isDigit(c); })
-        .getAsInteger(10, smVersion);
+    if (smVersionString.empty() ||
+        smVersionString.getAsInteger(10, smVersion) ||
+        smVersion > (std::numeric_limits<unsigned>::max() - suffix) / 10)
+      return std::nullopt;
 
-    return smVersion * 10 + (isAA ? 3 : 0) + (isFS ? 2 : 0);
+    return smVersion * 10 + suffix;
   }
 
   static bool isMinimumSMVersion(unsigned fullSmVersion) {
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index c6ed0a840a1b0..b8c4868206761 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -7148,9 +7148,14 @@ LogicalResult NVVMTargetAttr::verifyTarget(Operation *gpuModule) {
                      "NVVM target attribute must be attached to a GPU module");
   }
 
-  const unsigned targetFullSmVersion =
+  std::optional<unsigned> targetFullSmVersion =
       NVVMCheckSMVersion::getTargetFullSmVersionFromStr(getChip());
-  if (!NVVMCheckSMVersion::isMinimumSMVersion(targetFullSmVersion)) {
+  if (!targetFullSmVersion)
+    return emitError(gpuModule->getLoc())
+           << "invalid NVVM target chip \"" << getChip()
+           << "\", expected sm_<version>[a|f]";
+
+  if (!NVVMCheckSMVersion::isMinimumSMVersion(*targetFullSmVersion)) {
     return emitError(gpuModule->getLoc(),
                      "Minimum NVVM target SM version is sm_20");
   }
@@ -7160,7 +7165,7 @@ LogicalResult NVVMTargetAttr::verifyTarget(Operation *gpuModule) {
             if (auto reqOp = llvm::dyn_cast<NVVM::RequiresSMInterface>(op)) {
               const NVVMCheckSMVersion requirement =
                   reqOp.getRequiredMinSMVersion();
-              if (!requirement.isCompatibleWith(targetFullSmVersion)) {
+              if (!requirement.isCompatibleWith(*targetFullSmVersion)) {
                 op->emitOpError() << "is not supported on " << getChip();
                 return WalkResult::interrupt();
               }
diff --git a/mlir/test/Dialect/LLVMIR/nvvm-target-invalid.mlir b/mlir/test/Dialect/LLVMIR/nvvm-target-invalid.mlir
index 0f20e66dd2019..992b17564d35e 100644
--- a/mlir/test/Dialect/LLVMIR/nvvm-target-invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/nvvm-target-invalid.mlir
@@ -17,3 +17,11 @@ module attributes {gpu.container_module} {
   gpu.module @nvvm_target_invalid_opt_level [#nvvm.target<chip = "sm_90", O = 4>] {
   }
 }
+
+// -----
+
+module attributes {gpu.container_module} {
+  // expected-error @+1 {{invalid NVVM target chip "not-an-sm"}}
+  gpu.module @nvvm_target_invalid_chip [#nvvm.target<chip = "not-an-sm">] {
+  }
+}