SWDEV-489571 - Fix ihipGraphAddMemsetNode to allow memset of 3d portions of an array

* When hipMemset3dAsync is captured, a 3d extent can set be as a parameter (depth > 1). That worked on nvidia, but on amd wrong portion of array was filled because when creating Memset3D command, extent dimensions were used to create pitchedPtr, instead of original array width and height.
* Also, when capturing hipMemset3dAsync, nvidia allows any of the extent dimension to be 0, and in that case, no work should be done.

Change-Id: I46a605bf9ae801cd3348e98d528c21263a8eefce

hipamd/src/hip_graph.cpp

@@ -171,8 +171,8 @@ hipError_t ihipGraphAddMemcpyNode1D(hip::GraphNode** pGraphNode, hip::Graph* gra
 
 hipError_t ihipGraphAddMemsetNode(hip::GraphNode** pGraphNode, hip::Graph* graph,
                                   hip::GraphNode* const* pDependencies, size_t numDependencies,
-                                  const hipMemsetParams* pMemsetParams,
-                                  bool capture = true, size_t depth = 1) {
+                                  const hipMemsetParams* pMemsetParams, bool capture = true,
+                                  size_t depth = 1, size_t arrWidth = 1, size_t arrHeight = 1) {
   if (pGraphNode == nullptr || graph == nullptr || pMemsetParams == nullptr ||
       (numDependencies > 0 && pDependencies == nullptr) || pMemsetParams->height == 0) {
     return hipErrorInvalidValue;
@@ -208,7 +208,7 @@ hipError_t ihipGraphAddMemsetNode(hip::GraphNode** pGraphNode, hip::Graph* graph
   if (status != hipSuccess) {
     return status;
   }
-  *pGraphNode = new hip::GraphMemsetNode(pMemsetParams, depth);
+  *pGraphNode = new hip::GraphMemsetNode(pMemsetParams, depth, arrWidth, arrHeight);
   status = ihipGraphAddNode(*pGraphNode, graph, pDependencies, numDependencies, capture);
   return status;
 }
@@ -747,10 +747,16 @@ hipError_t capturehipMemset3DAsync(hipStream_t& stream, hipPitchedPtr& pitchedDe
                                    hipExtent& extent) {
   ClPrint(amd::LOG_INFO, amd::LOG_API, "[hipGraph] Current capture node Memset3D on stream : %p",
           stream);
-  hipMemsetParams memsetParams = {0};
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
+
+  // In this case no work should be done
+  if (extent.width == 0 || extent.height == 0 || extent.depth == 0) {
+    return hipSuccess;
+  }
+
+  hipMemsetParams memsetParams = {0};
   memsetParams.dst = pitchedDevPtr.ptr;
   memsetParams.value = value;
   memsetParams.width = extent.width;
@@ -761,7 +767,8 @@ hipError_t capturehipMemset3DAsync(hipStream_t& stream, hipPitchedPtr& pitchedDe
   hip::GraphNode* pGraphNode;
   hipError_t status =
       ihipGraphAddMemsetNode(&pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(),
-                             s->GetLastCapturedNodes().size(), &memsetParams, true, extent.depth);
+                             s->GetLastCapturedNodes().size(), &memsetParams, true, extent.depth,
+                             pitchedDevPtr.xsize, pitchedDevPtr.ysize);
   if (status != hipSuccess) {
     return status;
   }

hipamd/src/hip_graph_internal.hpp

@@ -1979,11 +1979,16 @@ class GraphMemcpyNodeToSymbol : public GraphMemcpyNode1D {
 class GraphMemsetNode : public GraphNode {
   hipMemsetParams memsetParams_;
   size_t depth_ = 1;
+  size_t arrWidth_ = 1;
+  size_t arrHeight_ = 1;
  public:
-  GraphMemsetNode(const hipMemsetParams* pMemsetParams, size_t depth = 1)
+  GraphMemsetNode(const hipMemsetParams* pMemsetParams, size_t depth = 1, size_t arrWidth = 1,
+                  size_t arrHeight = 1)
       : GraphNode(hipGraphNodeTypeMemset, "solid", "invtrapezium", "MEMSET") {
     memsetParams_ = *pMemsetParams;
     depth_ = depth;
+    arrWidth_ = arrWidth;
+    arrHeight_ = arrHeight;
     size_t sizeBytes = 0;
     if (memsetParams_.height == 1) {
       sizeBytes = memsetParams_.width * memsetParams_.elementSize;
@@ -1997,6 +2002,8 @@ class GraphMemsetNode : public GraphNode {
   GraphMemsetNode(const GraphMemsetNode& memsetNode) : GraphNode(memsetNode) {
     memsetParams_ = memsetNode.memsetParams_;
     depth_ = memsetNode.depth_;
+    arrWidth_ = memsetNode.arrWidth_;
+    arrHeight_ = memsetNode.arrHeight_;
   }
 
   GraphNode* clone() const override {
@@ -2040,15 +2047,15 @@ class GraphMemsetNode : public GraphNode {
     if (status != hipSuccess) {
       return status;
     }
-    if (memsetParams_.height == 1) {
+    if (memsetParams_.height == 1 && depth_ == 1) {
       size_t sizeBytes = memsetParams_.width * memsetParams_.elementSize;
       hipError_t status = ihipMemsetCommand(commands_, memsetParams_.dst, memsetParams_.value,
                                             memsetParams_.elementSize, sizeBytes, stream);
     } else {
       hipError_t status = ihipMemset3DCommand(
           commands_,
-          {memsetParams_.dst, memsetParams_.pitch, memsetParams_.width * memsetParams_.elementSize,
-           memsetParams_.height},
+          {memsetParams_.dst, memsetParams_.pitch, arrWidth_ * memsetParams_.elementSize,
+           arrHeight_},
           memsetParams_.value,
           {memsetParams_.width * memsetParams_.elementSize, memsetParams_.height, depth_}, stream,
           memsetParams_.elementSize);