Home
last modified time | relevance | path

Searched refs:batch_group_count (Results 1 – 25 of 46) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dconvolution_group_converter.cc204 int64 batch_group_count = convolution->batch_group_count(); in HandleBatchGroupCount() local
206 if (batch_group_count == 1) { in HandleBatchGroupCount()
210 VLOG(2) << "Dealing with batch_group_count " << batch_group_count in HandleBatchGroupCount()
233 if (output_feature != batch_group_count || input_batch != batch_group_count) { in HandleBatchGroupCount()
238 input_sizes[input_batch_dimension] /= batch_group_count; in HandleBatchGroupCount()
240 batch_group_count); in HandleBatchGroupCount()
257 kernel_sizes[kernel_output_feature_dimension] /= batch_group_count; in HandleBatchGroupCount()
259 batch_group_count); in HandleBatchGroupCount()
294 window_dim->set_base_dilation(batch_group_count); in HandleBatchGroupCount()
295 window_dim->set_size(batch_group_count); in HandleBatchGroupCount()
[all …]
Doperand_upcaster.cc34 instruction->feature_group_count(), instruction->batch_group_count(), in MaybeInferShape()
Ddot_as_convolution_util.cc145 /*batch_group_count=*/conv.batch_group_count(), window, conv_dnums, in CreateShardedConvForDotGeneralConvolution()
151 /*batch_group_count=*/conv.batch_group_count(), window, conv_dnums, in CreateShardedConvForDotGeneralConvolution()
Dshape_inference.cc1617 int64 batch_group_count, const Window& window, in InferConvolveShape() argument
1629 if (batch_group_count <= 0) { in InferConvolveShape()
1632 batch_group_count); in InferConvolveShape()
1635 if (batch_group_count > 1 && feature_group_count > 1) { in InferConvolveShape()
1639 batch_group_count, feature_group_count); in InferConvolveShape()
1750 if (kernel_output_features % batch_group_count != 0) { in InferConvolveShape()
1755 kernel_output_features, batch_group_count, ShapeUtil::HumanString(lhs), in InferConvolveShape()
1789 if (input_batch % batch_group_count != 0) { in InferConvolveShape()
1795 input_batch, batch_group_count, ShapeUtil::HumanString(lhs), in InferConvolveShape()
1821 dimensions[dnums.output_batch_dimension()] = input_batch / batch_group_count; in InferConvolveShape()
Dhlo_creation_utils.cc95 int64 batch_group_count, const Window& window, in MakeConvolveHlo() argument
104 lhs->shape(), rhs->shape(), feature_group_count, batch_group_count, in MakeConvolveHlo()
107 convolve_shape, lhs, rhs, feature_group_count, batch_group_count, window, in MakeConvolveHlo()
Dtranspose_folding.cc208 convolution.batch_group_count(), convolution.window(), new_dnums, in FoldTransposeIntoConvolution()
Dhlo_creation_utils.h66 int64 batch_group_count, const Window& window,
Dhlo_instructions.h1228 int64 feature_group_count, int64 batch_group_count, const Window& window,
1247 int64 batch_group_count() const { return batch_group_count_; } in batch_group_count() function
1449 void set_batch_group_count(int64 batch_group_count) { in set_batch_group_count() argument
1450 batch_group_count_ = batch_group_count; in set_batch_group_count()
1458 int64 batch_group_count() const { return batch_group_count_; } in batch_group_count() function
Dshape_inference.h113 int64 batch_group_count, const Window& window,
Dhlo_parser.cc1610 optional<int64> batch_group_count; in ParseInstructionRhs() local
1617 &batch_group_count}; in ParseInstructionRhs()
1631 if (!batch_group_count) { in ParseInstructionRhs()
1632 batch_group_count = 1; in ParseInstructionRhs()
1646 *feature_group_count, *batch_group_count, *window, *dnums, in ParseInstructionRhs()
1654 feature_group_count.value(), batch_group_count.value(), *window, in ParseInstructionRhs()
2241 optional<int64> batch_group_count; in ParseInstructionRhs() local
2257 &batch_group_count}; in ParseInstructionRhs()
2351 if (batch_group_count.has_value()) { in ParseInstructionRhs()
2352 custom_call_instr->set_batch_group_count(*batch_group_count); in ParseInstructionRhs()
Dhlo_instruction.cc518 std::max<int64>(proto.batch_group_count(), 1), proto.window(), in CreateFromProto()
596 std::max(static_cast<int64>(proto.batch_group_count()), int64{1})); in CreateFromProto()
985 int64 feature_group_count, int64 batch_group_count, const Window& window, in CreateConvolve() argument
989 shape, lhs, rhs, feature_group_count, batch_group_count, window, in CreateConvolve()
4181 int64 HloInstruction::batch_group_count() const { in batch_group_count() function in xla::HloInstruction
4183 return convolution->batch_group_count(); in batch_group_count()
4185 return Cast<HloCustomCallInstruction>(this)->batch_group_count(); in batch_group_count()
4188 void HloInstruction::set_batch_group_count(int64 batch_group_count) { in set_batch_group_count() argument
4190 return convolution->set_batch_group_count(batch_group_count); in set_batch_group_count()
4193 batch_group_count); in set_batch_group_count()
Dhlo_instruction.h586 int64 feature_group_count, int64 batch_group_count, const Window& window,
1850 int64 batch_group_count() const;
1852 void set_batch_group_count(int64 batch_group_count);
Dhlo_evaluator_typed_visitor.h1123 const int64 batch_group_count = conv->batch_group_count();
1128 batch_group_count](const absl::Span<const int64> out_index) {
1145 const int64 batch_group_size = input_batch_size / batch_group_count;
1163 batch_group_count > 1 ? output_z_size / input_batch_size : 1;
1293 conv->batch_group_count(), window, dnums,
/external/tensorflow/tensorflow/compiler/mlir/hlo/tests/
Dlhlo_gpu_ops.mlir51 batch_group_count = 1,
71 batch_group_count = 1 : i64,
101 batch_group_count = 1 : i64,
133 batch_group_count = 1 : i64,
164 batch_group_count = 1 : i64,
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dconvolution_handler.cc50 if (original_hlo->batch_group_count() == 1 || in PartitionConvolutionWithBatchGroupCount()
51 original_hlo->batch_group_count() < num_partitions) { in PartitionConvolutionWithBatchGroupCount()
62 original_hlo->batch_group_count() != input_batch_size) { in PartitionConvolutionWithBatchGroupCount()
292 if (original_hlo->batch_group_count() > 1 && in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()
602 if (original_hlo->batch_group_count() > 1 && in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()
829 if (original_hlo->batch_group_count() > 1) { in PartitionConvolutionBaseCase()
943 int64 batch_group_count = conv.batch_group_count(); in CreateShardedConvConvolution() local
944 if (batch_group_count > 1) { in CreateShardedConvConvolution()
945 batch_group_count = in CreateShardedConvConvolution()
953 feature_group_count, batch_group_count, window, conv_dnums, in CreateShardedConvConvolution()
[all …]
/external/tensorflow/tensorflow/compiler/xla/client/
Dxla_builder.h536 int64 batch_group_count = 1,
543 int64 feature_group_count = 1, int64 batch_group_count = 1,
550 int64 feature_group_count = 1, int64 batch_group_count = 1,
558 int64 feature_group_count = 1, int64 batch_group_count = 1,
568 int64 feature_group_count = 1, int64 batch_group_count = 1,
578 int64 feature_group_count, int64 batch_group_count,
589 int64 feature_group_count, int64 batch_group_count,
600 int64 feature_group_count, int64 batch_group_count,
610 int64 feature_group_count, int64 batch_group_count,
621 int64 feature_group_count, int64 batch_group_count,
[all …]
Dxla_builder.cc1375 int64 feature_group_count, int64 batch_group_count, in Conv() argument
1381 feature_group_count, batch_group_count, precision_config, in Conv()
1388 int64 feature_group_count, int64 batch_group_count, in ConvWithGeneralPadding() argument
1393 feature_group_count, batch_group_count, precision_config, in ConvWithGeneralPadding()
1400 int64 feature_group_count, int64 batch_group_count, in ConvWithGeneralDimensions() argument
1430 batch_group_count, precision_config, in ConvWithGeneralDimensions()
1439 int64 feature_group_count, int64 batch_group_count, in ConvGeneral() argument
1444 batch_group_count, precision_config, in ConvGeneral()
1453 int64 feature_group_count, int64 batch_group_count, in ConvGeneralDilated() argument
1477 *lhs_shape, *rhs_shape, feature_group_count, batch_group_count, in ConvGeneralDilated()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/translate/
Dlayouts_and_names.mlir9 batch_group_count = 1 : i64,
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_conv_rewriter.cc71 int64 num_groups = conv->batch_group_count(); in ConvertBatchGroupedToFeatureGroupedConvolution()
634 conv->batch_group_count(), conv->metadata()); in CreateCustomCallHelper()
678 if (conv->batch_group_count() > 1) { in CreateCustomCallHelper()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dconv_op_helpers.cc434 batch_group_count = in MakeXlaBackpropFilterConvOp() local
533 /*batch_group_count=*/batch_group_count, precision_config, in MakeXlaBackpropFilterConvOp()
540 /*batch_group_count=*/batch_group_count, precision_config); in MakeXlaBackpropFilterConvOp()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dconv_canonicalization.cc136 hlo->batch_group_count(), hlo->window(), new_dnums, in Run()
/external/tensorflow/tensorflow/compiler/mlir/xla/ir/
Dmlir_hlo_builder.h120 int64 feature_group_count, int64 batch_group_count,
Dmlir_hlo_builder.cc99 int64 feature_group_count, int64 batch_group_count, in ConvGeneralDilatedInternal() argument
115 builder_.getI64IntegerAttr(batch_group_count), config_attr); in ConvGeneralDilatedInternal()
/external/tensorflow/tensorflow/compiler/mlir/xla/tests/hlo_to_lhlo_with_xla/
Dhlo_text_to_lhlo_no_opt.hlotxt283 // CHECK-SAME: batch_group_count = 1 : i64
320 // CHECK-SAME: batch_group_count = 1 : i64
359 // CHECK-SAME: batch_group_count = 1 : i64
/external/tensorflow/tensorflow/compiler/mlir/xla/experimental/conv_emitter/
Dconv_emitter.cc593 if (conv->feature_group_count() != 1 || conv->batch_group_count() != 1) { in ConvIsImplemented()

12