add support for mem_banks for --add-pskernel

src/runtime_src/tools/xclbinutil/KernelUtilities.cxx

@@ -1,5 +1,6 @@
 /**
  * Copyright (C) 2021-2022 Xilinx, Inc
+ * Copyright (C) 2023 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License"). You may
  * not use this file except in compliance with the License. A copy of the
@@ -288,15 +289,30 @@ void addArgsToMemoryConnections(const unsigned int ipLayoutIndexID,
     const auto& ptArg = arg.second;
 
     // Determine if there should be a memory connection, if so add it
+    //
     static const std::string NOT_DEFINED = "<not defined>";
     auto memoryConnection = ptArg.get<std::string>("memory-connection", NOT_DEFINED);
 
-    // An empty memory connection indicates that we should connect it to a yet to be define PS kernel memory entry
+    // if memory-connection is not defined, skip this arg
+    // if memory-connection is empty we should connect the arg to a yet to be define PS kernel memory entry
+    // if memory-connection is not empty
+    //   1. for --add-pskernel, user can specify the memory bank indices in
+    //      the value
+    //   2. for --add-kernel, user can specify the memory bank tag in the 
+    //      input file
+    if (memoryConnection == NOT_DEFINED) {
+      ++argIndexID;
+      continue;
+    }
+
+    std::vector<std::string> vMemBanks;
+
     if (memoryConnection.empty()) {
       // Look for entry.  If it doesn't exist create one
       auto iter = std::find_if(memTopology.begin(), memTopology.end(),
                                [](const boost::property_tree::ptree& pt)
       {return pt.get<std::string>("m_type", "") == "MEM_PS_KERNEL";});
+
       if (iter == memTopology.end()) {
         XUtil::TRACE("MEM Entry of PS Kernel memory not found, creating one.");
         boost::property_tree::ptree ptMemData;
@@ -306,26 +322,67 @@ void addArgsToMemoryConnections(const unsigned int ipLayoutIndexID,
         ptMemData.put("m_base_address", "0x0");
         memTopology.push_back(ptMemData);
       }
-
-      memoryConnection = "MEM_PS_KERNEL";
-    }
-
-    // Do we have a connection to perform?
-    if (memoryConnection != NOT_DEFINED) {
-      auto iter = std::find_if(memTopology.begin(), memTopology.end(),
+      // memoryConnection = "MEM_PS_KERNEL";
+
+      // the newly added mem bank is always at the last of memTopology
+      const unsigned int memIndex = static_cast<unsigned int>(memTopology.size()) - 1;
+      vMemBanks.push_back(std::to_string(memIndex));
+    } else {
+      // memory-connection could be
+      //   a comma separated memory bank indices, for --add-pskernel
+      //   a string which matches to m_tag, for --add-kernel
+      std::vector<int> viMemBanks;
+      auto isNumeric = [&viMemBanks](const std::string& str) {
+        std::istringstream ss(str);
+        std::string token;
+        while(std::getline(ss, token, ',')) {
+          int index;
+          try {
+            index = std::stoi(token);
+          } catch (const std::exception&) {
+            return false;
+          }
+          viMemBanks.push_back(index);
+        }
+        return true;
+      };
+
+      if (isNumeric(memoryConnection)) {
+        // make sure the specified index exists
+        for (const auto& memBankIndex : viMemBanks) {
+          if (memBankIndex >= 0 && memBankIndex < static_cast<int>(memTopology.size())) {
+            vMemBanks.push_back(std::to_string(memBankIndex));
+            std::cout << "\tconnecting arg " << std::to_string(argIndexID) << " to mem bank " << std::to_string(memBankIndex) << std::endl;
+          }
+          else
+            throw std::runtime_error("Specifїed mem_banks is not valid");
+        }
+      } else {
+        auto iter = std::find_if(memTopology.begin(), memTopology.end(),
                                [memoryConnection](const boost::property_tree::ptree& pt)
-      {return pt.get<std::string>("m_tag", "") == memoryConnection;});
+        {return pt.get<std::string>("m_tag", "") == memoryConnection;});
 
-      if (iter == memTopology.end())
-        throw std::runtime_error("Error: Memory tag '" + memoryConnection + "' not found in the MEM_TOPOLOGY section.");
+        if (iter == memTopology.end())
+          throw std::runtime_error("Error: Memory tag '" + memoryConnection + "' not found in the MEM_TOPOLOGY section.");
 
-      // We have a match
-      uint64_t memIndex = std::distance(memTopology.begin(), iter);
+        // We have a match
+        uint64_t memIndex = std::distance(memTopology.begin(), iter);
+        vMemBanks.push_back(std::to_string(memIndex));
+      }
+    }
+
+    for (const auto& memBank : vMemBanks) {
       boost::property_tree::ptree ptEntry;
       ptEntry.put("arg_index", std::to_string(argIndexID));
       ptEntry.put("m_ip_layout_index", std::to_string(ipLayoutIndexID));
-      ptEntry.put("mem_data_index", std::to_string(memIndex));
+      ptEntry.put("mem_data_index", memBank);
       connectivity.push_back(ptEntry);
+
+      // update the corresponding memtopology
+      int index = std::stoi(memBank);
+      boost::property_tree::ptree& ptMemData = memTopology[index];
+      if (ptMemData.get<int>("m_used", 0) != 1)
+        ptMemData.put("m_used", "1");
     }
     ++argIndexID;
   }
@@ -561,7 +618,8 @@ XclBinUtilities::validateFunctions(const std::string& kernelLibrary, const boost
 
 
 void
-XclBinUtilities::createPSKernelMetadata(unsigned long numInstances,
+XclBinUtilities::createPSKernelMetadata(const std::string& memBanks,
+                                        unsigned long numInstances,
                                         const boost::property_tree::ptree& ptFunctions,
                                         const std::string& kernelLibrary,
                                         boost::property_tree::ptree& ptPSKernels)
@@ -605,7 +663,7 @@ XclBinUtilities::createPSKernelMetadata(unsigned long numInstances,
       if (addrQualifier == "GLOBAL") {
         byteSize = 16;
         if (entry.get<int>("use-id", 1))
-          ptArg.put("memory-connection", "");
+          ptArg.put("memory-connection", memBanks);
       }
 
       ptArg.put("byte-size", byteSize);

src/runtime_src/tools/xclbinutil/KernelUtilities.h

@@ -26,6 +26,6 @@ namespace XclBinUtilities {
 void addKernel(const boost::property_tree::ptree& ptKernel, bool isFixedPS, boost::property_tree::ptree& ptEmbeddedData);
 void addKernel(const boost::property_tree::ptree& ptKernel, boost::property_tree::ptree& ptMemTopology, boost::property_tree::ptree& ptIPLayout, boost::property_tree::ptree& ptConnectivity);
 void validateFunctions(const std::string& kernelLibrary, const boost::property_tree::ptree& ptFunctions);
-void createPSKernelMetadata(unsigned long numInstances, const boost::property_tree::ptree& ptFunctions, const std::string& kernelLibrary, boost::property_tree::ptree& ptPSKernels);
+void createPSKernelMetadata(const std::string& memBank, unsigned long numInstances, const boost::property_tree::ptree& ptFunctions, const std::string& kernelLibrary, boost::property_tree::ptree& ptPSKernels);
 };
 #endif

src/runtime_src/tools/xclbinutil/XclBinClass.cxx

@@ -1669,12 +1669,13 @@ XclBin::reportInfo(std::ostream& _ostream, const std::string& _sInputFile, bool
 
 static void
 parsePSKernelString(const std::string& encodedString,
+                    std::string& mem_banks,
                     std::string& symbol_name,
                     unsigned long& num_instances,
                     std::string& path_to_library)
 // Line being parsed:
-//   Syntax: <symbol_name>:<instances>:<path_to_shared_library>
-//   Example: myKernel:3:./data/mylib.so
+//   Syntax: <mem_banks>:<symbol_name>:<instances>:<path_to_shared_library>
+//   Example: 0,1:myKernel:3:./data/mylib.so
 //
 // Note: A file name can contain a colen (e.g., C:\test)
 {
@@ -1687,7 +1688,7 @@ parsePSKernelString(const std::string& encodedString,
   std::vector<std::string> tokens;
 
   // Parse the string until the entire string has been parsed or MAX_TOKENS tokens have been found
-  constexpr size_t maxTokens = 3;
+  constexpr size_t maxTokens = 4;
   while ((lastPos < encodedString.length() + 1) &&
          (tokens.size() < maxTokens)) {
     pos = encodedString.find_first_of(delimiters, lastPos);
@@ -1723,7 +1724,18 @@ parsePSKernelString(const std::string& encodedString,
   // -- [2]: Symbolic name --
   symbol_name = (tokens.size() > 2) ? tokens[2] : "";
 
-  XUtil::TRACE(boost::format("PSKernel command arguments: symbol_name='%s'; num_instances=%d; library='%s'") % symbol_name % num_instances % path_to_library);
+  // -- [3]: Mem banks --
+  mem_banks = (tokens.size() > 3) ? tokens[3] : "";
+
+  // add check for leading and trailing ;
+  if (!mem_banks.empty()) {
+    if (mem_banks.front() == ',' || mem_banks.back() == ',' )
+      throw std::runtime_error("Specifїed mem_banks is not valid");
+
+    std::cout << "Specifying memory banks in --add-pskernel is an adanced feature" << std::endl;
+  }
+
+  XUtil::TRACE(boost::format("PSKernel command arguments: mem_banks='%s', symbol_name='%s'; num_instances=%d; library='%s'") % mem_banks % symbol_name % num_instances % path_to_library);
 }
 
 void getSectionPayload(const XclBin* pXclBin,
@@ -1786,15 +1798,17 @@ updateKernelSections(const std::vector<boost::property_tree::ptree> &kernels,
   }
 }
 
+// --add-pskernel
 void
 XclBin::addPsKernel(const std::string& encodedString)
 {
   XUtil::TRACE("Adding PSKernel");
   // Get the PS Kernel metadata from the encoded string
+  std::string memBanks;
   std::string symbolicName;
   std::string kernelLibrary;
   unsigned long numInstances = 0;
-  parsePSKernelString(encodedString, symbolicName, numInstances, kernelLibrary);
+  parsePSKernelString(encodedString, memBanks, symbolicName, numInstances, kernelLibrary);
 
   // Examine the PS library data mining the function and its arguments
   // Convert the function signatures into something useful.
@@ -1804,7 +1818,7 @@ XclBin::addPsKernel(const std::string& encodedString)
 
   // Create the same schema that is used for kernels
   boost::property_tree::ptree ptPSKernels;
-  XUtil::createPSKernelMetadata(numInstances, ptFunctions, kernelLibrary, ptPSKernels);
+  XUtil::createPSKernelMetadata(memBanks, numInstances, ptFunctions, kernelLibrary, ptPSKernels);
 
   // Update the EMBEDDED_METADATA, MEM_TOPOLOGY, IP_LAYOUT, and CONNECTIVITY sections
   const boost::property_tree::ptree ptEmpty;  
@@ -1870,7 +1884,7 @@ XclBin::addPsKernel(const std::string& encodedString)
   }
 }
 
-
+// --add-kernel
 void
 XclBin::addKernels(const std::string& jsonFile)
 {

src/runtime_src/tools/xclbinutil/XclBinUtilMain.cxx

@@ -165,7 +165,7 @@ int main_(int argc, const char** argv) {
   po::options_description desc("Options");
   desc.add_options()
       ("add-merge-section", boost::program_options::value<decltype(sectionsToAddMerge)>(&sectionsToAddMerge)->multitoken(), "Section name to add or merge.  Format: <section>:<format>:<file>")
-      ("add-pskernel", boost::program_options::value<decltype(addPsKernels)>(&addPsKernels)->multitoken(), "Helper option to add PS kernels.  Format: <symbol_name>:<instances>:<path_to_shared_library>")
+      ("add-pskernel", boost::program_options::value<decltype(addPsKernels)>(&addPsKernels)->multitoken(), "Helper option to add PS kernels.  Format: [<mem_banks>]:[<symbol_name>]:[<instances>]:<path_to_shared_library>")
       ("add-replace-section", boost::program_options::value<decltype(sectionsToAddReplace)>(&sectionsToAddReplace)->multitoken(), "Section name to add or replace.  Format: <section>:<format>:<file>")
       ("add-section", boost::program_options::value<decltype(sectionsToAdd)>(&sectionsToAdd)->multitoken(), "Section name to add.  Format: <section>:<format>:<file>")
       ("add-signature", boost::program_options::value<decltype(sSignature)>(&sSignature), "Adds a user defined signature to the given xclbin image.")
@@ -502,11 +502,11 @@ int main_(int argc, const char** argv) {
   }
 
   // -- Add PS Kernels
-  for (const auto &psKernel : addPsKernels) 
+  for (const auto &psKernel : addPsKernels)
     xclBin.addPsKernel(psKernel);
 
   // -- Add Fixed Kernels files
-  for (const auto &kernel : addKernels) 
+  for (const auto &kernel : addKernels)
     xclBin.addKernels(kernel);
 
   // -- Post Section Processing --

src/runtime_src/tools/xclbinutil/unittests/PSKernel/PSKernel.py

@@ -54,8 +54,8 @@ def main():
 
   step = "1a) Create shared ps kernel library (compile objects)"
 
-  # Note: This hex image was created from the pskernel.cpp file and converted
-  # to hex via the command
+  # Note: This hex image was created by first compiling the pskernel.cpp file
+  # into pskernel.so, then converting the so to hex via the command
   # xxd -p pskernel.so | tr -d ' \n' > pskernel.hex
 
   psKernelSharedLibraryHex = os.path.join(args.resource_dir, "pskernel.hex")
@@ -71,7 +71,6 @@ def main():
       file.write(binImage)
 
   # ---------------------------------------------------------------------------
-
   step = "2a) Read in a PS kernel, updated and validate the sections"
 
   inputPSKernelLib = psKernelSharedLibrary
@@ -96,7 +95,7 @@ def main():
                      "--dump-section", "CONNECTIVITY:JSON:" + outputConnectivity,
                      "--dump-section", "MEM_TOPOLOGY:JSON:" + outputMemTopology,
                      "--output", outputXCLBIN, 
-                     "--force"
+                     "--force", "--trace"
                      ]
   execCmd(step, cmd)
 
@@ -123,7 +122,46 @@ def main():
   jsonFileCompare(expectedKernelJSON, outputKernelJSON)
 
   # ---------------------------------------------------------------------------
-  step = "3) Validate adding just a soft kernel"
+  step = "3a) Read in a PS kernel with specified mem banks, updated and validate the sections"
+
+  inputPSKernelLib = psKernelSharedLibrary
+  # mb: mem-bank
+  outputEmbeddedMetadata = "embedded_metadata_mb_updated.xml"
+  expectedEmbeddedMetadata = os.path.join(args.resource_dir, "embedded_metadata_expected.xml")
+
+  outputIpLayout = "ip_layout_mb_updated.json"
+  expectedIpLayout = os.path.join(args.resource_dir, "ip_layout_expected.json")
+
+  # this should be different
+  outputConnectivity = "connectivity_mb_updated.json"
+  expectedConnectivity = os.path.join(args.resource_dir, "connectivity_mb_expected.json")
+
+  # this should be different
+  outputMemTopology = "mem_topology_mb_updated.json"
+  expectedMemTopology = os.path.join(args.resource_dir, "mem_topology_mb_expected.json")
+
+  outputXCLBIN = "pskernel_output_mb.xclbin"
+
+  # connect the PS kernel to mem banks 0 and 1
+  cmd = [xclbinutil, "--input", workingXCLBIN,
+                     "--add-pskernel", "0,1:::" + inputPSKernelLib,
+                     "--dump-section", "EMBEDDED_METADATA:RAW:" + outputEmbeddedMetadata,
+                     "--dump-section", "IP_LAYOUT:JSON:" + outputIpLayout,
+                     "--dump-section", "CONNECTIVITY:JSON:" + outputConnectivity,
+                     "--dump-section", "MEM_TOPOLOGY:JSON:" + outputMemTopology,
+                     "--output", outputXCLBIN, 
+                     "--force"
+                     ]
+  execCmd(step, cmd)
+
+  # Validate the contents of the various sections
+  textFileCompare(outputEmbeddedMetadata, expectedEmbeddedMetadata)
+  jsonFileCompare(outputIpLayout, expectedIpLayout)
+  jsonFileCompare(outputConnectivity, expectedConnectivity)
+  jsonFileCompare(outputMemTopology, expectedMemTopology)
+
+  # ---------------------------------------------------------------------------
+  step = "4) Validate adding just a soft kernel"
 
   outputOnlyPSKernelXclbin = "only_pskernel.xclbin"
   outputPSKEmbeddedMetadata = "embedded_metadata_psk.xml"

src/runtime_src/tools/xclbinutil/unittests/PSKernel/connectivity_mb_expected.json

@@ -0,0 +1,62 @@
+{
+    "connectivity": {
+        "m_count": "11",
+        "m_connection": [
+            {
+                "arg_index": "0",
+                "m_ip_layout_index": "0",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "1",
+                "m_ip_layout_index": "0",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "2",
+                "m_ip_layout_index": "0",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "0",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "0",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "1"
+            },
+            {
+                "arg_index": "1",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "1",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "1"
+            },
+            {
+                "arg_index": "2",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "2",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "1"
+            },
+            {
+                "arg_index": "6",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "0"
+            },
+            {
+                "arg_index": "6",
+                "m_ip_layout_index": "1",
+                "mem_data_index": "1"
+            }
+        ]
+    }
+}