How to lower Ops that depend on other Ops for context

MLIR
1

Hi everyone,

I want to build a hiearchical IR where Ops depend on surrounding Ops for context. To give an example, I would have an op MyDialect.open_file that gets a filename as an Attribute and has a region attached to it. Within this region (directly or further nested), the Ops MyDialect.write and MyDialect.read can appear which don’t have a file pointer argument or attribute, instead they automatically write into the file that was opened in the enclosing open_file region.

Pseudo-IR:

module {
func.func @hello() -> () {
  mydialect.open_file () {"filename" = "test.txt"} {
    affine.for %arg0 = 0 to 10 {
      mydialect.write () {"text" = "hello"} 
    }
  }
}
}

I want to lower this to something like this:

module {
func.func @hello() -> () {
  //mydialect.open_file () {"filename" = "test.txt"} */
  %fileid = lowlevel.open ()  {"filename" = "test.txt"}
  affine.for %arg0 = 0 to 10 {
    //mydialect.write () {"text" = "hello"} 
    lowlevel.write (%fileid) {"text" = "hello"} 
  }
  //}
}
}

My question is about the approach to lowering such a construct. My first idea was using a RewritePattern for each of the mydialect Ops and run a Conversion Pass with these.
When using a RewritePattern that matches mydialect.write first, I can find the surrounding open_file Op using getParent() repeatedly but the fileid which I need for the lowering to lowlevel.write has not materialized yet.
When using a RewritePattern that matches mydialect.open_file first, I can materialize the fileid by lowering to lowlevel.open but I cannot pass the information about which SSA Value contains the fileid to the other RewritePattern that will be called afterwards to lower mydialect.write.

The other idea would be to use a single RewritePattern that matches mydialect.open_file. There, I lower to lowlevel.open and then walk the attached Region, manually converting/lowering all the other mydialect Ops that I encounter using the materialized fileid.

The second approach seems to work fine but I feel like there must be a more idiomatic way to do this (maybe some way of passing a state to a Pattern?). Or maybe the whole concept for my IR is flawed :smiley: . Thank you for any hints and advice!

2

It is possible to conceptualize this “high-level” IR via side effects: open_file has a write effect on some “implicit file name” resource, and write has a “read” effect on the same resource. Although this design feels like it tries to do a “language” choice that makes it easier for a human to write and harder for a compiler to reason about, it is not invalid per se.

I would suggest doing the conversion outside of the pattern infrastructure, patterns have weak support for anything that concerns regions. Just a preorder walk of the IR that injects lowlevel.open in front of each mydialect.open_file, and uses the result of the closest lowlevel.open in write conversion.

It is technically possible to have a shared state between patterns by having all of them take a reference to the state in their constructor. This is used, e.g, for LLVMTypeConverter. However, you don’t really control the order of pattern application by the rewriting infra and should not rely on it.

3

That’s a good point. The file example was my attempt to simplify my application to the basic principle of having to look at other Ops for context during lowering. I’ll try to elaborate a little more and would be interested if you think that concern still applies.
I’m trying to create a more general context Op (instead of open_file) where the resources made available to the Op’s Body Region could be statically allocated MemRefs but theoretically also file pointers or other things, but every instance of the Op would be allocated the same amount. Then an Op like mydialect.write (ins size: i64) would always write from the Block’s MemRef to the Block’s file pointer. The idea is that Ops could by design only use the resources allocated to their enclosing context Op. When optimizing the interior of the region, it would (supposedly) be easier to reason about MemRef or general resource access because there is always the predefined amount of resources available within the Region. When optimizing the code outside of the context Op, the compiler could treat the whole context Region as a black box with the previously allocated resource footprint.

I’ll try that, thank you!

4

I’d consider modeling this through explicit block arguments rather than side effects. Something like:

my.context_defining_operation() attributes { resource_kinds = [a, b, c] } {
^bb0(%resource_a: !my.a, %resource_b: !my.b, %resource_c: !my.c):
  my.access_resource %resource_a, %something_else
}

This could provide both the opaque reasoning about resource use from the outside, and precise information on which resource is needed where from the inside.