Serialize VM dialect to textual form (C/C++, etc)

[benvanik]

Similar to #34 (lowering VM dialect to LLVM IR dialect), the idea here is to emit the VM IR as C/C++. This would allow us to - without needing to involve LLVM IR and deal with target compilation/linkage issues - ahead-of-time compile the scheduler side of the VM operations (or the reference VMLA backend). The benefit here is that when targeting weird environments (which either don't have or may have incomplete or practically binary-only LLVM distributions) we aren't doing anything fancy. This makes debugging easier (we can emit #line directives to get source mapping), profiling easier (as we can use normal platform tools like vtune without special hacks), and the whole process more seamless for initial bringup (gcc module.c should be close to the scope of the required configuration).

The design of the VM C API should make this fairly straightforward (and if not we should fix the API :). The emitted code would implement the iree_vm_module_t interface and then could be plugged into an iree_vm_context_t just as any other custom module or the interpreter bytecode_module is plugged in.

This means that when walking a vm.module with vm.funcs you could emit the module struct and C functions and then emit the lookup functions (like get_function/lookup_function). It'd look similar to what the bytecode_module does (as I did that before writing all the C++ goo that we should avoid because it only makes it easier to do handwritten stuff and adds additional dependencies).

Some example IR snippets exist in the compiler/Dialect/VM/**/tests/ paths, like:
https://github.com/google/iree/blob/master/iree/compiler/Dialect/VM/Transforms/test/global_initialization.mlir

As an example, an input like this:

vm.module @my_module {
  vm.func @sub_i32(%arg0 : i32, %arg1 : i32) -> i32 {
    %0 = vm.sub.i32 %arg0, %arg1 : i32
    vm.return %0 : i32
  }
}

may end up as something like:

typedef struct {
  iree_vm_module_t interface;
} my_module_t;

typedef struct {
  // global variable storage, iree_vm_function-t import table, etc
} my_module_state_t;

iree_status_t my_module_alloc_state(...) {
  // allocate my_module_state_t, initialize variables, etc
}

// vm.func @sub_i32(%arg0 : i32, %arg1 : i32) -> i32
iree_status_t my_module_sub_i32(int32_t arg0, int32_t arg1, int32_t* out_ret0) {
  // vm.sub.i32 %arg0, %arg1 : i32
  int32_t v0 = arg0 - arg1;
  // vm.return %0 : i32
  *out_ret0 = v0;
  return IREE_STATUS_OK;
}

iree_status_t my_module_lookup_function(...) {
  if (function_name == "sub_i32") return ...;
}

iree_status_t my_module_create(iree_allocator_t allocator, iree_vm_module_t** out_module) {
  // alloc iree_vm_module_t
  // set function pointers
  module->interface.alloc_state = my_module_alloc_state;
  module->interface.lookup_function = my_module_lookup_function;
  ...
}

Because we are lowering from SSA IR we can simply emit the registers as is. vm.ref types in the IR can use the iree_vm_ref_* methods to manage the refs, and branches can be implemented with gotos (as we can emit the proper ref management when needed). Things like globals can be implemented with named values in the state struct (so vm.global.i32 @g2 in the IR becomes i32_t g2; as a struct member, etc).

