[XLA:CPU] [XLA:GPU] Adds compiler support for C64 primitive type, including relevant elementwise unary and binary op lowering for CPU and GPU.

We use a named LLVM struct "complex64", laid out the same as std::complex<float>. This named struct is accessed via the llvm::Module, which required changes to accessors of PrimitiveTypeToIrType & friends.

Ops that require atan2 (in particular, angle and log) are only supported on GPU at this point. LLVM lacks a CPU intrinsic for atan or atan2, whereas libdevice provides this for GPU.

PiperOrigin-RevId: 173676849

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/hlo_opcode.cc b/tensorflow/compiler/xla/service/hlo_opcode.cc
index e9000a8462..2f2263f70d 100644
--- a/tensorflow/compiler/xla/service/hlo_opcode.cc
+++ b/tensorflow/compiler/xla/service/hlo_opcode.cc
@@ -33,6 +33,8 @@ string HloOpcodeString(HloOpcode opcode) {
       return "abs";
     case HloOpcode::kAdd:
       return "add";
+    case HloOpcode::kAtan2:
+      return "atan2";
     case HloOpcode::kBatchNormTraining:
       return "batch-norm-training";
     case HloOpcode::kBatchNormInference:
@@ -47,6 +49,8 @@ string HloOpcodeString(HloOpcode opcode) {
       return "call";
     case HloOpcode::kClamp:
       return "clamp";
+    case HloOpcode::kComplex:
+      return "complex";
     case HloOpcode::kConcatenate:
       return "concatenate";
     case HloOpcode::kConstant:
@@ -87,6 +91,8 @@ string HloOpcodeString(HloOpcode opcode) {
       return "get-tuple-element";
     case HloOpcode::kGt:
       return "greater-than";
+    case HloOpcode::kImag:
+      return "imag";
     case HloOpcode::kIndex:
       return "index";
     case HloOpcode::kInfeed:
@@ -125,6 +131,8 @@ string HloOpcodeString(HloOpcode opcode) {
       return "parameter";
     case HloOpcode::kPower:
       return "power";
+    case HloOpcode::kReal:
+      return "real";
     case HloOpcode::kRecv:
       return "recv";
     case HloOpcode::kReduce: