[Mlir-commits] [mlir] [mlir][GPUDialect] Add cmdOption suffix consumer in GpuModuleToBinary Pass (PR #127646)

[Zichen Lu]

https://github.com/MikaOvO created https://github.com/llvm/llvm-project/pull/127646

Add cmdOption suffix consumer function in GpuModuleToBinary Pass, which can tokenize and remove a specific suffix of cmdOption.

>From a123e2336780a767641eaec82694abc435023ad2 Mon Sep 17 00:00:00 2001
From: Zichen Lu <mikaovo2000 at gmail.com>
Date: Tue, 18 Feb 2025 22:51:58 +0800
Subject: [PATCH] [mlir][GPUDialect] Add cmdOption suffix consumer in
 GpuModuleToBinary Pass

---
 .../mlir/Dialect/GPU/IR/CompilationInterfaces.h      |  6 ++++++
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp               | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
index 9a890ae24d8fc..139360f8bd3fc 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
@@ -79,6 +79,12 @@ class TargetOptions {
   std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
   tokenizeCmdOptions() const;
 
+  /// Returns a tokenization of the substr of the command line options that
+  /// starts with `startsWith` and ends with end of the command line options and
+  /// consumes it.
+  std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
+  tokenizeAndRemoveSuffixCmdOptions(llvm::StringRef startsWith);
+
   /// Returns the compilation target.
   CompilationTarget getCompilationTarget() const;
 
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 1bdeb3e356f4b..976432ea37120 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2591,6 +2591,18 @@ TargetOptions::tokenizeCmdOptions() const {
   return tokenizeCmdOptions(cmdOptions);
 }
 
+std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
+TargetOptions::tokenizeAndRemoveSuffixCmdOptions(llvm::StringRef startsWith) {
+  size_t startPos = cmdOptions.find(startsWith);
+  if (startPos == std::string::npos)
+    return {llvm::BumpPtrAllocator(), SmallVector<const char *>()};
+
+  auto tokenized =
+      tokenizeCmdOptions(cmdOptions.substr(startPos + startsWith.size()));
+  cmdOptions.resize(startPos);
+  return tokenized;
+}
+
 MLIR_DEFINE_EXPLICIT_TYPE_ID(::mlir::gpu::TargetOptions)
 
 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.cpp.inc"