[Mlir-commits] [mlir] [MLIR][EmitC] Rewrite globals before moving function body (PR #217474)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/217474

WalkPatternRewriter only permits erasing matched operations or their descendants. Replace get_global operations while they still belong to the matched function before transferring its body into the new class method.

Assisted-by: Codex

>From 7d52afec5d79e55f53febb1311e789c4ecacbe9f Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 19 Aug 2026 07:52:36 -0700
Subject: [PATCH] [MLIR][EmitC] Rewrite globals before moving function body

WalkPatternRewriter only permits erasing matched operations or their
descendants. Replace get_global operations while they still belong to the
matched function before transferring its body into the new class method.

Assisted-by: Codex
---
 .../EmitC/Transforms/WrapFuncInClass.cpp        | 17 +++++++++--------
 mlir/test/Dialect/EmitC/wrap-func-in-class.mlir |  2 --
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp b/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp
index 4754ba19f3f78..2b24dcac32b8b 100644
--- a/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp
+++ b/mlir/lib/Dialect/EmitC/Transforms/WrapFuncInClass.cpp
@@ -114,6 +114,15 @@ class WrapFuncInClass : public OpRewritePattern<FuncOp> {
     Location loc = funcOp.getLoc();
     FuncOp newFuncOp = FuncOp::create(rewriter, loc, (funcName), funcType);
 
+    // Rewrite globals while they are still descendants of the matched op.
+    funcOp.walk([&](GetGlobalOp getGlobalOp) {
+      rewriter.setInsertionPoint(getGlobalOp);
+      GetFieldOp getFieldOp =
+          GetFieldOp::create(rewriter, getGlobalOp.getLoc(),
+                             getGlobalOp.getType(), getGlobalOp.getNameAttr());
+      rewriter.replaceOp(getGlobalOp, getFieldOp);
+    });
+
     rewriter.createBlock(&newFuncOp.getBody());
     newFuncOp.getBody().takeBody(funcOp.getBody());
 
@@ -135,14 +144,6 @@ class WrapFuncInClass : public OpRewritePattern<FuncOp> {
     if (failed(newFuncOp.eraseArguments(argsToErase)))
       newFuncOp->emitOpError("failed to erase all arguments using BitVector");
 
-    newFuncOp.walk([&](GetGlobalOp getGlobalOp) {
-      rewriter.setInsertionPoint(getGlobalOp);
-      GetFieldOp getFieldOp =
-          GetFieldOp::create(rewriter, getGlobalOp.getLoc(),
-                             getGlobalOp.getType(), getGlobalOp.getNameAttr());
-      rewriter.replaceOp(getGlobalOp, getFieldOp);
-    });
-
     rewriter.replaceOp(funcOp, newClassOp);
     return success();
   }
diff --git a/mlir/test/Dialect/EmitC/wrap-func-in-class.mlir b/mlir/test/Dialect/EmitC/wrap-func-in-class.mlir
index 9c7f042466c28..91010689e1e65 100644
--- a/mlir/test/Dialect/EmitC/wrap-func-in-class.mlir
+++ b/mlir/test/Dialect/EmitC/wrap-func-in-class.mlir
@@ -2,8 +2,6 @@
 // RUN: mlir-opt %s -wrap-emitc-func-in-class=func-name=execute -split-input-file | FileCheck %s --check-prefixes=EXECUTE
 // RUN: mlir-opt %s -wrap-emitc-func-in-class=class-name-format=Custom_{} -split-input-file | FileCheck %s --check-prefixes=CLASS-NAME-FORMAT
 
-// XFAIL: mlir-expensive-checks
-
 emitc.func @foo(%arg0 : !emitc.array<1xf32>) {
   emitc.call_opaque "bar" (%arg0) : (!emitc.array<1xf32>) -> ()
   emitc.return