[LoopSchedule][Calyx] PipelineToCalyx Pass Produces Incorrect Results on Simple Designs

[andrewb1999]

I have been experimenting with the AffineToPipeline and PipelineToCalyx passes recently and discovered a bug with the PipelineToCalyx pass when testing a simple vectorized add.

The input program in affine (A, B, and O are the top level IO but are left as allocs because of how simulation in native Calyx works for now):

func.func @main() {                                                          
  %A = memref.alloc() : memref<8xi32>   
  %B = memref.alloc() : memref<8xi32>   
  %O = memref.alloc() : memref<8xi32>
  affine.for %arg2 = 0 to 8 step 1 {
    %7 = affine.load %A[%arg2] : memref<8xi32>
    %8 = affine.load %B[%arg2] : memref<8xi32>
    %9 = arith.addi %7, %8 : i32
    affine.store %9, %O[%arg2] : memref<8xi32>
  }
  return
}

I then ran this program through AffineToPipeline and PipelineToCalyx, exported to native calyx, and simulated the design to get the following output:

{
  "cycles": 69,
  "memories": {
    "mem_0": [
      1,
      3,
      7,
      15,
      31,
      63,
      127,
      255
    ],
    "mem_1": [
      1,
      1,
      1,
      1,
      1,
      1,
      1,
      1
    ],
    "mem_2": [
      256,
      2,
      4,
      8,
      16,
      32,
      64,
      128
    ]
  }
}

mem_2 in this case corresponds to %O in this case, the other memory values are set as inputs. Notice that the correct values are there but stored in the wrong address, including the last value wrapping around to address 0.

I believe this is happening because when we lower the pipeline to calyx we do not properly consider the timing of when the loop iter arg value is updated. Namely the same loop iter arg value is used in each stage when it also needs to be pipelined to ensure each stage has the correct value.

This has become a blocking issue for me in exploring improvements to scheduling for Calyx, so I'm willing to put in a decent amount of engineering effort for a fix here. However, I am unsure of where and how it is best to make this fix.

For completeness, here is the program after being lowered to a pipeline:

