Back to PBT run

test.mlir

// RUN: loom-raise-opt --loom-lower-for-to-graph --mlir-print-op-generic %s | FileCheck %s

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


func.func @host_container(%target: memref<4xi32>, %value: i32) {
  return
}

dataflow.thread private @thread_0 domain(#dataflow.thread_domain<dense>)(%mem: memref<4xi32>, %val: i32, %flag: i1) ctrl (%start: none) {
    %r0 = "loom.spatial_region"(%val, %flag)
        <{operandSegmentSizes = array<i32: 2, 0, 0, 0>,
          resultSegmentSizes = array<i32: 1, 0>}> ({
      ^bb0(%a: i32, %c: i1):
        %sel = scf.if %c -> (i32) {
          %t = arith.addi %a, %a : i32
          scf.yield %t : i32
        } else {
          %e = arith.muli %a, %a : i32
          scf.yield %e : i32
        }
        "loom.spatial_yield"(%sel)
            <{operandSegmentSizes = array<i32: 1, 0>}> : (i32) -> ()
    }) {graph_name = "graph_0", source_maps = []} : (i32, i1) -> i32
  dataflow.thread.yield
}

}

// CHECK: "builtin.module"() ({
// CHECK-NEXT:   "func.func"() <{function_type = (i32) -> i32, sym_name = "helper_double"}> ({
// CHECK-NEXT:   ^bb0(%arg9: i32):
// CHECK-NEXT:     %8 = "arith.addi"(%arg9, %arg9) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
// CHECK-NEXT:     "func.return"(%8) : (i32) -> ()
// CHECK-NEXT:   }) : () -> ()
// CHECK-NEXT:   "func.func"() <{function_type = (memref<4xi32>, i32) -> (), sym_name = "host_container"}> ({
// CHECK-NEXT:   ^bb0(%arg7: memref<4xi32>, %arg8: i32):
// CHECK-NEXT:     "func.return"() : () -> ()
// CHECK-NEXT:   }) : () -> ()
// CHECK-NEXT:   "dataflow.thread"() <{domain = #dataflow.thread_domain<dense>, function_type = (memref<4xi32>, i32, i1) -> (), sym_name = "thread_0", sym_visibility = "private"}> ({
// CHECK-NEXT:   ^bb0(%arg3: memref<4xi32>, %arg4: i32, %arg5: i1, %arg6: none):
// CHECK-NEXT:     %7:2 = "dataflow.graph.launch"(%arg6, %arg4, %arg5) <{callee = @graph_0, operandSegmentSizes = array<i32: 1, 2, 0, 0, 0>, resultSegmentSizes = array<i32: 1, 0, 1>, source_maps = []}> : (none, i32, i1) -> (i32, none)
// CHECK-NEXT:     "dataflow.thread.yield"(%7#1) : (none) -> ()
// CHECK-NEXT:   }) : () -> ()
// CHECK-NEXT:   "dataflow.graph"() <{function_type = (i32, i1) -> i32, input_segments = array<i32: 2, 0, 0>, result_segments = array<i32: 1, 0, 0>, sym_name = "graph_0", sym_visibility = "private"}> ({
// CHECK-NEXT:   ^bb0(%arg0: none, %arg1: i32, %arg2: i1):
// CHECK-NEXT:     %0:2 = "dataflow.demux"(%arg2, %arg0) : (i1, none) -> (none, none)
// CHECK-NEXT:     %1:2 = "dataflow.demux"(%arg2, %arg1) : (i1, i32) -> (i32, i32)
// CHECK-NEXT:     %2 = "arith.muli"(%1#0, %1#0) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
// CHECK-NEXT:     %3 = "arith.addi"(%1#1, %1#1) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
// CHECK-NEXT:     %4 = "dataflow.mux"(%arg2, %2, %3) : (i1, i32, i32) -> i32
// CHECK-NEXT:     %5 = "dataflow.mux"(%arg2, %0#0, %0#1) : (i1, none, none) -> none
// CHECK-NEXT:     %6:2 = "dataflow.sync"(%5, %4) : (none, i32) -> (none, i32)
// CHECK-NEXT:     "dataflow.graph.return"(%6#1, %6#0) <{operandSegmentSizes = array<i32: 1, 0, 0, 1>}> : (i32, none) -> ()
// CHECK-NEXT:   }) : () -> ()
// CHECK-NEXT: }) : () -> ()
// CHECK-EMPTY: