samples/custom_dispatch/cpu/mlp_plugin/mlp_plugin.c - 3p/openxla/iree - Git at Google

 // Copyright 2024 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 // Demonstrates an mlp example with the implementation of MLP provided
 // using system linked plugin exporting a single `mlp_external`
 // function.  See samples/custom_dispatch/cpu/plugin/system_plugin.c
 // for more information about system plugins and their caveats.

 #include <inttypes.h>
 #include <stdio.h>

 // The only header required from IREE:
 #include "iree/hal/local/executable_plugin.h"

 // Stateful plugin instance.
 // There may be multiple of these in a process at a time, each with its own
 // load/unload pairing. We pass a pointer to this to all import calls via the
 // context argument.
 typedef struct {
   iree_hal_executable_plugin_allocator_t host_allocator;
   FILE* file;
 } mlp_plugin_t;

 // Helper function to resolve index [i][j] into location for given strides.
 size_t get_index(size_t i, size_t j, size_t offset, size_t stride0,
                  size_t stride1) {
   return offset + i * stride0 + j * stride1;
 }

 // `ret = mlp(lhs, rhs)`
 //
 // Conforms to ABI:
 // #hal.pipeline.layout<push_constants = 1, sets = [
 //   <0, bindings = [
 //       <0, storage_buffer, ReadOnly>,
 //       <1, storage_buffer, ReadOnly>,
 //       <2, storage_buffer>
 //   ]>
 // ]>
 // With a workgroup size of 64x1x1.
 //
 // |context| is whatever was set in out_fn_contexts. This could point to shared
 // state or each import can have its own context (pointer into some JIT lookup
 // table, etc). In this sample we pass the sample plugin pointer to all imports.
 //
 // |params_ptr| points to a packed struct of all results followed by all args
 // using native arch packing/alignment rules. Results should be set before
 // returning.
 //
 // Expects a return of 0 on success and any other value indicates failure.
 // Try not to fail!
 static int mlp_external(void* params_ptr, void* context, void* reserved) {
   mlp_plugin_t* plugin = (mlp_plugin_t*)context;
   typedef struct {
     const float* restrict lhs;
     const float* restrict lhs_aligned;
     size_t lhs_offset;
     size_t lhs_size0;
     size_t lhs_size1;
     size_t lhs_stride0;
     size_t lhs_stride1;
     const float* restrict rhs;
     const float* restrict rhs_aligned;
     size_t rhs_offset;
     size_t rhs_size0;
     size_t rhs_size1;
     size_t rhs_stride0;
     size_t rhs_stride1;
     float* restrict result;
     float* restrict result_aligned;
     size_t result_offset;
     size_t result_size0;
     size_t result_size1;
     size_t result_stride0;
     size_t result_stride1;
     int32_t M;
     int32_t N;
     int32_t K;
   } params_t;
   const params_t* params = (const params_t*)params_ptr;
   fprintf(plugin->file, "[Plugin]: M = %d, N = %d, K = %d\n", params->M,
           params->N, params->K);
   for (int32_t i = 0; i < params->M; i++) {
     for (int32_t j = 0; j < params->N; j++) {
       float curr_result = 0.0;
       for (int32_t k = 0; k < params->K; k++) {
         size_t lhs_index = get_index(i, k, params->lhs_offset,
                                      params->lhs_stride0, params->lhs_stride1);
         size_t rhs_index = get_index(k, j, params->rhs_offset,
                                      params->rhs_stride0, params->rhs_stride1);
         curr_result += params->lhs[lhs_index] * params->rhs[rhs_index];
       }
       curr_result = curr_result < 0.0 ? 0.0 : curr_result;
       size_t result_index =
           get_index(i, j, params->result_offset, params->result_stride0,
                     params->result_stride1);
       params->result[result_index] = curr_result;
     }
   }
   return 0;
 }

 // Called once for each plugin load and paired with a future call to unload.
 // Even in standalone mode we could allocate using environment->host_allocator,
 // set an out_self pointer, and parse parameters but here in system mode we can
 // do whatever we want.
 //
 // If any state is required it should be allocated and stored in |out_self|.
 // This self value will be passed to all future calls related to the particular
 // instance. Note that there may be multiple instances of a plugin in any
 // particular process and this must be thread-safe.
 static iree_hal_executable_plugin_status_t mlp_plugin_load(
     const iree_hal_executable_plugin_environment_v0_t* environment,
     size_t param_count, const iree_hal_executable_plugin_string_pair_t* params,
     void** out_self) {
   // Allocate the plugin state.
   mlp_plugin_t* plugin = NULL;
   iree_hal_executable_plugin_status_t status =
       iree_hal_executable_plugin_allocator_malloc(
           environment->host_allocator, sizeof(*plugin), (void**)&plugin);
   if (status) return status;
   plugin->host_allocator = environment->host_allocator;

   // "Open standard out" simulating us doing some syscalls or other expensive
   // stateful/side-effecting things.
   plugin->file = stdout;

   // Pass back the plugin instance that'll be passed to resolve.
   *out_self = plugin;
   return iree_hal_executable_plugin_ok_status();
 }

 // Called to free any plugin state allocated in load.
 static void mlp_plugin_unload(void* self) {
   mlp_plugin_t* plugin = (mlp_plugin_t*)self;
   iree_hal_executable_plugin_allocator_t host_allocator =
       plugin->host_allocator;

   // "Close standard out" simulating us doing some syscalls and other expensive
   // stateful/side-effecting things.
   fflush(plugin->file);
   plugin->file = NULL;

   // Free the plugin state using the same allocator it came from.
   iree_hal_executable_plugin_allocator_free(host_allocator, plugin);
 }

 // Called to resolve one or more imports by symbol name.
 // See the plugin API header for more information. Note that some of the
 // functions may already be resolved and some may be optional.
 static iree_hal_executable_plugin_status_t mlp_plugin_resolve(
     void* self, const iree_hal_executable_plugin_resolve_params_v0_t* params,
     iree_hal_executable_plugin_resolution_t* out_resolution) {
   mlp_plugin_t* plugin = (mlp_plugin_t*)self;
   *out_resolution = 0;
   bool any_required_not_found = false;
   for (size_t i = 0; i < params->count; ++i) {
     if (params->out_fn_ptrs[i]) continue;
     const char* symbol_name = params->symbol_names[i];
     bool is_optional =
         iree_hal_executable_plugin_import_is_optional(symbol_name);
     if (is_optional) ++symbol_name;
     if (iree_hal_executable_plugin_strcmp(symbol_name, "mlp_external") == 0) {
       params->out_fn_ptrs[i] = mlp_external;
       params->out_fn_contexts[i] =
           plugin;  // passing plugin to each import call
     } else {
       if (is_optional) {
         *out_resolution |=
             IREE_HAL_EXECUTABLE_PLUGIN_RESOLUTION_MISSING_OPTIONAL;
       } else {
         any_required_not_found = true;
       }
     }
   }
   return any_required_not_found
              ? iree_hal_executable_plugin_status_from_code(
                    IREE_HAL_EXECUTABLE_PLUGIN_STATUS_NOT_FOUND)
              : iree_hal_executable_plugin_ok_status();
 }

 // Exported on the shared library and used by the runtime to query the plugin
 // interface. When statically linking the plugin this is just a function that
 // can be called and can have any name to allow for multiple plugins. When
 // dynamically linking the exported symbol must be exactly this with no C++
 // name mangling.
 IREE_HAL_EXECUTABLE_PLUGIN_EXPORT const iree_hal_executable_plugin_header_t**
 iree_hal_executable_plugin_query(
     iree_hal_executable_plugin_version_t max_version, void* reserved) {
   static const iree_hal_executable_plugin_header_t header = {
       // Declares what library version is present: newer runtimes may support
       // loading older plugins but newer plugins cannot load on older runtimes.
       .version = IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST,
       // Name and description are used for tracing/logging/diagnostics.
       .name = "sample_system",
       .description =
           "system plugin sample "
           "(custom_dispatch/cpu/plugin/mlp_plugin.c)",
       .features = 0,
       // Let the runtime know what sanitizer this plugin was compiled with.
       .sanitizer = IREE_HAL_EXECUTABLE_PLUGIN_SANITIZER_KIND,
   };
   static const iree_hal_executable_plugin_v0_t plugin = {
       .header = &header,
       .load = mlp_plugin_load,
       .unload = mlp_plugin_unload,
       .resolve = mlp_plugin_resolve,
   };
   return max_version <= IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST
              ? (const iree_hal_executable_plugin_header_t**)&plugin
              : NULL;
 }
	// Copyright 2024 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	// Demonstrates an mlp example with the implementation of MLP provided
	// using system linked plugin exporting a single `mlp_external`
	// function. See samples/custom_dispatch/cpu/plugin/system_plugin.c
	// for more information about system plugins and their caveats.

	#include <inttypes.h>
	#include <stdio.h>

	// The only header required from IREE:
	#include "iree/hal/local/executable_plugin.h"

	// Stateful plugin instance.
	// There may be multiple of these in a process at a time, each with its own
	// load/unload pairing. We pass a pointer to this to all import calls via the
	// context argument.
	typedef struct {
	iree_hal_executable_plugin_allocator_t host_allocator;
	FILE* file;
	} mlp_plugin_t;

	// Helper function to resolve index [i][j] into location for given strides.
	size_t get_index(size_t i, size_t j, size_t offset, size_t stride0,
	size_t stride1) {
	return offset + i * stride0 + j * stride1;
	}

	// `ret = mlp(lhs, rhs)`
	//
	// Conforms to ABI:
	// #hal.pipeline.layout<push_constants = 1, sets = [
	// <0, bindings = [
	// <0, storage_buffer, ReadOnly>,
	// <1, storage_buffer, ReadOnly>,
	// <2, storage_buffer>
	// ]>
	// ]>
	// With a workgroup size of 64x1x1.
	//
	// \|context\| is whatever was set in out_fn_contexts. This could point to shared
	// state or each import can have its own context (pointer into some JIT lookup
	// table, etc). In this sample we pass the sample plugin pointer to all imports.
	//
	// \|params_ptr\| points to a packed struct of all results followed by all args
	// using native arch packing/alignment rules. Results should be set before
	// returning.
	//
	// Expects a return of 0 on success and any other value indicates failure.
	// Try not to fail!
	static int mlp_external(void* params_ptr, void* context, void* reserved) {
	mlp_plugin_t* plugin = (mlp_plugin_t*)context;
	typedef struct {
	const float* restrict lhs;
	const float* restrict lhs_aligned;
	size_t lhs_offset;
	size_t lhs_size0;
	size_t lhs_size1;
	size_t lhs_stride0;
	size_t lhs_stride1;
	const float* restrict rhs;
	const float* restrict rhs_aligned;
	size_t rhs_offset;
	size_t rhs_size0;
	size_t rhs_size1;
	size_t rhs_stride0;
	size_t rhs_stride1;
	float* restrict result;
	float* restrict result_aligned;
	size_t result_offset;
	size_t result_size0;
	size_t result_size1;
	size_t result_stride0;
	size_t result_stride1;
	int32_t M;
	int32_t N;
	int32_t K;
	} params_t;
	const params_t* params = (const params_t*)params_ptr;
	fprintf(plugin->file, "[Plugin]: M = %d, N = %d, K = %d\n", params->M,
	params->N, params->K);
	for (int32_t i = 0; i < params->M; i++) {
	for (int32_t j = 0; j < params->N; j++) {
	float curr_result = 0.0;
	for (int32_t k = 0; k < params->K; k++) {
	size_t lhs_index = get_index(i, k, params->lhs_offset,
	params->lhs_stride0, params->lhs_stride1);
	size_t rhs_index = get_index(k, j, params->rhs_offset,
	params->rhs_stride0, params->rhs_stride1);
	curr_result += params->lhs[lhs_index] * params->rhs[rhs_index];
	}
	curr_result = curr_result < 0.0 ? 0.0 : curr_result;
	size_t result_index =
	get_index(i, j, params->result_offset, params->result_stride0,
	params->result_stride1);
	params->result[result_index] = curr_result;
	}
	}
	return 0;
	}

	// Called once for each plugin load and paired with a future call to unload.
	// Even in standalone mode we could allocate using environment->host_allocator,
	// set an out_self pointer, and parse parameters but here in system mode we can
	// do whatever we want.
	//
	// If any state is required it should be allocated and stored in \|out_self\|.
	// This self value will be passed to all future calls related to the particular
	// instance. Note that there may be multiple instances of a plugin in any
	// particular process and this must be thread-safe.
	static iree_hal_executable_plugin_status_t mlp_plugin_load(
	const iree_hal_executable_plugin_environment_v0_t* environment,
	size_t param_count, const iree_hal_executable_plugin_string_pair_t* params,
	void** out_self) {
	// Allocate the plugin state.
	mlp_plugin_t* plugin = NULL;
	iree_hal_executable_plugin_status_t status =
	iree_hal_executable_plugin_allocator_malloc(
	environment->host_allocator, sizeof(plugin), (void*)&plugin);
	if (status) return status;
	plugin->host_allocator = environment->host_allocator;

	// "Open standard out" simulating us doing some syscalls or other expensive
	// stateful/side-effecting things.
	plugin->file = stdout;

	// Pass back the plugin instance that'll be passed to resolve.
	*out_self = plugin;
	return iree_hal_executable_plugin_ok_status();
	}

	// Called to free any plugin state allocated in load.
	static void mlp_plugin_unload(void* self) {
	mlp_plugin_t* plugin = (mlp_plugin_t*)self;
	iree_hal_executable_plugin_allocator_t host_allocator =
	plugin->host_allocator;

	// "Close standard out" simulating us doing some syscalls and other expensive
	// stateful/side-effecting things.
	fflush(plugin->file);
	plugin->file = NULL;

	// Free the plugin state using the same allocator it came from.
	iree_hal_executable_plugin_allocator_free(host_allocator, plugin);
	}

	// Called to resolve one or more imports by symbol name.
	// See the plugin API header for more information. Note that some of the
	// functions may already be resolved and some may be optional.
	static iree_hal_executable_plugin_status_t mlp_plugin_resolve(
	void* self, const iree_hal_executable_plugin_resolve_params_v0_t* params,
	iree_hal_executable_plugin_resolution_t* out_resolution) {
	mlp_plugin_t* plugin = (mlp_plugin_t*)self;
	*out_resolution = 0;
	bool any_required_not_found = false;
	for (size_t i = 0; i < params->count; ++i) {
	if (params->out_fn_ptrs[i]) continue;
	const char* symbol_name = params->symbol_names[i];
	bool is_optional =
	iree_hal_executable_plugin_import_is_optional(symbol_name);
	if (is_optional) ++symbol_name;
	if (iree_hal_executable_plugin_strcmp(symbol_name, "mlp_external") == 0) {
	params->out_fn_ptrs[i] = mlp_external;
	params->out_fn_contexts[i] =
	plugin; // passing plugin to each import call
	} else {
	if (is_optional) {
	*out_resolution \|=
	IREE_HAL_EXECUTABLE_PLUGIN_RESOLUTION_MISSING_OPTIONAL;
	} else {
	any_required_not_found = true;
	}
	}
	}
	return any_required_not_found
	? iree_hal_executable_plugin_status_from_code(
	IREE_HAL_EXECUTABLE_PLUGIN_STATUS_NOT_FOUND)
	: iree_hal_executable_plugin_ok_status();
	}

	// Exported on the shared library and used by the runtime to query the plugin
	// interface. When statically linking the plugin this is just a function that
	// can be called and can have any name to allow for multiple plugins. When
	// dynamically linking the exported symbol must be exactly this with no C++
	// name mangling.
	IREE_HAL_EXECUTABLE_PLUGIN_EXPORT const iree_hal_executable_plugin_header_t**
	iree_hal_executable_plugin_query(
	iree_hal_executable_plugin_version_t max_version, void* reserved) {
	static const iree_hal_executable_plugin_header_t header = {
	// Declares what library version is present: newer runtimes may support
	// loading older plugins but newer plugins cannot load on older runtimes.
	.version = IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST,
	// Name and description are used for tracing/logging/diagnostics.
	.name = "sample_system",
	.description =
	"system plugin sample "
	"(custom_dispatch/cpu/plugin/mlp_plugin.c)",
	.features = 0,
	// Let the runtime know what sanitizer this plugin was compiled with.
	.sanitizer = IREE_HAL_EXECUTABLE_PLUGIN_SANITIZER_KIND,
	};
	static const iree_hal_executable_plugin_v0_t plugin = {
	.header = &header,
	.load = mlp_plugin_load,
	.unload = mlp_plugin_unload,
	.resolve = mlp_plugin_resolve,
	};
	return max_version <= IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST
	? (const iree_hal_executable_plugin_header_t**)&plugin
	: NULL;
	}