[XLA:CPU] Enable F64 convolutions

There's no good reason for disabling it completely. Right now this falls back
to the slow generic implementation. We could make it use Eigen if we want to.

PiperOrigin-RevId: 235548224

tensorflow/compiler/xla/service/cpu/ir_emission_utils.cc

@@ -72,7 +72,8 @@ bool PotentiallyImplementedAsEigenConvolution(
   CHECK(
       ShapeUtil::SameElementTypeIgnoringFpPrecision(input_shape, kernel_shape));
   // TODO(b/65408531): Explore using Eigen dot for complex64 type.
-  if (ShapeUtil::ElementIsComplex(input_shape)) {
+  PrimitiveType primitive_type = input_shape.element_type();
+  if (primitive_type != F16 && primitive_type != F32) {
     return false;
   }
   if (window_util::HasWindowReversal(convolution.window())) {

tensorflow/compiler/xla/service/cpu/ir_emitter.cc

@@ -1090,7 +1090,7 @@ Status IrEmitter::HandleConvolution(HloInstruction* convolution) {
   auto rhs = convolution->operand(1);
   TF_RETURN_IF_ERROR(ElementTypesSameAndSupported(
       /*instruction=*/*convolution, /*operands=*/{lhs, rhs},
-      /*supported_types=*/{F16, F32, C64, C128}));
+      /*supported_types=*/{F16, F32, F64, C64, C128}));
 
   // TODO(tonywy): Add PotentiallyImplementedAsMKLCovolution to support
   // different data layouts.

tensorflow/compiler/xla/tests/convolution_test.cc

@@ -1945,7 +1945,7 @@ XLA_TEST_F(ConvolutionTest, ConvolveF32BackwardInputGroupedConvolution) {
 
 class ConvolutionHloTest : public HloTestBase {};
 
-XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64Forward)) {
+XLA_TEST_F(ConvolutionHloTest, ConvolveF64Forward) {
   constexpr char kHlo[] = R"(
 HloModule TestModule
 
@@ -1957,7 +1957,7 @@ ENTRY Test {
   EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001}));
 }
 
-XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF32ForwardReversed)) {
+XLA_TEST_F(ConvolutionHloTest, ConvolveF32ForwardReversed) {
   constexpr char kHlo[] = R"(
 HloModule TestModule
 
@@ -1969,7 +1969,7 @@ ENTRY Test {
   EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001}));
 }
 
-XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64BackwardFilter)) {
+XLA_TEST_F(ConvolutionHloTest, ConvolveF64BackwardFilter) {
   constexpr char kHlo[] = R"(
 HloModule TestModule
 
@@ -1981,7 +1981,7 @@ ENTRY Test {
   EXPECT_TRUE(RunAndCompare(kHlo, ErrorSpec{0.001}));
 }
 
-XLA_TEST_F(ConvolutionHloTest, DISABLED_ON_CPU(ConvolveF64BackwardInput)) {
+XLA_TEST_F(ConvolutionHloTest, ConvolveF64BackwardInput) {
   constexpr char kHlo[] = R"(
 HloModule TestModule