feat: add webassembly compilation" (#12)

* feat: working wasm emission

* feat: output companion .js file

* wip: better interfacing

* feat: working listing of devices for a driver

* feat: add bytecode loader class

* wip: call bytecode

* wip

* merge

* feat: first working draft of wasm execution

* wip

* fix: return typed arrays

* fix: double-free on context

* chore: cleanup comments

* refactor: modularize cmakelists

* feat: compile fewer things for wasm host tools

.github/workflows/embedded_devices.yml

@@ -70,9 +70,9 @@ jobs:
         if: steps.cache-iree-dir.outputs.cache-hit != 'true'
         run: make clone_iree
       - name: Compile Host
-        run: BUILD_IREE_RUNTIME=true ./scripts/build_and_package.sh --target=host
+        run: ./scripts/build_and_package.sh --target=host
       - name: Compile
-        run: BUILD_IREE_RUNTIME=false ./scripts/build_and_package.sh --target=${{ matrix.build_target }}
+        run: ./scripts/build_and_package.sh --target=${{ matrix.build_target }}
       - name: Release Artifacts
         uses: softprops/action-gh-release@v2
         if: startsWith(github.ref, 'refs/tags/')

Makefile

@@ -17,36 +17,54 @@ PRIV_DIR = $(MIX_APP_PATH)/priv
 ifeq ($(NX_IREE_PREFER_PRECOMPILED), true)
 all: nx_iree
 else
-all: install_runtime nx_iree
+all: clone_iree install_runtime nx_iree
 endif
 
 .PHONY: clone_iree
 clone_iree: $(NX_IREE_SOURCE_DIR)
+	@echo "Cloned IREE into $(NX_IREE_SOURCE_DIR)"
 
-.PHONY: iree_source_dir
-iree_source_dir:
+$(NX_IREE_SOURCE_DIR):
 	./scripts/clone_iree.sh $(IREE_GIT_REV) $(NX_IREE_SOURCE_DIR)
 
 IREE_CMAKE_BUILD_DIR ?= $(abspath iree-runtime/iree-build)
 IREE_RUNTIME_INCLUDE_PATH := $(abspath $(NX_IREE_SOURCE_DIR)/runtime/src/iree)
 IREE_RUNTIME_BUILD_DIR ?= $(abspath iree-runtime/build)
 IREE_INSTALL_DIR ?= $(abspath iree-runtime/host/install)
 
-IREE_CMAKE_CONFIG ?= Release
-
 IREE_BUILD_TARGET ?= host
 
 BUILD_TARGET_FLAGS = -S $(abspath cmake)
 
 CUDA_PRESENT := $(shell command -v nvcc >/dev/null 2>&1 && echo true || echo false)
 
 ifeq ($(CUDA_PRESENT), true)
-	CFLAGS += -DCUDA_ENABLED
-	CMAKE_CXX_FLAGS += -DCUDA_ENABLED
+	ifeq ($(IREE_BUILD_TARGET), host)
+		CFLAGS += -DCUDA_ENABLED
+		CMAKE_CXX_FLAGS += -DCUDA_ENABLED
+	endif
 endif
 
-# flags for xcode 15.4
+BUILD_HOST_COMPILER ?= OFF
+BUILD_HOST_COMPILER_FLAGS := ""
+
+ifeq ($(BUILD_HOST_COMPILER), ON)
+	BUILD_HOST_COMPILER_FLAGS += \
+		-DIREE_BUILD_COMPILER=ON \
+		-DIREE_INPUT_TORCH=OFF \
+		-DIREE_INPUT_TOSA=OFF \
+		-DIREE_BUILD_SAMPLES=OFF \
+		-DIREE_BUILD_TESTS=OFF \
+		-DIREE_HAL_DRIVER_DEFAULTS=OFF \
+		-DIREE_BUILD_PYTHON_BINDINGS=OFF \
+		-DIREE_BUILD_BINDINGS_TFLITE=OFF
+endif
+
+# apple target flags specified for xcode 15.4
 ifeq ($(IREE_BUILD_TARGET), host)
+else ifeq ($(IREE_BUILD_TARGET), webassembly)
+  BUILD_TARGET_FLAGS += \
+		-DIREE_HOST_BIN_DIR=$(abspath $(IREE_HOST_BIN_DIR))
 else ifeq ($(IREE_BUILD_TARGET), ios)
 	BUILD_TARGET_FLAGS += \
 		-DCMAKE_SYSTEM_NAME=iOS\
