MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / CopyDataFromInput

Function CopyDataFromInput

tensorflow/compiler/xla/service/hlo_evaluator.cc:1118–1158  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1116//
1117template <typename InputType>
1118bool CopyDataFromInput(const Literal& input_literal, int64 input_start,
1119 int64 fft_rank, FftType fft_type, int64 fft_size,
1120 const absl::Span<const int64> fft_lengths,
1121 const absl::Span<const int64> fft_strides,
1122 const absl::Span<const int64> input_lengths,
1123 const absl::Span<const int64> input_strides,
1124 absl::Span<complex128> data) {
1125 CHECK_GE(data.size(), fft_size);
1126
1127 const bool input_is_truncated = fft_type == FftType::IRFFT;
1128
1129 // Recursively visit each transform dimension to copy input values to the
1130 // working data set. The base case handles inputs along the X axis.
1131 bool input_is_zero = true;
1132 const InputType* input_data = input_literal.data<InputType>().data();
1133 auto base_case = [&](int64 axis, int64 dst_index, int64 src_index,
1134 bool within_src_bounds) {
1135 if (axis == 0) {
1136 // For IRFFT, the negative frequencies are only needed for the sweep along
1137 // the X axis, which is performed last. Leave this part of the working set
1138 // uninitialized until then.
1139 const int64 length = fft_lengths[axis];
1140 const int64 ub = input_is_truncated ? (length / 2) + 1 : length;
1141 for (int64 i = 0; i < ub; i++) {
1142 complex128 value = InputType(0);
1143 // Read input value only if the index is within bounds.
1144 if (within_src_bounds && i < input_lengths[axis]) {
1145 value = GetAs<complex128, InputType>(
1146 input_data[src_index + i * input_strides[axis]]);
1147 input_is_zero &= value == complex128(0.0, 0.0);
1148 }
1149 data[dst_index + i * fft_strides[axis]] = value;
1150 }
1151 return true;
1152 }
1153 return false;
1154 };
1155 GenerateIndices(fft_lengths, fft_strides, input_lengths, input_strides,
1156 fft_rank, 0, input_start, base_case);
1157 return input_is_zero;
1158}
1159
1160// Copies the result of the transform to the literal output. The sizes of the
1161// transform and output must match.

Callers 1

HandleFftMethod · 0.85

Calls 3

GenerateIndicesFunction · 0.85
sizeMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected