search

llvm / llvm-project

The LLVM Project is a collection of modular and reusable compiler and toolchain technologies.
http://llvm.org
Other
28.37k stars 11.72k forks source link

tensor.expand_shape cannot expand to multiple dynamic dimensions #51564

Open silvasean opened 2 years ago

silvasean commented 2 years ago
Bugzilla Link 52222
Version unspecified
OS Linux
CC @cathyzhyi,@joker-eph,@nicolasvasilache,@MaheshRavishankar

Extended Description

This came up in torch-mlir: we are lowering an (B1, B2, M, K) x (B1, B2, K, N) matmul to linalg.batch_matmul. This requires collapsing B1,B2 -> B and doing the matmul and then expanding back B -> B1,B2.

linalg.tensor_expand_shape cannot do such a reshaping. The workaround is to use a linalg.generic.

Maybe linalg.batch_matmul could allow variadic leading dimensions?

Still, it seems useful to support arbitrary expansions in linalg.tensor_expand_shape

MaheshRavishankar commented 2 years ago

The use of higher-D tensors (or reasoning about operations at a higher dimensional iteration space) really helped for fusion purposes, especially fusing with reshapes. I see this case as being slightly different. This is a named op, where every dimension carries some significance. So having multiple batch dimensions muddles that a bit. I cant think of a reasonable code-gen strategy that will not just fold all the batch dimensions to expose the most amount of parallelism.

My preference would be to leave the batch_matmul as a 4D iterations space (batch, M, N and K). THe issue then is having linalg.tensor_expand_shape work for dynamic dimensions. Its a known issue, with a known solution. We need extra operand to the op to specify the shape of the result. This operand would mostly be generated using linalg.init_tensor operation. This is similar to what is done for linalg.generic operations where sometimes the outs is specified only for its shape.

silvasean commented 2 years ago

IIRC, we want to move more towards the direction of higher-D tensors natively, so maybe the linalg.tensor_expand_shape thing isn't so pressing if we can make sure we handle higher-D inputs better.

silvasean commented 2 years ago

assigned to @nicolasvasilache

nicolasvasilache commented 2 years ago

We are planning to go significantly beyond the current capabilities, see: https://discourse.llvm.org/t/rfc-structured-codegen-beyond-rectangular-arrays/64707

In the meantime we could make the semantics more powerful for tensor.expand_shape and memref.expand_shape if this is still useful?

silvasean commented 2 years ago

Yes this is still useful for us.