| /external/tensorflow/tensorflow/compiler/xla/service/ |
| D | while_loop_simplifier.cc | 126 const HloInstruction* while_op) { in GetLoopInductionVarTupleIdx() argument
127 CHECK_EQ(while_op->opcode(), HloOpcode::kWhile); in GetLoopInductionVarTupleIdx()
129 << while_op->ToShortString(); in GetLoopInductionVarTupleIdx()
137 auto* while_cond = while_op->while_condition(); in GetLoopInductionVarTupleIdx()
155 auto* while_body = while_op->while_body(); in GetLoopInductionVarTupleIdx()
183 auto* while_init = while_op->operand(0); in GetLoopInductionVarTupleIdx()
195 static optional<int64> GetLoopTripCount(HloInstruction* while_op) { in GetLoopTripCount() argument
196 CHECK_EQ(while_op->opcode(), HloOpcode::kWhile); in GetLoopTripCount()
197 VLOG(2) << "Getting trip count for loop " << while_op->ToString(); in GetLoopTripCount()
204 optional<int64> indvar_tuple_idx = GetLoopInductionVarTupleIdx(while_op); in GetLoopTripCount()
[all …]
|
| D | while_loop_simplifier_test.cc | 128 auto* while_op = computation->root_instruction(); in TEST_F() local
129 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); in TEST_F()
130 auto* true_op = while_op->while_body()->AddInstruction( in TEST_F()
133 while_op->while_body()->root_instruction())); in TEST_F()
144 auto* while_op = computation->root_instruction(); in TEST_F() local
145 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); in TEST_F()
146 auto* while_body = while_op->while_body(); in TEST_F()
157 auto* while_op = computation->root_instruction(); in TEST_F() local
158 ASSERT_EQ(while_op->opcode(), HloOpcode::kWhile); in TEST_F()
159 auto* while_body = while_op->while_body(); in TEST_F()
[all …]
|
| D | hlo_scheduling_test.cc | 74 HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( in TEST_F() local
87 module_sequence[entry_computation] = {iter, data, tuple, while_op}; in TEST_F()
|
| D | heap_simulator_test.cc | 442 HloInstruction* while_op = builder.AddInstruction(HloInstruction::CreateWhile( in TEST_F() local
447 {param, while_op, body_param, cond_param, cond_iter, cond_data, cond_lt}); in TEST_F()
453 {kAlloc, tracker.BufferAt(while_op, {})}, in TEST_F()
454 {kAlloc, tracker.BufferAt(while_op, {0})}, in TEST_F()
455 {kAlloc, tracker.BufferAt(while_op, {1})}, in TEST_F()
481 {kFree, tracker.BufferAt(while_op, {})}, in TEST_F()
482 {kFree, tracker.BufferAt(while_op, {0})}, in TEST_F()
483 {kFree, tracker.BufferAt(while_op, {1})}, in TEST_F()
|
| D | buffer_assignment_test.cc | 701 auto while_op = builder.AddInstruction(HloInstruction::CreateWhile( in TEST_F() local
705 const std::vector<const HloInstruction*> level0 = GetInstructions(while_op); in TEST_F()
738 while_op->shape(), in TEST_F()
739 [this, &buffers, while_op, body_root](const Shape& /*subshape*/, in TEST_F()
741 auto while_op_allocation = GetAllocation(*buffers, while_op, index); in TEST_F()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | BUILD | 96 ":while_op",
142 name = "while_op",
143 srcs = ["while_op.cc"],
144 hdrs = ["while_op.h"],
|
| /external/tensorflow/tensorflow/compiler/tf2xla/ |
| D | functionalize_control_flow_test.cc | 251 auto while_op = in TEST() local
254 auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); in TEST()
354 auto while_op = in TEST() local
496 auto while_op = in TEST() local
499 auto sink_x = ops::Identity(scope.WithOpName("sink_x"), while_op[0]); in TEST()
500 auto sink_y = ops::Identity(scope.WithOpName("sink_y"), while_op[1]); in TEST()
736 auto while_op = ops::XlaWhile(scope.WithOpName("outer/LoopCond"), in TEST() local
739 auto sink = ops::Identity(scope.WithOpName("sink"), while_op[0]); in TEST()
792 auto while_op = in TEST() local
802 while_op[0].op(), while_op[1].op()}), in TEST()
|