Speedup RewriteStackPtr by avoiding one mod.walk call (#4803)

Timings are obtained from compiling flex attn backward kernel:

`1.0694 (  6.3%)    1.0694 ( 10.7%)  TritonIntelGPURewriteStackPtr`
vs
 `0.8938 (  5.2%)    0.8938 (  8.8%)  TritonIntelGPURewriteStackPtr`


Relates to #4062

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

third_party/intel/lib/TritonIntelGPUTransforms/RewriteStackPtr.cpp

@@ -38,11 +38,9 @@ struct TritonIntelGPURewriteStackPtrPass
 
     // 1: Process function arguments for root functions
     if (!usePoison) {
-      mod.walk([&](FunctionOpInterface funcOp) {
-        if (allocation.isRoot(funcOp)) {
-          insertFuncArguments(funcOp, ptrTy);
-        }
-      });
+      for (auto &root : allocation.getRoots()) {
+        insertFuncArguments(root, ptrTy);
+      }
     }
 
     // 2: Collect all AddressOfOp that need updating