MLIR one-to-many function signature conversion (due to vector unrolling)

Hi, I have a WIP implementation for one-to-many function signature conversion (due to vector unrolling). The purpose is to convert the vector types to the ones that are supported by the SPIR-V dialect.

Could someone please take a look and let me know whether you think I am on the right track or if I should do it differently? Here is the URL to my branch. Thanks!

Example Input

func.func @simple_ret(%arg0: vector<8xi32>) -> vector<8xi32> {
  return %arg0 : vector<8xi32>
}

Example Output

func.func @simple_ret(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> (vector<4xi32>, vector<4xi32>) {
    %cst = arith.constant dense<0> : vector<8xi32>
    %0 = vector.insert_strided_slice %arg0, %cst {offsets = [0], strides = [1]} : vector<4xi32> into vector<8xi32>
    %1 = vector.insert_strided_slice %arg1, %0 {offsets = [4], strides = [1]} : vector<4xi32> into vector<8xi32>
    %2 = vector.extract_strided_slice %cst {offsets = [0], strides = [1]} : vector<8xi32> to vector<4xi32>
    %3 = vector.extract_strided_slice %cst {offsets = [4], strides = [1]} : vector<8xi32> to vector<4xi32>
    return %2, %3 : vector<4xi32>,  vector<4xi32>
  }

@angelz913 could you show an example of input IR and the IR you would like to produce?

One thing I noticed is that you are using a greedy pattern rewrite instead of the 1-to-N dialect conversion. Not saying that that’s a bad choice. Both a valid implementation strategies. A greedy pattern rewrite is definitely “simpler” than a dialect conversion because there is less “magic” happening.

One benefit of the 1:N conversion is that you can reuse the infrastructure to convert func.func, func.return, func.call ops.

1 Like

More details have been added to the original post.

@ftynse Hi Alex, would you like to take a look if you have some time? Thank you!

From a cursory glance, looks alright to me. Given that the rest of the pass seems to use the greedy driver, it seems okay to also use it here, although one could consider the 1:N dialect conversion instead.

1 Like