Vectorization on complex numbers in linalg dialect

MLIR
1

Hi, all!

I’m working with complex number computation using Linalg dialect, but currently, the vectorize transform in Linalg doesn’t support complex values, as shown in the test:

llvm-project/mlir/test/Dialect/Linalg/vectorization.mlir

// CHECK-LABEL: func @test_do_not_vectorize_unsupported_element_types
func.func @test_do_not_vectorize_unsupported_element_types(%A : memref<8x16xcomplex>, %arg0 : complex) {
// CHECK-NOT: vector.broadcast
// CHECK-NOT: vector.transfer_write
linalg.generic {
indexing_maps = [affine_map<(m, n) → ()>, affine_map<(m, n) → (m, n)>],
iterator_types = [“parallel”, “parallel”]}
ins(%arg0 : complex)
outs(%A: memref<8x16xcomplex>) {
^bb(%0: complex, %1: complex) :
linalg.yield %0 : complex
}
return
}

I also tried the sparse tensor dialect’s vectorizer, which neither supports complex values. I’ve made the vectorization work by lowering the complex dialect when the computation is represented using affine loops.

I wonder, is there a solution to vectorize the complex values in the Linalg dialect?

2

Hi @yifeihe007!

Good point! I’m not very familiar with the representation and lowering of the Complex type but I’m assuming that vectorizing a tensor of complex elements would be like vectorizing an array of structs, where the struct has two elements. It’s likely that you want to lower the Complex type before vectorization, as you did, to generate something friendlier to vectorization. Would this be a long term solution for you? It would be possible to add support for the Complex type to the vectorizer but I’m not sure if the complexity is worth it.

Hopefully that helps!