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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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.

Found by Coverity.

Assisted-by: Codex

---
Full diff: https://github.com/llvm/llvm-project/pull/220873.diff


3 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h (+22-8) 
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+8-3) 
- (modified) mlir/test/Dialect/LLVMIR/nvvm-target-invalid.mlir (+8) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
index fa7eb4194dfac..99c5e0634eef2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMRequiresSMTraits.h
@@ -15,7 +15,9 @@
 
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StorageUniquerSupport.h"
-#include "llvm/ADT/StringExtras.h"
+
+#include <limits>
+#include <optional>
 
 namespace mlir {
 
@@ -62,16 +64,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("sm_"))
+      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 843344703d0a5..89f7672803509 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -7151,9 +7151,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");
   }
@@ -7163,7 +7168,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">] {
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/220873