1 /* Copyright 2017 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 #include "tensorflow/compiler/xla/service/gpu/partition_assignment.h" 17 18 #include <ostream> 19 #include <string> 20 21 #include "tensorflow/compiler/xla/map_util.h" 22 #include "tensorflow/compiler/xla/ptr_util.h" 23 #include "tensorflow/compiler/xla/service/hlo_computation.h" 24 #include "tensorflow/compiler/xla/service/hlo_opcode.h" 25 #include "tensorflow/compiler/xla/shape_util.h" 26 #include "tensorflow/compiler/xla/types.h" 27 #include "tensorflow/compiler/xla/util.h" 28 #include "tensorflow/core/lib/core/bits.h" 29 #include "tensorflow/core/lib/strings/stringprintf.h" 30 #include "tensorflow/core/platform/logging.h" 31 32 namespace se = ::perftools::gputools; 33 34 namespace xla { 35 namespace gpu { 36 37 std::ostream& operator<<(std::ostream& out, 38 const LaunchDimensions& launch_dims) { 39 out << tensorflow::strings::Printf("[block: %lld, thread: %lld]", 40 launch_dims.block_count(), 41 launch_dims.threads_per_block()); 42 return out; 43 } 44 45 // Calculates the launch dimensions used to invoke `hlo`. 46 LaunchDimensions CalculateLaunchDimensions( 47 const Shape& shape, const se::DeviceDescription& device_desc) { 48 int64 num_elements = ShapeUtil::ElementsIn(shape); 49 if (num_elements <= 1) { 50 return LaunchDimensions(); 51 } 52 53 // Since we don't do any inter-warp communication, we're free to choose any 54 // block size we want, subject to hardware constraints. We choose the 55 // smallest block size that allows the GPU to reach full occupancy (assuming 56 // the kernel uses sufficiently few registers). This gives us max performance 57 // when the kernel uses few registers, and lets us scale down gracefully as 58 // the kernel uses more registers. 59 // 60 // Specifically, we choose the number of threads per block such that 61 // 62 // <num threads per block> * <max blocks per core> = <max threads per core> 63 64 auto threads_per_core = device_desc.threads_per_core_limit(); 65 auto blocks_per_core = device_desc.blocks_per_core_limit(); 66 int64 threads_per_block; 67 if (threads_per_core != 0 && blocks_per_core != 0) { 68 threads_per_block = device_desc.threads_per_core_limit() / 69 device_desc.blocks_per_core_limit(); 70 } else { 71 static std::atomic<int64> log_count{0}; 72 if (log_count.fetch_add(1) < 8) { 73 LOG(WARNING) << "Attempting to calculate launch dimensions for GPU " 74 "without full information about its capabilities. " 75 "StreamExecutor's PopulateDeviceDescription should be " 76 "updated for this device."; 77 } 78 threads_per_block = device_desc.threads_per_warp(); 79 if (threads_per_block == 0) { 80 // Fall back to *something* if we can't even get num threads per warp. 81 threads_per_block = 32; 82 } 83 } 84 85 if (num_elements < threads_per_block) { 86 threads_per_block = num_elements; 87 VLOG(2) << "Update # of threads per block to the element count (" 88 << threads_per_block << ") because the latter is smaller."; 89 } 90 91 int64 block_count = CeilOfRatio(num_elements, threads_per_block); 92 VLOG(2) << tensorflow::strings::Printf( 93 "Initialized the block count to ceil(# of elements / threads per " 94 "block) = ceil(%lld/%lld) = %lld", 95 num_elements, threads_per_block, block_count); 96 97 return LaunchDimensions(block_count, threads_per_block); 98 } 99 100 } // namespace gpu 101 } // namespace xla 102
