Home
last modified time | relevance | path

Searched refs:dnums (Results 1 – 25 of 58) sorted by relevance

123

/external/tensorflow/tensorflow/compiler/xla/tests/
Dconvolution_test.cc254 ConvolutionDimensionNumbers dnums; in XLA_TEST_F() local
255 dnums.set_input_batch_dimension(0); in XLA_TEST_F()
256 dnums.set_output_batch_dimension(0); in XLA_TEST_F()
257 dnums.add_input_spatial_dimensions(1); in XLA_TEST_F()
258 dnums.add_output_spatial_dimensions(1); in XLA_TEST_F()
259 dnums.add_input_spatial_dimensions(2); in XLA_TEST_F()
260 dnums.add_output_spatial_dimensions(2); in XLA_TEST_F()
261 dnums.add_input_spatial_dimensions(3); in XLA_TEST_F()
262 dnums.add_output_spatial_dimensions(3); in XLA_TEST_F()
263 dnums.set_input_feature_dimension(4); in XLA_TEST_F()
[all …]
Dconvolution_variants_test.cc980 ConvolutionDimensionNumbers dnums; in XLA_TEST_F() local
982 dnums.set_input_batch_dimension(0); in XLA_TEST_F()
983 dnums.set_output_batch_dimension(0); in XLA_TEST_F()
984 dnums.add_input_spatial_dimensions(1); in XLA_TEST_F()
985 dnums.add_output_spatial_dimensions(1); in XLA_TEST_F()
986 dnums.add_input_spatial_dimensions(2); in XLA_TEST_F()
987 dnums.add_output_spatial_dimensions(2); in XLA_TEST_F()
988 dnums.set_input_feature_dimension(3); in XLA_TEST_F()
989 dnums.set_output_feature_dimension(3); in XLA_TEST_F()
992 dnums.add_kernel_spatial_dimensions(0); in XLA_TEST_F()
[all …]
Dconvolution_test_1d.cc93 ConvolutionDimensionNumbers dnums; in TestImpl() local
94 dnums.set_input_batch_dimension(0); in TestImpl()
95 dnums.set_output_batch_dimension(0); in TestImpl()
96 dnums.add_input_spatial_dimensions(1); in TestImpl()
97 dnums.add_output_spatial_dimensions(1); in TestImpl()
98 dnums.set_input_feature_dimension(2); in TestImpl()
99 dnums.set_output_feature_dimension(2); in TestImpl()
100 dnums.add_kernel_spatial_dimensions(0); in TestImpl()
101 dnums.set_kernel_input_feature_dimension(1); in TestImpl()
102 dnums.set_kernel_output_feature_dimension(2); in TestImpl()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/
Dattribute_importer.cc49 const xla::GatherDimensionNumbers& dnums, mlir::Builder* builder) { in ConvertGatherDimensionNumbers() argument
50 std::vector<int64_t> offset_dims(dnums.offset_dims().begin(), in ConvertGatherDimensionNumbers()
51 dnums.offset_dims().end()); in ConvertGatherDimensionNumbers()
53 dnums.collapsed_slice_dims().begin(), dnums.collapsed_slice_dims().end()); in ConvertGatherDimensionNumbers()
54 std::vector<int64_t> start_index_map(dnums.start_index_map().begin(), in ConvertGatherDimensionNumbers()
55 dnums.start_index_map().end()); in ConvertGatherDimensionNumbers()
59 builder->getI64IntegerAttr(dnums.index_vector_dim()), in ConvertGatherDimensionNumbers()
64 const xla::ScatterDimensionNumbers& dnums, mlir::Builder* builder) { in ConvertScatterDimensionNumbers() argument
65 std::vector<int64_t> update_window_dims(dnums.update_window_dims().begin(), in ConvertScatterDimensionNumbers()
66 dnums.update_window_dims().end()); in ConvertScatterDimensionNumbers()
[all …]
Dattribute_importer.h34 const xla::GatherDimensionNumbers& dnums, mlir::Builder* builder);
38 const xla::ScatterDimensionNumbers& dnums, mlir::Builder* builder);
42 const DotDimensionNumbers& dnums, mlir::Builder* builder);
46 const xla::ConvolutionDimensionNumbers& dnums, mlir::Builder* builder);
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dconv_canonicalization_test.cc70 ConvolutionDimensionNumbers dnums; in TEST_F() local
71 dnums.set_input_batch_dimension(1); in TEST_F()
72 dnums.set_output_batch_dimension(1); in TEST_F()
73 dnums.add_input_spatial_dimensions(2); in TEST_F()
74 dnums.add_output_spatial_dimensions(2); in TEST_F()
75 dnums.add_input_spatial_dimensions(3); in TEST_F()
76 dnums.add_output_spatial_dimensions(3); in TEST_F()
77 dnums.set_input_feature_dimension(0); in TEST_F()
78 dnums.set_output_feature_dimension(0); in TEST_F()
79 dnums.add_kernel_spatial_dimensions(2); in TEST_F()
[all …]
Dir_emission_utils.cc83 const ConvolutionDimensionNumbers& dnums = in PotentiallyImplementedAsEigenConvolution() local
87 const int64 num_spatial_dims = dnums.output_spatial_dimensions_size(); in PotentiallyImplementedAsEigenConvolution()
93 if (dnums.input_spatial_dimensions(i) != i + 1) { in PotentiallyImplementedAsEigenConvolution()
96 if (dnums.kernel_spatial_dimensions(i) != i) { in PotentiallyImplementedAsEigenConvolution()
99 if (dnums.output_spatial_dimensions(i) != i + 1) { in PotentiallyImplementedAsEigenConvolution()
104 return dnums.input_batch_dimension() == 0 && in PotentiallyImplementedAsEigenConvolution()
105 dnums.input_feature_dimension() == input_shape.dimensions_size() - 1 && in PotentiallyImplementedAsEigenConvolution()
106 dnums.output_batch_dimension() == 0 && in PotentiallyImplementedAsEigenConvolution()
107 dnums.output_feature_dimension() == in PotentiallyImplementedAsEigenConvolution()
109 dnums.kernel_input_feature_dimension() == in PotentiallyImplementedAsEigenConvolution()
[all …]
Dconv_canonicalization.cc39 const ConvolutionDimensionNumbers& dnums = in Run() local
41 auto input_batch_dim = dnums.input_batch_dimension(); in Run()
42 auto input_feature_dim = dnums.input_feature_dimension(); in Run()
43 auto kernel_input_feature_dim = dnums.kernel_input_feature_dimension(); in Run()
44 auto kernel_output_feature_dim = dnums.kernel_output_feature_dimension(); in Run()
46 const int64 num_spatial_dims = dnums.output_spatial_dimensions_size(); in Run()
65 new_input_dim_order[i + 1] = dnums.input_spatial_dimensions(i); in Run()
67 input->shape().dimensions(dnums.input_spatial_dimensions(i)); in Run()
84 new_kernel_dim_order[i] = dnums.kernel_spatial_dimensions(i); in Run()
86 kernel->shape().dimensions(dnums.kernel_spatial_dimensions(i)); in Run()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Deinsum.cc139 EinsumDimensionNumbers dnums; in GetEinsumDimensionNumbers() local
144 dnums.lhs.emplace_back(i); in GetEinsumDimensionNumbers()
146 dnums.lhs_out.emplace_back(i, out_index->second); in GetEinsumDimensionNumbers()
148 dnums.lhs_rhs.emplace_back(i, rhs_index->second); in GetEinsumDimensionNumbers()
150 dnums.lhs_rhs_out.emplace_back(i, rhs_index->second, out_index->second); in GetEinsumDimensionNumbers()
158 dnums.rhs.emplace_back(i); in GetEinsumDimensionNumbers()
160 dnums.rhs_out.emplace_back(i, out_index->second); in GetEinsumDimensionNumbers()
172 return dnums; in GetEinsumDimensionNumbers()
187 const Location& loc, EinsumDimensionNumbers& dnums, Value* lhs, Value* rhs, in transposeForBatchMatmul() argument
192 lhs_transpose.reserve(dnums.lhs_rhs_out.size() + dnums.lhs_out.size() + in transposeForBatchMatmul()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc55 const auto& dnums = original_hlo->convolution_dimension_numbers(); in PartitionConvolutionWithBatchGroupCount() local
58 lhs.base_shape().dimensions(dnums.input_batch_dimension()); in PartitionConvolutionWithBatchGroupCount()
60 rhs.base_shape().dimensions(dnums.kernel_output_feature_dimension()); in PartitionConvolutionWithBatchGroupCount()
68 rhs_to_lhs_indices[dnums.kernel_output_feature_dimension()] = in PartitionConvolutionWithBatchGroupCount()
69 dnums.input_batch_dimension(); in PartitionConvolutionWithBatchGroupCount()
70 rhs_to_lhs_indices[dnums.kernel_input_feature_dimension()] = in PartitionConvolutionWithBatchGroupCount()
71 dnums.input_feature_dimension(); in PartitionConvolutionWithBatchGroupCount()
72 for (int64 i = 0; i < dnums.input_spatial_dimensions_size(); ++i) { in PartitionConvolutionWithBatchGroupCount()
73 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithBatchGroupCount()
74 dnums.input_spatial_dimensions(i); in PartitionConvolutionWithBatchGroupCount()
[all …]
Dgather_scatter_handler.cc139 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in PartitionIndexOnlyPartition() local
142 (dnums.index_vector_dim() == indices.base_shape().rank() || in PartitionIndexOnlyPartition()
143 indices.sharding().tile_assignment().dim(dnums.index_vector_dim()) == in PartitionIndexOnlyPartition()
149 indices.hlo()->shape(), dnums, in PartitionIndexOnlyPartition()
157 int64 indices_batch_dim = i < dnums.index_vector_dim() ? i : i + 1; in PartitionIndexOnlyPartition()
185 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in ParititonPassthroughOperand() local
200 pshape, operand.hlo(), indices.hlo(), dnums, pslice_sizes, in ParititonPassthroughOperand()
220 GatherDimensionNumbers dnums = gather->gather_dimension_numbers(); in ParititonTrivialIndexedOperandDimension() local
221 std::vector<int64> start_index_map(dnums.start_index_map().begin(), in ParititonTrivialIndexedOperandDimension()
222 dnums.start_index_map().end()); in ParititonTrivialIndexedOperandDimension()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dstream_executor_util.cc49 StreamExecutorConvLayoutsToXlaLayouts(const ConvolutionDimensionNumbers& dnums, in StreamExecutorConvLayoutsToXlaLayouts() argument
55 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
56 input_layout.push_back(dnums.input_feature_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
58 dnums.input_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts()
59 dnums.input_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts()
62 input_layout.push_back(dnums.input_batch_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
64 dnums.input_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts()
65 dnums.input_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts()
66 input_layout.push_back(dnums.input_feature_dimension()); in StreamExecutorConvLayoutsToXlaLayouts()
71 ConvolutionDimensionNumbersToString(dnums)); in StreamExecutorConvLayoutsToXlaLayouts()
[all …]
Dcudnn_pad_for_convolutions.cc182 const auto& dnums = conv->convolution_dimension_numbers(); in TryResolvePaddedShapesForTensorCore() local
227 new_input_shape->dimensions(dnums.input_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
229 new_output_shape->dimensions(dnums.output_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
231 new_input_shape->set_dimensions(dnums.input_feature_dimension(), 4); in TryResolvePaddedShapesForTensorCore()
232 new_filter_shape->set_dimensions(dnums.kernel_input_feature_dimension(), 4); in TryResolvePaddedShapesForTensorCore()
237 pad_dim(new_input_shape, dnums.input_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
238 pad_dim(new_filter_shape, dnums.kernel_input_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
239 pad_dim(new_filter_shape, dnums.kernel_output_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
240 pad_dim(new_output_shape, dnums.output_feature_dimension()); in TryResolvePaddedShapesForTensorCore()
299 const auto& dnums = conv->convolution_dimension_numbers(); in TryResolvePaddedShapesForIntegerConvolution() local
[all …]
Dgpu_conv_rewriter.cc43 const ConvolutionDimensionNumbers& dnums, in CreateGpuConv() argument
62 custom_call->set_convolution_dimension_numbers(dnums); in CreateGpuConv()
122 const ConvolutionDimensionNumbers& dnums = in CanImplementAsGpuForwardConv() local
124 if (dnums.input_spatial_dimensions_size() > 3) { in CanImplementAsGpuForwardConv()
136 if (dnums.input_spatial_dimensions_size() == 2 in CanImplementAsGpuForwardConv()
344 ConvolutionDimensionNumbers dnums = conv->convolution_dimension_numbers(); in MatchBackwardInput() local
364 absl::c_is_permutation(dnums.kernel_spatial_dimensions(), in MatchBackwardInput()
398 const auto& input_spatial_dims = dnums.input_spatial_dimensions(); in MatchBackwardInput()
399 const auto& output_spatial_dims = dnums.output_spatial_dimensions(); in MatchBackwardInput()
505 dnums.set_kernel_input_feature_dimension( in MatchBackwardInput()
[all …]
Dgpu_conv_runner.cc314 const ConvolutionDimensionNumbers& dnums = desc.dnums; in GetGpuConvConfig() local
328 VLOG(3) << "Dim nums: { " << dnums.ShortDebugString() << " }"; in GetGpuConvConfig()
344 CHECK_EQ(num_dimensions, dnums.input_spatial_dimensions_size()) in GetGpuConvConfig()
346 CHECK_EQ(num_dimensions, dnums.kernel_spatial_dimensions_size()) in GetGpuConvConfig()
348 CHECK_EQ(num_dimensions, dnums.output_spatial_dimensions_size()) in GetGpuConvConfig()
371 dnums, input_shape.layout(), filter_shape.layout(), in GetGpuConvConfig()
378 input_shape.dimensions(dnums.input_feature_dimension())) in GetGpuConvConfig()
379 .set_count(input_shape.dimensions(dnums.input_batch_dimension())); in GetGpuConvConfig()
384 input_shape.dimensions(dnums.input_spatial_dimensions(dim))); in GetGpuConvConfig()
391 filter_shape.dimensions(dnums.kernel_input_feature_dimension())) in GetGpuConvConfig()
[all …]
/external/tensorflow/tensorflow/compiler/xla/service/
Ddot_as_convolution_util.cc157 dot_as_convolution_util::DotConvolutionDimsInfo dnums; in ParseDotGeneralFromDot() local
159 dnums.batch_dims.emplace_back(); in ParseDotGeneralFromDot()
160 dnums.batch_dims.back().lhs = dot_dim_numbs.lhs_batch_dimensions(i); in ParseDotGeneralFromDot()
161 dnums.batch_dims.back().rhs = dot_dim_numbs.rhs_batch_dimensions(i); in ParseDotGeneralFromDot()
162 dnums.batch_dims.back().output = i; in ParseDotGeneralFromDot()
163 dnums.batch_dims.back().spatial_dim = -1; in ParseDotGeneralFromDot()
167 dnums.contracting_dims.emplace_back(); in ParseDotGeneralFromDot()
168 dnums.contracting_dims.back().lhs = in ParseDotGeneralFromDot()
170 dnums.contracting_dims.back().rhs = in ParseDotGeneralFromDot()
172 dnums.contracting_dims.back().output = -1; in ParseDotGeneralFromDot()
[all …]
Dtranspose_folding_test.cc231 auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(); in TEST_F() local
241 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F()
245 /*batch_group_count=*/1, window, dnums, in TEST_F()
250 /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, in TEST_F()
268 EXPECT_EQ(dnums.kernel_input_feature_dimension(), in TEST_F()
271 EXPECT_EQ(dnums.kernel_output_feature_dimension(), in TEST_F()
288 auto dnums = XlaBuilder::CreateDefaultConvDimensionNumbers(); in TEST_F() local
298 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F()
302 /*batch_group_count=*/1, window, dnums, in TEST_F()
307 /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums, in TEST_F()
[all …]
Dhlo_sharding_util.cc505 const GatherDimensionNumbers& dnums = hlo->gather_dimension_numbers(); in GatherOutputSharding() local
508 if (absl::c_binary_search(dnums.offset_dims(), i)) { in GatherOutputSharding()
512 index_dim >= dnums.index_vector_dim() ? index_dim + 1 : index_dim; in GatherOutputSharding()
544 const GatherDimensionNumbers& dnums = hlo->gather_dimension_numbers(); in GatherIndexSharding() local
547 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherIndexSharding()
557 index_tile_assignment_dims.begin() + dnums.index_vector_dim(), 1); in GatherIndexSharding()
593 const GatherDimensionNumbers& dnums = hlo.gather_dimension_numbers(); in GatherEffectiveOutputSharding() local
597 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherEffectiveOutputSharding()
628 if (!absl::c_binary_search(dnums.offset_dims(), i)) { in GatherEffectiveOutputSharding()
645 const ScatterDimensionNumbers& dnums = hlo->scatter_dimension_numbers(); in ScatterIndexSharding() local
[all …]
Ddot_decomposer.cc180 const DotDimensionNumbers& dnums = instruction->dot_dimension_numbers(); in Run() local
183 if (dnums.lhs_contracting_dimensions_size() != 1) { in Run()
189 if (dnums.lhs_batch_dimensions_size() + 2 < in Run()
191 dnums.rhs_batch_dimensions_size() + 2 < in Run()
196 if (dnums.lhs_batch_dimensions().empty() && in Run()
197 dnums.lhs_contracting_dimensions().empty()) { in Run()
201 if (dnums.lhs_batch_dimensions().empty()) { in Run()
205 dnums.lhs_batch_dimensions_size()); in Run()
207 if (!absl::c_equal(dnums.lhs_batch_dimensions(), canonical_batch_dims) || in Run()
208 !absl::c_equal(dnums.rhs_batch_dimensions(), canonical_batch_dims)) { in Run()
Dshape_inference_test.cc403 ConvolutionDimensionNumbers dnums; in TEST_F() local
407 dnums.set_input_batch_dimension(0); in TEST_F()
408 dnums.set_output_batch_dimension(0); in TEST_F()
409 dnums.set_input_feature_dimension(1); in TEST_F()
410 dnums.set_output_feature_dimension(1); in TEST_F()
411 dnums.add_input_spatial_dimensions(2); in TEST_F()
412 dnums.add_output_spatial_dimensions(2); in TEST_F()
413 dnums.add_input_spatial_dimensions(3); in TEST_F()
414 dnums.add_output_spatial_dimensions(3); in TEST_F()
418 dnums.set_kernel_input_feature_dimension(2); in TEST_F()
[all …]
Dsharding_propagation.cc132 const auto& dnums = instruction->convolution_dimension_numbers(); in IsConvolutionKernelSmall() local
133 for (int64 i = 0; i < dnums.input_spatial_dimensions().size(); ++i) { in IsConvolutionKernelSmall()
135 rhs->shape().dimensions(dnums.kernel_spatial_dimensions(i)); in IsConvolutionKernelSmall()
137 instruction->shape().dimensions(dnums.output_spatial_dimensions(i)); in IsConvolutionKernelSmall()
338 const dot_as_convolution_util::DotConvolutionDimsInfo& dnums, in InferDotShardingFromOperands() argument
347 contracting_dims.reserve(dnums.contracting_dims.size()); in InferDotShardingFromOperands()
348 for (const auto& dim : dnums.contracting_dims) { in InferDotShardingFromOperands()
354 ? dnums.rhs_non_contracting_dims in InferDotShardingFromOperands()
355 : dnums.lhs_non_contracting_dims) { in InferDotShardingFromOperands()
366 for (const auto& dim : dnums.batch_dims) { in InferDotShardingFromOperands()
[all …]
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dconv_op_helpers.cc327 xla::ConvolutionDimensionNumbers dnums; in MakeXlaBackpropInputConvOp() local
328 dnums.set_input_batch_dimension(batch_dim); in MakeXlaBackpropInputConvOp()
329 dnums.set_output_batch_dimension(batch_dim); in MakeXlaBackpropInputConvOp()
330 dnums.set_input_feature_dimension(feature_dim); in MakeXlaBackpropInputConvOp()
331 dnums.set_output_feature_dimension(feature_dim); in MakeXlaBackpropInputConvOp()
335 dnums.set_kernel_input_feature_dimension(attrs.num_spatial_dims + 1); in MakeXlaBackpropInputConvOp()
336 dnums.set_kernel_output_feature_dimension(attrs.num_spatial_dims); in MakeXlaBackpropInputConvOp()
356 dnums.add_input_spatial_dimensions(dim); in MakeXlaBackpropInputConvOp()
357 dnums.add_kernel_spatial_dimensions(i); in MakeXlaBackpropInputConvOp()
358 dnums.add_output_spatial_dimensions(dim); in MakeXlaBackpropInputConvOp()
[all …]
Dreverse_sequence_op.cc100 xla::GatherDimensionNumbers dnums; in Compile() local
101 dnums.set_index_vector_dim(2); in Compile()
104 dnums.add_start_index_map(batch_dim_); in Compile()
105 dnums.add_start_index_map(seq_dim_); in Compile()
111 dnums.add_offset_dims(i); in Compile()
113 dnums.add_collapsed_slice_dims(i); in Compile()
122 xla::Gather(input, start_indices, dnums, slice_sizes)); in Compile()
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder_test.cc783 ConvolutionDimensionNumbers dnums; in TEST_F() local
784 dnums.set_input_batch_dimension(0); in TEST_F()
785 dnums.set_output_batch_dimension(0); in TEST_F()
786 dnums.add_input_spatial_dimensions(1); in TEST_F()
787 dnums.add_output_spatial_dimensions(1); in TEST_F()
788 dnums.add_input_spatial_dimensions(2); in TEST_F()
789 dnums.add_output_spatial_dimensions(2); in TEST_F()
790 dnums.set_input_feature_dimension(3); in TEST_F()
791 dnums.set_output_feature_dimension(3); in TEST_F()
792 dnums.add_kernel_spatial_dimensions(0); in TEST_F()
[all …]
/external/tensorflow/tensorflow/compiler/tests/
Dxla_ops_test.py136 dnums = xla_data_pb2.ConvolutionDimensionNumbers()
138 dnums.input_batch_dimension = 0
139 dnums.input_feature_dimension = 1
140 dnums.output_batch_dimension = 0
141 dnums.output_feature_dimension = 1
142 dnums.kernel_output_feature_dimension = 0
143 dnums.kernel_input_feature_dimension = 1
144 dnums.input_spatial_dimensions.extend(range(2, 2 + num_spatial_dims))
145 dnums.kernel_spatial_dimensions.extend(range(2, 2 + num_spatial_dims))
146 dnums.output_spatial_dimensions.extend(range(2, 2 + num_spatial_dims))
[all …]

123