use IntegerRangeAnalysis and update launchOp::inferResultRanges.

Author: linuxlonelyeagle

mlir/include/mlir/Dialect/Affine/Analysis/LoopAnalysis.h

@@ -43,9 +43,10 @@ void getTripCountMapAndOperands(AffineForOp forOp, AffineMap *map,
 /// constant trip count in non-trivial cases.
 std::optional<uint64_t> getConstantTripCount(AffineForOp forOp);
 
-/// In some scenarios, such as GPU, the number of trip of each thread in the
-/// loop is inconsistent. This function returns the maximum number of trip.
-std::optional<uint64_t> getMaxConstantTripCount(AffineForOp forOp);
+/// Returns the maximum trip count when the operand of forOp has a range. If the
+/// operand of forOp is a constant, the return value is the same as
+/// `getConstantTripCount`.
+std::optional<uint64_t> getUpperBoundOnTripCount(AffineForOp forOp);
 
 /// Returns the greatest known integral divisor of the trip count. Affine
 /// expression analysis is used (indirectly through getTripCount), and

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -1035,12 +1035,6 @@ def GPU_LaunchOp : GPU_Op<"launch", [
     static StringRef getNumWorkgroupAttributionsAttrName() {
       return "workgroup_attributions";
     }
-
-    /// Find BlockSize via the BlockArgument of gpu.launch.
-    Value getBlockSizeOnAxis(Value threadId);
-
-    /// Find BlockSize via the Dimension Information.
-    Value getBlockSizeOnAxis(Dimension dimension);
   }];
 
   let hasCanonicalizer = 1;

mlir/lib/Dialect/Affine/Analysis/LoopAnalysis.cpp

@@ -12,13 +12,15 @@
 
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/IntegerRangeAnalysis.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Dialect/Affine/Analysis/NestedMatcher.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Interfaces/FunctionInterfaces.h"
 #include "llvm/Support/MathExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
@@ -31,6 +33,7 @@
 
 using namespace mlir;
 using namespace mlir::affine;
+using namespace mlir::dataflow;
 
 #define DEBUG_TYPE "affine-loop-analysis"
 
@@ -85,48 +88,54 @@ void mlir::affine::getTripCountMapAndOperands(
                             tripCountValueMap.getOperands().end());
 }
 
