Back to PBT run

input.mlir

module {
func.func @helper_double(%a: i32) -> i32 {
  %d = arith.addi %a, %a : i32
  return %d : i32
}

llvm.func @extern_scale(i32) -> i32

func.func @host_container(%target: memref<4xi32>, %value: i32) {
  %hz = arith.constant 0 : index
  %hn = arith.constant 4 : index
  %hs = arith.constant 1 : index
  scf.for %hi = %hz to %hn step %hs {
    memref.store %value, %target[%hi] : memref<4xi32>
  }
  return
}

dataflow.thread private @thread_0 domain(#dataflow.thread_domain<dense>)(%mem: memref<4xi32>, %val: i32, %flag: i1) ctrl (%start: none) {
  %lo = arith.constant 0 : index
  %hi = arith.constant 2 : index
  %stp = arith.constant 1 : index
  scf.for %iv = %lo to %hi step %stp {
    "loom.spatial_region"(%val, %mem)
        <{operandSegmentSizes = array<i32: 1, 0, 1, 0>,
          resultSegmentSizes = array<i32: 0, 0>}> ({
      ^bb0(%payload: i32, %memory: memref<4xi32>):
        %z = arith.constant 0 : index
        memref.store %payload, %memory[%z] : memref<4xi32>
        "loom.spatial_yield"()
            <{operandSegmentSizes = array<i32: 0, 0>}> : () -> ()
    }) {graph_name = "graph_0", source_maps = []} : (i32, memref<4xi32>) -> ()
  }
    "loom.spatial_region"(%val, %mem)
        <{operandSegmentSizes = array<i32: 1, 0, 1, 0>,
          resultSegmentSizes = array<i32: 0, 0>}> ({
      ^bb0(%payload: i32, %memory: memref<4xi32>):
        %z = arith.constant 0 : index
        memref.store %payload, %memory[%z] : memref<4xi32>
        "loom.spatial_yield"()
            <{operandSegmentSizes = array<i32: 0, 0>}> : () -> ()
    }) {graph_name = "graph_1", source_maps = []} : (i32, memref<4xi32>) -> ()
  dataflow.thread.yield
}

dataflow.thread private @thread_1 domain(#dataflow.thread_domain<dense>)(%mem: memref<4xi32>, %val: i32, %flag: i1) ctrl (%start: none) {
  %lo = arith.constant 0 : index
  %hi = arith.constant 2 : index
  %stp = arith.constant 1 : index
  %icall = func.call @helper_double(%val) : (i32) -> i32
    "loom.spatial_region"(%val, %mem)
        <{operandSegmentSizes = array<i32: 1, 0, 1, 0>,
          resultSegmentSizes = array<i32: 0, 0>}> ({
      ^bb0(%payload: i32, %memory: memref<4xi32>):
        %z = arith.constant 0 : index
        memref.store %payload, %memory[%z] : memref<4xi32>
        "loom.spatial_yield"()
            <{operandSegmentSizes = array<i32: 0, 0>}> : () -> ()
    }) {graph_name = "graph_2", source_maps = []} : (i32, memref<4xi32>) -> ()
  scf.if %flag {
    %r3 = "loom.spatial_region"(%val)
        <{operandSegmentSizes = array<i32: 1, 0, 0, 0>,
          resultSegmentSizes = array<i32: 1, 0>}> ({
      ^bb0(%a: i32):
        %called = llvm.call @extern_scale(%a) : (i32) -> i32
        "loom.spatial_yield"(%called)
            <{operandSegmentSizes = array<i32: 1, 0>}> : (i32) -> ()
    }) {graph_name = "graph_3", source_maps = []} : (i32) -> i32
  }
  dataflow.thread.yield
}

}