SparseDataFlowAnalysis don't visit linalg.generic operation


I have a custom dataflow analysis based on the SparseDataFlowAnalysis class. My pass is loading both the mlir::dataflow::DeadCodeAnalysis and my custom analysis.

In an example MLIR (below) using the linalg.generic operation, the operation doesn’t seem to be visited by the visitOperation method, and the reason seems to be that the block of the generic operation has been marked as non-live by the DeadCodeAnalysis.

Is there a reason for the DeadCodeAnalysis to do so, or is it something that isn’t supported yet?

#map0 = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0) -> (0)>
func.func @main(%arg0: tensor<2x!FHE.eint<7>>, %arg1: tensor<2xi8>, %acc:
!FHE.eint<7>) -> !FHE.eint<7> {
  %tacc = tensor.from_elements %acc : tensor<1x!FHE.eint<7>>
  %2 = linalg.generic 
  {indexing_maps = [#map0, #map0, #map1], iterator_types = ["reduction"]}
  ins(%arg0, %arg1 : tensor<2x!FHE.eint<7>>, tensor<2xi8>)
  outs(%tacc : tensor<1x!FHE.eint<7>>)
  ^bb0(%arg2: !FHE.eint<7>, %arg3: i8, %arg4: !FHE.eint<7>):  // no predecessors
    %4 = "FHE.mul_eint_int"(%arg2, %arg3) : (!FHE.eint<7>, i8) -> !FHE.eint<7>
    %5 = "FHE.add_eint"(%4, %arg4) : (!FHE.eint<7>, !FHE.eint<7>) -> !FHE.eint<7> 
    linalg.yield %5 : !FHE.eint<7>
  } -> tensor<1x!FHE.eint<7>>
  %c0 = arith.constant 0 : index
  %ret = tensor.extract %2[%c0] : tensor<1x!FHE.eint<7>>
  return %ret : !FHE.eint<7>

Linalg operations don’t implement region-based control flow interfaces, so the dead code analysis dataflow analysis doesn’t know that the control can enter the body of a linalg.generic and considers the operation instead it “non-live”. They don’t implement the interface because the kind of flow Linalg operations actually have cannot be expressed in terms those interfaces require. Specifically, the interface expects that data flows unaltered from operands of the parent operation to a block argument (and similarly from terminator operands to the result of the parent operation), which isn’t correct for linalg.generic where only a subset of data flows, moreover a different one on each iteration. As a matter of fact, I removed the interface from this operation in [mlir] remove RegionBranchOpInterface from linalg ops · llvm/llvm-project@2469cdd · GitHub precisely because it was producing unexpected/wrong results in (downstream) dataflow analyses.

Is the operation itself not visited, or its body?
Also please double-check if you are using a sufficiently recent version.

1 Like

The operation itself isn’t visited. So no call to visitiOperation with the GenericOp. If the operation was visited, I could at least try to compute the value of the analysis for it, without requiring the analysis to go through the body.

Currently running on [clangd] Publish diagnostics with stale preambles · llvm/llvm-project@465ee9b · GitHub which I don’t think is considered as a recent version.

You can try cherry-picking my patch into your tree. Without it, the linalg.generic operation is treated as a control-flow only operation and therefore is not visited by itself.

Cherry-picking the patch works for me, thanks Alex!

1 Like