Fix TTIR to TTNN conversion for all gather

[gfengTT]

This fixes the incorrect conversion of the ttnn all_gather op. It should not use the output of an EmptyOp as the input to all_gather.

Testing with:

./build/bin/ttmlir-opt --ttir-to-ttnn-backend-pipeline test/ttmlir/Dialect/TTNN/ccl/all_gather.mlir

Before:

  func.func @forward(%arg0: tensor<1x1x32x32xbf16, #layout>) -> tensor<1x1x32x128xbf16, #layout1> {
    %0 = "ttnn.get_device"() <{mesh_shape = #ttnn<mesh_shape 1x1>}> : () -> !tt.device<#device>
    %1 = "ttnn.to_device"(%arg0, %0) <{memory_config = #ttnn.memory_config<<interleaved>, <dram>, <<1x1>>>}> : (tensor<1x1x32x32xbf16, #layout>, !tt.device<#device>) -> tensor<1x1x32x32xbf16, #layout2>
    %2 = "ttnn.to_layout"(%1) <{layout = #ttnn.layout<tile>}> : (tensor<1x1x32x32xbf16, #layout2>) -> tensor<1x1x32x32xbf16, #layout2>
    "ttnn.dealloc"(%2) : (tensor<1x1x32x32xbf16, #layout2>) -> ()
    "ttnn.dealloc"(%1) : (tensor<1x1x32x32xbf16, #layout2>) -> ()
    %3 = "ttnn.empty"(%0) <{dtype = #tt.supportedDataTypes<bf16>, layout = #ttnn.layout<tile>, memory_config = #ttnn.memory_config<<interleaved>, <dram>, <<1x1>>>, shape = #ttnn.shape<1x1x32x32>}> : (!tt.device<#device>) -> tensor<1x1x32x32xbf16, #layout2>
    %4 = "ttnn.all_gather"(%3) <{dim = 3 : si32, num_links = 1 : si32}> : (tensor<1x1x32x32xbf16, #layout2>) -> tensor<1x1x32x128xbf16, #layout3>
    "ttnn.dealloc"(%3) : (tensor<1x1x32x32xbf16, #layout2>) -> ()
    %5 = "ttnn.from_device"(%4) : (tensor<1x1x32x128xbf16, #layout3>) -> tensor<1x1x32x128xbf16, #layout1>
    "ttnn.dealloc"(%4) : (tensor<1x1x32x128xbf16, #layout3>) -> ()
    %6 = "ttnn.to_layout"(%5) <{layout = #ttnn.layout<row_major>}> : (tensor<1x1x32x128xbf16, #layout1>) -> tensor<1x1x32x128xbf16, #layout1>
    "ttnn.dealloc"(%5) : (tensor<1x1x32x128xbf16, #layout1>) -> ()
    return %6 : tensor<1x1x32x128xbf16, #layout1>
  }

New:

  func.func @forward(%arg0: tensor<1x1x32x32xbf16, #layout>) -> tensor<1x1x32x128xbf16, #layout1> {
    %0 = "ttnn.get_device"() <{mesh_shape = #ttnn<mesh_shape 1x1>}> : () -> !tt.device<#device>
    %1 = "ttnn.to_device"(%arg0, %0) <{memory_config = #ttnn.memory_config<<interleaved>, <dram>, <<1x1>>>}> : (tensor<1x1x32x32xbf16, #layout>, !tt.device<#device>) -> tensor<1x1x32x32xbf16, #layout2>
    %2 = "ttnn.to_layout"(%1) <{layout = #ttnn.layout<tile>}> : (tensor<1x1x32x32xbf16, #layout2>) -> tensor<1x1x32x32xbf16, #layout2>
    "ttnn.dealloc"(%1) : (tensor<1x1x32x32xbf16, #layout2>) -> ()
    %3 = "ttnn.all_gather"(%2) <{dim = 3 : si32, num_links = 1 : si32}> : (tensor<1x1x32x32xbf16, #layout2>) -> tensor<1x1x32x128xbf16, #layout3>
    "ttnn.dealloc"(%2) : (tensor<1x1x32x32xbf16, #layout2>) -> ()
    %4 = "ttnn.from_device"(%3) : (tensor<1x1x32x128xbf16, #layout3>) -> tensor<1x1x32x128xbf16, #layout1>
    "ttnn.dealloc"(%3) : (tensor<1x1x32x128xbf16, #layout3>) -> ()
    %5 = "ttnn.to_layout"(%4) <{layout = #ttnn.layout<row_major>}> : (tensor<1x1x32x128xbf16, #layout1>) -> tensor<1x1x32x128xbf16, #layout1>
    "ttnn.dealloc"(%4) : (tensor<1x1x32x128xbf16, #layout1>) -> ()
    return %5 : tensor<1x1x32x128xbf16, #layout1>
  }

[nsmithtt]

Hold off on landing this, I need to understand this better. I think we probably want to keep all gather dps.

[nsmithtt]

nvm, OK looks good!

[nsmithtt]

FYI @wooseokTT