Moved getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td to add static shape verification

[• inho9606]

Summary

It is to use the methods in LinalgInterfaces.cpp for additional static shape verification to match the shaped operands and loop on linalgOps. If I used the existing methods, I would face circular dependency linking issue. Now we can use them as methods of LinalgOp.

[ftynse]

Please provide a rationale for making this change in the commit description. Why is it necessary? Is there any op that will customize the behavior and why can't it do so by overloading the functions instead?

[hanchung]

In D98163#2610701, @ftynse wrote:

Please provide a rationale for making this change in the commit description. Why is it necessary? Is there any op that will customize the behavior and why can't it do so by overloading the functions instead?

Inho is going to add static shape verification for Linalg structured ops. E.g., verify if the shapes match the loops for conv op. He is going to use these methods in the change. To reuse methods in
verifyStructuredOpInterface, we have to move methods out of Utils.h. Otherwise, there will be circular dependency. We already have some similar interface methods like getShapesToLoopsMap, etc. I think it makes sense to move static version to interface methods.

Inho, could you provide more context to commit description and fix the title? Please also fix the build.

The below link has an overview of writing a good change descriptions, please take a look: https://google.github.io/eng-practices/review/developer/cl-descriptions.html

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

197–234

I think they are not for "loop types handling". Maybe move them to "Linalg generalization hooks."?

[• inho9606]

1. Updating D98163: #Enter a commit message. # Moved fectored getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment.

Moved getStaticShape and getStaticLoopRanges method to Linalg Generlization Hooks part.n LinalgInterfaces.td
This is to use these methods in LinalgInterfaces.cpp for additional static shape verification to match the shapes and loops. If I used existing methods, circular dependency issue happened. Now we can use the methods with linalgOp.getStaticLoopRanges().

[• inho9606]

In D98163#2610701, @ftynse wrote:

Please provide a rationale for making this change in the commit description. Why is it necessary? Is there any op that will customize the behavior and why can't it do so by overloading the functions instead?

Sorry for confusing commit message. It was my first change request here, so the interface was not familiar for me. I didn't know how to update the revision at the moment, now I updated it.
Thanks

[• inho9606]

Uh,, Sorry I removed @ftynse from subscribers,, but don't know why it happened..

[• inho9606]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

197–234

I moved it, could you look at it again?

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

1129

Let's follow other places -- adding a blank line right before the comment.

1132

Please also update the comment. In the past, we did return -1 which indicated to dynamic dim. Now we have ShapeType::kDynamicSize, so let's use it in the comment instead.

1139–1145

Do we need to get the operation and cast to LinalgOp? I thought getShapedOperands is a class method which can be access directly?

1152

ditto

1160

ditto

[• inho9606]

1. Updating D98163: Moved getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td to add static shape verification #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment.

Reflected request changes

[• inho9606]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

1139–1145

Oh right, I had to check it beforehand,, thanks. And also I think linalgOp is not used in getStaticLoopRanges method, so I removed it.

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

1113

I think this -1 should also be ShapeType::kDynamicSize?

[• inho9606]

1. Updating D98163: Moved getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td to add static shape verification #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment

Changed -1 to ShapedType::kDynamicSize in comment

[hanchung]

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

1113–1114

It is too long, we have to format it a bit.

Returns the statically-known loop ranges. Composes
`getShapesToLoopsMap()` with the result of `getStaticShape`.
Returns None if `getShapesToLoopsMap()` fails. Returns
ShapeType::kDynamicSize for non-statically-known loop ranges.

[• inho9606]

1. Updating D98163: Moved getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td to add static shape verification #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment.

reformating

[• inho9606]

1. Updating D98163: Moved getStaticLoopRanges and getStaticShape methods to LinalgInterfaces.td to add static shape verification #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment.

