fix bug for single iteration

small fix
bytedance · Jun 24, 2024 · 9dcf5d7 · 9dcf5d7
1 parent c4d3b35
commit 9dcf5d7
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 5 deletions.
diff --git a/compiler/lib/Dialect/Linalg/Transforms/LinalgPromotion.cpp b/compiler/lib/Dialect/Linalg/Transforms/LinalgPromotion.cpp
@@ -136,11 +136,12 @@ LogicalResult copyWorkgroupMemoryToGlobalMemory(OpBuilder &b, Value src,
   // get the only scf.for op inside the scf.forall op.
   scf::ForallOp forallOp = op->getParentOfType<scf::ForallOp>();
   auto forOps = llvm::to_vector(forallOp.getOps<scf::ForOp>());
-  if (forOps.size() != 1)
-    return forallOp.emitError("expected a single scf.for op");
 
   // copyWorkgroupMemoryToGlobalMemory after gemm compute ends.
-  b.setInsertionPointAfter(forOps[0]);
+  if (forOps.size() == 1)
+    b.setInsertionPointAfter(forOps[0]);
+  if (forOps.size() > 1)
+    return failure();
   b.create<gpu::BarrierOp>(src.getLoc());
   Operation *copyOp = b.create<linalg::CopyOp>(src.getLoc(), src, dst);
   setLinalgTransformationMarker(copyOp,
@@ -291,7 +292,10 @@ struct LinalgPromotionPass : public LinalgPromotionBase<LinalgPromotionPass> {
     // As we know linalg.matmul is in a scf.for, and the subview promotionImpl
     // inserts should be in the scf.forall op.
     auto forOp = linalgContractOp->getParentOfType<scf::ForOp>();
-    builder.setInsertionPoint(forOp); // before forOp
+    if (forOp)
+      builder.setInsertionPoint(forOp); // before forOp
+    else
+      builder.setInsertionPoint(linalgContractOp); // before linalgContractOp
     (void)promotionImpl<MatmulOperands::C>(builder, linalgContractOp);
 
     // The linalg.copy should be fused with its consumer linalg.generic.

diff --git a/compiler/lib/Pipelines/GPU/GemmCodegen.cpp b/compiler/lib/Pipelines/GPU/GemmCodegen.cpp
@@ -308,6 +308,29 @@ void mlir::createGPUAddGemmCodegenLoweringConfigTransform(
 }
 
 namespace {
+
+int numIterations(scf::ForOp forOp) {
+  Value lowerBound = forOp.getLowerBound();
+  Value upperBound = forOp.getUpperBound();
+  Value step = forOp.getStep();
+
+  // get def constant value
+  auto defLowerBound = lowerBound.getDefiningOp<arith::ConstantOp>();
+  auto defUpperBound = upperBound.getDefiningOp<arith::ConstantOp>();
+  auto defStep = step.getDefiningOp<arith::ConstantOp>();
+
+  if (defLowerBound && defUpperBound && defStep) {
+    auto lowerBoundValue = defLowerBound.getValue();
+    auto upperBoundValue = defUpperBound.getValue();
+    auto stepValue = defStep.getValue();
+
+    auto lowerBoundInt = cast<IntegerAttr>(lowerBoundValue).getInt();
+    auto upperBoundInt = cast<IntegerAttr>(upperBoundValue).getInt();
+    auto stepInt = cast<IntegerAttr>(stepValue).getInt();
+    return (upperBoundInt - lowerBoundInt) / stepInt;
+  }
+  return -1;
+}
 void createGPUPipeliningTransformImpl(OpPassManager &pm,
                                       const std::string &anchor,
                                       const std::string &prefix) {
@@ -318,7 +341,21 @@ void createGPUPipeliningTransformImpl(OpPassManager &pm,
 
   config.opFilter = [=](Operation *op) {
     if (auto forallOp = llvm::dyn_cast_or_null<scf::ForallOp>(op)) {
-      return isMappedToGPUBlocks(forallOp);
+      if (!isMappedToGPUBlocks(forallOp)) {
+        return false;
+      }
+      func::FuncOp funcOp = forallOp->getParentOfType<func::FuncOp>();
+      auto pipelineStageOptional = getGemmPipelineDepth(funcOp);
+      if (!pipelineStageOptional) {
+        return false;
+      }
+      SmallVector<scf::ForOp> forOps;
+      forallOp.walk([&](scf::ForOp forOp) { forOps.push_back(forOp); });
+      if (forOps.size() != 1)
+        return false;
+      scf::ForOp forOp = forOps[0];
+      if (numIterations(forOp) <= pipelineStageOptional.value())
+        return false;
     }
     return false;
   };