Add an option for insert gpu allocs to treat allocations as gpu-native

Author: zhczhong

cmake/imex-version.txt

@@ -1 +1 @@
-6c2e414a953b9a118bce6adac21cf9d42630e674
+20f2eef4f6c10fcbd68d358591c8b3ef4d1b97d2

include/gc/Transforms/Passes.h

@@ -20,7 +20,6 @@ namespace func {
 class FuncOp;
 } // namespace func
 
-
 namespace LLVM {
 class LLVMDialect;
 }
@@ -116,7 +115,8 @@ void populateFrontendPasses(mlir::OpPassManager &);
 void populateCPUPipeline(mlir::OpPassManager &);
 
 #ifdef GC_USE_IMEX
-void populateGPUPipeline(mlir::OpPassManager &);
+struct GPUPipelineOption;
+void populateGPUPipeline(mlir::OpPassManager &, const GPUPipelineOption &);
 #endif
 
 #define GEN_PASS_DECL

lib/gc/Transforms/GPU/Pipeline.cpp

@@ -38,7 +38,17 @@
 
 namespace mlir::gc {
 
-void populateGPUPipeline(mlir::OpPassManager &pm) {
+struct GPUPipelineOption : PassPipelineOptions<GPUPipelineOption> {
+  PassOptions::Option<bool> isUsmArgs{
+      *this, "is-usm-args",
+      llvm::cl::desc("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"),
+      llvm::cl::init(true)};
+};
+
+void populateGPUPipeline(mlir::OpPassManager &pm,
+                         const GPUPipelineOption &pipelineOption) {
   pm.addNestedPass<func::FuncOp>(createIterativeTilingAndFusion());
 
   pm.addPass(bufferization::createEmptyTensorEliminationPass());
@@ -76,7 +86,11 @@ void populateGPUPipeline(mlir::OpPassManager &pm) {
   pm.addNestedPass<func::FuncOp>(createGpuMapParallelLoopsPass());
   pm.addNestedPass<func::FuncOp>(createParallelLoopToGpuPass());
 
-  pm.addNestedPass<func::FuncOp>(imex::createInsertGPUAllocsPass("opencl"));
+  imex::InsertGPUAllocsOptions insertGPUAllocsOption{
+      /*clientAPI*/ "opencl", /*inRegions*/ false,
+      /*isUsmArgs*/ pipelineOption.isUsmArgs.getValue()};
+  pm.addNestedPass<func::FuncOp>(
+      imex::createInsertGPUAllocsPass(insertGPUAllocsOption));
   pm.addPass(createGpuKernelOutliningPass());
   pm.addPass(createCanonicalizerPass());
   pm.addPass(imex::createSetSPIRVCapabilitiesPass());
@@ -112,9 +126,9 @@ void populateGPUPipeline(mlir::OpPassManager &pm) {
 }
 
 void registerGPUPipeline() {
-  PassPipelineRegistration<>("gc-gpu-pipeline",
-                             "The GPU pipeline for Graph Compiler with IMEX",
-                             populateGPUPipeline);
+  PassPipelineRegistration<GPUPipelineOption>(
+      "gc-gpu-pipeline", "The GPU pipeline for Graph Compiler with IMEX",
+      populateGPUPipeline);
 }
 
 } // namespace mlir::gc