fix deepTileMatmul

Author: zhczhong

lib/gc/Transforms/DeepTileContractionNamedOp.cpp

@@ -322,11 +322,10 @@ static unsigned getOprandDim(linalg::LinalgOp &linalgOp, unsigned iteratorPos,
   return linalgOp.getShape(linalgOp.getDpsInputOperand(operandIdx))[dimPos];
 }
 
-static LogicalResult setStaticSizeForExtractSliceOp(RewriterBase &rewriter,
-                                                    Operation *op,
-                                                    bool isExtract,
-                                                    SmallVector<int64_t> size,
-                                                    int shrinDimNum = 0) {
+static void setStaticSizeForExtractSliceOp(RewriterBase &rewriter,
+                                           Operation *op, bool isExtract,
+                                           SmallVector<int64_t> size,
+                                           int shrinDimNum = 0) {
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(op);
   if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(op)) {
@@ -348,15 +347,12 @@ static LogicalResult setStaticSizeForExtractSliceOp(RewriterBase &rewriter,
           extractSlice, extractSlice.getSource(), mixedOffsets, mixedSizes,
           mixedStrides);
     }
-  } else {
-    return failure();
   }
-  return mlir::success();
 }
 
-static LogicalResult setStaticSizeForInsertSliceOp(RewriterBase &rewriter,
-                                                   Operation *op, Value source,
-                                                   SmallVector<int64_t> size) {
+static void setStaticSizeForInsertSliceOp(RewriterBase &rewriter, Operation *op,
+                                          Value source,
+                                          SmallVector<int64_t> size) {
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(op);
   if (auto insertSlice = dyn_cast<tensor::InsertSliceOp>(op)) {
@@ -369,10 +365,7 @@ static LogicalResult setStaticSizeForInsertSliceOp(RewriterBase &rewriter,
     rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
         insertSlice, source, insertSlice.getDest(), mixedOffsets, mixedSizes,
         mixedStrides);
-  } else {
-    return failure();
   }
-  return success();
 }
 
 using InnermostFullResultCallBackFn = std::function<FailureOr<linalg::LinalgOp>(
@@ -691,7 +684,6 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
                                     linalg::LinalgOp originOp,
                                     linalg::LinalgOp currentOp,
                                     innerBodyGenerationOption &option) const {
-
     mlir::easybuild::EasyBuilder eb{rewriter, originOp.getLoc()};
     auto operandDimTypes = getOprandDimType(originOp);
     auto cfg = MatmulConfigAnalysis(originOp.getOperation()).getConfig();
@@ -744,6 +736,7 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
       CInnermostDims =
           SmallVector<int64_t>{cfg.innerMostMBlock, cfg.innerMostNBlock};
     }
+
     if (NDimNum > 1) {
       firstN = true;
       firstK = true;
@@ -780,21 +773,17 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
 
     // update the extractSlice to static size, replace it with
     // useBlockedLayout when
-    if (failed(setStaticSizeForExtractSliceOp(
-            rewriter, currentOp.getDpsInits()[0].getDefiningOp(), true,
-            CInnermostDims, MDimNum > 1 ? 2 : 0)) ||
-        failed(setStaticSizeForExtractSliceOp(
-            rewriter, currentOp.getDpsInputs()[1].getDefiningOp(), true,
-            BInnermostDims, NDimNum > 1)) ||
-        failed(setStaticSizeForExtractSliceOp(
-            rewriter, currentOp.getDpsInputs()[0].getDefiningOp(), true,
-            AInnermostDims, MDimNum > 1)) ||
-        (currentOp.getDpsInits().size() > 1 &&
-         failed(setStaticSizeForExtractSliceOp(
-             rewriter, currentOp.getDpsInits()[1].getDefiningOp(), true,
-             CInnermostDims, MDimNum > 1 ? 2 : 0)))) {
-      return failure();
+    setStaticSizeForExtractSliceOp(rewriter,
+                                   currentOp.getDpsInputs()[1].getDefiningOp(),
+                                   true, BInnermostDims, NDimNum > 1);
+    setStaticSizeForExtractSliceOp(rewriter,
+                                   currentOp.getDpsInputs()[0].getDefiningOp(),
+                                   true, AInnermostDims, MDimNum > 1);
+    for (auto init : currentOp.getDpsInits()) {
+      setStaticSizeForExtractSliceOp(rewriter, init.getDefiningOp(), true,
+                                     CInnermostDims, MDimNum > 1 ? 2 : 0);
     }
+
     // View the tensor to brgemm required format
     Value dataOprand = tensorViewRankedTensor(
         rewriter,
@@ -841,10 +830,7 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
 
     // Insert the result back to the original tensor
     for (Operation *user : currentOp->getResult(0).getUsers()) {
-      if (failed(setStaticSizeForInsertSliceOp(rewriter, user, result,
-                                               CInnermostDims))) {
-        return failure();
-      }
+      setStaticSizeForInsertSliceOp(rewriter, user, result, CInnermostDims);
     }
 
     if (option.needLowPrecisionCast) {
@@ -869,10 +855,8 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
       auto ifOp = eb.getLastOperaion();
       // set static size for the insertSliceOp of copyOp
       for (Operation *user : currentOp->getResult(1).getUsers()) {
-        if (failed(setStaticSizeForInsertSliceOp(
-                rewriter, user, ifOp->getResult(0), CInnermostDims))) {
-          return failure();
-        }
+        setStaticSizeForInsertSliceOp(rewriter, user, ifOp->getResult(0),
+                                      CInnermostDims);
       }
       rewriter.replaceOp(currentOp, {matmul->getResult(0), ifOp->getResult(0)});
     } else {