[Pass] Sibling fusion strategy requires a specific memref

[chhzh123]

The following test program fails.

def test_compute_at_no_dep():
    hcl.init()
    A = hcl.compute((10, 10), lambda y, x: y + x, "A")
    B = hcl.compute((10, 10), lambda y, x: y - x, "B")
    s = hcl.create_schedule([A, B])
    s[A].compute_at(s[B], B.axis[1])
    f = hcl.build(s)

It generates correct IR but cannot fuse the loops together.

module {
  func @top() -> (memref<10x10xi32>, memref<10x10xi32>) attributes {extra_itypes = "", extra_otypes = "ss"} {
    %0 = hcl.create_loop_handle "y" : !hcl.LoopHandle
    %1 = hcl.create_loop_handle "x" : !hcl.LoopHandle
    %2 = memref.alloc() {name = "A"} : memref<10x10xi32>
    affine.for %arg0 = 0 to 10 {
      affine.for %arg1 = 0 to 10 {
        %8 = arith.addi %arg0, %arg1 : index
        %9 = arith.index_cast %8 : index to i32
        affine.store %9, %2[%arg0, %arg1] {to = "A"} : memref<10x10xi32>
      } {loop_name = "x"}
    } {loop_name = "y", stage_name = "A"}
    %3 = hcl.create_stage_handle "A" : !hcl.StageHandle
    %4 = hcl.create_loop_handle "y" : !hcl.LoopHandle
    %5 = hcl.create_loop_handle "x" : !hcl.LoopHandle
    %6 = memref.alloc() {name = "B"} : memref<10x10xi32>
    affine.for %arg0 = 0 to 10 {
      affine.for %arg1 = 0 to 10 {
        %8 = arith.subi %arg0, %arg1 : index
        %9 = arith.index_cast %8 : index to i32
        affine.store %9, %6[%arg0, %arg1] {to = "B"} : memref<10x10xi32>
      } {loop_name = "x"}
    } {loop_name = "y", stage_name = "B"}
    %7 = hcl.create_stage_handle "B" : !hcl.StageHandle
    return %2, %6 : memref<10x10xi32>, memref<10x10xi32>
  }
}

Got the following error.

python3: /scratch/users/hc676/llvm-project/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h:79: mlir::FusionStrategy::FusionStrategy(mlir::FusionStrategy::StrategyEnum): Assertion `strategy != Sibling && "Sibling fusion strategy requires a specific memref"' failed.
 #0 0x00007feeac49a8ef PrintStackTraceSignalHandler(void*) Signals.cpp:0:0
 #1 0x00007feeac498319 SignalHandler(int) Signals.cpp:0:0
 #2 0x00007feec917e630 __restore_rt sigaction.c:0:0
 #3 0x00007feec8dd7387 raise (/lib64/libc.so.6+0x36387)
 #4 0x00007feec8dd8a78 abort (/lib64/libc.so.6+0x37a78)
 #5 0x00007feec8dd01a6 __assert_fail_base (/lib64/libc.so.6+0x2f1a6)
 #6 0x00007feec8dd0252 (/lib64/libc.so.6+0x2f252)
 #7 0x00007feeac42d688 mlir::FusionStrategy::FusionStrategy(mlir::FusionStrategy::StrategyEnum) /scratch/users/hc676/llvm-project/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h:80:0
 #8 0x00007feeac4213f8 mlir::hcl::runComputeAt(mlir::FuncOp&, mlir::hcl::ComputeAtOp&) /home/hc676/hcl-dialect-prototype/lib/Transforms/LoopTransformations.cpp:813:0
 #9 0x00007feeac426fb8 mlir::hcl::applyLoopTransformationOnSingleFunction(mlir::FuncOp&) /home/hc676/hcl-dialect-prototype/lib/Transforms/LoopTransformations.cpp:1639:0
#10 0x00007feeac42756a mlir::hcl::applyLoopTransformation(mlir::ModuleOp&) /home/hc676/hcl-dialect-prototype/lib/Transforms/LoopTransformations.cpp:1689:0
#11 0x00007feeac3e3009 loopTransformation(MlirModule&) //home/hc676/hcl-dialect-prototype/lib/Bindings/Python/HCLModule.cpp:75:0