feat(rocr/aie): resolve a full ELF's external symbols against defined variables

projects/rocr-runtime/docs/aie-hsaco-format.md

@@ -281,24 +281,25 @@ The symbol name gives the meaning of the patch site:
 | --- | --- | --- |
 | Starts with `.pdi` | The device address of the PDI in the section with this name. If more than one section has this name, their contents must be the same, or ROCR refuses the ELF. | 8 (`Address64`) |
 | Decimal digits only, 0 to 4095 | The address of kernel argument N. | 5 (`ShimDma48`) |
-| Any other name | Not supported (for example scratch pads or control packets). ROCR refuses the ELF. | — |
+| Any other name | The address of an external variable that the application defines before it loads the code object, for example the control scratchpad, `scratch-pad-ctrl` (section 9.3). `st_size` is the number of bytes that the control code reads there. | 8 (`Address64`) |
 
 ROCR applies these limits:
 
 - A kernel can have more than one PDI patch site, and the sites can name different PDI sections. For example, the control code from MLIR-AIR loads an empty PDI before and after the design runs. ROCR copies each different PDI into device memory one time for each kernel. Sites that name PDIs with the same contents use one copy.
 - Each PDI patch site offset must not be 0, must be a multiple of 4, and `offset + 8` must not be more than the control-code size.
 - Each argument patch site offset must be a multiple of 4, and `offset + 12` must not be more than the control-code size.
+- Each external variable patch site offset must be a multiple of 4, and `offset + 8` must not be more than the control-code size. All the symbols that name one external variable in a kernel must have the same `st_size`.
 - An argument can have more than one patch site. An argument index can have no patch site.
 - The number of kernel arguments is the highest argument index plus 1.
 - The loader refuses a kernel that has no PDI patch site.
 
-Preemption save and restore sections, control packets and scalar arguments are not supported.
+Preemption save and restore sections, control packets and scalar arguments are not supported. A relocation that names a control packet looks like an external variable to ROCR, so the load fails unless the application defines a variable with that name.
 
 ### 7.5 Patch schemes
 
-The loader applies the PDI patches once, at load time, to its copy of the control code. The driver applies the argument patches at dispatch time, on a new copy of that control code. It never patches the loader's copy, because `ShimDma48` adds to the existing value.
+The loader applies the PDI and external variable patches once, at load time, to its copy of the control code. The driver applies the argument patches at dispatch time, on a new copy of that control code. It never patches the loader's copy, because `ShimDma48` adds to the existing value.
 
-**`Address64` (8), for a PDI.** The loader writes the 64-bit device address of the PDI at each patch site of that PDI. The low word is first. This is a store, not an addition.
+**`Address64` (8), for a PDI or an external variable.** The loader writes the 64-bit device address of the PDI, or the device address of the external variable plus the addend, at each patch site. The low word is first. This is a store, not an addition.
 
 **`ShimDma48` (5), for arguments.** The patch site is three 32-bit words, `w[0]`, `w[1]` and `w[2]`. The driver does these steps, with `addr = argument address + addend`:
 
@@ -364,6 +365,26 @@ buffer arguments.
 
 The kernel descriptor is a host structure. It is not part of the file format. Its first word is a version (`kAieKernelDescriptorVersion`, now 1). The driver examines this word to find a handle that is not valid.
 
