From cb79dc7486520e3e4b2db4359a667f55aebb913d Mon Sep 17 00:00:00 2001
From: Ajay <ajay.gunashekar@amd.com>
Date: Fri, 21 Jan 2022 23:50:18 +0000
Subject: [PATCH] SWDEV-313082 - hipGraphMemsetNodeSetParams invalid params

Change-Id: Ibe57df8d7d583daf9ba9c7a86863eacca4868f3a
---
 hipamd/src/hip_graph.cpp          |  4 ++++
 hipamd/src/hip_graph_helper.hpp   |  2 ++
 hipamd/src/hip_graph_internal.hpp | 17 +++++++++++++++--
 hipamd/src/hip_memory.cpp         | 14 ++++++++++++++
 4 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp
index 4fd87a4791..3f6c971521 100644
--- a/hipamd/src/hip_graph.cpp
+++ b/hipamd/src/hip_graph.cpp
@@ -122,6 +122,10 @@ hipError_t ihipGraphAddMemsetNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
     return hipErrorInvalidValue;
   }
   hipError_t status;
+  status = ihipGraphMemsetParams_validate(pMemsetParams);
+  if (status != hipSuccess) {
+    return status;
+  }
   if (pMemsetParams->height == 1) {
     status =
         ihipMemset_validate(pMemsetParams->dst, pMemsetParams->value, pMemsetParams->elementSize,
diff --git a/hipamd/src/hip_graph_helper.hpp b/hipamd/src/hip_graph_helper.hpp
index ffd97e01df..873f66c31c 100644
--- a/hipamd/src/hip_graph_helper.hpp
+++ b/hipamd/src/hip_graph_helper.hpp
@@ -86,3 +86,5 @@ hipError_t ihipMemcpyAtoHValidate(hipArray* srcArray, void* dstHost, amd::Coord3
                                   amd::Coord3D& dstOrigin, amd::Coord3D& copyRegion,
                                   size_t dstRowPitch, size_t dstSlicePitch, amd::Image*& srcImage,
                                   amd::BufferRect& dstRect);
+
+hipError_t ihipGraphMemsetParams_validate(const hipMemsetParams* pNodeParams);
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 7b457eab40..7c7da322e3 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -811,9 +811,22 @@ class hipGraphMemsetNode : public hipGraphNode {
     std::memcpy(params, pMemsetParams_, sizeof(hipMemsetParams));
   }
   hipError_t SetParams(const hipMemsetParams* params) {
+
     hipError_t hip_error = hipSuccess;
-    hip_error = ihipMemset_validate(params->dst, params->value, params->elementSize,
-                                    params->width * params->elementSize);
+    hip_error = ihipGraphMemsetParams_validate(params);
+    if (hip_error != hipSuccess) {
+      return hip_error;
+    }
+    if (params->height == 1) {
+      hip_error = ihipMemset_validate(params->dst, params->value, params->elementSize,
+                                      params->width * params->elementSize);
+    } else {
+      auto sizeBytes = params->width * params->height * 1;
+      hip_error = ihipMemset3D_validate(
+          {params->dst, params->pitch, params->width, params->height},
+          params->value, {params->width, params->height, 1}, sizeBytes);
+    }
+
     if (hip_error != hipSuccess) {
       return hip_error;
     }
diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index c4dc8ebff1..89b553a0a2 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -2276,6 +2276,20 @@ hipError_t ihipMemset_validate(void* dst, int64_t value, size_t valueSize,
   return hipSuccess;
 }
 
+hipError_t ihipGraphMemsetParams_validate(const hipMemsetParams* pNodeParams) {
+  if (pNodeParams == nullptr) {
+    return hipErrorInvalidValue;
+  }
+  if (pNodeParams->elementSize != 1 && pNodeParams->elementSize != 2 && pNodeParams->elementSize != 4) {
+    return hipErrorInvalidValue;
+  }
+
+  if (pNodeParams->height <= 0) {
+    return hipErrorInvalidValue;
+  }
+  return hipSuccess;
+}
+
 hipError_t ihipMemsetCommand(std::vector<amd::Command*>& commands, void* dst, int64_t value,
                              size_t valueSize, size_t sizeBytes, amd::HostQueue* queue) {
   hipError_t hip_error = hipSuccess;