From f2b0542442fe1c5022e84fa301b012fc835af7f4 Mon Sep 17 00:00:00 2001
From: Jaydeep Patel <jaydeepkumar.patel@amd.com>
Date: Mon, 26 Dec 2022 09:28:34 +0000
Subject: [PATCH] SWDEV-374370 - Propogate element size to blit kernel.

Change-Id: I8a3355c62947265be34e6a2befd26f86b3dc6cd2
---
 hipamd/src/hip_graph_helper.hpp   |  2 +-
 hipamd/src/hip_graph_internal.hpp | 12 ++++++------
 hipamd/src/hip_memory.cpp         |  6 +++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/hipamd/src/hip_graph_helper.hpp b/hipamd/src/hip_graph_helper.hpp
index fd8836a08b..4af57c6900 100644
--- a/hipamd/src/hip_graph_helper.hpp
+++ b/hipamd/src/hip_graph_helper.hpp
@@ -34,7 +34,7 @@ hipError_t ihipMemsetCommand(std::vector<amd::Command*>& commands, void* dst, in
                              size_t valueSize, size_t sizeBytes, amd::HostQueue* queue);
 
 hipError_t ihipMemset3DCommand(std::vector<amd::Command*>& commands, hipPitchedPtr pitchedDevPtr,
-                               int value, hipExtent extent, amd::HostQueue* queue);
+                               int value, hipExtent extent, amd::HostQueue* queue, size_t elementSize = 1);
 
 hipError_t ihipMemcpySymbol_validate(const void* symbol, size_t sizeBytes, size_t offset,
                                      size_t& sym_size, hipDeviceptr_t& device_ptr);
diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp
index 8f0d6b9bcf..da64616793 100644
--- a/hipamd/src/hip_graph_internal.hpp
+++ b/hipamd/src/hip_graph_internal.hpp
@@ -1438,7 +1438,7 @@ class hipGraphMemsetNode : public hipGraphNode {
     if (pMemsetParams_->height == 1) {
       sizeBytes = pMemsetParams_->width * pMemsetParams_->elementSize;
     } else {
-      sizeBytes = pMemsetParams_->width * pMemsetParams_->height;
+      sizeBytes = pMemsetParams_->width * pMemsetParams_->height * pMemsetParams_->elementSize;
     }
   }
 
@@ -1468,7 +1468,7 @@ class hipGraphMemsetNode : public hipGraphNode {
       if (pMemsetParams_->height == 1) {
         sizeBytes = pMemsetParams_->width * pMemsetParams_->elementSize;
       } else {
-        sizeBytes = pMemsetParams_->width * pMemsetParams_->height;
+        sizeBytes = pMemsetParams_->width * pMemsetParams_->height * pMemsetParams_->elementSize;
       }
       label = std::to_string(GetID()) + "\n" + label_ + "\n(" +
           std::to_string(pMemsetParams_->value) + "," + std::to_string(sizeBytes) + ")";
@@ -1496,9 +1496,9 @@ class hipGraphMemsetNode : public hipGraphNode {
     } else {
       hipError_t status = ihipMemset3DCommand(
           commands_,
-          {pMemsetParams_->dst, pMemsetParams_->pitch, pMemsetParams_->width,
+          {pMemsetParams_->dst, pMemsetParams_->pitch, pMemsetParams_->width * pMemsetParams_->elementSize,
            pMemsetParams_->height},
-          pMemsetParams_->value, {pMemsetParams_->width, pMemsetParams_->height, 1}, queue);
+          pMemsetParams_->value, {pMemsetParams_->width * pMemsetParams_->elementSize, pMemsetParams_->height, 1}, queue, pMemsetParams_->elementSize);
     }
     return status;
   }
@@ -1520,8 +1520,8 @@ class hipGraphMemsetNode : public hipGraphNode {
     } else {
       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);
+          ihipMemset3D_validate({params->dst, params->pitch, params->width * params->elementSize, params->height},
+                                params->value, {params->width * params->elementSize, 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 4131b367e7..219b5fedab 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -2920,12 +2920,12 @@ hipError_t ihipMemset3D_validate(hipPitchedPtr pitchedDevPtr, int value, hipExte
 }
 
 hipError_t ihipMemset3DCommand(std::vector<amd::Command*> &commands, hipPitchedPtr pitchedDevPtr,
-                               int value, hipExtent extent, amd::HostQueue* queue) {
+                               int value, hipExtent extent, amd::HostQueue* queue, size_t elementSize = 1) {
   size_t offset = 0;
   auto sizeBytes = extent.width * extent.height * extent.depth;
   amd::Memory* memory = getMemoryObject(pitchedDevPtr.ptr, offset);
   if (pitchedDevPtr.pitch == extent.width) {
-    return ihipMemsetCommand(commands, pitchedDevPtr.ptr, value, sizeof(int8_t),
+    return ihipMemsetCommand(commands, pitchedDevPtr.ptr, value, elementSize,
                                   static_cast<size_t>(sizeBytes), queue);
   }
   // Workaround for cases when pitch > row until fill kernel will be updated to support pitch.
@@ -2943,7 +2943,7 @@ hipError_t ihipMemset3DCommand(std::vector<amd::Command*> &commands, hipPitchedP
   amd::FillMemoryCommand* command;
   command = new amd::FillMemoryCommand(
       *queue, CL_COMMAND_FILL_BUFFER, amd::Command::EventWaitList{}, *memory->asBuffer(),
-      &value, sizeof(int8_t), origin, region, surface);
+      &value, elementSize, origin, region, surface);
   commands.push_back(command);
   return hipSuccess;
 }