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

[Renato Golin]

================
@@ -502,6 +503,307 @@ 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());
+    std::optional<unsigned> splitLoopDim;
+    // How many inputs the concat ops have. 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.
+    unsigned numPartitions = 0;
+    // 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 numPartitions in this case is 2 and partitionSizes is [3, 4].
+    // Same as above, we limit the partitionSizes to be the same for different
+    // concat ops.
+    SmallVector<int64_t> partitionSizes;
+
+    for (auto [index, operand] :
+         llvm::enumerate(genericOp.getDpsInputOperands())) {
+      auto operandType = dyn_cast<RankedTensorType>(operand->get().getType());
+      if (!operandType)
+        continue;
+
+      auto concatOp = operand->get().getDefiningOp<tensor::ConcatOp>();
+      if (!concatOp)
+        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;
+
+      if (!numPartitions) {
+        numPartitions = 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");
+          partitionSizes.push_back(size);
+        }
+      } else {
+        if (concatOp.getInputs().size() != numPartitions)
+          return rewriter.notifyMatchFailure(
+              genericOp, "concat inputs have different partition counts");
+        for (auto [input, expectedSize] :
+             llvm::zip_equal(concatOp.getInputs(), partitionSizes)) {
+          int64_t size = cast<RankedTensorType>(input.getType())
+                             .getDimSize(concatOp.getDim());
+          if (size != expectedSize)
+            return rewriter.notifyMatchFailure(
+                genericOp, "concat inputs have different partition sizes");
+        }
+      }
+      concatOps[index] = concatOp;
+    }
+
+    if (!splitLoopDim)
+      return rewriter.notifyMatchFailure(genericOp, "has no concat input");
+
+    // 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 =
----------------
rengolin wrote:

usually we build a list of work items that have been verified and don't need extra checks. Here you have two loops over the same operands in the same way with slightly different checks, which is confusing.

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