[Mlir-commits] [mlir] [mlir] Change `tensor.extract/insert` to take static/dynamic indices. (PR #104488)

[Jacques Pienaar]

jpienaar wrote:

I don't see this as proposing a change to MLIR core. This is a change to the Tensor dialect, making the insert and extract ops more uniform within the dialect itself.

I currently have inputs where I have O(40k) scalar constant operations whose only reason for existing is indexing. Its inefficient, its more difficult to test and cumbersome to write (using a little lazy cache per isolated from above context to avoid ending at O(400k) scalar constant operations that get CSE'd later). Its not special to indexing ops that one could capture constant values, but it is very wasteful for them (I forget, its 10x higher storage than just parameter).

We could also just introduce new ops: tensor.insert\_static and tensor.extract\_static that is all static. That is an option. It introduces different set of conditional "expansion" (e.g., where insert is expected today, one would need to consider the static insert too).

Now wrt API, indeed this change doesn't change those. An option, is to add method like `std::optional<int64_t> getConstantDim(i)` which would then either look at constant property or match on operand to try and extract constant value, and `std::optional<Value> getDynamicDim(i)` which would just return operand else nullopt. Still results in conditionals, but for static that's not new (match is already a conditional) and unseen by user, while for dynamic I actually don't know how its being used except for codegen later and there it is most likely post checking if constant already :) This also doesn't do any automatic promotion to static args, so its only if created such.

https://github.com/llvm/llvm-project/pull/104488