Some python binding plumbing to enable the new VM2 module loading.
* Just adds enough of the new hal types to get a default device.
* Restricted to vulkan for the moment.
* Verifies that with the new hal module, the VM2 compiled simple_mul module loads.
* Disabled vm_test for TAP due to missing vulkan deps on test hosts.
PiperOrigin-RevId: 286590561
diff --git a/bindings/python/pyiree/vm.cc b/bindings/python/pyiree/vm.cc
index 72af468..c2e59fc 100644
--- a/bindings/python/pyiree/vm.cc
+++ b/bindings/python/pyiree/vm.cc
@@ -17,10 +17,13 @@
#include "absl/types/optional.h"
#include "bindings/python/pyiree/status_utils.h"
#include "iree/base/api.h"
+#include "iree/modules/hal/hal_module.h"
namespace iree {
namespace python {
+namespace {
+
RtModule CreateModuleFromBlob(std::shared_ptr<OpaqueBlob> blob) {
iree_rt_module_t* module;
auto free_fn = OpaqueBlob::CreateFreeFn(blob);
@@ -31,6 +34,16 @@
return RtModule::CreateRetained(module);
}
+VmModule CreateHalModule(HalDevice* device) {
+ iree_vm_module_t* module;
+ CheckApiStatus(
+ iree_hal_module_create(device->raw_ptr(), IREE_ALLOCATOR_SYSTEM, &module),
+ "Error creating hal module");
+ return VmModule::CreateRetained(module);
+}
+
+} // namespace
+
//------------------------------------------------------------------------------
// VmInstance
//------------------------------------------------------------------------------
@@ -115,9 +128,15 @@
}
void SetupVmBindings(pybind11::module m) {
+ CHECK_EQ(IREE_STATUS_OK, iree_vm_register_builtin_types());
+ CHECK_EQ(IREE_STATUS_OK, iree_hal_module_register_types());
+
// Deprecated: VM1 module.
m.def("create_module_from_blob", CreateModuleFromBlob);
+ // Built-in module creation.
+ m.def("create_hal_module", &CreateHalModule);
+
py::enum_<iree_vm_function_linkage_t>(m, "Linkage")
.value("INTERNAL", IREE_VM_FUNCTION_LINKAGE_INTERNAL)
.value("IMPORT", IREE_VM_FUNCTION_LINKAGE_IMPORT)
