diff --git a/runtime/cmake/quidditch_module.cmake b/runtime/cmake/quidditch_module.cmake index 1317397..00be9ae 100644 --- a/runtime/cmake/quidditch_module.cmake +++ b/runtime/cmake/quidditch_module.cmake @@ -73,7 +73,7 @@ function(quidditch_module) set(_EXTRA_DEPENDS ${_RULE_DEPENDS}) if (_RULE_LLVM) list(APPEND _COMPILER_ARGS "--iree-hal-target-backends=llvm-cpu") - list(APPEND _COMPILER_ARGS "--iree-llvmcpu-debug-symbols=false") + list(APPEND _COMPILER_ARGS "--iree-llvmcpu-debug-symbols=true") list(APPEND _COMPILER_ARGS "--iree-llvmcpu-target-triple=riscv32-unknown-elf") list(APPEND _COMPILER_ARGS "--iree-llvmcpu-target-cpu=generic-rv32") list(APPEND _COMPILER_ARGS "--iree-llvmcpu-target-cpu-features=+m,+f,+d,+zfh") @@ -121,3 +121,24 @@ function(quidditch_module) LINKER_LANGUAGE C ) endfunction() + +# Use iree-turbine to convert a PyTorch model to MLIR. +# 'SRC' should be the path to the python file, 'DTYPE' one of "f32" or "F64" and +# 'DST' the path to the output file. +# The python script should take one positional argument, which is the 'DST' +# argument and output the MLIR file there. Additionally the 'DTYPE' flag is +# communicated via a `--dtype=` flag. +macro(iree_turbine) + cmake_parse_arguments(_RULE "" "SRC;DTYPE;DST" "" ${ARGN}) + + cmake_path(GET _RULE_SRC STEM filename) + cmake_path(ABSOLUTE_PATH _RULE_SRC BASE_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} OUTPUT_VARIABLE source_path) + + add_custom_command( + OUTPUT ${_RULE_DST} + COMMAND ${Python3_EXECUTABLE} ${source_path} ${_RULE_DST} --dtype=${_RULE_DTYPE} + WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} + DEPENDS ${_RULE_SRC} + COMMENT "Translating ${filename} using iree-turbine" + ) +endmacro() diff --git a/runtime/samples/nsnet2/CMakeLists.txt b/runtime/samples/nsnet2/CMakeLists.txt index e5655b1..5c2904c 100644 --- a/runtime/samples/nsnet2/CMakeLists.txt +++ b/runtime/samples/nsnet2/CMakeLists.txt @@ -1,12 +1,5 @@ -add_custom_command( - OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/nsnet2.mlirbc - COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/NsNet2.py nsnet2.mlirbc --dtype=f32 - WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} - DEPENDS NsNet2.py - COMMENT "Translating NsNet2.py using iree-turbine" -) - +iree_turbine(SRC NsNet2.py DST ${CMAKE_CURRENT_BINARY_DIR}/nsnet2.mlirbc DTYPE "f32") quidditch_module(SRC ${CMAKE_CURRENT_BINARY_DIR}/nsnet2.mlirbc LLVM) add_executable(NsNet2LLVM main.c)