Home
last modified time | relevance | path

Searched refs:BufferAllocation (Results 1 – 25 of 50) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcudnn_batchnorm_thunk.h49 CudnnBatchNormForwardInferenceThunk(const BufferAllocation::Slice& operand,
50 const BufferAllocation::Slice& scale,
51 const BufferAllocation::Slice& offset,
52 const BufferAllocation::Slice& mean,
53 const BufferAllocation::Slice& variance,
55 const BufferAllocation::Slice& output,
68 BufferAllocation::Slice operand_;
69 BufferAllocation::Slice scale_;
70 BufferAllocation::Slice offset_;
71 BufferAllocation::Slice mean_;
[all …]
Dcudnn_batchnorm_thunk.cc79 const BufferAllocation::Slice& operand, in CudnnBatchNormForwardInferenceThunk()
80 const BufferAllocation::Slice& scale, const BufferAllocation::Slice& offset, in CudnnBatchNormForwardInferenceThunk()
81 const BufferAllocation::Slice& mean, in CudnnBatchNormForwardInferenceThunk()
82 const BufferAllocation::Slice& variance, float epsilon, int64 feature_index, in CudnnBatchNormForwardInferenceThunk()
83 const BufferAllocation::Slice& output, const HloInstruction* hlo) in CudnnBatchNormForwardInferenceThunk()
136 const BufferAllocation::Slice& operand, in CudnnBatchNormForwardTrainingThunk()
137 const BufferAllocation::Slice& scale, const BufferAllocation::Slice& offset, in CudnnBatchNormForwardTrainingThunk()
139 const BufferAllocation::Slice& output_data, in CudnnBatchNormForwardTrainingThunk()
140 const BufferAllocation::Slice& output_mean, in CudnnBatchNormForwardTrainingThunk()
141 const BufferAllocation::Slice& output_inv_stddev, in CudnnBatchNormForwardTrainingThunk()
[all …]
Dconvolution_thunk.h47 std::vector<BufferAllocation::Slice> operand_slices,
48 BufferAllocation::Slice result_slice,
49 BufferAllocation::Slice scratch_slice,
50 BufferAllocation::Slice tuple_result_slice);
62 std::vector<BufferAllocation::Slice> operand_buffers_;
63 BufferAllocation::Slice result_buffer_;
64 BufferAllocation::Slice scratch_buffer_;
65 BufferAllocation::Slice tuple_result_buffer_;
Dbuffer_allocations.cc35 void BufferAllocations::Builder::RegisterBuffer(BufferAllocation::Index index, in RegisterBuffer()
47 for (BufferAllocation::Index i = 0; i < num_buffers; ++i) { in Build()
48 const BufferAllocation& allocation = buffer_assignment->GetAllocation(i); in Build()
109 for (BufferAllocation::Index i = 0; i < num_buffers; ++i) { in Build()
134 for (BufferAllocation::Index i = 0; i < num_buffers; ++i) { in TearDown()
135 const BufferAllocation& allocation = buffer_assignment_->GetAllocation(i); in TearDown()
154 BufferAllocation::Index buffer_index) const { in GetDeviceAddress()
161 const BufferAllocation::Slice& buffer_slice) const { in GetDeviceAddress()
170 void BufferAllocations::SetBuffer(BufferAllocation::Index buffer_index, in SetBuffer()
Dcopy_thunk.h37 const BufferAllocation::Slice& destination_buffer,
49 const BufferAllocation::Slice destination_buffer_;
59 DeviceToDeviceCopyThunk(const BufferAllocation::Slice& source_buffer,
60 const BufferAllocation::Slice& destination_buffer,
72 const BufferAllocation::Slice source_buffer_;
73 const BufferAllocation::Slice destination_buffer_;
Dcholesky_thunk.h45 BufferAllocation::Slice a_buffer,
46 BufferAllocation::Slice workspace_buffer,
47 BufferAllocation::Slice info_buffer,
61 const BufferAllocation::Slice a_buffer_;
62 const BufferAllocation::Slice workspace_buffer_;
63 const BufferAllocation::Slice info_buffer_;
Dgemm_thunk.h40 GemmThunk(const BufferAllocation::Slice& lhs_buffer,
41 const BufferAllocation::Slice& rhs_buffer,
42 const BufferAllocation::Slice& output_buffer,
65 const BufferAllocation::Slice lhs_buffer_;
66 const BufferAllocation::Slice rhs_buffer_;
67 const BufferAllocation::Slice output_buffer_;
Dbuffer_allocations.h44 void RegisterBuffer(BufferAllocation::Index index,
56 absl::flat_hash_map<BufferAllocation::Index, se::DeviceMemoryBase>
72 BufferAllocation::Index buffer_index) const;
77 const BufferAllocation::Slice& buffer_slice) const;
86 BufferAllocations(BufferAllocation::Index buffer_count, int device_ordinal, in BufferAllocations()
95 void SetBuffer(BufferAllocation::Index buffer_index,
Dnccl_all_reduce_thunk.h44 const BufferAllocation::Slice& source_buffer,
45 const BufferAllocation::Slice& destination_buffer,
55 const BufferAllocation::Slice source_buffer_;
56 const BufferAllocation::Slice destination_buffer_;
Dtuple_thunk.h37 TupleThunk(absl::Span<const BufferAllocation::Slice> tuple_element_buffers, in TupleThunk()
38 const BufferAllocation::Slice& dest_buffer, in TupleThunk()
53 const std::vector<BufferAllocation::Slice> tuple_element_buffers_;
54 const BufferAllocation::Slice dest_buffer_;
Dconditional_thunk.h46 const BufferAllocation::Slice& branch_index_buffer_index,
47 absl::Span<const BufferAllocation::Slice> branch_operand_buffer_indexes,
62 BufferAllocation::Slice branch_index_buffer_index_;
63 std::vector<BufferAllocation::Slice> branch_operand_buffer_indexes_;
Dmemset_thunk.h35 explicit MemzeroThunk(const BufferAllocation::Slice& dest, in MemzeroThunk()
44 const BufferAllocation::Slice dest_;
52 const BufferAllocation::Slice& dest, in Memset32BitValueThunk()
62 const BufferAllocation::Slice dest_;
Dtriangular_solve_thunk.h42 const BufferAllocation::Slice& a_buffer,
43 const BufferAllocation::Slice& b_buffer,
61 const BufferAllocation::Slice a_buffer_;
62 const BufferAllocation::Slice b_buffer_;
Dconvolution_thunk.cc34 std::vector<BufferAllocation::Slice> operand_slices, in ConvolutionThunk()
35 BufferAllocation::Slice result_slice, BufferAllocation::Slice scratch_slice, in ConvolutionThunk()
36 BufferAllocation::Slice tuple_result_slice) in ConvolutionThunk()
Dfft_thunk.h66 const BufferAllocation::Slice& input_buffer,
67 const BufferAllocation::Slice& output_buffer,
87 const BufferAllocation::Slice input_buffer_;
88 const BufferAllocation::Slice output_buffer_;
Dcopy_thunk.cc26 const BufferAllocation::Slice& destination_buffer, uint64 mem_size, in HostToDeviceCopyThunk()
44 const BufferAllocation::Slice& source_buffer, in DeviceToDeviceCopyThunk()
45 const BufferAllocation::Slice& destination_buffer, uint64 mem_size, in DeviceToDeviceCopyThunk()
Doutfeed_thunk.cc26 OutfeedThunk::OutfeedThunk(ShapeTree<BufferAllocation::Slice> outfeed_slices, in OutfeedThunk()
54 BufferAllocation::Slice slice = outfeed_slices_.element(index); in ExecuteOnStream()
64 BufferAllocation::Slice tuple_slice = in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/xla/service/
Dbuffer_assignment.h65 class BufferAllocation {
71 BufferAllocation(Index index, int64 size, LogicalBuffer::Color color) in BufferAllocation() function
73 ~BufferAllocation() {} in ~BufferAllocation()
165 Slice(const BufferAllocation* allocation, int64 offset, int64 size) in Slice()
168 const BufferAllocation* allocation() const { return allocation_; } in allocation()
201 const BufferAllocation* allocation_ = nullptr;
262 bool operator==(const BufferAllocation& other) const {
265 bool operator!=(const BufferAllocation& other) const {
268 bool operator<(const BufferAllocation& other) const {
346 std::ostream& operator<<(std::ostream& out, const BufferAllocation& s);
[all …]
Dbuffer_assignment.cc222 string BufferAllocation::Slice::ToString() const { in ToString()
227 BufferAllocation::Slice BufferAllocation::GetSlice( in GetSlice()
233 void BufferAllocation::AddAssignment(const LogicalBuffer& buffer, int64 offset, in AddAssignment()
253 BufferAllocationProto BufferAllocation::ToProto() const { in ToProto()
284 string BufferAllocation::ToString() const { in ToString()
326 std::ostream& operator<<(std::ostream& out, const BufferAllocation& buffer) { in operator <<()
331 std::ostream& operator<<(std::ostream& out, const BufferAllocation::Slice& s) { in operator <<()
346 const BufferAllocation& BufferAssignment::GetAssignedAllocation( in GetAssignedAllocation()
352 BufferAllocation* BufferAssignment::GetMutableAssignedAllocation( in GetMutableAssignedAllocation()
354 return const_cast<BufferAllocation*>(&GetAssignedAllocation(buffer)); in GetMutableAssignedAllocation()
[all …]
Dbuffer_assignment_test.cc114 const BufferAllocation& alloc, in RunBufferAssignmentNoBuffersReuseForAdd()
242 const BufferAllocation& GetAssignedInputAllocation( in GetAssignedInputAllocation()
245 const BufferAllocation& buffer = in GetAssignedInputAllocation()
253 const BufferAllocation& GetAssignedOutputAllocation( in GetAssignedOutputAllocation()
256 const BufferAllocation& buffer = GetTopLevelAllocation(buffers, hlo); in GetAssignedOutputAllocation()
261 const BufferAllocation& GetAllocation(const BufferAssignment& buffers, in GetAllocation()
266 const BufferAllocation& GetTopLevelAllocation(const BufferAssignment& buffers, in GetTopLevelAllocation()
314 absl::flat_hash_set<BufferAllocation::Slice> a_slices; in BuffersDistinct()
449 BufferAllocation paramscalar_buffer = in TEST_F()
451 BufferAllocation param0_buffer = GetAssignedInputAllocation(*buffers, param0); in TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/llvm_ir/
Dalias_analysis.cc30 static const BufferAllocation* kParameterAllocation = new BufferAllocation(
36 BufferAllocation::Slice buffer_slice; in AddAliasingInformationToIrArray()
41 buffer_slice = BufferAllocation::Slice(kParameterAllocation, 0, 0); in AddAliasingInformationToIrArray()
43 const std::set<BufferAllocation::Slice> slices = in AddAliasingInformationToIrArray()
109 const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain) { in GetAliasScopeMetadataForBuffer()
124 const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain, in GetNoaliasMetadataForBuffer()
165 std::set<BufferAllocation::Slice> buffers; in GetNoaliasMetadataForBuffer()
172 const BufferAllocation::Slice noalias_slice = in GetNoaliasMetadataForBuffer()
198 for (const BufferAllocation::Slice noalias_slice : buffers) { in GetNoaliasMetadataForBuffer()
Dalias_analysis.h49 const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain);
60 const BufferAllocation::Slice& buffer_slice, llvm::MDNode* domain,
79 absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*>
84 absl::flat_hash_map<BufferAllocation::Slice, llvm::MDNode*> noalias_metadata_;
Dbuffer_assignment_util.cc22 const BufferAllocation& allocation) { in InstrForConstantBufferAllocation()
58 const BufferAllocation& allocation) { in ConstantBufferAllocationToGlobalName()
68 const BufferAllocation& allocation) { in LiteralForConstantAllocation()
Dbuffer_assignment_util.h30 string ConstantBufferAllocationToGlobalName(const BufferAllocation& allocation);
34 const Literal& LiteralForConstantAllocation(const BufferAllocation& allocation);
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dir_emitter.h191 BufferAllocation::Slice GetAllocationSlice(
262 llvm::Value* EmitGlobalBufferPointer(const BufferAllocation::Slice& slice,
267 const BufferAllocation::Slice& slice, const Shape& target_shape);
270 llvm::Value* EmitBufferPointer(const BufferAllocation::Slice& slice,
434 std::map<std::pair<llvm::Function*, BufferAllocation::Slice>, llvm::Value*>
446 BufferAllocation::Slice computation_root_allocation_;
451 absl::flat_hash_map<BufferAllocation::Index, int64>
580 absl::flat_hash_map<BufferAllocation::Index, llvm::Constant*>

12