+### 9.3 External variables
+
+A full ELF can name a buffer that it does not contain (section 7.4). The control scratchpad is an example: the host writes runtime parameters into it, and the command processor copies them to the AIE array when the control code runs `create_scratchpad`.
+
+The application allocates the buffer and defines it before it loads the code object:
+
+```cpp
+void* scratchpad = nullptr;
+hsa_amd_memory_pool_allocate(dev_pool, size, 0, &scratchpad);  // the AIE agent's device-SVM pool
+hsa_executable_agent_global_variable_define(executable, aie_agent, "scratch-pad-ctrl", scratchpad);
+hsa_executable_load_agent_code_object(executable, aie_agent, reader, nullptr, nullptr);
+```
+
+This is the same mechanism that a GPU code object uses for an undefined symbol.
+
+- The variable must be defined for the AIE agent before the load. If it is not, the load fails with `HSA_STATUS_ERROR_VARIABLE_UNDEFINED`.
+- The `st_size` bytes from the defined address must be in one allocation that the AIE agent can reach. If they are not, the load fails with `HSA_STATUS_ERROR_INVALID_ALLOCATION`. For the control scratchpad, the allocation must come from the device-SVM pool. The command processor cannot read the scratchpad from other memory, and the dispatch fails.
+- The loader writes the device address of the variable into the control code. The variable must not be freed before the executable is destroyed.
+- At each dispatch, the driver flushes the `st_size` bytes from the CPU cache and keeps the allocation resident. The host can change the contents between dispatches with ordinary stores.
+
 ## 10. Load procedure in ROCR
 
 `hsa_executable_load_agent_code_object` does these steps for an AIE agent:
@@ -376,7 +397,7 @@ The kernel descriptor is a host structure. It is not part of the file format. It
    1. Refuse a `kind` of 2 or more.
    2. Refuse `num_cols` of 0, or more than the agent columns.
    3. `PdiInsts`: refuse an entry with no PDI. Copy the instruction blob and the PDI blob into device memory. Get the BO handles.
-   4. `FullElf`: parse the full ELF (one time for each distinct blob). Find the kernel by name. Examine `kernarg_size`. Copy each different PDI into device memory and get its BO handle. Keep the control code and the patch sites in host memory.
+   4. `FullElf`: parse the full ELF (one time for each distinct blob). Find the kernel by name. Examine `kernarg_size`. Copy each different PDI into device memory and get its BO handle. Find each external variable among the variables that the application defined (section 9.3), write its device address into the control code and get its BO handle. Keep the control code and the patch sites in host memory.
 6. Publish the symbols. Before this step, a failure releases all device memory from steps 5.3 and 5.4.
 
 Device memory for blobs comes from the device SVM region (the device heap). Each buffer has 64-byte alignment. The loader copies each distinct blob one time. It identifies a blob by its address and size in the hsaco buffer. This does not apply to the PDIs of a full-ELF kernel: the loader copies them out of the nested ELF and places them once per kernel, so two kernels that load the same PDI each get their own copy. After the copy, the loader flushes the CPU cache for the buffer one time.
@@ -416,6 +437,8 @@ The table shows which tool examines each rule. "—" means that the tool does no
 | Full ELF groups without `.ctrltext` are not kernels | Yes | Yes | — |
 | Full ELF patch sites in range | Yes | — | — |
 | Full ELF PDI section names are not ambiguous | Yes | — | — |
+| Full ELF external variables use `Address64` and one `st_size` | Yes | — | — |
+| Full ELF external variables are defined, in reachable memory | Yes (at load) | — | — |
 | Full ELF kernel names are unique | Yes | — | — |
 
 ## 12. The packer command line

projects/rocr-runtime/rocrtst/suites/aie/CMakeLists.txt

@@ -262,6 +262,15 @@ function(add_aie_kernel name source)
   set(${name}_INSTS_BIN "${insts}" PARENT_SCOPE)
   set(${name}_DESIGN_PDI "${pdi}" PARENT_SCOPE)
 
+  add_aie_full_elf(${name} "${source}" ${ARGN})
+  set(${name}_FULL_ELF "${${name}_FULL_ELF}" PARENT_SCOPE)
+  set(${name}_FULL_ELF_TARGET "${${name}_FULL_ELF_TARGET}" PARENT_SCOPE)
+endfunction()
+
+# add_aie_full_elf(<name> <source.py>) is the full-ELF half of add_aie_kernel, for a design
+# that only has a full-ELF build. It sets <name>_FULL_ELF and <name>_FULL_ELF_TARGET as above.
+function(add_aie_full_elf name source)
+  get_filename_component(source_name "${source}" NAME)
   set(elf_build_dir "${CMAKE_CURRENT_BINARY_DIR}/kernel_full_elf_${name}")
   set(elf_mlir "${elf_build_dir}/aie.mlir")
   set(elf "${elf_build_dir}/aie.elf")
