Home
last modified time | relevance | path

Searched refs:reduction_dimensions (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dtree_reduction_rewriter.cc72 ReductionDimensions reduction_dimensions = in RewriteReduction() local
95 return RewriteBatchDimensionLargerThanTile(hlo, reduction_dimensions, in RewriteReduction()
99 bool is_row_reduction = reduction_dimensions.is_row_reduction; in RewriteReduction()
201 HloInstruction *hlo, const ReductionDimensions &reduction_dimensions, in RewriteBatchDimensionLargerThanTile() argument
207 CHECK(reduction_dimensions.is_row_reduction); in RewriteBatchDimensionLargerThanTile()
Dir_emission_utils.cc134 const ReductionDimensions& reduction_dimensions, in GetReductionTiling() argument
137 if (reduction_dimensions.is_row_reduction) { in GetReductionTiling()
138 int64 tile_z = std::min(reduction_dimensions.dimensions[0], int64{8}); in GetReductionTiling()
139 if (reduction_dimensions.dimensions[1] == 1) { in GetReductionTiling()
140 CHECK_EQ(reduction_dimensions.dimensions[0], 1); in GetReductionTiling()
143 if (reduction_dimensions.dimensions[2] % (kWarpSize * kWarpSize * 64) == in GetReductionTiling()
284 ReductionDimensions reduction_dimensions = in IsReductionFromOrToContiguousDimensions() local
287 if (reduction_dimensions.is_row_reduction) { in IsReductionFromOrToContiguousDimensions()
291 return reduction_dimensions.dimensions[2] >= kWarpSize; in IsReductionFromOrToContiguousDimensions()
297 return reduction_dimensions.dimensions[1] >= kWarpSize; in IsReductionFromOrToContiguousDimensions()
[all …]
Dir_emission_utils.h202 const ReductionDimensions& reduction_dimensions,
Dir_emitter_unnested.cc5297 ReductionDimensions reduction_dimensions = in ComputeReductionCodegenInfo() local
5299 VLOG(10) << "is_row_reduction " << reduction_dimensions.is_row_reduction in ComputeReductionCodegenInfo()
5300 << " " << reduction_dimensions.dimensions[0] << " " in ComputeReductionCodegenInfo()
5301 << reduction_dimensions.dimensions[1] << " " in ComputeReductionCodegenInfo()
5302 << reduction_dimensions.dimensions[2]; in ComputeReductionCodegenInfo()
5317 GetReductionTiling(reduction_dimensions, smallest_input_dtype_bits, in ComputeReductionCodegenInfo()
5320 int64 num_threads_y = reduction_dimensions.is_row_reduction ? 1 : kWarpSize; in ComputeReductionCodegenInfo()
5322 if (reduction_dimensions.is_row_reduction) { in ComputeReductionCodegenInfo()
5336 RoundUpToNearest(CeilOfRatio(reduction_dimensions.dimensions[2], in ComputeReductionCodegenInfo()
5343 bool tile_fit = reduction_dimensions.dimensions[kDimX] % in ComputeReductionCodegenInfo()
[all …]
/external/tensorflow/tensorflow/compiler/mlir/xla/transforms/
Dlegalize_tf.cc3875 DenseIntElementsAttr reduction_dimensions = in matchAndRewrite() local
3880 llvm::ArrayRef<Value>(init_values), reduction_dimensions); in matchAndRewrite()