iree-org / iree

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

Dequantization + Extract Slice Problems #17642

Open IanWood1 opened 2 weeks ago

IanWood1 commented 2 weeks ago

17455 uncovered several problems when dealing with tensor.extract_slice that consume the results of dequantization-like linalg.generic ops. See this gist for an mlir example.

What should happen

The dequantization + ExtactSliceOp + consumer will be placed in the same dispatch AND bufferization will convert the ExtractSliceOp into a view instead of allocating an entirely new high bitwidth tensor.

dequantization + ExtactSliceOp + consumer
-> (make slice continuous via transpose)
dequantization + transpose + ExtactSliceOp + consumer
-> (move transpose before dequant)
transpose + dequantization + ExtactSliceOp + consumer
-> dequant and extractsliceop cloned into consumer

What is currently happening

Although dequantization + ExtactSliceOp + consumer gets placed in the same dispatch, bufferiation cannot handle the tensor.extract_slice since it is extracting on the innermost dim. Full logs

Problems