[Mlir-commits] [mlir] [mlir][tensor] Move ExtractFromTensorGenerate pattern to ExtractOp (NFC) (PR #223201)

[llvmlistbot at llvm.org]

https://github.com/pranavsmoorthy created https://github.com/llvm/llvm-project/pull/223201

This patch moves `ExtractFromTensorGenerate` from `GenerateOp::getCanonicalizationPatterns` to `ExtractOp::getCanonicalizationPatterns`, resolving an explicit TODO in `TensorOps.cpp`.

Because `ExtractFromTensorGenerate` matches and rewrites `tensor::ExtractOp`, registering it under `ExtractOp` adheres to MLIR's convention of grouping canonicalization patterns by their root operation.

>From 3a1430d0e0955394c8167561f001db0147750484 Mon Sep 17 00:00:00 2001
From: snowpeakapps <pranavsmoorthy at gmail.com>
Date: Sat, 12 Sep 2026 22:42:17 -0500
Subject: [PATCH] [mlir][tensor] Move ExtractFromTensorGenerate pattern to
 ExtractOp (NFC)

This patch moves ExtractFromTensorGenerate from
GenerateOp::getCanonicalizationPatterns to
ExtractOp::getCanonicalizationPatterns, resolving an explicit TODO in
TensorOps.cpp.

Because ExtractFromTensorGenerate matches and rewrites tensor::ExtractOp,
registering it under ExtractOp adheres to MLIR canonicalization pattern
organization.
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 72 ++++++++++++------------
 1 file changed, 35 insertions(+), 37 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 43bb5cabb0f4f..6cbf250eb2e1c 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -1365,6 +1365,39 @@ struct ExtractFromCollapseShape : public OpRewritePattern<tensor::ExtractOp> {
   }
 };
 
+/// Canonicalizes the pattern of the form
+///
+/// %tensor = tensor.generate %x {
+///   ^bb0(%arg0: index):
+///   <computation>
+///   yield %1 : index
+/// } : tensor<?xindex>
+/// %extracted_element = tensor.extract %tensor[%c0] : tensor<?xi32>
+///
+/// to just <computation> with %arg0 replaced by %c0. We only do this if the
+/// tensor.generate operation has no side-effects.
+struct ExtractFromTensorGenerate : public OpRewritePattern<tensor::ExtractOp> {
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp extract,
+                                PatternRewriter &rewriter) const final {
+    auto tensorFromElements = extract.getTensor().getDefiningOp<GenerateOp>();
+    if (!tensorFromElements || !wouldOpBeTriviallyDead(tensorFromElements))
+      return failure();
+
+    IRMapping mapping;
+    Block *body = &tensorFromElements.getBody().front();
+    mapping.map(body->getArguments(), extract.getIndices());
+    for (auto &op : body->without_terminator())
+      rewriter.clone(op, mapping);
+
+    auto yield = cast<YieldOp>(body->getTerminator());
+
+    rewriter.replaceOp(extract, mapping.lookupOrDefault(yield.getValue()));
+    return success();
+  }
+};
+
 } // namespace
 
 void ExtractOp::getAsmResultNames(
@@ -1452,7 +1485,7 @@ OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
 
 void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<ExtractFromTensorCast>(context);
+  results.add<ExtractFromTensorCast, ExtractFromTensorGenerate>(context);
 }
 
 void mlir::tensor::populateFoldCollapseExtractPatterns(
@@ -1767,46 +1800,11 @@ struct StaticTensorGenerate : public OpRewritePattern<GenerateOp> {
     return success();
   }
 };
-
-/// Canonicalizes the pattern of the form
-///
-/// %tensor = tensor.generate %x {
-///   ^bb0(%arg0: index):
-///   <computation>
-///   yield %1 : index
-/// } : tensor<?xindex>
-/// %extracted_element = tensor.extract %tensor[%c0] : tensor<?xi32>
-///
-/// to just <computation> with %arg0 replaced by %c0. We only do this if the
-/// tensor.generate operation has no side-effects.
-struct ExtractFromTensorGenerate : public OpRewritePattern<tensor::ExtractOp> {
-  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tensor::ExtractOp extract,
-                                PatternRewriter &rewriter) const final {
-    auto tensorFromElements = extract.getTensor().getDefiningOp<GenerateOp>();
-    if (!tensorFromElements || !wouldOpBeTriviallyDead(tensorFromElements))
-      return failure();
-
-    IRMapping mapping;
-    Block *body = &tensorFromElements.getBody().front();
-    mapping.map(body->getArguments(), extract.getIndices());
-    for (auto &op : body->without_terminator())
-      rewriter.clone(op, mapping);
-
-    auto yield = cast<YieldOp>(body->getTerminator());
-
-    rewriter.replaceOp(extract, mapping.lookupOrDefault(yield.getValue()));
-    return success();
-  }
-};
-
 } // namespace
 
 void GenerateOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
-  // TODO: Move extract pattern to tensor::ExtractOp.
-  results.add<ExtractFromTensorGenerate, StaticTensorGenerate>(context);
+  results.add<StaticTensorGenerate>(context);
 }
 
 //===----------------------------------------------------------------------===//