| 69 | } |
| 70 | |
| 71 | static void BuildOperator(const Operator& op, raw_ostream* output) { |
| 72 | auto& os = *output; |
| 73 | os << " auto& value_map = *lowering_context.values;\n" |
| 74 | << " auto result = xla_op.getResult();\n"; |
| 75 | |
| 76 | // Build a conversion for each of the arguments. |
| 77 | int operand_number = 0; |
| 78 | for (int index : llvm::seq<int>(0, op.getNumArgs())) { |
| 79 | auto arg = op.getArg(index); |
| 80 | |
| 81 | // Emit an argument for an operand. |
| 82 | if (auto* operand_cst = arg.dyn_cast<NamedTypeConstraint*>()) { |
| 83 | // Handle a non-variadic operand. |
| 84 | if (!operand_cst->isVariadic()) { |
| 85 | os << " auto xla_arg_" << index |
| 86 | << " = value_map[*xla_op.getODSOperands(" << operand_number++ |
| 87 | << ").begin()];\n"; |
| 88 | continue; |
| 89 | } |
| 90 | |
| 91 | // Otherwise, this is a varidiac operand list. |
| 92 | os << " std::vector<xla::XlaOp> xla_arg_" << index << ";\n" |
| 93 | << " for (auto operand : xla_op.getODSOperands(" << operand_number++ |
| 94 | << "))\n xla_arg_" << index |
| 95 | << ".push_back(value_map[operand]);\n"; |
| 96 | continue; |
| 97 | } |
| 98 | |
| 99 | // Otherwise, this is an attribute. |
| 100 | auto named_attr = arg.get<NamedAttribute*>(); |
| 101 | os << " auto xla_arg_" << index << " = " |
| 102 | << GetDefaultAttrExport(*named_attr) << "(xla_op." |
| 103 | << op.getArgName(index) << "());\n"; |
| 104 | } |
| 105 | |
| 106 | // Emit call to client API |
| 107 | os << " auto xla_result = xla::" << GetClientBuilder(op) << "("; |
| 108 | |
| 109 | // If all operands are variadic, then pass the builder explicitly to xla |
| 110 | // client API call |
| 111 | if (op.getNumOperands() == op.getNumVariadicOperands()) { |
| 112 | os << "lowering_context.builder"; |
| 113 | if (op.getNumArgs() != 0) os << ", "; |
| 114 | } |
| 115 | |
| 116 | // Emit each of the arguments. |
| 117 | interleaveComma(llvm::seq<int>(0, op.getNumArgs()), os, |
| 118 | [&](int i) { os << "Unwrap(xla_arg_" << i << ')'; }); |
| 119 | os << ");\n"; |
| 120 | |
| 121 | os << " value_map[result] = xla_result;\n"; |
| 122 | os << " return mlir::success();\n"; |
| 123 | } |
| 124 | |
| 125 | // The function below has a non-constant reference as that is required by LLVM's |
| 126 | // TableGenMain. |
no test coverage detected