Home
last modified time | relevance | path

Searched refs:custom_call (Results 1 – 21 of 21) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcusolver_rewriter.cc122 HloInstruction* custom_call = in CreateCholesky() local
125 custom_call->set_metadata(metadata); in CreateCholesky()
126 TF_RETURN_IF_ERROR(custom_call->set_backend_config(options)); in CreateCholesky()
127 return custom_call; in CreateCholesky()
141 HloInstruction * custom_call, in RunOnInstruction()
146 << custom_call->ToString(); in RunOnInstruction()
152 custom_call, 0))); in RunOnInstruction()
Dcudnn_conv_rewriter.cc58 HloInstruction* custom_call = computation->AddInstruction( in CreateCudnnConv() local
60 custom_call->set_window(window); in CreateCudnnConv()
61 custom_call->set_convolution_dimension_numbers(dnums); in CreateCudnnConv()
62 custom_call->set_feature_group_count(feature_group_count); in CreateCudnnConv()
63 custom_call->set_metadata(metadata); in CreateCudnnConv()
64 return custom_call; in CreateCudnnConv()
501 HloInstruction* custom_call = [&]() -> HloInstruction* { in RunOnInstruction() local
534 if (custom_call == nullptr) { in RunOnInstruction()
539 custom_call->set_backend_config(GetDefaultBackendConfig())); in RunOnInstruction()
542 << custom_call->ToString(); in RunOnInstruction()
[all …]
Dir_emitter_unnested.cc373 Status IrEmitterUnnested::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument
380 if (custom_call->custom_call_target() == in HandleCustomCall()
382 const HloInstruction* epsilon = custom_call->operand(5); in HandleCustomCall()
386 const HloInstruction* feature_index = custom_call->operand(6); in HandleCustomCall()
392 /*operand=*/GetAllocationSlice(*custom_call->operand(0)), in HandleCustomCall()
393 /*scale=*/GetAllocationSlice(*custom_call->operand(1)), in HandleCustomCall()
394 /*offset=*/GetAllocationSlice(*custom_call->operand(2)), in HandleCustomCall()
395 /*mean=*/GetAllocationSlice(*custom_call->operand(3)), in HandleCustomCall()
396 /*variance=*/GetAllocationSlice(*custom_call->operand(4)), in HandleCustomCall()
399 /*output=*/GetAllocationSlice(*custom_call), in HandleCustomCall()
[all …]
Dcudnn_conv_rewriter_test.cc309 const HloInstruction* custom_call = in TEST_F() local
312 const WindowDimension& window_dim = custom_call->window().dimensions(i); in TEST_F()
444 const HloInstruction* custom_call = in TEST_F() local
447 const WindowDimension& window_dim = custom_call->window().dimensions(i); in TEST_F()
Dir_emitter.h99 Status HandleCustomCall(HloInstruction* custom_call) override;
Dir_emitter_unnested.h164 Status HandleCustomCall(HloInstruction* custom_call) override;
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/
Dcompiler.cc53 HloInstruction* custom_call, absl::Span<const Literal*> operands) { in HandleEvaluatorCustomCall() argument
56 void* target_fn = registry->Lookup(custom_call->custom_call_target()); in HandleEvaluatorCustomCall()
59 custom_call->custom_call_target()); in HandleEvaluatorCustomCall()
68 auto output = Literal::CreateFromShape(custom_call->shape()); in HandleEvaluatorCustomCall()
/external/tensorflow/tensorflow/compiler/xla/service/
Dhlo_evaluator.h141 HloInstruction* custom_call, absl::Span<const Literal*> operands)>;
148 std::function<StatusOr<Literal>(HloInstruction* custom_call, in set_custom_call_handler() argument
241 Status HandleCustomCall(HloInstruction* custom_call) override;
346 std::function<StatusOr<Literal>(HloInstruction* custom_call,
Dlayout_assignment_test.cc1208 const HloInstruction* custom_call = in TEST_F() local
1210 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F()
1211 ExpectLayoutIs(custom_call->operand(0)->shape(), {0, 1}); in TEST_F()
1212 ExpectLayoutIs(custom_call->operand(1)->shape(), {1, 0}); in TEST_F()
1234 const HloInstruction* custom_call = in TEST_F() local
1236 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F()
1268 const HloInstruction* custom_call = in TEST_F() local
1270 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F()
1271 ExpectTupleLayoutIs(custom_call->operand(0)->shape(), {{1, 0}, {0, 1}}); in TEST_F()
1299 const HloInstruction* custom_call = FindInstruction(m.get(), "custom-call"); in TEST_F() local
[all …]
Dhlo_verifier.cc545 const HloCustomCallInstruction* custom_call = in HandleCustomCall() local
547 TF_RET_CHECK(custom_call != nullptr); in HandleCustomCall()
548 if (custom_call->layout_constrained()) { in HandleCustomCall()
552 TF_RET_CHECK(LayoutUtil::HasLayout(custom_call->shape())); in HandleCustomCall()
553 TF_RET_CHECK(custom_call->operand_count() == in HandleCustomCall()
554 custom_call->operand_shapes_with_layout().size()); in HandleCustomCall()
555 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in HandleCustomCall()
557 custom_call->operand_shapes_with_layout()[i]; in HandleCustomCall()
558 TF_RET_CHECK(ShapeUtil::Compatible(custom_call->operand(i)->shape(), in HandleCustomCall()
560 << custom_call->operand(i)->shape().ToString() << " operand " in HandleCustomCall()
Dlayout_assignment.cc421 const HloCustomCallInstruction* custom_call = in IsLayoutConstrainedCustomCall() local
423 return custom_call != nullptr && custom_call->layout_constrained(); in IsLayoutConstrainedCustomCall()
468 const HloCustomCallInstruction* custom_call = in AddMandatoryConstraints() local
471 constraints->SetInstructionLayout(custom_call->shape(), custom_call)); in AddMandatoryConstraints()
472 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in AddMandatoryConstraints()
474 custom_call->operand_shapes_with_layout()[i], custom_call, i)); in AddMandatoryConstraints()
674 const HloCustomCallInstruction* custom_call = in CheckCustomCallLayout() local
676 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in CheckCustomCallLayout()
678 custom_call->operand(i)->shape(), in CheckCustomCallLayout()
679 custom_call->operand_shapes_with_layout()[i])); in CheckCustomCallLayout()
Ddfs_hlo_visitor_with_default.h145 Status HandleCustomCall(HloInstructionPtr custom_call) override { in HandleCustomCall() argument
146 return DefaultAction(custom_call); in HandleCustomCall()
Dhlo_evaluator.cc1576 Status HloEvaluator::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument
1579 return DefaultAction(custom_call); in HandleCustomCall()
1584 operands.reserve(custom_call->operand_count()); in HandleCustomCall()
1585 for (const HloInstruction* operand : custom_call->operands()) { in HandleCustomCall()
1591 auto output, custom_call_handler_(custom_call, absl::MakeSpan(operands))); in HandleCustomCall()
1593 evaluated_[custom_call] = std::move(output); in HandleCustomCall()
Dbuffer_assignment_test.cc1366 auto custom_call = builder.AddInstruction(HloInstruction::CreateCustomCall( in TEST_F() local
1376 GetAllocation(*assignment, custom_call, /*index=*/{}).maybe_live_out()); in TEST_F()
1378 GetAllocation(*assignment, custom_call, /*index=*/{0}).maybe_live_out()); in TEST_F()
1380 GetAllocation(*assignment, custom_call, /*index=*/{1}).maybe_live_out()); in TEST_F()
2360 HloInstruction* custom_call = main->GetInstructionWithName("custom_call"); in TEST_F() local
2367 EXPECT_EQ(GetAllocation(*buffers, custom_call, {}), in TEST_F()
2369 EXPECT_EQ(GetAllocation(*buffers, custom_call, {0}), in TEST_F()
2371 EXPECT_EQ(GetAllocation(*buffers, custom_call, {1}), in TEST_F()
Dhlo_cost_analysis.h94 Status HandleCustomCall(const HloInstruction* custom_call) override;
Dhlo_evaluator_test.cc3211 [](HloInstruction* custom_call, absl::Span<const Literal*> operands) { in TEST_F() argument
3236 [](HloInstruction* custom_call, absl::Span<const Literal*> operands) { in TEST_F() argument
3237 EXPECT_EQ(HloOpcode::kCustomCall, custom_call->opcode()); in TEST_F()
3238 EXPECT_EQ("_my_custom_call", custom_call->custom_call_target()); in TEST_F()
3239 EXPECT_EQ(2, custom_call->operand_count()); in TEST_F()
3241 auto output = Literal::CreateFromShape(custom_call->shape()); in TEST_F()
Dhlo_instruction.cc3443 if (auto custom_call = DynCast<HloCustomCallInstruction>(this)) { in convolution_dimension_numbers() local
3444 return custom_call->convolution_dimension_numbers(); in convolution_dimension_numbers()
3453 } else if (auto custom_call = DynCast<HloCustomCallInstruction>(this)) { in set_convolution_dimension_numbers() local
3454 custom_call->set_convolution_dimension_numbers(dnums); in set_convolution_dimension_numbers()
Dhlo_parser_test.cc1281 R"(HloModule custom_call in CreateTestCases()
1293 R"(HloModule custom_call in CreateTestCases()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dcustom_call_test.cc175 auto custom_call = b.AddInstruction(HloInstruction::CreateCustomCall( in XLA_TEST_F() local
178 custom_call->CloneWithNewOperands(r2f32_dim0_major, {custom_call})); in XLA_TEST_F()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.cc2229 Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument
2230 absl::Span<HloInstruction* const> operands(custom_call->operands()); in HandleCustomCall()
2265 custom_call->custom_call_target(), in HandleCustomCall()
2272 TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call)); in HandleCustomCall()
2274 if (custom_call->shape().IsTuple()) { in HandleCustomCall()
2276 for (int i = 0; i < ShapeUtil::TupleElementCount(custom_call->shape()); in HandleCustomCall()
2279 ShapeUtil::GetTupleElementShape(custom_call->shape(), i); in HandleCustomCall()
2282 assignment_.GetUniqueSlice(custom_call, {i})); in HandleCustomCall()
2286 llvm_ir::EmitTuple(GetIrArrayFor(custom_call), base_ptrs, &b_); in HandleCustomCall()
2289 PointerCast(GetEmittedValueFor(custom_call), i8_ptr_type); in HandleCustomCall()
Dir_emitter.h177 Status HandleCustomCall(HloInstruction* custom_call) override;