Addressing PR comments

Author: dhernandez0

mlir/include/mlir/Dialect/Rock/utility/loweringUtils.h

@@ -262,9 +262,6 @@ FailureOr<VectorDimInfo> getVectorDim(Location loc, Value matrix, Type elemType,
 // Get the LDS size of the memref
 std::optional<int64_t> getWorkgroupMemorySize(MemRefType type);
 
-// Return trip count for scf::ForOp
-std::optional<int64_t> getConstantLoopTripCount(scf::ForOp loopOp);
-
 } // end namespace rock
 } // end namespace mlir
 #endif

mlir/lib/Dialect/Rock/Transforms/GridwiseGemmToBlockwise.cpp

@@ -2262,14 +2262,15 @@ struct GridwiseAttentionAccelRewritePattern
       // LDS Barrier (issue 1811): some threads might be loading from LDS
       // while others are in the next iteration (here), writing to LDS. This
       // barrier prevents that.
-      std::optional<uint64_t> mLoopIters = std::nullopt;
+      std::optional<APInt> mLoopIters = std::nullopt;
       // mLoopOp can be a dynamic loop if we are using KV Cache or Causal
       // masking. If that's the case, we can't know the number of iterations
       // at compile time.
       if (!dynamicMLoop)
-        mLoopIters = getConstantLoopTripCount(mLoopOp);
+        mLoopIters = mLoopOp.getStaticTripCount();
 
-      bool mIterOneIter = mLoopIters.has_value() && mLoopIters.value() == 1;
+      bool mIterOneIter =
+          mLoopIters.has_value() && mLoopIters.value().getSExtValue() == 1;
       if (!mIterOneIter) {
         LLVM_DEBUG(llvm::dbgs() << "adding a barrier in the first gemm loop\n");
         LDSBarrierOp::create(rewriter, loc);

mlir/lib/Dialect/Rock/utility/loweringUtils.cpp

@@ -1184,21 +1184,3 @@ std::optional<int64_t> mlir::rock::getWorkgroupMemorySize(MemRefType type) {
   }
   return std::nullopt;
 }
-
-std::optional<int64_t> mlir::rock::getConstantLoopTripCount(scf::ForOp loopOp) {
-  Value lowerBound = loopOp.getLowerBound();
-  Value upperBound = loopOp.getUpperBound();
-  Value step = loopOp.getStep();
-  auto lbConst = getConstantIntValue(lowerBound);
-  auto ubConst = getConstantIntValue(upperBound);
-  auto stepConst = getConstantIntValue(step);
-  if (lbConst && ubConst && stepConst) {
-    int64_t tripCount =
-        (ubConst.value() - lbConst.value() + stepConst.value() - 1) /
-        stepConst.value();
-    // Handle the case where tripCount might be negative
-    tripCount = std::max(tripCount, 0L);
-    return tripCount;
-  }
-  return std::nullopt;
-}

mlir/test/e2e/AttentionDirectToLDS.cfg

@@ -1,4 +1,4 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True
 
 if not 'direct_to_lds_32b' in config.features and not 'direct_to_lds_128b' in config.features:

mlir/test/e2e/AttentionSchedule.cfg

@@ -1,2 +1,2 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True

mlir/test/e2e/ConvElementwiseGemmDirectToLDS.cfg

@@ -1,4 +1,4 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True
 
 if not 'direct_to_lds_32b' in config.features and not 'direct_to_lds_128b' in config.features:

mlir/test/e2e/ConvElementwiseGemmSchedule.cfg

@@ -1,2 +1,2 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True

mlir/test/e2e/GemmElementwiseGemmDirectToLDS.cfg

@@ -1,4 +1,4 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True
 
 if not 'direct_to_lds_32b' in config.features and not 'direct_to_lds_128b' in config.features:

mlir/test/e2e/GemmElementwiseGemmSchedule.cfg

@@ -1,2 +1,2 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True

mlir/test/e2e/PrAttentionBF16.cfg

@@ -1,2 +1,2 @@
-if (not config.arch_support_mfma) and (not config.arch_support_wmma):
+if not (config.arch_support_mfma or config.arch_support_wmma):
   config.unsupported = True