[spirv] Add canonicalization patterns for ops

[antiagainst]

We don't have much canonicalization patterns for SPIR-V ops thus far. It would be nice to gradually build them up. Interesting ones include but not limited to:

• Combine chained spv.AccessChains into one spv.AccessChain
• Convert spv.selection to spv.Select if possible
• Constant fold arithmetic / logical / cast / etc. ops
• Other peephole optimizations like converting chained spv.Bitcast ops into one, turning spv.LogicalNot (spv.IEqual ...) into spv.INotEqual, etc.
• ...

[antiagainst]

Ping @denis0x0D to see if you are potentially interested in helping. :)

[denis0x0D]

@antiagainst it would be great for me to take some of these tasks, thanks a lot!

[denis0x0D]

@antiagainst @MaheshRavishankar is it ok, if I take

• Convert spv.selection to spv.Select if possible

[antiagainst]

Sure, please feel free to take it! Thanks for the help as always, Denis! :)

[denis0x0D]

@antiagainst @MaheshRavishankar can you please specify regarding to "convert spv.selection to spv.Select", how actually deep analysis should be?

For example:

func @selection(%cond: i1) -> () {
...
  %zero = spv.constant 0: i32
  %one = spv.constant 1: i32
  %two = spv.constant 2: i32
  %x = spv.Variable init(%zero) : !spv.ptr<i32, Function>
  spv.selection {
   spv.BranchConditional %cond, ^then, ^else
  ^then:
    spv.Store "Function" %x, %one : i32
    spv.Branch ^merge
  ^else:
    spv.Store "Function" %x, %two : i32
    spv.Branch ^merge
  ^merge:
    spv._merge
  }
 ...
}

this code could be simplified as follows:

func @selection(%arg0: i1) -> () {
...
    %0 = spv.constant 0 : i32
    %1 = spv.constant 1 : i32
    %2 = spv.constant 2 : i32
    %3 = spv.Variable init(%0) : !spv.ptr<i32, Function>
    %4 = spv.Select %arg0, %1, %2 : i1, i32
    spv.Store "Function" %3, %4 : i32
  ...
}

this pattern is simple to detect and therefore to simplify the code, but there could situations that require more deep analysis for example

1. globalVariable modifiication, should it be analyzed and simplified if possible?
2. multiple load/store, should it be outlined from the region and then perform multiple spv.Select on it? Thanks.

[MaheshRavishankar]

This is just how I would approach it. You see the statements within the true and false branch. If you can find a statement that could be hoisted out of the two branches you do that. This might uncover more statements that can be hoisted out. So you have an outer loop that tries to hoist statements as long as there are no more statements that can be hoisted. The decision of what can be hoisted can be specified for each statement. For example, and arithmatic operation in both the branches where the value is coming from outside of the selection construct can be hoisted out. We can extend them for more statements as this pass is developed.

[denis0x0D]

@MaheshRavishankar thanks for answer! To clarify, in our approach, we could have a situation when some statements are not valid to hoist out of some branch, and therefore svp.selection is not fully replaced with a sequence of spv.Select? Do I get you right? Does this task require a full replacement of spv.selection by spv.Select or not?

By the way, I was thinking to analyze branches at first, for example: for each branch collect a set of "variables":

1. Defined out of the branch (spv.globalVariable, spv.Variable)
2. Modified inside the branch (by spv.Store)

then we can compare sets and if the sets are equal we can try to perform replacement of spv.selection to a sequence of spv.Select/spv.Store, but I may miss something here.

Thanks.

[antiagainst]

My suggestion is that we start small by covering the case suggested by Denis https://github.com/tensorflow/mlir/issues/165#issuecomment-540093667. :) This issue is about canonicalization pattern, which is typically simple and straightforward. Hoisting is nice to have but I think it can be its own pass (so also covering spv.loop) and it should compose well with the spv.selection canonicalization here. What do you all think?

[MaheshRavishankar]

To clarify, in our approach, we could have a situation when some statements are not valid to hoist out of some branch, and therefore svp.selection is not fully replaced with a sequence of spv.Select? Do I get you right?

Thats right. I dont think you can do it for all statements. After repeated transformation, if the spv.selection construct is empty, we can delete that construct.

@denis0x0D : It might be simpler to start with arithmatic instructions. With store you need to reason about memory semantics (like handling Volatile, etc.) and make sure it is valid to hoist the store which is more involved. Better place to start would be arithmatic operations.

[MaheshRavishankar]

My suggestion is that we start small by covering the case suggested by Denis #165 (comment). :) This issue is about canonicalization pattern, which is typically simple and straightforward. Hoisting is nice to have but I think it can be its own pass (so also covering spv.loop) and it should compose well with the spv.selection canonicalization here. What do you all think?

Agreed. the analysis needed for deciding if an instruction can be hoisted or not can be shared (maybe it can even be something generic within MLIR itself since most, if not all, of the logic involved would not be specific to SPIR-V dialect).

Also agree that of the canonicalizations the selection -> Select is more involved. The other ones listed are simpler and provide a lot of value.

[denis0x0D]

@antiagainst @MaheshRavishankar thanks for the answers, sounds great to me about simple patterns, since "canonicalization" applies as part of the operation simplification https://github.com/tensorflow/mlir/blob/master/lib/Transforms/Canonicalizer.cpp#L50 Thanks!

[denis0x0D]

@antiagainst @MaheshRavishankar is it ok, if I take "Combine chained spv.AccessChains into one spv.AccessChain" ?

[antiagainst]

Sure, please go ahead. Thanks for the help as usual, @denis0x0D! :)

[denis0x0D]

@antiagainst @MaheshRavishankar does anyone work on "converting chained spv.Bitcast ops into one"?

[antiagainst]

Nope. Feel free to pick it up if you are interested! Thanks! :)

[denis0x0D]

@antiagainst @MaheshRavishankar does anyone work on "turning spv.LogicalNot (spv.IEqual ...) into spv.INotEqual, etc." ?

[antiagainst]

I don't think so. Feel free to take it. Thanks! :)