sese_loop_tests.sil

// RUN: %target-sil-opt -tf-xla-cfg-canonicalize -tf-ensure-single-loop-exit -assume-parsing-unqualified-ownership-sil %s -o /dev/null | %FileCheck %s


import Builtin
import Swift
import TensorFlow

//--------------------------------------------------------------------------
// A test to make sure that SESE canonicalization works correctly when
// any of the backedges or exit edges are critical edges.
sil private @$loopWithCriticalEdge : $@callee_owned () -> () {
bb0:
  br bb1                                          // id: %0

bb1:                                              // Preds: bb1 bb0
  %1 = graph_op "tfc.RecvFromHost"() : $Builtin.Int1 // user: %2
  cond_br %1, bb1, bb2                            // id: %2

bb2:                                              // Preds: bb1
  %3 = tuple ()                                   // user: %4
  return %3 : $()                                 // id: %4
} // end sil function '$loopWithCriticalEdge'

// CHECK-LABEL: --- XLA CFG Canonicalize: $loopWithCriticalEdge
// CHECK: [sequence
// CHECK:   <while Preheader: {{bb[0-9]+}}, Header: {{bb[0-9]+}}, exit: [[EXIT:bb[0-9]+]]
// CHECK:     [sequence
// CHECK:       {condition Header: {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}}
// CHECK:       block {{bb[0-9]+}}]>
// CHECK:   block [[EXIT]]]
// CHECK: --- XLA CFG Canonicalize end

//--------------------------------------------------------------------------
// Following tests Check that the loop canonicalization is done if header has
// communication with the host even if the loop is in the desired form.

// receiveFromHostGraphOp
sil private @receiveFromHostGraphOp : $@convention(thin) @callee_owned () -> () {
bb0:
  br bb2                                          // id: %0

bb1:                                              // Preds: bb2
  br bb2                                          // id: %1

bb2:                                              // Preds: bb1 bb0
  %2 = graph_op "tfc.RecvFromHost"() {tensorId: i32 0, __device: "/device:CPU:0"} : $Builtin.Int32 // user: %4
  %3 = integer_literal $Builtin.Int32, 0          // user: %4
  %4 = builtin "cmp_eq_Int32"(%2 : $Builtin.Int32, %3 : $Builtin.Int32) : $Builtin.Int1 // user: %5
  cond_br %4, bb1, bb3                            // id: %5

bb3:                                              // Preds: bb2
  %6 = tuple ()                                   // user: %7
  return %6 : $()                                 // id: %7
} // end sil function 'receiveFromHostGraphOp'

// CHECK-LABEL: --- XLA CFG Canonicalize: {{.*}}receiveFromHostGraphOp{{.*}}
// CHECK: [sequence
// CHECK:   <while Preheader: {{bb[0-9]+}}, Header: {{bb[0-9]+}}, exit: [[EXIT:bb[0-9]+]]
// CHECK:     [sequence
// CHECK:       {condition Header: {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}}
// CHECK:       block {{bb[0-9]+}}]>
// CHECK:   block [[EXIT]]]

// sendToHost
sil private @sendToHost : $@convention(thin) @callee_owned () -> Builtin.Int32 {
bb0:
  %0 = integer_literal $Builtin.Int32, 0          // user: %3
  %1 = integer_literal $Builtin.Int32, 10         // user: %6
  %2 = integer_literal $Builtin.Int32, 1          // user: %5
  br bb1(%0 : $Builtin.Int32)                     // id: %3

// %4                                             // user: %5
bb1(%4 : $Builtin.Int32):                         // Preds: bb2 bb0
  %5 = builtin "sadd_with_overflow_Int32"(%4 : $Builtin.Int32, %2 : $Builtin.Int32) : $Builtin.Int32 // users: %10, %9, %6
  %6 = builtin "cmp_slt_Int32"(%5 : $Builtin.Int32, %1 : $Builtin.Int32) : $Builtin.Int1 // users: %8, %7
  %7 = graph_op "tfc.SendToHost"(%6 : $Builtin.Int1) {tensorId: i32 1, __device: "/device:CPU:0"} : $()
  cond_br %6, bb2, bb3                            // id: %8

bb2:                                              // Preds: bb1
  br bb1(%5 : $Builtin.Int32)                     // id: %9

bb3:                                              // Preds: bb1
  return %5 : $Builtin.Int32                      // id: %10
} // end sil function 'sendToHost'
// CHECK-LABEL: --- XLA CFG Canonicalize: {{.*}}sendToHost{{.*}}
// CHECK: [sequence
// CHECK:   <while Preheader: {{bb[0-9]+}}, Header: {{bb[0-9]+}}, exit: [[EXIT:bb[0-9]+]]
// CHECK:     [sequence
// CHECK:       {condition Header: {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}
// CHECK:         block {{bb[0-9]+}}}
// CHECK:       block {{bb[0-9]+}}]>
// CHECK:   block [[EXIT]]]