add an option in insertGpuAllocs to skip the func args copy

Author: zhczhong

include/imex/Transforms/Passes.h

@@ -19,12 +19,15 @@
 #include <memory>
 
 namespace imex {
+struct InsertGPUAllocsOptions;
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 std::unique_ptr<mlir::Pass> createSerializeSPIRVPass();
 std::unique_ptr<mlir::Pass>
 createInsertGPUAllocsPass(const char *clientAPI = "vulkan");
+std::unique_ptr<mlir::Pass>
+createInsertGPUAllocsPass(const InsertGPUAllocsOptions &);
 std::unique_ptr<mlir::Pass> createSetSPIRVCapabilitiesPass();
 std::unique_ptr<mlir::Pass>
 createSetSPIRVAbiAttributePass(const char *clientAPI = "vulkan");

include/imex/Transforms/Passes.td

@@ -41,7 +41,11 @@ def InsertGPUAllocs : Pass<"insert-gpu-allocs", "::mlir::func::FuncOp"> {
     Option<"clientAPI", "client-api", "std::string", /*default=*/"\"opencl\"",
            "The client API to use for inserting gpu allocs">,
     Option<"inRegions", "in-regions", "bool", "false",
-           "Add gpu allocs only for memref.AllocOps within GPU regions">
+           "Add gpu allocs only for memref.AllocOps within GPU regions">,
+        Option<"isUsmArgs", "is-usm-args", "bool", "false",
+           "Whether to use USM(unified shared memory) func args, in which the "
+           "host and device could access the same buffer and there is no need "
+           "to add memcpy explicitly">
   ];
 }
 

lib/Transforms/InsertGPUAllocs.cpp

@@ -47,6 +47,12 @@ class InsertGPUAllocsPass final
   explicit InsertGPUAllocsPass() : m_clientAPI("vulkan") {}
   explicit InsertGPUAllocsPass(const mlir::StringRef &clientAPI)
       : m_clientAPI(clientAPI) {}
+  explicit InsertGPUAllocsPass(const imex::InsertGPUAllocsOptions &options)
+      : InsertGPUAllocsBase<InsertGPUAllocsPass>(options) {
+    if (clientAPI == "opencl") {
+      m_clientAPI = "opencl";
+    }
+  }
 
   mlir::LogicalResult
   initializeOptions(mlir::StringRef options,
@@ -540,15 +546,17 @@ class InsertGPUAllocsPass final
     // This is the case where the inputs are passed as arguments to the
     // function. This code will add the IR for memory allocation on the device
     // with gpu.alloc and insert a memref.copy from host to device
-    for (const auto &it : gpuBufferParams) {
-      auto param = block.getArgument(it.first);
-      if (isGpuAddrSpace(param))
-        continue;
-      auto access = getAccessType(param);
-      access.hostRead = true;
-      access.hostWrite = true;
-      builder.setInsertionPointToStart(&block);
-      add_gpu_alloc(builder, param, access, term);
+    if (!isUsmArgs.getValue()) {
+      for (const auto &it : gpuBufferParams) {
+        auto param = block.getArgument(it.first);
+        if (isGpuAddrSpace(param))
+          continue;
+        auto access = getAccessType(param);
+        access.hostRead = true;
+        access.hostWrite = true;
+        builder.setInsertionPointToStart(&block);
+        add_gpu_alloc(builder, param, access, term);
+      }
     }
 
     // CallOp Case: This is the case where the memref producer is coming
@@ -580,4 +588,8 @@ namespace imex {
 std::unique_ptr<mlir::Pass> createInsertGPUAllocsPass(const char *clientAPI) {
   return std::make_unique<InsertGPUAllocsPass>(clientAPI);
 }
+std::unique_ptr<mlir::Pass>
+createInsertGPUAllocsPass(const InsertGPUAllocsOptions &option) {
+  return std::make_unique<InsertGPUAllocsPass>(option);
+}
 } // namespace imex