[mlir][Interface] Factor out common IndexingMapOpInterface behavior in a new generic interface

[nicolasvasilache]

Refactor the verifiers to make use of the common bits and make vector.contract also use this interface.
In the process, the confusingly named getStaticShape has disappeared.

Note: the verifier for IndexingMapOpInterface is currently called manually from other verifiers as it was unclear how to avoid it taking precedence over more meaningful error messages

[llvmbot]

@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-vector

Author: Nicolas Vasilache (nicolasvasilache)

Changes

…contract

---

Full diff: https://github.com/llvm/llvm-project/pull/145313.diff

1 Files Affected:

• (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+33)

diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 926a92eff2ebb..7ac661d8bbf69 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -207,6 +207,39 @@ def Vector_ContractionOp :
               .template getAsValueRange<IteratorTypeAttr, IteratorType>();
       return {range.begin(), range.end()};
     }
+
+    //===------------------------------------------------------------------===//
+    // The code below is shared with LinalgStructuredInterface.
+    // vector.contract is really a linalg.generic on vectors without region.
+    // TODO: factor out in a common interface to inherit from ince identified.
+    //===------------------------------------------------------------------===//
+    ArrayRef<int64_t> getShape(OpOperand * opOperand) {
+      assert(opOperand->getOwner() == this->getOperation());
+      Type t = opOperand->get().getType();
+      return cast<VectorType>(t).getShape();
+    }
+
+    AffineMap getLoopsToShapesMap() {
+      auto maps = getIndexingMapsArray();
+      return concatAffineMaps(maps);
+    }
+
+    AffineMap getShapesToLoopsMap() {
+      return inversePermutation(getLoopsToShapesMap());
+    }
+
+    SmallVector<int64_t> getStaticShape(){
+      SmallVector<int64_t> res;
+      for (OpOperand &opOperand : this->getOperation()->getOpOperands())
+        llvm::append_range(res, getShape(&opOperand));
+      return res;
+    }
+
+    SmallVector<int64_t> getStaticLoopRanges() {
+      SmallVector<int64_t> viewSizes = getStaticShape();
+      AffineMap invertedMap = getShapesToLoopsMap();
+      return invertedMap.compose(viewSizes);
+    }
   }];
 
   let hasCanonicalizer = 1;

[fabianmcg]

LGTM, but add comments pointing each member to the corresponding LinalgStructuredInterface method, otherwise the methods appear undocumented and it's unclear what they do without looking at the header.

NIT, fix description.

[Groverkss]

I think a better implementation would to add these methods to IndexingMapOpInterface and make vector.contract implement it.

[fabianmcg]

I think a better implementation would to add these methods to IndexingMapOpInterface and make vector.contract implement it.

I agree, but that's a can of worms in RFCs, as we don't have a structured directory to put the interface in (we should have one).

[Groverkss]

I think a better implementation would to add these methods to IndexingMapOpInterface and make vector.contract implement it.

I agree, but that's a can of worms in RFCs, as we don't have a structured directory to put the interface in (we should have one).

I don't think it is. The Interface lives in lib/Interfaces right?

[fabianmcg]

I don't think it is. The Interface lives in lib/Interfaces right?

Nope: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L225

[banach-space]

I'm OK with this change in principle, but we should either add dedicated tests or ensure there are in-tree uses of this interface. Otherwise, it becomes "dead code," and there's nothing stopping it from being removed in the future by someone upstream.

[nicolasvasilache]

@rengolin @Groverkss @banach-space I have a WIP followup to put stuff in lib/Interface but I think I am limited atm.

TL;DR:

1. getShape should not be the same implementation for Linalg and Vector (i.e. Linalg<memref> is a useful thing to have and thus vector would be an "elemental type" for Linalg but a "ShapedType" for vector)
2. the other methods to factor out in IndexingMapOpInterface want to make use of getShape, so it needs a default impl to be seen by the IndexingMapOpInterface
3. it seems there is a limitation on interface inheritance, where an interface cannot override, atm I have the following case:
   a. IndexingMapOpInterface -> Vector::ContractOp seems fine
   b. IndexingMapOpInterface -> LinalgStructuredInterface -> Linalg::XXXOp does not seem to be fine
4. I don't want to have to individually override in each Linalg::XXXOp

I can put up a followup PR and we can discuss there

[nicolasvasilache]

I'm OK with this change in principle, but we should either add dedicated tests or ensure there are in-tree uses of this interface. Otherwise, it becomes "dead code," and there's nothing stopping it from being removed in the future by someone upstream.

good point !

[nicolasvasilache]

#145332 if someone sees a way to do this better

[Groverkss]

#145332 if someone sees a way to do this better

I like that PR much more. I think we should go with that instead.

[Groverkss]

LGTM! This is really nice and will be good for other operations as well.

I think there is an issue with dps interface implementation and indexingmapopinterface being in opposite files, but rest lgtm.

Also, would be good to change the pr title.

[nicolasvasilache]

#145332 if someone sees a way to do this better

I like that PR much more. I think we should go with that instead.

@ftynse helped me figure out what was wrong, it was an instance of the infamous /*methodBody=*/"" vs /*defaultImplementation=*/[{}] case ...