Stepper Looplet and Pass

[nullplay]

This PR introduces an operation finch.stepper and a corresponding pass finch-looplet-stepper.
It is a reimplemented version of https://github.com/finch-tensor/Finch.jl/blob/main/src/looplets/steppers.jl in Finch.jl

1. finch.stepper operation

• defined in include/Finch/FinchOps.td

This operation defines the Stepper Looplet, which represents repetitive patterns in child looplets.
Stepper takes four function arguments that always return with finch.return:

1. seek: coordinate:index -> position:index
   Returns the position of the child looplet whose range includes the given input coordinate.

2. stop: position:index -> coordinate:index
   Returns the coordinate at which the child looplet at the given position ends.

3. body: position:index -> child:looplet
   Returns the child looplet located at the specified position.

4. next: position:index -> position:index
   Returns the subsequent position following the given position.

finch.stepper seek(crd):pos, stop(pos):crd, body(pos):looplet, next(pos):pos can be thought as a function

p = initpos
f(x) = body(p) (if x < stop(p))
       body(next(p)) (if stop(p) <= x < stop(next(p)))
       body(next(next(p))) (if stop(next(p)) <= x < stop(next(next(p))))
       ...

Example mlir :

%step = finch.stepper
    seek={
      ^bb0(%idx : index):
        finch.return %idx : index
    }
    stop={
      ^bb(%pos : index):
        %crd = arith.addi %pos, %c2 : index
        finch.return %crd : index
    }
    body={
      ^bb(%pos : index):
        %run = finch.run %f1 : (f32) -> (!finch.looplet)
        finch.return %run : !finch.looplet
    }
    next={
      ^bb0(%pos : index):
        %nextpos = arith.addi %pos, %c1 : index
        finch.return %nextpos : index
    }

                   stop(initpos)            stop(next(initpos))
                         v                              v
[-----body(initpos)-----)[-----body(next(initpos))-----)[----....

2. finch-looplet-stepper pass

• defined in lib/Finch/FinchStepperPass.cpp

Lowers finch.access of finch.stepper inside scf.for

if we iterate f(x) over [st, en),
p = seek(st)
f(x) = body(p) (if x in [st, min(en, stop(p))) )
       body(next(p)) (if x in [min(en, stop(p)), min(en, stop(next(p)))))
       ...

A pseudocode of lowering a single stepper looks like :

%0 = finch.stepper
  seek = {
    ^bb(%crd:index):
      ...
  }
  stop = {
    ^bb(%pos:index):
      ...
  }
  body = {
    ^bb(%pos:index):
      ...
  }
  next = {
    ^bb(%pos:index):
      ...
  }
scf.for %idx = %st to %en step %c1 {
    %1 = finch.access %0, %idx : f32
}

to

%initpos = stepper.seek(%st)
%1:2 = scf.while (%pos=%initpos, %idx=%st) : (index,index) -> (index,index) {
  %cmp = arith.cmpi ult, %idx, %en : index
  scf.condition(%cmp) %pos, %idx : index, index
} do {
  ^bb(%pos: index, %idx: index):
    %stop = stepper.stop(%pos)
    %body = stepper.body(%pos)
    %end = arith.minui %2, %en : index
    scf.for %idx2 = %idx to %end step %c1 {
      %1 = finch.access %body, %idx2 : f32
    }
    
    %nextpos = stepper.next(%pos)
    scf.yield %nextpos, %end
}

3. Test cases

You can find tests at https://github.com/finch-tensor/Finch-mlir/blob/b453a57d9193e73ee64f55d60987be51fd30bec0/test/Finch/looplet-stepper.mlir

cmake --build build --target check-finch will run this test automatically, but if you want to run manually, do :
./build/bin/finch-opt ./test/Finch/looplet-stepper.mlir --finch-looplet-stepper

[hameerabbasi]

LGTM -- IIUC the stepper pass has been moved and slightly edited, correct?

[nullplay]

Yes, there was a minor bug in the stepper pass. I fixed it while moving it.