Properly set the upper bound of gpu.lane_id when rewriteForallToLanes. (iree-org#20513)

If subgroup size is available, properly sets the upper bound of
gpu.lane_id when it's created.

Fixes: iree-org#20385

Signed-off-by: Lin, Peiyong <[email protected]>

Author: lpy

compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_distribute_forall.mlir

@@ -60,7 +60,7 @@ func.func @distribute_lane_forall(%out : memref<?xi32>)
 }
 
 // CHECK-LABEL: func @distribute_lane_forall
-//       CHECK:   %[[LANEID:.+]] = gpu.lane_id
+//       CHECK:   %[[LANEID:.+]] = gpu.lane_id upper_bound 32
 //       CHECK:   memref.store {{.*}}[%[[LANEID]]]
 
 // -----

compiler/src/iree/compiler/Codegen/Dialect/GPU/TransformExtensions/test/distribute_lane_forall.mlir

@@ -28,7 +28,7 @@ module attributes { transform.with_named_sequence } {
 // CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (d0 * 16)>
 
 // CHECK-LABEL: func @distribute_lane_forall
-//       CHECK:   %[[LANE_ID:.+]] = gpu.lane_id
+//       CHECK:   %[[LANE_ID:.+]] = gpu.lane_id upper_bound 64
 //   CHECK-NOT:   scf.forall
 //       CHECK:   affine.delinearize_index %[[LANE_ID]] into (4, 16) : index, index
 //       CHECK:   linalg.copy

compiler/src/iree/compiler/Codegen/Dialect/GPU/Transforms/Transforms.cpp

@@ -1469,7 +1469,15 @@ static void rewriteForallToLanes(RewriterBase &rewriter, scf::ForallOp forallOp,
   Location loc = forallOp->getLoc();
   assert(isLaneMappableForall(forallOp) && "mapping non-lane forall op");
 
-  Value laneId = rewriter.create<gpu::LaneIdOp>(loc, /*upperBound=*/nullptr);
+  auto upperBounds = forallOp.getLoopUpperBounds();
+  std::optional<IntegerAttr> upperBound;
+  if (upperBounds && upperBounds->size() > 0) {
+    if (auto upperBoundAttr = (*upperBounds)[0].dyn_cast<Attribute>()) {
+      upperBound = dyn_cast<IntegerAttr>(upperBoundAttr);
+    }
+  }
+  Value laneId = rewriter.create<gpu::LaneIdOp>(
+      loc, upperBound ? rewriter.getIndexAttr(upperBound->getInt()) : nullptr);
   rewriter.eraseOp(forallOp.getTerminator());
   rewriter.setInsertionPoint(forallOp);
   rewriter.inlineBlockBefore(forallOp.getBody(), forallOp, {laneId});