[CINN] support dcu get_value_in_kernel_args with sycl and hip (#74266)

Author: ZhangX-21

paddle/cinn/backends/codegen_device_util.cc

@@ -44,13 +44,12 @@ ir::Module CreateSwitchWithBroadcastConditionModule(
       ir::Argument(kernel_args, ir::Argument::IO::kOutput),
       ir::Argument(kernel_args_num, ir::Argument::IO::kInput),
       ir::Argument(tensor_shape_args, ir::Argument::IO::kOutput)};
-
   const auto &symbolic_arg_define = [&]() -> std::vector<ir::Expr> {
     std::vector<ir::Expr> arg_defs;
     for (const auto &item : symbolic_shape_var_index) {
       ir::Expr call_get_value_in_kernel_args =
           ir::Call::Make(Int(64),
-                         runtime::intrinsic::get_value_in_cuda_kernel_args,
+                         runtime::intrinsic::get_value_in_kernel_args,
                          {kernel_args, ir::Expr(item.first)},
                          {},
                          ir::CallType::Extern,
@@ -384,7 +383,7 @@ void detail::CollectBucketStrategyHostFunctionVisitor::ProcessArgs(
     if (args[i].is_var()) {
       ir::Expr call_get_value_in_kernel_args =
           ir::Call::Make(Int(64),
-                         runtime::intrinsic::get_value_in_cuda_kernel_args,
+                         runtime::intrinsic::get_value_in_kernel_args,
                          {kernel_args_, ir::Expr(i)},
                          {},
                          ir::CallType::Extern,

paddle/cinn/backends/codegen_invoke_module.h

@@ -68,7 +68,7 @@ class CodeGenSwitchHost : public CodeGenInvokeModule {
       : CodeGenInvokeModule(m, b, vars) {}
   // only support call of args get function and inner case host function call
   llvm::Value *Visit(const ir::Call *op) override {
-    if (op->name == runtime::intrinsic::get_value_in_cuda_kernel_args) {
+    if (op->name == runtime::intrinsic::get_value_in_kernel_args) {
       return CodeGenLLVM::Visit(op);
     } else {
       return LowerInnerCaseCall(op);

paddle/cinn/runtime/cpu/host_intrinsics.cc

@@ -280,6 +280,12 @@ inline int64_t FN_INT64(logical_right_shift)(int64_t x, int64_t y) {
 }
 
 #undef FN_INT64
+
+int64_t cinn_get_value_in_kernel_args(void* v_args, int idx) {
+  cinn_pod_value_t* args = static_cast<cinn_pod_value_t*>(v_args);
+  return args[idx].operator int64_t();
+}
+
 }  // extern "C"
 
 CINN_REGISTER_HELPER(host_intrinsics) {
@@ -469,5 +475,11 @@ CINN_REGISTER_HELPER(host_intrinsics) {
       .AddInputType<int>()
       .End();
 
+  REGISTER_EXTERN_FUNC_HELPER(cinn_get_value_in_kernel_args, host_target)
+      .SetRetType<int64_t>()
+      .AddInputType<void*>()  // args
+      .AddInputType<int>()    // index
+      .End();
+
   return true;
 }

paddle/cinn/runtime/cpu/host_intrinsics.h

@@ -121,4 +121,6 @@ inline float FN_FP32(cbrt)(float x);
 inline double FN_FP64(cbrt)(double x);
 
 #undef FN_FP64
+
+int64_t cinn_get_value_in_kernel_args(void* v_args, int idx);
 }

paddle/cinn/runtime/cuda/cuda_intrinsics.cc

@@ -431,14 +431,6 @@ CINN_REGISTER_HELPER(cuda_intrinsics) {
 }
 
 CINN_REGISTER_HELPER(cinn_cuda_host_api) {
-  using cinn::runtime::cuda::cinn_get_value_in_cuda_kernel_args;
-  REGISTER_EXTERN_FUNC_HELPER(cinn_get_value_in_cuda_kernel_args,
-                              cinn::common::DefaultHostTarget())
-      .SetRetType<int64_t>()
-      .AddInputType<void *>()  // args
-      .AddInputType<int>()     // index
-      .End();
-
   using cinn::runtime::cuda::cinn_get_item_in_cuda_kernel_args;
   REGISTER_EXTERN_FUNC_HELPER(cinn_get_item_in_cuda_kernel_args,
                               cinn::common::DefaultHostTarget())

paddle/cinn/runtime/cuda/cuda_util.cc

@@ -78,11 +78,6 @@ class CublasHandle {
   cublasHandle_t cuhandle;
 };
 
-int64_t cinn_get_value_in_cuda_kernel_args(void *v_args, int idx) {
-  cinn_pod_value_t *args = static_cast<cinn_pod_value_t *>(v_args);
-  return args[idx].operator int64_t();
-}
-
 void *cinn_get_item_in_cuda_kernel_args(void *v_args, int idx) {
   cinn_pod_value_t *args = static_cast<cinn_pod_value_t *>(v_args);
   return static_cast<void *>(&args[idx]);

paddle/cinn/runtime/cuda/cuda_util.h

@@ -89,7 +89,6 @@ void cinn_call_cuda_memcpy(void* v_args,
                            size_t count,
                            void* stream = nullptr);
 
-int64_t cinn_get_value_in_cuda_kernel_args(void* v_args, int idx);
 void* cinn_get_item_in_cuda_kernel_args(void* v_args, int idx);
 
 void infer_shape_set_value(int row, int col, int64_t value, int64_t** v);

paddle/cinn/runtime/hip/hip_intrinsics.cc

@@ -17,12 +17,12 @@ using cinn::backends::GlobalSymbolRegistry;
 using cinn::runtime::hip::HIPBackendAPI;
 #include "paddle/cinn/backends/extern_func_jit_register.h"
 #include "paddle/cinn/runtime/hip/hip_util.h"
-using cinn::runtime::hip::cinn_call_hip_kernel;
 
 CINN_REGISTER_HELPER(cinn_hip_host_api) {
   GlobalSymbolRegistry::Global().RegisterFn(
       "backend_api.hip", reinterpret_cast<void *>(HIPBackendAPI::Global()));
 
+  using cinn::runtime::hip::cinn_call_hip_kernel;
   REGISTER_EXTERN_FUNC_HELPER(cinn_call_hip_kernel,
                               cinn::common::DefaultHostTarget())
       .SetRetType<void>()

paddle/cinn/runtime/intrinsic.h

@@ -113,8 +113,7 @@ static const char* call_sycl_kernel = "cinn_call_sycl_kernel";
 
 static const char* call_cuda_memset = "cinn_call_cuda_memset";
 
-static const char* get_value_in_cuda_kernel_args =
-    "cinn_get_value_in_cuda_kernel_args";
+static const char* get_value_in_kernel_args = "cinn_get_value_in_kernel_args";
 
 static const char* get_item_in_cuda_kernel_args =
     "cinn_get_item_in_cuda_kernel_args";

paddle/cinn/runtime/sycl/sycl_util.cc

@@ -139,7 +139,6 @@ void cinn_call_sycl_memcpy(void *v_args,
 }
 
 #ifdef CINN_WITH_CNNL
-
 class CnnlHandle {
  public:
   CnnlHandle(const CnnlHandle &) = delete;