| /external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
| D | replica_id_thunk.h | 31 ReplicaOrPartitionIdThunk(Kind kind, ThunkInfo thunk_info, in ReplicaOrPartitionIdThunk() argument
33 : Thunk(kind, thunk_info), dest_(dest) {} in ReplicaOrPartitionIdThunk()
41 ReplicaIdThunk(ThunkInfo thunk_info, const BufferAllocation::Slice& dest) in ReplicaIdThunk() argument
42 : ReplicaOrPartitionIdThunk(Kind::kReplicaId, thunk_info, dest) {} in ReplicaIdThunk()
47 PartitionIdThunk(ThunkInfo thunk_info, const BufferAllocation::Slice& dest) in PartitionIdThunk() argument
48 : ReplicaOrPartitionIdThunk(Kind::kPartitionId, thunk_info, dest) {} in PartitionIdThunk()
|
| D | memset_thunk.h | 35 explicit MemzeroThunk(ThunkInfo thunk_info, in MemzeroThunk() argument
37 : Thunk(Kind::kMemzero, thunk_info), dest_(dest) {} in MemzeroThunk()
49 explicit Memset32BitValueThunk(ThunkInfo thunk_info, uint32 value, in Memset32BitValueThunk() argument
51 : Thunk(Kind::kMemset32BitValue, thunk_info), in Memset32BitValueThunk()
|
| D | copy_thunk.cc | 25 ThunkInfo thunk_info, const void* source_address, in HostToDeviceCopyThunk() argument
27 : Thunk(Kind::kCopy, thunk_info), in HostToDeviceCopyThunk()
42 ThunkInfo thunk_info, const BufferAllocation::Slice& source_buffer, in DeviceToDeviceCopyThunk() argument
44 : Thunk(Kind::kCopy, thunk_info), in DeviceToDeviceCopyThunk()
|
| D | cudnn_batchnorm_thunk.cc | 35 ThunkInfo thunk_info, CudnnBatchNormConfig config, in CudnnBatchNormForwardInferenceThunk() argument
41 : Thunk(Thunk::Kind::kCudnnBatchNormForwardInference, thunk_info), in CudnnBatchNormForwardInferenceThunk()
74 ThunkInfo thunk_info, CudnnBatchNormConfig config, in CudnnBatchNormForwardTrainingThunk() argument
80 : Thunk(Thunk::Kind::kCudnnBatchNormForwardTraining, thunk_info), in CudnnBatchNormForwardTrainingThunk()
118 ThunkInfo thunk_info, CudnnBatchNormConfig config, in CudnnBatchNormBackwardThunk() argument
126 : Thunk(Thunk::Kind::kCudnnBatchNormBackward, thunk_info), in CudnnBatchNormBackwardThunk()
|
| D | thunk.h | 82 explicit Thunk(Kind kind, ThunkInfo thunk_info) in Thunk() argument
84 profile_index_(thunk_info.profile_index), in Thunk()
85 profile_annotation_(thunk_info.profile_annotation) {} in Thunk()
|
| D | cudnn_batchnorm_thunk.h | 50 CudnnBatchNormForwardInferenceThunk(ThunkInfo thunk_info,
79 ThunkInfo thunk_info, CudnnBatchNormConfig config,
107 ThunkInfo thunk_info, CudnnBatchNormConfig config,
|
| D | sequential_thunk.cc | 27 SequentialThunk::SequentialThunk(ThunkInfo thunk_info, in SequentialThunk() argument
29 : Thunk(Kind::kSequential, thunk_info), thunks_(std::move(thunks)) {} in SequentialThunk()
|
| D | copy_thunk.h | 36 HostToDeviceCopyThunk(ThunkInfo thunk_info, const void* source_address,
57 DeviceToDeviceCopyThunk(ThunkInfo thunk_info,
|
| D | for_thunk.cc | 26 ForThunk::ForThunk(ThunkInfo thunk_info, const int64 loop_limit, in ForThunk() argument
29 : Thunk(Kind::kWhile, thunk_info), in ForThunk()
|
| D | convolution_thunk.cc | 34 ThunkInfo thunk_info, GpuConvConfig config, in ConvolutionThunk() argument
37 : Thunk(Kind::kConvolution, thunk_info), in ConvolutionThunk()
|
| D | custom_call_thunk.cc | 30 CustomCallThunk::CustomCallThunk(ThunkInfo thunk_info, void* call_target, in CustomCallThunk() argument
34 : Thunk(Thunk::kCustomCall, thunk_info), in CustomCallThunk()
|
| D | infeed_thunk.cc | 29 InfeedThunk::InfeedThunk(ThunkInfo thunk_info, in InfeedThunk() argument
31 : Thunk(Kind::kInfeed, thunk_info), dest_slices_(std::move(dest_slices)) {} in InfeedThunk()
|
| D | while_thunk.cc | 27 ThunkInfo thunk_info, in WhileThunk() argument
33 : Thunk(Kind::kWhile, thunk_info), in WhileThunk()
|
| D | outfeed_thunk.cc | 28 OutfeedThunk::OutfeedThunk(ThunkInfo thunk_info, in OutfeedThunk() argument
30 : Thunk(Kind::kOutfeed, thunk_info), in OutfeedThunk()
|
| D | conditional_thunk.cc | 48 ThunkInfo thunk_info, ConditionalThunkConfig config, in ConditionalThunk() argument
51 : Thunk(Kind::kConditional, thunk_info), in ConditionalThunk()
|
| D | nccl_all_gather_thunk.cc | 62 ThunkInfo thunk_info, mlir::lmhlo::AllGatherOp op, int64 replica_count, in NcclAllGatherThunk() argument
64 : NcclCollectiveThunk(Thunk::kNcclAllGather, thunk_info), in NcclAllGatherThunk()
|
| D | cholesky_thunk.cc | 34 CholeskyThunk::CholeskyThunk(ThunkInfo thunk_info, in CholeskyThunk() argument
40 : Thunk(Kind::kCholesky, thunk_info), in CholeskyThunk()
|
| D | kernel_thunk.cc | 36 KernelThunk::KernelThunk(ThunkInfo thunk_info, in KernelThunk() argument
39 : Thunk(Kind::kKernel, thunk_info), in KernelThunk()
|
| D | nccl_all_to_all_thunk.cc | 54 ThunkInfo thunk_info, mlir::lmhlo::AllToAllOp op, int64 replica_count, in NcclAllToAllThunk() argument
56 : NcclCollectiveThunk(Thunk::kNcclAllToAll, thunk_info), in NcclAllToAllThunk()
|
| D | triangular_solve_thunk.cc | 35 ThunkInfo thunk_info, const TriangularSolveOptions& options, in TriangularSolveThunk() argument
40 : Thunk(Kind::kTriangularSolve, thunk_info), in TriangularSolveThunk()
|
| D | nccl_all_reduce_thunk.cc | 111 ThunkInfo thunk_info, mlir::lmhlo::AllReduceOp op, int64 replica_count, in NcclAllReduceThunk() argument
113 : NcclCollectiveThunk(Thunk::kNcclAllReduce, thunk_info), in NcclAllReduceThunk()
|
| D | ir_emitter_unnested.h | 62 Thunk::ThunkInfo thunk_info; member
642 absl::string_view name, Thunk::ThunkInfo thunk_info,
656 absl::string_view name, Thunk::ThunkInfo thunk_info,
662 Thunk::ThunkInfo thunk_info, absl::optional<MlirBufferSlice> extra_slice,
666 mlir::Operation* op, Thunk::ThunkInfo thunk_info,
|
| D | infeed_thunk.h | 37 InfeedThunk(ThunkInfo thunk_info, std::vector<ShapedSlice> dest_slices);
|
| D | outfeed_thunk.h | 35 OutfeedThunk(ThunkInfo thunk_info, std::vector<ShapedSlice> source_slices);
|
| D | collective_permute_thunk.h | 31 ThunkInfo thunk_info,
|