Rebase patch: Modify insert-gpu-allocs pass to also consider xegpu.store/load

Author: Mei, Yijie

lib/Transforms/InsertGPUAllocs.cpp

@@ -28,6 +28,7 @@
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/GPU/Transforms/Passes.h>
 #include <mlir/Dialect/MemRef/IR/MemRef.h>
+#include <mlir/Dialect/XeGPU/IR/XeGPU.h>
 #include <mlir/Dialect/SCF/IR/SCF.h>
 #include <mlir/Pass/Pass.h>
 
@@ -158,6 +159,9 @@ class InsertGPUAllocsPass final
       } else if (auto init_tile =
                      mlir::dyn_cast<imex::xetile::InitTileOp>(op)) {
         return {{init_tile.getSource()}};
+      } else if (auto init_xedesc =
+                     mlir::dyn_cast<mlir::xegpu::CreateNdDescOp>(op)) {
+        return {{init_xedesc.getSource()}};
       } else {
         op->emitError("Uhhandled mem op in gpu region");
         return std::nullopt;
@@ -187,6 +191,9 @@ class InsertGPUAllocsPass final
         // Only handle the case where the tile source is a memref
         return init_tile.isSourceMemRef();
       }
+      if (auto init_xedesc = mlir::dyn_cast<mlir::xegpu::CreateNdDescOp>(op)) {
+        return true;
+      }
       return false;
     };
 
@@ -259,6 +266,36 @@ class InsertGPUAllocsPass final
       return;
     }
 
+    // walk over the users and find xegpu.load/store ops
+    std::function<void(mlir::Operation*, bool, AccessType&)> findXeGPULoadStore;
+    findXeGPULoadStore = [&](mlir::Operation *use, bool onDevice, AccessType& ret) {
+      if (auto tile_update = mlir::dyn_cast<mlir::xegpu::UpdateNdOffsetOp>(use)) {
+        auto res = tile_update->getResult(0);
+        for (auto u : res.getUsers()) {
+          findXeGPULoadStore(u, onDevice, ret);
+        }
+      }
+      if (auto tile_for = mlir::dyn_cast<::mlir::scf::ForOp>(use)) {
+        for (size_t idx=0; idx<tile_for.getInits().size(); idx++) {
+          auto a = tile_for.getRegionIterArg(idx);
+          for (auto u : a.getUsers()) {
+            findXeGPULoadStore(u, onDevice, ret);
+          }
+        }
+      }
+      if (auto tile_load =
+              mlir::dyn_cast<mlir::xegpu::LoadNdOp>(use)) {
+        (onDevice ? ret.deviceRead : ret.hostRead) = true;
+      }
+      else if (auto tile_prefetch =
+                    mlir::dyn_cast<mlir::xegpu::PrefetchNdOp>(use)) {
+        (onDevice ? ret.deviceRead : ret.hostRead) = true;
+      } else if (auto tile_store =
+                    mlir::dyn_cast<mlir::xegpu::StoreNdOp>(use)) {
+        (onDevice ? ret.deviceWrite : ret.hostWrite) = true;
+      }
+    };
+
     // Checks the access type of the OP under consideration.
     auto getAccessType = [&](mlir::Value memref) {
       AccessType ret;
@@ -298,6 +335,15 @@ class InsertGPUAllocsPass final
             continue;
           }
 
+          if (auto init_xedesc = mlir::dyn_cast<mlir::xegpu::CreateNdDescOp>(user)) {
+            bool onDevice = user->getParentOfType<mlir::gpu::LaunchOp>();
+            auto res = init_xedesc->getResult(0);
+            for (auto use : res.getUsers()) {
+              findXeGPULoadStore(use, onDevice, ret);
+            }
+            continue;
+          }
+
           if (mlir::isa<mlir::func::ReturnOp>(user)) {
             ret.hostRead = true;
             ret.hostWrite = true;