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

[Amir Bishara]

================
@@ -610,26 +628,103 @@ LogicalResult mlir::tensor::getExpandedExtractSliceInfo(
   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)) {
+      // When the size is dynamic, We will handle a simple case where there
+      // is a single dimension in the reassociation group that has a size
+      // not equal to 1, which guarantees it is a contiguous slice.
       int nonUnitSizeCount = 0;
+      SmallVector<OpFoldResult> tentativeSizes, tentativeOffsets;
       for (int64_t expandedShapeIdx : reassocIndices) {
         if (expandedShape[expandedShapeIdx] != 1) {
           nonUnitSizeCount++;
-          expandedSizes.push_back(collapsedSize);
-          expandedOffsets.push_back(collapsedOffset);
+          tentativeSizes.push_back(collapsedSize);
+          tentativeOffsets.push_back(collapsedOffset);
           continue;
         }
 
-        expandedSizes.push_back(b.getIndexAttr(1));
-        expandedOffsets.push_back(b.getIndexAttr(0));
+        tentativeSizes.push_back(b.getIndexAttr(1));
+        tentativeOffsets.push_back(b.getIndexAttr(0));
+      }
+
+      if (nonUnitSizeCount == 1) {
+        // Only one non-unit dimension, slice is contiguous.
+        expandedSizes.append(tentativeSizes);
+        expandedOffsets.append(tentativeOffsets);
+        continue;
       }
 
-      if (nonUnitSizeCount != 1) {
+      std::optional<int64_t> staticSize = getConstantIntValue(collapsedSize);
+      if (!staticSize.has_value())
         return failure();
+
+      // If the size is statically 1, it means we're extracting a
+      // single element. In this case, we can always handle this by
+      // delinearizing the dynamic offset and get a contiguous slice.
+      if (staticSize.value() == 1) {
+        SmallVector<int64_t> basis;
+        for (int64_t expandedShapeIdx : reassocIndices) {
+          expandedSizes.push_back(b.getIndexAttr(1));
+          basis.push_back(expandedShape[expandedShapeIdx]);
+        }
+
+        Value offsetVal = getValueOrCreateConstantIndexOp(b, sliceOp.getLoc(),
+                                                          collapsedOffset);
+        auto delinearizeOp = affine::AffineDelinearizeIndexOp::create(
----------------
amirBish wrote:

Right, fixed Thanks.

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