[compiler][runtime] Add DMA core pointer to executable format (#51)

Once dma core code starts to be generated, we'll be required to find and
call the code using the DMA core. While not implemented yet, we already
use the field to distinguish between kernels generated by LLVM and xDSL
to perform the different workgroup and subgroup distributions.

codegen/compiler/src/Quidditch/Target/LibraryBuilder.cpp

@@ -182,18 +182,19 @@ makeStageLocationTableType(llvm::LLVMContext &context) {
   return type;
 }
 
-// %struct.iree_hal_executable_export_table_v0_t = type {
+// %struct.quidditch_executable_export_table_v0_t = type {
 //   i32,
-//   i32*,
+//   %struct.iree_hal_executable_dispatch_v0_t**,
 //   %struct.iree_hal_executable_dispatch_attrs_v0_t*,
 //   i8**,
 //   i8**,
 //   %struct.iree_hal_executable_source_location_v0_t*,
 //   %struct.iree_hal_executable_stage_location_table_v0_t*,
+//   %struct.iree_hal_executable_dispatch_v0_t*,
 // }
 static llvm::StructType *makeExportTableType(llvm::LLVMContext &context) {
   if (auto *existingType = llvm::StructType::getTypeByName(
-          context, "iree_hal_executable_export_table_v0_t")) {
+          context, "quidditch_executable_export_table_v0_t")) {
     return existingType;
   }
   auto *i32Type = llvm::IntegerType::getInt32Ty(context);
@@ -212,8 +213,9 @@ static llvm::StructType *makeExportTableType(llvm::LLVMContext &context) {
           i8PtrType->getPointerTo(),
           sourceLocationType->getPointerTo(),
           stageLocationTableType->getPointerTo(),
+          dispatchFunctionType->getPointerTo()->getPointerTo(),
       },
-      "iree_hal_executable_export_table_v0_t",
+      "quidditch_executable_export_table_v0_t",
       /*isPacked=*/false);
   return type;
 }
@@ -493,7 +495,7 @@ LibraryBuilder::buildLibraryV0ExportTable(std::string libraryName) {
   // iree_hal_executable_export_table_v0_t::ptrs
   SmallVector<llvm::Constant *> exportPtrValues;
   for (auto dispatch : exports)
-    exportPtrValues.push_back(dispatch.func);
+    exportPtrValues.push_back(dispatch.compute_func);
   llvm::Constant *exportPtrs = createArrayConstant(
       libraryName + "_funcs", dispatchFunctionType->getPointerTo(),
       exportPtrValues, module);
@@ -618,6 +620,20 @@ LibraryBuilder::buildLibraryV0ExportTable(std::string libraryName) {
     }
   }
 
+  // quidditch_executable_export_table_v0_t::dma_core_ptrs
+  SmallVector<llvm::Constant *> dmaCorePtrs;
+  for (Dispatch &dispatch : exports) {
+    if (!dispatch.dma_func) {
+      dmaCorePtrs.push_back(
+          llvm::Constant::getNullValue(dispatchAttrsType->getPointerTo()));
+      continue;
+    }
+    dmaCorePtrs.push_back(dispatch.dma_func);
+  }
+  llvm::Constant *dmaCorePtrsArray = createArrayConstant(
+      libraryName + "_dma_core_ptrs", dispatchAttrsType->getPointerTo(),
+      dmaCorePtrs, module);
+
   return llvm::ConstantStruct::get(
       exportTableType, {
                            // count=
@@ -634,6 +650,7 @@ LibraryBuilder::buildLibraryV0ExportTable(std::string libraryName) {
                            exportSourceLocations,
                            // stage_locations=
                            exportStageLocations,
+                           dmaCorePtrsArray,
                        });
 }
 

codegen/compiler/src/Quidditch/Target/LibraryBuilder.h

@@ -19,7 +19,7 @@ namespace Quidditch {
 // Defines an `iree_hal_executable_library_v0_t` and builds the runtime metadata
 // structures and query functions.
 //
-// See iree/hal/local/executable_library.h for more information.
+// See Quidditch/executable/executable_library.h for more information.
 //
 // Usage:
 //  LibraryBuilder builder(&module);
@@ -121,9 +121,10 @@ class LibraryBuilder {
   void addExport(llvm::StringRef name, SourceLocation sourceLocation,
                  llvm::SmallVector<SourceLocation> stageLocations,
                  llvm::StringRef tag, DispatchAttrs attrs,
-                 llvm::Function *func) {
+                 llvm::Function *compute_func, llvm::Function *dmaFunc) {
     exports.push_back({name.str(), std::move(sourceLocation),
-                       std::move(stageLocations), tag.str(), attrs, func});
+                       std::move(stageLocations), tag.str(), attrs,
+                       compute_func, dmaFunc});
   }
 
   // Defines a source file embedded in the library.
