/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | cusolver_rewriter.cc | 91 HloInstruction* custom_call = in CreateCholesky() local 94 custom_call->set_metadata(metadata); in CreateCholesky() 95 TF_RETURN_IF_ERROR(custom_call->set_backend_config(options)); in CreateCholesky() 97 HloInstruction::CreateGetTupleElement(a_shape, custom_call, 0)); in CreateCholesky() 99 HloInstruction::CreateGetTupleElement(info_shape, custom_call, 2)); in CreateCholesky() 139 HloInstruction * custom_call, in RunOnInstruction() 144 << custom_call->ToString(); in RunOnInstruction() 147 instruction->parent()->ReplaceInstruction(instruction, custom_call)); in RunOnInstruction()
|
D | gpu_conv_rewriter.cc | 59 HloInstruction* custom_call = computation->AddInstruction( in CreateGpuConv() local 61 custom_call->set_window(window); in CreateGpuConv() 62 custom_call->set_convolution_dimension_numbers(dnums); in CreateGpuConv() 63 custom_call->set_feature_group_count(feature_group_count); in CreateGpuConv() 64 custom_call->set_metadata(metadata); in CreateGpuConv() 65 return custom_call; in CreateGpuConv() 695 TF_ASSIGN_OR_RETURN(HloInstruction * custom_call, in RunOnInstruction() 697 if (custom_call == nullptr) { in RunOnInstruction() 702 custom_call->set_backend_config(GetDefaultBackendConfig())); in RunOnInstruction() 705 << custom_call->ToString(); in RunOnInstruction() [all …]
|
D | gpu_conv_rewriter_test.cc | 312 const HloInstruction* custom_call = in TEST_F() local 315 const WindowDimension& window_dim = custom_call->window().dimensions(i); in TEST_F() 451 const HloInstruction* custom_call = in TEST_F() local 454 const WindowDimension& window_dim = custom_call->window().dimensions(i); in TEST_F()
|
D | ir_emitter.h | 96 Status HandleCustomCall(HloInstruction* custom_call) override;
|
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
D | mhlo_to_lhlo_with_xla.cc | 636 TF_ASSIGN_OR_RETURN(auto custom_call, in EmitCustomCallOp() 639 custom_call.call_target_nameAttr( in EmitCustomCallOp() 641 custom_call.backend_configAttr( in EmitCustomCallOp() 645 custom_call->setAttr(lmhlo::CustomCallOp::getOperandSegmentSizeAttr(), in EmitCustomCallOp() 647 return custom_call.getOperation(); in EmitCustomCallOp() 651 const HloCustomCallInstruction* custom_call) { in EmitCholesky() argument 653 CreateOpWithoutAttrs<lmhlo_gpu::CholeskyOp>(custom_call)); in EmitCholesky() 655 custom_call->backend_config<xla::CholeskyOptions>()); in EmitCholesky() 661 const HloCustomCallInstruction* custom_call) { in EmitGemm() argument 664 custom_call->backend_config<xla::gpu::GemmBackendConfig>()); in EmitGemm() [all …]
|
D | mhlo_to_lhlo_with_xla.h | 67 const xla::HloCustomCallInstruction* custom_call); 69 const xla::HloCustomCallInstruction* custom_call); 71 const xla::HloCustomCallInstruction* custom_call); 73 const xla::HloCustomCallInstruction* custom_call);
|
/external/tensorflow/tensorflow/compiler/xla/service/interpreter/ |
D | compiler.cc | 54 HloInstruction* custom_call, absl::Span<const Literal*> operands) { in HandleEvaluatorCustomCall() argument 57 void* target_fn = registry->Lookup(custom_call->custom_call_target(), "Host"); in HandleEvaluatorCustomCall() 60 custom_call->custom_call_target()); in HandleEvaluatorCustomCall() 69 auto output = Literal::CreateFromShape(custom_call->shape()); in HandleEvaluatorCustomCall()
|
/external/tensorflow/tensorflow/compiler/xla/service/ |
D | logical_buffer_analysis.cc | 192 Status LogicalBufferAnalysis::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument 193 auto ccall = Cast<HloCustomCallInstruction>(custom_call); in HandleCustomCall() 201 NewLogicalBuffer(custom_call, index); in HandleCustomCall()
|
D | memory_space_assignment_utils.cc | 93 if (auto* custom_call = in IsValueAllowedInAlternateMemory() local 95 for (const auto& pair : custom_call->output_to_operand_aliasing()) { in IsValueAllowedInAlternateMemory()
|
D | hlo_evaluator.h | 147 HloInstruction* custom_call, absl::Span<const Literal*> operands)>; 154 std::function<StatusOr<Literal>(HloInstruction* custom_call, in set_custom_call_handler() argument 265 Status HandleCustomCall(HloInstruction* custom_call) override; 370 std::function<StatusOr<Literal>(HloInstruction* custom_call,
|
D | hlo_verifier.cc | 793 const HloCustomCallInstruction* custom_call = in HandleCustomCall() local 795 TF_RET_CHECK(custom_call != nullptr); in HandleCustomCall() 796 if (custom_call->layout_constrained()) { in HandleCustomCall() 800 TF_RET_CHECK(LayoutUtil::HasLayout(custom_call->shape())); in HandleCustomCall() 801 TF_RET_CHECK(custom_call->operand_count() == in HandleCustomCall() 802 custom_call->operand_shapes_with_layout().size()); in HandleCustomCall() 803 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in HandleCustomCall() 805 custom_call->operand_shapes_with_layout()[i]; in HandleCustomCall() 806 TF_RET_CHECK(ShapeUtil::Compatible(custom_call->operand(i)->shape(), in HandleCustomCall() 808 << custom_call->operand(i)->shape().ToString() << " operand " in HandleCustomCall() [all …]
|
D | layout_assignment_test.cc | 1202 const HloInstruction* custom_call = in TEST_F() local 1204 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F() 1205 ExpectLayoutIs(custom_call->operand(0)->shape(), {0, 1}); in TEST_F() 1206 ExpectLayoutIs(custom_call->operand(1)->shape(), {1, 0}); in TEST_F() 1228 const HloInstruction* custom_call = in TEST_F() local 1230 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F() 1262 const HloInstruction* custom_call = in TEST_F() local 1264 ExpectLayoutIs(custom_call->shape(), {3, 2, 0, 1}); in TEST_F() 1265 ExpectTupleLayoutIs(custom_call->operand(0)->shape(), {{1, 0}, {0, 1}}); in TEST_F() 1293 const HloInstruction* custom_call = FindInstruction(m.get(), "custom-call"); in TEST_F() local [all …]
|
D | tuple_points_to_analysis.cc | 478 Status TuplePointsToAnalysis::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument 479 auto ccall = Cast<HloCustomCallInstruction>(custom_call); in HandleCustomCall() 480 PointsToSet& points_to_set = CreateEmptyPointsToSet(custom_call); in HandleCustomCall() 490 logical_buffer_analysis_->GetBuffer(custom_call, index), index); in HandleCustomCall() 504 points_to_set.add_tuple_source({}, custom_call); in HandleCustomCall()
|
D | layout_assignment.cc | 431 const HloCustomCallInstruction* custom_call = in IsLayoutConstrainedCustomCall() local 433 return custom_call != nullptr && custom_call->layout_constrained(); in IsLayoutConstrainedCustomCall() 482 const HloCustomCallInstruction* custom_call = in AddMandatoryConstraints() local 485 constraints->SetInstructionLayout(custom_call->shape(), custom_call)); in AddMandatoryConstraints() 486 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in AddMandatoryConstraints() 488 custom_call->operand_shapes_with_layout()[i], custom_call, i)); in AddMandatoryConstraints() 691 const HloCustomCallInstruction* custom_call = in CheckCustomCallLayout() local 693 for (int64 i = 0; i < custom_call->operand_count(); ++i) { in CheckCustomCallLayout() 695 LayoutsInShapesEqual(custom_call->operand(i)->shape(), in CheckCustomCallLayout() 696 custom_call->operand_shapes_with_layout()[i])); in CheckCustomCallLayout()
|
D | logical_buffer_analysis.h | 71 Status HandleCustomCall(HloInstruction* custom_call) override;
|
D | dfs_hlo_visitor_with_default.h | 164 Status HandleCustomCall(HloInstructionPtr custom_call) override { in HandleCustomCall() argument 165 return DefaultAction(custom_call); in HandleCustomCall()
|
D | hlo_dataflow_analysis.cc | 436 HloInstruction* custom_call) { in UpdateCustomCallValueSet() argument 437 CHECK_EQ(custom_call->opcode(), HloOpcode::kCustomCall); in UpdateCustomCallValueSet() 439 for (const auto& aliasing : Cast<HloCustomCallInstruction>(custom_call) in UpdateCustomCallValueSet() 442 custom_call->operand(aliasing.second.first), aliasing.second.second); in UpdateCustomCallValueSet() 443 HloValueSet& value_set = GetValueSet(custom_call, aliasing.first); in UpdateCustomCallValueSet()
|
D | hlo_dataflow_analysis.h | 219 bool UpdateCustomCallValueSet(HloInstruction* custom_call);
|
D | tuple_points_to_analysis.h | 259 Status HandleCustomCall(HloInstruction* custom_call) override;
|
D | buffer_assignment_test.cc | 1538 auto custom_call = builder.AddInstruction(HloInstruction::CreateCustomCall( in TEST_F() local 1548 GetAllocation(*assignment, custom_call, /*index=*/{}).maybe_live_out()); in TEST_F() 1550 GetAllocation(*assignment, custom_call, /*index=*/{0}).maybe_live_out()); in TEST_F() 1552 GetAllocation(*assignment, custom_call, /*index=*/{1}).maybe_live_out()); in TEST_F() 2507 HloInstruction* custom_call = main->GetInstructionWithName("custom_call"); in TEST_F() local 2514 EXPECT_EQ(GetAllocation(*buffers, custom_call, {}), in TEST_F() 2516 EXPECT_EQ(GetAllocation(*buffers, custom_call, {0}), in TEST_F() 2518 EXPECT_EQ(GetAllocation(*buffers, custom_call, {1}), in TEST_F()
|
D | hlo_cost_analysis.h | 103 Status HandleCustomCall(const HloInstruction* custom_call) override;
|
/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/mlir/hlo/tests/ |
D | hlo-legalize-to-lhlo.mlir | 598 // CHECK-LABEL: func @custom_call 600 func @custom_call(%arg0: tensor<2x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<4x4xf16> { 601 …// CHECK: "lmhlo.custom_call"([[ARG0]], [[ARG1]], %{{.*}}) {backend_config = "", call_target_name … 602 %result = "mhlo.custom_call"(%arg0, %arg1) 613 …// CHECK: "lmhlo.custom_call"([[ARG0]], [[ARG1]], %{{.*}}, %{{.*}}) {backend_config = "", call_tar… 614 %temp:2 = "mhlo.custom_call"(%arg0, %arg1)
|
/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
D | ir_emitter.cc | 2295 Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) { in HandleCustomCall() argument 2296 if (custom_call->custom_call_target() == "PadToStatic") { in HandleCustomCall() 2297 return HandlePadToStatic(custom_call); in HandleCustomCall() 2299 if (custom_call->custom_call_target() == "SliceToDynamic") { in HandleCustomCall() 2300 return HandleSliceToDynamic(custom_call); in HandleCustomCall() 2302 if (custom_call->custom_call_target() == "TopK") { in HandleCustomCall() 2303 return HandleTopK(custom_call); in HandleCustomCall() 2305 absl::Span<HloInstruction* const> operands(custom_call->operands()); in HandleCustomCall() 2332 TF_RETURN_IF_ERROR(EmitTargetAddressForOp(custom_call)); in HandleCustomCall() 2334 if (custom_call->shape().IsTuple()) { in HandleCustomCall() [all …]
|
/external/tensorflow/tensorflow/compiler/mlir/xla/ |
D | hlo_function_importer.cc | 319 auto custom_call = Cast<HloCustomCallInstruction>(instruction); in ImportInstructionImpl() local 322 builder_->getStringAttr(custom_call->custom_call_target()))); in ImportInstructionImpl() 325 builder_->getBoolAttr(custom_call->custom_call_has_side_effect()))); in ImportInstructionImpl() 328 builder_->getStringAttr(custom_call->raw_backend_config_string()))); in ImportInstructionImpl()
|