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

 // Copyright 2023 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 a system linked plugin exporting a single `simple_mul_workgroup`
 // function that also prints to stdout. This is not a great idea but shows how
 // plugins can have side-effecting behavior - even if in most cases a standalone
 // plugin can be used with much smaller code size and portability.
 //
 // The major use-case for a system linked plugin is JITs that may compile
 // imports on-demand. Such plugins could either JIT everything on load time
 // or defer JITting to the first call to a particular import. Performing JIT at
 // load time is strongly preferred as it keeps all of the expensive work in one
 // place before the program starts scheduling execution. Deferring will
 // introduce first-run delays and require warmup steps. Since only the imports
 // used by the program are present and most programs use all imports it's almost
 // always going to be better to do things ahead of time.
 //
 // NOTE: when using the system loader all unsafe behavior is allowed: TLS,
 // threads, mutable globals, syscalls, etc. Doing any of those things will
 // likely break in interesting ways as the import functions are called from
 // arbitrary threads concurrently. Be very careful and prefer standalone plugins
 // instead except when debugging/profiling.

 #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;
 } system_plugin_t;

 // `ret = lhs * rhs`
 //
 // Conforms to ABI:
 // #hal.pipeline.layout<constants = 1, bindings = [
 //   #hal.pipeline.binding<storage_buffer, ReadOnly>,
 //   #hal.pipeline.binding<storage_buffer, ReadOnly>,
 //   #hal.pipeline.binding<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 simple_mul_workgroup(void* params_ptr, void* context,
                                 void* reserved) {
   system_plugin_t* plugin = (system_plugin_t*)context;
   typedef struct {
     const float* restrict binding0;
     size_t binding0_offset;
     const float* restrict binding1;
     size_t binding1_offset;
     float* restrict binding2;
     size_t binding2_offset;
     size_t size;
     size_t tid;
     uint32_t processor_id;
     const uint64_t* restrict processor_data;
   } params_t;
   const params_t* params = (const params_t*)params_ptr;
   fprintf(plugin->file, "processor_id=%u\n", params->processor_id);
   if (params->processor_data) {
     fprintf(plugin->file, "processor_data[0]=%" PRIX64 "\n",
             params->processor_data[0]);
   }
   // The operation `iree_codegen.ukernel.generic` always operates
   // on a slice of the inputs to produce a slice of the output,
   // so the loop here just needs to iterate from `0` to `size`,
   // where `size` is the size of the slice to be executed by this call.
   for (size_t i = 0; i < params->size; ++i) {
     params->binding2[params->binding2_offset + i] =
         params->binding0[params->binding0_offset + i] *
         params->binding1[params->binding2_offset + i];
     fprintf(plugin->file, "mul[%zu:%zu](%g * %g = %g)\n", params->tid, i,
             params->binding0[params->binding0_offset + i],
             params->binding1[params->binding1_offset + i],
             params->binding2[params->binding2_offset + i]);
   }
   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 system_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.
   system_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 system_plugin_unload(void* self) {
   system_plugin_t* plugin = (system_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 system_plugin_resolve(
     void* self, const iree_hal_executable_plugin_resolve_params_v0_t* params,
     iree_hal_executable_plugin_resolution_t* out_resolution) {
   system_plugin_t* plugin = (system_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,
                                           "simple_mul_workgroup") == 0) {
       params->out_fn_ptrs[i] = simple_mul_workgroup;
       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/system_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 = system_plugin_load,
       .unload = system_plugin_unload,
       .resolve = system_plugin_resolve,
   };
   return max_version <= IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST
              ? (const iree_hal_executable_plugin_header_t**)&plugin
              : NULL;
 }
	// Copyright 2023 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 a system linked plugin exporting a single `simple_mul_workgroup`
	// function that also prints to stdout. This is not a great idea but shows how
	// plugins can have side-effecting behavior - even if in most cases a standalone
	// plugin can be used with much smaller code size and portability.
	//
	// The major use-case for a system linked plugin is JITs that may compile
	// imports on-demand. Such plugins could either JIT everything on load time
	// or defer JITting to the first call to a particular import. Performing JIT at
	// load time is strongly preferred as it keeps all of the expensive work in one
	// place before the program starts scheduling execution. Deferring will
	// introduce first-run delays and require warmup steps. Since only the imports
	// used by the program are present and most programs use all imports it's almost
	// always going to be better to do things ahead of time.
	//
	// NOTE: when using the system loader all unsafe behavior is allowed: TLS,
	// threads, mutable globals, syscalls, etc. Doing any of those things will
	// likely break in interesting ways as the import functions are called from
	// arbitrary threads concurrently. Be very careful and prefer standalone plugins
	// instead except when debugging/profiling.

	#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;
	} system_plugin_t;

	// `ret = lhs * rhs`
	//
	// Conforms to ABI:
	// #hal.pipeline.layout<constants = 1, bindings = [
	// #hal.pipeline.binding<storage_buffer, ReadOnly>,
	// #hal.pipeline.binding<storage_buffer, ReadOnly>,
	// #hal.pipeline.binding<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 simple_mul_workgroup(void* params_ptr, void* context,
	void* reserved) {
	system_plugin_t* plugin = (system_plugin_t*)context;
	typedef struct {
	const float* restrict binding0;
	size_t binding0_offset;
	const float* restrict binding1;
	size_t binding1_offset;
	float* restrict binding2;
	size_t binding2_offset;
	size_t size;
	size_t tid;
	uint32_t processor_id;
	const uint64_t* restrict processor_data;
	} params_t;
	const params_t* params = (const params_t*)params_ptr;
	fprintf(plugin->file, "processor_id=%u\n", params->processor_id);
	if (params->processor_data) {
	fprintf(plugin->file, "processor_data[0]=%" PRIX64 "\n",
	params->processor_data[0]);
	}
	// The operation `iree_codegen.ukernel.generic` always operates
	// on a slice of the inputs to produce a slice of the output,
	// so the loop here just needs to iterate from `0` to `size`,
	// where `size` is the size of the slice to be executed by this call.
	for (size_t i = 0; i < params->size; ++i) {
	params->binding2[params->binding2_offset + i] =
	params->binding0[params->binding0_offset + i] *
	params->binding1[params->binding2_offset + i];
	fprintf(plugin->file, "mul[%zu:%zu](%g * %g = %g)\n", params->tid, i,
	params->binding0[params->binding0_offset + i],
	params->binding1[params->binding1_offset + i],
	params->binding2[params->binding2_offset + i]);
	}
	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 system_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.
	system_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 system_plugin_unload(void* self) {
	system_plugin_t* plugin = (system_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 system_plugin_resolve(
	void* self, const iree_hal_executable_plugin_resolve_params_v0_t* params,
	iree_hal_executable_plugin_resolution_t* out_resolution) {
	system_plugin_t* plugin = (system_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,
	"simple_mul_workgroup") == 0) {
	params->out_fn_ptrs[i] = simple_mul_workgroup;
	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/system_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 = system_plugin_load,
	.unload = system_plugin_unload,
	.resolve = system_plugin_resolve,
	};
	return max_version <= IREE_HAL_EXECUTABLE_PLUGIN_VERSION_LATEST
	? (const iree_hal_executable_plugin_header_t**)&plugin
	: NULL;
	}