/external/tensorflow/tensorflow/compiler/xla/service/ |
D | convolution_group_converter.cc | 201 auto dim_numbers = convolution->convolution_dimension_numbers(); in HandleBatchGroupCount() local 217 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in HandleBatchGroupCount() 218 const int64 input_feature_dimension = dim_numbers.input_feature_dimension(); in HandleBatchGroupCount() 220 int64 output_batch_dimension = dim_numbers.output_batch_dimension(); in HandleBatchGroupCount() 221 int64 output_feature_dimension = dim_numbers.output_feature_dimension(); in HandleBatchGroupCount() 224 dim_numbers.kernel_input_feature_dimension(); in HandleBatchGroupCount() 226 dim_numbers.kernel_output_feature_dimension(); in HandleBatchGroupCount() 242 for (auto& d : *dim_numbers.mutable_input_spatial_dimensions()) { in HandleBatchGroupCount() 247 dim_numbers.add_input_spatial_dimensions(input_batch_dimension); in HandleBatchGroupCount() 248 dim_numbers.set_input_batch_dimension(input_batch_dimension + 1); in HandleBatchGroupCount() [all …]
|
D | gather_expander.cc | 111 HloInstruction* index_vector, const GatherDimensionNumbers& dim_numbers, in ExpandIndexVectorIntoOperandSpace() argument 132 int64 index_vector_dim_index = FindIndex(dim_numbers.start_index_map(), i); in ExpandIndexVectorIntoOperandSpace() 133 if (index_vector_dim_index != dim_numbers.start_index_map_size()) { in ExpandIndexVectorIntoOperandSpace() 153 const GatherDimensionNumbers& dim_numbers = gather.gather_dimension_numbers(); in GatherLoopBody() local 161 dim_numbers.index_vector_dim() == in GatherLoopBody() 197 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, in GatherLoopBody() 207 AsInt64Slice(dim_numbers.collapsed_slice_dims()))); in GatherLoopBody() 235 const GatherDimensionNumbers& dim_numbers) { in CreateGatherLoopAccumulatorInitValue() argument 240 if (!absl::c_binary_search(dim_numbers.collapsed_slice_dims(), i)) { in CreateGatherLoopAccumulatorInitValue() 276 const GatherDimensionNumbers& dim_numbers = in GatherLoopTripCount() local [all …]
|
D | scatter_expander.cc | 133 HloInstruction* index_vector, const ScatterDimensionNumbers& dim_numbers, in ExpandIndexVectorIntoOperandSpace() argument 154 FindIndex(dim_numbers.scatter_dims_to_operand_dims(), i); in ExpandIndexVectorIntoOperandSpace() 156 dim_numbers.scatter_dims_to_operand_dims_size()) { in ExpandIndexVectorIntoOperandSpace() 222 const ScatterDimensionNumbers& dim_numbers = in ScatterLoopBody() local 258 ExpandIndexVectorIntoOperandSpace(index_vector, dim_numbers, in ScatterLoopBody() 279 AsInt64Slice(dim_numbers.inserted_window_dims()))); in ScatterLoopBody() 333 const ScatterDimensionNumbers& dim_numbers = in ScatterTripCount() local 337 if (i != dim_numbers.index_vector_dim()) { in ScatterTripCount() 366 const ScatterDimensionNumbers& dim_numbers = in ExpandInstruction() local 389 scatter_indices, dim_numbers.index_vector_dim())); in ExpandInstruction() [all …]
|
D | indexed_array_analysis.cc | 254 const Shape& shape, const GatherDimensionNumbers& dim_numbers, in ComputeArrayForGather() argument 256 if (dim_numbers.index_vector_dim() != indices->shape().dimensions_size()) { in ComputeArrayForGather() 261 CHECK_EQ(dim_numbers.start_index_map_size(), 1); in ComputeArrayForGather() 266 if (dim_numbers.collapsed_slice_dims_size() != 1 || in ComputeArrayForGather() 267 dim_numbers.collapsed_slice_dims(0) != dim_numbers.start_index_map(0)) { in ComputeArrayForGather() 279 if (i != dim_numbers.collapsed_slice_dims(0) && in ComputeArrayForGather() 285 << dim_numbers.collapsed_slice_dims(0); in ComputeArrayForGather() 290 int64 source_dim = dim_numbers.start_index_map(0); in ComputeArrayForGather() 293 if (!absl::c_binary_search(dim_numbers.offset_dims(), i)) { in ComputeArrayForGather() 1033 const Shape& shape, const DotDimensionNumbers& dim_numbers, in ComputeArrayForDotWithIndexedLhs() argument [all …]
|
D | batch_dot_simplification.cc | 44 const DotDimensionNumbers& dim_numbers = batch_dot->dot_dimension_numbers(); in ElideDegenerateBatchDimensionFromBatchDot() local 52 if (dim_numbers.lhs_contracting_dimensions_size() != 1) { in ElideDegenerateBatchDimensionFromBatchDot() 57 for (int64 batch_dim : dim_numbers.lhs_batch_dimensions()) { in ElideDegenerateBatchDimensionFromBatchDot() 72 DotDimensionNumbers new_dim_numbers = dim_numbers; in ElideDegenerateBatchDimensionFromBatchDot() 76 for (int64 i = 0, e = dim_numbers.lhs_batch_dimensions_size() - in ElideDegenerateBatchDimensionFromBatchDot()
|
D | shape_inference.cc | 3213 const GatherDimensionNumbers& dim_numbers) { in ValidateGatherDimensionNumbers() argument 3214 if (!absl::c_is_sorted(dim_numbers.offset_dims())) { in ValidateGatherDimensionNumbers() 3217 StrJoin(dim_numbers.offset_dims(), ", ")); in ValidateGatherDimensionNumbers() 3220 if (absl::c_adjacent_find(dim_numbers.offset_dims()) != in ValidateGatherDimensionNumbers() 3221 dim_numbers.offset_dims().end()) { in ValidateGatherDimensionNumbers() 3224 StrJoin(dim_numbers.offset_dims(), ", ")); in ValidateGatherDimensionNumbers() 3227 const int64 output_offset_dim_count = dim_numbers.offset_dims_size(); in ValidateGatherDimensionNumbers() 3231 for (int i = 0; i < dim_numbers.offset_dims_size(); ++i) { in ValidateGatherDimensionNumbers() 3232 int64 offset_dim = dim_numbers.offset_dims(i); in ValidateGatherDimensionNumbers() 3242 if (dim_numbers.start_index_map_size() != in ValidateGatherDimensionNumbers() [all …]
|
D | space_to_batch_converter.cc | 80 ConvolutionDimensionNumbers& dim_numbers); 117 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers, 162 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers, 278 ConvolutionDimensionNumbers dim_numbers = in IsConvSuitableForSpaceToBatch() local 282 if (dim_numbers.input_spatial_dimensions_size() < 1) { in IsConvSuitableForSpaceToBatch() 298 const ConvDetails c = GetConvolutionDetails(convolution, dim_numbers); in IsConvSuitableForSpaceToBatch() 323 int64 activations_batch_dim = dim_numbers.input_batch_dimension(); in IsConvSuitableForSpaceToBatch() 479 HloInstruction* activations, ConvolutionDimensionNumbers& dim_numbers, in BringSpaceNextToBatch() argument 486 ConvolutionDimensionNumbers new_dim_numbers = dim_numbers; in BringSpaceNextToBatch() 503 if (i == dim_numbers.kernel_output_feature_dimension()) { in BringSpaceNextToBatch() [all …]
|
D | triangular_solve_expander.cc | 77 GatherDimensionNumbers dim_numbers; in DiagonalBlocks() local 79 dim_numbers.add_offset_dims(i); in DiagonalBlocks() 80 dim_numbers.add_start_index_map(i); in DiagonalBlocks() 84 dim_numbers.add_offset_dims(ndims - 1); in DiagonalBlocks() 85 dim_numbers.add_offset_dims(ndims); in DiagonalBlocks() 86 dim_numbers.add_start_index_map(ndims - 2); in DiagonalBlocks() 87 dim_numbers.add_start_index_map(ndims - 1); in DiagonalBlocks() 88 dim_numbers.set_index_vector_dim(1); in DiagonalBlocks() 89 diag_blocks = Gather(a, start_indices, dim_numbers, slice_sizes); in DiagonalBlocks()
|
D | indexed_array_analysis.h | 265 const Shape& shape, const GatherDimensionNumbers& dim_numbers, 269 const Shape& shape, const DotDimensionNumbers& dim_numbers, 274 const Shape& shape, const DotDimensionNumbers& dim_numbers, 279 const DotDimensionNumbers& dim_numbers,
|
D | hlo_cost_analysis_test.cc | 977 GatherDimensionNumbers dim_numbers; in TEST_F() local 978 dim_numbers.add_offset_dims(1); in TEST_F() 979 dim_numbers.add_collapsed_slice_dims(0); in TEST_F() 980 dim_numbers.add_start_index_map(0); in TEST_F() 981 dim_numbers.set_index_vector_dim(1); in TEST_F() 982 Gather(operand, indices, dim_numbers, {1, 3}); in TEST_F() 1009 ScatterDimensionNumbers dim_numbers; in TEST_F() local 1010 dim_numbers.set_index_vector_dim(1); in TEST_F() 1011 dim_numbers.add_update_window_dims(1); in TEST_F() 1012 dim_numbers.add_inserted_window_dims(0); in TEST_F() [all …]
|
D | hlo_evaluator.cc | 383 const DotDimensionNumbers& dim_numbers, in EvaluateDotOp() argument 393 lhs.shape(), rhs.shape(), dim_numbers, in EvaluateDotOp() 398 dim_numbers, precision_config); in EvaluateDotOp() 1451 const Shape& output_shape, const GatherDimensionNumbers& dim_numbers) { in IterationSpaceForOutputBatchIndices() argument 1458 !absl::c_binary_search(dim_numbers.offset_dims(), i); in IterationSpaceForOutputBatchIndices() 1470 const GatherDimensionNumbers& dim_numbers) { in IterationSpaceForOutputOffsetIndices() argument 1476 absl::c_binary_search(dim_numbers.offset_dims(), i); in IterationSpaceForOutputOffsetIndices() 1478 while (absl::c_binary_search(dim_numbers.collapsed_slice_dims(), in IterationSpaceForOutputOffsetIndices() 1500 const GatherDimensionNumbers* dim_numbers, const Shape& input_shape, in OutputBatchIndexToInputIndex() argument 1502 : dim_numbers_(*dim_numbers), start_indices_(*start_indices) { in OutputBatchIndexToInputIndex() [all …]
|
D | hlo_evaluator_typed_visitor.h | 2082 const Shape& updates_shape, const ScatterDimensionNumbers& dim_numbers) { 2088 !absl::c_binary_search(dim_numbers.update_window_dims(), i); 2101 const Shape& updates_shape, const ScatterDimensionNumbers& dim_numbers) { 2107 absl::c_binary_search(dim_numbers.update_window_dims(), i); 2129 const ScatterDimensionNumbers* dim_numbers, const Shape& input_shape, 2131 : dim_numbers_(*dim_numbers), scatter_indices_(*scatter_indices) { 2259 const ScatterDimensionNumbers& dim_numbers, const Shape& input_shape, 2264 if (absl::c_binary_search(dim_numbers.update_window_dims(), i)) { 2273 if (absl::c_binary_search(dim_numbers.inserted_window_dims(), i)) { 2331 const ScatterDimensionNumbers& dim_numbers = [all …]
|
D | elemental_ir_emitter.cc | 1914 const GatherDimensionNumbers& dim_numbers = hlo->gather_dimension_numbers(); in EmitElementalGather() local 1932 if (absl::c_binary_search(dim_numbers.collapsed_slice_dims(), i)) { in EmitElementalGather() 1935 int64 output_window_dim = dim_numbers.offset_dims(operand_index_dim++); in EmitElementalGather() 1945 if (!absl::c_binary_search(dim_numbers.offset_dims(), i)) { in EmitElementalGather() 1954 dim_numbers.index_vector_dim(), in EmitElementalGather() 1971 int64 operand_dim = dim_numbers.start_index_map(dim); in EmitElementalGather() 2000 if (indices_shape.dimensions_size() == dim_numbers.index_vector_dim()) { in EmitElementalGather() 2008 indices_shape.dimensions(dim_numbers.index_vector_dim()); in EmitElementalGather() 2010 gather_index_index_components[dim_numbers.index_vector_dim()] = in EmitElementalGather() 2177 const DotDimensionNumbers& dim_numbers = hlo->dot_dimension_numbers(); in EmitElementalDot() local [all …]
|
D | hlo_creation_utils.cc | 288 const DotDimensionNumbers& dim_numbers, in MakeDotHlo() argument 295 ShapeInference::InferDotOpShape(lhs->shape(), rhs->shape(), dim_numbers, in MakeDotHlo() 298 dot_shape, lhs, rhs, dim_numbers, precision_config)); in MakeDotHlo()
|
D | hlo_creation_utils.h | 139 const DotDimensionNumbers& dim_numbers,
|
D | hlo_evaluator.h | 127 StatusOr<Literal> EvaluateDotOp(const DotDimensionNumbers& dim_numbers,
|
/external/tensorflow/tensorflow/compiler/tf2xla/lib/ |
D | scatter.cc | 138 xla::ScatterDimensionNumbers dim_numbers; in XlaScatter() local 139 dim_numbers.set_index_vector_dim(indices_are_vectors in XlaScatter() 165 dim_numbers.add_update_window_dims(i); in XlaScatter() 170 dim_numbers.add_inserted_window_dims(i); in XlaScatter() 171 dim_numbers.add_scatter_dims_to_operand_dims(i); in XlaScatter() 193 VLOG(3) << " index_vector_dim: " << dim_numbers.index_vector_dim(); in XlaScatter() 195 << absl::StrJoin(dim_numbers.update_window_dims(), ",") << "]"; in XlaScatter() 197 << absl::StrJoin(dim_numbers.inserted_window_dims(), ",") << "]"; in XlaScatter() 199 << absl::StrJoin(dim_numbers.scatter_dims_to_operand_dims(), ",") in XlaScatter() 203 dim_numbers); in XlaScatter()
|
/external/tensorflow/tensorflow/compiler/xla/tests/ |
D | dynamism_inference_test.cc | 281 GatherDimensionNumbers dim_numbers; in TEST_F() local 282 dim_numbers.add_offset_dims(1); in TEST_F() 283 dim_numbers.add_start_index_map(0); in TEST_F() 284 dim_numbers.set_index_vector_dim(1); in TEST_F() 285 auto gather = Gather(operand1, indices, dim_numbers, {1}); in TEST_F() 301 GatherDimensionNumbers dim_numbers; in TEST_F() local 302 dim_numbers.add_offset_dims(1); in TEST_F() 303 dim_numbers.add_start_index_map(0); in TEST_F() 304 dim_numbers.set_index_vector_dim(1); in TEST_F() 305 auto gather = Gather(data_operand, indices, dim_numbers, {1}); in TEST_F() [all …]
|
D | gather_operation_test.cc | 757 GatherDimensionNumbers dim_numbers; in XLA_TEST_F() local 758 dim_numbers.add_offset_dims(1); in XLA_TEST_F() 759 dim_numbers.add_collapsed_slice_dims(0); in XLA_TEST_F() 760 dim_numbers.add_start_index_map(0); in XLA_TEST_F() 761 dim_numbers.set_index_vector_dim(1); in XLA_TEST_F() 762 Gather(operand, indices, dim_numbers, {1, 3}); in XLA_TEST_F()
|
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
D | gather_op.cc | 122 xla::GatherDimensionNumbers dim_numbers; in XlaGather() local 128 dim_numbers.add_collapsed_slice_dims(i); in XlaGather() 137 dim_numbers.add_offset_dims(i); in XlaGather() 141 dim_numbers.add_offset_dims(i + indices_rank - num_index_dims); in XlaGather() 145 dim_numbers.set_index_vector_dim(indices_are_nd ? (indices_shape.dims() - 1) in XlaGather() 148 dim_numbers.add_start_index_map(i); in XlaGather() 151 *gather_output = xla::Gather(input, indices, dim_numbers, slice_sizes); in XlaGather()
|
/external/tensorflow/tensorflow/compiler/xla/client/lib/ |
D | matrix.cc | 149 xla::GatherDimensionNumbers dim_numbers; in GetMatrixDiagonalViaGather() local 155 dim_numbers.add_collapsed_slice_dims(i); in GetMatrixDiagonalViaGather() 156 dim_numbers.add_start_index_map(i); in GetMatrixDiagonalViaGather() 159 dim_numbers.add_offset_dims(i); in GetMatrixDiagonalViaGather() 165 dim_numbers.set_index_vector_dim(1); in GetMatrixDiagonalViaGather() 167 return Gather(x, start_indices, dim_numbers, slice_sizes, in GetMatrixDiagonalViaGather()
|
/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
D | gpu_conv_rewriter.cc | 72 auto dim_numbers = conv->convolution_dimension_numbers(); in ConvertBatchGroupedToFeatureGroupedConvolution() local 76 int64 input_batch_dimension = dim_numbers.input_batch_dimension(); in ConvertBatchGroupedToFeatureGroupedConvolution() 79 int64 input_feature_dimension = dim_numbers.input_feature_dimension(); in ConvertBatchGroupedToFeatureGroupedConvolution() 117 new_conv->set_convolution_dimension_numbers(dim_numbers); in ConvertBatchGroupedToFeatureGroupedConvolution()
|
D | ir_emission_utils.cc | 112 const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers(); in IsMatrixMultiplication() local 117 dim_numbers.lhs_batch_dimensions_size())) { in IsMatrixMultiplication() 121 CHECK_EQ(lhs_shape.dimensions(dim_numbers.lhs_contracting_dimensions(0)), in IsMatrixMultiplication() 122 rhs_shape.dimensions(dim_numbers.rhs_contracting_dimensions(0))); in IsMatrixMultiplication()
|
/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
D | dot_op_emitter.cc | 1092 Status ValidateDotDimensionNumbers(const DotDimensionNumbers& dim_numbers) { in ValidateDotDimensionNumbers() argument 1095 TF_RET_CHECK(dim_numbers.lhs_contracting_dimensions_size() == 1); in ValidateDotDimensionNumbers() 1096 std::vector<int64> batch_dim_numbers(dim_numbers.lhs_batch_dimensions_size()); in ValidateDotDimensionNumbers() 1099 absl::c_equal(batch_dim_numbers, dim_numbers.lhs_batch_dimensions())); in ValidateDotDimensionNumbers() 1101 absl::c_equal(batch_dim_numbers, dim_numbers.rhs_batch_dimensions())); in ValidateDotDimensionNumbers()
|
/external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
D | legalize_to_linalg.cc | 1201 mhlo::DotDimensionNumbers dim_numbers = op.dot_dimension_numbers(); in matchAndRewrite() local 1203 Extract1DVector(dim_numbers.lhs_batching_dimensions()); in matchAndRewrite() 1205 Extract1DVector(dim_numbers.rhs_batching_dimensions()); in matchAndRewrite() 1207 Extract1DVector(dim_numbers.lhs_contracting_dimensions()); in matchAndRewrite() 1209 Extract1DVector(dim_numbers.rhs_contracting_dimensions()); in matchAndRewrite()
|