| 552 | base_attrs_(base_attrs) {} |
| 553 | |
| 554 | Maybe<void> UserOpExpr::Init(const std::shared_ptr<const UserOpExpr>& self) { |
| 555 | const auto& op_type_name = op_proto_.op_type_name(); |
| 556 | const auto* registry = user_op::UserOpRegistryMgr::Get().GetOpRegistryResult(op_type_name); |
| 557 | CHECK_NOTNULL_OR_RETURN(registry); |
| 558 | logical_tensor_desc_infer_fn_ = registry->logical_tensor_desc_infer_fn; |
| 559 | CHECK_OR_RETURN(static_cast<bool>(logical_tensor_desc_infer_fn_)) |
| 560 | << Error::RuntimeError() << "registry->logical_tensor_desc_infer_fn failed."; |
| 561 | physical_tensor_desc_infer_fn_ = registry->physical_tensor_desc_infer_fn; |
| 562 | CHECK_OR_RETURN(static_cast<bool>(physical_tensor_desc_infer_fn_)) |
| 563 | << Error::RuntimeError() << "registry->logical_tensor_desc_infer_fn failed."; |
| 564 | dtype_infer_fn_ = registry->data_type_infer_fn; |
| 565 | CHECK_OR_RETURN(static_cast<bool>(dtype_infer_fn_)) |
| 566 | << Error::RuntimeError() << "registry->data_type_infer_fn failed."; |
| 567 | if (registry->device_and_stream_infer_fn) { |
| 568 | device_and_stream_infer_fn_ = registry->device_and_stream_infer_fn; |
| 569 | } |
| 570 | local_tensor_infer_cache_.reset(new LocalTensorInferCache(self)); |
| 571 | global_tensor_infer_cache_.reset(new GlobalTensorInferCache(self)); |
| 572 | const auto& indexed_input_pairs = this->indexed_input_pairs(); |
| 573 | for (int32_t i = 0; i < indexed_input_pairs.size(); ++i) { |
| 574 | const auto& input_pair = JUST(VectorAt(indexed_input_pairs, i)); |
| 575 | if (user_op::UserOpHostMemoryInputRegistry::Get().IsHostMemoryInput4Op( |
| 576 | op_type_name, input_pair.first, input_pair.second)) { |
| 577 | host_memory_input_ids_.emplace_back(i); |
| 578 | } |
| 579 | } |
| 580 | return Maybe<void>::Ok(); |
| 581 | } |
| 582 | |
| 583 | /* static */ Maybe<UserOpExpr> UserOpExpr::New(const std::string& op_name, UserOpConf&& op_proto, |
| 584 | const std::vector<std::string>& indexed_ibns, |
no test coverage detected