Reverted the test file

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
       /*methodName=*/"hasSingleReductionLoop",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         auto iters = $_op.iterator_types();
         return iters.size() == 1 &&
                getNumIterators(getReductionIteratorTypeName(), iters) == 1;
       }]>,
     //===------------------------------------------------------------------===//
     // Num input/output arguments handling.
     //===------------------------------------------------------------------===//
     // `inputs` must be defined by each op that wants to implement the
     // LinalgStructuredInterface.
     InterfaceMethod<
       /*desc=*/[{
         Return the input shape operands.
       }],
       /*retTy=*/"ValueRange",
       /*methodName=*/"inputs",
       /*args=*/(ins)
     >,
     // These special methods rely on `inputs` and `outputs` being defined by
     // each op that wants to implement the LinalgStructuredInterface.
     InterfaceMethod<
       /*desc=*/[{
         Return the number of inputs.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumInputs",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.inputs().size();
       }]
     >,
     // `outputs` must be defined by each op that wants to implement the
     // LinalgStructuredInterface.
     InterfaceMethod<
       /*desc=*/[{
         Return the output shape operands.
       }],
       /*retTy=*/"ValueRange",
       /*methodName=*/"outputs",
       /*args=*/(ins)
     >,
     InterfaceMethod<
       /*desc=*/[{
         Return the number of outputs.
       }],
       /*retTy=*/"unsigned",
       /*methodName=*/"getNumOutputs",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
InterfaceMethod<
       /*args=*/(ins "unsigned":$resultIdx, "unsigned":$dim),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         if (resultIdx >= getNumOutputs()) return {};
         return getOperandDimPositionInLoopsToShapeMap(
             getNumInputs() + resultIdx, dim);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        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 ShapeType::kDynamicSize otherwise.
