Make sure memory copies get eliminated after bufferization

Bufferization of gml_st.set_yield will introduce a copy from source
into output. Add a canonicalizer pass after it to identify copies into
itself, and eliminate those.

PiperOrigin-RevId: 484208007

Author: gflegar

tensorflow/compiler/xla/mlir_hlo/lib/Transforms/hlo_to_gpu_pipeline.cc

@@ -62,9 +62,6 @@ void mlir::createHloToGpuPipeline(OpPassManager& pm,
     pm.addNestedPass<FuncOp>(gml_st::createTilingSoftmaxPass(
         /*distribute=*/true, options.warpTileDim, "warp"));
 
-    // Collapse all materialize ops.
-    pm.addPass(gml_st::createCollapseMaterializeOpsPass());
-
     // GPU-specific tiling for ops on the warp level.
     pm.addNestedPass<FuncOp>(gml_st::createTilingGPUWarpPass());
     pm.addNestedPass<FuncOp>(createScalarizationPass());
@@ -106,6 +103,8 @@ void mlir::createHloToGpuPipeline(OpPassManager& pm,
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
   pm.addNestedPass<FuncOp>(bufferization::createBufferDeallocationPass());
+  // Canonicalize away memory copies into itself
+  pm.addPass(createCanonicalizerPass());
 
   // Linalg + GmlSt -> GPU
   pm.addNestedPass<FuncOp>(createGmlStToGpuPass());