Move vectorize_copy.cc and copy_removal.cc passes out of gml_st.

They don't use gml_st dialect.

PiperOrigin-RevId: 585099478

Author: pifon2a

BUILD

@@ -1224,11 +1224,13 @@ cc_library(
         "transforms/detensorize_scf_ops.cc",
         "transforms/generic_host_to_llvm.cc",
         "transforms/lower_index_cast_pass.cc",
+        "transforms/naive_copy_removal.cc",
         "transforms/propagate_static_shapes_to_kernel.cc",
         "transforms/test_hlo_transform_dialect_interpreter.cc",
         "transforms/tile_loops_pass.cc",
         "transforms/unbufferize_pass.cc",
         "transforms/unroll_loops.cc",
+        "transforms/vectorize_copy.cc",
     ],
     hdrs = [
         "transforms/passes.h",
@@ -1286,6 +1288,7 @@ cc_library(
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:MemRefToLLVM",
         "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:MemRefUtils",
         "@llvm-project//mlir:NVVMDialect",
         "@llvm-project//mlir:PDLDialect",
         "@llvm-project//mlir:Pass",
@@ -1481,7 +1484,6 @@ cc_library(
         "gml_st/transforms/collapse_shape/collapse_shape.cc",
         "gml_st/transforms/collect_stats/collect_stats.cc",
         "gml_st/transforms/compose_extract_insert_slice/compose_extract_insert_slice.cc",
-        "gml_st/transforms/copy_removal/copy_removal.cc",
         "gml_st/transforms/cpu_tiling/cpu_tiling_pipeline.cc",
         "gml_st/transforms/cpu_tiling/fusion_outlining.cc",
         "gml_st/transforms/cpu_tiling/fusion_planning_for_cpu.cc",
@@ -1505,7 +1507,6 @@ cc_library(
         "gml_st/transforms/transforms.h",
         "gml_st/transforms/vectorization/lower_vectors.cc",
         "gml_st/transforms/vectorization/vectorization.cc",
-        "gml_st/transforms/vectorization/vectorize_copy.cc",
         "gml_st/transforms/vectorization/vectorize_for_cpu.cc",
         "gml_st/utils/linalg_utils.cc",
         "gml_st/utils/tensor_utils.cc",

gml_st/transforms/CMakeLists.txt

@@ -32,7 +32,6 @@ add_mlir_library(GmlStPasses
   collapse_shape/collapse_shape.cc
   collect_stats/collect_stats.cc
   compose_extract_insert_slice/compose_extract_insert_slice.cc
-  copy_removal/copy_removal.cc
   cpu_tiling/cpu_tiling_pipeline.cc
   cpu_tiling/fusion_outlining.cc
   cpu_tiling/fusion_planning_for_cpu.cc
@@ -54,7 +53,6 @@ add_mlir_library(GmlStPasses
   tiling_softmax/tiling_softmax.cc
   vectorization/lower_vectors.cc
   vectorization/vectorization.cc
-  vectorization/vectorize_copy.cc
   vectorization/vectorize_for_cpu.cc
 
   DEPENDS

gml_st/transforms/passes.h

@@ -61,13 +61,6 @@ createComposeExtractInsertSlicePass();
 std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeForCPUPass(
     int64_t numElementsThreshold = 1024);
 
-/// Pass to vectorize `memref.copy`.
-std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeCopyPass(
-    int64_t numElementsThreshold = 8);
-
-/// Pass to remove redundant `memref.copy` ops.
-std::unique_ptr<OperationPass<func::FuncOp>> createNaiveCopyRemovalPass();
-
 /// Pass to gradually lower vector ops to SCF.
 std::unique_ptr<OperationPass<func::FuncOp>> createLowerVectorsPass(
     bool enableAVX2 = true, bool flatten = false);

gml_st/transforms/passes.td

@@ -70,27 +70,6 @@ def VectorizeForCPUPass : Pass<"vectorize-for-cpu", "mlir::func::FuncOp"> {
   ];
 }
 
-def VectorizeCopyPass : Pass<"vectorize-copy", "mlir::func::FuncOp"> {
-  let summary = "Pass to vectorize `memref.copy`.";
-  let constructor = "::mlir::gml_st::createVectorizeCopyPass()";
-  let dependentDialects = [
-    "scf::SCFDialect",
-    "vector::VectorDialect",
-  ];
-  let options = [
-    Option<"numElementsThreshold", "num-elements-threshold", "int64_t",
-           /*default=*/"8",
-           "Max number of elements in src and dst memref for a copy to be "
-           "vectorized.">,
-  ];
-}
-
-def NaiveCopyRemovalPass : Pass<"naive-copy-removal", "mlir::func::FuncOp"> {
-  let summary = "Pass to remove redundant `memref.copy` ops.";
-  let constructor = "::mlir::gml_st::createNaiveCopyRemovalPass()";
-  let dependentDialects = ["::mlir::memref::MemRefDialect"];
-}
-
 def LowerVectorsPass : Pass<"lower-vectors", "mlir::func::FuncOp"> {
   let summary = "Pass to lower vector operations progressively.";
   let constructor = "::mlir::gml_st::createLowerVectorsPass()";

tests/Dialect/gml_st/simplify_dead_copy.mlir renamed to tests/naive_copy_removal.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-hlo-opt %s --vectorize-copy="num-elements-threshold=8" --split-input-file | FileCheck %s
+// RUN: mlir-hlo-opt %s --vectorize-copy --split-input-file | FileCheck %s
 
 func.func @vectorize_copy(%arg: memref<2x2xf32>) -> memref<2x2xf32> {
   %subview = memref.subview %arg[0, 0] [2, 2] [1, 1] : memref<2x2xf32> to memref<2x2xf32, strided<[16, 1]>>

tests/Dialect/gml_st/vectorize_copy.mlir renamed to tests/vectorize_copy.mlir

@@ -31,9 +31,11 @@ add_mlir_library(MLIRBufferTransforms
   detensorize_scf_ops.cc
   generic_host_to_llvm.cc
   lower_index_cast_pass.cc
+  naive_copy_removal.cc
   propagate_static_shapes_to_kernel.cc
   test_hlo_transform_dialect_interpreter.cc
   tile_loops_pass.cc
+  vectorize_copy.cc
   unbufferize_pass.cc
   unroll_loops.cc
 

transforms/CMakeLists.txt

@@ -1,4 +1,4 @@
-/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -16,16 +16,18 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
-#include "gml_st/transforms/passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "transforms/passes.h"
 
-namespace mlir::gml_st {
+namespace mlir {
 namespace {
 
 #define GEN_PASS_DEF_NAIVECOPYREMOVALPASS
-#include "gml_st/transforms/passes.h.inc"
+#include "transforms/passes.h.inc"
 
 /// Remove memref::CopyOp whose target (can be either a memref::SubViewOp or
 /// memref::AllocOp) has no other users.
@@ -88,4 +90,4 @@ std::unique_ptr<OperationPass<func::FuncOp>> createNaiveCopyRemovalPass() {
   return std::make_unique<NaiveCopyRemovalPass>();
 }
 
-}  // namespace mlir::gml_st
+}  // namespace mlir

gml_st/transforms/copy_removal/copy_removal.cc renamed to transforms/naive_copy_removal.cc

@@ -49,6 +49,7 @@ using BufferizePatternsCallback = std::function<void(
 #define GEN_PASS_DECL_PROPAGATESTATICSHAPESTOKERNELPASS
 #define GEN_PASS_DECL_TILELOOPSPASS
 #define GEN_PASS_DECL_GENERICHOSTTOLLVMPASS
+#define GEN_PASS_DECL_VECTORIZECOPYPASS
 #include "transforms/passes.h.inc"
 
 /// Creates a pass that merges smaller buffer into bigger buffer to optimize
@@ -99,6 +100,12 @@ std::unique_ptr<OperationPass<func::FuncOp>> createTileLoopsPass(
 // and scf.if.
 std::unique_ptr<OperationPass<func::FuncOp>> createDetensorizeScfOpsPass();
 
+/// Pass to remove redundant `memref.copy` ops.
+std::unique_ptr<OperationPass<func::FuncOp>> createNaiveCopyRemovalPass();
+
+/// Pass to vectorize `memref.copy`.
+std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeCopyPass();
+
 /// Registers the test pass for erasing transform dialect ops.
 void registerTestHloTransformDialectEraseSchedulePass();
 

transforms/passes.h

@@ -172,4 +172,16 @@ def AllocToArgPass : Pass<"alloc-to-arg", "mlir::func::FuncOp"> {
   let constructor = "hlo::createAllocToArgPass()";
 }
 
+def NaiveCopyRemovalPass : Pass<"naive-copy-removal", "mlir::func::FuncOp"> {
+  let summary = "Pass to remove redundant `memref.copy` ops.";
+  let constructor = "createNaiveCopyRemovalPass()";
+  let dependentDialects = ["memref::MemRefDialect"];
+}
+
+def VectorizeCopyPass : Pass<"vectorize-copy", "mlir::func::FuncOp"> {
+  let summary = "Pass to vectorize `memref.copy`.";
+  let constructor = "createVectorizeCopyPass()";
+  let dependentDialects = ["scf::SCFDialect", "vector::VectorDialect"];
+}
+
 #endif // TENSORFLOW_COMPILER_MLIR_HLO_TRANSFORMS_PASSES

transforms/passes.td

@@ -14,24 +14,21 @@ limitations under the License.
 ==============================================================================*/
 
 #include <algorithm>
-#include <limits>
 #include <memory>
-#include <optional>
 #include <utility>
 
-#include "gml_st/transforms/passes.h"
-#include "gml_st/transforms/vectorization/vectorization.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
-namespace gml_st {
 namespace {
 
 #define GEN_PASS_DEF_VECTORIZECOPYPASS
-#include "gml_st/transforms/passes.h.inc"
+#include "transforms/passes.h.inc"
 
 /// Transforms a big non-contiguous `memref.copy` into a loop over smaller
 /// copies that are either contiguous or can be vectorized.
@@ -217,7 +214,7 @@ struct VectorizeCopyPass
 
     RewritePatternSet patterns(ctx);
     patterns.add<TileCopyPattern, CopyVectorizationPattern>(
-        ctx, numElementsThreshold);
+        ctx, /*numElementsThreshold = */ 8);
     if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns)))) {
       return signalPassFailure();
     }
@@ -226,12 +223,8 @@ struct VectorizeCopyPass
 
 }  // namespace
 
-std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeCopyPass(
-    int64_t numElementsThreshold) {
-  VectorizeCopyPassOptions opts;
-  opts.numElementsThreshold = numElementsThreshold;
-  return std::make_unique<VectorizeCopyPass>(opts);
+std::unique_ptr<OperationPass<func::FuncOp>> createVectorizeCopyPass() {
+  return std::make_unique<VectorizeCopyPass>();
 }
 
-}  // namespace gml_st
 }  // namespace mlir