Fix crash due to complex types not being considered in `KernelDispatch.cpp`

iree-org / iree

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

http://iree.dev/

Apache License 2.0

2.85k stars 622 forks source link

Fix crash due to complex types not being considered in `KernelDispatch.cpp` #19185

Closed giacs-epic closed 1 week ago

giacs-epic commented 1 week ago

When querying the element type's bitwidth via Type::getIntOrFloatBitwidth() the underlying assumption is that Type is IntegerType or FloatType. However ComplexType is also a valid type in this situation.

Adds a computation of the complex element type size.

hanhanW commented 1 week ago

I'm not aware that we have complex types as inputs. Can you elaborate how you discover the case and how does it get to the CPU codegen inputs?

I just learned from the other comment that we should support complex types. This looks good to me, but a test is needed. Perhaps we can let @IanWood1 land his PR with the change, and add you to co-author. What do you think?

IanWood1 commented 1 week ago

Fixed with https://github.com/iree-org/iree/pull/19172 and added as coauthor

giacs-epic commented 6 days ago

@hanhanW @IanWood1 Thanks for taking care of this, absolutely fine with how this was landed. Just to add a further example, this comes from implementing torch.stft https://pytorch.org/docs/stable/generated/torch.stft.html which has complex outputs. In my implementation the output is returned from a torch.prim.Loop and this triggered the crash.