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

[Chuanqi Xu]

================
@@ -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");
+        }
----------------
ChuanqiXu9 wrote:

Done by changing partitionSizes to `SmallVector<SmallVector<int64_t>>` and then we check `if (!llvm::all_equal(partitionSizes))` 

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