The intent was to mirror the VM stack (iree_vm_stack_t) such that it is valid during execution. This will let us get nice stack walking on errors, implement a debugger for modules that can inspect module state, etc. It's not strictly required, though, except for coming into/out of the execute calls. I've got a bit of cleanup to do there and can try to make it lighter for this use case (#1172).

Imports (calls to vm.import ops) can be resolved through iree_vm_function_t. If the imports are listed in the lookup_function method (like hal.allocator.allocate, which would be generated from the list of vm.import ops in the IR) the runtime will call resolve_import for each one and then the iree_vm_module_t::execute method on the target function can be used to call it (ala https://github.com/google/iree/blob/master/iree/vm/bytecode_dispatch.c#L735-L736).

We can use this to clean up the interface, reduce rough edges, add more tests, etc. The nice thing about this is that the same interface is used by the HAL module and any custom module added (tensor lists, etc), so improvements to docs/performance/etc help everything.

To enable this all we can add a new translation registration like -iree-mlir-to-vm-c-module or something here: https://github.com/google/iree/blob/master/iree/compiler/Translation/IREEVM.cpp#L107 - note that the only difference between going to a flatbuffer with bytecode or C (or LLVM IR/etc) would be the final translation call: https://github.com/google/iree/blob/master/iree/compiler/Translation/IREEVM.cpp#L93

The translation then lives under VM/Target/:
https://github.com/google/iree/tree/master/iree/compiler/Dialect/VM/Target

[benvanik]

Let me know if this all sounds crazy or how it matches up with what you were thinking :) The intent is to make things faster/leaner/simpler vs. the flatbuffers+bytecode approach (still needed for dynamic module loading/deployment) and make our lives easier in testing/debugging/etc - if it starts to feel like that's not the case we can figure out if it's worth it.

[simon-camp]

I've looked into the bytecode module and put something together. Some pieces are missing / unclear to me though. If this goes in the right direction I'll update and cleanup PR #2842 to emit the corresponding code.
I think we should get something simple like the above example running end to end first so we can test this and add the missing ops afterwards.

typedef struct {
  iree_vm_module_t interface;
} MODULE_NAME_t;

typedef struct {
  // global variable storage, iree_vm_function-t import table, etc
} MODULE_NAME_state_t;

void MODULE_NAME_destroy(void *self) {
  MODULE_NAME_t *module = (MODULE_NAME_t *)self;

  return iree_allocator_free(module->allocator, module);
}

iree_string_view_t MODULE_NAME_name(void *self) {
  return iree_string_view_t{MODULE_NAME, strlen(MODULE_NAME)};
}

iree_vm_module_signature_t MODULE_NAME_signature(void *self) {
  //???
}

iree_status_t MODULE_NAME_get_function(
    void *self, iree_vm_function_linkage_t linkage, int32_t ordinal,
    iree_vm_function_t *out_function, iree_string_view_t *out_name,
    iree_vm_function_signature_t *out_signature) {
  //???
}

iree_status_t MODULE_NAME_lookup_function(void *self,
                                          iree_vm_function_linkage_t linkage,
                                          iree_string_view_t name,
                                          iree_vm_function_t *out_function) {
  if (!out_function) return IREE_STATUS_INVALID_ARGUMENT;
  memset(out_function, 0, sizeof(iree_vm_function_t));

  if (!name.data || !name.size) return IREE_STATUS_INVALID_ARGUMENT;

  MODULE_NAME_t *module = (MODULE_NAME_t *)self;

  if (!strcmp(name.data, "MODULE_NAME_FUNC_A")) {
    out_function->module = &module->interface;
    out_function->linkage = NULL;  //???
    out_function->ordinal = NULL;  //???
    return IREE_STATUS_OK;
  }
  // repeat for each function

  return IREE_STATUS_NOT_FOUND;
}

iree_status_t MODULE_NAME_alloc_state(
    void *self, iree_allocator_t allocator,
    iree_vm_module_state_t **out_module_state) {
  if (!out_module_state) return IREE_STATUS_INVALID_ARGUMENT;
  *out_module_state = NULL;

  MODULE_NAME_t *module = (MODULE_NAME_t *)self;

  MODULE_NAME_state_t *state = NULL;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc(
      allocator, sizeof(MODULE_NAME_state_t), (void **)&state));
  state->allocator = allocator;

  //???

  *out_module_state = (iree_vm_module_state_t *)state;
  return IREE_STATUS_OK;
}

void MODULE_NAME_free_state(void *self, iree_vm_module_state_t *module_state) {
  if (!module_state) return;

  MODULE_NAME_state_t *state = (MODULE_NAME_state_t *)module_state;

  //???

  iree_allocator_free(state->allocator, module_state);
}

iree_status_t MODULE_NAME_resolve_import(void *self,
                                         iree_vm_module_state_t *module_state,
                                         int32_t ordinal,
                                         iree_vm_function_t function) {
  //???
}

iree_status_t MODULE_NAME_begin_call(void *self, iree_vm_stack_t *stack,
                                     const iree_vm_function_call_t *call,
                                     iree_vm_execution_result_t *out_result) {
  //???
}

