Method HandleReduce

tensorflow/compiler/xla/service/gpu/ir_emitter.cc:715–822

}

Status IrEmitter::HandleReduce(HloInstruction* instr) {
  const HloReduceInstruction* reduce = Cast<HloReduceInstruction>(instr);
  const Shape& out_shape = reduce->shape();
  bool returns_tuple = !out_shape.IsArray();
  int accumulators_count = 1;
  if (returns_tuple) {
    CHECK(out_shape.IsTuple());
    accumulators_count = out_shape.tuple_shapes_size();
  }

  auto arg = reduce->operand(0);
  absl::Span<const int64> dimensions(reduce->dimensions());
  HloComputation* function = reduce->to_apply();
  return EmitTargetElementLoop(
      *reduce,
      [=](const llvm_ir::IrArray::Index& index) -> StatusOr<llvm::Value*> {
        std::vector<llvm::Value*> accumulator_addrs;
        std::vector<llvm::Type*> accumulator_types;

        // Initialize accumulators with initial values.
        for (int i = 0; i < accumulators_count; i++) {
          auto init_value = reduce->init_values()[i];
          const Shape& element_shape =
              returns_tuple ? out_shape.tuple_shapes(i) : out_shape;
          PrimitiveType accumulator_type = element_shape.element_type();
          llvm::Type* accumulator_llvm_type =
              llvm_ir::PrimitiveTypeToIrType(accumulator_type, module_);
          llvm::AllocaInst* accumulator_addr = Alloca(accumulator_llvm_type);
          Store(Load(GetBasePointer(*init_value)), accumulator_addr);
          accumulator_addrs.push_back(accumulator_addr);
          accumulator_types.push_back(accumulator_llvm_type);
        }

        // The enclosing loops go over all the target elements. Now we have to
        // compute the actual target element. For this, we build a new loop nest
        // to iterate over all the reduction dimensions in the argument.
        // AddLoopsForShapeOnDimensions will return an Index where induction
        // Value*s are placed for each dimension in dimensions, and all the rest
        // are nullptrs.
        llvm_ir::ForLoopNest loops(IrName(reduce, "inner"), &b_);
        std::vector<llvm::Value*> input_multi_index =
            loops.AddLoopsForShapeOnDimensions(arg->shape(), dimensions,
                                               "reduction_dim");

        SetToFirstInsertPoint(loops.GetInnerLoopBodyBasicBlock(), &b_);

        // Build a full index for the input argument, using reduced_dims_index
        // as the base. In reduced_dims_index only the reduction dimensions are
        // filled in. We fill in the rest of the dimensions with induction
        // Value*s taken from 'index' which iterates over the target array.
        // See the high-level description in the XLA documentation for details.
        llvm_ir::IrArray::Index::const_iterator it = index.begin();

        for (auto& i : input_multi_index) {
          if (i == nullptr) {
            i = *it++;
          }
        }