[Feature][IR][Dataflow] Keep `meta_if` / `meta_else` in IR to preserve 'singel df.kernel' program structure

[Fangtangtang]

Motivation and Expected Behavior

This is a pre-task for #510.

Currently, Allo dataflow IR creates multiple code variants after lowering meta_if / meta_elif / meta_else. However, Shardy enforces a strict SPMD model. Directly lowering meta control flow in Allo IR before integrating Shardy would break this, leading to compatibility issues.
To integrate Shardy into Allo dataflow, we need to preserve meta control flow in the IR first.

We may introduce new MLIR operations in Allo to support meta control flow.

Related resources

related dialects:

• https://github.com/openxla/shardy
• https://mlir.llvm.org/docs/Dialects/Shard/

[Fangtangtang]

We can directly use sdy.mesh to construct mesh under the top module

sdy.mesh @mesh = <["a"=4]> // 4 tiles
sdy.mesh @mesh_2D = <["a"=2, "b"=2]> // 2x2 tiles

🔴Design choice for lowering df.get_pid()

Note: We need to preserve df.get_pid() in the IR. Its value is not only used for meta meta_if / meta_else, but also for stream selection and certain computations.

Option 1: use process_multi_index in shard dialect

%pid_0 = shard.process_multi_index on @mesh axes = [0] : index
%pid:2 = shard.process_multi_index on @mesh_2D axes = [0, 1] : index

Pros:

• Directly reuses an in-tree dialect operation.

Cons:

• axes does not fully match sdy.mesh.
• When multiple pids are needed (e.g., from a 2D mesh), the result may be awkward to use directly inside meta control flow.

Option 2: introduce new operation in allo

For example,

%pid_0 = allo.mesh_coord @mesh["a"] : index

%pid_0 = allo.mesh_coord @mesh_2D["a"] : index
%pid_1 = allo.mesh_coord @mesh_2D["b"] : index

We are free to customize it, but introduing new operation will add burden to allo IR.

🔴Design choice for lowering meta control flow structure

Option 1: nested meta_if and meta_else

allo.meta_if %p0 { // with allo.meta_if(%p0)
  // branch for %p0
  ... 
} else {
  allo.meta_if %p1 { // with allo.meta_elif(%p0)
    // branch for %p1
    ...
  } else { // with allo.meta_else()
    // else branch
    ...
  }
}

Option 2: switch style

allo.meta_if {
  allo.meta_case %p0 { // with allo.meta_if(%p0)
    // branch for %p0
    ...
  }
  allo.meta_case %p1 { // with allo.meta_elif(%p0)
    // branch for %p1
    ...
  }
  allo.meta_default { // with allo.meta_else()
    // else branch
    ...
  }
}

🔴Design choice for lowering meta control flow conditions

The conditions are compile time evaluable constants.

Option 1: treat conditions as normal branch conditions

Lower these conditions as regular branch conditions using compare operations. For example:

%c0 = arith.constant 0 : index
%p0 = arith.cmpi eq, %pid_0, %c0 : index

An optional attribute can be added to mark the condition %p0 as compile-time evaluable.

Option 2: introduce a new operation in Allo for predicates

Introduce a dedicated Allo operation to represent meta control flow conditions (predicates) symbolically. For example:

%p0 = allo.predicate "pid_0 == 0"

directly using the symbolic representation of the condition.

Option 3: more 'aggressive' version of option 2

Combine the symbolic predicate directly into the meta control flow IR. For example:

allo.meta_if "pid_0 == 0" { // with allo.meta_if(%pid_0 == 0)
  // branch for %p0
  ... 
} else {  // with allo.meta_else()
  // else branch
}

or

allo.meta_if {
  allo.meta_case "pid_0 == 0" { // with allo.meta_if(%pid_0 == 0)
    // branch for %p0
    ...
  }
  allo.meta_default { // with allo.meta_else()
    // else branch
    ...
  }
}

[Fangtangtang]

Option 2: switch style

allo.meta_if {
  allo.meta_case %p0 { // with allo.meta_if(%p0)
    // branch for %p0
    ...
  }
  allo.meta_case %p1 { // with allo.meta_elif(%p0)
    // branch for %p1
    ...
  }
  allo.meta_default { // with allo.meta_else()
    // else branch
    ...
  }
}

The default block can be omitted.

[chhzh123]

We should try to avoid introducing extra operations. I think option 1 is better for all the cases.

[Fangtangtang]

Then should we directly lower the meta control flow structure to scf.if?
Also, reusing existing operations place much more burden on analysis, we may eventually turn to attaching many attributes to the operations to simplify analysis. I think attributes are even worse than new operaitons

[zhangzhiru]