iree_status_t MODULE_NAME_get_function_reflection_attr(
    void *self, iree_vm_function_linkage_t linkage, int32_t ordinal,
    int32_t index, iree_string_view_t *key, iree_string_view_t *value) {
  //???
}

iree_status_t MODULE_NAME_create(iree_allocator_t allocator,
                                 iree_vm_module_t **out_module) {
  if (!out_module) return IREE_STATUS_INVALID_ARGUMENT;
  *out_module = NULL;

  MODULE_NAME_t *module = NULL;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc(allocator, sizeof(MODULE_NAME_t),
                                             (void **)&module));
  module->allocator = allocator;

  iree_vm_module_initialize(&module->interface, module);
  module->interface.destroy = MODULE_NAME_destroy;
  module->interface.name = MODULE_NAME_name;
  module->interface.signature = MODULE_NAME_signature;
  module->interface.get_function = MODULE_NAME_get_function;
  module->interface.lookup_function = MODULE_NAME_lookup_function;
  module->interface.alloc_state = MODULE_NAME_alloc_state;
  module->interface.free_state = MODULE_NAME_free_state;
  module->interface.resolve_import = MODULE_NAME_resolve_import;
  module->interface.begin_call = MODULE_NAME_begin_call;
  module->interface.get_function_reflection_attr =
      MODULE_NAME_get_function_reflection_attr;

  *out_module = &module->interface;
  return IREE_STATUS_OK;
}

[benvanik]

Oo a timely question! I was just looking at something similar for revamping the native modules (like the HAL and custom modules) so that they shared more code. We may be able to share the same thing with this, too, which will keep things easier to modify and reduce total binary size.

My plan was to have an iree/vm/native_module.h+.c that had all of the general stuff for reflection and method dispatch and then generate an inline include file containing a few tables of information that were used per module. This is similar to how the current C++ wrapper in module_abi_cc.h is done (with the table here https://github.com/google/iree/blob/e734f9a8a4d7f17f23dd0d838f964f89a32f73b4/iree/vm/module_abi_cc.h#L240 and the rest of the code in that file using that table for information).

For purely native modules (custom/HAL/etc) they'd then just need to define the functions they export and put them in one or more tables. Here's a super rough sketch I did last night: https://gist.github.com/benvanik/2d7929373181e23511319aa0ed7b1ad9 - of note is that users provide the fn and then the fn_call shims are generated that marshal arguments into C arguments. Users can have any other additional C functions that use whatever convention they want internally but the exported ones are derived from the vm.import ops for example in https://github.com/google/iree/blob/9624bbbcaa1f491f067aad500d6679b538db3403/iree/compiler/Dialect/HAL/hal.imports.mlir#L28-L33.

What I'm thinking is that you could generate the same kind of thing here; if you emitted your C functions then all you'd need to do is also emit the dispatch table and some module metadata (like its name) and leave the rest like function lookup to the hand-written native_module code. I was thinking that the emitter for the tables would take in an IREE::VM::ModuleOp with the ImportOps and then emit the tables, so if you took your exported IREE::VM::FuncOps and turned them into ImportOps you'd be able to just pass those along. Then we have just one place taking the ops and producing the function signatures so they won't get out of sync.

If this sounds reasonable, I can get the first pass on this done today (PST) - or at least the header with the table types - and then we can iterate from there. Or, you can do what you have above and we could always reconcile it later - I don't want to block you, but this may help take care of a lot of goo for you that isn't core to what the actual emitc work is doing :)

[benvanik]

Here's a rough sketch of the header:

typedef struct {
  iree_string_view_t key;
  iree_string_view_t value;
} iree_vm_reflection_attr_t;

// Describes an imported native function in a native module.
// All of this information is assumed read-only and will be referenced for the
// lifetime of any module created with the descriptor.
typedef struct {
  // Fully-qualified function name (for example, 'other_module.foo').
  iree_string_view_t full_name;
} iree_vm_native_import_descriptor_t;

