[XLA:CPU] Include gml-tile-by-one pass in CPU pipeline before scalarization

Remaining tileable ops will be tiled to size 1 and scalarized.

PiperOrigin-RevId: 507850026

Author: frgossen

tensorflow/compiler/xla/mlir_hlo/gml_st/transforms/cpu_tiling/cpu_tiling_pipeline.cc

@@ -43,7 +43,11 @@ void addCPUTilingPipeline(OpPassManager& pm,
 
   pm.addNestedPass<FuncOp>(createComposeExtractInsertSlicePass());
   pm.addNestedPass<FuncOp>(createVectorizeForCPUPass());
+
+  // Tile remaining ops by size one and scalarize what we can.
+  pm.addNestedPass<FuncOp>(createTileByOnePass());
   pm.addNestedPass<FuncOp>(createScalarizationPass());
+
   pm.addNestedPass<FuncOp>(createRewriteVectorContractPass());
 }
 

tensorflow/compiler/xla/mlir_hlo/tests/Dialect/gml_st/cpu_tiling/map_matmul.mlir

@@ -54,8 +54,12 @@ func.func @map_matmul(%arg0: tensor<?x?xf32>,
 // CHECK:        scf.for
 // CHECK:          linalg.matmul
 // CHECK:          scf.yield
-// CHECK:        linalg.map
-// CHECK:        linalg.map
+// CHECK:        scf.for
+// CHECK:          scf.for
+// CHECK:            math.absf
+// CHECK:        scf.for
+// CHECK:          scf.for
+// CHECK:            arith.addf
 // CHECK:        gml_st.set_yield
 
 // CHECK:      gml_st.parallel
@@ -65,6 +69,10 @@ func.func @map_matmul(%arg0: tensor<?x?xf32>,
 // CHECK:        scf.for
 // CHECK:          linalg.matmul
 // CHECK:          scf.yield
-// CHECK:        linalg.map
-// CHECK:        linalg.map
+// CHECK:        scf.for
+// CHECK:          scf.for
+// CHECK:            math.absf
+// CHECK:        scf.for
+// CHECK:          scf.for
+// CHECK:            arith.addf
 // CHECK:        gml_st.set_yield