[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, 14 insertions, 0 deletions

diff --git a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
index 5b1dbf439c..adaff90913 100644
--- a/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
+++ b/tensorflow/compiler/xla/service/dfs_hlo_visitor.h
@@ -85,6 +85,10 @@ class DfsHloVisitor {
   virtual Status HandleCopy(HloInstruction* copy) {
     return HandleElementwiseUnary(copy);
   }
+  virtual Status HandleComplex(HloInstruction* complex, HloInstruction* real,
+                               HloInstruction* imag) {
+    return HandleElementwiseBinary(complex);
+  }
   virtual Status HandleMultiply(HloInstruction* multiply, HloInstruction* lhs,
                                 HloInstruction* rhs) {
     return HandleElementwiseBinary(multiply);
@@ -122,6 +126,10 @@ class DfsHloVisitor {
   virtual Status HandleAbs(HloInstruction* abs, HloInstruction* operand) {
     return HandleElementwiseUnary(abs);
   }
+  virtual Status HandleAtan2(HloInstruction* atan2, HloInstruction* y,
+                             HloInstruction* x) {
+    return HandleElementwiseBinary(atan2);
+  }
   virtual Status HandleRound(HloInstruction* round) {
     return HandleElementwiseUnary(round);
   }
@@ -152,6 +160,12 @@ class DfsHloVisitor {
   virtual Status HandleTanh(HloInstruction* tanh, HloInstruction* operand) {
     return HandleElementwiseUnary(tanh);
   }
+  virtual Status HandleReal(HloInstruction* real, HloInstruction* operand) {
+    return HandleElementwiseUnary(real);
+  }
+  virtual Status HandleImag(HloInstruction* imag, HloInstruction* operand) {
+    return HandleElementwiseUnary(imag);
+  }
   virtual Status HandleIsFinite(HloInstruction* is_finite,
                                 HloInstruction* operand) {
     return HandleElementwiseUnary(is_finite);