// Describes an exported native function in a native module.
// All of this information is assumed read-only and will be referenced for the
// lifetime of any module created with the descriptor.
typedef struct {
  // Module-local function name (for example, 'foo' for function 'module.foo').
  iree_string_view_t local_name;

  // TODO(#1979): move register info to iree_vm_function_signature_t.
  // Total number of valid i32 registers used by the function.
  uint16_t i32_register_count;
  // Total number of valid ref registers used by the function.
  uint16_t ref_register_count;

  // An optional list of function-level reflection attributes.
  iree_host_size_t reflection_attr_count;
  const iree_vm_reflection_attr_t* reflection_attrs;
} iree_vm_native_export_descriptor_t;

// Describes an native module implementation by way of descriptor tables.
// All of this information is assumed read-only and will be referenced for the
// lifetime of any module created with the descriptor.
//
// The common native module code will use this descriptor to return metadata on
// query, lookup exported functions, and call module-provided implementation
// functions for state and call management.
typedef struct {
  // Name of the module prefixed on all exported functions.
  iree_string_view_t module_name;

  // All imported function descriptors.
  // interface.resolve_import will be called for each import.
  iree_host_size_t import_count;
  const iree_vm_native_import_descriptor_t* imports;

  // All exported function descriptors.
  iree_host_size_t export_count;
  const iree_vm_native_export_descriptor_t* exports;

  // An optional list of module-level reflection attributes.
  iree_host_size_t reflection_attr_count;
  const iree_vm_reflection_attr_t* reflection_attrs;
} iree_vm_native_module_descriptor_t;

// Creates a new native module with the metadata tables in |descriptor|.
// These tables will be used for reflection and function lookup, and the
// provided function pointers will be called when state needs to be managed or
// exported functions need to be called.
//
// An implementation |interface| providing functions for state management and
// function calls can be provided to override default implementations of
// functions. The structure will be copied.
//
// The provided |descriptor| will be referenced by the created module and must
// be kept live for the lifetime of the module.
IREE_API_EXPORT iree_status_t IREE_API_CALL iree_vm_native_module_create(
    const iree_vm_module_t* interface,
    const iree_vm_native_module_descriptor_t* module_descriptor,
    iree_allocator_t allocator, iree_vm_module_t** out_module);

So you could generate some readonly tables like:

static const iree_vm_native_import_descriptor_t my_module_imports_[] = {
  {iree_make_cstring_view("other.foo")},
  {iree_make_cstring_view("other.bar")},
};
static const iree_vm_reflection_attr_t fn_b_attrs_[] = {
  {iree_make_cstring_view("key"), iree_make_cstring_view("value")},
};
static const iree_vm_native_export_descriptor_t my_module_exports_[] = {
  {iree_make_cstring_view("fn_a"), 8, 12, 0, NULL},
  {iree_make_cstring_view("fn_b"), 8, 12, IREE_ARRAYSIZE(fn_b_attrs_), fn_b_attrs_},
};

static const iree_vm_native_module_descriptor_t my_module_descriptor_ = {
  iree_make_cstring_view("my_module"),
  IREE_ARRAYSIZE(my_module_imports_), my_module_imports_,
  IREE_ARRAYSIZE(my_module_exports_), my_module_exports_,
  0, NULL,
};

And provide them when you create your module:

iree_vm_module_t interface = {0};
interface.alloc_state = my_module_alloc_state;
interface.free_state = my_module_free_state;
interface.begin_call = my_module_begin_call;
iree_vm_native_module_create(&interface, my_module_descriptor_, ...);

Besides a lot of the common code (like a binary search on exported function names in function lookup, reflection attribute queries and error validation, etc) being shared, which is nice, the only thing your generated module code has to do is have some compile-time generated tables and provide your state alloc, free, and function call implementations.

[benvanik]

Here's a WIP PR: #2866
native_module_test.h shows what some modules look like and is probably the most interesting. In this model you'd be emitting a .h/.c with:

• the import/export/module descriptors
• a custom module_t for caching types (like hal.buffer, etc)
• a custom module_state_t with any module globals (if needed)
• the function table itself pointing at your generated functions
• a creation function that sets up the type table

