1 /* Copyright 2018 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_COMPILER_XLA_SERVICE_GPU_MEMSET_THUNK_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_GPU_MEMSET_THUNK_H_ 18 19 #include "tensorflow/compiler/xla/service/buffer_assignment.h" 20 #include "tensorflow/compiler/xla/service/gpu/hlo_execution_profiler.h" 21 #include "tensorflow/compiler/xla/service/gpu/thunk.h" 22 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 23 #include "tensorflow/compiler/xla/status.h" 24 #include "tensorflow/stream_executor/stream_executor.h" 25 26 // This file contains thunks that set a buffer's elements to a particular value. 27 // This can be faster than emitting a kernel to set the elements. 28 29 namespace xla { 30 namespace gpu { 31 32 // Thunk that zeroes out a given chunk of memory. 33 class MemzeroThunk : public Thunk { 34 public: MemzeroThunk(ThunkInfo thunk_info,const BufferAllocation::Slice & dest)35 explicit MemzeroThunk(ThunkInfo thunk_info, 36 const BufferAllocation::Slice& dest) 37 : Thunk(Kind::kMemzero, thunk_info), dest_(dest) {} 38 39 Status ExecuteOnStream(const ExecuteParams& params) override; 40 41 private: 42 const BufferAllocation::Slice dest_; 43 }; 44 45 // Thunk that sets a given chunk of memory to a particular 32-bit value. The 46 // destination chunk must have size divisible by 32 bits. 47 class Memset32BitValueThunk : public Thunk { 48 public: Memset32BitValueThunk(ThunkInfo thunk_info,uint32 value,const BufferAllocation::Slice & dest)49 explicit Memset32BitValueThunk(ThunkInfo thunk_info, uint32 value, 50 const BufferAllocation::Slice& dest) 51 : Thunk(Kind::kMemset32BitValue, thunk_info), 52 value_(value), 53 dest_(dest) {} 54 55 Status ExecuteOnStream(const ExecuteParams& params) override; 56 57 private: 58 const uint32 value_; 59 const BufferAllocation::Slice dest_; 60 }; 61 62 } // namespace gpu 63 } // namespace xla 64 65 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_GPU_MEMSET_THUNK_H_ 66