dialects: (affine) No custom format for `affine.for`

[zhen8838]

I am trying to use xdsl to parse a mlir file:

module {
  func.func @main(%arg0: memref<8x128x384xf32>, %arg1: memref<8x384x512xf32>, %arg2: memref<8x128x512xf32>, %arg3: memref<8x512x64xf32>, %arg4: memref<8x128x64xf32>) {
    affine.for %arg5 = 0 to 8 {
      affine.for %arg6 = 0 to 128 {
        affine.for %arg7 = 0 to 512 {
          affine.for %arg8 = 0 to 384 {
            %0 = affine.load %arg0[%arg5, %arg6, %arg8] : memref<8x128x384xf32>
            %1 = affine.load %arg1[%arg5, %arg8, %arg7] : memref<8x384x512xf32>
            %2 = affine.load %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>
            %3 = arith.mulf %0, %1 : f32
            %4 = arith.addf %2, %3 : f32
            affine.store %4, %arg2[%arg5, %arg6, %arg7] : memref<8x128x512xf32>
          }
        }
      }
    }
}

here is my code:

mctx = MLContext()
mctx.load_dialect(Affine)
mctx.load_dialect(Builtin)
mctx.load_dialect(MemRef)
mctx.load_dialect(Func)
mod = None
p = None
with open("test1.mlir") as f:
    p = Parser(mctx, f.read())
    p.parse_module()

but I got error:

ParseError: <unknown>:3:15
    affine.for %arg5 = 0 to 8 {
               ^^^^^
               Operation affine.for does not have a custom format.

I found the xdsl needs to override the parse method manually. How can I use the mlir ODS to generate this method implement?

[superlopuh]

The easiest way to start is usually using the assembly format API, our assembly format closely follows the MLIR one, so it's sometimes possible to copy/paste (without the operation name), and it often works. If it doesn't, then we usually use the parse and print methods, as in the scf.for operation, which should have a similar structure. A PR implementing the custom format would be fantastic, if you don't mind implementing and testing it. Alternatively, we should be able to parse the generic format already, you can get mlir-opt to use it by adding the --mlir-print-op-generic argument.

[superlopuh]

I just checked the documentation, and it looks like there's no assembly format for this operation, and a custom parse and print functions are the way to go.

[zhen8838]

thanks for your answer. I will try to customize those functions.

[superlopuh]

Let's keep the issue around until the fix is merged, just to keep track of it. Can you please comment with #fixes #2673 when/if you open the PR?

[zhen8838]

Alright, I have encountered problems when I try to implement the affine. load parse function, can you help me?

the affine load's syntax:

operation ::= ssa-id `=` `affine.load` ssa-use `[` multi-dim-affine-map-of-ssa-ids `]` `:` memref-type

test.mlir:

%0 = affine.load %arg0[%arg5 + 3, %arg6, %arg8] : memref<8x128x384xf32>

my question is how to parse multi-dim-affine-map-of-ssa-ids, I think the mlir needs to parse the index expression and use it to construct the affine map attribute. But in the xdsl, it seems like no parse function can do this stuff.

here is my code and the error message:

    @classmethod
    def parse(cls: Self, parser: Parser) -> Self:
        memref = parser.parse_operand()

        def parse_index():
            return parser.parse_operand() # <-- need to parse operand expreesion
        indices = parser.parse_comma_separated_list(
            Parser.Delimiter.SQUARE, parse_index)
        parser.parse_punctuation(':')
        rtype: MemRefType = parser.parse_type()
        return cls(memref, indices, None, rtype.get_element_type())

xdsl.utils.exceptions.ParseError: <unknown>:7:41
            %0 = affine.load %arg0[%arg5 + 3, %arg6, %arg8] : memref<8x128x384xf32>
                                         ^
                                         Expected ']'

[superlopuh]

A search for multi-dim-affine-map-of-ssa-ids in the LLVM project got 0 hits apart from the documentation, I'm guessing it's stale. Similarly, parseAffineLoad yields nothing. I'm not sure where else to look. Without the MLIR implementation as reference, I'd say that the best bet is to look at our AffineParser, and to borrow the _parse_multi_affine_expr implementation as a private helper in your file. I'd copy/paste the code, and replace the bits that parse identifiers with parsers for SSA values. I think copy/pasting here would be better than making the existing parser parametric, but that's somewhat subjective. If you happen to find the way MLIR does this that would be very helpful.

[superlopuh]

It might be the case that the affine expressions supported in affine custom syntax are much simpler than what affine maps support, in which case writing a new parsing function from scratch might be easier.

[zhen8838]

ok, I will be writing a new parsing function to deal with this.