[mlir][Vector] Turn Vector Linearization from a conversion to an IR rewrite?

[dcaballe]

Hey,

I've been looking into the Vector Linearization transformation and hit a few issues. The current implementation uses the Conversion pattern driver and mostly makes every operation with a multi-dimensional vector type illegal using this target method. I see a couple of issues with this:

1. If an operation is set as illegal but we don't have a conversion pattern for it or existing conversion patterns fail for some reason, the whole linearization conversion will fail. This will happen even when using applyPartialConversion because the target method mentioned above will explicit set operations as illegal. To give you a concrete example, I'm currently hitting a conversion error because scalar vector.insert operations are not supported:

test2.mlir:13:12: error: failed to legalize operation 'vector.insert' that was explicitly marked illegal
      %8 = vector.insert %7, %cst [0, 0] : f32 into vector<16x8xf32>

I hit more issues with other ops. Note that not linearizing an n-D operation shouldn't be a big deal as we have lowering patterns to LLVM that can handle multi-dimensional vectors so it should be fine to continue the compilation regardless or any implementation gaps in this pass.

2. The linearization "conversion" is similar to other transformations that we have in the Vector dialect and are implemented as simple IR rewrites which can be applied using the greedy pattern rewriter driver. However, since we are using the conversion driver for vector linearization, it can't be composed with these IR rewrites. I think the idea of using the type converter and target to model the n-D -> 1-D vector type conversion is fine but this kind of vector "reshape" transformations are common in different vector rewrites and we don't use the conversion infrastructure for them.

All in all, I'm really excited to see the progress we are making in this transformation. I think it's reaching a point where we can start using it more broadly and make it more composable with other existing vector rewrites. As such, we may want to consider moving the current Conversion based implementation to a simpler IR rewrite. I'd love to hear what others think!

cc: @newling, @Hardcode84, @banach-space

[llvmbot]

@llvm/issue-subscribers-mlir

Author: Diego Caballero (dcaballe)

Hey,

I've been looking into the Vector Linearization transformation and hit a few issues. The current implementation uses the Conversion pattern driver and mostly makes every operation with a multi-dimensional vector type illegal using this target method. I see a couple of issues with this:

1. If an operation is set as illegal but we don't have a conversion pattern for it or existing conversion patterns fail for some reason, the whole linearization conversion will fail. This will happen even when using applyPartialConversion because the target method mentioned above will explicit set operations as illegal. To give you a concrete example, I'm currently hitting a conversion error because scalar vector.insert operations are not supported:

test2.mlir:13:12: error: failed to legalize operation 'vector.insert' that was explicitly marked illegal
      %8 = vector.insert %7, %cst [0, 0] : f32 into vector<16x8xf32>

I hit more issues with other ops. Note that not linearizing an n-D operation shouldn't be a big deal as we have lowering patterns to LLVM that can handle multi-dimensional vectors so it should be fine to continue the compilation regardless or any implementation gaps in this pass.

2. The linearization "conversion" is similar to other transformations that we have in the Vector dialect and are implemented as simple IR rewrites which can be applied using the greedy pattern rewriter driver. However, since we are using the conversion driver for vector linearization, it can't be composed with these IR rewrites. I think the idea of using the type converter and target to model the n-D -> 1-D vector type conversion is fine but this kind of vector "reshape" transformations are common in different vector rewrites and we don't use the conversion infrastructure for them.

All in all, I'm really excited to see the progress we are making in this transformation. I think it's reaching a point where we can start using it more broadly and make it more composable with other existing vector rewrites. As such, we may want to consider moving the current Conversion based implementation to a simpler IR rewrite. I'd love to hear what others think!

cc: @newling, @Hardcode84, @banach-space

[dcaballe]

Also @akroviakov

[Hardcode84]

My motivation for using dialect conversion #136581 (comment)

[newling]

I'm in favour of changing to use simple greedy patterns. I think the structural ops can still be converted without too much difficulty in a separate conversion phase to achieve this.

[banach-space]

+1 to consistently using simple rewriters + greedy pattern rewriter

[dcaballe]

Yes, is it possible to do the scf conversion in a separate step? That would be helpful.

I couldn't find a way to use the transformation with the current implementation. Simple ops like scalar vector insert/extract ops are not supported and even if we ignore the LogicalResult returned by applyPartialConversion, we still get errors like the following dumped to stderr:

error: failed to legalize operation 'vector.extract' that was explicitly marked illegal

[newling]

Yes, is it possible to do the scf conversion in a separate step? That would be helpful.

I'm fairly confident it can be. Flow would be something like

input:

func.func @test_linearize_across_for(%arg0 : vector<4xi8>) -> vector<4xi8> {
  %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<2x2xi8>
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c4 = arith.constant 4 : index
  %1 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg1 = %0) -> (vector<2x2xi8>) {
    %2 = arith.addi %arg1, %0 : vector<2x2xi8>
    scf.yield %2 : vector<2x2xi8>
  }
  %3 = vector.shape_cast %1 : vector<2x2xi8> to vector<4xi8>
  return %3 : vector<4xi8>
}

run basic patterns to reduce rank:

func.func @test_linearize_across_for(%arg0 : vector<4xi8>) -> vector<4xi8> {
  %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<2x2xi8>
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c4 = arith.constant 4 : index
  %1 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg1 = %0) -> (vector<2x2xi8>) {
    %flat_arg1 = vector.shape_cast %arg1 : vector<2x2xi8> to vector<4xi8>
    %flat_0 = vector.shape_cast %0 : vector<2x2xi8> to vector<4xi8>
    %2 = arith.addi %flat_arg1, %flat_0 : vector<4xi8>
    %reshaped = vector.shape_cast %2 : vector<4xi8> to vector<2x2xi8>
    scf.yield %reshaped : vector<2x2xi8>
  }
  %3 = vector.shape_cast %1 : vector<2x2xi8> to vector<4xi8>
  return %3 : vector<4xi8>
}

legalize structured ops, fold shape_casts:

func.func @test_linearize_across_for(%arg0 : vector<4xi8>) -> vector<4xi8> {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c4 = arith.constant 4 : index
  %1 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg1 = %arg0) -> (vector<4xi8>) {
    %2 = arith.addi %arg1, %arg0 : vector<4xi8>
    scf.yield %2 : vector<4xi8>
  }
  return %1 : vector<4xi8>
}

I couldn't find a way to use the transformation with the current implementation. Simple ops like scalar vector insert/extract ops are not supported and even if we ignore the LogicalResult returned by applyPartialConversion, we still get errors like the following dumped to stderr: error: failed to legalize operation 'vector.extract' that was explicitly marked illegal

I've hit that scalar edge case before too :/. One solution would be to add an additional legalization to the target to say that vector.extract to scalar is legal. Like done in: https://github.com/llvm/llvm-project/blame/da8e03a560fc0850b280b39cb82e0df8e86a4fce/mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp#L908 But we should also get a fix for the scalar case in at some point (before/after any major changes which seem to be on the horizon)

[dcaballe]

I have implemented the scalar cases. I'll create a PR soon. Anyone else that can help with this issue? :)

[newling]

I have implemented the scalar cases. I'll create a PR soon.

Thanks!

Anyone else that can help with this issue? :)

I would like to help on this issue. I have made some changes in this direction while working on changes to bitwise logic in: https://discourse.llvm.org/t/rfc-mlir-vector-modify-remove-bit-width-specific-flattening-linearization/86876/8

If there is no immediate requirement for this change, I propose we touch base here in a week or so. Hopefully I'll be nearly done by then, if not I might ask for additional help. Does that sound like a reasonable plan?