search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.85k stars 614 forks source link

dynamic_gather does not support properties on Windows #17385

Open bmacadam-sfu opened 6 months ago

bmacadam-sfu commented 6 months ago

What happened?

I exported a model from JAX to MLIR using shape polymorphism in the batch index, but ran into the following issue when I tried to compile on a Windows machine:

PS C:\Users\ben.macadam> iree-compile.exe --iree-llvmcpu-target-triple=x86_64-pc-windows-msvc --iree-input-type=auto --iree-hal-target-backends=llvm-cpu --iree-llvmcpu-link-embedded=false --iree-llvmcpu-link-static --output-format=vm-bytecode --iree-llvmcpu-static-library-output-path="C:\Users\ben.macadam\Downloads\47cee9104002533701f7f0a6ed9f934d.o" -o "C:\Users\ben.macadam\Downloads\47cee9104002533701f7f0a6ed9f934d.vmfb" "C:\Users\ben.macadam\Downloads\dyn_guidance.mlir"
C:\Users\ben.macadam\Downloads\dyn_guidance.mlir:532:12: error: invalid properties {dimension_numbers = #stablehlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true} for op stablehlo.dynamic_gather: this operation does not support properties
    %182 = "stablehlo.dynamic_gather"(%178, %179, %181) <{dimension_numbers = #stablehlo.gather<offset_dims = [0, 1], collapsed_slice_dims = [1], start_index_map = [1]>, indices_are_sorted = true}> : (tensor<?x33x2xf32>, tensor<1xi32>, tensor<3xi32>) -> tensor<?x2xf32> loc(#loc148)
           ^

This worked on the original linux machine using the python API and basic compilation settings, so I'm a bit confused why it is failing in this case.

Steps to reproduce your issue

  1. Go to '...'
  2. Click on '....'
  3. Scroll down to '....'
  4. See error

What component(s) does this issue relate to?

No response

Version information

IREE compiler version 20240410.859 @ b4273a4bfc66ba6dd8f62f6483d74d42a7b936f1

Additional context

It does seem to work with default settings in python on my original linux machine, so perhaps it has to do with the build target.

ScottTodd commented 6 months ago

This shouldn't be operating system specific. I suspect you have a different version of iree-compile or JAX on Linux.

(If we switch IREE and its deps to use VHLO instead of StableHLO this class of problems may go away)

bmacadam-sfu commented 6 months ago

Ok - I tried compiling everything on the same Windows machine and got a more useful error. I'll try to write a minimal reproduction in the morning.