From 1af38b5848c2c2153ed8b4b8a603c90bdf4433d3 Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Wed, 14 Jun 2023 20:20:58 -0700 Subject: [PATCH 1/2] Add VmModule.mmap() to Python API. We really need page aligned flatbuffer blobs vs normal malloc alignment. The best way to be loading a file is via mmap, so just make that available as an API. This could be done in Python by the caller but is error-prone. The public API will make this more robust. Provides the mechanism to fix #13887 --- runtime/bindings/python/tests/vm_test.py | 23 +++++++++++++++++++++++ runtime/bindings/python/vm.cc | 18 ++++++++++++++++++ runtime/bindings/python/vm.h | 1 + 3 files changed, 42 insertions(+) diff --git a/runtime/bindings/python/tests/vm_test.py b/runtime/bindings/python/tests/vm_test.py index bde6ae966510..85e944d15926 100644 --- a/runtime/bindings/python/tests/vm_test.py +++ b/runtime/bindings/python/tests/vm_test.py @@ -8,6 +8,7 @@ import logging import numpy as np +import tempfile import unittest import iree.compiler @@ -106,6 +107,28 @@ def test_add_scalar_new_abi(self): logging.info("result: %s", result) self.assertEqual(result, 11) + def test_mmap(self): + binary = iree.compiler.compile_str( + """ + func.func @add_scalar(%arg0: i32, %arg1: i32) -> i32 { + %0 = arith.addi %arg0, %arg1 : i32 + return %0 : i32 + } + """, + target_backends=iree.compiler.core.DEFAULT_TESTING_BACKENDS + ) + with tempfile.NamedTemporaryFile() as tf: + tf.write(binary) + tf.flush() + m = iree.runtime.VmModule.mmap(self.instance, tf.name) + context = iree.runtime.VmContext(self.instance, + modules=[self.hal_module, m]) + f = m.lookup_function("add_scalar") + finv = iree.runtime.FunctionInvoker(context, self.device, f, tracer=None) + result = finv(5, 6) + logging.info("result: %s", result) + self.assertEqual(result, 11) + def test_synchronous_dynamic_shape_invoke_function_new_abi(self): m = create_simple_dynamic_abs_module(self.instance) context = iree.runtime.VmContext(self.instance, diff --git a/runtime/bindings/python/vm.cc b/runtime/bindings/python/vm.cc index b72ecdea1989..9a3a4fdf2f7e 100644 --- a/runtime/bindings/python/vm.cc +++ b/runtime/bindings/python/vm.cc @@ -150,6 +150,23 @@ VmModule VmModule::ResolveModuleDependency(VmInstance* instance, return py_module; } +VmModule VmModule::MMap(VmInstance* instance, std::string filepath) { + IREE_TRACE_SCOPE_NAMED("VmModule::MMap"); + auto mmap_module = py::module::import("mmap"); + auto open_func = py::module::import("io").attr("open"); + auto file_obj = open_func(filepath, "r+b"); + auto flags = py::cast(mmap_module.attr("MAP_SHARED")); + // MAP_POPULATE isn't available on all versions/platforms. + if (py::hasattr(mmap_module, "MAP_POPULATE")) { + flags |= py::cast(mmap_module.attr("MAP_POPULATE")); + } + auto prot = py::cast(mmap_module.attr("PROT_READ")); + auto mapped_file = + mmap_module.attr("mmap")(file_obj.attr("fileno")(), 0, flags, prot); + mapped_file.attr("madvise")(mmap_module.attr("MADV_RANDOM")); + return FromFlatbufferBlob(instance, mapped_file); +} + VmModule VmModule::FromFlatbufferBlob(VmInstance* instance, py::object flatbuffer_blob_object) { IREE_TRACE_SCOPE_NAMED("VmModule::FromFlatbufferBlob"); @@ -661,6 +678,7 @@ void SetupVmBindings(pybind11::module m) { .def_static("resolve_module_dependency", &VmModule::ResolveModuleDependency) .def_static("from_flatbuffer", &VmModule::FromFlatbufferBlob) + .def_static("mmap", &VmModule::MMap) .def_property_readonly("name", &VmModule::name) .def_property_readonly("version", [](VmModule& self) { diff --git a/runtime/bindings/python/vm.h b/runtime/bindings/python/vm.h index 4b097e7e35e1..9bf25f462ba9 100644 --- a/runtime/bindings/python/vm.h +++ b/runtime/bindings/python/vm.h @@ -133,6 +133,7 @@ class VmModule : public ApiRefCounted { const std::string& name, uint32_t minimum_version); + static VmModule MMap(VmInstance* instance, std::string filepath); static VmModule FromFlatbufferBlob(VmInstance* instance, py::object flatbuffer_blob_object); From ff086159340967c82834a6a7b55e4f9ae9ef285c Mon Sep 17 00:00:00 2001 From: Stella Laurenzo Date: Wed, 14 Jun 2023 20:26:12 -0700 Subject: [PATCH 2/2] Format --- runtime/bindings/python/tests/vm_test.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/runtime/bindings/python/tests/vm_test.py b/runtime/bindings/python/tests/vm_test.py index 85e944d15926..7ff760028374 100644 --- a/runtime/bindings/python/tests/vm_test.py +++ b/runtime/bindings/python/tests/vm_test.py @@ -115,14 +115,13 @@ def test_mmap(self): return %0 : i32 } """, - target_backends=iree.compiler.core.DEFAULT_TESTING_BACKENDS - ) + target_backends=iree.compiler.core.DEFAULT_TESTING_BACKENDS) with tempfile.NamedTemporaryFile() as tf: tf.write(binary) tf.flush() m = iree.runtime.VmModule.mmap(self.instance, tf.name) context = iree.runtime.VmContext(self.instance, - modules=[self.hal_module, m]) + modules=[self.hal_module, m]) f = m.lookup_function("add_scalar") finv = iree.runtime.FunctionInvoker(context, self.device, f, tracer=None) result = finv(5, 6)