Write a MLIR pass for dumping MLIR Function to .dot file

[Fare9]

Currently MLIR implements a ViewOpGraph that writes a Module as a .dot graph, since it is generic for all the the MLIR dialects the representation is based on connecting statements by their parameters and result values, so those statements which are not using values from other instructions are disconnected on the graph, and the output is a little bit weird.

For example, for the next module:

module @"LSimple;" {
  func.func @my_add(%arg0: i32, %arg1: i32) -> i32 {
    %0 = arith.addi %arg0, %arg1 : i32
    %1 = arith.muli %arg0, %0 : i32
    %2 = arith.divsi %1, %arg1 : i32
    %3 = MjolnIR.loadfield @"LSimple;" -> @test_field(0) : i16
    MjolnIR.storefield %3 : i16, @"LSimple;" -> @test_field(0)
    MjolnIR.nop
    %4 = arith.cmpi ne, %0, %2 : i32
    cf.cond_br %4, ^bb2(%2 : i32), ^bb1(%0 : i32)
  ^bb1(%5: i32):  // pred: ^bb0
    %6 = "MjolnIR.move"(%5) : (i32) -> i32
    cf.br ^bb3(%6 : i32)
  ^bb2(%7: i32):  // pred: ^bb0
    MjolnIR.nop
    MjolnIR.fallthrough ^bb3(%7 : i32)
  ^bb3(%8: i32):  // 2 preds: ^bb1, ^bb2
    MjolnIR.nop
    return %8 : i32
  }
}

The next .dot is dumped:

digraph G {
  compound = true;
  subgraph cluster_1 {
    v2 [label = " ", shape = plain];
    label = "builtin.module : ())\n\nsym_name: \"LSimple;\"";
    subgraph cluster_3 {
      v4 [label = " ", shape = plain];
      label = "";
      subgraph cluster_5 {
        v6 [label = " ", shape = plain];
        label = "func.func : ())\n\nfunction_type: (i32, i32) -> i32\nsym_name: \"my_add\"";
        subgraph cluster_7 {
          v8 [label = " ", shape = plain];
          label = "";
          v9 [label = "arg0", shape = ellipse];
          v10 [label = "arg1", shape = ellipse];
          v11 [label = "arith.addi : (i32))\n", shape = ellipse];
          v12 [label = "arith.muli : (i32))\n", shape = ellipse];
          v13 [label = "arith.divsi : (i32))\n", shape = ellipse];
          v14 [label = "MjolnIR.loadfield : (i16))\n\nfieldClass: @\"LSimple;\"\nfieldName: @test_field\nfieldRef: 0 : ui32", shape = ellipse];
          v15 [label = "MjolnIR.storefield : ())\n\nfieldClass: @\"LSimple;\"\nfieldName: @test_field\nfieldRef: 0 : ui32", shape = ellipse];
          v16 [label = "MjolnIR.nop : ())\n", shape = ellipse];
          v17 [label = "arith.cmpi : (i1))\n\npredicate: 1 : i64", shape = ellipse];
          v18 [label = "cf.cond_br : ())\n\noperand_segment_sizes: array<i32: 1, 1, 1>", shape = ellipse];
        }
        subgraph cluster_19 {
          v20 [label = " ", shape = plain];
          label = "";
          v21 [label = "arg0", shape = ellipse];
          v22 [label = "MjolnIR.move : (i32))\n", shape = ellipse];
          v23 [label = "cf.br : ())\n", shape = ellipse];
        }
        subgraph cluster_24 {
          v25 [label = " ", shape = plain];
          label = "";
          v26 [label = "arg0", shape = ellipse];
          v27 [label = "MjolnIR.nop : ())\n", shape = ellipse];
          v28 [label = "MjolnIR.fallthrough : ())\n", shape = ellipse];
        }
        subgraph cluster_29 {
          v30 [label = " ", shape = plain];
          label = "";
          v31 [label = "arg0", shape = ellipse];
          v32 [label = "MjolnIR.nop : ())\n", shape = ellipse];
          v33 [label = "func.return : ())\n", shape = ellipse];
        }
      }
    }
  }
  v9 -> v11 [label = "0", style = solid];
  v10 -> v11 [label = "1", style = solid];
  v9 -> v12 [label = "0", style = solid];
  v11 -> v12 [label = "1", style = solid];
  v12 -> v13 [label = "0", style = solid];
  v10 -> v13 [label = "1", style = solid];
  v14 -> v15 [label = "", style = solid];
  v11 -> v17 [label = "0", style = solid];
  v13 -> v17 [label = "1", style = solid];
  v17 -> v18 [label = "0", style = solid];
  v13 -> v18 [label = "1", style = solid];
  v11 -> v18 [label = "2", style = solid];
  v21 -> v22 [label = "", style = solid];
  v22 -> v23 [label = "", style = solid];
  v26 -> v28 [label = "", style = solid];
  v31 -> v33 [label = "", style = solid];
}

