[xla-next][mlir][sparse] add mhlo sparsity rewriting to pipeline

PiperOrigin-RevId: 511567163
vwbaker · Feb 22, 2023 · b4d47d3 · b4d47d3
1 parent 0ae88ee
commit b4d47d3
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/xla/service/cpu/hlo_xla_runtime_pipeline.cc b/xla/service/cpu/hlo_xla_runtime_pipeline.cc
@@ -118,6 +118,12 @@ static Status CreateHloXlaPipeline(
   pm.addPass(mlir::createCSEPass());
   pm.addPass(mlir::createCanonicalizerPass());
 
+  // Some early sparse rewriting rules.
+  if (options.sparse_bufferization) {
+    pm.addNestedPass<mlir::func::FuncOp>(
+        mlir::mhlo::createSparseRewritingPass());
+  }
+
   // Transform HLO operations to Linalg.
   pm.addNestedPass<mlir::func::FuncOp>(mlir::mhlo::createLegalizeSortPass());
   pm.addNestedPass<mlir::func::FuncOp>(