[Mlir-commits] [mlir] [mlir][tensor]-Handle Dynamic Offset in BubbleUpSliceOpThroughCollapse (PR #178921)

[ofri frishman]

================
@@ -602,116 +767,48 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
     return failure();
   }
 
+  using ReassocGroupResult =
+      std::pair<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>>;
+  SmallVector<ReassocGroupResult> groupResults;
+
   // Compute new offsets, sizes, and strides for tensor.extract_slice.
   // The new tensor.extract_slice will work on a tensor that has has a rank
   // equal to the rank of the src of the collapse_shape. In each iteration of
   // the loop, the offsets and sizes will be computed per reassociation group.
-  expandedStrides.resize(expandedShape.size(), b.getIndexAttr(1));
   for (auto [collapsedSize, collapsedOffset, reassocIndices] :
        llvm::zip_equal(collapsedSizes, collapsedOffsets, reassociation)) {
-    // CASE #1 - size and/or offset are dynamic.
-    // In this case, the slice can be represented as a contiguous slice only
-    // if there is a single dimension in the reassociation group that has a
-    // size not equal to 1.
-    if (isa<Value>(collapsedSize) || isa<Value>(collapsedOffset)) {
-      int nonUnitSizeCount = 0;
-      for (int64_t expandedShapeIdx : reassocIndices) {
-        if (expandedShape[expandedShapeIdx] != 1) {
-          nonUnitSizeCount++;
-          expandedSizes.push_back(collapsedSize);
-          expandedOffsets.push_back(collapsedOffset);
-          continue;
-        }
-
-        expandedSizes.push_back(b.getIndexAttr(1));
-        expandedOffsets.push_back(b.getIndexAttr(0));
-      }
-
-      if (nonUnitSizeCount != 1) {
-        return failure();
-      }
-      continue;
-    }
 
-    // CASE #2 = size and offset are static.
-    // Verify that the slice can be represented as a contiguous slice of the
-    // src of the collapse_shape.
-    // Checking this is done on order of most internal dimensions first,
-    // so traversal is done in reverse order of the reassociation group.
-    // If the expected slice shape is [1, 1, ..., 1, Sk, Ak + 1, Ak + 2,
-    // ...,An] then we first find the size and offset for n...k+1 then for k
-    // and then for k-1...0.
-
-    // currentCollapsedsize and currentCollapsedOffset are initialized with
-    // the original collapsed size and offset and divided by the expanded
-    // shape size in each dimension as we go along the reassociation group.
-    // In essence we are spreading the original collapsed size and offset over
-    // the various expanded slice dimensions.
-    // The variables are used both to check the validity of the slice and to
-    // compute the expanded sizes and offsets.
-    int64_t currentCollapsedsize = getConstantIntValue(collapsedSize).value();
-    int64_t currentCollapsedOffset =
-        getConstantIntValue(collapsedOffset).value();
-    SmallVector<OpFoldResult> groupExpandedSizes, groupExpandedOffsets;
-    ReassociationIndices reversedReassocIndices(reassocIndices.rbegin(),
-                                                reassocIndices.rend());
-    int64_t idx = 0;
-    int64_t reassocGroupSize = reassocIndices.size();
-
-    // First handle the trailing dimensions where the slice size should be
-    // equal to the tensor shape and the offset should be 0 (n...k+1).
-    for (; idx < reassocGroupSize; ++idx) {
-      int64_t expandedShapeSize = expandedShape[reversedReassocIndices[idx]];
-
-      if (currentCollapsedsize < expandedShapeSize)
-        break;
-
-      // We need to make sure that the slice size can be set to the shape size
-      // and the offset to 0.
-      if ((currentCollapsedsize % expandedShapeSize) != 0 ||
-          (currentCollapsedOffset % expandedShapeSize) != 0) {
-        return failure();
-      }
-
-      groupExpandedSizes.push_back(b.getIndexAttr(expandedShapeSize));
-      groupExpandedOffsets.push_back(b.getIndexAttr(0));
+    SmallVector<OpFoldResult> groupSizes;
+    SmallVector<OpFoldResult> groupOffsets;
+    LogicalResult result = computeExpandedSliceInfoForReassocGroup(
+        b, collapsedSize, collapsedOffset, reassocIndices, expandedShape,
+        groupSizes, groupOffsets);
+    if (failed(result))
+      return failure();
+    groupResults.emplace_back(std::make_pair(groupSizes, groupOffsets));
+  }
 
-      currentCollapsedsize /= expandedShapeSize;
-      currentCollapsedOffset /= expandedShapeSize;
-    }
+  expandedStrides.resize(expandedShape.size(), b.getIndexAttr(1));
+  for (auto [groupIdx, reassocIndices] : llvm::enumerate(reassociation)) {
+    auto &[sizes, offsets] = groupResults[groupIdx];
+    expandedSizes.append(sizes);
 
-    // Now handle the first dim where slicing occurs on (k).
-    if (idx < reassocGroupSize) {
-      int64_t expandedShapeSize = expandedShape[reversedReassocIndices[idx]];
-      int64_t offsetInDim = currentCollapsedOffset % expandedShapeSize;
-      // We need to make sure that the slice size in this dim + offset will
-      // not exceed the shape size.
-      if ((currentCollapsedsize + offsetInDim) >= expandedShapeSize) {
-        return failure();
-      }
-      groupExpandedSizes.push_back(b.getIndexAttr(currentCollapsedsize));
-      groupExpandedOffsets.push_back(b.getIndexAttr(offsetInDim));
-      currentCollapsedOffset /= expandedShapeSize;
+    if (!offsets.empty()) {
----------------
ofri-frishman wrote:

it seems counter intuitive that sometimes computeExpandedSliceInfoForReassocGroup returns offsets and sometimes it doesn't. I suggest changing it so it always returns them or always doesn't.

https://github.com/llvm/llvm-project/pull/178921