[mlir] Fix build & test of mlir python bindings on Windows

[stella.stamenova]

Summary

There are a couple of issues with the python bindings on Windows:

• create_symlink requires special permissions on Windows - using copy_if_different instead allows the build to complete and then be usable
• the path to the python_executable is likely to contain spaces if python is installed in Program Files. llvm's python substitution adds extra quotes in order to account for this case, but mlir's own python substitution does not
• the location of the shared libraries is different on windows
• if the type is not specified for numpy arrays, they appear to be treated as strings

I've implemented the smallest possible changes for each of these in the patch, but I would actually prefer a slightly more comprehensive fix for the python_executable and the shared libraries.

For the python substitution, I think it makes sense to leverage the existing %python instead of adding %PYTHON and instead add a new variable for the case when preloading is needed. This would also make it clearer which tests are which and should be skipped on platforms where the preloading won't work.

For the shared libraries, I think it would make sense to pass the correct path and extension (possibly even the names) to the python script since these are known by lit and don't have to be hardcoded in the test at all.

[mehdi_amini]

mlir/test/python/dialects/shape.py

26

I didn't quite get the issue here?

[stella.stamenova]

mlir/test/python/dialects/shape.py

26

This was much fun (?) to investigate. The test expects the produced shape to look like:

shape.const_shape [10, 20] : tensor<2xindex>

Instead, on Windows, we get:

shape.const_shape [85899345930, 85899345930] : tensor<2xindex>

As far as I can tell, what's happening is that on Windows, unless the type is explicitly specified, the values are treated as a string when passing from the python layer to the C++ layer. Then the DenseElementsAttr on Windows is treated as splat resulting in the 85899345930 value (85899345930 is 0x140000000A; 0x14 is 20 and 0xA is 10). If the type is specified, everything works as expected. There are some other tests that also explicitly specify the type, so I opted to do so and did not investigate further.

[stellaraccident]

Thanks for the fixes. I'm out today but will have a detailed look this weekend. Over time, I've tried to keep these things working for Windows but we must be building with a Windows setup that doesn't tickle the weird bits.

One request: can we special case the copy_if_different thing for Windows only. Symlinks are pretty convenient for dev because you can just edit/run, and I'd prefer to not regress that too account for this weird Windows-ism.

[stella.stamenova]

In D125122#3497606, @stellaraccident wrote:

Thanks for the fixes. I'm out today but will have a detailed look this weekend. Over time, I've tried to keep these things working for Windows but we must be building with a Windows setup that doesn't tickle the weird bits.

One request: can we special case the copy_if_different thing for Windows only. Symlinks are pretty convenient for dev because you can just edit/run, and I'd prefer to not regress that too account for this weird Windows-ism.

I updated it to use symlinks when not on Windows. This is the smallest set of changes needed to work on Windows and once it's in, I'd like to turn on the build of the python bindings on the buildbot as well.

[stellaraccident]

In D125122#3497845, @stella.stamenova wrote:

In D125122#3497606, @stellaraccident wrote:

Thanks for the fixes. I'm out today but will have a detailed look this weekend. Over time, I've tried to keep these things working for Windows but we must be building with a Windows setup that doesn't tickle the weird bits.

One request: can we special case the copy_if_different thing for Windows only. Symlinks are pretty convenient for dev because you can just edit/run, and I'd prefer to not regress that too account for this weird Windows-ism.

I updated it to use symlinks when not on Windows. This is the smallest set of changes needed to work on Windows and once it's in, I'd like to turn on the build of the python bindings on the buildbot as well.

Thanks. I am open to more refactorings once we do that. Thank you for the help.