Personally, I like option 2 better as it facilitates deeper analysis. We could have a lower pass later to turn them into other common operations when we are done with those analyses/transformations.

[Fangtangtang]

had an offline discussion with @chhzh123 and currently converged on the following design choices.

🔴 Design choice for lowering df.get_pid()

Use process_multi_index from shard dialect.
to avoid introducing new operations.

🔴 Design choice for lowering meta control-flow structures

Use nested meta_if.

Introduce only one new operation: allo.meta_if (following the design of scf.if)

Sample code:

allo.meta_if %p0 { // allo.meta_if(%p0)
  // branch for %p0
  ...
} else {
  allo.meta_if %p1 { // allo.meta_elif(%p1)
    // branch for %p1
    ...
  } else { // allo.meta_else()
    // else branch
    ...
  }
}

We choose this because it matches the frontend meta_if / meta_elif / meta_else interface more naturally.

🔴 Design choice for lowering meta control-flow conditions

Treat meta conditions as normal branch conditions.
and avoid attaching additional attributes.

---

During lowering:

We can first replace process_multi_index with constant assignments.
For example: lower

%pid_0 = shard.process_multi_index on @mesh axes = [0] : index

to:

%pid_0 = arith.constant 0 : i32

then rely on constant propagation and constant folding (which should be reusable from existing MLIR passes).
The propagated results are consumed by a dedicated pass that lowers meta_if.

the analysis complexity should remains acceptable.

[Fangtangtang]

Use nested meta_if.
Introduce only one new operation: allo.meta_if (following the design of scf.if)

Since we have ConstExpr, we can directly use if, else and enforce the condition to be ConstExpr (expressions on constants and pids). After instantiate instance for each df.kernel (after some lowering pass), we can use constant propagation followed by DCE to remove the dead branches.
(This would remove some dead if block as well, but should be fine as long as no one want to keep some dead code in the compiled program)

[Fangtangtang]

🔴 Design choice for lowering df.get_pid()

Use process_multi_index from shard dialect.
to avoid introducing new operations.

process_multi_index from shard dialect seems incompatible with mesh from shardy. It only takes the mesh (/ grid) in its own dialect, so some verification would fail.
ShardOps.td

[Fangtangtang]

since we renamed df.kernel to work, maybe get_pid should also be renamed to get_wid.

---

💡another proposal for handling wid.

for each work, we internally insert a call __allo__.get_wid() statement at the entry of the function body. This is a builtin function that is completely transparent to users. It returns len(mapping) symbolic values formatted as f__wid_{i}. I plan to treat these names as reserved words, meaning users should not use identifiers of the form f__wid_{i} in their source code.

The frontend interface remains unchanged. When a user writes:

wi, wj = get_wid()

this does not generate an actual function call. Instead, wi and wj are redirected to __wid_0 and __wid_1, respectively.

At the IR level, the first operation in the work function body becomes:

%wid0, %wid1 = call @__allo__get_wid()

All uses of the get_wid() results in the source code are rewritten to use the results of this inserted call. There will be no remaining call get_wid() operations originating from the user code.

Then, when unrolling work instances, all wid values become per-instance "constexprs". We can simply replace all uses of the call @__allo__.get_wid() results with corresponding arith.constant values for each instance.

Pros:

• constant replacement and propagation during unrolling become more straightforward.
• No need to integrate a new dialect or add a glue op to the existing Allo dialect.

Cons:

