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