Sync View Support for Split Operations

[Ha-Eyz]

I am currently working on a tool that visualizes IR code at different compilation stages, allowing users to compare the IR across these stages. However, I encountered an issue in scenarios where a single operation is split into multiple operations during transformation.

Here is a concrete example with Tosa:

From this original IR:

#loc2 = loc("/app/example.mlir":1:28)
module {
  func.func @test_abs_scalar(%arg0: tensor<f32> loc("/app/example.mlir":1:28)) -> tensor<f32> {
    %0 = tosa.abs %arg0 : (tensor<f32>) -> tensor<f32> loc(#loc3)
    return %0 : tensor<f32> loc(#loc4)
  } loc(#loc1)
} loc(#loc)
#loc = loc("/app/example.mlir":0:0)
#loc1 = loc("/app/example.mlir":1:1)
#loc3 = loc("/app/example.mlir":3:8)
#loc4 = loc("/app/example.mlir":5:2)

To this transformed IR:

#loc2 = loc("/app/example.mlir":1:28)
#loc3 = loc("/app/example.mlir":3:8)
#map = affine_map<() -> ()>
module {
  func.func @test_abs_scalar(%arg0: tensor<f32> loc("/app/example.mlir":1:28)) -> tensor<f32> {
    %0 = tensor.empty() : tensor<f32> loc(#loc3)
    %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = []} ins(%arg0 : tensor<f32>) outs(%0 : tensor<f32>) {
    ^bb0(%in: f32 loc("/app/example.mlir":1:28), %out: f32 loc("/app/example.mlir":3:8)):
      %2 = math.absf %in : f32 loc(#loc3)
      linalg.yield %2 : f32 loc(#loc3)
    } -> tensor<f32> loc(#loc3)
    return %1 : tensor<f32> loc(#loc4)
  } loc(#loc1)
} loc(#loc)
#loc = loc("/app/example.mlir":0:0)
#loc1 = loc("/app/example.mlir":1:1)
#loc4 = loc("/app/example.mlir":5:2)

(Here is a link to Compiler Explorer with the example: https://godbolt.org/z/YbhrP8hsb)

In this example, the tosa.abs operation at #loc3 is split into three separate operations in the transformed IR. To accurately showcase these changes in the sync view, it would require associating one operation in the original IR with multiple operations in the transformed IR. However, there doesn’t appear to be support for linking one operation to multiple operations. As far as I understand, sync view currently only allows a synchronised view between: a node to another node, a layer to another layer, a node to a layer, and a layer to a node.

Another example would be cases, where multiple operations are fused into a single operation.

So I wanted to ask, if sync view is capable of displaying such scenarios and, if it is not capable of doing so, if there are plans to implement a feature capable of displaying this scenario, or expand sync view to display it.

[pkgoogle]

Hi @Ha-Eyz, if you ever get this tool done -- I think integrating with model-explorer would be amazing. Have you tried Custom node matching?

https://github.com/google-ai-edge/model-explorer/wiki/2.-User-Guide#synchronize-navigation

You can use tuples as keys and values... let me know if that works.

[Ha-Eyz]

Hi @pkgoogle, I apologize for the confusion.

I have experimented with the custom node matching, but unfortunately, I couldn’t achieve the behavior I was aiming for. Specifically, I wasn’t able to sync multiple "keys" or multiple "values." As mentioned in my initial post, I was only able to achieve the following mappings:

• 1 Node to 1 other Node
• 1 Layer to 1 other Layer
• 1 Node to 1 Layer
• 1 Layer to 1 Node

I tried various configurations in the JSON file for mapping, but none of them worked as intended. Therefore, I wanted to reiterate my question: is this type of functionality possible at all, or are there any plans to introduce a feature that would support such a use case?

In case that it is possible, I wanted to ask, if you could expand the user guide with an example on how to do it.

Regarding the integration with Model Explorer, I regret to inform you that it’s currently intended as an internal tool and there are no immediate plans to make it publicly available, as far as I know.

[pkgoogle]

No worries, thanks for the extra information, I just wanted to check first to ensure if we haven't implemented this already. @jinjingforever, can you please take a look? Thanks.