From 50e3514a8221d02fd99a0fa766079fda694d6d00 Mon Sep 17 00:00:00 2001
From: Satyanvesh Dittakavi <Satyanvesh.Dittakavi@amd.com>
Date: Mon, 25 Sep 2023 14:21:32 +0000
Subject: [PATCH] SWDEV-421415 - Add hipRTC test with device-side malloc

Change-Id: Ia1c3f2ab703432915bda0c62cae584349e535463
---
 catch/unit/rtc/CMakeLists.txt        |   1 +
 catch/unit/rtc/hipRTCDeviceMalloc.cc | 129 +++++++++++++++++++++++++++
 2 files changed, 130 insertions(+)
 create mode 100644 catch/unit/rtc/hipRTCDeviceMalloc.cc

diff --git a/catch/unit/rtc/CMakeLists.txt b/catch/unit/rtc/CMakeLists.txt
index bd66868564..7b60e3754e 100644
--- a/catch/unit/rtc/CMakeLists.txt
+++ b/catch/unit/rtc/CMakeLists.txt
@@ -5,6 +5,7 @@ set(TEST_SRC
     hipRtcFunctional.cc
     hipStreamCaptureRtc.cc
     includepath.cc
+    hipRTCDeviceMalloc.cc
 )
 
 # AMD only tests
diff --git a/catch/unit/rtc/hipRTCDeviceMalloc.cc b/catch/unit/rtc/hipRTCDeviceMalloc.cc
new file mode 100644
index 0000000000..9262f35aff
--- /dev/null
+++ b/catch/unit/rtc/hipRTCDeviceMalloc.cc
@@ -0,0 +1,129 @@
+#include <hip_test_common.hh>
+
+#include <hip/hiprtc.h>
+#include <hip/hip_runtime.h>
+
+
+#include <cassert>
+#include <cstddef>
+#include <memory>
+#include <iostream>
+#include <iterator>
+#include <vector>
+
+static constexpr auto NUM_THREADS{32};
+static constexpr auto NUM_BLOCKS{10};
+
+// Tests device side malloc and free functions within a hiprtc kernel using simple add operation
+static constexpr auto devicemalloc{
+    R"(
+extern "C"
+__global__
+void devicemalloc(float* x, float* y, float* out, float** px, float** py, size_t n)
+{
+  size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid == 0) {
+    *px = (float*)malloc(sizeof(float) * n);
+    *py = (float*)malloc(sizeof(float) * n);
+     for (int i = 0; i < n; i++) {
+       (*px)[i] = x[i];
+       (*py)[i] = y[i];
+       out[i] = (*px)[i] + (*py)[i];
+     }
+     free(*px);
+     free(*py);
+  }
+}
+)"};
+
+TEST_CASE("Unit_hiprtc_devicemalloc") {
+  using namespace std;
+  hiprtcProgram prog;
+  hiprtcCreateProgram(&prog,       // prog
+                      devicemalloc,       // buffer
+                      "devicemalloc.cu",  // name
+                      0, nullptr, nullptr);
+  hipDeviceProp_t props;
+  int device = 0;
+  HIP_CHECK(hipGetDeviceProperties(&props, device));
+#ifdef __HIP_PLATFORM_AMD__
+  std::string sarg = std::string("--gpu-architecture=") + props.gcnArchName;
+#else
+  std::string sarg = std::string("--fmad=false");
+#endif
+  const char* options[] = {sarg.c_str()};
+  hiprtcResult compileResult{hiprtcCompileProgram(prog, 1, options)};
+  size_t logSize;
+  HIPRTC_CHECK(hiprtcGetProgramLogSize(prog, &logSize));
+  if (logSize) {
+    string log(logSize, '\0');
+    HIPRTC_CHECK(hiprtcGetProgramLog(prog, &log[0]));
+    std::cout << log << '\n';
+  }
+  REQUIRE(compileResult == HIPRTC_SUCCESS);
+  size_t codeSize;
+  HIPRTC_CHECK(hiprtcGetCodeSize(prog, &codeSize));
+
+  vector<char> code(codeSize);
+  HIPRTC_CHECK(hiprtcGetCode(prog, code.data()));
+
+  HIPRTC_CHECK(hiprtcDestroyProgram(&prog));
+
+  // Do hip malloc first so that we donot need to do a cuInit manually before calling hipModule APIs
+  size_t n = NUM_THREADS * NUM_BLOCKS;
+  size_t bufferSize = n * sizeof(float);
+
+  float *dX, *dY, *dOut;
+  HIP_CHECK(hipMalloc(&dX, bufferSize));
+  HIP_CHECK(hipMalloc(&dY, bufferSize));
+  HIP_CHECK(hipMalloc(&dOut, bufferSize));
+
+  hipModule_t module;
+  hipFunction_t kernel;
+  HIP_CHECK(hipModuleLoadData(&module, code.data()));
+  HIP_CHECK(hipModuleGetFunction(&kernel, module, "devicemalloc"));
+
+  unique_ptr<float[]> hX{new float[n]};
+  unique_ptr<float[]> hY{new float[n]};
+  unique_ptr<float[]> hOut{new float[n]};
+  for (size_t i = 0; i < n; ++i) {
+    hX[i] = static_cast<float>(i);
+    hY[i] = static_cast<float>(i * 2);
+  }
+
+  HIP_CHECK(hipMemcpy(dX, hX.get(), bufferSize, hipMemcpyHostToDevice));
+  HIP_CHECK(hipMemcpy(dY, hY.get(), bufferSize, hipMemcpyHostToDevice));
+
+  float **pA, **pB;
+  HIP_CHECK(hipMalloc((float***)&pA, sizeof(float*)));
+  HIP_CHECK(hipMalloc((float***)&pB, sizeof(float*)));
+
+  struct {
+    float* b_;
+    float* c_;
+    float* d_;
+    float** e_;
+    float** f_;
+    size_t g_;
+  } args{dX, dY, dOut, pA, pB, n};
+
+  auto size = sizeof(args);
+  void* config[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, &args, HIP_LAUNCH_PARAM_BUFFER_SIZE, &size,
+                    HIP_LAUNCH_PARAM_END};
+
+  HIP_CHECK(hipModuleLaunchKernel(kernel, NUM_BLOCKS, 1, 1, NUM_THREADS, 1, 1, 0, nullptr, nullptr, config));
+
+  HIP_CHECK(hipMemcpy(hOut.get(), dOut, bufferSize, hipMemcpyDeviceToHost));
+
+  HIP_CHECK(hipFree(dX));
+  HIP_CHECK(hipFree(dY));
+  HIP_CHECK(hipFree(dOut));
+
+  HIP_CHECK(hipModuleUnload(module));
+
+  for (size_t i = 0; i < n; ++i) {
+    INFO("For " << i << " Value: " << fabs(hX[i] + hY[i] - hOut[i])
+                << " with: " << (fabs(hOut[i] * 1.0f) * 1e-6));
+    REQUIRE(fabs(hX[i] + hY[i] - hOut[i]) <= fabs(hOut[i]) * 1e-6);
+  }
+}