Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcusolver_rewriter.cc91 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()
Dgpu_conv_rewriter.cc59 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 …]
Dgpu_conv_rewriter_test.cc312 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()
Dir_emitter.h96 Status HandleCustomCall(HloInstruction* custom_call) override;
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dmhlo_to_lhlo_with_xla.cc636 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 …]
Dmhlo_to_lhlo_with_xla.h67 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/
Dcompiler.cc54 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/
Dlogical_buffer_analysis.cc192 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()
Dmemory_space_assignment_utils.cc93 if (auto* custom_call = in IsValueAllowedInAlternateMemory() local
95 for (const auto& pair : custom_call->output_to_operand_aliasing()) { in IsValueAllowedInAlternateMemory()
Dhlo_evaluator.h147 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,
Dhlo_verifier.cc793 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 …]
Dlayout_assignment_test.cc1202 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 …]
Dtuple_points_to_analysis.cc478 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()
Dlayout_assignment.cc431 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()
Dlogical_buffer_analysis.h71 Status HandleCustomCall(HloInstruction* custom_call) override;
Ddfs_hlo_visitor_with_default.h164 Status HandleCustomCall(HloInstructionPtr custom_call) override { in HandleCustomCall() argument
165 return DefaultAction(custom_call); in HandleCustomCall()
Dhlo_dataflow_analysis.cc436 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()
Dhlo_dataflow_analysis.h219 bool UpdateCustomCallValueSet(HloInstruction* custom_call);
Dtuple_points_to_analysis.h259 Status HandleCustomCall(HloInstruction* custom_call) override;
Dbuffer_assignment_test.cc1538 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()
Dhlo_cost_analysis.h103 Status HandleCustomCall(const HloInstruction* custom_call) override;
/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/mlir/hlo/tests/
Dhlo-legalize-to-lhlo.mlir598 // 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/
Dir_emitter.cc2295 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/
Dhlo_function_importer.cc319 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()

12