@@ -167,7 +168,8 @@ class LibraryBuilder {
     llvm::SmallVector<SourceLocation> stageLocations;
     std::string tag;
     DispatchAttrs attrs;
-    llvm::Function *func;
+    llvm::Function *compute_func;
+    llvm::Function *dma_func;
   };
   std::vector<Dispatch> exports;
 

codegen/compiler/src/Quidditch/Target/QuidditchTarget.cpp

@@ -334,11 +334,21 @@ class QuidditchTargetBackend final : public IREE::HAL::TargetBackend {
     auto align16 = llvm::Attribute::getWithAlignment(context, llvm::Align(16));
     for (auto exportOp :
          variantOp.getBlock().getOps<IREE::HAL::ExecutableExportOp>()) {
+      llvm::Function *dmaPointer = nullptr;
       // Find the matching function in the LLVM module.
       auto *llvmFunc =
           llvmModule->getFunction((exportOp.getName() + "$iree_to_xdsl").str());
-      if (!llvmFunc)
+      if (!llvmFunc) {
+        // LLVMCPU kernel rather than xDSL.
         llvmFunc = llvmModule->getFunction(exportOp.getName());
+      } else {
+        // xDSL kernel.
+
+        // TODO: Replace with real DMA pointer. For now just a place holder that
+        //  is not used in the runtime but makes the dispatch recognizeable as
+        //  an xDSL dispatch.
+        dmaPointer = llvmFunc;
+      }
       if (!llvmFunc)
         continue;
 
@@ -365,7 +375,8 @@ class QuidditchTargetBackend final : public IREE::HAL::TargetBackend {
       libraryBuilder.addExport(
           exportOp.getName(), std::move(sourceLocation),
           std::move(stageLocations), /*tag=*/"",
-          Quidditch::LibraryBuilder::DispatchAttrs{localMemorySize}, llvmFunc);
+          Quidditch::LibraryBuilder::DispatchAttrs{localMemorySize}, llvmFunc,
+          dmaPointer);
     }
     auto *queryLibraryFunc =
         libraryBuilder.build("quidditch_" + libraryName + "_library_query");

runtime/runtime/src/Quidditch/dispatch/dispatch.c

@@ -56,10 +56,11 @@ int quidditch_dispatch_enter_worker_loop() {
     snrt_cluster_hw_barrier();
     if (exit) break;
 
+    // Cores not in use may spuriously wake up at any time and need to reenter
+    // the barrier asap.
     iree_hal_executable_workgroup_state_v0_t* workgroupState =
         &configuredWorkgroupState[snrt_cluster_core_idx()];
     if (workgroupState->processor_id == -1) {
-      snrt_cluster_hw_barrier();
       continue;
     }
 

runtime/runtime/src/Quidditch/executable/CMakeLists.txt

@@ -4,8 +4,10 @@ iree_cc_library(
     executable
     HDRS
     "executable.h"
+    "executable_library.h"
     SRCS
     "executable.c"
+    "executable_library.c"
     DEPS
     snRuntimeInterface
     Quidditch::dispatch::dispatch

runtime/runtime/src/Quidditch/executable/executable.c

@@ -65,6 +65,19 @@ static int quidditch_executable_import_thunk_v0(
   return fn_ptr(params, context, reserved);
 }
 
+static const iree_hal_executable_export_table_v0_t*
+quidditch_executable_get_exports(quidditch_executable_t* executable) {
+  if (executable->is_llvm_cpu_executable)
+    return &executable->library.llvmcpu_v0->exports;
+  return (const iree_hal_executable_export_table_v0_t*)&executable->library
+      .quidditch_v0->exports;
+}
+
+static const iree_hal_executable_dispatch_attrs_v0_t*
+quidditch_executable_get_attrs(quidditch_executable_t* executable) {
+  return quidditch_executable_get_exports(executable)->attrs;
+}
+
 iree_status_t quidditch_executable_create(
     const iree_hal_executable_params_t* executable_params,
     const iree_hal_executable_library_header_t** library_header,
@@ -93,11 +106,11 @@ iree_status_t quidditch_executable_create(
                                     &executable->layouts[0], host_allocator,
                                     executable);
     executable->library.header = library_header;
+    executable->is_llvm_cpu_executable = !iree_string_view_equal(
+        executable_params->executable_format, IREE_SV("snitch"));
     executable->identifier = iree_make_cstring_view((*library_header)->name);
-    executable->dispatch_attrs = executable->library.v0->exports.attrs;
+    executable->dispatch_attrs = quidditch_executable_get_attrs(executable);
   }
-  executable->is_llvm = !iree_string_view_equal(
-      executable_params->executable_format, IREE_SV("snitch"));
 
   // Copy executable constants so we own them.
   if (iree_status_is_ok(status) && executable_params->constant_count > 0) {
@@ -111,23 +124,26 @@ iree_status_t quidditch_executable_create(
     executable->environment.constants = target_constants;
   }
 
+  // TODO: Might need to do the below for Quidditch executables as well.
+
   // Resolve imports, if any.
-  if (iree_status_is_ok(status)) {
+  if (iree_status_is_ok(status) && executable->is_llvm_cpu_executable) {
     status = iree_hal_executable_library_initialize_imports(
         &executable->environment, import_provider,
-        &executable->library.v0->imports, quidditch_executable_import_thunk_v0,
-        host_allocator);
+        &executable->library.llvmcpu_v0->imports,
+        quidditch_executable_import_thunk_v0, host_allocator);
   }
 
   // Verify that the library matches the executable params.
-  if (iree_status_is_ok(status)) {
+  if (iree_status_is_ok(status) && executable->is_llvm_cpu_executable) {
     status = iree_hal_executable_library_verify(executable_params,
-                                                executable->library.v0);
+                                                executable->library.llvmcpu_v0);
   }
 
   // Publish the executable sources with the tracing infrastructure.
-  if (iree_status_is_ok(status)) {
-    iree_hal_executable_library_publish_source_files(executable->library.v0);
+  if (iree_status_is_ok(status) && executable->is_llvm_cpu_executable) {
+    iree_hal_executable_library_publish_source_files(
+        executable->library.llvmcpu_v0);
   }
 
   if (iree_status_is_ok(status)) {
@@ -171,19 +187,26 @@ iree_status_t quidditch_executable_issue_dispatch_inline(
   workgroup_state.local_memory = local_memory.data;
   workgroup_state.local_memory_size = (size_t)local_memory.data_length;
 
-  const iree_hal_executable_library_v0_t* library = executable->library.v0;
+  const iree_hal_executable_export_table_v0_t* exports =
+      quidditch_executable_get_exports(executable);
 
-  if (IREE_UNLIKELY(ordinal >= library->exports.count)) {
+  if (IREE_UNLIKELY(ordinal >= exports->count)) {
     return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                             "entry point ordinal out of bounds");
   }
-  iree_hal_executable_dispatch_v0_t kernel = library->exports.ptrs[ordinal];
+  iree_hal_executable_dispatch_v0_t kernel = exports->ptrs[ordinal];
 
   quidditch_dispatch_set_kernel(kernel, &executable->environment,
                                 dispatch_state);
 
+  bool compute_cores_are_workgroups = false;
+  if (executable->is_llvm_cpu_executable ||
+      !((quidditch_executable_export_table_v0_t*)exports)
+           ->dma_core_ptrs[ordinal])
+    compute_cores_are_workgroups = true;
+
   read_csr(mcycle);
-  if (executable->is_llvm) {
+  if (compute_cores_are_workgroups) {
     // LLVM distributes workgroups to compute cores.
     for (uint32_t z = 0; z < workgroup_count_z; ++z) {
       workgroup_state.workgroup_id_z = z;

runtime/runtime/src/Quidditch/executable/executable.h

@@ -8,6 +8,7 @@
 
 #include <iree/hal/local/executable_loader.h>
 
+#include "executable_library.h"
 #include "iree/base/api.h"
 #include "iree/hal/api.h"
 #include "iree/hal/local/executable_library.h"
@@ -44,10 +45,11 @@ typedef struct quidditch_executable_t {
 
   union {
     const iree_hal_executable_library_header_t** header;
-    const iree_hal_executable_library_v0_t* v0;
+    const iree_hal_executable_library_v0_t* llvmcpu_v0;
+    const quidditch_executable_library_v0_t* quidditch_v0;
   } library;
 
-  bool is_llvm;
+  bool is_llvm_cpu_executable;
 
   iree_hal_pipeline_layout_t* layouts[];
 } quidditch_executable_t;