• may introduce many additional function definitions in the IR (e.g., per-work's builtin functions if the number of returned wids differs).

[Fangtangtang]

💡another proposal for handling wid.

An example:

source code

@spmw.unit()
def top(A: int32[1024], B: int32[1024]):
    @spmw.work(mapping=[4], inputs=[A], outputs=[B])
    def core(local_A: int32[1024] @ [S(0)], local_B: int32[1024] @ [S(0)]):
        pi: ConstExpr[index] = spmw.get_wid() # can also be: pi = spmw.get_wid()
        local_B[:] = local_A + pi

• the results of spmw.get_wid() are ConstExpr[index].
• since it can have multiple return values, the type annotations can be omitted.

preprocessed AST

@spmw.unit()
def top(A: __allo__[i32, (1024,), None], B: __allo__[i32, (1024,), None]):
    __top__core(A, B)
    return 

@spmw.work(mapping=[4], inputs=[A], outputs=[B])
def __top__core(local_A: __allo__[i32, (256,), [S(0)]], local_B: __allo__[i32, (256,), [S(0)]]):
    __wid_0 = __allo__.get_wid() // insert builtin function instead
    local_B[0:256:1]: __allo__[i32, (256,), None] = __allo__.Add(__allo__.broadcast(local_A, (), __allo__[i32, (256,), None]), __allo__.broadcast(__allo__.cast_index(__wid_0, __allo__[i32, (), None]), (0,), __allo__[i32, (256,), None]), __allo__[i32, (256,), None])
    return 

• __wid_0 = __allo__.get_wid() is the internally inserted builtin function. (it is used to setup some per-work meta data, we can further extend such kind of builtin functions for future needs.)
• the use of pi is replaced with __wid_0

constructed IR

module {
  func.func @top(%arg0: memref<1024xi32>, %arg1: memref<1024xi32>) {
    %0 = bufferization.to_tensor %arg0 : memref<1024xi32> to tensor<1024xi32>
    %1 = sdy.manual_computation(%0) in_shardings=[<@__top__core_mesh, [{"0"}]>] out_shardings=[<@__top__core_mesh, [{"0"}]>] manual_axes={"0"} (%arg2: tensor<256xi32>) {
      %3 = bufferization.to_buffer %arg2 : tensor<256xi32> to memref<256xi32>
      %alloc = memref.alloc() : memref<256xi32>
      func.call @__top__core(%3, %alloc) : (memref<256xi32>, memref<256xi32>) -> ()
      %4 = bufferization.to_tensor %alloc : memref<256xi32> to tensor<256xi32>
      sdy.return %4 : tensor<256xi32>
    } : (tensor<1024xi32>) -> tensor<1024xi32>
    %2 = bufferization.to_buffer %1 : tensor<1024xi32> to memref<1024xi32>
    memref.copy %2, %arg1 : memref<1024xi32> to memref<1024xi32>
    return
  }
  sdy.mesh @__top__core_mesh = <["0"=4]>
  func.func @__top__core(%arg0: memref<256xi32>, %arg1: memref<256xi32>) {
    %0 = call @__top__core_mesh_get_wid() : () -> index // call private placeholder function
    %alloc = memref.alloc() : memref<256xi32>
    linalg.broadcast ins(%arg0 : memref<256xi32>) outs(%alloc : memref<256xi32>) dimensions = [] 
    %1 = arith.index_cast %0 : index to i32
    %alloc_0 = memref.alloc() : memref<256xi32>
    linalg.fill ins(%1 : i32) outs(%alloc_0 : memref<256xi32>)
    %alloc_1 = memref.alloc() : memref<256xi32>
    linalg.add ins(%alloc, %alloc_0 : memref<256xi32>, memref<256xi32>) outs(%alloc_1 : memref<256xi32>)
    %subview = memref.subview %arg1[0] [256] [1] : memref<256xi32> to memref<256xi32, strided<[1]>>
    memref.copy %alloc_1, %subview : memref<256xi32> to memref<256xi32, strided<[1]>>
    return
  }
  func.func private @__top__core_mesh_get_wid() -> index
}

• the builtin function is lowered to a private function func.func private @__top__core_mesh_get_wid() -> index. This is a placeholder for the work's meta data.

[Fangtangtang]

Do we still need meta_if as frontend interface? We can directly use if / elif / else in the source code, and eliminate the dead branch if the condition is a compile time constant. The DCE seems pretty easy with MLIR's existing passes like canonicalize.

Example

source code

@spmw.unit()
def top(A: int32[1024], B: int32[1024]):
    @spmw.work(mapping=[4], inputs=[A], outputs=[B])
    def core(local_A: int32[1024] @ [S(0)], local_B: int32[1024] @ [S(0)]):
        pi = spmw.get_wid()
        if pi > 1:
            local_B[:] = local_A + pi
        else:
            local_B[:] = local_A - pi

constructed IR

module {
  func.func @top(%arg0: memref<1024xi32>, %arg1: memref<1024xi32>) {
    %0 = bufferization.to_tensor %arg0 : memref<1024xi32> to tensor<1024xi32>
    %1 = sdy.manual_computation(%0) in_shardings=[<@__top__core_mesh, [{"0"}]>] out_shardings=[<@__top__core_mesh, [{"0"}]>] manual_axes={"0"} (%arg2: tensor<256xi32>) {
      %3 = bufferization.to_buffer %arg2 : tensor<256xi32> to memref<256xi32>
      %alloc = memref.alloc() : memref<256xi32>
      func.call @__top__core(%3, %alloc) : (memref<256xi32>, memref<256xi32>) -> ()
      %4 = bufferization.to_tensor %alloc : memref<256xi32> to tensor<256xi32>
      sdy.return %4 : tensor<256xi32>
    } : (tensor<1024xi32>) -> tensor<1024xi32>
    %2 = bufferization.to_buffer %1 : tensor<1024xi32> to memref<1024xi32>
    memref.copy %2, %arg1 : memref<1024xi32> to memref<1024xi32>
    return
  }
  sdy.mesh @__top__core_mesh = <["0"=4]>
  func.func @__top__core(%arg0: memref<256xi32>, %arg1: memref<256xi32>) {
    %0 = call @__top__core_mesh_get_wid() : () -> index
    %1 = arith.index_cast %0 {unsigned} : index to i32
    %c1_i32 = arith.constant 1 : i32
    %2 = arith.cmpi sgt, %1, %c1_i32 : i32
    scf.if %2 {
      %3 = arith.index_cast %0 : index to i32
      %alloc = memref.alloc() : memref<256xi32>
      linalg.fill ins(%3 : i32) outs(%alloc : memref<256xi32>)
      %alloc_0 = memref.alloc() : memref<256xi32>
      linalg.add ins(%arg0, %alloc : memref<256xi32>, memref<256xi32>) outs(%alloc_0 : memref<256xi32>)
      %subview = memref.subview %arg1[0] [256] [1] : memref<256xi32> to memref<256xi32, strided<[1]>>
      memref.copy %alloc_0, %subview : memref<256xi32> to memref<256xi32, strided<[1]>>
    } else {
      %3 = arith.index_cast %0 : index to i32
      %alloc = memref.alloc() : memref<256xi32>
      linalg.fill ins(%3 : i32) outs(%alloc : memref<256xi32>)
      %alloc_0 = memref.alloc() : memref<256xi32>
      linalg.sub ins(%arg0, %alloc : memref<256xi32>, memref<256xi32>) outs(%alloc_0 : memref<256xi32>)
      %subview = memref.subview %arg1[0] [256] [1] : memref<256xi32> to memref<256xi32, strided<[1]>>
      memref.copy %alloc_0, %subview : memref<256xi32> to memref<256xi32, strided<[1]>>
    }
    return
  }
  func.func private @__top__core_mesh_get_wid() -> index
}

The close-to-source IR. works are are not instantiated on each grid point. wid is not a constant for the constructed function so we keep the branches.

'Unrolled' works

after instantiating work on the grid point, wid are constant. we can replace the results of xxx_get_wid with arith.constant.

The work on grid point (0,) is

  func.func @__top__core_0(%arg0: memref<256xi32>, %arg1: memref<256xi32>) {
    %c0 = arith.constant 0 : index
    %0 = arith.index_cast %c0 {unsigned} : index to i32
    %c1_i32 = arith.constant 1 : i32
    %1 = arith.cmpi sgt, %0, %c1_i32 : i32
    scf.if %1 {
      %2 = arith.index_cast %c0 : index to i32
      %alloc = memref.alloc() : memref<256xi32>
      linalg.fill ins(%2 : i32) outs(%alloc : memref<256xi32>)
      %alloc_0 = memref.alloc() : memref<256xi32>
      linalg.add ins(%arg0, %alloc : memref<256xi32>, memref<256xi32>) outs(%alloc_0 : memref<256xi32>)
      %subview = memref.subview %arg1[0] [256] [1] : memref<256xi32> to memref<256xi32, strided<[1]>>
      memref.copy %alloc_0, %subview : memref<256xi32> to memref<256xi32, strided<[1]>>
    } else {
      %2 = arith.index_cast %c0 : index to i32
      %alloc = memref.alloc() : memref<256xi32>
      linalg.fill ins(%2 : i32) outs(%alloc : memref<256xi32>)
      %alloc_0 = memref.alloc() : memref<256xi32>
      linalg.sub ins(%arg0, %alloc : memref<256xi32>, memref<256xi32>) outs(%alloc_0 : memref<256xi32>)
      %subview = memref.subview %arg1[0] [256] [1] : memref<256xi32> to memref<256xi32, strided<[1]>>
      memref.copy %alloc_0, %subview : memref<256xi32> to memref<256xi32, strided<[1]>>
    }
    return
  }

Then run canonicalize pass, we can eliminate the dead branch

  func.func @__top__core_0(%arg0: memref<256xi32>, %arg1: memref<256xi32>) {
    %c0_i32 = arith.constant 0 : i32
    %alloc = memref.alloc() : memref<256xi32>
    linalg.fill ins(%c0_i32 : i32) outs(%alloc : memref<256xi32>)
    %alloc_0 = memref.alloc() : memref<256xi32>
    linalg.sub ins(%arg0, %alloc : memref<256xi32>, memref<256xi32>) outs(%alloc_0 : memref<256xi32>)
    memref.copy %alloc_0, %arg1 : memref<256xi32> to memref<256xi32>
    return
  }

[zhangzhiru]

Do we still need meta_if as frontend interface? We can directly use if / elif / else in the source code, and eliminate the dead branch if the condition is a compile time constant. The DCE seems pretty easy with MLIR's existing passes like canonicalize.

This sounds better than using meta_if. The code also looks cleaner.