[MLIR][LINALG] Add canonicalization pattern in `linalg.generic` op for static shape inference.

[gprateek93]

Summary

This commit adds canonicalization pattern in linalg.generic op
for static shape inference. If any of the inputs or outputs have
static shape or is casted from a tensor of static shape, then
shapes of all the inputs and outputs can be inferred by using the
affine map of the static shape input/output.

Signed-Off-By: Prateek Gupta <[e-mail]>

[mravishankar]

Thanks for helping push on this. I don think we can simplify the implementation here by a lot.

@nicolasvasilache might have better ways of doing this for arbitrary affine exprs in indexing maps. In general you need to solve a set of linear equations to compute static dimensions. There might be a way to formulate that in affine map "arithmetic". That seems very heavy wieght (and I havent worked out how to express than using affine map computations). I am thinking of solving a subset of a common case when affine maps are all projected permutations.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

912

General note. I prefer not to use lambdas this way. Its much more readable to have a static method outside of the function.

916

Do you need to pass a copy of SmallVector<OpOperand *> here? Can you use an ArrayRef?

918

If you are going to use a lambda, I would put the loop outside of the lambda and simplify the signature of the lambda.

924

I'd refactor this a bit. First get the shape value

• If operand is result of a tensor.cast get the shape from source of the tensor.cast. If not use the shape of the operand.
• Then iterate over the shape you get in step 1 and check if individual dimensions are static or dynamic. If static, then add it to sourceMap.

929

It is not necessary that the result be an AffineDimExpr. You need to check that the it is and then bail if it is not.

939

I think the else part can be dropped after the above suggestion.

966

Same comment here as above. Move loop out of the lambda.

[gprateek93]

Refactored the code.

[gprateek93]

Added check for the operands to be RankedTensorType.

[mravishankar]

Thanks for the updates. I have one more round of comments. I think after that this should be ready to land.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

904

The startIdx part is confusing to me in the way this is written. There seems to be some handshake between source and operands here. Maybe this could work,

static void populateMap(GenericOp genericOp, ArrayRef<OpOperand *> operands,
                        ValueRange source,
                        llvm::DenseMap<AffineExpr, int64_t> &affineExprToSize) {
  for (auto it : llvm::zip(operands, source)) {
    OpOperand *opOperand = std::get<0>(it);
    Value src = std::get<1>(it);
    ...
  }
  ...
}

In any case, looking at the uses, do you even need to use source. Looks like operands is all you need.

907

Could you add an assert here to check sourceMap.isProjectedPermutation() ? I am assuming the caller verifies that the map is a projected permutation and doesnt call this function if that is the case.

920

nit: LLVM style is to not have { ..} for 1-line scopes.

927

Please do not use -1. Use ShapedType::isDynamicDim instead.

944

Same as above, you probably dont need source and startIdx. Also please move affineExprToSize "before" the parameter that is the return by reference (resultTypeVector I guess). I think there is a simpler way to do this. Ill leave some comments below where this function is called.

966

Please dont use -1 explicitly. Use ShapedType::isDynamicDim instead.

1007

You can replace lines 1007 - 1024 with

if (!genericOp.hasTensorSemantics())
  return failure();

1008

You also need to check that all maps of the genericOp are projected permutations. Something like

if (llvm::any_of(genericOp.getIndexingMaps(), [](AffineMap map) { return !map.isProjectedPermutation(); }))
  return failure();

1032

I think both of these calls can be combined into a single call to populateMap.

1042

I think this split of updateOperands is unnecessary and confusing.

For example for the inputs newResultTypes is unnecessary. We can have separation of concerns and do it in two steps.

SmallVector<Value> newOperands;
SmallVector<Type> resultTypes;
bool noChangeNeeded = true;
newOperands.reserve(genericOp.getNumInputsAndOutputs());
resultTypes.reserve(genericOp.getNumOutputs());
for (OpOperand *operand : genericOp.getInputAndOutputOpOperands()) {
  // Check for type of operand changing and insert casts if needed
  noChangeNeeded = ... // set to true if a change is needed.
  newOperands.push_back(newOperand);
  if (genericOp->isOutputTensor(arg)) {
    resultType.push_back(newOperand.getType());
  }
}
if (noChangeNeeded) return failure();
auto *newOp = ... // Create new generic op.
SmallVector<Value> replacements;
replacements.reserve(newOp->getNumResults());
for (auto it : llvm::zip(genericOp->getResults(), newOp->getResults())) {
   if (std::get<0>(it).getType() != std::get<1>(it).getType()) {
     // Insert cast.
     replacements.push_back(castResult);
   } else {
     replacements.push_back(std::get<1>(it));
   }
}
rewriter.replaceOp(genericOp, replacements);

