[Mlir-commits] [mlir] 27bb2a3 - [mlir][Pass] Handle escaped pipline option values (#97667)

[llvmlistbot at llvm.org]

Author: Nikhil Kalra
Date: 2024-07-05T09:49:34+02:00
New Revision: 27bb2a36401b89fdce68e0e427493d645c42a218

URL: https://github.com/llvm/llvm-project/commit/27bb2a36401b89fdce68e0e427493d645c42a218
DIFF: https://github.com/llvm/llvm-project/commit/27bb2a36401b89fdce68e0e427493d645c42a218.diff

LOG: [mlir][Pass] Handle escaped pipline option values (#97667)

The PassRegistry parser properly handles escape tokens (', ", {}) when
parsing pass options from string but then does not strip the escape
tokens when providing the values back to the caller.

This change updates the parser such that escape tokens are properly
removed and whitespace is trimmed when extracting option values.

Added: 
    

Modified: 
    mlir/include/mlir/Pass/PassOptions.h
    mlir/lib/Pass/PassRegistry.cpp
    mlir/test/Pass/pipeline-options-parsing.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Pass/PassOptions.h b/mlir/include/mlir/Pass/PassOptions.h
index b99d91f07dd1cc..6bffa84f7b16b4 100644
--- a/mlir/include/mlir/Pass/PassOptions.h
+++ b/mlir/include/mlir/Pass/PassOptions.h
@@ -60,6 +60,20 @@ template <typename ParserT>
 static void printOptionValue(raw_ostream &os, const bool &value) {
   os << (value ? StringRef("true") : StringRef("false"));
 }
+template <typename ParserT>
+static void printOptionValue(raw_ostream &os, const std::string &str) {
+  // Check if the string needs to be escaped before writing it to the ostream.
+  const size_t spaceIndex = str.find_first_of(' ');
+  const size_t escapeIndex =
+      std::min({str.find_first_of('{'), str.find_first_of('\''),
+                str.find_first_of('"')});
+  const bool requiresEscape = spaceIndex < escapeIndex;
+  if (requiresEscape)
+    os << "{";
+  os << str;
+  if (requiresEscape)
+    os << "}";
+}
 template <typename ParserT, typename DataT>
 static std::enable_if_t<has_stream_operator<DataT>::value>
 printOptionValue(raw_ostream &os, const DataT &value) {

diff  --git a/mlir/lib/Pass/PassRegistry.cpp b/mlir/lib/Pass/PassRegistry.cpp
index f8149673a40939..483cbe80faba6a 100644
--- a/mlir/lib/Pass/PassRegistry.cpp
+++ b/mlir/lib/Pass/PassRegistry.cpp
@@ -218,6 +218,20 @@ parseNextArg(StringRef options) {
   auto extractArgAndUpdateOptions = [&](size_t argSize) {
     StringRef str = options.take_front(argSize).trim();
     options = options.drop_front(argSize).ltrim();
+    // Handle escape sequences
+    if (str.size() > 2) {
+      const auto escapePairs = {std::make_pair('\'', '\''),
+                                std::make_pair('"', '"'),
+                                std::make_pair('{', '}')};
+      for (const auto &escape : escapePairs) {
+        if (str.front() == escape.first && str.back() == escape.second) {
+          // Drop the escape characters and trim.
+          str = str.drop_front().drop_back().trim();
+          // Don't process additional escape sequences.
+          break;
+        }
+      }
+    }
     return str;
   };
   // Try to process the given punctuation, properly escaping any contained

diff  --git a/mlir/test/Pass/pipeline-options-parsing.mlir b/mlir/test/Pass/pipeline-options-parsing.mlir
index 34313cd5566033..50f08241ee5cfa 100644
--- a/mlir/test/Pass/pipeline-options-parsing.mlir
+++ b/mlir/test/Pass/pipeline-options-parsing.mlir
@@ -7,6 +7,10 @@
 // RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(func.func(test-options-pass{string-list=a list=1,2,3,4 string-list=b,c list=5 string-list=d string=nested_pipeline{arg1=10 arg2=" {} " arg3=true}}))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_1 %s
 // RUN: mlir-opt %s -verify-each=false -test-options-pass-pipeline='list=1 string-list=a,b enum=one' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_2 %s
 // RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{enum=one list=1,2,3,4})))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_3 %s
+// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{enum=one list=1,2,3,4 string="foobar"})))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_4 %s
+// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{enum=one list=1,2,3,4 string="foo bar baz"})))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_5 %s
+// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{enum=one list=1,2,3,4 string={foo bar baz}})))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_5 %s
+// RUN: mlir-opt %s -verify-each=false -pass-pipeline='builtin.module(builtin.module(func.func(test-options-pass{list=3}), func.func(test-options-pass{enum=one list=1,2,3,4 string=foo"bar"baz})))' -dump-pass-pipeline 2>&1 | FileCheck --check-prefix=CHECK_6 %s
 
 // CHECK_ERROR_1: missing closing '}' while processing pass options
 // CHECK_ERROR_2: no such option test-option
@@ -17,3 +21,6 @@
 // CHECK_1: test-options-pass{enum=zero list=1,2,3,4,5 string=nested_pipeline{arg1=10 arg2=" {} " arg3=true} string-list=a,b,c,d}
 // CHECK_2: test-options-pass{enum=one list=1 string= string-list=a,b}
 // CHECK_3: builtin.module(builtin.module(func.func(test-options-pass{enum=zero list=3 string= }),func.func(test-options-pass{enum=one list=1,2,3,4 string= })))
+// CHECK_4: builtin.module(builtin.module(func.func(test-options-pass{enum=zero list=3 string= }),func.func(test-options-pass{enum=one list=1,2,3,4 string=foobar })))
+// CHECK_5: builtin.module(builtin.module(func.func(test-options-pass{enum=zero list=3 string= }),func.func(test-options-pass{enum=one list=1,2,3,4 string={foo bar baz} })))
+// CHECK_6: builtin.module(builtin.module(func.func(test-options-pass{enum=zero list=3 string= }),func.func(test-options-pass{enum=one list=1,2,3,4 string=foo"bar"baz })))