ukernel/importer.c

/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "iree/builtins/ukernel/api.h"
#include "iree/hal/local/executable_loader.h"
#include "iree/hal/local/executable_plugin.h"

// Importer entry point wrapping the actual ukernel.
static int iree_uk_importer_mmt4d(void* params_ptr, void* context,
                                  void* reserved) {
  iree_uk_mmt4d((const iree_uk_mmt4d_params_t*)params_ptr);
  return 0;
}

// Called to resolve the iree_uk_mmt4d import by symbol name.
static iree_status_t iree_uk_importer_resolve(
    void* self, iree_host_size_t count, const char* const* symbol_names,
    void** out_fn_ptrs, void** out_fn_contexts,
    iree_hal_executable_import_resolution_t* out_resolution) {
  *out_resolution = 0;
  bool any_required_not_found = true;
  for (size_t i = 0; i < count; ++i) {
    if (out_fn_ptrs[i]) continue;
    bool found = iree_hal_executable_plugin_strcmp(symbol_names[i],
                                                   "iree_uk_mmt4d") == 0;
    if (found) {
      out_fn_ptrs[i] = (void*)iree_uk_importer_mmt4d;
      out_fn_contexts[i] = NULL;
      any_required_not_found = false;
    } else {
      break;
    }
  }
  return any_required_not_found ? iree_make_status(IREE_STATUS_NOT_FOUND)
                                : iree_ok_status();
}

iree_hal_executable_import_provider_t importer_query(void) {
  static iree_hal_executable_import_provider_t importer = {
      .resolve = iree_uk_importer_resolve,
  };
  return importer;
}