| 2382 | } |
| 2383 | |
| 2384 | Status HloEvaluator::HandleReduce(HloInstruction* instr) { |
| 2385 | HloReduceInstruction* reduce = Cast<HloReduceInstruction>(instr); |
| 2386 | int64 num_args = reduce->inputs().size(); |
| 2387 | absl::Span<const int64> dimensions_to_reduce(reduce->dimensions()); |
| 2388 | HloComputation* function = reduce->to_apply(); |
| 2389 | |
| 2390 | absl::InlinedVector<const Shape*, 1> operand_shapes; |
| 2391 | for (const HloInstruction* operand : reduce->operands()) { |
| 2392 | operand_shapes.push_back(&operand->shape()); |
| 2393 | } |
| 2394 | TF_ASSIGN_OR_RETURN(auto inferred_return_shape, |
| 2395 | ShapeInference::InferReduceShape( |
| 2396 | operand_shapes, dimensions_to_reduce, |
| 2397 | /*to_apply=*/function->ComputeProgramShape())); |
| 2398 | TF_RET_CHECK(ShapeUtil::CompatibleIgnoringFpPrecision(reduce->shape(), |
| 2399 | inferred_return_shape)) |
| 2400 | << "return shape is set to: " << ShapeUtil::HumanString(reduce->shape()) |
| 2401 | << " but is inferred to be: " |
| 2402 | << ShapeUtil::HumanString(inferred_return_shape); |
| 2403 | |
| 2404 | absl::InlinedVector<const Literal*, 1> input_args(num_args); |
| 2405 | absl::InlinedVector<const Literal*, 1> init_values(num_args); |
| 2406 | for (int64 i = 0; i < num_args; ++i) { |
| 2407 | input_args[i] = &GetEvaluatedLiteralFor(reduce->inputs()[i]); |
| 2408 | VLOG(3) << "HandleReduce arg_literal: " << input_args[i]->ToString(); |
| 2409 | init_values[i] = &GetEvaluatedLiteralFor(reduce->init_values()[i]); |
| 2410 | VLOG(3) << "HandleReduce init_literal: " << init_values[i]->ToString(); |
| 2411 | TF_RET_CHECK(ShapeUtil::IsScalar(init_values[i]->shape())); |
| 2412 | } |
| 2413 | |
| 2414 | // All args and results have the same dimensions, so pick an arbitrary one. |
| 2415 | const Shape& arg_shape = input_args[0]->shape(); |
| 2416 | const Shape& out_shape = inferred_return_shape; |
| 2417 | bool is_tuple = out_shape.IsTuple(); |
| 2418 | const Shape& output_shape = inferred_return_shape.IsTuple() |
| 2419 | ? inferred_return_shape.tuple_shapes(0) |
| 2420 | : inferred_return_shape; |
| 2421 | |
| 2422 | absl::Span<const int64> arg_dimensions = AsInt64Slice(arg_shape.dimensions()); |
| 2423 | |
| 2424 | // All increments are set to 0. |
| 2425 | std::vector<int64> arg_dim_steps(arg_dimensions.size()); |
| 2426 | |
| 2427 | // All counts are set to 0. |
| 2428 | std::vector<int64> arg_dim_counts(arg_dimensions.size()); |
| 2429 | |
| 2430 | // Set steps and counts for reduced dimensions. |
| 2431 | // This avoids iterating over non-reduced dimensions, as their step |
| 2432 | // and count is set to zero. |
| 2433 | for (const int64 dim : dimensions_to_reduce) { |
| 2434 | arg_dim_steps[dim] = 1; |
| 2435 | arg_dim_counts[dim] = arg_dimensions[dim]; |
| 2436 | } |
| 2437 | |
| 2438 | // Map each dimension in the result to a dimension in arg that isn't |
| 2439 | // being reduced. |
| 2440 | std::vector<int64> result_to_arg_index; |
| 2441 | for (int64 i = 0; i < arg_dimensions.size(); ++i) { |
nothing calls this directly
no test coverage detected