[Mlir-commits] [mlir] [mlir][linalg] Split elementwise ops with concat inputs (PR #213630)

[Renato Golin]

================
@@ -502,6 +503,292 @@ class FuseElementwiseOps : public OpRewritePattern<GenericOp> {
 private:
   ControlFusionFn controlFn;
 };
+
+/// Split an elementwise operation at the boundaries of its `tensor.concat`
+/// inputs. This exposes the producers of the concat inputs to the elementwise
+/// fusion patterns.
+///
+///   elementwise(concat(x0, x1), concat(y0, y1))
+///
+/// becomes
+///
+///   concat(elementwise(x0, y0), elementwise(x1, y1))
+///
+/// This pattern is intentionally expressed on `linalg.generic`: tensor
+/// elementwise operations such as `arith.addf` are converted to that form by
+/// `-convert-elementwise-to-linalg`, before this pattern runs as a preamble to
+/// Linalg elementwise fusion.
+///
+/// A partition is one input of a concat, viewed as a contiguous interval of
+/// the concat dimension. All concat inputs must have matching partitions: the
+/// same number of partitions with the same static size at each index.
+///
+/// All concat operands must partition the same iteration-space dimension into
+/// the same statically-sized pieces. Inputs that do not use that iteration
+/// dimension (for example, broadcast inputs) can be shared by all pieces.
+class SplitElementwiseOpWithConcatInputs : public OpRewritePattern<GenericOp> {
+public:
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp genericOp,
+                                PatternRewriter &rewriter) const override {
+    if (!genericOp.hasPureTensorSemantics() || !isElementwise(genericOp) ||
+        genericOp.hasIndexSemantics())
+      return failure();
+
+    SmallVector<tensor::ConcatOp> concatOps(genericOp.getNumDpsInputs());
+    SmallVector<OpOperand *> nonConcatInputs;
+    std::optional<unsigned> splitLoopDim;
+    // Now we limit the concat ops to have the same number of inputs for
+    // simplicity.
+    // TODO: technically, elementwise(concat(x0, x1), concat(y0, y1, y2)) ->
+    // concat(elementwise(...), elementwise(...), elementwise(...)) may be
+    // fine too. But that may require we create new slices, which might be
+    // more complex.
+    // The size in the concat dimension of different inputs. For example,
+    //
+    //  x0: tensor<2x3xf32>
+    //  x1: tensor<2x4xf32>
+    //  x:  tensor<2x7xf32>
+    //  %x = tensor.concat dim(1) %x0, %x1
+    //
+    // The partition sizes in this case are [3, 4]. Same as above, we limit the
+    // partition sizes to be the same for different concat ops.
+    SmallVector<SmallVector<int64_t>> partitionSizes;
+
+    for (auto [index, operand] :
+         llvm::enumerate(genericOp.getDpsInputOperands())) {
+      auto concatOp = operand->get().getDefiningOp<tensor::ConcatOp>();
+      if (!concatOp) {
+        nonConcatInputs.push_back(operand);
+        continue;
+      }
+
+      // Rewriting a concat that has other consumers could increase the amount
+      // of live computation instead of just exposing fusion opportunities.
+      if (llvm::any_of(concatOp.getResult().getUses(), [&](OpOperand &use) {
+            return use.getOwner() != genericOp.getOperation();
+          }))
+        return rewriter.notifyMatchFailure(genericOp,
+                                           "concat input has another consumer");
+
+      AffineMap inputMap = genericOp.getMatchingIndexingMap(operand);
+      auto concatDimExpr =
+          dyn_cast<AffineDimExpr>(inputMap.getResult(concatOp.getDim()));
+      if (!concatDimExpr)
+        return rewriter.notifyMatchFailure(
+            genericOp, "concat dimension does not map to a loop dimension");
+
+      unsigned currentSplitLoopDim = concatDimExpr.getPosition();
+      if (splitLoopDim && *splitLoopDim != currentSplitLoopDim)
+        return rewriter.notifyMatchFailure(
+            genericOp, "concat inputs partition different loop dimensions");
+      splitLoopDim = currentSplitLoopDim;
+
+      SmallVector<int64_t> currentPartitionSizes;
+      currentPartitionSizes.reserve(concatOp.getInputs().size());
+      for (Value input : concatOp.getInputs()) {
+        int64_t size = cast<RankedTensorType>(input.getType())
+                           .getDimSize(concatOp.getDim());
+        if (ShapedType::isDynamic(size))
+          return rewriter.notifyMatchFailure(
+              genericOp, "concat partition size is dynamic");
+        currentPartitionSizes.push_back(size);
+      }
+      partitionSizes.push_back(std::move(currentPartitionSizes));
+      concatOps[index] = concatOp;
+    }
+
+    if (!splitLoopDim)
+      return rewriter.notifyMatchFailure(genericOp, "has no concat input");
+    if (!llvm::all_equal(partitionSizes))
+      return rewriter.notifyMatchFailure(
+          genericOp, "concat inputs have different partition sizes");
+
+    // A tensor input that varies along the split dimension must itself be a
+    // compatible concat. Inputs that are invariant along that dimension can be
+    // reused by every split operation.
+    AffineExpr splitDimExpr =
+        getAffineDimExpr(*splitLoopDim, genericOp.getContext());
+    for (OpOperand *operand : nonConcatInputs) {
+      Type operandType = operand->get().getType();
+      // Scalars do not vary along an iteration-space dimension and can be
+      // reused in every partition. Other shaped inputs must be ranked tensors
+      // so that their indexing map can be checked below.
+      if (!isa<RankedTensorType>(operandType)) {
+        if (!isa<ShapedType>(operandType))
----------------
rengolin wrote:

This check doesn't need to be inside the one above. Then split the comment above to match each `if` check.

Also, not sure `not ShapedType === ScalarType`. Token types, dialect types and other things can break this assumption.

https://github.com/llvm/llvm-project/pull/213630