[GmlSt] Do not transform ops that had been already tiles by another pass.

PiperOrigin-RevId: 495372204

Author: olegshyshkov

xla/mlir_hlo/gml_st/transforms/transform_map_for_cpu/transform_map_for_cpu.cc

@@ -47,6 +47,12 @@ struct TileMapPattern : public OpRewritePattern<linalg::MapOp> {
 
   LogicalResult matchAndRewrite(linalg::MapOp op,
                                 PatternRewriter &rewriter) const override {
+    if (hasLabel(op, kMapTransformedLabel)) return failure();
+
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(op->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          op, "has already been tiled by another pass.");
+
     auto fuseFilterFn = [](Operation *op) {
       return isa<linalg::BroadcastOp, linalg::MapOp>(op);
     };

xla/mlir_hlo/gml_st/transforms/transform_matmul_for_cpu/transform_matmul_for_cpu.cc

@@ -519,6 +519,10 @@ struct MatmulTransformPattern : public OpRewritePattern<linalg::MatmulOp> {
       return rewriter.notifyMatchFailure(matmulOp,
                                          "has already been transformed.");
 
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(matmulOp->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          matmulOp, "has already been tiled by another pass.");
+
     SmallVector<Operation *> fusionCluster = getFusionCluster(matmulOp);
 
     // First element of the cluster is always the root for tiling.

xla/mlir_hlo/gml_st/transforms/transform_matmul_for_triton/transform_matmul_for_triton.cc

@@ -76,6 +76,10 @@ struct MatmulTransformPattern : public OpRewritePattern<linalg::MatmulOp> {
       return rewriter.notifyMatchFailure(matmulOp,
                                          "has already been transformed.");
 
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(matmulOp->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          matmulOp, "has already been tiled by another pass.");
+
     // First level tiling: parallel dimensions.
     SmallVector<int64_t> parallelDimsTileSizes{lhsParallelDimTileSize,
                                                rhsParallelDimTileSize, 0};

xla/mlir_hlo/gml_st/transforms/transform_reduce_for_cpu/transform_reduce_for_cpu.cc

@@ -102,6 +102,10 @@ struct Reduce1DTransformPattern : public OpRewritePattern<linalg::ReduceOp> {
       return rewriter.notifyMatchFailure(reduceOp,
                                          "has already been transformed.");
 
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(reduceOp->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          reduceOp, "has already been tiled by another pass.");
+
     if (failed(validateOp(reduceOp, rewriter, /*expectedRank=*/1)))
       return failure();
 

xla/mlir_hlo/gml_st/transforms/transform_sort_for_cpu/transform_sort_for_cpu.cc

@@ -49,6 +49,10 @@ struct TileSortPattern : public OpRewritePattern<SortOp> {
                                 PatternRewriter &rewriter) const override {
     if (hasLabel(op, kSortTransformedLabel)) return failure();
 
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(op->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          op, "has already been tiled by another pass.");
+
     auto tilingResult =
         tile(options, rewriter, cast<TilingInterface>(op.getOperation()));
     if (failed(tilingResult)) return failure();

xla/mlir_hlo/gml_st/transforms/transform_transpose_for_cpu/transform_transpose_for_cpu.cc

@@ -50,6 +50,10 @@ struct TileTransposePattern : public OpRewritePattern<linalg::TransposeOp> {
                                 PatternRewriter &rewriter) const override {
     if (hasLabel(op, kTransposeTransformedLabel)) return failure();
 
+    if (isa<gml_st::ParallelOp, gml_st::ForOp>(op->getParentOp()))
+      return rewriter.notifyMatchFailure(
+          op, "has already been tiled by another pass.");
+
     auto tilingResult =
         tile(options, rewriter, cast<TilingInterface>(op.getOperation()));
     if (failed(tilingResult)) return failure();