Refactor unfusing and remove unrelated changes

Author: Ferdinand Lemaire

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

@@ -4096,39 +4096,38 @@ structured_op: !LinalgStructuredOpConfig
     name: I
     kind: input_tensor
     type_var: T1
-    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2
-      * s3 + s4 * s5, s6 * s7 + s8 * s9)>
+    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s9, s1 *
+     s2 + s3 * s4, s5 * s6 + s7 * s8)>
   - !LinalgOperandDefConfig
     name: K
     kind: input_tensor
     type_var: T2
-    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s1, s4, s8)>
+    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s9, s3, s7)>
   - !LinalgOperandDefConfig
     name: O
     kind: output_tensor
     type_var: U
-    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s1, s2,
-      s6)>
+    shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s0, s9, s1, s5)>
   - !LinalgOperandDefConfig
     name: strides
     kind: index_attr
-    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s3,
-      s7)>
+    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s2,
+      s6)>
     default_indices:
     - 1
     - 1
   - !LinalgOperandDefConfig
     name: dilations
     kind: index_attr
-    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s5,
-      s9)>
+    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4,
+      s8)>
     default_indices:
     - 1
     - 1
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d3, d1 * s3 + d4 * s5, d2 * s7 + d5 * s9)>
+      -> (d0, d3, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
       -> (d3, d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]

mlir/lib/Dialect/Linalg/Transforms/Unfuse.cpp

@@ -651,12 +651,9 @@ struct GlobalAveragePool2DLowering : OpRewritePattern<GlobalAveragePool2DOp> {
   }
 };
 
-/// Torch MLIR does a similar lowering for their Linear operator to lin alg
-/// here we implement the same so we can run tests using the unfused version
-struct LinearLowering : OpRewritePattern<LinearOp> {
-  using OpRewritePattern<LinearOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(LinearOp op,
-                                PatternRewriter &rewriter) const override {
+template <class Linear>
+static Value unfuseLinear(Linear &op, PatternRewriter &rewriter) {
+    
     Location loc = op.getLoc();
     Value input = op.getOperand(0);
     Value weights = op.getOperand(1);
@@ -691,48 +688,29 @@ struct LinearLowering : OpRewritePattern<LinearOp> {
             ->getResult(0);
 
     // Create the matmul operation that does the multiplcation and addition
-    rewriter.replaceOpWithNewOp<MatmulOp>(op, output.getType(),
-                                          ValueRange{input, transposeWeightsOp},
-                                          broadcastBiasOp);
-
+    auto newOp = rewriter.create<MatmulOp>(loc, outputType, ValueRange{op.getOperand(0), transposeWeightsOp},
+                  broadcastBiasOp).getResult(0);
+    return newOp;
+}
+/// Torch MLIR does a similar lowering for their Linear operator to lin alg
+/// here we implement the same so we can run tests using the unfused version
+struct LinearLowering : OpRewritePattern<LinearOp> {
+  using OpRewritePattern<LinearOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(LinearOp op,
+                                PatternRewriter &rewriter) const override {
+    Value matmul = unfuseLinear<LinearOp>(op, rewriter);
+    rewriter.replaceOp(op, matmul);
     return success();
   }
 };
 
+
 struct LinearReluLowering : OpRewritePattern<LinearReluOp> {
   using OpRewritePattern<LinearReluOp>::OpRewritePattern;
   LogicalResult matchAndRewrite(LinearReluOp op,
                                 PatternRewriter &rewriter) const override {
-    Location loc = op.getLoc();
-    Value weights = op.getOperand(1);
-    Value bias = op.getOperand(2);
-
-    auto weightsType = weights.getType().cast<RankedTensorType>();
-    auto biasType = bias.getType().cast<RankedTensorType>();
-    auto outputType = op->getResult(0).getType().cast<RankedTensorType>();
-
-    // Create a linalg op that transposes the weights tensor
-    // The transposedWeights is simply used to describe the output shape.
-    llvm::ArrayRef<int64_t> weightsShape = weightsType.getShape();
-    Value transposedWeights = rewriter.create<tensor::EmptyOp>(
-        loc,
-        ArrayRef<int64_t>{weightsShape[1], weightsShape[0]},
-        weightsType.getElementType());
-    Value transposeWeightsOp =
-        rewriter.create<Transpose2DOp>(loc, weights, transposedWeights)
-            ->getResult(0);
-
-    // Create a linalg op that broadcasts the 1D bias values across
-    // the 2nd dimension
-    Value broadcastedBias = rewriter.create<tensor::EmptyOp>(
-        loc, outputType.getShape(), biasType.getElementType());
-    Value broadcastBiasOp =
-        rewriter.create<Broadcast1DTo2DOp>(loc, bias, broadcastedBias)
-            ->getResult(0);
 
-    auto linearResult = rewriter.create<MatmulOp>(loc,
-                  outputType, ValueRange{op.getOperand(0), transposeWeightsOp},
-                  broadcastBiasOp).getResult(0);
+    Value linearResult = unfuseLinear<LinearReluOp>(op, rewriter);
 
     rewriter.replaceOpWithNewOp<Relu2DNchwOp>(
         op,