Refactor LoopFuseSiblingOp and support parallel fusion

[srcarroll]

This patch refactors code related to LoopFuseSiblingOp transform in attempt to reduce duplicate common code. The aim is to refactor as much as possible to a functions on LoopLikeOpInterfaces, but this is still a work in progress. A full refactor will require more additions to the LoopLikeOpInterface.

In addition, scf.parallel fusion support has been added.

[Hardcode84]

I'm concerned about difference of semantics for fuse_sibling transform, which do only the structural checks for the scf.for and scf.forall but will try to do an aliasing checks for scf.parallel. We probably should split fuseIfLegal func into 3:

• checkFusionStructuralLegality - doing only the structural checks (interation domain and dominace), used both in transform and in naivelyFuseParallelOps
• fuseParallel - fuses 2 scf.parallel without any checks, also used both in transform and in naivelyFuseParallelOps
• checkParallelAliasing - nested parallel, haveNoReadsAfterWriteExceptSameIndex and aliasing checks, used only in naivelyFuseParallelOps - also you can keep this func in ParallelLoopFusion.cpp and avoid introducing memref dependency to utils.

[srcarroll]

I'm concerned about difference of semantics for fuse_sibling transform, which do only the structural checks for the scf.for and scf.forall but will try to do an aliasing checks for scf.parallel. We probably should split fuseIfLegal func into 3:

• checkFusionStructuralLegality - doing only the structural checks (interation domain and dominace), used both in transform and in naivelyFuseParallelOps
• fuseParallel - fuses 2 scf.parallel without any checks, also used both in transform and in naivelyFuseParallelOps
• checkParallelAliasing - nested parallel, haveNoReadsAfterWriteExceptSameIndex and aliasing checks, used only in naivelyFuseParallelOps - also you can keep this func in ParallelLoopFusion.cpp and avoid introducing memref dependency to utils.

sounds good. this was just a first atttempt with minimal changes but I agree with you and will refactor per your suggestions. I'll let you know if i come across problems with the suggestion

[srcarroll]

so this is where the mess i mentioned really becomes a problem. there are static helpers all over the place, when a lot of these should be a part of libraries to be reused. for example there is a isSibling that is used for checking dominance in SCFTransformOps.cpp, which is much more complex than what is checked for the pass

  for (Operation *user : firstPloop->getUsers())
    if (!dom.properlyDominates(secondPloop, user, /*enclosingOpOk*/ false))
      return;

so should I move isSibling to some library and use for both?

I'll likely have a lot of questions like this. Surely I can figure it out on my own, but will take a while since this is the first time seeing a lot of this code. So will appreciate feedback on these kinds of choices.

[Hardcode84]

so should I move isSibling to some library and use for both?

Not all the checks in isSibling are needed for ParallelLoopFusion.cpp as algorithm there guarantees they will be always two separate ops with predefined order. I think, for now you can use isSibling for the transform and keep a simpler version in ParallelLoopFusion.cpp as static function.

[Hardcode84]

Side note: not many folks will see this PR while it's in draft state and I'm not very familiar with the transform dialect.

[srcarroll]

Just want to run some organizational details through you. This is what i currently think about what header and library the functions you mentioned should live in.

checkFusionStructuralLegality - declare in SCF/Utils/Utils.h, implemented in SCF/Utils/Utils.cpp which will make it part of the MLIRSCFUtils lib
fuseParallel - declare in SCF/Transforms/Transforms.h to be part of the MLIRSCFTransforms. not sure what source this should be implemented in yet
checkParallelAliasing - i suppose for now this could be static in ParallelLoopFusion.cpp but i imagine this can be reused eventually. just dont have much of an opinion on it yet.

Any thoughts?

Side note: not many folks will see this PR while it's in draft state and I'm not very familiar with the transform dialect.

i'll undraft after implementing your suggestions and ping @ftynse to take a look if he's available

[Hardcode84]

• checkFusionStructuralLegality - keep as static function for now in ParallelLoopFusion.cpp, use exisitng isSibling for transform.
• fuseParallel (fuseIndependentSiblingParallelLoops) - Utils.h/cpp as fuseIndependentSiblingForallLoops and fuseIndependentSiblingForLoops are already there
• checkParallelAliasing - public function in ParallelLoopFusion.cpp

[srcarroll]

checkFusionStructuralLegality - keep as static function for now in ParallelLoopFusion.cpp, use exisitng isSibling for transform.

as far as i can tell, this change wouldn't really benefit anything if it's not shared by pass and transform. it would just be replacing inline code with a static function. am i missing something?

fuseParallel (fuseIndependentSiblingParallelLoops) - Utils.h/cpp as fuseIndependentSiblingForallLoops and fuseIndependentSiblingForLoops are already there

ah ok, i thought originally you were suggesting a new fuseParallel, but i see now that's not necessary. i already have fuseIndependentSiblingParallelLoops in utils

as a related side note: i do think the iteration domain check should surely be shared. In which case i think having it in Utils makes sense. would you agree?

[Hardcode84]

