Method GetKernel

tensorflow/stream_executor/rocm/rocm_gpu_executor.cc:233–282

}

port::Status GpuExecutor::GetKernel(const MultiKernelLoaderSpec& spec,
                                    KernelBase* kernel) {
  GpuKernel* rocm_kernel = AsGpuKernel(kernel);
  hipModule_t module = nullptr;
  const string* kernelname;

  const OnDiskKernelLoaderSpec* on_disk_spec = nullptr;
  bool has_cubin = spec.has_cuda_cubin_on_disk();
  if (has_cubin) {
    on_disk_spec = &spec.cuda_cubin_on_disk();
  }

  if (on_disk_spec != nullptr) {
    return port::InternalError(
        "Loading ROCM kernel from disk is not supported");
  } else if (spec.has_cuda_cubin_in_memory()) {
    kernelname = &spec.cuda_cubin_in_memory().kernelname();

    const char* hsaco = spec.cuda_cubin_in_memory().bytes();
    absl::MutexLock lock{&in_memory_modules_mu_};
    module = in_memory_modules_[hsaco];

    if (module == nullptr) {
      if (!GpuDriver::LoadHsaco(context_, hsaco, &module)) {
        return port::InternalError("Failed to load HSACO");
      }
    }
    kernel_to_gpu_binary_[kernel] = hsaco;
  } else {
    return port::InternalError("No method of loading ROCM kernel provided");
  }

  VLOG(2) << "getting function " << *kernelname << " from module " << module;
  if (!GpuDriver::GetModuleFunction(context_, module, kernelname->c_str(),
                                    rocm_kernel->gpu_function_ptr())) {
    return port::InternalError("Failed getting module function");
  }

  // We have to trust the kernel loader spec arity because there doesn't appear
  // to be a way to reflect on the number of expected arguments w/the ROCM API.
  rocm_kernel->set_arity(spec.arity());

  KernelMetadata kernel_metadata;
  if (!GetKernelMetadata(rocm_kernel, &kernel_metadata)) {
    LOG(WARNING) << "Unable to get metadata for kernel " << kernelname;
  }
  kernel->set_metadata(kernel_metadata);
  kernel->set_name(*kernelname);
  return port::Status::OK();
}

bool GpuExecutor::GetKernelMetadata(GpuKernel* rocm_kernel,
                                    KernelMetadata* kernel_metadata) {