more clean up

Signed-off-by: Eric Feng <Eric.Feng@amd.com>

Author: efric

compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp

@@ -966,43 +966,24 @@ static LogicalResult populateCanonicalOffsetsSizesAndStrides(
 
   // Build broadcastIndex from unused delinearize results. Unused results
   // (not mapped by dimToVtid) represent broadcast lanes that see duplicate
-  // data. When a splitDim exists, we linearize these unused components into
-  // a single index in [0, broadcastFactor) to differentiate those lanes.
+  // data. When a splitDim exists, we use the single unused component as an
+  // index in [0, broadcastFactor) to differentiate those lanes.
   Value broadcastIndex;
   if (splitDim) {
     llvm::SmallDenseSet<size_t, 4> usedResults(dimToVtid.begin(),
                                                dimToVtid.end());
-
-    // Collect unused basis entries and compute overflow size.
-    int64_t unusedBasisProduct = 1;
-    SmallVector<Value> unusedLaneIdComponents;
-    SmallVector<int64_t> unusedBasisSizes;
+    // Find the single unused delinearize result representing broadcast lanes.
+    std::optional<size_t> unusedResultIdx;
     for (size_t i = 1, e = vtidBasis.size(); i <= e; ++i) {
       if (!usedResults.contains(i)) {
-        unusedBasisProduct *= vtidBasis[i - 1];
-        unusedLaneIdComponents.push_back(splitLaneId.getResult(i));
-        unusedBasisSizes.push_back(vtidBasis[i - 1]);
+        assert(!unusedResultIdx && "expected exactly one unused basis entry");
+        unusedResultIdx = i;
       }
     }
-    int64_t overflowSize = broadcastFactor / unusedBasisProduct;
-
-    // Collect in most-significant-first order: overflow, then basis entries.
-    if (overflowSize > 1) {
-      unusedLaneIdComponents.insert(unusedLaneIdComponents.begin(),
-                                    splitLaneId.getResult(0));
-      unusedBasisSizes.insert(unusedBasisSizes.begin(), overflowSize);
-    }
-
-    assert(llvm::product_of(unusedBasisSizes) == broadcastFactor &&
-           "unused basis sizes product must equal broadcast factor");
-
-    if (unusedLaneIdComponents.size() == 1) {
-      broadcastIndex = unusedLaneIdComponents[0];
-    } else if (unusedLaneIdComponents.size() > 1) {
-      broadcastIndex = affine::AffineLinearizeIndexOp::create(
-          builder, loc, unusedLaneIdComponents, unusedBasisSizes,
-          /*disjoint=*/true);
-    }
+    assert(unusedResultIdx && "expected one unused basis entry for broadcast");
+    assert(vtidBasis[*unusedResultIdx - 1] == broadcastFactor &&
+           "unused basis size must equal broadcast factor");
+    broadcastIndex = splitLaneId.getResult(*unusedResultIdx);
   }
 
   // Each thread grabs `element` contiguous data, so the vtid needs to be
@@ -1026,17 +1007,10 @@ static LogicalResult populateCanonicalOffsetsSizesAndStrides(
     if (splitDim && dimIdx == *splitDim) {
       // offset = vtid * element + broadcastIndex * perLaneElement.
       int64_t perLaneElement = element / broadcastFactor;
-      if (perLaneElement != 1) {
-        vtid = affine::AffineLinearizeIndexOp::create(
-            builder, loc, ValueRange{vtid, broadcastIndex, cZero},
-            ArrayRef<int64_t>{vtidLen, broadcastFactor, perLaneElement},
-            /*disjoint=*/true);
-      } else {
-        vtid = affine::AffineLinearizeIndexOp::create(
-            builder, loc, ValueRange{vtid, broadcastIndex},
-            ArrayRef<int64_t>{vtidLen, broadcastFactor},
-            /*disjoint=*/true);
-      }
+      vtid = affine::AffineLinearizeIndexOp::create(
+          builder, loc, ValueRange{vtid, broadcastIndex, cZero},
+          ArrayRef<int64_t>{vtidLen, broadcastFactor, perLaneElement},
+          /*disjoint=*/true);
     } else if (element != 1) {
       vtid = affine::AffineLinearizeIndexOp::create(
           builder, loc, ValueRange{vtid, cZero},