iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.

Home Page:http://iree.dev/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

dynamic_gather does not support properties on Windows

bmacadam-sfu opened this issue · comments

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 @ b4273a4

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.

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)

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.