Visualize stablehlo/vhlo graphs in tt-explorer

[tapspatel]

Requesting additional feature in tt-explorer to be able to visualize stablehlo graphs. Reason for request is when bringing up models from tt-torch or tt-xla, they go through sub dialect steps, and it much easier to be able to visualize the decomposition of the graph between all the different dialects

Sample vhlo graph

module @jit_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  vhlo.func_v1 @main(%arg0: !vhlo.tensor_v1<1x128x!vhlo.f32_v1>, %arg1: !vhlo.tensor_v1<128x64x!vhlo.f32_v1>, %arg2: !vhlo.tensor_v1<1x64x!vhlo.f32_v1>, %arg3: !vhlo.tensor_v1<1x64x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<1x64x!vhlo.f32_v1>) {
    %0 = "vhlo.dot_general_v2"(%arg0, %arg1) <{accumulation_type = #vhlo.type_v1<!vhlo.none_v1>, allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>, lhs_batching_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>, lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>, lhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>, lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>, num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>, precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>, rhs_batching_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>, rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>, rhs_contracting_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>, rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>}> : (!vhlo.tensor_v1<1x128x!vhlo.f32_v1>, !vhlo.tensor_v1<128x64x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x64x!vhlo.f32_v1>
    %1 = "vhlo.multiply_v1"(%0, %arg2) : (!vhlo.tensor_v1<1x64x!vhlo.f32_v1>, !vhlo.tensor_v1<1x64x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x64x!vhlo.f32_v1>
    %2 = "vhlo.add_v1"(%1, %arg3) : (!vhlo.tensor_v1<1x64x!vhlo.f32_v1>, !vhlo.tensor_v1<1x64x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x64x!vhlo.f32_v1>
    "vhlo.return_v1"(%2) : (!vhlo.tensor_v1<1x64x!vhlo.f32_v1>) -> ()
  } {arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"mhlo.sharding"> = #vhlo.string_v1<"{replicated}">}>, #vhlo.dict_v1<{#vhlo.string_v1<"mhlo.sharding"> = #vhlo.string_v1<"{replicated}">}>, #vhlo.dict_v1<{#vhlo.string_v1<"mhlo.sharding"> = #vhlo.string_v1<"{replicated}">}>, #vhlo.dict_v1<{#vhlo.string_v1<"mhlo.sharding"> = #vhlo.string_v1<"{replicated}">}>]>, res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"jax.result_info"> = #vhlo.string_v1<"">}>]>, sym_visibility = #vhlo.string_v1<"public">}
}

Sample shlo graph

module @jit_module attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<1x128xf32> {mhlo.sharding = "{replicated}"}, %arg1: tensor<128x64xf32> {mhlo.sharding = "{replicated}"}, %arg2: tensor<1x64xf32> {mhlo.sharding = "{replicated}"}, %arg3: tensor<1x64xf32> {mhlo.sharding = "{replicated}"}) -> (tensor<1x64xf32> {jax.result_info = ""}) {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x128xf32>, tensor<128x64xf32>) -> tensor<1x64xf32>
    %1 = stablehlo.multiply %0, %arg2 : tensor<1x64xf32>
    %2 = stablehlo.add %1, %arg3 : tensor<1x64xf32>
    return %2 : tensor<1x64xf32>
  }
}

[odjuricicTT]

@tapspatel This is something that is planned and requested by forge-fe as well. I agree that it would be really useful, not just for frontends, but for debugging tt-mlir internal passes as well.

The main purpose of tt-explorer is to be a tool used for model performance tuning so unfortunately this is not a high priority for us. We don't have enough resources to tackle this any time soon. If you or someone else has the time, all contributions are welcome :)