1054

You can just use cast<linalg::LinalgOp>(genericOp.getOperation()).clone(...) from here https://github.com/llvm/llvm-project/blob/be1cca1b0d7ce01dff2a3cf97438f437576a1fa0/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td#L959

[gprateek93]

Refactored the code.

[mravishankar]

Sorry, one more use of -1 and also a if-then-else style nit. Otherwise looks good!

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

975

Nit: I think LLVM style is to do

for (...) {
  if (...) {
     ...
     continue;
  }

  // previously else-if part.
  if (...) {
    ....
    continue;
  }

  // previously else part
  ...
}

985

Please dont use the -1 directly. Used ShapedType::isDynamicSize or equivalent. In any case, I dont think you need this check here.

[gprateek93]

Refactored the code and updated one of the reshape-fusion test case.

[bondhugula]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

960–1013

This whole logic is too long (about 100 lines), with too much nesting, appears messy with generic boolean variables like changeNeeded. A couple of suggestions:

1. document changeNeeded above,
2. refactor it (for eg. the part iterating over operands) into an appropriately named method.

985

Please use is_contained instead of comparing with end().

mlir/test/Dialect/Linalg/canonicalize.mlir

604

Nit: Add a comment on what these are testing -- the function names themselves aren't descriptive enough.

[gprateek93]

Refactored code.

[gprateek93]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

985

I am not sure how is_contained can be used with map iterators. Can you please elaborate? Thanks!

[bondhugula]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

985

I see - find() is fine.

[gprateek93]

Rebasing into main.

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

LogicalResult matchAndRewrite(GenericOp genericOp,
     }
 
     if (returnedArgs.size() != genericOp->getNumResults())
       return failure();
     rewriter.replaceOp(genericOp, returnedArgs);
     return success();
   }
 };
