[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) {
----------------
rengolin wrote:

Can't you just call `!hasOneUse()`? `concatOp` is the defining op of an operand that was extracted from this generic op, so if it has one user, it clearly is the generic.

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