Improper parsing of memref syntax

[3ulalia]

Hello,

I'm trying to do some operations on the MLIR generated for a simple C function:

void test(int *x, int *y) {
  *x = *y + 2;
}

Polygeist yields the following MLIR for this:

❯ ./bin/cgeist test/first.c -function=test -S
module { // (module attributes elided for space, and because mlir-opt chokes on them for some reason)
  func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %0 = affine.load %arg1[0] : memref<?xi32>
    %1 = arith.addi %0, %c2_i32 : i32
    affine.store %1, %arg0[0] : memref<?xi32>
    return
  }
}

Which can then be lowered by mlir-opt:

❯ mlir-opt --lower-affine first.mlir
module {
  func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %c0 = arith.constant 0 : index
    %0 = memref.load %arg1[%c0] : memref<?xi32>
    %1 = arith.addi %0, %c2_i32 : i32
    %c0_0 = arith.constant 0 : index
    memref.store %1, %arg0[%c0_0] : memref<?xi32>
    return
  }
}

However, attempting to parse this with pymlir makes it choke on the braces after the memref.load argument:

❯ python3
Python 3.12.9 (main, Feb  4 2025, 14:38:38) [GCC 14.2.1 20241116] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import mlir
>>> mlir.parse_string("""module {
...   func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
...     %c2_i32 = arith.constant 2 : i32
...     %c0 = arith.constant 0 : index
...     %0 = memref.load %arg1[%c0] : memref<?xi32>
...     %1 = arith.addi %0, %c2_i32 : i32
...     %c0_0 = arith.constant 0 : index
...     memref.store %1, %arg0[%c0_0] : memref<?xi32>
...     return
...   }
... }""")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 133, in parse_string
    return parser.parse(code)
           ^^^^^^^^^^^^^^^^^^
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 93, in parse
    tree = self.parser.parse(code)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/lark.py", line 655, in parse
    return self.parser.parse(text, start=start, on_error=on_error)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parser_frontends.py", line 104, in parse
    return self.parser.parse(stream, chosen_start, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/earley.py", line 280, in parse
    to_scan = self._parse(lexer, columns, to_scan, start_symbol)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 152, in _parse
    to_scan = scan(i, to_scan)
              ^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 125, in scan
    raise UnexpectedCharacters(stream, i, text_line, text_column, {item.expect.name for item in to_scan},
lark.exceptions.UnexpectedCharacters: No terminal matches '[' in the current parser context, at line 5 col 27

    %0 = memref.load %arg1[%c0] : memref<?xi32>
                          ^
Expected one of:
	* __ANON_4
	* __ANON_8
	* __ANON_0
	* COMMA
	* HASH
	* COLON
	* __ANON_7

Is this me being silly or is it an actual parsing error? I haven't spent enough time with pymlir to be able to tell.

Thanks!

[3ulalia]

Oh, interesting - it does work when I try to parse the affine version (output from Polygeist)

[3ulalia]

Eliminating the indices on the first argument to memref.load and memref.store makes it parse, but there's no guarantee how correctly.

That is, this IR works fine:

module {
  func.func @test(%arg0: memref<?xi32>, %arg1: memref<?xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
    %c2_i32 = arith.constant 2 : i32
    %c0 = arith.constant 0 : index
    %0 = memref.load %arg1 : memref<?xi32>
    %1 = arith.addi %0, %c2_i32 : i32
    %c0_0 = arith.constant 0 : index
    memref.store %1, %arg0 : memref<?xi32>
    return
  }
}

In this case, it'd probably be fine, because both of the indices are zero, but it appears that memref expects to be able to reference by index into a memory address and that PyMLIR doesn't agree.

[3ulalia]

I'll stop spamming this issue until someone who knows more about it than I do can weigh in :) As far as I can tell, this is probably an issue with the parsing logic somewhere, but I can't quite figure out where. Something about properly escaping square braces as terminals or something.

[tbennun]

Have you tried printing the ops in generic mode? It looks like dialect specific syntactic sugar that you can add a specialized parser for, as we do for some ops. However, for the most part it’s simpler to canonicalize the MLIR and parse the result.

[3ulalia]

Tried that, unfortunately. mlir-opt yields this MLIR:

❯ mlir-opt --lower-affine a | mlir-opt --mlir-print-op-generic
"builtin.module"() ({
  "func.func"() <{function_type = (memref<?xi32>, memref<?xi32>) -> (), sym_name = "test"}> ({
  ^bb0(%arg0: memref<?xi32>, %arg1: memref<?xi32>):
    %0 = "arith.constant"() <{value = 2 : i32}> : () -> i32
    %1 = "arith.constant"() <{value = 0 : index}> : () -> index
    %2 = "memref.load"(%arg1, %1) : (memref<?xi32>, index) -> i32
    %3 = "arith.addi"(%2, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
    %4 = "arith.constant"() <{value = 0 : index}> : () -> index
    "memref.store"(%3, %arg0, %4) : (i32, memref<?xi32>, index) -> ()
    "func.return"() : () -> ()
  }) {llvm.linkage = #llvm.linkage<external>} : () -> ()
}) : () -> ()

which pymlir also can't handle:

>>> mlir.parse_string(""""builtin.module"() ({
...   "func.func"() <{function_type = (memref<?xi32>, memref<?xi32>) -> (), sym_name = "test"}> ({
...   ^bb0(%arg0: memref<?xi32>, %arg1: memref<?xi32>):
...     %0 = "arith.constant"() <{value = 2 : i32}> : () -> i32
...     %1 = "arith.constant"() <{value = 0 : index}> : () -> index
...     %2 = "memref.load"(%arg1, %1) : (memref<?xi32>, index) -> i32
...     %3 = "arith.addi"(%2, %0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
...     %4 = "arith.constant"() <{value = 0 : index}> : () -> index
...     "memref.store"(%3, %arg0, %4) : (i32, memref<?xi32>, index) -> ()
...     "func.return"() : () -> ()
...   }) {llvm.linkage = #llvm.linkage<external>} : () -> ()
... }) : () -> ()""")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 133, in parse_string
    return parser.parse(code)
           ^^^^^^^^^^^^^^^^^^
  File "/nix/store/s0d7723ifv9s3i57rgysj3b4sj1c2lc6-python3.12-pymlir-0.5/lib/python3.12/site-packages/mlir/parser.py", line 93, in parse
    tree = self.parser.parse(code)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/lark.py", line 655, in parse
    return self.parser.parse(text, start=start, on_error=on_error)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parser_frontends.py", line 104, in parse
    return self.parser.parse(stream, chosen_start, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/earley.py", line 280, in parse
    to_scan = self._parse(lexer, columns, to_scan, start_symbol)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 152, in _parse
    to_scan = scan(i, to_scan)
              ^^^^^^^^^^^^^^^^
  File "/nix/store/41f9yi425bgkr4n49132whi04fn909d7-python3.12-lark-1.2.2/lib/python3.12/site-packages/lark/parsers/xearley.py", line 125, in scan
    raise UnexpectedCharacters(stream, i, text_line, text_column, {item.expect.name for item in to_scan},
lark.exceptions.UnexpectedCharacters: No terminal matches '<' in the current parser context, at line 2 col 17

  "func.func"() <{function_type = (memref<?xi32>, memref
                ^
Expected one of:
	* COLON
	* LSQB
	* LBRACE
	* LPAR

Seems like a different issue, but I'm not as familiar with genericized forms of MLIR, so I couldn't say for sure.

[tbennun]

This is a rather new syntax for func. I don’t think I’ve seen it before (the documentation shows another syntax: https://mlir.llvm.org/docs/Dialects/Func/#funcfunc-funcfuncop)

Which version/revision of MLIR are you using?

[3ulalia]

Good question.

❯ mlir-opt --version
LLVM (http://llvm.org/):
  LLVM version 19.1.7
  Optimized build.

I presume that's what you're after?

[tbennun]

This tag syntax (e.g. %1 = "arith.constant"() <{value = 0 : index}> : () -> index is definitely new (newer than what is supported in pyMLIR at least). If you say that it is released with the latest version of LLVM, then we must fix it.

Thanks for letting us know!

[3ulalia]

Absolutely! I'm not super familiar with the tag syntax, but glad to have pointed out the issue at least.

As a workaround in the meantime, how could I go about writing a specialized parser for the syntactic sugar in the specific (i.e. non-generic) version of the code?

[tbennun]

It should be straightforward, actually. We have some documentation on how it works here: https://github.com/spcl/pymlir/blob/master/doc/custom_dialect.rst

And here is a similar example from one of the built in dialects:

pymlir/mlir/dialects/standard.py

Line 78 in ab32731

_syntax_ = [

[3ulalia]

What's weird is that in that same file, it looks like the syntax of LoadOperation should support the braces:

pymlir/mlir/dialects/standard.py

Lines 99 to 104 in ab32731

	@dataclass
	class LoadOperation(DialectOp):
	arg: SsaUse
	index: List[SsaUse]
	type: mast.MemRefType
	_syntax_ = 'load {arg.ssa_use} [ {index.ssa_use_list} ] : {type.memref_type}'

Am I misreading it?

[3ulalia]

Aha! It's because the operations are being referenced with their qualified names (memref.load vs load). Easy fix, I'll make a PR. (the new genericized version should also be added!)

[tbennun]

This is exactly right. Before there were builtin dialects, there were more builtin ops. The fix should be to modernize the call syntax there to use the dialect.