Enable split-k Fusions

[umangyadav]

DOR (Definition of Ready)

Ready

Description

rocMLIR can support some output fusions if perfConfig is a split-k. https://github.com/ROCm/rocMLIR-internal/issues/1664

rocMLIR can fuse these kind of functions: f(x+y) = f(x) + f(y). That would mean multiplication only. However, internally rocMLIR applies a simple trick to allow addition: f(x)=x+c, we can convert the constant to c/splitk. Therefore, rocMLIR can do following fusions with split-k.

add x, constant -> add x, constant/splitk mul x, constant -> mul x, constant add x, y (x!=y) -> add x, y/splitk add x, x -> add x, x -> mul x, 2 -> mul x, constant mul x, y (x!=y) -> mul x, y (if x==y this is not allowed)

Note that normalization transformation for addition constant/splitk happens inside rocMLIR already. Therefore MIGraphX should just create the fused module and pass it to rocMLIR for compilation.

Allowing addition and multiplication means it also supports division, subtraction and negation.

It also supports type conversions e.g. convert/cast oprations.

ReduceSum and ReduceMean are also supported.

https://github.com/ROCm/rocMLIR-internal/issues/1681

Task is for the MIGraphX to implement these fusions for the split-k.

DOD (Definition of Done)

Fusions are implemeneted. Tests are added for different dtypes. Tests are added with different kinds of fusions.