as far as i can tell, this change wouldn't really benefit anything if it's not shared by pass and transform. it would just be replacing inline code with a static function. am i missing something?

Ok, if you can pull isOpSibling from transforms into somewhere like Utils.hpp you can use it instead.

as a related side note: i do think the iteration domain check should surely be shared. In which case i think having it in Utils makes sense. would you agree?

Ideally, it should be somewhere in utils, operating on LoopLikeInterface instead of concrete ops types, but looks like LoopLikeOpInterface currently doesn't support multiple steps/upper/lower bounds so it won't work for multidim scf.forall and scf.parallel.

Updating LoopLikeOpInterface to support them will also be a good refactoring, but it should go in separate PR.

[srcarroll]

Updating LoopLikeOpInterface to support them will also be a good refactoring, but it should go in separate PR.

yah that's what i was thinking. maybe i should do that first

[srcarroll]

here's the PR for the additional interface methods #94516

[Hardcode84]

LGTM

[srcarroll]

LGTM

thanks for the review. i need to resolve conflicts and then i'll wait for @ftynse 's approval since he requested a bunch of non-trivial changes i want to make i addressed.

[qedawkins]

This patch appears to be introducing some dominance errors, e.g.

#map = affine_map<(d0) -> (d0 * 32)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
module {
  func.func @loop_sibling_fusion(%arg0: tensor<128xf32>, %arg1: tensor<128x128xf16>, %arg2: tensor<128x64xf32>, %arg3: tensor<128x128xf32>) {
    %0 = scf.forall (%arg4) in (4) shared_outs(%arg5 = %arg0) -> (tensor<128xf32>) {
      %3 = affine.apply #map(%arg4)
      %extracted_slice = tensor.extract_slice %arg3[%3, 0] [32, 1] [1, 1] : tensor<128x128xf32> to tensor<32xf32>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %extracted_slice into %arg5[%3] [32] [1] : tensor<32xf32> into tensor<128xf32>
      }   
    } {mapping = [#gpu.warp<linear_dim_0>]}
    %1 = tensor.empty() : tensor<128x128xf16>
    %2 = scf.forall (%arg4) in (4) shared_outs(%arg5 = %arg1) -> (tensor<128x128xf16>) {
      %3 = affine.apply #map(%arg4)
      %extracted_slice = tensor.extract_slice %arg3[%3, 0] [32, 128] [1, 1] : tensor<128x128xf32> to tensor<32x128xf32>
      %extracted_slice_0 = tensor.extract_slice %1[%3, 0] [32, 128] [1, 1] : tensor<128x128xf16> to tensor<32x128xf16>
      %4 = linalg.generic {indexing_maps = [#map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice : tensor<32x128xf32>) outs(%extracted_slice_0 : tensor<32x128xf16>) {
      ^bb0(%in: f32, %out: f16):
        %5 = arith.truncf %in : f32 to f16 
        linalg.yield %5 : f16 
      } -> tensor<32x128xf16>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %4 into %arg5[%3, 0] [32, 128] [1, 1] : tensor<32x128xf16> into tensor<128x128xf16>
      }   
    } {mapping = [#gpu.warp<linear_dim_0>]}
    return
  }
}

module attributes { transform.with_named_sequence } { 
  transform.named_sequence @__transform_main(%root: !transform.any_op) {
    %loops = transform.structured.match ops{["scf.forall"]} in %root : (!transform.any_op) -> !transform.any_op
    %loop1, %loop2 = transform.split_handle %loops : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
    %loop3 = transform.loop.fuse_sibling %loop1 into %loop2 : (!transform.any_op, !transform.any_op) -> !transform.any_op
    transform.yield
  }
}

with command mlir-opt repro.mlir -transform-interpreter is giving the following output for me

repro.mlir:16:28: error: operand #0 does not dominate this use
      %extracted_slice_0 = tensor.extract_slice %1[%3, 0] [32, 128] [1, 1] : tensor<128x128xf16> to tensor<32x128xf16>
                           ^
repro.mlir:16:28: note: see current operation: %6 = "tensor.extract_slice"(%1, %4) <{operandSegmentSizes = array<i32: 1, 1, 0, 0>, static_offsets = array<i64: -9223372036854775808, 0>, static_sizes = array<i64: 32, 128>, static_strides = array<i64: 1, 1>}> : (tensor<128x128xf16>, index) -> tensor<32x128xf16>
repro.mlir:12:10: note: operand defined here (op in a parent region)
    %1 = tensor.empty() : tensor<128x128xf16>
         ^

Dumping the generic IR reveals that the implicitly captured tensor.empty is somehow ending up after the fused loop: https://gist.github.com/qedawkins/11a10aeefef6168829f6bffdf949fe1e

I'm guessing there is an insertion point bug introduced somewhere in this refactor. (Initially observed downstream here)

[srcarroll]

@qedawkins thanks for quickly bringing this to my attention. i will revert the patch, fix, and add that test case. thanks!

[srcarroll]

@qedawkins done #97523

[qedawkins]

Nice, thanks for the quick revert!

[srcarroll]

no problem. i'll ping you when i reland the PR