Page MenuHomePhabricator

[mlir][Affine] Add affine.delinearize_index operation
ClosedPublic

Authored by christopherbate on Aug 16 2022, 4:07 PM.

Details

Summary

This change adds a new AffineDelinearizeIndexOp to the affine dialect.
The operation accepts an index type as well as a basis (array of index
values) representing how the index should be decomposed into a
multi-index. The decomposition obeys a canonical semantic that treats
the final basis element as "fastest varying" and the first basis element
as "slowest varying". A naive lowering of the operation using a sequence
of AffineApplyOps is given.

RFC was discussed on discourse here: https://discourse.llvm.org/t/rfc-tensor-extracting-slices-from-tensor-collapse-shape/64034

Diff Detail

Event Timeline

christopherbate requested review of this revision.Aug 16 2022, 4:07 PM
bondhugula requested changes to this revision.Aug 16 2022, 6:10 PM
bondhugula added inline comments.
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
1078

It's perhaps more readable if we added a keyword like "into" or "to" between %linear_index and the parenthetical list.

mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
6–12

We should just be using affine.apply ops here, which are a higher and a better level abstraction than arith ops; the former will enable better optimization: for eg., the delinearize_index results will typically flow into affine load/store ops and affine.apply ops will compose into them. And the affine.apply ops will anyway be converted to the arith dialect ones with lower-affine.

arith.divui -> floordiv
arith.remui -> mod

This revision now requires changes to proceed.Aug 16 2022, 6:10 PM

Address comments

Thanks @christopherbate, I consider the comments unambiguously addressed.
Please update the commit message that still mentions arith ops to use affine_apply instead.

LGTM predicated on the commit message update!

mlir/test/Dialect/Affine/affine-expand-index-ops.mlir
6–12

+1 to avoid jumping the abstraction gap and using affine.apply.

mlir/test/Dialect/Affine/invalid.mlir
504

newline

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
1078

the asm format was updated as requested but not the doc, please udpate it too before landing

bondhugula accepted this revision.Aug 18 2022, 6:51 AM

LGTM - but several minor comments.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
1104

Nit: drop extra blank line.

mlir/include/mlir/Dialect/Affine/Utils.h
307

Comments to be terminated with a period - here and below.

313

Nit: Extra space before and.

313

a -> lhs, b -> rhs?

316

basis is absent.

mlir/lib/Dialect/Affine/Transforms/AffineExpandIndexOps.cpp
14–20

Can combine these together.

56

Empty private section; please drop.

mlir/lib/Dialect/Affine/Utils/Utils.cpp
1844

Use ... e = ...; i < e form.

1856

You can simply use dimSizes.drop_front(i).

1863

results.reserve(divisors.size());

This revision is now accepted and ready to land.Aug 18 2022, 6:51 AM
christopherbate marked 15 inline comments as done.

Address comments

christopherbate edited the summary of this revision. (Show Details)Aug 19 2022, 8:08 AM