[AMD] Fix program id/count range in RangeAnalysis (#8103)

The program id/count starts from 0 albeit it is a signed integer.
This change fixes some places where don't take into account.

Author: yangshuxin

test/TritonGPU/amd/amd-range-analysis.mlir

@@ -35,7 +35,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
   tt.func @assumepid(%arg0: !tt.ptr<f32>) -> tensor<1024xf32> {
     %c0 = arith.constant 0 : i32
     %c1024_i32 = arith.constant 1024 : i32
-    // expected-remark@+2 {{unsigned : [0, 2147483647] signed : [0, 1024]}}
+    // expected-remark@+2 {{unsigned : [0, 1024] signed : [0, 1024]}}
     // expected-remark@+1 {{non-neg}}
     %pid = tt.get_program_id x : i32
     // expected-remark@+2 {{unsigned : [1, 1] signed : [-1, -1]}}
@@ -519,7 +519,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
     // expected-remark@+1 {{non-neg}}
     %0 = tt.get_program_id x : i32
     %c65535_i32 = arith.constant 65535 : i32
-    %cmpule_pid = arith.cmpi ule, %0, %c65535_i32 : i32
+    %cmpule_pid = arith.cmpi sle, %0, %c65535_i32 : i32
     llvm.intr.assume %cmpule_pid : i1
     // expected-remark@+2 {{unsigned : [0, 8388480] signed : [0, 8388480]}}
     // expected-remark@+1 {{non-neg}}

third_party/amd/lib/Analysis/RangeAnalysis.cpp

@@ -193,6 +193,12 @@ std::optional<ConstantIntRanges> maybeGetAssumedRange(Operation *assumption,
     APInt min, max;
     if (isSigned) {
       min = APInt::getSignedMinValue(bitWidth);
+      if (llvm::isa_and_nonnull<mlir::triton::GetProgramIdOp,
+                                mlir::triton::GetNumProgramsOp>(
+              anchor.getDefiningOp())) {
+        min = APInt::getZero(bitWidth);
+      } else
+        min = APInt::getSignedMinValue(bitWidth);
       max = APInt::getSignedMaxValue(bitWidth);
     } else {
       min = APInt::getMinValue(bitWidth);
@@ -297,6 +303,12 @@ TritonIntegerRangeAnalysis::maybeGetAssumedRange(Value anchor) const {
   unsigned bitWidth = ConstantIntRanges::getStorageBitwidth(anchor.getType());
   assert(bitWidth > 0 && "expected non-zero bitwidth");
   ConstantIntRanges constIntRange = ConstantIntRanges::maxRange(bitWidth);
+  if (llvm::isa_and_nonnull<GetProgramIdOp>(anchor.getDefiningOp())) {
+    constIntRange = ConstantIntRanges::range(
+        APInt::getZero(bitWidth),
+        APInt(bitWidth, kDefaultMaxPrograms - 1, true), true);
+  }
+
   for (auto assumption : matchingAssumptions) {
     if (auto constIntRange_ = ::maybeGetAssumedRange(assumption, anchor))
       constIntRange = constIntRange.intersection(*constIntRange_);