[mlir][Linalg] Add utility function that return static loop bounds of Linalg ops

[asaadaldien]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

120

nit: Maybe a comment to explain what is the shape of an op ?

[hanchung]

few nits, LGTM!

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

162

nit: I think we don't need template because this is not a template function?

162

If it is getStaticShape, do we want to check if shape.hasStaticShape()?

[mravishankar]

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

120

It says its the static version of the getShape method above. That explains what shape is.

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

162

The comment in the header files says it returns -1 for the dynamic shapes. Might worth returning that in case clients want to check which dim is dynamic.

[mravishankar]

Address comments

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 /// Applying the inverse, concatenated loopToOperandRangeMaps to this list
 /// allows the derivation of loop ranges for any linalgOp.
 SmallVector<Value, 8> getShape(OpBuilder &builder, LinalgOp linalgOp);
 template <typename ConcreteOpTy>
 SmallVector<Value, 8> getShape(OpBuilder &builder, ConcreteOpTy linalgOp) {
   return getShape(builder, cast<linalg::LinalgOp>(linalgOp.getOperation()));
 }
 
+/// Like `getShape`, but only returns statically-known information, without
+/// generating any new IR. For each shape dimension, returns >=0 if that
+/// dimension is statically known, or -1 otherwise.
+SmallVector<int64_t, 8> getStaticShape(LinalgOp linalgOp);
+
 /// Returns the loop ranges of the `linalgOp`. Applies the inverse of the
 /// concatenated indexing maps to the result of `getShape`. Returns None if
 /// the bounds computation fails.
 Optional<SmallVector<Value, 4>> getLoopRanges(OpBuilder &builder,
                                               LinalgOp linalgOp);
 
+/// Returns the statically-known loop ranges of the `linalgOp`. Applies the
+/// inverse of the concatenated indexing maps to the result of `getStaticShape`.
+/// Returns None if inverting the concatenated indexing map fails. Returns -1
+/// for non-statically-known loop ranges.
+Optional<SmallVector<int64_t, 4>> getStaticLoopRanges(LinalgOp linalgOp);
+
 /// Returns the values obtained by applying `map` to the list of values.
 SmallVector<Value, 4> applyMapToValues(OpBuilder &b, Location loc,
                                        AffineMap map, ValueRange values);
 
 /// Apply the permutation defined by `permutation` to `inVec`.
 /// Element `i` in `inVec` is mapped to location `j = permutation[i]`.
 /// E.g.: for an input vector `inVec = ['a', 'b', 'c']` and a permutation vector
 /// `permutation = [2, 0, 1]`, this function leaves `inVec = ['c', 'a', 'b']`.

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

if (attr) {
     res.reserve(res.size() + ranks.size() * numSymb);
     for (unsigned idx = 0, s = ranks.size(); idx < s; ++idx)
       for (unsigned idx2 = 0; idx2 < numSymb; ++idx2)
         res.push_back(res[symbolsPos + idx2]);
   }
   return res;
 }
 
+SmallVector<int64_t, 8> getStaticShape(LinalgOp linalgOp) {
+  SmallVector<int64_t, 8> res;
+  for (Value v : linalgOp.getShapedOperands()) {
+    auto shape = v.getType().cast<ShapedType>().getShape();
+    res.append(shape.begin(), shape.end());
+  }
+  return res;
+}
+
 Optional<SmallVector<Value, 4>> getLoopRanges(OpBuilder &builder,
                                               LinalgOp linalgOp) {
   SmallVector<Value, 8> viewSizes = getShape(builder, linalgOp);
   AffineMap invertedMap =
       inversePermutation(concatAffineMaps(linalgOp.getIndexingMaps()));
   if (!invertedMap)
     return {};
   return applyMapToValues(builder, linalgOp.getLoc(), invertedMap, viewSizes);
 }
 
+Optional<SmallVector<int64_t, 4>> getStaticLoopRanges(LinalgOp linalgOp) {
+  SmallVector<int64_t, 8> viewSizes = getStaticShape(linalgOp);
+  AffineMap invertedMap =
+      inversePermutation(concatAffineMaps(linalgOp.getIndexingMaps()));
+  if (!invertedMap)
+    return {};
+  return invertedMap.compose(viewSizes);
+}
+
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
     ArrayRef<Range> loopRanges, ValueRange iterArgInitValues,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(ValueRange, ValueRange)> bodyBuilderFn,
     Optional<LinalgLoopDistributionOptions> distributionOptions) {
   // Create procInfo so it dominate loops, if appropriate.