[mlir][Pass] Handle spaces in pipeline strings

[boschmitt]

Summary

An user might want to add extra spaces for better readability, e.g:

mypm = pm.PassManager.parse(f"""builtin.module(
    mypass1,
    func.func(mypass2,mypass3)
)""")

GitHub issue #59151

The parser was not taking into account the possibility of spaces after )or }

[boschmitt]

Rebase. Fix typo in comments.

[myhsu]

Could you add a test (not a multi-line one, just a pipeline string with white spaces) in test/Pass/pipeline-parsing.mlir as well? It should be easy

[boschmitt]

Add another test and rebase

[boschmitt]

In D142821#4089010, @myhsu wrote:

Could you add a test (not a multi-line one, just a pipeline string with white spaces) in test/Pass/pipeline-parsing.mlir as well? It should be easy

Thanks for the review. I'm fairly new to the codebase, so I didn't know where else this was being tested.

[myhsu]

LGTM thanks :-)

[rkayaith]

There's still some cases where whitespace causes the parsing to fail:

" builtin.module()" - error: can't run ' builtin.module' pass manager on 'builtin.module' op
"builtin.module ()" - error: can't run 'builtin.module ' pass manager on 'builtin.module' op
"builtin.module() " - error: expected pass pipeline to be wrapped with the anchor operation type, e.g. 'builtin.module(...)'

[boschmitt]

Handlle more cases

[boschmitt]

Small fix to avoid unecessary work

[myhsu]

Let us know if you need helps pushing the patching

mlir/test/Pass/pipeline-parsing.mlir

6

super nit: could you move this one line up so that CHECK_ERROR_X tests are put together

[boschmitt]

Small fix

[boschmitt]

Yes, I need help pushing the patch as I don't have commit access.

[myhsu]

In D142821#4091312, @boschmitt wrote:

Yes, I need help pushing the patch as I don't have commit access.

Cool, please provide the email you want to use. btw I can help you to push D142818 as well

[boschmitt]

In D142821#4091666, @myhsu wrote:

Cool, please provide the email you want to use. btw I can help you to push D142818 as well

Thank you for the help. The e-mail is [e-mail]

mlir/lib/Pass/PassRegistry.cpp

 //===- PassRegistry.cpp - Pass Registration Utilities ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include <utility>
 #include <optional>
+#include <utility>
 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/ManagedStatic.h"
if (sep == '{') {
       if (close == StringRef::npos) {
         return errorHandler(
             /*rawLoc=*/text.data() - 1,
             "missing closing '}' while processing pass options");
       }
       pipeline.back().options = text.substr(0, close);
       text = text.substr(close + 1);
 
+      // Consume space characters that an user might add for readability.
+      text = text.ltrim();
+
       // Skip checking for '(' because nested pipelines cannot have options.
     } else if (sep == '(') {
       text = text.substr(1);
 
       // Push the inner pipeline onto the stack to continue processing.
       pipelineStack.push_back(&pipeline.back().innerPipeline);
       continue;
     }
 
     // When handling the close parenthesis, we greedily consume them to avoid
     // empty strings in the pipeline.
     while (text.consume_front(")")) {
       // If we try to pop the outer pipeline we have unbalanced parentheses.
       if (pipelineStack.size() == 1)
         return errorHandler(/*rawLoc=*/text.data() - 1,
                             "encountered extra closing ')' creating unbalanced "
                             "parentheses while parsing pipeline");
 
       pipelineStack.pop_back();
+      // Consume space characters that an user might add for readability.
+      text = text.ltrim();
     }
 
     // Check if we've finished parsing.
     if (text.empty())
       break;
 
     // Otherwise, the end of an inner pipeline always has to be followed by
     // a comma, and then we can continue.