@@ -102,39 +120,43 @@ else
 endif
 
 .PHONY: install_runtime
-install_runtime: iree_host $(IREE_INSTALL_DIR)
-
+ifneq ($(strip $(IREE_HOST_BUILD_DIR)),)
+install_runtime: $(IREE_HOST_INSTALL_DIR)/bin/iree-flatcc-cli $(IREE_INSTALL_DIR)
+else
+install_runtime: $(IREE_INSTALL_DIR)
+endif
 
-CMAKE_SOURCES = $(abspath cmake/src/runtime.cc) $(abspath cmake/src/runtime.h)
+CMAKE_SOURCES = $(wildcard cmake/src/*.cc cmake/src/*.h)
 
-$(IREE_INSTALL_DIR): iree_source_dir $(CMAKE_SOURCES)
-	cmake -G Ninja -B $(IREE_CMAKE_BUILD_DIR) \
+$(IREE_INSTALL_DIR): $(NX_IREE_SOURCE_DIR) $(CMAKE_SOURCES)
+	$(EMCMAKE) cmake -G Ninja -B $(IREE_CMAKE_BUILD_DIR) \
 		-DCMAKE_BUILD_TYPE=$(IREE_CMAKE_CONFIG)\
 		-DIREE_BUILD_COMPILER=OFF\
 		-DIREE_RUNTIME_BUILD_DIR=$(IREE_RUNTIME_BUILD_DIR)\
 		-DIREE_RUNTIME_INCLUDE_PATH=$(IREE_RUNTIME_INCLUDE_PATH)\
 		-DNX_IREE_SOURCE_DIR=$(NX_IREE_SOURCE_DIR) \
 		-DCMAKE_CXX_FLAGS=$(CMAKE_CXX_FLAGS) \
 		$(BUILD_TARGET_FLAGS)
-
 	cmake --build $(IREE_CMAKE_BUILD_DIR) --config $(IREE_CMAKE_CONFIG)
 	cmake --install $(IREE_CMAKE_BUILD_DIR) --config $(IREE_CMAKE_CONFIG) --prefix $(IREE_INSTALL_DIR)
 
 .PHONY: iree_host
-ifeq ($(BUILD_IREE_RUNTIME), true)
+ifneq ($(strip $(IREE_HOST_BUILD_DIR)),)
+iree_host: $(IREE_HOST_BUILD_DIR)/bin/iree-flatcc-cli
+else
 iree_host:
-	@echo "Building IREE runtime host binaries at $(IREE_HOST_BUILD_DIR)."
+	@echo "IREE_HOST_BUILD_DIR not set. Skipping host binaries build."
+endif
+
+$(IREE_HOST_INSTALL_DIR)/bin/iree-flatcc-cli: $(NX_IREE_SOURCE_DIR) $(CMAKE_SOURCES)
+	@echo "Building IREE runtime host binaries at `$(IREE_HOST_BUILD_DIR)`."
 	cmake -G Ninja -B $(IREE_HOST_BUILD_DIR) \
 		-DCMAKE_INSTALL_PREFIX=$(IREE_HOST_INSTALL_DIR) \
-		-DIREE_BUILD_COMPILER=OFF\
+		-DIREE_BUILD_COMPILER=$(BUILD_HOST_COMPILER) $(BUILD_HOST_COMPILER_FLAGS) \
 		-DCMAKE_BUILD_TYPE=$(IREE_CMAKE_CONFIG) \
 		-DCMAKE_CXX_FLAGS=$(CMAKE_CXX_FLAGS) \
 		-S $(NX_IREE_SOURCE_DIR)
 	cmake --build $(IREE_HOST_BUILD_DIR) --target install
-else
-iree_host:
-	@echo "Not building IREE runtime host binaries. Skipping."
-endif
 
 ### NxIREE Runtime NIF library
 
@@ -146,11 +168,12 @@ NX_IREE_RUNTIME_LIB = cache/iree-runtime
 NX_IREE__IREE_RUNTIME_INCLUDE_PATH = $(NX_IREE_RUNTIME_LIB)/include
 NX_IREE_RUNTIME_SO ?= $(MIX_APP_PATH)/priv/libnx_iree_runtime.so
 
-CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -I$(NX_IREE__IREE_RUNTIME_INCLUDE_PATH) -Wall -Wno-sign-compare \
-	 -Wno-unused-parameter -Wno-missing-field-initializers -Wno-comment \
-	 -std=c++17 -w
+CFLAGS = -fPIC -I$(ERTS_INCLUDE_DIR) -I$(NX_IREE__IREE_RUNTIME_INCLUDE_PATH) -Wall -std=c++17 -w
+
+IREE_CMAKE_CONFIG ?= Release
 
 ifdef DEBUG
+	IREE_CMAKE_CONFIG = RelWithDebInfo
 	CFLAGS += -g
 else
 	CFLAGS += -O3
@@ -215,14 +238,20 @@ $(NX_IREE_SO): $(NX_IREE_CACHE_SO)
 $(NX_IREE__IREE_RUNTIME_INCLUDE_PATH):
 	cp -r iree-runtime/host/install $(dir $@)
 
-cache/objs/%.o: c_src/%.cc
+cache/objs/%.o: c_src/%.cc $(CMAKE_SOURCES)
 	@ mkdir -p $(dir $@)
 	$(CXX) $(CFLAGS) -o $@ -c $<
 
 # Print IREE Dir
+.PHONY: PTD
 PTD:
 	@ echo $(NX_IREE_SOURCE_DIR)
 
+.PHONY: wasm_build
+webassembly:
+	./scripts/build_and_package.sh --target=host --build-compiler
+	./scripts/build_and_package.sh --target=webassembly
+
 clean:
 	rm -rf cache/objs
 	rm -rf $(TARGET_DIR)

c_src/nx_iree.cc

@@ -309,42 +309,6 @@ std::string iree_type_to_nx_type(iree_hal_element_type_t type) {
   return "invalid_type";
 }
 
-iree_hal_element_type_t nx_type_to_iree_type(std::string type) {
-  using type_enum = iree_hal_element_types_t;
-
-  if (type == "i8") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_INT_8;
-  } else if (type == "i16") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_INT_16;
-  } else if (type == "i32") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_INT_32;
-  } else if (type == "i64") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_INT_64;
-  } else if (type == "ui8") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_8;
-  } else if (type == "ui16") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_16;
-  } else if (type == "ui32") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_32;
-  } else if (type == "ui64") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_UINT_64;
-  } else if (type == "bf16") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_BFLOAT_16;
-  } else if (type == "f16") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_FLOAT_16;
-  } else if (type == "f32") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_FLOAT_32;
-  } else if (type == "f64") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_FLOAT_64;
-  } else if (type == "c64") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_64;
-  } else if (type == "c128") {
-    return type_enum::IREE_HAL_ELEMENT_TYPE_COMPLEX_FLOAT_128;
-  }
-
-  return type_enum::IREE_HAL_ELEMENT_TYPE_NONE;
-}
-
 DECLARE_NIF(read_buffer_nif) {
   iree_hal_device_t** device;
   iree::runtime::IREETensor** input;

cmake/CMakeLists.txt

@@ -14,72 +14,31 @@ set(IREE_INPUT_STABLEHLO ON)
 set(IREE_BUILD_TESTS OFF)
 set(IREE_BUILD_SAMPLES OFF)
 
-set(IREE_HAL_DRIVER_DEFAULTS ON)
-set(IREE_HAL_DRIVER_LOCAL_SYNC ON)
-set(IREE_HAL_EXECUTABLE_LOADER_DEFAULTS OFF)
-set(IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF ON)
-
 if(CMAKE_BUILD_TYPE MATCHES MinSizeRel)
   set(IREE_SIZE_OPTIMIZED ON)
 endif()
 
 file(GLOB iree_runtime_sources CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/src/*.cc" "${CMAKE_CURRENT_SOURCE_DIR}/src/*.h")
-add_library(${_NAME} SHARED ${iree_runtime_sources})
-
-set(__BUILD_DIR "${IREE_RUNTIME_BUILD_DIR}")
-
-add_subdirectory("${NX_IREE_SOURCE_DIR}" ${__BUILD_DIR} EXCLUDE_FROM_ALL)
 
-install(
-    TARGETS ${_NAME}
-    DESTINATION "."
-)
-
-set_target_properties(${_NAME} PROPERTIES SUFFIX ".so")
+if("$ENV{IREE_BUILD_TARGET}" STREQUAL "webassembly")
+  add_subdirectory(modules/emscripten)
+else()
+  add_subdirectory(modules/default)
+endif()
 
 set_target_properties(${_NAME} PROPERTIES
     INSTALL_RPATH_USE_LINK_PATH TRUE
     BUILD_WITH_INSTALL_RPATH TRUE
 )
 
-if($ENV{DEBUG})
-  set(CMAKE_BUILD_TYPE Debug)
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g")
-endif()
+set(__BUILD_DIR "${IREE_RUNTIME_BUILD_DIR}")
 
-if(NOT APPLE)
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -shared -fvisibility=default")
-  set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--export-dynamic")
-else()
-  # Although the compiler complains about not using these,
-  # things only work with them set
-  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -undefined dynamic_lookup -fvisibility=default")
-  set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,-export_dynamic")
-  check_c_compiler_flag("-arch arm64" ARM64_SUPPORTED)
-  if(ARM64_SUPPORTED)
-      set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DMAC_ARM64")
-  endif()
-endif()
+add_subdirectory("${NX_IREE_SOURCE_DIR}" ${__BUILD_DIR} EXCLUDE_FROM_ALL)
 
 # Ensure visibility of all symbols
 set(CMAKE_CXX_VISIBILITY_PRESET default)
 set(CMAKE_VISIBILITY_INLINES_HIDDEN OFF)
 
 target_compile_options(${_NAME} PUBLIC ${IREE_DEFAULT_COPTS})
 
-add_definitions(-DLLVM_DISABLE_ABI_BREAKING_CHECKS_ENFORCING=1)
-
-target_link_libraries(${_NAME} iree_runtime_runtime)
-target_link_libraries(${_NAME} iree_tooling_context_util)
-
-# Install the header files - this will make it easier to copy them over
-# to the final bundle.
-
-install(DIRECTORY "${IREE_RUNTIME_INCLUDE_PATH}"
-        DESTINATION "include"
-        FILES_MATCHING PATTERN "*.h")
-
-file(GLOB HEADER_FILES "${CMAKE_CURRENT_SOURCE_DIR}/src/*.h")
-
-install(FILES ${HEADER_FILES}
-        DESTINATION "include/nx_iree")
+add_definitions(-DLLVM_DISABLE_ABI_BREAKING_CHECKS_ENFORCING=1)

cmake/modules/default/CMakeLists.txt

@@ -0,0 +1,50 @@
+set(IREE_HAL_DRIVER_DEFAULTS ON)
+set(IREE_HAL_DRIVER_LOCAL_SYNC ON)
+set(IREE_HAL_EXECUTABLE_LOADER_DEFAULTS OFF)
+set(IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF ON)
+
+list(REMOVE_ITEM iree_runtime_sources
+  "${CMAKE_SOURCE_DIR}/src/emscripten_api.cc"
+  "${CMAKE_SOURCE_DIR}/src/emscripten_api.h"
+)
+
+add_library(${_NAME} SHARED ${iree_runtime_sources})
+
+set_target_properties(${_NAME} PROPERTIES
+  RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+  ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+  SUFFIX ".so"
+)
+
+install(
+  TARGETS ${_NAME}
+  DESTINATION "."
+)
+
+target_link_libraries(${_NAME} iree_runtime_runtime)
+target_link_libraries(${_NAME} iree_tooling_context_util)
+
+# Install the header files - this will make it easier to copy them over
+# to the final bundle.
+install(DIRECTORY "${IREE_RUNTIME_INCLUDE_PATH}"
+        DESTINATION "include"
+        FILES_MATCHING PATTERN "*.h")
+
+file(GLOB HEADER_FILES "${CMAKE_SOURCE_DIR}/src/*.h")
+
+install(FILES ${HEADER_FILES}
+        DESTINATION "include/nx_iree")
+
+if(APPLE)
+  # Although the compiler complains about not using these,
+  # things only work with them set
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -undefined dynamic_lookup -fvisibility=default")
+  set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,-export_dynamic")
+  check_c_compiler_flag("-arch arm64" ARM64_SUPPORTED)
+  if(ARM64_SUPPORTED)
+      set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DMAC_ARM64")
+  endif()
+else()
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -shared -fvisibility=default")
+  set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--export-dynamic")
+endif()