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

[ofri frishman]

================
@@ -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(
+            b, sliceOp.getLoc(), offsetVal, basis, /*hasOuterBound=*/true);
+        for (auto result : delinearizeOp.getResults())
+          expandedOffsets.push_back(result);
+
+        continue;
       }
+
+      // If the size is greater than 1, we will take a more
+      // restrictive case where both the offset and the innermost
----------------
ofri-frishman wrote:

I think you could extend this so it won't only be for cases where innermost dimension is divisible by the size. Basically if there is any dimension such that the multiplication of all dimensions up to it (including it) is divisible by the size then that + the condition about offset is enough to guarantee contiguity 

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