/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
H A D | cudnn_vectorize_convolutions_test.cc | 87 ASSERT_EQ(dnums.kernel_spatial_dimensions().size(), 2); in TEST_F() 95 EXPECT_EQ(dnums.kernel_spatial_dimensions()[0], 0); in TEST_F() 96 EXPECT_EQ(dnums.kernel_spatial_dimensions()[1], 1); in TEST_F() 160 ASSERT_EQ(dnums.kernel_spatial_dimensions().size(), 2); in TEST_F() 170 EXPECT_EQ(dnums.kernel_spatial_dimensions()[0], 3); in TEST_F() 171 EXPECT_EQ(dnums.kernel_spatial_dimensions()[1], 4); in TEST_F() 213 ASSERT_EQ(dnums.kernel_spatial_dimensions().size(), 2); in TEST_F() 222 EXPECT_EQ(dnums.kernel_spatial_dimensions()[0], 2); in TEST_F() 223 EXPECT_EQ(dnums.kernel_spatial_dimensions()[1], 3); in TEST_F() 266 ASSERT_EQ(dnums.kernel_spatial_dimensions().size(), 2); in TEST_F() [all …]
|
H A D | stream_executor_util.cc | 107 dnums.kernel_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts() 108 dnums.kernel_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts() 115 dnums.kernel_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts() 116 dnums.kernel_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts() 122 dnums.kernel_spatial_dimensions().begin(), in StreamExecutorConvLayoutsToXlaLayouts() 123 dnums.kernel_spatial_dimensions().end()); in StreamExecutorConvLayoutsToXlaLayouts() 305 dnums.kernel_spatial_dimensions()), in FindVectorizedFeatureDims()
|
H A D | cudnn_support_utils.cc | 71 dnums.kernel_spatial_dimensions().size() != 2 || in CudnnSupportsOptimizedIntegerConvolution() 97 int64_t R = kernel_shape.dimensions(dnums.kernel_spatial_dimensions()[0]); in CudnnSupportsOptimizedIntegerConvolution() 98 int64_t S = kernel_shape.dimensions(dnums.kernel_spatial_dimensions()[1]); in CudnnSupportsOptimizedIntegerConvolution()
|
H A D | gpu_conv_rewriter.cc | 107 auto kernel_spatial_dims = conv_dnums.kernel_spatial_dimensions(); in MatchBackwardFilter() 312 absl::c_is_permutation(dnums.kernel_spatial_dimensions(), in MatchBackwardInput() 477 dnums.kernel_spatial_dimensions())); in MatchBackwardInput() 480 dnums.kernel_spatial_dimensions())); in MatchBackwardInput()
|
H A D | gpu_conv_padding_legalization.cc | 137 for (size_t i = 0; i < conv_dnums.kernel_spatial_dimensions().size(); ++i) { in MaybePaddedKernel() 138 int64_t dim = conv_dnums.kernel_spatial_dimensions(i); in MaybePaddedKernel() 175 conv->convolution_dimension_numbers().kernel_spatial_dimensions(i))); in CanonicalizeForwardConvolution()
|
H A D | cudnn_simplify_padding.cc | 78 for (int64_t d : dnums.kernel_spatial_dimensions()) { in FindKernelVectCDim()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | transpose_folding_test.cc | 247 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F() 305 transpose_y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F() 336 EXPECT_EQ(dnums.kernel_spatial_dimensions(1), in TEST_F() 341 new_conv->convolution_dimension_numbers().kernel_spatial_dimensions(0)); in TEST_F() 343 dnums.kernel_spatial_dimensions(0), in TEST_F() 344 new_conv->convolution_dimension_numbers().kernel_spatial_dimensions(1)); in TEST_F() 368 dim->set_size(y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F() 437 dim->set_size(y->shape().dimensions(dnums.kernel_spatial_dimensions(i))); in TEST_F()
|
H A D | convolution_4d_expander.cc | 77 int64_t kernel_spatial_dim = dim_nums.kernel_spatial_dimensions(i); in ExpandInstruction() 153 for (int64_t i = 0; i < new_dim_nums.kernel_spatial_dimensions().size(); in ExpandInstruction() 157 new_dim_nums.kernel_spatial_dimensions(i))); in ExpandInstruction()
|
H A D | dot_as_convolution_util.cc | 97 int64_t rhs = conv_dims.kernel_spatial_dimensions(i); in ParseConvolutionDimsInfo() 163 conv_dnums.kernel_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvForDotGeneralConvolution()
|
H A D | space_to_batch_converter.cc | 385 const int64_t kernel_spatial_dim = dim_numbers.kernel_spatial_dimensions( in IsForwardWindowDilatedConv() 676 auto it = absl::c_find(dim_numbers.kernel_spatial_dimensions(), i); in BringSpaceNextToBatch() 677 if (it != dim_numbers.kernel_spatial_dimensions().end()) { in BringSpaceNextToBatch() 678 int64_t j = it - dim_numbers.kernel_spatial_dimensions().begin(); in BringSpaceNextToBatch() 1097 if (dim_numbers.kernel_spatial_dimensions( in CanPropagate() 1102 for (int i = 0; i < dim_numbers.kernel_spatial_dimensions().size(); in CanPropagate() 1104 if (dim_numbers.kernel_spatial_dimensions(i) == in CanPropagate() 1106 dim_numbers.kernel_spatial_dimensions(i) == in CanPropagate() 3012 original_conv_dims.kernel_spatial_dimensions( in PropagateOnBackpropFilterConv() 3149 original_conv_dims.kernel_spatial_dimensions(i))); in PropagateOnBackpropFilterConv() [all …]
|
H A D | hlo_cse.cc | 186 absl::MakeSpan(conv_dimension_numbers.kernel_spatial_dimensions()), in AbslHashValue()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
H A D | convolution_handler.cc | 73 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithBatchGroupCount() 164 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithFeatureGroupCount() 246 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() 322 int64_t rhs_dimension = dnums.kernel_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() 359 int64_t rhs_dimension = dnums.kernel_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() 465 int64_t dim = dnums.kernel_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() 538 rhs_to_lhs_indices[dnums.kernel_spatial_dimensions(i)] = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() 550 reversed_rhs_dims.push_back(dnums.kernel_spatial_dimensions(i)); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() 631 int64_t rhs_dimension = dnums.kernel_spatial_dimensions(i); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() 924 conv_dnums.kernel_spatial_dimensions(dim.spatial_dim))); in CreateShardedConvConvolution()
|
H A D | dot_handler.cc | 180 if (conv_dnums.kernel_spatial_dimensions(i) == rhs_concat_dim && in GenNewWindow() 274 std::vector<int64_t> kernel_spatial_dimensions( in GenNewConvDNums() local 275 dnums.kernel_spatial_dimensions().begin(), in GenNewConvDNums() 276 dnums.kernel_spatial_dimensions().end()); in GenNewConvDNums() 284 for (int64_t i = 0; i < kernel_spatial_dimensions.size(); ++i) { in GenNewConvDNums() 285 if (rhs_concat_dim <= kernel_spatial_dimensions[i]) { in GenNewConvDNums() 286 kernel_spatial_dimensions[i]++; in GenNewConvDNums() 289 kernel_spatial_dimensions.push_back(rhs_concat_dim); in GenNewConvDNums() 293 kernel_spatial_dimensions.push_back(dot_rhs->shape().dimensions_size() - 1); in GenNewConvDNums() 330 for (auto dim : kernel_spatial_dimensions) { in GenNewConvDNums()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/ |
H A D | attribute_importer.cc | 94 dnums.kernel_spatial_dimensions().begin(), in ConvertConvDimensionNumbers() 95 dnums.kernel_spatial_dimensions().end()); in ConvertConvDimensionNumbers() 105 arrayref(dnums.kernel_spatial_dimensions()), in ConvertConvDimensionNumbers()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/python/ |
H A D | attributes.py | 100 kernel_spatial_dimensions=[0, 3], 111 assert attr.kernel_spatial_dimensions == [0, 3]
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/ |
H A D | reference_util.cc | 424 if (dnums.kernel_spatial_dimensions(0) > dnums.kernel_spatial_dimensions(1)) { in ConvArray4DGeneralDimensionsDilated() 437 rhs_literal.shape().dimensions(dnums.kernel_spatial_dimensions(0)); in ConvArray4DGeneralDimensionsDilated() 439 rhs_literal.shape().dimensions(dnums.kernel_spatial_dimensions(1)); in ConvArray4DGeneralDimensionsDilated() 450 rhs_literal.shape().dimensions(dnums.kernel_spatial_dimensions(0))); in ConvArray4DGeneralDimensionsDilated() 460 rhs_literal.shape().dimensions(dnums.kernel_spatial_dimensions(1))); in ConvArray4DGeneralDimensionsDilated()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/mhlo/ |
H A D | verifier_conv_op.mlir | 17 kernel_spatial_dimensions = [0, 1], 44 kernel_spatial_dimensions = [0, 1], 202 kernel_spatial_dimensions = [0, 1], 230 kernel_spatial_dimensions = [0, 0], 258 kernel_spatial_dimensions = [0, 1], 517 kernel_spatial_dimensions = [0, 1], 545 kernel_spatial_dimensions = [0, 1],
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/stablehlo/tests/ |
H A D | verify_conv.mlir | 17 kernel_spatial_dimensions = [0, 1], 44 kernel_spatial_dimensions = [0, 1], 202 kernel_spatial_dimensions = [0, 1], 230 kernel_spatial_dimensions = [0, 0], 258 kernel_spatial_dimensions = [0, 1], 517 kernel_spatial_dimensions = [0, 1], 545 kernel_spatial_dimensions = [0, 1],
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/cpu/ |
H A D | conv_canonicalization.cc | 86 new_kernel_dim_order[i] = dnums.kernel_spatial_dimensions(i); in Run() 88 kernel->shape().dimensions(dnums.kernel_spatial_dimensions(i)); in Run()
|
H A D | ir_emission_utils.cc | 94 if (dnums.kernel_spatial_dimensions(i) != i) { in PotentiallyImplementedAsEigenConvolution()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/ |
H A D | xla_client.py | 519 self.kernel_spatial_dimensions = [] 565 dimension_numbers.kernel_spatial_dimensions.extend(range(2, 2 + nd)) 578 dimension_numbers.kernel_spatial_dimensions.extend(
|
H A D | xla_client.pyi | 145 kernel_spatial_dimensions: List[int]
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/ |
H A D | layouts_and_names.mlir | 16 kernel_spatial_dimensions = [1, 2],
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/mlir_hlo/tests/Dialect/lhlo_gpu/ |
H A D | lhlo_gpu_ops.mlir | 20 kernel_spatial_dimensions = [2, 3],
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/ |
H A D | xla_builder.cc | 1439 dimension_numbers.kernel_spatial_dimensions())); in VerifyConvolution() 1496 rhs_shape->dimensions(dimension_numbers.kernel_spatial_dimensions(i)); in ConvWithGeneralDimensions() 1545 rhs_shape->dimensions(dimension_numbers.kernel_spatial_dimensions(i)); in ConvGeneralDilated() 1580 rhs_shape->dimensions(dimension_numbers.kernel_spatial_dimensions(i)); in DynamicConvInstruction() 3881 dnum.kernel_spatial_dimensions(0), in Validate() 3882 dnum.kernel_spatial_dimensions(1)}) in Validate() 3889 dnum.kernel_spatial_dimensions(0), dnum.kernel_spatial_dimensions(1)); in Validate()
|