I'll need to check when I'm back online, but I could have sworn we had this enabled on a bot somewhere (I've broken it before). But in any case, +1 on more coverage.

[stella.stamenova]

I'll need to check when I'm back online, but I could have sworn we had this enabled on a bot somewhere (I've broken it before). But in any case, +1 on more coverage.

The python bindings are enabled on several different linux/ubuntu bots, but not on the windows bot (https://lab.llvm.org/buildbot/#/builders/13). I sent the change for that as well (https://reviews.llvm.org/D125134), but I'm not planning on committing it before this weekend even if I commit the change to the source because I don't want to deal with potential breaks over the weekend :).

[mehdi_amini]

mlir/test/python/dialects/shape.py

26

I'm concerned about fixing the test instead of fixing the infrastructure: this seems like the kind of things that will hit users.

[stellaraccident]

LGTM modulo nits and open discussion

mlir/cmake/modules/AddMLIRPython.cmake

235

Cmake style isn't really a thing, but I do generally try to be consistent within a macro/file. In this case, "local" variables are named _like_this. Can we just call this _link_or_copy?

mlir/test/lit.cfg.py

97

So. Quites to quotes

mlir/test/python/dialects/shape.py

26

In this specific case, I'd be open to an issue and references to it in the code. These kind of platform specific corners are hard to get a complete lock on in one step, and I would bias towards landing this so that we get bot coverage, and then fix from there.

mlir/test/python/execution_engine.py

348

This could be made generic with sysconfig.get_config_var('SHLIB_SUFFIX') (and I think there is a corresponding prefix). I'm a little uncertain how/if this is working on macos, which is also different. Certainly, if we had three cases here, I would generalize it in some way.

I'm ok to leave as is for now since Windows is different in many details here.

[stella.stamenova]

mlir/cmake/modules/AddMLIRPython.cmake

235

I updated it to match this file. I had originally followed the pattern used in llvm (addllvm.cmake, etc.). I wish there was a consistent style within llvm-project :(.

mlir/test/python/dialects/shape.py

26

My goal with these changes is to get things working enough on Windows, so we can add it to the bot. I think this won't be a simple change to fix and I already spent quite a bit of time trying to track it down.

mlir/test/python/execution_engine.py

348

We would actually need three different things: prefix, suffix and path (bin on Windows, lib most everywhere else). This would be needed for a number of the other tests that are currently disabled on Windows as well. I also noticed that some of the python tests set the environment before making the call to their respective scripts and the same could be done here. My preference would be to fix this locally in the file for the moment and then have a more comprehensive change to normalize how the shared libraries are set among all of the tests.

[stellaraccident]

mlir/test/python/dialects/shape.py

26

I agree. Can you file an issue. Feel free to assign it to me. Since the test fix is subtle, I don't have confidence that I'll remember the repro workout details.

mlir/test/python/execution_engine.py

348

Yeah. I think the prefix and suffix can come from the environment but the location is an llvm thing. Fine with me to do this as you have it. Then we can have a better look.

I can't say I love all this open coded execution engine library integration, and I've also made point fixes to it over time (but have never looked at the whole).

[mehdi_amini]

mlir/test/python/dialects/shape.py

26

Perfectly fine with me to track this as an issue to fix later: I just wanted to make sure we agree there is an issue to be fixed (and we track it) :)

mlir/cmake/modules/AddMLIRPython.cmake

if(_source_type STREQUAL "pure")
         endif()
         set(_src_path "${_python_root_dir}/${_source_relative_path}")
         set(_dest_path "${ARG_ROOT_PREFIX}/${_dest_relative_path}")
 
         get_filename_component(_dest_dir "${_dest_path}" DIRECTORY)
         get_filename_component(_install_path "${ARG_INSTALL_PREFIX}/${_dest_relative_path}" DIRECTORY)
 
         file(MAKE_DIRECTORY "${_dest_dir}")
+
+        # On Windows create_symlink requires special permissions. Use copy_if_different instead.
+        if(CMAKE_HOST_WIN32)
+          set(_link_or_copy copy_if_different)
+        else()
+          set(_link_or_copy create_symlink)
+        endif()
+
         add_custom_command(
           TARGET ${modules_target} PRE_BUILD
           COMMENT "Copying python source ${_src_path} -> ${_dest_path}"
           DEPENDS "${_src_path}"
           BYPRODUCTS "${_dest_path}"
-          COMMAND "${CMAKE_COMMAND}" -E create_symlink
+          COMMAND "${CMAKE_COMMAND}" -E ${_link_or_copy}
               "${_src_path}" "${_dest_path}"
         )
         install(
           FILES "${_src_path}"
           DESTINATION "${_install_path}"
           COMPONENT ${modules_target}
         )
       endforeach()
function(add_mlir_python_modules name)
 
   # Build the modules target.
   add_custom_target(${name} ALL DEPENDS ${_depends})
   foreach(sources_target ${_flat_targets})
     _process_target(${name} ${sources_target})
   endforeach()
 
   # Create an install target.
-  if (NOT LLVM_ENABLE_IDE)
+  if(NOT LLVM_ENABLE_IDE)
     add_llvm_install_targets(
       install-${name}
       DEPENDS ${name}
       COMPONENT ${name})
   endif()
 endfunction()
 
 # Function: declare_mlir_dialect_python_bindings

 # Build python extension
 ################################################################################
 function(add_mlir_python_extension libname extname)
   cmake_parse_arguments(ARG
   ""
   "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
   "SOURCES;LINK_LIBS"
   ${ARGN})
-  if (ARG_UNPARSED_ARGUMENTS)
+  if(ARG_UNPARSED_ARGUMENTS)
     message(FATAL_ERROR " Unhandled arguments to add_mlir_python_extension(${libname}, ... : ${ARG_UNPARSED_ARGUMENTS}")
   endif()
-  if ("${ARG_SOURCES}" STREQUAL "")
+  if("${ARG_SOURCES}" STREQUAL "")
     message(FATAL_ERROR " Missing SOURCES argument to add_mlir_python_extension(${libname}, ...")
   endif()
 
   # The actual extension library produces a shared-object or DLL and has
   # sources that must be compiled in accordance with pybind11 needs (RTTI and
   # exceptions).
   pybind11_add_module(${libname}
     ${ARG_SOURCES}
if(WIN32)
     # control where the .dll gets written.
     set_target_properties(
       ${libname} PROPERTIES
       RUNTIME_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
       ARCHIVE_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
     )
   endif()
 
-  # Python extensions depends *only* on the public API and LLVMSupport unless
-  # if further dependencies are added explicitly.
   target_link_libraries(${libname}
     PRIVATE
     ${ARG_LINK_LIBS}
-    ${PYEXT_LIBADD}
   )
 
   target_link_options(${libname}
     PRIVATE
       # On Linux, disable re-export of any static linked libraries that
       # came through.
       $<$<PLATFORM_ID:Linux>:LINKER:--exclude-libs,ALL>
   )
 
   ################################################################################
   # Install
   ################################################################################
-  if (ARG_INSTALL_DIR)
+  if(ARG_INSTALL_DIR)
     install(TARGETS ${libname}
       COMPONENT ${ARG_INSTALL_COMPONENT}
       LIBRARY DESTINATION ${ARG_INSTALL_DIR}
       ARCHIVE DESTINATION ${ARG_INSTALL_DIR}
       # NOTE: Even on DLL-platforms, extensions go in the lib directory tree.
       RUNTIME DESTINATION ${ARG_INSTALL_DIR}
     )
   endif()
 endfunction()

mlir/test/lit.cfg.py

tools.extend([
     ToolSubst('%mlir_integration_test_dir', config.mlir_integration_test_dir, unresolved='ignore'),
 ])
 
 python_executable = config.python_executable
 # Python configuration with sanitizer requires some magic preloading. This will only work on clang/linux.
 # TODO: detect Darwin/Windows situation (or mark these tests as unsupported on these platforms).
 if "asan" in config.available_features and "Linux" in config.host_os:
   python_executable = f"LD_PRELOAD=$({config.host_cxx} -print-file-name=libclang_rt.asan-{config.host_arch}.so) {config.python_executable}"
+# On Windows the path to python could contains spaces in which case it needs to be provided in quotes.
+# This is the equivalent of how %python is setup in llvm/utils/lit/lit/llvm/config.py.
+elif "Windows" in config.host_os:
+  python_executable = '"%s"' % (python_executable)
 tools.extend([
   ToolSubst('%PYTHON', python_executable, unresolved='ignore'),
 ])
 
 llvm_config.add_tool_substitutions(tools, tool_dirs)
 
 
 # FileCheck -enable-var-scope is enabled by default in MLIR test

mlir/test/python/dialects/shape.py

def testConstShape():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f32 = F32Type.get()
     with InsertionPoint(module.body):
       @func.FuncOp.from_py_func(
           RankedTensorType.get((12, -1), f32))
       def const_shape_tensor(arg):
         return shape.ConstShapeOp(
-          DenseElementsAttr.get(np.array([10, 20]), type=IndexType.get()))
+          DenseElementsAttr.get(np.array([10, 20], dtype=np.int64), type=IndexType.get()))
 
     # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
     # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
     print(module)

mlir/test/python/execution_engine.py

module = Module.parse("""
       }
       func.func private @printMemrefF32(memref<*xf32>) attributes { llvm.emit_c_interface }
      } """)
     arg0 = np.array([0.0]).astype(np.float32)
 
     arg0_memref_ptr = ctypes.pointer(
         ctypes.pointer(get_ranked_memref_descriptor(arg0)))
 
-    execution_engine = ExecutionEngine(
-        lowerToLLVM(module),
-        opt_level=3,
-        shared_libs=[
+    if sys.platform == 'win32':
+        shared_libs = [
+            "../../../../bin/mlir_runner_utils.dll",
+            "../../../../bin/mlir_c_runner_utils.dll"
+        ]
+    else:
+        shared_libs = [
             "../../../../lib/libmlir_runner_utils.so",
             "../../../../lib/libmlir_c_runner_utils.so"
-        ])
+        ]
+
+    execution_engine = ExecutionEngine(
+        lowerToLLVM(module),
+        opt_level=3,
+        shared_libs=shared_libs)
     execution_engine.invoke("main", arg0_memref_ptr)
     # CHECK: Unranked Memref
     # CHECK-NEXT: [42]
 
 
 run(testSharedLibLoad)
 
 
module = Module.parse("""
         %u_memref = memref.cast %memref : memref<1xi64> to memref<*xi64>
         call @printMemrefI64(%u_memref) : (memref<*xi64>) -> ()
         return
       }
       func.func private @nanoTime() -> i64 attributes { llvm.emit_c_interface }
       func.func private @printMemrefI64(memref<*xi64>) attributes { llvm.emit_c_interface }
     }""")
 
-    execution_engine = ExecutionEngine(
-        lowerToLLVM(module),
-        opt_level=3,
-        shared_libs=[
+    if sys.platform == 'win32':
+        shared_libs = [
+            "../../../../bin/mlir_runner_utils.dll",
+            "../../../../bin/mlir_c_runner_utils.dll"
+        ]
+    else:
+        shared_libs = [
             "../../../../lib/libmlir_runner_utils.so",
             "../../../../lib/libmlir_c_runner_utils.so"
-        ])
+        ]
+
+    execution_engine = ExecutionEngine(
+        lowerToLLVM(module),
+        opt_level=3,
+        shared_libs=shared_libs)
     execution_engine.invoke("main")
     # CHECK: Unranked Memref
     # CHECK: [{{.*}}]
 
 
 run(testNanoTime)