Home
last modified time | relevance | path

Searched refs:GetUniqueSlice (Results 1 – 10 of 10) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/
Dbuffer_assignment_test.cc264 return *buffers.GetUniqueSlice(hlo, index).ConsumeValueOrDie().allocation(); in GetAllocation()
1544 assignment->GetUniqueSlice(tuple_param0, /*index=*/{0}) in TEST_F()
1546 assignment->GetUniqueSlice(tuple_param1, /*index=*/{0}) in TEST_F()
1959 EXPECT_EQ(assignment->GetUniqueSlice(input0, {}).ConsumeValueOrDie(), in TEST_F()
1960 assignment->GetUniqueSlice(while0, {0}).ConsumeValueOrDie()); in TEST_F()
1962 EXPECT_EQ(assignment->GetUniqueSlice(weights0, {}).ConsumeValueOrDie(), in TEST_F()
1963 assignment->GetUniqueSlice(while0, {1}).ConsumeValueOrDie()); in TEST_F()
1965 EXPECT_EQ(assignment->GetUniqueSlice(while0, {2}).ConsumeValueOrDie(), in TEST_F()
1966 assignment->GetUniqueSlice(while1, {0}).ConsumeValueOrDie()); in TEST_F()
1968 EXPECT_EQ(assignment->GetUniqueSlice(weights1, {}).ConsumeValueOrDie(), in TEST_F()
[all …]
Dbuffer_assignment.cc410 StatusOr<BufferAllocation::Slice> BufferAssignment::GetUniqueSlice( in GetUniqueSlice() function in xla::BufferAssignment
444 return GetUniqueSlice(instruction, /*index=*/{}); in GetUniqueTopLevelSlice()
450 return GetUniqueSlice(hlo_a, shape_index_a).ConsumeValueOrDie() == in SharesSliceAtIndex()
451 GetUniqueSlice(hlo_b, shape_index_b).ConsumeValueOrDie(); in SharesSliceAtIndex()
Dbuffer_assignment.h397 StatusOr<BufferAllocation::Slice> GetUniqueSlice(
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dir_emitter_unnested.cc417 auto output_data = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); in HandleCustomCall()
418 auto output_mean = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); in HandleCustomCall()
419 auto output_inv_stddev = assn.GetUniqueSlice(custom_call, {2}).ValueOrDie(); in HandleCustomCall()
447 auto output_grad_data = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); in HandleCustomCall()
448 auto output_grad_scale = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); in HandleCustomCall()
450 assn.GetUniqueSlice(custom_call, {2}).ValueOrDie(); in HandleCustomCall()
475 auto conv_result_slice = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); in HandleCustomCall()
476 auto scratch_slice = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); in HandleCustomCall()
500 auto a_buffer = assn.GetUniqueSlice(custom_call, {0}).ValueOrDie(); in HandleCustomCall()
501 auto workspace_buffer = assn.GetUniqueSlice(custom_call, {1}).ValueOrDie(); in HandleCustomCall()
[all …]
Dir_emitter.h141 .GetUniqueSlice(&hlo, index)
Dhlo_to_ir_bindings.cc104 buffer_assignment_->GetUniqueSlice(non_io_hlo, index); in EmitBasePointersForHlos()
Dgpu_executable.cc312 this->assignment_->GetUniqueSlice(src_hlo, sources[0]->index())); in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_executable.cc236 this->assignment_->GetUniqueSlice(src, buffer_source->index())); in CreateResultShapedBuffer()
Dir_emitter.cc338 assignment_.GetUniqueSlice(infeed, {0})); in HandleInfeed()
341 assignment_.GetUniqueSlice(infeed, {1})); in HandleInfeed()
356 assignment_.GetUniqueSlice(infeed, {0, i})); in HandleInfeed()
1335 assignment_.GetUniqueSlice(crs, {i})); in HandleAllReduce()
2282 assignment_.GetUniqueSlice(custom_call, {i})); in HandleCustomCall()
2311 assignment_.GetUniqueSlice(a, index).ConsumeValueOrDie(); in HandleWhile()
2313 assignment_.GetUniqueSlice(b, index).ConsumeValueOrDie(); in HandleWhile()
2963 assignment_.GetUniqueSlice(target_op, {i})); in EmitTargetElementLoop()
Dir_emitter.h193 return assignment_.GetUniqueSlice(&hlo, index).ConsumeValueOrDie();