+
+/// Static shapes for the operands can be inferred if any one of the operands
+/// have a static shape. This can be done by referring to the affine dim
+/// expressions for the operand.
+struct InferStaticShapeOfOperands : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    ValueRange inputs = genericOp.inputs();
+    ValueRange outputs = genericOp.outputs();
+    unsigned numInputs = inputs.size();
+
+    // Maps affine dim expressions to the static size of that dimension.
+    llvm::DenseMap<AffineExpr, int64_t> affineExprToSize;
+    Location loc = genericOp.getLoc();
+
+    // For each of the operand in `operands` this function maps the static sizes
+    // of dimensions to their affine dim expressions. `source` is either the
+    // inputs or the output operands. `startIdx` is the start of inputs/outputs
+    // in the operand list of the op.
+    auto populateMap = [&](SmallVector<OpOperand *> operands, ValueRange source,
+                           unsigned startIdx) {
+      for (auto opOperand : operands) {
+        unsigned index = opOperand->getOperandNumber() - startIdx;
+        Value src = source[index];
+        auto sourceType = src.getType().cast<RankedTensorType>();
+        auto sourceMap = genericOp.getTiedIndexingMap(opOperand);
+
+        // If the source type has a static shape, then for all the dimensions,
+        // map the affine dim expression to the known static size.
+        if (sourceType.hasStaticShape()) {
+          ArrayRef<int64_t> sourceShape = sourceType.getShape();
+          for (unsigned i = 0; i < sourceShape.size(); i++) {
+            AffineExpr affineDimExpr = sourceMap.getResult(i);
+            affineExprToSize.try_emplace(affineDimExpr, sourceShape[i]);
+          }
+        } else {
+          // If the source is a result of a `tensor.cast` operation and its
+          // source has a static shape, then map the affine dim expressions to
+          // those sizes.
+          auto parentOp = src.getDefiningOp();
+          if (!parentOp)
+            continue;
+          if (auto castOp = dyn_cast<tensor::CastOp>(parentOp)) {
+            Value castSource = castOp.source();
+            auto castSourceType = castSource.getType().cast<RankedTensorType>();
+            if (castSourceType.hasStaticShape()) {
+              ArrayRef<int64_t> castSourceShape = castSourceType.getShape();
+              for (unsigned i = 0; i < castSourceShape.size(); i++) {
+                AffineExpr affineDimExpr = sourceMap.getResult(i);
+                affineExprToSize.try_emplace(affineDimExpr, castSourceShape[i]);
+              }
+            }
+          }
+        }
+      }
+    };
+
+    // For each of the `operands` that do not have a static shape, assign them
+    // static shape if present in the `affineExprToSize` map. `source` is either
+    // the inputs or the output operands. `startIdx` is the start of
+    // inputs/outputs in the operand list of the op. All the final values with
+    // static shapes is stored in `dest`. If `source` is output, then
+    // `updateResultType` will be true, and the updated result types are stored
+    // in `resulTypeVector`.
+    auto updateOperands = [&](SmallVector<OpOperand *> operands,
+                              ValueRange source, unsigned startIdx,
+                              SmallVector<Value> &dest, bool updateResultType,
+                              SmallVector<Type> &resulTypeVector) {
+      bool result = false;
+      for (auto opOperand : operands) {
+        unsigned index = opOperand->getOperandNumber() - startIdx;
+        Value src = source[index];
+        auto sourceType = src.getType().cast<RankedTensorType>();
+        Type resultType = genericOp->getResult(index).getType();
+        if (!sourceType.hasStaticShape()) {
+          ArrayRef<int64_t> sourceShape = sourceType.getShape();
+          AffineMap sourceMap = genericOp.getTiedIndexingMap(opOperand);
+          SmallVector<int64_t> newShape;
+
+          // If operand is updated with new shape, `newOperandNeeded` will be
+          // true.
+          bool newOperandNeeded = false;
+          for (unsigned i = 0; i < sourceShape.size(); i++) {
+            int64_t dimShape = sourceShape[i];
+            AffineExpr dimExpr = sourceMap.getResult(i);
+            if (dimShape == -1 &&
+                affineExprToSize.find(dimExpr) != affineExprToSize.end()) {
+              // Dimension has a dynamic shape and corresponding affine dim
+              // expression is present in the map. So assign the size for the
+              // given affine dim expression to the dimension.
+              newShape.push_back(affineExprToSize[dimExpr]);
+              newOperandNeeded = true;
+            } else {
+              newShape.push_back(dimShape);
+            }
+          }
+          if (newOperandNeeded) {
+            resultType =
+                RankedTensorType::get(newShape, sourceType.getElementType());
+
+            // Get the new operand value given its size and element type by
+            // casting it.
+            Value newOperand =
+                rewriter.create<tensor::CastOp>(loc, resultType, src);
+            dest[index] = newOperand;
+            result = true;
+          } else {
+            dest[index] = src;
+          }
+        } else {
+          dest[index] = src;
+        }
+        if (updateResultType)
+          resulTypeVector[index] = resultType;
+      }
+      return result;
+    };
+
+    // For each of the affine dim expression, check if the size is known. If
+    // known add that in the map.
+    populateMap(genericOp.getInputOperands(), inputs, 0);
+    populateMap(genericOp.getOutputOperands(), outputs, numInputs);
+
+    // Iterate over all the operands and update the static sizes.
+    SmallVector<Value> newInputs(numInputs);
+    unsigned numOutputs = outputs.size();
+    SmallVector<Value> newOutputs(numOutputs);
+    SmallVector<Type> newResultTypes(numOutputs);
+    bool updateInputs =
+        updateOperands(genericOp.getInputOperands(), inputs, /*startIdx=*/0,
+                       newInputs, /*updateResultType=*/false, newResultTypes);
+    bool updateOutputs =
+        updateOperands(genericOp.getOutputOperands(), outputs, numInputs,
+                       newOutputs, /*updateResultType=*/true, newResultTypes);
+    if (!updateInputs && !updateOutputs)
+      return failure();
+
+    // Create new op with the `newInputs`, `newOutputs` and `newResultTypes`.
+    auto newOp = rewriter.create<GenericOp>(
+        loc, newResultTypes, newInputs, newOutputs, genericOp.indexing_maps(),
+        genericOp.iterator_types(), genericOp.docAttr(),
+        genericOp.library_callAttr());
+
+    // Copy the payload as it is from `genericOp`.
+    rewriter.inlineRegionBefore(genericOp.region(), newOp.region(),
+                                newOp.region().begin());
+
+    rewriter.replaceOp(genericOp, newOp->getResults());
+    return success();
+  }
+};
 } // namespace
 
 void GenericOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp>(context);