I think it's psosible to write a MLIR pass that given a MjolnIR module, it dumps a .dot file for each function with the shape of a Control Flow Graph (CFG).

[Fare9]

Started the MLIR Pass for writing the .dot file, the code can be found here: https://github.com/Fare9/KUNAI-static-analyzer/blob/refactoring/kunai-lib/MjolnIR/lib/Transforms/MjolnIRToOpGraph.cpp. An example of the output given is the next one:

digraph G {
  compound = true;
  subgraph cluster_1 {
    v2[label = " ", shape = plain];
    label = "func.func : ()\n\nfunction_type: (i32, i32) -> i32\nsym_name: \"my_add\"";
    subgraph cluster_3 {
      v4[label = " ", shape = plain];
      label = "BB";
      v5[label = "arith.addi : (i32)\n", shape = ellipse];
      v6[label = "arith.muli : (i32)\n", shape = ellipse];
      v5 -> v6;
      v7[label = "arith.divsi : (i32)\n", shape = ellipse];
      v6 -> v7;
      v8[label = "MjolnIR.loadfield : (i16)\n\nfieldClass: @\"LSimple;\"\nfieldName: @test_field\nfieldRef: 0 : ui32", shape = ellipse];
      v7 -> v8;
      v9[label = "MjolnIR.storefield : ()\n\nfieldClass: @\"LSimple;\"\nfieldName: @test_field\nfieldRef: 0 : ui32", shape = ellipse];
      v8 -> v9;
      v10[label = "MjolnIR.nop : ()\n", shape = ellipse];
      v9 -> v10;
      v11[label = "arith.cmpi : (i1)\n\npredicate: 1 : i64", shape = ellipse];
      v10 -> v11;
      v12[label = "cf.cond_br : ()\n\noperand_segment_sizes: array<i32: 1, 1, 1>", shape = ellipse];
      v11 -> v12;
    }
    subgraph cluster_13 {
      v14[label = " ", shape = plain];
      label = "BB";
      v15[label = "MjolnIR.move : (i32)\n", shape = ellipse];
      v16[label = "cf.br : ()\n", shape = ellipse];
      v15 -> v16;
    }
    subgraph cluster_17 {
      v18[label = " ", shape = plain];
      label = "BB";
      v19[label = "MjolnIR.nop : ()\n", shape = ellipse];
      v20[label = "MjolnIR.fallthrough : ()\n", shape = ellipse];
      v19 -> v20;
    }
    subgraph cluster_21 {
      v22[label = " ", shape = plain];
      label = "BB";
      v23[label = "MjolnIR.nop : ()\n", shape = ellipse];
      v24[label = "func.return : ()\n", shape = ellipse];
      v23 -> v24;
    }
  }
  v12 -> v19[style="solid,bold",color=green,weight=10,constraint=true];
  v12 -> v15[style="solid,bold",color=red,weight=10,constraint=true];
  v16 -> v23[style="solid,bold",color=blue,weight=10,constraint=true];
  v20 -> v23[style="solid,bold",color=black,weight=10,constraint=true];
}