@@ -316,6 +325,14 @@ add_aie_kernel(vsadd "${CMAKE_CURRENT_SOURCE_DIR}/vector_scalar_add.py")
 # value is compiled into the kernel and into the assertion that checks it.
 set(VSMUL_SCALE 3)
 add_aie_kernel(vsmul "${CMAKE_CURRENT_SOURCE_DIR}/vector_scalar_mul.py" "--scale=${VSMUL_SCALE}")
+# vector_scalar_add_param: like vsadd, but adds a value the host writes into the control
+# scratchpad. Full ELF and npu2 only, which the scratchpad needs.
+if(NPU_TARGET STREQUAL "npu2")
+  add_aie_full_elf(vsparam "${CMAKE_CURRENT_SOURCE_DIR}/vector_scalar_add_param.py")
+else()
+  set(vsparam_FULL_ELF "${CMAKE_CURRENT_BINARY_DIR}/kernel_full_elf_vsparam/aie.elf")
+  set(vsparam_FULL_ELF_TARGET "")
+endif()
 
 # ---------------------------------------------------------------------------
 # Unified hsaco packaging
@@ -428,6 +445,23 @@ if(AIE_FULL_ELF_SUPPORTED)
   add_custom_target(vsmul_elf_hsaco_artifact DEPENDS "${vsmul_ELF_HSACO}")
 endif()
 
+# The scratchpad design's full ELF, packaged the same way, for the external-variable tests.
+set(vsparam_ELF_HSACO "${CMAKE_CURRENT_BINARY_DIR}/vsparam_elf.hsaco")
+if(vsparam_FULL_ELF_TARGET)
+  add_custom_command(
+    OUTPUT "${vsparam_ELF_HSACO}"
+    COMMAND "${AIE_HSACO}"
+            --hsaco "${vsparam_ELF_HSACO}"
+            --arch "${AIE_HSACO_ARCH}"
+            --kernel "elf:${vsparam_FULL_ELF}:0:1"
+    # Depend on the target, not its output files: see the vsadd_HSACO rule above for why.
+    DEPENDS ${vsparam_FULL_ELF_TARGET}
+    COMMENT "Packaging vsparam full-ELF artifact into a unified AIE hsaco"
+    VERBATIM
+  )
+  add_custom_target(vsparam_elf_hsaco_artifact DEPENDS "${vsparam_ELF_HSACO}")
+endif()
+
 # ---------------------------------------------------------------------------
 # test_dispatch executable
 # ---------------------------------------------------------------------------
