| 196 | } |
| 197 | |
| 198 | StatusOr<mlir::Operation*> HloFunctionImporter::ImportInstruction( |
| 199 | HloInstruction* instruction, mlir::OpBuilder* func_builder) { |
| 200 | TF_ASSIGN_OR_RETURN(auto operands, GetOperands(instruction)); |
| 201 | TF_ASSIGN_OR_RETURN(auto result_type, ConvertShapeToType<RankedTensorType>( |
| 202 | instruction->shape(), *builder_)); |
| 203 | llvm::SmallVector<NamedAttribute, 10> attributes = {builder_->getNamedAttr( |
| 204 | "name", builder_->getStringAttr(instruction->name()))}; |
| 205 | mlir::Location loc = func_builder->getUnknownLoc(); |
| 206 | |
| 207 | switch (instruction->opcode()) { |
| 208 | case HloOpcode::kParameter: { |
| 209 | return nullptr; |
| 210 | } |
| 211 | case HloOpcode::kConstant: { |
| 212 | auto attr = CreateDenseAttrFromLiteral( |
| 213 | result_type.cast<mlir::TensorType>(), instruction->literal()); |
| 214 | if (!attr.ok()) return attr.status(); |
| 215 | mlir::Operation* new_operation = |
| 216 | func_builder->create<mlir::ConstantOp>(loc, attr.ValueOrDie()); |
| 217 | for (auto attr : attributes) { |
| 218 | new_operation->setAttr(attr.first, attr.second); |
| 219 | } |
| 220 | return new_operation; |
| 221 | } |
| 222 | case HloOpcode::kIota: { |
| 223 | return func_builder |
| 224 | ->create<mlir::xla_hlo::IotaOp>( |
| 225 | loc, result_type, |
| 226 | func_builder->getI64IntegerAttr( |
| 227 | Cast<HloIotaInstruction>(instruction)->iota_dimension())) |
| 228 | .getOperation(); |
| 229 | } |
| 230 | #define MakeAndReturn(mlir_op) \ |
| 231 | { \ |
| 232 | mlir::Operation* new_operation = \ |
| 233 | func_builder->create<mlir::xla_hlo::mlir_op>(loc, result_type, \ |
| 234 | operands, attributes); \ |
| 235 | return new_operation; \ |
| 236 | } |
| 237 | case HloOpcode::kBroadcast: { |
| 238 | // Note that the HLO broadcast is more powerful than the XLA broadcast op. |
| 239 | // BroadcastInDim offers a superset of the HLO op's functionality. |
| 240 | if (!instruction->dimensions().empty()) { |
| 241 | attributes.push_back(builder_->getNamedAttr( |
| 242 | "broadcast_dimensions", |
| 243 | ConvertDimensions(instruction->dimensions()))); |
| 244 | } |
| 245 | MakeAndReturn(BroadcastInDimOp); |
| 246 | } |
| 247 | #define MakeAndReturnBatchNormOp(batch_norm_op) \ |
| 248 | { \ |
| 249 | attributes.push_back(builder_->getNamedAttr( \ |
| 250 | "epsilon", builder_->getF32FloatAttr(instruction->epsilon()))); \ |
| 251 | attributes.push_back(builder_->getNamedAttr( \ |
| 252 | "feature_index", \ |
| 253 | builder_->getI64IntegerAttr(instruction->feature_index()))); \ |
| 254 | MakeAndReturn(batch_norm_op); \ |
| 255 | } |
nothing calls this directly
no test coverage detected