| 156 | } |
| 157 | |
| 158 | Status CgmodeCompileOp::Compile(OpKernelContext* ctx, tstring& compiled_key) { |
| 159 | auto gpu_device = dynamic_cast<BaseGPUDevice*>(ctx->device()); |
| 160 | CudaGraphGPUBFCAllocator* device_allocator = |
| 161 | reinterpret_cast<CudaGraphGPUBFCAllocator*>( |
| 162 | gpu_device->GetAllocator(AllocatorAttributes())); |
| 163 | if (gpu_device == nullptr) { |
| 164 | return errors::Internal("BaseGPUDevice not found"); |
| 165 | } |
| 166 | se::Stream* default_stream = gpu_device->GetDefaultTFStream(); |
| 167 | cudaStream_t cu_stream; |
| 168 | cu_stream = *(static_cast<cudaStream_t*>(gpu_device->GetStream())); |
| 169 | |
| 170 | FunctionLibraryRuntime::Handle handle; |
| 171 | Status s_flib_instantiate = flib_->Instantiate( |
| 172 | function_.name(), AttrSlice(&function_.attr()), &handle); |
| 173 | if (!s_flib_instantiate.ok()) { |
| 174 | return s_flib_instantiate; |
| 175 | } |
| 176 | |
| 177 | const FunctionBody* fbody; |
| 178 | fbody = flib_->GetFunctionBody(handle); |
| 179 | VLOG(1) << "CgmodeCompileOp function def: " << DebugString(fbody->fdef); |
| 180 | |
| 181 | FunctionLibraryRuntime::Options opts; |
| 182 | std::vector<Tensor> in; |
| 183 | AllocatorAttributes attr; |
| 184 | attr.set_gpu_compatible(true); |
| 185 | for (int i = 0; i < ctx->num_inputs(); i++) { |
| 186 | // validate the original input is on device |
| 187 | PersistentTensor arg_persistent; |
| 188 | Tensor* out_tensor = nullptr; |
| 189 | OP_REQUIRES_OK_RETURN(ctx, errors::Internal("Failed to allocate persistent memory"), |
| 190 | ctx->allocate_persistent(ctx->input(i).dtype(), ctx->input(i).shape(), |
| 191 | &arg_persistent, &out_tensor, attr)); |
| 192 | persist_args_.emplace_back(arg_persistent); |
| 193 | cudaError_t e_cu_copy = cudaMemcpyAsync( |
| 194 | out_tensor->data(), ctx->input(i).data(), |
| 195 | ctx->input(i).TotalBytes(), cudaMemcpyDefault, cu_stream); |
| 196 | in.emplace_back(*out_tensor); |
| 197 | } |
| 198 | |
| 199 | Status s_dry_run; |
| 200 | Notification done_dry_run; |
| 201 | std::vector<Tensor> out_dry_run(fbody->ret_types.size()); |
| 202 | device_allocator->EnableCudaGraphModeMem(); |
| 203 | flib_->Run(opts, handle, in, &out_dry_run, |
| 204 | [&s_dry_run, &done_dry_run](const Status& s) { |
| 205 | s_dry_run = s; |
| 206 | done_dry_run.Notify(); |
| 207 | }); |
| 208 | done_dry_run.WaitForNotification(); |
| 209 | cudaStreamSynchronize(cu_stream); |
| 210 | if (!s_dry_run.ok()) { |
| 211 | return errors::Internal("CgmodeCompile failed" + s_dry_run.ToString()); |
| 212 | } |
| 213 | |
| 214 | device_allocator->DisableCudaGraphModeMem(); |
| 215 | Status s_compile; |
no test coverage detected