@@ -464,6 +498,7 @@ target_compile_definitions(test_dispatch
     MUL_HSACO_PATH=${vsmul_HSACO}
     MUL_HSACO_KERNEL_NAME="${VSMUL_HSACO_KERNEL_NAME}"
     MUL_ELF_HSACO_PATH=${vsmul_ELF_HSACO}
+    PARAM_ELF_HSACO_PATH=${vsparam_ELF_HSACO}
 )
 target_link_libraries(test_dispatch
   PRIVATE
@@ -482,6 +517,9 @@ endif()
 if(AIE_FULL_ELF_SUPPORTED)
   add_dependencies(test_dispatch vsadd_elf_hsaco_artifact vsmul_elf_hsaco_artifact)
 endif()
+if(TARGET vsparam_elf_hsaco_artifact)
+  add_dependencies(test_dispatch vsparam_elf_hsaco_artifact)
+endif()
 
 gtest_discover_tests(test_dispatch)
 
@@ -513,6 +551,9 @@ target_link_libraries(elf_parse_test PRIVATE GTest::gtest_main hsa-runtime64::hs
 if(vsadd_FULL_ELF_TARGET)
   add_dependencies(elf_parse_test ${vsadd_FULL_ELF_TARGET})
 endif()
+if(vsparam_FULL_ELF_TARGET)
+  add_dependencies(elf_parse_test ${vsparam_FULL_ELF_TARGET})
+endif()
 add_test(NAME elf_parse_test COMMAND elf_parse_test WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
 
 # ---------------------------------------------------------------------------

projects/rocr-runtime/rocrtst/suites/aie/elf_parse_test.cc

@@ -6,6 +6,7 @@
 
 #include <gtest/gtest.h>
 
+#include <algorithm>
 #include <cstring>
 #include <elf.h>
 #include <fstream>
@@ -489,6 +490,121 @@ TEST(AieElfParse, FindsStringTablesThroughSymbolTableLinks) {
   }
 }
 
+// The scratchpad design: like kElf, but its control code also takes the address of the control
+// scratchpad, which the ELF names but does not carry.
+constexpr const char* kParamElf = "kernel_full_elf_vsparam/aie.elf";
+constexpr const char* kScratchpad = "scratch-pad-ctrl";
+
+/**
+ * @brief Finds the .rela.dyn relocations whose symbol has a given name.
+ *
+ * @param image ELF contents
+ * @param name symbol name
+ * @param st_size receives the symbol's st_size, if any relocation names it
+ * @return the indices of those relocations in .rela.dyn
+ */
+std::vector<std::uint32_t> RelocationsNaming(const std::vector<std::uint8_t>& image,
+                                             const std::string& name, std::uint32_t* st_size) {
+  std::vector<std::uint32_t> found;
+  const std::uint32_t rela = FindSection(image, ".rela.dyn");
+  const std::uint32_t dynsym = FindSection(image, ".dynsym");
+  const std::uint32_t dynstr = FindSection(image, ".dynstr");
+  if (rela == 0 || dynsym == 0 || dynstr == 0) return found;
+  const Elf32_Shdr rela_sh = SectionHeader(image, rela);
+  const Elf32_Shdr sym_sh = SectionHeader(image, dynsym);
+  const Elf32_Shdr str_sh = SectionHeader(image, dynstr);
+  for (std::uint32_t r = 0; r < rela_sh.sh_size / sizeof(Elf32_Rela); ++r) {
+    Elf32_Rela rel{};
+    std::memcpy(&rel, image.data() + rela_sh.sh_offset + r * sizeof(Elf32_Rela), sizeof(rel));
+    Elf32_Sym sym{};
+    std::memcpy(&sym, image.data() + sym_sh.sh_offset + ELF32_R_SYM(rel.r_info) * sizeof(Elf32_Sym),
+                sizeof(sym));
+    const auto* sym_name =
+        reinterpret_cast<const char*>(image.data() + str_sh.sh_offset + sym.st_name);
+    if (name != sym_name) continue;
+    found.push_back(r);
+    *st_size = sym.st_size;
+  }
+  return found;
+}
+
+TEST(AieElfParse, ParsesScratchpadAsExternalSymbol) {
+  // The scratchpad is neither an argument nor a PDI. The reader records its sites and size for the
+  // loader to resolve.
+  const auto image = ReadFile(kParamElf);
+  if (image.empty()) GTEST_SKIP() << "scratchpad full-ELF artifact not built";
+
+  std::uint32_t st_size = 0;
+  const auto relocs = RelocationsNaming(image, kScratchpad, &st_size);
+  ASSERT_FALSE(relocs.empty()) << "artifact has no " << kScratchpad << " relocation";
+
+  std::map<std::string, rocr::AMD::aie_elf::Kernel> kernels;
+  std::string error;
+  ASSERT_EQ(ParseImage(image, &kernels, &error), HSA_STATUS_SUCCESS) << error;
+  const auto& k = kernels.begin()->second;
+
+  ASSERT_EQ(k.external_symbols.size(), 1u);
+  const auto it = k.external_symbols.find(kScratchpad);
+  ASSERT_NE(it, k.external_symbols.end());
+  EXPECT_EQ(it->second.size, st_size);
+  ASSERT_EQ(it->second.sites.size(), relocs.size());
+  for (std::size_t i = 0; i < relocs.size(); ++i) {
+    EXPECT_EQ(it->second.sites[i].offset, RelocationOffset(image, relocs[i]));
+  }
+  // Its sites are not PDI sites, and the argument count is unchanged.
+  for (const auto& pdi : k.pdis) {
+    for (const auto& site : it->second.sites) {
+      EXPECT_EQ(std::count(pdi.patch_offsets.begin(), pdi.patch_offsets.end(), site.offset), 0);
+    }
+  }
+  EXPECT_EQ(k.num_args(), 2u);
+}
+
+TEST(AieElfParse, RejectsExternalSymbolWithAnotherScheme) {
+  // Only Address64 is defined for an external symbol.
+  auto image = ReadFile(kParamElf);
+  if (image.empty()) GTEST_SKIP() << "scratchpad full-ELF artifact not built";
+
+  std::uint32_t st_size = 0;
+  const auto relocs = RelocationsNaming(image, kScratchpad, &st_size);
+  ASSERT_FALSE(relocs.empty());
+  const std::uint32_t rela = FindSection(image, ".rela.dyn");
+  Elf32_Rela rel{};
+  std::memcpy(&rel,
+              image.data() + SectionHeader(image, rela).sh_offset + relocs[0] * sizeof(Elf32_Rela),
+              sizeof(rel));
+  // Same symbol, ShimDma48 instead of Address64.
+  ASSERT_TRUE(MutateRelocation(image, relocs[0], ELF32_R_SYM(rel.r_info), 5));
+
+  std::map<std::string, rocr::AMD::aie_elf::Kernel> kernels;
+  std::string error;
+  EXPECT_NE(ParseImage(image, &kernels, &error), HSA_STATUS_SUCCESS)
+      << "an external symbol with a ShimDma48 patch was accepted";
+  EXPECT_NE(error.find("external symbol"), std::string::npos)
+      << "rejected for the wrong reason: " << error;
+}
+
+TEST(AieElfParse, RejectsExternalSymbolPatchSiteOutOfRange) {
+  // The 8-byte store must lie inside the control code.
+  auto image = ReadFile(kParamElf);
+  if (image.empty()) GTEST_SKIP() << "scratchpad full-ELF artifact not built";
+
+  std::map<std::string, rocr::AMD::aie_elf::Kernel> kernels;
+  std::string error;
+  ASSERT_EQ(ParseImage(image, &kernels, &error), HSA_STATUS_SUCCESS) << error;
+  const std::size_t ctrl_code_size = kernels.begin()->second.ctrl_code.size();
+
+  std::uint32_t st_size = 0;
+  const auto relocs = RelocationsNaming(image, kScratchpad, &st_size);
+  ASSERT_FALSE(relocs.empty());
+  ASSERT_TRUE(
+      SetRelocationOffset(image, relocs[0], static_cast<std::uint32_t>(ctrl_code_size - 4)));
+  EXPECT_NE(ParseImage(image, &kernels, &error), HSA_STATUS_SUCCESS)
+      << "an out-of-range external symbol patch site was accepted";
+  EXPECT_NE(error.find("external symbol patch site"), std::string::npos)
+      << "rejected for the wrong reason: " << error;
+}
+
 TEST(AieElfPatch, ShimDma48IsAdditive) {
   // Two applications must not equal one: this is why the driver patches a pristine copy.
   std::uint32_t once[3] = {0, 0x1000, 0};