The rest of the code for reflection/function lookups/etc are in native_module.c and shared. It's not a lot of code, but it is nice to have a consistent implementation and when 10-20+ modules are built into a binary not duplicating it all across all emit-c modules, custom user modules, etc.

The thing still missing in this PR is the stuff from https://gist.github.com/benvanik/2d7929373181e23511319aa0ed7b1ad9 for marshaling args nicely, which I'll get finished tomorrow. Exports are easy (demonstrated in module_b.entry as an i32->i32), but the imports (module_a.*) require some nicer macros.

Hopefully this is useful and gives us something to iterate on!

[simon-camp]

Thank you :)
I've adapted your test case to the VM module from PR #2842 and got it working. https://gist.github.com/simon-camp/cb63d7d434c586e6552e28fd5bd07317 is what I've come up with.

So add_module_test.h needs to be generated by the CModuleTarget, right?

[benvanik]

Nice! That's the idea - then you'd just have an extern module_a_create (or whatever your module name is) and call that to create the module (vs iree_vm_bytecode_module_create or iree_hal_module_create). The nice thing about the way you have it separated there with add_mlir_generated.h and the add_module_test.h is that the goo setting up modules is isolated from the actual code that was translated, so as a user I could ignore the module setup and see that nice mapping from add.mlir -> add_mlir_generated.h.

Now I just need to hide all the boilerplate to marshal function args and it'll be even better :)
I didn't get to that today, but I think with what you have it shouldn't be blocking so I can land the native_module PR as-is if that helps!

[simon-camp]

Yeah landing your PR as-is would already be really useful. I'll then update my PR. I should be able to get something up soon.
Once this is done we can get to the more interesting stuff :)

[benvanik]

Cool - I'll get that landed today!

[simon-camp]

The code generation is now hooked up in the build system so things are tested properly. The generated code looks like this

vm.module @add_module {
  vm.func @add_1(%arg0 : i32, %arg1 : i32) -> (i32, i32) {
    %0 = vm.add.i32 %arg0, %arg1 : i32
    %1 = vm.add.i32 %0, %0 : i32
    vm.return %0, %1 : i32, i32
  }
  vm.export @add_1
}

#include "iree/vm/context.h"
#include "iree/vm/instance.h"
#include "iree/vm/native_module.h"
#include "iree/vm/ref.h"
#include "iree/vm/stack.h"

#include "iree/compiler/Dialect/VM/Target/C/vm_c_funcs.h"

//=============================================================================
// module "add_module"
//=============================================================================

iree_status_t add_module_add_1_impl(int32_t v1, int32_t v2, int32_t *out0, int32_t *out1) {
int32_t v3 = vm_add_i32(v1, v2);
int32_t v4 = vm_add_i32(v3, v3);
*out0 = v3;
*out1 = v4;
return iree_ok_status();
}

//=============================================================================
// The code below setups functions and lookup tables to implement the vm
// interface
//=============================================================================
//=============================================================================
// module "add_module"
//=============================================================================

static const iree_vm_native_export_descriptor_t add_module_exports_[] = {
{iree_make_cstring_view("add_1"), 0, 0, 0, NULL},
};

static const iree_vm_native_import_descriptor_t add_module_imports_[] = {
};

static const iree_vm_native_module_descriptor_t add_module_descriptor_ = {
iree_make_cstring_view("add_module"),
IREE_ARRAYSIZE(add_module_imports_),
add_module_imports_,
IREE_ARRAYSIZE(add_module_exports_),
add_module_exports_,
0,
NULL,
};

The rest of the setup is written by hand for now. I'd like to set up one or two more tests to run the translation on. Supporting vm.call would allow us to test imports I guess. Not sure how the translation would look like though.

[benvanik]

That's awesome @simon-camp!

For vm.call, you could handle only internal calls to start. These would be only to other functions within the same module so you don't have to worry about cross-module calling conventions. In that case, %1 = vm.call @foo(%0) would simply be IREE_RETURN_IF_ERROR(foo(v0, &v1));

Of other things to start testing next (building up to running full modules using the HAL):

• br/cond_br
• module state for global variables/initialization function
• iree_vm_ref_t support for ref<...> types

[marbre]

With what we already have landed, this is basically done! 🥳