-/// Replace thread_id with its maximum value, if `replaceWithZero` is true,
-/// thread_id will be replaced by its minimum value 0.
-static void replaceGPUOperands(AffineForOp forOp,
-                               SmallVectorImpl<Value> &operands,
-                               SmallVectorImpl<AffineExpr> &symReplacements,
-                               unsigned numDim, bool replaceWithZero = false) {
-  auto launchOp = forOp->getParentOfType<gpu::LaunchOp>();
-  if (!launchOp)
+/// By running `IntegerRangeAnalysis` to get the ranges of operand, then fill
+/// the `symReplacements` with range. If `replaceByMin` is set to true,
+/// construct `replacement` using the smallest value.By default, the largest
+/// value will be used for constructing `replacement`.
+static void replaceOperandByRange(AffineForOp forOp,
+                                  SmallVectorImpl<Value> &operands,
+                                  SmallVectorImpl<AffineExpr> &symReplacements,
+                                  unsigned numDim, bool replaceByMin = false) {
+  DataFlowSolver solver;
+  solver.load<DeadCodeAnalysis>();
+  solver.load<IntegerRangeAnalysis>();
+  if (failed(solver.initializeAndRun(
+          forOp->getParentOfType<FunctionOpInterface>())))
     return;
 
-  // `b` is only used to create `AffineExpr`.
+  // `b` is used to create affineExpr
   Builder b(forOp.getContext());
-  unsigned idx = 0;
-
   for (unsigned i = numDim, e = operands.size(); i < e; ++i) {
     Value operand = operands[i];
-    if (Value blockSize = launchOp.getBlockSizeOnAxis(operand)) {
-      operands[i] = blockSize;
-      if (!replaceWithZero)
-        symReplacements.push_back(b.getAffineSymbolExpr(idx++) - 1);
-      else
-        symReplacements.push_back(b.getAffineConstantExpr(0));
+    auto lattice =
+        solver.lookupState<dataflow::IntegerValueRangeLattice>(operand);
+    if (!lattice) {
+      symReplacements.push_back(b.getAffineSymbolExpr(i - numDim));
       continue;
     }
 
-    Operation *defOp = operand.getDefiningOp();
-    if (!defOp) {
-      ++idx;
+    if (lattice->getValue().isUninitialized()) {
+      symReplacements.push_back(b.getAffineSymbolExpr(i - numDim));
       continue;
     }
 
-    if (auto threadIdOp = mlir::dyn_cast<gpu::ThreadIdOp>(defOp)) {
-      gpu::Dimension dimension = threadIdOp.getDimension();
-      operands[i] = launchOp.getBlockSizeOnAxis(dimension);
-      if (!replaceWithZero)
-        symReplacements.push_back(b.getAffineSymbolExpr(idx++) - 1);
-      else
-        symReplacements.push_back(b.getAffineConstantExpr(0));
+    ConstantIntRanges range = lattice->getValue().getValue();
+    APInt max = range.smax();
+    APInt min = range.smin();
+    unsigned bitNums = max.getBitWidth();
+
+    if (APInt::getSignedMaxValue(bitNums) == max &&
+        APInt::getSignedMinValue(bitNums) == min) {
+      symReplacements.push_back(b.getAffineSymbolExpr(i - numDim));
       continue;
     }
-    ++idx;
+
+    if (!replaceByMin)
+      symReplacements.push_back(b.getAffineConstantExpr(max.getZExtValue()));
+    else
+      symReplacements.push_back(b.getAffineConstantExpr(min.getZExtValue()));
   }
+  return;
 }
 
 /// Take the min if all trip counts are constant.
@@ -158,32 +167,28 @@ std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
   if (!map)
     return std::nullopt;
   SmallVector<AffineExpr, 4> symReplacements;
-  replaceGPUOperands(forOp, operands, symReplacements, map.getNumDims());
+  replaceOperandByRange(forOp, operands, symReplacements, map.getNumDims());
   map = map.replaceDimsAndSymbols({}, symReplacements, map.getNumDims(),
                                   map.getNumSymbols());
-  affine::AffineValueMap valueMap(map, operands);
-  (void)valueMap.canonicalize();
-  map = valueMap.getAffineMap();
   return getConstantTripCountFromAffineMap(map);
 }
 
-/// In some scenarios, such as GPU, the number of trip of each thread in the
-/// loop is inconsistent. This function returns the maximum number of trip.
+/// Returns the maximum trip count when the operand of forOp has a range. If the
+/// operand of forOp is a constant, the return value is the same as
+/// `getConstantTripCount`.
 std::optional<uint64_t>
-mlir::affine::getMaxConstantTripCount(AffineForOp forOp) {
+mlir::affine::getUpperBoundOnTripCount(AffineForOp forOp) {
   SmallVector<Value, 4> operands;
   AffineMap map;
   getTripCountMapAndOperands(forOp, &map, &operands);
 
   if (!map)
     return std::nullopt;
   SmallVector<AffineExpr, 4> symReplacements;
-  replaceGPUOperands(forOp, operands, symReplacements, map.getNumDims(), true);
+  replaceOperandByRange(forOp, operands, symReplacements, map.getNumDims(),
+                        true);
   map = map.replaceDimsAndSymbols({}, symReplacements, map.getNumDims(),
                                   map.getNumSymbols());
-  affine::AffineValueMap valueMap(map, operands);
-  (void)valueMap.canonicalize();
-  map = valueMap.getAffineMap();
   return getConstantTripCountFromAffineMap(map);
 }
 
@@ -198,12 +203,9 @@ uint64_t mlir::affine::getLargestDivisorOfTripCount(AffineForOp forOp) {
   if (!map)
     return 1;
   SmallVector<AffineExpr, 4> symReplacements;
-  replaceGPUOperands(forOp, operands, symReplacements, map.getNumDims());
+  replaceOperandByRange(forOp, operands, symReplacements, map.getNumDims());
   map = map.replaceDimsAndSymbols({}, symReplacements, map.getNumDims(),
                                   map.getNumSymbols());
-  affine::AffineValueMap valueMap(map, operands);
-  (void)valueMap.canonicalize();
-  map = valueMap.getAffineMap();
   // The largest divisor of the trip count is the GCD of the individual largest
   // divisors.
   assert(map.getNumResults() >= 1 && "expected one or more results");

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

@@ -17,7 +17,6 @@
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/IRMapping.h"
@@ -118,7 +117,7 @@ static void replaceIterArgsAndYieldResults(AffineForOp forOp) {
 /// was known to have a single iteration.
 LogicalResult mlir::affine::promoteIfSingleIteration(AffineForOp forOp) {
   std::optional<uint64_t> tripCount = getConstantTripCount(forOp);
-  std::optional<uint64_t> maxTripCount = getMaxConstantTripCount(forOp);
+  std::optional<uint64_t> maxTripCount = getUpperBoundOnTripCount(forOp);
   if (!tripCount || *tripCount != 1 || !maxTripCount || *maxTripCount != 1)
     return failure();
 
@@ -888,7 +887,7 @@ void mlir::affine::getTileableBands(
 LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
   std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   std::optional<uint64_t> maxMayBeConstantTripCount =
-      getMaxConstantTripCount(forOp);
+      getUpperBoundOnTripCount(forOp);
 
   if (!mayBeConstantTripCount.has_value() &&
       !maxMayBeConstantTripCount.has_value())
@@ -1025,7 +1024,7 @@ LogicalResult mlir::affine::loopUnrollByFactor(
 
   std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
   std::optional<uint64_t> maxMayBeConstantTripCount =
-      getMaxConstantTripCount(forOp);
+      getUpperBoundOnTripCount(forOp);
   if (unrollFactor == 1) {
     if (mayBeConstantTripCount && *mayBeConstantTripCount == 1 &&
         maxMayBeConstantTripCount && *maxMayBeConstantTripCount == 1 &&

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

@@ -799,26 +799,6 @@ std::optional<KernelDim3> LaunchOp::getClusterSizeOperandValues() {
   return KernelDim3{operands[6], operands[7], operands[8]};
 }
 
-Value LaunchOp::getBlockSizeOnAxis(Dimension dimension) {
-  if (dimension == Dimension::x)
-    return getBlockSizeX();
-  else if (dimension == Dimension::y)
-    return getBlockSizeY();
-  else
-    return getBlockSizeZ();
-}
-
-Value LaunchOp::getBlockSizeOnAxis(Value threadId) {
-  KernelDim3 threadIds = getThreadIds();
-  if (threadIds.x == threadId)
-    return getBlockSizeX();
-  else if (threadIds.y == threadId)
-    return getBlockSizeY();
-  else if (threadIds.z == threadId)
-    return getBlockSizeZ();
-  return {};
-}
-
 LogicalResult LaunchOp::verify() {
   if (!(hasClusterSize()) &&
       (getClusterSizeX() || getClusterSizeY() || getClusterSizeZ()))

mlir/lib/Dialect/GPU/IR/InferIntRangeInterfaceImpls.cpp

@@ -250,26 +250,34 @@ void SubgroupSizeOp::inferResultRanges(ArrayRef<ConstantIntRanges>,
 void LaunchOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                  SetIntRangeFn setResultRange) {
   auto setRange = [&](const ConstantIntRanges &argRange, Value dimResult,
-                      Value idxResult) {
+                      Value idxResult, Value size) {
     if (argRange.umin().getBitWidth() != IndexType::kInternalStorageBitWidth)
       return;
-    ConstantIntRanges dimRange =
-        argRange.intersection(getIndexRange(1, kMaxDim));
-    setResultRange(dimResult, dimRange);
-    ConstantIntRanges idxRange =
-        getIndexRange(0, dimRange.umax().getZExtValue() - 1);
-    setResultRange(idxResult, idxRange);
+    APInt sizeInt;
+    if (matchPattern(size, m_ConstantInt(&sizeInt))) {
+      ConstantIntRanges dimRange = ConstantIntRanges::constant(sizeInt);
+      setResultRange(dimResult, dimRange);
+      ConstantIntRanges idxRange = getIndexRange(0, sizeInt.getZExtValue() - 1);
+      setResultRange(idxResult, idxRange);
+    } else {
+      ConstantIntRanges dimRange =
+          argRange.intersection(getIndexRange(1, kMaxDim));
+      setResultRange(dimResult, dimRange);
+      ConstantIntRanges idxRange =
+          getIndexRange(0, dimRange.umax().getZExtValue() - 1);
+      setResultRange(idxResult, idxRange);
+    }
   };
 
   argRanges = argRanges.drop_front(getAsyncDependencies().size());
   KernelDim3 gridDims = getGridSize();
   KernelDim3 blockIds = getBlockIds();
-  setRange(argRanges[0], gridDims.x, blockIds.x);
-  setRange(argRanges[1], gridDims.y, blockIds.y);
-  setRange(argRanges[2], gridDims.z, blockIds.z);
+  setRange(argRanges[0], gridDims.x, blockIds.x, getGridSizeX());
+  setRange(argRanges[1], gridDims.y, blockIds.y, getGridSizeY());
+  setRange(argRanges[2], gridDims.z, blockIds.z, getGridSizeZ());
   KernelDim3 blockDims = getBlockSize();
   KernelDim3 threadIds = getThreadIds();
-  setRange(argRanges[3], blockDims.x, threadIds.x);
-  setRange(argRanges[4], blockDims.y, threadIds.y);
-  setRange(argRanges[5], blockDims.z, threadIds.z);
+  setRange(argRanges[3], blockDims.x, threadIds.x, getBlockSizeX());
+  setRange(argRanges[4], blockDims.y, threadIds.y, getBlockSizeY());
+  setRange(argRanges[5], blockDims.z, threadIds.z, getBlockSizeZ());
 }