+      }],
+      /*retTy=*/"SmallVector<int64_t, 8>",
+      /*methodName=*/"getStaticShape",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        SmallVector<int64_t, 8> res;
+        for (Value v : getShapedOperands()) {
+          auto shape = v.getType().cast<ShapedType>().getShape();
+          res.append(shape.begin(), shape.end());
+        }
+        return res;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Returns the statically-known loop ranges. Composes
+        `getShapesToLoopsMap()` with the result of `getStaticShape`.
+        Returns None if `getShapesToLoopsMap()` fails. Returns
+        ShapeType::kDynamicSize for non-statically-known loop ranges.
+      }],
+      /*retTy=*/"Optional<SmallVector<int64_t, 4>>",
+      /*methodName=*/"getStaticLoopRanges",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        SmallVector<int64_t, 8> viewSizes = getStaticShape();
+        AffineMap invertedMap = getShapesToLoopsMap();
+        if (!invertedMap)
+          return {};
+        return invertedMap.compose(viewSizes);
+      }]
+    >,
 
     //===------------------------------------------------------------------===//
     // Other static interface methods.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/[{
         Clone the current operation with the given location and operands. This
         is used to abstract away the optional underlying region creation. This
         does not change the balance between input, output_buffer and
         init_tensors operands.
       }],
       /*retTy=*/"Operation *",
       /*methodName=*/"clone",
       (ins "OpBuilder &":$b, "Location":$loc, "TypeRange":$resultTypes,
            "ValueRange":$operands),
       [{
         BlockAndValueMapping bvm;
         OperationState state(
           loc, ConcreteOp::getOperationName(), operands, resultTypes,
           $_op->getAttrs());
         for (Region &r : $_op->getRegions())
           r.cloneInto(state.addRegion(), bvm);
         return b.createOperation(state);
       }]
     >,
     StaticInterfaceMethod<
       /*desc=*/[{
         Returns the region builder for constructing the body for linalg.generic.
         Returns a null function if this named op does not define a region
         builder.
       }],
       /*retTy=*/"std::function<void(Block &, ValueRange)>",
       /*methodName=*/"getRegionBuilder",
       (ins),
       [{ return ConcreteOp::getRegionBuilder(); }]
     >
   ];
 
   let extraClassDeclaration = [{
     /// Return the flat list of all operand dimension sizes in the order they
     /// appear in the operands.

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

Optional<FusionInfo> fuseProducerOfTensor(OpBuilder &b,
                                           OpResult producerOpResult,
                                           OpOperand &consumerOpOperand);
 
 /// Fuse linalg operation on tensors, with the producer of the operand at
 /// position `consumerIdx` of the consumer.
 Optional<SmallVector<Value, 1>> fuseTensorOps(PatternRewriter &rewriter,
                                               OpOperand &consumerOpOperand);
 
-/// 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 statically-known loop ranges of the `linalgOp`. Composes
-/// `linalgOp.getShapesToLoopsMap()` with the result of `getStaticShape`.
-/// Returns None if `linalgOp.getShapesToLoopsMap()` fails. Returns -1
-/// for non-statically-known loop ranges.
-Optional<SmallVector<int64_t, 4>> getStaticLoopRanges(LinalgOp linalgOp);
-
 /// 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']`.
 template <typename T, unsigned N>
 void applyPermutationToVector(SmallVector<T, N> &inVec,
                               ArrayRef<unsigned> permutation) {
   SmallVector<T, N> auxVec(inVec.size());

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

LogicalResult ExpansionInfo::compute(LinalgOp linalgOp,
                                      unsigned fusedTensorIndex,
                                      ArrayRef<AffineMap> reassociationMaps,
                                      ArrayRef<int64_t> expandedShape) {
   if (reassociationMaps.empty())
     return failure();
   AffineMap fusedIndexMap = linalgOp.getIndexingMap(fusedTensorIndex);
 
   Optional<SmallVector<int64_t, 4>> originalLoopRange =
-      getStaticLoopRanges(linalgOp);
+      linalgOp.getStaticLoopRanges();
   if (!originalLoopRange)
     return linalgOp.emitError("unable to find loop range for operation");
 
   reassociation.clear();
   expandedShapeMap.clear();
   // Compute the number of dimension in the expanded op that correspond to each
   // dimension of the original op.
   SmallVector<unsigned, 4> numExpandedDims(fusedIndexMap.getNumDims(), 1);

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

for (Range range : ranges) {
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
 namespace mlir {
 namespace linalg {
 
-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<int64_t, 4>> getStaticLoopRanges(LinalgOp linalgOp) {
-  SmallVector<int64_t, 8> viewSizes = getStaticShape(linalgOp);
-  AffineMap invertedMap = linalgOp.getShapesToLoopsMap();
-  if (!invertedMap)
-    return {};
-  return invertedMap.compose(viewSizes);
-}
-
 /// If `size` comes from an AffineMinOp and one of the values of AffineMinOp
 /// is a constant then return a new value set to the smallest such constant.
 /// Otherwise returngetSmallestBoundingIndex nullptr.
 IntegerAttr getSmallestBoundingIndex(Value size) {
   Optional<int64_t> boundingConst = {};
   if (auto affineMinOp = size.getDefiningOp<AffineMinOp>()) {
     for (auto e : affineMinOp.getAffineMap().getResults())
       if (auto cst = e.dyn_cast<AffineConstantExpr>())

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

 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel"]}
 // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
 // CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>)
 
 // CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
 // CHECK-NEXT:      %[[MUL:.+]] = mulf %[[BBARG0]], %[[BBARG1]] : f32
 // CHECK-NEXT:      %[[ADD:.+]] = addf %[[BBARG2]], %[[MUL]] : f32
 // CHECK-NEXT:      linalg.yield %[[ADD]] : f32
-
-// -----
-
-func @pooling_nhwc_sum(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
-  linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
-    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
-    outs(%init: memref<?x?x?x?xf32>)
-  return
-}
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
-// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
-
-// CHECK: func @pooling_nhwc_sum
-
-// CHECK: linalg.generic
-// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
-// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
-// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
-// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
-
-// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
-// CHECK-NEXT:      %[[RES:.+]] = addf %[[BBARG2]], %[[BBARG0]] : f32
-// CHECK-NEXT:      linalg.yield %[[RES]] : f32
-
-// -----
-
-func @pooling_nhwc_max(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
-  linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
-    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
-    outs(%init: memref<?x?x?x?xf32>)
-  return
-}
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 3 + d5, d3)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
-// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
-
-// CHECK: func @pooling_nhwc_max
-
-// CHECK: linalg.generic
-// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
-// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
-// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
-// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
-
-// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
-// CHECK-NEXT:      %[[CMP:.+]] = cmpf ogt, %[[BBARG0]], %[[BBARG2]] : f32
-// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
-// CHECK-NEXT:      linalg.yield %[[RES]] : f32
-
-// -----
-
-func @pooling_nhwc_min(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
-  linalg.pooling_nhwc_min {dilations = dense<3> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
-    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
-    outs(%init: memref<?x?x?x?xf32>)
-  return
-}
-
-// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4 * 3, d2 * 2 + d5 * 3, d3)>
-// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
-// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
-
-// CHECK: func @pooling_nhwc_min
-
-// CHECK: linalg.generic
-// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
-// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
-// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
-// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
-
-// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
-// CHECK-NEXT:      %[[CMP:.+]] = cmpf olt, %[[BBARG0]], %[[BBARG2]] : f32
-// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
-// CHECK-NEXT:      linalg.yield %[[RES]] : f32