Skip to content

Commit

Permalink
[mlir][Pass] Handle escaped pipline option values (llvm#97667)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
nikalra authored Jul 5, 2024
1 parent a14a53e commit 27bb2a3
Show file tree
Hide file tree
Showing 3 changed files with 35 additions and 0 deletions.
14 changes: 14 additions & 0 deletions mlir/include/mlir/Pass/PassOptions.h
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
14 changes: 14 additions & 0 deletions mlir/lib/Pass/PassRegistry.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 7 additions & 0 deletions mlir/test/Pass/pipeline-options-parsing.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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 })))

0 comments on commit 27bb2a3

Please sign in to comment.