[RFC] EnumAttr for iterator types in Linalg

MLIR
1

Linalg uses “parallel” ,“reduction” and “window” strings to label dimensions of structured ops. This is mostly because when Linalg dialect started, there was no special attribute for enums.

I suggest to replace the strings with

def IteratorTypeParallel : I32EnumAttrCase<"parallel", 0>;
def IteratorTypeSequential : I32EnumAttrCase<"sequential", 1>;

def IteratorTypeEnum : I32EnumAttr<
    "IteratorType",
    "Iterator type",
    [IteratorTypeParallel, IteratorTypeSequential]> {
  let cppNamespace = "::mlir::linalg";
}

In this snippet I replaced “reduction” type with “sequential” because I see that very often it is used to indicate exactly that, e.g. SortOp::getLoopIteratorTypes.

@nicolasvasilache do we actually need “window”?

2

Fantastic Alex, please do!
Re. Window it is not yet load bearing but the idea is to use the information in heuristics to compute e.g. tile sizes.

If you are able to remove without pain go ahead and we can reintroduce when we really need it.

Thanks much!

1 Like
3

This is great - thank you.

4

I was looking into this refactoring and got stuck, because of a few moments. Not sure if it’s because of my limited MLIR knowledge or a genuine limitations.

  1. Print-parsing format of EnumAttr doesn’t work as a drop-in replacement for the existing array of String attributes, especially when wrapped into TypedArrayAttrBase. I was able to get two formats to work without implementing a custom logic:
iterator_types = [#linalg.iterator_type<parallel>, #linalg.iterator_type<sequential>];
iterator_types = [0 : i32, 1 : i32];

The first one is obviously better, but still far from perfect. The desired option would be to have

iterator_types = [parallel, sequential]; without double quotes

but for some reason it doesn’t get along with TypedArrayAttrBase.

  1. The same format of iterator_types is also used in Vector dialect: VectorOps.td and there is a code in a shared library that expects those two iterator_types to conform: include/mlir/Dialect/Utils/StructuredOpsUtils.h. I’m not sure if this change can be limited to only Linalg without touching Vector and StructuredOpsUtils.

We could place this IteratorType enum in a shared place, but I’m not sure what is this place should be.

@nicolasvasilache @pifon2a any suggestions? Am I missing something?

1 Like
5

@nicolasvasilache Iterator types seem to be used by the VectorDialect as well. It would make sense to put the attribute to a common place, e.g. Dialect/Utils/IR/utils_enums.td. EnumAttr requires a dialect name. Can we use utils as the dialect name?

@mehdi_amini, WDYT?

6

Sorry I missed this …

I would split the 2 and duplicate: the iterator types for Linalg and Vector will diverge as we introduce more on the Linalg side.

1 Like
7

There is also StructuredGenerator in StructuredOpsUtils. Should we move it to LinalgOps and to VectorOps?

8

I would try to avoid duplicating the template helper unless absolutely necessary: the iterator types will soon have differences but the general underlying mechanisms should not AFAICS.

9

This can achieved by adding a custom array attribute whose element type is the enum attribute. In an op assembly format, it will produce the nice syntax you desire.

1 Like