This is an archive of the discontinued LLVM Phabricator instance.

[mlir][OpDSL] Add type function attributes.
ClosedPublic

Authored by gysit on Feb 14 2022, 7:07 AM.

Details

Summary

Previously, OpDSL operation used hardcoded type conversion operations (cast or cast_unsigned). Supporting signed and unsigned casts thus meant implementing two different operations. Type function attributes allow us to define a single operation that has a cast type function attribute which at operation instantiation time may be set to cast or cast_unsigned. We may for example, defina a matmul operation with a cast argument:

@linalg_structured_op
def matmul(A=TensorDef(T1, S.M, S.K), B=TensorDef(T2, S.K, S.N), C=TensorDef(U, S.M, S.N, output=True),
    cast=TypeFnAttrDef(default=TypeFn.cast)):
  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])

When instantiating the operation the attribute may be set to the desired cast function:

linalg.matmul(lhs, rhs, outs=[out], cast=TypeFn.cast_unsigned)

The revsion introduces a enum in the Linalg dialect that maps one-by-one to the type functions defined by OpDSL.

Diff Detail

Event Timeline

gysit created this revision.Feb 14 2022, 7:07 AM
gysit requested review of this revision.Feb 14 2022, 7:07 AM
gysit updated this revision to Diff 408430.Feb 14 2022, 8:06 AM

Improve tests.

gysit added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
40

This is needed since the operation has a cast member function to access the cast attribute.

gysit updated this revision to Diff 409174.Feb 16 2022, 1:34 AM

Update OperandKind enum to match coding conventions.

gysit updated this revision to Diff 409694.Feb 17 2022, 9:27 AM

Rebase and cleaner test.

aartbik added inline comments.Feb 18 2022, 10:47 AM
mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
4

why is this line removed?

553

this, too, seems unrelated to the revision at large?

gysit updated this revision to Diff 409999.Feb 18 2022, 12:34 PM

Attempt to fix autoformatting

gysit marked 2 inline comments as done.Feb 18 2022, 12:37 PM

At least for now the auto formatting noise seems gone...

gysit updated this revision to Diff 410240.Feb 21 2022, 1:00 AM

Use custom enum format.

LG for the pytaco part

Aart, would you mind doing a review of the entire patch? (I've got a lot of latency this week and would also like to expand the reviewer pool here -- and you are an active user of this)?

gysit updated this revision to Diff 410795.Feb 23 2022, 6:12 AM

Fix enum mapping and add addtional check to avoid attribute name clashes.

did a bit more thorough pass over the other parts too

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
63

Although I assume you want the common case to be shorter, there is something to say to use cast_signed and cast_unsigned to make the distinction explicit

65

would a cast_nop make sense
(i.e. use the OpsDSL for cases where there is *no* cast in the instantation, even though pattern defines one)

mlir/lib/Dialect/Linalg/IR/LinalgDialect.cpp
158

Thanks for moving that here. Having this at the wrong place can be very confusing
(even though it typically only defines the destructor)

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py
172

attached is a bit strange term here, how about specified, like L164

554

typo

669

clashes -> conflicts

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py
138

typo: attribute

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir
42

Verifies that ....

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp
191–193

_vals reads weird now that we have _fn, can we rename?

gysit marked 5 inline comments as done.Feb 24 2022, 3:26 AM
gysit added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
63

I would suggest to do this in a separate revision together with max -> max_signed. I will change this in OpDSL as well and many places will be affected. I think OpDSL historically was meant to be signed by default and we added unsigned later. I think now that we can use attributes in most places going for the longer names in the default case should be fine.

65

I an a bit unsure about this one. Both cast and cast_unsigned are nops if the types match. A cast_nop/no_cast would result in a verification failure after op construction and we may emit a warning. It may not be easy to understand though since the region of named operations is hidden. I am thus a bit unsure if this would be a good addition.

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py
172

Yes, I switched to "not registered with an op" since this used in most other places. The check in fact makes sure a specific operand has been registered with an operation. Otherwise it has no name yet.

gysit updated this revision to Diff 411059.Feb 24 2022, 3:28 AM
gysit marked an inline comment as done.

Address comments.

aartbik accepted this revision.Feb 24 2022, 10:32 AM
aartbik added inline comments.
mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
65

Fair enough. I was just thinking out loud here.

This revision is now accepted and ready to land.Feb 24 2022, 10:32 AM
gysit added inline comments.Feb 24 2022, 11:07 AM
mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td
65

Thanks for the review and especially for thinking out loud! A cast_nop is IMO a good idea! It is just not so easy to implement in the current framework.

This revision was automatically updated to reflect the committed changes.