Method Emit

tensorflow/compiler/xla/service/cpu/dot_op_emitter.cc:374–424

}

Status DotOpEmitter::Emit() {
  // The dot operation performs a sum of products over dimension 0 of the left
  // hand side operand and dimension 1 of the right hand side operand.
  //
  // Let the shapes of lhs and rhs be defined as below:
  //
  //   lhs = [L{n-1} x L{n-2} x ... L{0}]
  //   rhs = [R{m-1} x R{m-2} x ... R{0}]
  //
  // The sum-of-products dimension in the lhs has size L{0} and the dimension in
  // the rhs has size R{1}. Necessarily, then:
  //
  //   L{0} == R{1}
  //
  // The output of the operation has the following shape:
  //
  //   output = [L{n-1} x L{n-2} x ... L{1} x R{m-1} x R{m-2} x ... R{2} x R{0}]
  //
  // To perform the operation we construct a loop nest with one for-loop for
  // each dimension of the output. Inside this loop nest is another for-loop
  // which performs the sum-of-products (the reduction loop) before storing
  // the result in the output buffer.

  const Shape& lhs_shape = lhs_array_.GetShape();
  const Shape& rhs_shape = rhs_array_.GetShape();

  if (ShapeUtil::IsScalar(lhs_shape) || ShapeUtil::IsScalar(rhs_shape)) {
    // If the operands are scalar, don't emit any loops.
    TF_RET_CHECK(ShapeUtil::IsScalar(lhs_shape) &&
                 ShapeUtil::IsScalar(rhs_shape));
    return EmitScalarDot();
  }

  switch (GetDotImplementationStrategy(hlo_module_config_, dot_info_,
                                       target_machine_features_)) {
    case DotImplementationStrategy::kNaiveLlvmIr:
      EmitNaiveLlvmIrGemm();
      return Status::OK();

    case DotImplementationStrategy::kTiledLlvmIrGemv:
      EmitTiledLlvmIrGemv();
      return Status::OK();

    case DotImplementationStrategy::kTiledLlvmIrGemm:
      EmitTiledLlvmIrGemm();
      return Status::OK();

    case DotImplementationStrategy::kEigen:
      return EmitCallToRuntime();
  }
}

void DotOpEmitter::EmitNaiveLlvmIrGemm() {
  CHECK_EQ(addend_array_, nullptr);