1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #ifndef TENSORFLOW_STREAM_EXECUTOR_GPU_ASM_COMPILER_H_
17 #define TENSORFLOW_STREAM_EXECUTOR_GPU_ASM_COMPILER_H_
18
19 #include <vector>
20
21 #include "absl/container/flat_hash_map.h"
22 #include "absl/types/span.h"
23 #include "tensorflow/core/platform/mutex.h"
24 #include "tensorflow/stream_executor/gpu/gpu_asm_opts.h"
25 #include "tensorflow/stream_executor/kernel.h"
26 #include "tensorflow/stream_executor/lib/statusor.h"
27 #include "tensorflow/stream_executor/platform/port.h"
28 #include "tensorflow/stream_executor/stream_executor_pimpl.h"
29 #if GOOGLE_CUDA
30 #include "tensorflow/stream_executor/cuda/cuda_driver.h"
31 #endif // GOOGLE_CUDA
32
33 namespace stream_executor {
34 namespace gpu {
35 class GpuContext;
36 }
37
38 // Compiles the given PTX string using ptxas and returns the resulting machine
39 // code (i.e. a cubin) as a byte array. The generated cubin matches the compute
40 // capabilities of the device associated with 'device_ordinal'.
41 //
42 // 'options' is used to query for the CUDA location in case it is
43 // customized in a passed flag, and for controlling ptxas optimizations.
44 port::StatusOr<std::vector<uint8>> CompileGpuAsm(int device_ordinal,
45 const char* ptx_contents,
46 GpuAsmOpts options);
47
48 // Compiles the given PTX string using ptxas and returns the resulting machine
49 // code (i.e. a cubin) as a byte array. The generated cubin matches the compute
50 // capabilities provided by 'cc_major' and 'cc_minor'.
51 //
52 // 'options' is used to query for the CUDA location in case it is
53 // customized in a passed flag, and for controlling ptxas optimizations.
54 port::StatusOr<std::vector<uint8>> CompileGpuAsm(int cc_major, int cc_minor,
55 const char* ptx_contents,
56 GpuAsmOpts options);
57
58 // Same as CompileGpuAsm, but caches the result, and returns unowned view of
59 // the compiled binary.
60 //
61 // A copy of the string provided in ptx will be made.
62 port::StatusOr<absl::Span<const uint8>> CompileGpuAsmOrGetCached(
63 int device_ordinal, const char* ptx, GpuAsmOpts compilation_options);
64
65 struct CubinOrPTXImage {
66 std::string profile;
67 std::vector<uint8> bytes;
68 };
69
70 // Bundles the GPU machine code (cubins) and PTX if requested and returns the
71 // resulting binary (i.e. a fatbin) as a byte array.
72 port::StatusOr<std::vector<uint8>> BundleGpuAsm(
73 std::vector<CubinOrPTXImage> images, GpuAsmOpts options);
74
75 struct HsacoImage {
76 std::string gfx_arch;
77 std::vector<uint8> bytes;
78 };
79
80 // Bundles the GPU machine code (HSA Code Object) and returns the resulting
81 // binary (i.e. a fatbin) as a byte array.
82 port::StatusOr<std::vector<uint8>> BundleGpuAsm(
83 std::vector<HsacoImage> images, const std::string rocm_root_dir);
84
85 // Links multiple relocatable GPU images (e.g. results of ptxas -c) into a
86 // single image.
87 port::StatusOr<std::vector<uint8>> LinkGpuAsm(
88 gpu::GpuContext* context, std::vector<CubinOrPTXImage> images);
89
90 #if GOOGLE_CUDA
91 // Maintains a cache of pointers to loaded kernels
92 template <typename... Args>
LoadKernelOrGetPtr(StreamExecutor * executor,absl::string_view kernel_name,absl::string_view ptx,absl::Span<const uint8> cubin_data)93 port::StatusOr<std::shared_ptr<TypedKernel<Args...>>> LoadKernelOrGetPtr(
94 StreamExecutor* executor, absl::string_view kernel_name,
95 absl::string_view ptx, absl::Span<const uint8> cubin_data) {
96 using KernelPtrCacheKey =
97 std::tuple<CUcontext, absl::string_view, absl::string_view>;
98
99 static tensorflow::mutex kernel_ptr_cache_mutex(
100 tensorflow::LINKER_INITIALIZED);
101 static auto& kernel_ptr_cache TF_GUARDED_BY(kernel_ptr_cache_mutex) =
102 *new absl::flat_hash_map<KernelPtrCacheKey,
103 std::shared_ptr<TypedKernel<Args...>>>();
104 CUcontext current_context = cuda::CurrentContextOrDie();
105 KernelPtrCacheKey kernel_ptr_cache_key{current_context, kernel_name, ptx};
106 tensorflow::mutex_lock lock(kernel_ptr_cache_mutex);
107
108 auto it = kernel_ptr_cache.find(kernel_ptr_cache_key);
109 if (it == kernel_ptr_cache.end()) {
110 TF_ASSIGN_OR_RETURN(
111 std::shared_ptr<TypedKernel<Args...>> loaded,
112 executor->CreateTypedKernel<Args...>(kernel_name, ptx, cubin_data));
113 it =
114 kernel_ptr_cache.emplace(kernel_ptr_cache_key, std::move(loaded)).first;
115 }
116
117 CHECK(it != kernel_ptr_cache.end());
118 return it->second;
119 }
120 #endif // GOOGLE_CUDA
121
122 } // namespace stream_executor
123
124 #endif // TENSORFLOW_STREAM_EXECUTOR_GPU_ASM_COMPILER_H_
125