Canonicalization of TTIR ops #1670
Conversation
- added InvolutionTrait - added IdempotentTrait - TODO: address other folders - TODO: refactor check into the separate function - TODO: added -canonicalize into the pipeline - TODO: add tests for involution and idempotence
- fixed some tests that failed due to canonicalization
- fixed test for fork-join: changed relu to gelu
- added a traits on few other ops as well - TODO: tests for ReverseOp canonicalization
- few op traits modifed - TODO: add tests for bitwise xor op canonicalization
- Broadcast noop folder
- op should be replaced instead of producer op
- tests for PermuteOp canonicalization
- some clean up
- improved comments in few places
azecevicTT
requested review from
nobradovictt,
odjuricicTT,
sdjordjevicTT,
svuckovicTT,
mtopalovicTT,
jserbedzijaTT,
jnie-TT and
nsmithtt
as code owners
| @@ -41,6 +43,28 @@ def TTIR_Dialect : Dialect { | |||
| //===----------------------------------------------------------------------===// | |||
|
|
|||
| class TTIR_Op<string mnemonic, list<Trait> traits = []> : | |||
| Op<TTIR_Dialect, mnemonic, !listconcat(traits, [Pure])>; | |||
| Op<TTIR_Dialect, mnemonic, !listconcat([Pure], traits)>; | |||
There was a problem hiding this comment.
I've changed the order of listconcatarguments everywhere, because of dependent trait list for traits like TTIR_Involution. Otherwise it fails long before compiling (it check for existence of this traits while still parsing .td file), so it doesn't see a trait/interface like DestinationStyleOpInterface on ops that has it, because it builds list in bottom-up fashion ([trait_of_op, trait_of_ops_class, trait_op_ops_class_class,...]).
From what I saw it doesn't make any difference in generated C++ code. @nsmithtt Did you have an experience with this?
| @@ -1,18 +1,5 @@ | |||
| // RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s | |||
| module { | |||
| func.func @linear_1d_1d(%arg0: tensor<128xbf16>, %arg1: tensor<128xbf16>) -> tensor<1xbf16> { | |||
There was a problem hiding this comment.
There have been a few tests that have failed due to canonicalization changing the graph. I have removed them and wrote tests that specifically check for the correctness of rewriting.
| @@ -27,15 +27,15 @@ module attributes {} { | |||
| // CHECK: %{{.*}} = "ttnn.relu"{{.*}} -> tensor<64x64xbf16, #[[LAYOUT_3]]> | |||
| %1 = "ttir.relu"(%arg0, %0) <{operandSegmentSizes = array<i32: 1, 1>}> : (tensor<64x64xbf16>, tensor<64x64xbf16>) -> tensor<64x64xbf16> | |||
| %2 = tensor.empty() : tensor<64x64xbf16> | |||
| %3 = "ttir.relu"(%1, %2) <{operandSegmentSizes = array<i32: 1, 1>}> : (tensor<64x64xbf16>, tensor<64x64xbf16>) -> tensor<64x64xbf16> | |||
| %3 = "ttir.gelu"(%1, %2) <{operandSegmentSizes = array<i32: 1, 1>}> : (tensor<64x64xbf16>, tensor<64x64xbf16>) -> tensor<64x64xbf16> | |||
There was a problem hiding this comment.
I've synced offline with @fbajraktariTT regarding this, due to RELU being idempotent op, two consecutive RELUs result in rewriting, so I've just replaced one RELU with GELU, as concrete op wasn't a concern in this test.
Canonicalization can be roughly divided into three forms:
The first one allows us to decleratively add folding for a large class of ops. While traits like
Involutionalready exists inMLIRinfrastructure it doesn't account for DPS, so I added a new one that takes care of it.The second one allows us in practice to define some simple graph rewritings where we replace the producing value of op with some existing value (or constant).
The third one gives us the most freedom, where we can do arbitrary graph rewritings, we usually use it when we have to create a new op during rewriting.
I added a
canonicalizepass both before and afterttir-to-ttir-decomposition-passinttir-to-ttnn-backend-pipeline, as there are many graph rewritings during that pass, so we might benefit form canonicalization both before and after.I plan to cover one big part of canonicalization in the future, and that's constant folding. I will also write a short document on adding a canonicalization for new and existing ops. While this PR covers a lot of patterns, it's definitely not an exhaustive list. This MLIR doc is already a great source of information, I just believe we might also benefit from the additional context of TTIR dialect.
This PR should cover a big part of #1264, as I said the missing part is constant folding.