[Warp Spec] Always use 1 buffer for SSA partition dependencies (#7686)

There are realistic cases where this should be more than 1, but let's
fix it to 1 for now since that's what happens in practice.

Author: Mogball

lib/Dialect/TritonGPU/Transforms/WarpSpecialization/RewritePartitionDependencies.cpp

@@ -48,7 +48,7 @@ struct UseInfo {
 int UseInfo::getMaxUseDistance(const Partition &partition) {
   int maxDistance = 0;
   for (auto [usePartition, distance] : llvm::make_first_range(consumers)) {
-    int dist = 2 + distance;
+    int dist = 1 + distance;
     maxDistance = std::max(maxDistance, dist);
   }
   return maxDistance;

test/TritonGPU/rewrite-partition-dependencies.mlir

@@ -10,7 +10,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
 // CHECK-LABEL: @two_consumers
 tt.func @two_consumers(%lb: i32, %ub: i32, %step: i32) {
   // CHECK: [[C0:%.*]] = arith.constant 0 : i32
-  // CHECK-NEXT: [[ABUF:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<2x1xi32, {{.*}}>
+  // CHECK-NEXT: [[ABUF:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<1x1xi32, {{.*}}>
   // CHECK-NEXT: [[AREF:%.*]] = nvws.aref.create [[ABUF]]
   scf.for %i = %lb to %ub step %step iter_args() -> () : i32 {
     %0 = "op_a"() {ttg.partition = 0} : () -> !ty
@@ -40,7 +40,7 @@ tt.func @two_consumers(%lb: i32, %ub: i32, %step: i32) {
 // CHECK-LABEL: @distance_one
 tt.func @distance_one(%lb: i32, %ub: i32, %step: i32) {
   // CHECK: [[C0:%.*]] = arith.constant 0 : i32
-  // CHECK: [[ABUF:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<2x1xi32, {{.*}}>
+  // CHECK: [[ABUF:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<1x1xi32, {{.*}}>
   // CHECK-NEXT: [[AREF:%.*]] = nvws.aref.create [[ABUF]]
   %cst = arith.constant dense<0> : !ty
   // CHECK: scf.for [[IV:%.*]] = [[LB:%.*]] to [[UB:%.*]] step [[STEP:%.*]] iter_args([[K:%.*]] = {{.*}})
@@ -63,9 +63,9 @@ tt.func @distance_one(%lb: i32, %ub: i32, %step: i32) {
 }
 
 tt.func @complex_case(%lb: i32, %ub: i32, %step: i32) {
-  // CHECK: [[ABUF1:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<2x1xi32, {{.*}}>
+  // CHECK: [[ABUF1:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<1x1xi32, {{.*}}>
   // CHECK-NEXT: [[AREF1:%.*]] = nvws.aref.create [[ABUF1]]
-  // CHECK-NEXT: [[ABUF2:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<2x1xi32, {{.*}}>
+  // CHECK-NEXT: [[ABUF2:%.*]] = ttg.local_alloc : () -> !ttg.memdesc<1x1xi32, {{.*}}>
   // CHECK-NEXT: [[AREF2:%.*]] = nvws.aref.create [[ABUF2]]
   %cst = arith.constant dense<0> : !ty
   // CHECK: scf.for [[IV:%.*]] = [[LB:%.*]] to [[UB:%.*]] step [[STEP:%.*]] iter_args([[K:%.*]] = {{.*}}, [[L:%.*]] = {{.*}})