module {
  func.func @main() {
    %alloc = memref.alloc() : memref<8xi32>
    %alloc_0 = memref.alloc() : memref<8xi32>
    %alloc_1 = memref.alloc() : memref<8xi32>
    %c0 = arith.constant 0 : index
    %c8 = arith.constant 8 : index
    %c1 = arith.constant 1 : index
    pipeline.while II =  1 trip_count =  8 iter_args(%arg0 = %c0) : (index) -> () {
      %0 = arith.cmpi ult, %arg0, %c8 : index
      pipeline.register %0 : i1
    } do {
      %0:3 = pipeline.while.stage start = 0 {
        %1 = memref.load %alloc[%arg0] : memref<8xi32>
        %2 = memref.load %alloc_0[%arg0] : memref<8xi32>
        %3 = arith.addi %arg0, %c1 : index
        pipeline.register %1, %2, %3 : i32, i32, index
      } : i32, i32, index
      pipeline.while.stage start = 1 {
        %1 = arith.addi %0#0, %0#1 : i32
        memref.store %1, %alloc_1[%arg0] : memref<8xi32>
        pipeline.register 
      }
      pipeline.terminator iter_args(%0#2), results() : (index) -> ()
    }
    return
  }
}

and the design after lowering to Calyx:

module attributes {calyx.entrypoint = "main"} {
  calyx.component @main(%clk: i1 {clk}, %reset: i1 {reset}, %go: i1 {go}) -> (%done: i1 {done}) {
    %c0_i32 = hw.constant 0 : i32
    %c8_i32 = hw.constant 8 : i32
    %c1_i32 = hw.constant 1 : i32
    %true = hw.constant true
    %std_slice_2.in, %std_slice_2.out = calyx.std_slice @std_slice_2 : i32, i3
    %std_slice_1.in, %std_slice_1.out = calyx.std_slice @std_slice_1 : i32, i3
    %std_slice_0.in, %std_slice_0.out = calyx.std_slice @std_slice_0 : i32, i3
    %std_add_1.left, %std_add_1.right, %std_add_1.out = calyx.std_add @std_add_1 : i32, i32, i32
    %std_add_0.left, %std_add_0.right, %std_add_0.out = calyx.std_add @std_add_0 : i32, i32, i32
    %std_lt_0.left, %std_lt_0.right, %std_lt_0.out = calyx.std_lt @std_lt_0 : i32, i32, i1
    %mem_2.addr0, %mem_2.write_data, %mem_2.write_en, %mem_2.clk, %mem_2.read_data, %mem_2.done = calyx.memory @mem_2 <[8] x 32> [3] {external = true} : i3, i32, i1, i1, i32, i1
    %mem_1.addr0, %mem_1.write_data, %mem_1.write_en, %mem_1.clk, %mem_1.read_data, %mem_1.done = calyx.memory @mem_1 <[8] x 32> [3] {external = true} : i3, i32, i1, i1, i32, i1
    %mem_0.addr0, %mem_0.write_data, %mem_0.write_en, %mem_0.clk, %mem_0.read_data, %mem_0.done = calyx.memory @mem_0 <[8] x 32> [3] {external = true} : i3, i32, i1, i1, i32, i1
    %stage_0_register_1_reg.in, %stage_0_register_1_reg.write_en, %stage_0_register_1_reg.clk, %stage_0_register_1_reg.reset, %stage_0_register_1_reg.out, %stage_0_register_1_reg.done = calyx.register @stage_0_register_1_reg : i32, i1, i1, i1, i32, i1
    %stage_0_register_0_reg.in, %stage_0_register_0_reg.write_en, %stage_0_register_0_reg.clk, %stage_0_register_0_reg.reset, %stage_0_register_0_reg.out, %stage_0_register_0_reg.done = calyx.register @stage_0_register_0_reg : i32, i1, i1, i1, i32, i1
    %while_0_arg0_reg.in, %while_0_arg0_reg.write_en, %while_0_arg0_reg.clk, %while_0_arg0_reg.reset, %while_0_arg0_reg.out, %while_0_arg0_reg.done = calyx.register @while_0_arg0_reg : i32, i1, i1, i1, i32, i1
    calyx.wires {
      calyx.group @assign_while_0_init_0 {
        calyx.assign %while_0_arg0_reg.in = %c0_i32 : i32
        calyx.assign %while_0_arg0_reg.write_en = %true : i1
        calyx.group_done %while_0_arg0_reg.done : i1
      }
      calyx.comb_group @bb0_0 {
        calyx.assign %std_lt_0.left = %while_0_arg0_reg.out : i32
        calyx.assign %std_lt_0.right = %c8_i32 : i32
      }
      calyx.group @bb0_1 {
        calyx.assign %std_slice_2.in = %while_0_arg0_reg.out : i32
        calyx.assign %mem_0.addr0 = %std_slice_2.out : i3
        calyx.assign %stage_0_register_0_reg.in = %mem_0.read_data : i32
        calyx.assign %stage_0_register_0_reg.write_en = %true : i1
        calyx.group_done %stage_0_register_0_reg.done : i1
      }
      calyx.group @bb0_2 {
        calyx.assign %std_slice_1.in = %while_0_arg0_reg.out : i32
        calyx.assign %mem_1.addr0 = %std_slice_1.out : i3
        calyx.assign %stage_0_register_1_reg.in = %mem_1.read_data : i32
        calyx.assign %stage_0_register_1_reg.write_en = %true : i1
        calyx.group_done %stage_0_register_1_reg.done : i1
      }
      calyx.group @bb0_3 {
        calyx.assign %std_add_0.left = %while_0_arg0_reg.out : i32
        calyx.assign %std_add_0.right = %c1_i32 : i32
        calyx.assign %while_0_arg0_reg.in = %std_add_0.out : i32
        calyx.assign %while_0_arg0_reg.write_en = %true : i1
        calyx.group_done %while_0_arg0_reg.done : i1
      }
      calyx.group @bb0_5 {
        calyx.assign %std_slice_0.in = %while_0_arg0_reg.out : i32
        calyx.assign %mem_2.addr0 = %std_slice_0.out : i3
        calyx.assign %mem_2.write_data = %std_add_1.out : i32
        calyx.assign %mem_2.write_en = %true : i1
        calyx.assign %std_add_1.left = %stage_0_register_0_reg.out : i32
        calyx.assign %std_add_1.right = %stage_0_register_1_reg.out : i32
        calyx.group_done %mem_2.done : i1
      }
    }
    calyx.control {
      calyx.seq {
        calyx.par {
          calyx.enable @assign_while_0_init_0
        }
        calyx.par {
          calyx.enable @bb0_1
          calyx.enable @bb0_2
          calyx.enable @bb0_3
        }
        calyx.while %std_lt_0.out with @bb0_0 {
          calyx.par {
            calyx.enable @bb0_1
            calyx.enable @bb0_2
            calyx.enable @bb0_3
            calyx.enable @bb0_5
          }
        } {bound = 7 : i64}
        calyx.par {
          calyx.enable @bb0_5
        }
      }
    }
  } {toplevel}
}

[andrewb1999]

@cgyurgyik Any ideas here?

[cgyurgyik]

At the Calyx abstraction, given DRF0 semantics, we have undefined behavior in @bb0_1, @bb0_2, @bb0_3: while_0_arg0_reg is both being written to and read from in the same par sequence. I know Cidr has some support for race detection; it might be a good first step to see if the debugger will catch this (and open a bug in Calyx otherwise).

At the Pipeline abstraction, this looks like a scheduling issue, as you've pointed out. My first suggestion would be to look at the PipelineScheduler, and iterate from there.

For extra context, the two passes SCFToCalyx and PipelineToCalyx used to be a single monolith pass. I also ran into another issue with PipelineToCalyx (#3311). It was difficult to iterate on this, so I opened #2988; I haven't come around to finishing this because of time constraints. I don't think will block you from making changes though.

[andrewb1999]

Thanks for the info. I guess my main question is whether the Pipeline is actually incorrect or if our interpretation of it in PipelineToCalyx is just incorrect. Would the following be a correct output of AffineToPipeline?

  module {
    func.func @main() {
      %alloc = memref.alloc() : memref<8xi32>
      %alloc_0 = memref.alloc() : memref<8xi32>
      %alloc_1 = memref.alloc() : memref<8xi32>
      %c0 = arith.constant 0 : index
      %c8 = arith.constant 8 : index
      %c1 = arith.constant 1 : index
      pipeline.while II =  1 trip_count =  8 iter_args(%arg0 = %c0) : (index) -> () {
        %0 = arith.cmpi ult, %arg0, %c8 : index
        pipeline.register %0 : i1
      } do {
        %0:3 = pipeline.while.stage start = 0 {
          %1 = memref.load %alloc[%arg0] : memref<8xi32>
          %2 = memref.load %alloc_0[%arg0] : memref<8xi32>
          pipeline.register %1, %2, %arg0 : i32, i32, index
        } : i32, i32, index
        %3 = pipeline.while.stage start = 1 {
          %1 = arith.addi %0#0, %0#1 : i32
          memref.store %1, %alloc_1[%0#2] : memref<8xi32>
          %2 = arith.addi %0#2, %c1 : index
          pipeline.register %2 : index
        } : index
        pipeline.terminator iter_args(%3), results() : (index) -> ()
      }
      return
    }
  }

If so, maybe we need to add verification to the pipeline dialect that ensures only stages that start at time 0 can use iter_args.

[cgyurgyik]

Ah, good question, I'll punt this one to the original author @mikeurbach.

[mikeurbach]

Thanks for the detailed investigation! Yes, I think this is unimplemented support for later stages using iter arg(s). Currently we only support trivial loops and hardcode the iter arg update to the first stage, so really any stage after the first that uses the iter arg would be susceptible to this. I've only gotten this flow working for a dot product example, which doesn't write back results in a later stage, but instead passes it through iter args directly, which is why it never came up before.

We had some discussion in the original issue on Discourse (https://discourse.llvm.org/t/rfc-representing-pipelined-loops/4171).

At one point I proposed exactly what is shown in the proposed snippet from @andrewb1999: the original value of iter arg(s) could be plumbed through pipeline registers to whatever stage they are consumed in by the AffineToPipeline pass.

Another option, which I think is actually where we settled last year, is to leave it implicit, so the PipelineToCalyx pass would have to detect the need for this, and insert the registers itself.

Thinking about it fresh, I sort of prefer making this a concern of AffineToPipeline, and explicitly registering iter_args needed moving forward through the pipeline. Implementing this at the pipeline dialect is simpler, and makes it more explicit to the backend.

[andrewb1999]

@mikeurbach thanks for the info! I would also much prefer making this a concern of AffineToPipeline. The PipelineToCalyx pass is already quite complicated and hard to debug. By making minor modifications to AffineToPipeline I was able to get quite a few more test cases working. I will open a draft PR with these changes so we can discuss further.

[mikeurbach]

Wow that's awesome, I was going to offer to take a look but sounds like you're way ahead of me. And yeah, PipelineToCalyx is annoying because that pass grew a lot of functionality organically and has been partly refactored... but not enough to remove all the accidental complexity.

[mikeurbach]

I will open a draft PR with these changes so we can discuss further.

Sounds good to me. One idea before you get too far: there is only really one induction variable per loop now by design, but someday, I would like to get this pass working for nested loops. So for the change you're making, I'd suggest whatever code plumbs through induction variables to later stages plumb through a SmallVector or some container of values, not just one. This might be a little more work, but should be more future proof for when we get multiple induction variables working.