LogicalResult mlir::parsePassPipeline(StringRef pipeline, OpPassManager &pm,
   };
   if (failed(pipelineParser.addToPipeline(pm, errorHandler)))
     return failure();
   return success();
 }
 
 FailureOr<OpPassManager> mlir::parsePassPipeline(StringRef pipeline,
                                                  raw_ostream &errorStream) {
+  pipeline = pipeline.trim();
   // Pipelines are expected to be of the form `<op-name>(<pipeline>)`.
   size_t pipelineStart = pipeline.find_first_of('(');
   if (pipelineStart == 0 || pipelineStart == StringRef::npos ||
       !pipeline.consume_back(")")) {
     errorStream << "expected pass pipeline to be wrapped with the anchor "
                    "operation type, e.g. 'builtin.module(...)'";
     return failure();
   }
 
-  StringRef opName = pipeline.take_front(pipelineStart);
+  StringRef opName = pipeline.take_front(pipelineStart).rtrim();
   OpPassManager pm(opName);
   if (failed(parsePassPipeline(pipeline.drop_front(1 + pipelineStart), pm,
                                errorStream)))
     return failure();
   return pm;
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/Pass/pipeline-parsing.mlir

 // RUN: mlir-opt %s -mlir-disable-threading -pass-pipeline='builtin.module(builtin.module(test-module-pass,func.func(test-function-pass)),func.func(test-function-pass),func.func(cse,canonicalize))' -verify-each=false -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck %s
 // RUN: mlir-opt %s -mlir-disable-threading -test-textual-pm-nested-pipeline -verify-each=false -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck %s --check-prefix=TEXTUAL_CHECK
 // RUN: mlir-opt %s -mlir-disable-threading -pass-pipeline='builtin.module(builtin.module(test-module-pass),any(test-interface-pass),any(test-interface-pass),func.func(test-function-pass),any(canonicalize),func.func(cse))' -verify-each=false -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck %s --check-prefix=GENERIC_MERGE_CHECK
+// RUN: mlir-opt %s -mlir-disable-threading -pass-pipeline=' builtin.module ( builtin.module( func.func( test-function-pass, print-op-stats{ json=false } ) ) ) ' -verify-each=false -mlir-timing -mlir-timing-display=tree 2>&1 | FileCheck %s --check-prefix=PIPELINE_STR_WITH_SPACES_CHECK
 // RUN: not mlir-opt %s -pass-pipeline='any(builtin.module(test-module-pass)' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_1 %s
 // RUN: not mlir-opt %s -pass-pipeline='builtin.module(test-module-pass))' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_2 %s
 // RUN: not mlir-opt %s -pass-pipeline='any(builtin.module()()' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_3 %s
 // RUN: not mlir-opt %s -pass-pipeline='any(,)' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_4 %s
 // RUN: not mlir-opt %s -pass-pipeline='func.func(test-module-pass)' 2>&1 | FileCheck --check-prefix=CHECK_ERROR_5 %s
 
 // CHECK_ERROR_1: encountered unbalanced parentheses while parsing pipeline
 // CHECK_ERROR_2: encountered extra closing ')' creating unbalanced parentheses while parsing pipeline
 // CHECK_ERROR_3: expected ',' after parsing pipeline
 // CHECK_ERROR_4: does not refer to a registered pass or pass pipeline

 // TEXTUAL_CHECK: Pipeline Collection : ['builtin.module', 'func.func']
 // TEXTUAL_CHECK-NEXT:   'func.func' Pipeline
 // TEXTUAL_CHECK-NEXT:     TestFunctionPass
 // TEXTUAL_CHECK-NEXT:   'builtin.module' Pipeline
 // TEXTUAL_CHECK-NEXT:     TestModulePass
 // TEXTUAL_CHECK-NEXT:     'func.func' Pipeline
 // TEXTUAL_CHECK-NEXT:       TestFunctionPass
 
+// PIPELINE_STR_WITH_SPACES_CHECK:   'builtin.module' Pipeline
+// PIPELINE_STR_WITH_SPACES_CHECK-NEXT:   'func.func' Pipeline
+// PIPELINE_STR_WITH_SPACES_CHECK-NEXT:     TestFunctionPass
+// PIPELINE_STR_WITH_SPACES_CHECK-NEXT:     PrintOpStats
+
 // Check that generic pass pipelines are only merged when they aren't
 // going to overlap with op-specific pipelines.
 // GENERIC_MERGE_CHECK:      Pipeline Collection : ['builtin.module', 'any']
 // GENERIC_MERGE_CHECK-NEXT:   'any' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:     TestInterfacePass
 // GENERIC_MERGE_CHECK-NEXT:   'builtin.module' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:     TestModulePass
 // GENERIC_MERGE_CHECK-NEXT: 'any' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:   TestInterfacePass
 // GENERIC_MERGE_CHECK-NEXT: 'func.func' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:   TestFunctionPass
 // GENERIC_MERGE_CHECK-NEXT: 'any' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:   Canonicalizer
 // GENERIC_MERGE_CHECK-NEXT: 'func.func' Pipeline
 // GENERIC_MERGE_CHECK-NEXT:   CSE
 // GENERIC_MERGE_CHECK-NEXT:     (A) DominanceInfo

mlir/test/python/pass_manager.py

else:
       log("Exception not produced")
 
     # A registered pass should parse successfully.
     pm = PassManager.parse("builtin.module(func.func(print-op-stats{json=false}))")
     # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
     log("Roundtrip: ", pm)
 run(testParseSuccess)
 
+# Verify successful round-trip.
+# CHECK-LABEL: TEST: testParseSpacedPipeline
+def testParseSpacedPipeline():
+  with Context():
+    # A registered pass should parse successfully even if has extras spaces for readability
+    pm = PassManager.parse("""builtin.module(
+        func.func( print-op-stats{ json=false } )
+    )""")
+    # CHECK: Roundtrip: builtin.module(func.func(print-op-stats{json=false}))
+    log("Roundtrip: ", pm)
+run(testParseSpacedPipeline)
+
 # Verify failure on unregistered pass.
 # CHECK-LABEL: TEST: testParseFail
 def testParseFail():
   with Context():
     try:
       pm = PassManager.parse("any(unknown-pass)")
     except ValueError as e:
       #      CHECK: ValueError exception: MLIR Textual PassPipeline Parser:1:1: error: