Adding transpose_twice test

[mtopalovicTT]

Adding transpose test

[rpavlovicTT]

Fixes

From this IR

%0 = tensor.empty
%1 = ttir.transpose(%arg1, %0)
%2 = tensor.empty
%3 = ttir.transpose(%0, %2)

To this IR

%1 = ttir.transpose(%arg1, %0)
%0 = tensor.empty
%2 = tensor.empty
%3 = ttir.transpose(%0, %2)

Also fixes how we update results types of ops. For example if we take initial IR from above, layout pass will update %0 attr with correct memory space and tile type. It will update all of it's uses to reflect correct attribute but it won't update %1 to have the same attribute as %0 which will lead to this error:

test/ttmlir/Dialect/TTNN/transpose/transpose_twice.mlir:7:10: error: 'ttir.transpose' op expected type of operand #1 ('tensor<128x64xbf16, #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<4x2x!tt.tile<32x32, bf16>, #tt.memory_space<dram>>>>') to match type of corresponding result ('tensor<128x64xbf16, #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<128x64xbf16, #tt.memory_space<system>>>>')
    %1 = "ttir.transpose"(%arg0, %0) <{dim0 = 0 : si32, dim1 = 1 : si32, operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<64x128xbf16>, tensor<128x64xbf16>) -> tensor<128x64xbf16>

And here is dump of whole IR before the error

#layout = #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<64x128xbf16, #system>>
#layout1 = #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<4x2x!tt.tile<32x32, bf16>, #dram>>
#layout2 = #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<2x4x!tt.tile<32x32, bf16>, #dram>>
#layout3 = #tt.layout<(d0, d1) -> (d0, d1), undef, <1x1>, memref<128x64xbf16, #system>>
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<64x128xbf16, #layout>) -> tensor<64x128xbf16, #layout>, sym_name = "forward"}> ({
  ^bb0(%arg0: tensor<64x128xbf16, #layout>):
    %0 = "tensor.empty"() : () -> tensor<128x64xbf16, #layout1>
    %1 = "tensor.empty"() : () -> tensor<64x128xbf16, #layout2>
    %2 = "ttir.to_layout"(%arg0, %1) : (tensor<64x128xbf16, #layout>, tensor<64x128xbf16, #layout2>) -> tensor<64x128xbf16, #layout2>
    %3 = "ttir.transpose"(%2, %0) <{dim0 = 0 : si32, dim1 = 1 : si32, operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<64x128xbf16, #layout2>, tensor<128x64xbf16, #layout1>) -> tensor<128x64xbf16, #layout3>
    %4 = "tensor.empty"() : () -> tensor<64x128xbf16, #layout2>
    %5 = "ttir.transpose"(%0, %4) <{dim0 = 1 : si32, dim1 = 0 : si32, operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<128x64xbf16, #layout1>, tensor<64x128xbf16, #layout2>) -> tensor<64x128xbf16, #layout2>
    %6 = "tensor.empty"() : () -> tensor<64x128xbf16, #layout>
    %7 = "ttir.to_layout"(%5, %6) : (tensor<64x128xbf16, #layout2>, tensor<64x128xbf16, #layout>) -> tensor<64x128xbf16, #layout>
    "func.return"(%7) : (tensor<64x128xbf16, #layout>) -> ()
  }) : () -> ()

You can see that %0 layout and %3 layout don't match.

I will experiment some more to see if we hit similar errors with different IR.

fixes #558

I don't get how is this correct ?

%1 = ttir.transpose(%arg1, %0)
%0 = tensor.empty

[mtopalovicTT]

@rpavlovicTT comment is misleading, updated

[mtopalovicTT]

@nsmithtt Fixed bugs... Thanks again for pointing out

[mtopalovicTT]

Need repro from FE team for resnet transpose which didn't work....