Init C buffer with easy builder

Author: zhczhong

lib/gc/Transforms/CMakeLists.txt

@@ -7,6 +7,7 @@ gc_set_mlir_link_components(MLIR_LINK_COMPONENTS
 add_mlir_library(GCPasses
   OneDNNGraphToLinalg.cpp
   DeepTileContractionNamedOp.cpp
+  Tiling.cpp
 
   ADDITIONAL_HEADER_DIRS
     ${PROJECT_SOURCE_DIR}/include

lib/gc/Transforms/DeepTileContractionNamedOp.cpp

@@ -11,6 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "./Tiling.hpp"
+#include "gc/Dialect/Arith/Utils/EasyBuild.h"
+#include "gc/IR/EasyBuild.h"
+#include "gc/IR/EasyBuildSCF.h"
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -179,6 +182,7 @@ struct OuterLoopGenerationResult {
   SmallVector<Operation *> tiledOps;
   /// The `scf.for` operations that iterate over the tiles.
   SmallVector<LoopLikeOpInterface> loops;
+  SmallVector<LoopLikeOpInterface> reductionLoops;
   /// Values to use as replacements for the untiled op. Is the same size as the
   /// number of results of the untiled op.
   SmallVector<Value> replacements;
@@ -192,6 +196,8 @@ generateOuterLoop(RewriterBase &b, linalg::LinalgOp linalgOp,
   auto nestedTileSizes = option.nestedTileSizes;
   auto loopType = option.loopType;
   auto loopDim = option.loopDim;
+  SmallVector<mlir::utils::IteratorType> iteratorTypes =
+      linalgOp.getIteratorTypesArray();
 
   if (loopType.size() != loopDim.size() ||
       loopDim.size() != nestedTileSizes.size()) {
@@ -228,6 +234,13 @@ generateOuterLoop(RewriterBase &b, linalg::LinalgOp linalgOp,
         return failure();
       b.replaceOp(currentOp, tilingResult->replacements);
       currentOp = dyn_cast<linalg::LinalgOp>(tilingResult->tiledOps.back());
+
+      for (auto [dim, loop] : llvm::zip(currentDim, tilingResult->loops)) {
+        if (iteratorTypes[dim] == mlir::utils::IteratorType::reduction) {
+          result.reductionLoops.push_back(loop);
+        }
+        result.loops.push_back(loop);
+      }
     } else if (type == OuterLoopGenerationOption::LoopType::ForallOp) {
       SmallVector<OpFoldResult> tileSizes(
           currentOp.getNumLoops(), getAsIndexOpFoldResult(b.getContext(), 0));
@@ -395,16 +408,16 @@ getOprandDimType(linalg::LinalgOp &linalgOp) {
 }
 
 /*
-forall([PM, PN]: [MThreads, NThreads) {
-  for(PK : KThreads) {
+matmul(A, B) -> C
+---------------->
+forall([PM, PN, PK]: [MThreads, NThreads, KThreads]) {
     CSlice = [KThreads, PM * MOuterBlock: (PM + 1) * MOuterBlock,
      PN * NOuterBlock: (PN + 1) * NOuterBlock]
     ASlice = A[PM * MOuterBlock: (PM + 1) * MOuterBlock, PK * KOuterBlock * (PK
 + 1) * KOuterBlock]
     BSlice = B[PK * KOuterBlock * (PK + 1) * KOuterBlock, PN *
 NOuterBlock: (PN + 1) * NOuterBlock] CSlice2 = CSlice[PK, PM * MOuterBlock: (PM
 + 1) * MOuterBlock, PN * NOuterBlock: (PN + 1) * NOuterBlock]
-
     MNumBlock = MOuterBlock / MBlock
     NNumBlock = NOuterBlock / NBlock
     KNumBlock = KOuterBlock / KBlovk
@@ -426,9 +439,8 @@ iin_block_: (in + 1) * iin_block_] (init with 0 when ok == 0)
 A=ASlice3, B=BSlice3, C=CSlice4, onlyUpdate=(ok!=0));
       }
     }
-  }
-  C = final_reduce(CSlice)
 }
+C = final_reduce(CSlice)
 */
 struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
   using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern;
@@ -508,12 +520,14 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
   struct innerBodyGenerationOption {
     bool hasFillOp = false;
     Value fillValue;
+    SmallVector<LoopLikeOpInterface> KLoopHandles;
   };
 
   LogicalResult
   innerBodyGeneration(RewriterBase &rewriter, linalg::LinalgOp originOp,
                       linalg::LinalgOp currentOp,
                       const innerBodyGenerationOption &option) const {
+    mlir::easybuild::EasyBuilder eb{rewriter, originOp.getLoc()};
     auto operandDimTypes = getOprandDimType(originOp);
     MatmulConfig cfg = getDefaultMatmulConfig(originOp);
     auto AShape = originOp.getShape(originOp.getDpsInputOperand(0));
@@ -656,18 +670,31 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
     currentOp = matmul;
 
     if (option.hasFillOp) {
-      // TODO: support partial K in sinsngle threads, control flow may need
-      // easy builder support
       rewriter.setInsertionPointAfter(currentOp);
-      auto fillOp = rewriter.create<linalg::FillOp>(
-          currentOp->getLoc(), option.fillValue, currentOp.getDpsInits()[0]);
-      IRMapping mapping;
-      mapping.map(currentOp.getDpsInits()[0], fillOp.getResult(0));
-      auto res = rewriter.clone(*(currentOp.getOperation()), mapping);
-      rewriter.replaceOp(currentOp, res);
-      currentOp = dyn_cast<linalg::LinalgOp>(res);
+
+      auto cond = eb(true);
+      for (auto loop : option.KLoopHandles) {
+        auto induceVar = eb.wrap<mlir::easybuild::EBUnsigned>(
+            loop.getLoopRegions().front()->front().getArgument(0));
+        auto currentCond = induceVar == eb.toIndex(0);
+        cond = cond & currentCond;
+      }
+      EB_scf_if(cond, {currentOp.getDpsInits()[0].getType()}) {
+        auto fillOp = rewriter.create<linalg::FillOp>(
+            currentOp->getLoc(), option.fillValue, currentOp.getDpsInits()[0]);
+        IRMapping mapping;
+        mapping.map(currentOp.getDpsInits()[0], fillOp.getResult(0));
+        auto res = rewriter.clone(*(currentOp.getOperation()), mapping);
+        eb.yield(res->getResult(0));
+      }
+      EB_else {
+        auto res = rewriter.clone(*(currentOp.getOperation()));
+        eb.yield(res->getResult(0));
+      }
+      auto ifOp = eb.getLastOperaion();
+      rewriter.replaceOp(currentOp, ifOp);
+      ifOp->getParentOfType<func::FuncOp>().dump();
     }
-    currentOp.getOperation()->getParentOfType<func::FuncOp>().dump();
     return success();
   }
 
@@ -685,7 +712,6 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
     // if-else block)
     bool hasFillOp = false;
     Value fillValue;
-    SmallVector<LoopLikeOpInterface> KLoopHandle;
     if (auto op = dyn_cast<linalg::FillOp>(
             linalgOp.getDpsInits()[0].getDefiningOp())) {
       hasFillOp = true;
@@ -707,7 +733,8 @@ struct deepTileMatmul : public OpInterfaceRewritePattern<linalg::LinalgOp> {
     // Step 3 inner loop generation, convert the linalg.generic to brgemm
     if (failed(innerBodyGeneration(
             rewriter, matmulOp, linalgOp,
-            innerBodyGenerationOption{hasFillOp, fillValue}))) {
+            innerBodyGenerationOption{hasFillOp, fillValue,
+                                      outerLoopResult->reductionLoops}))) {
       return failure();
     }
     return success();

lib/gc/Transforms/Tiling.cpp

@@ -844,7 +844,7 @@ tileAllUsingForall(RewriterBase &b, PartialReductionOpInterface op,
     return b.notifyMatchFailure(op, "if tile sizes are present it must have as "
                                     "many elements as number of threads");
 
-  if (redDims.front() >= numThreads.size())
+  if ((unsigned)redDims.front() >= numThreads.size())
     return b.notifyMatchFailure(
         op, "reduction dimension must be mapped to threads");
 
@@ -914,7 +914,7 @@ tileAllUsingForall(RewriterBase &b, PartialReductionOpInterface op,
         }
 
         auto nonZeroDimIdx = 0;
-        for (auto dim = 0; dim < numThreads.size(); dim++) {
+        for (auto dim = 0UL; dim < numThreads.size(); dim++) {
           if (!isConstantIntValue(numThreads[dim], 0)) {
             if (llvm::find(redDims, dim) != redDims.end())
               outOffsets[dim] = forallOp.getInductionVars()[nonZeroDimIdx];
@@ -973,7 +973,7 @@ tileAllUsingForall(RewriterBase &b, PartialReductionOpInterface op,
     int64_t offIdx = 0;
     int64_t sizeIdx = 0;
     int64_t nonZeroDimIdx = 0;
-    for (int64_t i = 0; i < numThreads.size(); ++i) {
+    for (auto i = 0UL; i < numThreads.size(); ++i) {
       if (llvm::find(redDims, i) != redDims.end()) {
         if (hasReductionThreads) {
           resultOffsetsRank.push_back(