+  results.add<DeduplicateGenericOpInputs, EraseIdentityGenericOp,
+              InferStaticShapeOfOperands>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // InitTensorOp
 //===----------------------------------------------------------------------===//
 
 void InitTensorOp::build(OpBuilder &b, OperationState &result,
                          ArrayRef<OpFoldResult> sizes, Type elementType,
                          ArrayRef<NamedAttribute> attrs) {

mlir/test/Dialect/Linalg/canonicalize.mlir

linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                   iterator_types = ["parallel", "parallel"]} 
     ins(%0 : memref<4x16xf32>) 
     outs(%0 : memref<4x16xf32>) {
       ^bb0(%arg4: f32, %arg5: f32):
         linalg.yield %arg4 : f32
     }
   return 
 }
+
+// The below test cases check the inference of static shapes in linalg.generic operation.
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_input_without_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+func @static_input_without_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %arg1 : tensor<2x3x4xf32>, tensor<?x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %5 = tensor.cast %4 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %5 : tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_input_with_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+func @static_input_with_cast(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<?x?x?xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
+  %5 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %4 : tensor<2x3x4xf32>, tensor<2x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %6 = tensor.cast %5 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %6: tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @static_output_with_cast
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<?x?x?xf32>, %[[ARG1:.*]]: tensor<?x?x?xf32>, %[[ARG2:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+func @static_output_with_cast(%arg0 : tensor<?x?x?xf32>, %arg1: tensor<?x?x?xf32>, %arg2: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg2, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg2, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg2, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %3 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  %5 = tensor.cast %arg1 : tensor<?x?x?xf32> to tensor<2x?x?xf32>
+  %6 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %5 : tensor<?x?x?xf32>, tensor<2x?x?xf32>)
+    outs(%4 : tensor<2x3x4xf32>) {
+  ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
+    %9 = arith.addf %arg3, %arg4 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<2x3x4xf32>)
+  return %6: tensor<2x3x4xf32>
+    //  CHECK:      %[[CAST_ARG0:.*]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[CAST_ARG1:.*]] = tensor.cast %[[ARG1]] : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+    //  CHECK-NEXT: %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[CAST_ARG0]], %[[CAST_ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}
+
+// -----
+
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-LABEL: func @cast_source
+// CHECK-SAME:  (%[[ARG0:.*]]: tensor<2x3x4xf32>, %[[ARG1:.*]]: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+func @cast_source(%arg0 : tensor<2x3x4xf32>, %arg1: tensor<2x3x4xf32>) -> tensor<2x3x4xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = tensor.dim %arg0, %c0 : tensor<2x3x4xf32>
+  %1 = tensor.dim %arg0, %c1 : tensor<2x3x4xf32>
+  %2 = tensor.dim %arg0, %c2 : tensor<2x3x4xf32>
+  %3 = linalg.init_tensor [%0, %1, %2] : tensor<?x?x?xf32>
+  %4 = tensor.cast %arg0 : tensor<2x3x4xf32> to tensor<?x?x?xf32> 
+  %5 = tensor.cast %arg1 : tensor<2x3x4xf32> to tensor<?x?x?xf32>
+  %6 = linalg.generic {
+    indexing_maps = [#map, #map, #map],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%4, %5 : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
+    outs(%3 : tensor<?x?x?xf32>) {
+  ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32):
+    %9 = arith.addf %arg2, %arg3 : f32
+    linalg.yield %9 : f32
+  } -> (tensor<?x?x?xf32>)
+  %7 = tensor.cast %6 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+  return %7: tensor<2x3x4xf32>
+    //  CHECK:      %[[GENERIC_OP:.*]] = linalg.generic
+    //  CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<2x3x4xf32>, tensor<2x3x4xf32>)
+    //  CHECK-SAME: outs({{.*}} : tensor<2x3x4xf32>)
+}