Error when using a permutation_map in vector.transfer_write

MLIR
1

Hello,

I am trying to use the permutation_map attribute of vector.transfer_write and I am stuck with an error I don’t really understand.

Here is the culprit:

%14 = vector.transfer_write %13, %arg8[%7, %9, %c0_3, %c0_4] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)> } : vector<8x4xf32>, tensor<?x?x4x8xf32>

It gives me the following error:

error: ‘vector.transfer_write’ op requires a permutation_map with input dims of the same rank as the source type

Looking at the code I found that “source” was a sort of misnomer (it’s not refering to the input vector but the output tensor): this error is triggered by llvm-project/mlir/lib/Dialect/Vector/VectorOps.cpp, line 2576:

if (permutationMap.getNumInputs() != shapedType.getRank())	
	return op->emitOpError("requires a permutation_map with input dims of the "
	"same rank as the source type");

My question is: why is this code comparing the permutation_map dimensions and the tensor dimension? I would have expected the permutation to happen with the input vector and be lowered to (according to the rewriting pattern in mlir/lib/Dialect/Vector/VectorTransferPermutationMapRewritePatterns.cpp.cpp:143):

%foo = vector.transpose %13 [1,0] : vector<4x8xf32>
%14 = vector.transfer_write %foo, %arg8[%7, %9, %c0_3, %c0_4] {in_bounds = [true, true] ,  permutation_map = affine_map<(d0, d1) -> (d0, d1)>} : vector<4x8xf32>, tensor<?x?x4x8xf32>

I would also expect that new permutation_map to be dropped at some point because it’s just an identity map.

Which would not require me to have to use a 4-dimension permutation_map.

I guess I am wrong with my assumptions and it’s important to keep a permutation_map with as many input dims as the output tensor, but I don’t fully understand why? I also tried to give the vector.transfer_write a 4-dims permutation map (i.e. (d0,d1,d2,d3) → (d0,d1,d3,d2)) but it fails with:

mlir::Attribute mlir::ArrayAttr::operator(unsigned int) const: Assertion `idx < size() && “index out of bounds”’ failed.

Because I imagine it’s trying to read too far into the input vector.

Can someone help me understand this?

Thanks

EDIT: After reading the code of VectorTransferPermutationMapRewritePatterns, I found out that doing (d0,d1,d2,d3) → (d3,d2) is working fine. I’m still not sure why I need 4 dimensions as input though.

2

cc @giuseros

3

Yes, this is an unfortunate remnant of early use of interface, a bit of renaming would be welcome, I’ll try to squeeze in on my stack.

It can be elided if it is a minor identity (as per the doc):
“The permutation map may be implicit and omitted from parsing and printing if it is the canonical minor identity map (i.e. if it does not permute any dimension)”

If it does not elide properly then that should be fixed.

You have 4 tensor dimensions and 2 vector dimensions; each vector dimension can map to any one of the tensor dimensions and you need a way to specify that:

// write vector<16x32x64xf32> into the slice
//   `%A[%i0, %i1:%i1+32, %i2:%i2+64, %i3:%i3+16]`
vector.transfer_write %val, %A[%i0, %i1, %i2, %i3]
  {permutation_map: (d0, d1, d2, d3) -> (d3, d1, d2)} :
   vector<16x32x64xf32>, memref<?x?x?x?xf32>

you could also e.g. want to match 32 on the tensor/memref dimension 0 instead and you would write:

// write vector<16x32x64xf32> into the slice
//   `%A[%i0, %i1:%i1+32, %i2:%i2+64, %i3:%i3+16]`
vector.transfer_write %val, %A[%i0, %i1, %i2, %i3]
  {permutation_map: (d0, d1, d2, d3) -> (d3, **d0**, d2)} :
   vector<16x32x64xf32>, memref<?x?x?x?xf32>
4

Thank you for your message!

Yes indeed and it works fine when using a correct map, no need to fix anything.

I think I was confused about this because for me it looked like the mapping defined by the permutation map was the wrong way around (4 dimensions mapped to 2 dimensions). I think it makes sense now though because if I wrote it as I was expecting it, I wouldn’t know what to put for the first 2 dimensions of the tensor: (d0,d1) → (?,?, d1,d0).

Thank you!