Also run canonicalization before ShapeToDescriptorsPass

This removes ConstWitnessOps which otherwise would block the
lowering of tf.IsInf.

PiperOrigin-RevId: 339864990
Change-Id: I9296481df4b4c029621e095ea7a7e13afa7b22ed

tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc

@@ -87,12 +87,14 @@ Status LowerTFtoGPU(mlir::ModuleOp module, bool gpu_binary_only,
     pm.addPass(::mlir::createBufferHoistingPass());
     pm.addPass(::mlir::createBufferDeallocationPass());
     pm.addNestedPass<mlir::FuncOp>(mlir::createCopyRemovalPass());
+    pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(mlir::kernel_gen::transforms::CreateShapeToDescriptorsPass());
   } else {
     pm.addPass(mlir::mhlo::createLegalizeTFPass(
         /*allow_partial_conversion=*/false, /*legalize_chlo=*/false));
     pm.addPass(mlir::createTransformUnrankedHloPass());
     pm.addPass(mlir::mhlo::createChloLegalizeToHloPass());
+    pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(mlir::kernel_gen::transforms::CreateShapeToDescriptorsPass());
     // Clean up the IR created above. In particular, operations on descriptors
     // are simplified here.

tensorflow/compiler/mlir/tools/kernel_gen/tests/isinf.mlir

+// RUN: tf-opt %s --test-tf-lower-tf --xla-legalize-tf | mlir-hlo-opt --transform-unranked-hlo | kernel-gen-opt -allow-unregistered-dialect --canonicalize --shape-to-descriptors --canonicalize --bufferize | FileCheck %s
+
+// Test whether all shape computations required for isinf can be lowered to
+// the standard dialect, scf and descriptors.
+// CHECK-LABEL: @isinf
+func @isinf(%arg0: tensor<?xf32>) -> tensor<?xi1> {
+  %0 = "tf.IsInf"(%arg0) : (tensor<?xf32>) -> tensor<?xi1>
+  return %0 : tensor<?xi1>
+}