/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | cusolver_rewriter.cc | 122 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()
|
D | cudnn_conv_rewriter.cc | 58 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 …]
|
D | ir_emitter_unnested.cc | 373 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 …]
|
D | cudnn_conv_rewriter_test.cc | 309 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()
|
D | ir_emitter.h | 99 Status HandleCustomCall(HloInstruction* custom_call) override;
|
D | ir_emitter_unnested.h | 164 Status HandleCustomCall(HloInstruction* custom_call) override;
|
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
D | compiler.cc | 53 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/ |
D | hlo_evaluator.h | 141 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,
|
D | layout_assignment_test.cc | 1208 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 …]
|
D | hlo_verifier.cc | 545 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()
|
D | layout_assignment.cc | 421 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()
|
D | dfs_hlo_visitor_with_default.h | 145 Status HandleCustomCall(HloInstructionPtr custom_call) override { in HandleCustomCall() argument 146 return DefaultAction(custom_call); in HandleCustomCall()
|
D | hlo_evaluator.cc | 1576 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()
|
D | buffer_assignment_test.cc | 1366 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()
|
D | hlo_cost_analysis.h | 94 Status HandleCustomCall(const HloInstruction* custom_call) override;
|
D | hlo_evaluator_test.cc | 3211 [](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()
|
D | hlo_instruction.cc | 3443 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()
|
D | hlo_parser_test.cc | 1281 R"(HloModule custom_call in CreateTestCases() 1293 R"(HloModule custom_call in CreateTestCases()
|
/external/tensorflow/tensorflow/compiler/xla/tests/ |
D | custom_call_test.cc | 175 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/ |
D | ir_emitter.cc | 2229 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()
|
D | ir_emitter.h | 177 Status HandleCustomCall(HloInstruction* custom_call) override;
|