[mlir][CSE] Introduce hoist-pure-ops logic to CSE pass

[linuxlonelyeagle]

This PR is based on this theory: if an Op is a Pure Op, we have the opportunity to hoist its position based on SSA dominance. This logic has now been incorporated into the CSE pass, now we can use it to further optimize the IR to achieve more concise code.
RFC: https://discourse.llvm.org/t/rfc-mlir-introduce-hoist-pure-ops-pass/88903

[linuxlonelyeagle]

It looks like the CI failed to build because it couldn't pull LLVM.

[github-actions]

🐧 Linux x64 Test Results

• 8534 tests passed
• 630 tests skipped

✅ The build succeeded and all tests passed.

[github-actions]

🪟 Windows x64 Test Results

• 8146 tests passed
• 690 tests skipped

✅ The build succeeded and all tests passed.

[linuxlonelyeagle]

Apologies, I'm running into some issues with Flang and MLIR Python bindings. I'm currently investigating.

[linuxlonelyeagle]

    hlfir.forall lb {
      hlfir.yield %c1_i32 : i32 
    } ub {
      hlfir.yield %c10_i32 : i32 
    }  (%arg0: i32) {
      hlfir.region_assign {
        %25 = fir.convert %arg0 : (i32) -> i64
        %26 = hlfir.designate %19#0 (%25)  typeparams %c1 : (!fir.ref<!fir.array<10x!fir.char<1>>>, i64, index) -> !fir.ref<!fir.char<1>>
        %27 = fir.embox %26 : (!fir.ref<!fir.char<1>>) -> !fir.box<!fir.ptr<!fir.char<1>>>
        %28 = fir.rebox %27 : (!fir.box<!fir.ptr<!fir.char<1>>>) -> !fir.class<!fir.ptr<none>>
        hlfir.yield %28 : !fir.class<!fir.ptr<none>> 
      } to {
        %25 = fir.convert %arg0 : (i32) -> i64
        %26 = hlfir.designate %9#0 (%25)  : (!fir.ref<!fir.array<10x!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>>, i64) -> !fir.ref<!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>
        %27 = hlfir.designate %26{"ptr"}   {fortran_attrs = #fir.var_attrs<pointer>} : (!fir.ref<!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>) -> !fir.ref<!fir.class<!fir.ptr<none>>>
        hlfir.yield %27 : !fir.ref<!fir.class<!fir.ptr<none>>> 
      }
    }

after cse of PR, %300 = "fir.convert"(%arg0) : (i32) -> i64 is hoisted.

    ^bb0(%arg0: i32):
      %300 = "fir.convert"(%arg0) : (i32) -> i64
      "hlfir.region_assign"() ({
        %303 = "hlfir.designate"(%294#0, %300, %272) <{is_triplet = array<i1: false>, operandSegmentSizes = array<i32: 1, 0, 1, 0, 0, 1>}> : (!fir.ref<!fir.array<10x!fir.char<1>>>, i64, index) -> !fir.ref<!fir.char<1>>
        %304 = "fir.embox"(%303) <{operandSegmentSizes = array<i32: 1, 0, 0, 0, 0>}> : (!fir.ref<!fir.char<1>>) -> !fir.box<!fir.ptr<!fir.char<1>>>
        %305 = "fir.rebox"(%304) <{operandSegmentSizes = array<i32: 1, 0, 0>}> : (!fir.box<!fir.ptr<!fir.char<1>>>) -> !fir.class<!fir.ptr<none>>
        "hlfir.yield"(%305) ({
        }) : (!fir.class<!fir.ptr<none>>) -> ()
      }, {
        %301 = "hlfir.designate"(%286#0, %300) <{is_triplet = array<i1: false>, operandSegmentSizes = array<i32: 1, 0, 1, 0, 0, 0>}> : (!fir.ref<!fir.array<10x!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>>, i64) -> !fir.ref<!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>
        %302 = "hlfir.designate"(%301) <{component = "ptr", fortran_attrs = #fir.var_attrs<pointer>, is_triplet = array<i1>, operandSegmentSizes = array<i32: 1, 0, 0, 0, 0, 0>}> : (!fir.ref<!fir.type<_QFforallpolymorphic3Tdt{ptr:!fir.class<!fir.ptr<none>>}>>) -> !fir.ref<!fir.class<!fir.ptr<none>>>
        "hlfir.yield"(%302) ({
        }) : (!fir.ref<!fir.class<!fir.ptr<none>>>) -> ()
      }, {
      }) : () -> ()

It report, Can we allow pure operations to exist within the 'hlfir.forall' op's body region?

error: 'hlfir.forall' op body region must only contain OrderedAssignmentTreeOpInterface operations or fir.end
    "hlfir.forall"() ({
    ^
c.mlir:34:5: note: see current operation: 
"hlfir.forall"() ({
  "hlfir.yield"(%3) ({
  }) : (i32) -> ()
}, {
  "hlfir.yield"(%0) ({
  }) : (i32) -> ()
}, {
}, {

cc: @tarunprabhu @banach-space @eugeneepshteyn @joker-eph @ftynse @matthias-springer @kuhar . I would appreciate any suggestions you might have. Thanks!

[linuxlonelyeagle]

module {
  irdl.dialect @myint {
    irdl.operation @constant {
      %0 = irdl.base "#builtin.integer" 
      irdl.attributes {"value" = %0}
      %1 = irdl.is i32 
      irdl.results(cst: %1)
    }
    irdl.operation @add {
      %0 = irdl.is i32 
      %1 = irdl.is i32 
      irdl.operands(lhs: %0, rhs: %1)
      %2 = irdl.is i32 
      irdl.results(res: %2)
    }
  }
}

after cse.I've gone through some test files, and it might make sense to add the IsIsolatedFromAbove trait to irdl.operation.
cc: @joker-eph @ftynse @matthias-springer @kuhar @math-fehr 😘

a:6:12: error: 'irdl.is' op expects parent op to be one of 'irdl.type, irdl.attribute, irdl.operation'
      %1 = irdl.is i32 
           ^
a:6:12: note: see current operation: %0 = "irdl.is"() <{expected = i32}> : () -> !irdl.attribute
// -----// IR Dump After CSE Failed (cse) //----- //
"builtin.module"() ({
  "irdl.dialect"() <{sym_name = "myint"}> ({
    %0 = "irdl.is"() <{expected = i32}> : () -> !irdl.attribute
    "irdl.operation"() <{sym_name = "constant"}> ({
      %1 = "irdl.base"() <{base_name = "#builtin.integer"}> : () -> !irdl.attribute
      "irdl.attributes"(%1) <{attributeValueNames = ["value"]}> : (!irdl.attribute) -> ()
      "irdl.results"(%0) <{names = ["cst"], variadicity = #irdl<variadicity_array[ single]>}> : (!irdl.attribute) -> ()
    }) : () -> ()
    "irdl.operation"() <{sym_name = "add"}> ({
      "irdl.operands"(%0, %0) <{names = ["lhs", "rhs"], variadicity = #irdl<variadicity_array[ single,  single]>}> : (!irdl.attribute, !irdl.attribute) -> ()
      "irdl.results"(%0) <{names = ["res"], variadicity = #irdl<variadicity_array[ single]>}> : (!irdl.attribute) -> ()
    }) : () -> ()
  }) : () -> ()
}) : 

[linuxlonelyeagle]

CSE has markAnalysesPreserved<DominanceInfo, PostDominanceInfo>(), is this still correct? Before we were just deleting redundant expression, now we're moving across blocks which likely invalidate it? We should detect when operations move across blocks and not preserve the analysis when it happens (or update it?).

Let me address this question now, markAnalysesPreserved<DominanceInfo, PostDominanceInfo>() remains correct.
You can see

llvm-project/mlir/lib/IR/Dominance.cpp

Line 299 in 682ae8b

// If the blocks are different, use DomTree to resolve the query.

We only rely on the Dominator Tree when determining dominance relationships between different blocks within the same region. In all other cases, dominance can be inferred from the IR's relative positions. The implementation of DominanceInfo leverages LLVM's DominatorTreeBase, which aligns perfectly with LLVM IR's design. This means that although I am hoisting operations here, it does not fundamentally change the relative positioning of blocks within the region. Therefore, DominanceInfo and PostDominanceInfo remain valid.

Also I'm not sure about test coverage, you had to update many tests, but it's not clear that we cover all the edge cases. Do we have tests covering hoisting blocked at IsIsolatedFromAbove in various conditions? I would suggest adding new minimal tests specifically targeted for this CSE hoisting.

I’ve added some test cases and would love to hear your thoughts. I’m definitely looking for more feedback on this part. 😉.

[joker-eph]

This means that although I am hoisting operations here, it does not fundamentally change the relative positioning of blocks within the region. Therefore, DominanceInfo and PostDominanceInfo remain valid.

Contrary to LLVM, we do have regions: when we CSE an operation with regions it means we're deleting blocks. Are we deleting the entry for these regions in the analysis cache?

[linuxlonelyeagle]

This means that although I am hoisting operations here, it does not fundamentally change the relative positioning of blocks within the region. Therefore, DominanceInfo and PostDominanceInfo remain valid.

Contrary to LLVM, we do have regions: when we CSE an operation with regions it means we're deleting blocks. Are we deleting the entry for these regions in the analysis cache?

Thanks for bringing this up. It really pushed me to think more deeply about this issue. To be honest, this issue still persists in the original CSE, such as in the cse_multiple_regions test. I need to investigate this further. Let me give it some more thought and get back to you.😉

[linuxlonelyeagle]

This means that although I am hoisting operations here, it does not fundamentally change the relative positioning of blocks within the region. Therefore, DominanceInfo and PostDominanceInfo remain valid.

Contrary to LLVM, we do have regions: when we CSE an operation with regions it means we're deleting blocks. Are we deleting the entry for these regions in the analysis cache?

In DominanceInfoBase, a DenseMap is used to store the dominator tree for each region. I think we should also remove the stale cache from this map when a region-carrying operation is deleted. We could explicitly add a method to DominanceInfoBase to erase a region's dominator tree; otherwise, the memory occupied by these trees won't be released until DominanceInfo itself is destroyed.

[linuxlonelyeagle]

Please feel free to continue the review when you have time. Thank you all!