Further refactor

Signed-off-by: Nirvedh Meshram <[email protected]>

Author: nirvedhmeshram

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

@@ -814,26 +814,64 @@ validateDynamicDimExpansion(LinalgOp linalgOp,
   }
   return success();
 }
+
+// Create an expanded transpose op.
+static Operation *
+createExpandedTransposeOp(PatternRewriter &rewriter, TransposeOp transposeOp,
+                          SmallVector<ReassociationIndices> reassociation,
+                          Value expandedInput, Value output) {
+  applyPermutationToVector(reassociation, transposeOp.getPermutation());
+  SmallVector<int64_t> newPerm;
+  for (auto reassoc : reassociation) {
+    for (auto dim : reassoc) {
+      newPerm.push_back(dim);
+    }
+  }
+  return rewriter.create<TransposeOp>(transposeOp.getLoc(), expandedInput,
+                                      output, newPerm);
+}
+
+// Create an expanded generic op.
+static Operation *createExpandedGenericOp(
+    PatternRewriter &rewriter, LinalgOp linalgOp, TypeRange resultTypes,
+    ArrayRef<Value> &expandedOpOperands, ArrayRef<Value> outputs,
+    ExpansionInfo &expansionInfo, ArrayRef<AffineMap> expandedOpIndexingMaps) {
+  // The iterator types of the expanded op are all parallel.
+  SmallVector<utils::IteratorType> iteratorTypes(
+      expansionInfo.getExpandedOpNumDims(), utils::IteratorType::parallel);
+
+  for (auto [i, type] : llvm::enumerate(linalgOp.getIteratorTypesArray()))
+    for (auto j : expansionInfo.getExpandedDims(i))
+      iteratorTypes[j] = type;
+
+  Operation *fused = rewriter.create<GenericOp>(
+      linalgOp.getLoc(), resultTypes, expandedOpOperands, outputs,
+      expandedOpIndexingMaps, iteratorTypes);
+
+  Region &fusedRegion = fused->getRegion(0);
+  Region &originalRegion = linalgOp->getRegion(0);
+  rewriter.cloneRegionBefore(originalRegion, fusedRegion, fusedRegion.begin());
+
+  // Update the index accesses after the expansion.
+  updateExpandedGenericOpRegion(rewriter, linalgOp.getLoc(), fusedRegion,
+                                expansionInfo);
+
+  return fused;
+}
+
 // Create an expanded fused op that retains the name for certain ops
 // such as fill, copy and transpose and produce a generic op for
 // rest of linalg ops.
-Operation *createFusedOpForReshapeByExpansion(
+static Operation *createExpandedOp(
     PatternRewriter &rewriter, LinalgOp linalgOp, TypeRange resultTypes,
     ArrayRef<Value> expandedOpOperands, ArrayRef<Value> outputs,
     ArrayRef<AffineMap> expandedOpIndexingMaps, ExpansionInfo &expansionInfo,
     SmallVector<ReassociationIndices> reassociation) {
 
   return TypeSwitch<Operation *, Operation *>(linalgOp.getOperation())
-      .Case<TransposeOp>([&](TransposeOp op) {
-        applyPermutationToVector(reassociation, op.getPermutation());
-        SmallVector<int64_t> newPerm;
-        for (auto reassoc : reassociation) {
-          for (auto dim : reassoc) {
-            newPerm.push_back(dim);
-          }
-        }
-        return rewriter.create<TransposeOp>(
-            linalgOp.getLoc(), expandedOpOperands[0], outputs[0], newPerm);
+      .Case<TransposeOp>([&](TransposeOp transposeOp) {
+        return createExpandedTransposeOp(rewriter, transposeOp, reassociation,
+                                         expandedOpOperands[0], outputs[0]);
       })
       .Case<FillOp, CopyOp>([&](Operation *op) {
         return clone(rewriter, linalgOp, resultTypes,
@@ -842,25 +880,9 @@ Operation *createFusedOpForReshapeByExpansion(
                          llvm::to_vector(outputs))));
       })
       .Default([&](Operation *op) {
-        // The iterator types of the expanded op are all parallel.
-        SmallVector<utils::IteratorType> iteratorTypes(
-            expansionInfo.getExpandedOpNumDims(),
-            utils::IteratorType::parallel);
-        for (auto [i, type] : llvm::enumerate(linalgOp.getIteratorTypesArray()))
-          for (auto j : expansionInfo.getExpandedDims(i))
-            iteratorTypes[j] = type;
-        Operation *fused = rewriter.create<GenericOp>(
-            linalgOp.getLoc(), resultTypes, expandedOpOperands, outputs,
-            expandedOpIndexingMaps, iteratorTypes);
-        Region &fusedRegion = fused->getRegion(0);
-        Region &originalRegion = linalgOp->getRegion(0);
-        rewriter.cloneRegionBefore(originalRegion, fusedRegion,
-                                   fusedRegion.begin());
-
-        // Update the index accesses after the expansion.
-        updateExpandedGenericOpRegion(rewriter, linalgOp.getLoc(), fusedRegion,
-                                      expansionInfo);
-        return fused;
+        return createExpandedGenericOp(rewriter, linalgOp, resultTypes,
+                                       expandedOpOperands, outputs,
+                                       expansionInfo, expandedOpIndexingMaps);
       });
 }
 
@@ -972,7 +994,7 @@ fuseWithReshapeByExpansion(LinalgOp linalgOp, Operation *reshapeOp,
   SmallVector<ReassociationIndices> reassociationBeforeExpansion =
       isExpanding ? expandingReshapeOp.getReassociationIndices()
                   : collapsingReshapeOp.getReassociationIndices();
-  Operation *fusedOp = createFusedOpForReshapeByExpansion(
+  Operation *fusedOp = createExpandedOp(
       rewriter, linalgOp, resultTypes, expandedOpOperands, outputs,
       expandedOpIndexingMaps, expansionInfo, reassociationBeforeExpansion);
   // Reshape the result values to their original shape if this is a collapsing