Unable to create linalg.generic op using Python bindings

MLIR
,
1

I could previously create a linalg.generic op using the Python bindings

from mlir.dialects import linalg

generic_op = linalg.GenericOp(
                [output_type],
                [input_arg],
                [output_arg],
                ir.ArrayAttr.get([ir.AffineMapAttr.get(ir.AffineMap.get_permutation([0, 1])),
                                  ir.AffineMapAttr.get(ir.AffineMap.get_permutation([0, 1]))]),
                ir.ArrayAttr.get([ir.StringAttr.get("parallel")]*2),
            )

However, after pulling the latest code on main, this no longer works. I get an error:

error: 'linalg.generic' op attribute 'iterator_types' failed to satisfy constraint: Iterator type should be an enum.

This appears to be related to changes made on 11/10 – Replace “string” iterator_types attr with enums in LinalgInterface. Instead of passing “parallel” as a string attr, it now wants an enum with the same meaning.

I see the object I need in C++ (mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.td), but I can’t figure out how to create or access that in Python.

If this is simply missing from the Python bindings, I’m happy to improve things if someone is willing to lend some guidance.

2

@ftynse for the guidance

3

It looks like this new attribute was not exposed to Python bindings. It can be added to C API and Python bindings following the example in llvm-project/DialectSparseTensor.cpp at 6198f3abf56f10bee6b3ba771c2de4a922429d4d · llvm/llvm-project · GitHub but in DialectLinalg.cpp (presumably, I don’t know how the DialectUtils thing is layered wrt dialects).

Alternatively, you can obtain it by parsing, Attribute.parse('#linalg.iterator_type<parallel>'), but this is a hack.

4

Thanks for the reply. I’ll try to follow the example and get this added to the Python bindings.

And the hack is helpful as well. I knew I could use Module.parse and write all of my MLIR as a string, but I didn’t realize individual elements like Attribute could also parse strings.

5

FWIW, there was some discussion on python bindings for MLIR enum attributes (and eventually generating them automatically) as part of this (abandoned) diff:
https://reviews.llvm.org/D128300
It seems to me like there should be an automated way to do this - @stellaraccident had some thoughts in the comments for the review above, but I never got around to trying to implement them.