[Mlir-commits] [mlir] [mlir][vector] Account for subview offset in gather lowering. (PR #195359)

[Lukas Sommer]

================
@@ -134,27 +141,64 @@ struct RemoveStrideFromGatherSource : OpRewritePattern<vector::GatherOp> {
     if (stridedLayoutAttr.getStrides()[0] != srcTrailingDim)
       return failure();
 
+    // The result memref's offset is the linearized position of the subview's
+    // first element within the source memref. Bail out on dynamic offsets so
+    // we don't have to materialize them; the conditional-load fallback will
+    // still produce correct code.
+    int64_t subviewOffset = stridedLayoutAttr.getOffset();
+    if (ShapedType::isDynamic(subviewOffset))
+      return failure();
+
     // 1. Collapse the input memref so that it's "flat".
     SmallVector<ReassociationIndices> reassoc = {{0, 1}};
     Value collapsed = memref::CollapseShapeOp::create(
         rewriter, op.getLoc(), subview.getSource(), reassoc);
 
-    // 2. Generate new gather indices that will model the
-    // strided access.
+    // 2. Generate new gather indices that will model the strided access.
+    // Take `memref<4xf32, strided<[3], offset: 1>>` and lane k as an example.
+    // For the rewrite to be correct, the flat positions must match:
+    //   new_off + new_idxs[k] = 1 + (base_off + idxs[k]) * 3
+    //                         = 1 + base_off * 3 + idxs[k] * 3
+    // So the newIdxs is scaled with the stride.
     IntegerAttr stride = rewriter.getIndexAttr(srcTrailingDim);
     VectorType vType = op.getIndices().getType();
     Value mulCst = arith::ConstantOp::create(
         rewriter, op.getLoc(), vType, DenseElementsAttr::get(vType, stride));
-
     Value newIdxs =
         arith::MulIOp::create(rewriter, op.getLoc(), op.getIndices(), mulCst);
 
-    // 3. Create an updated gather op with the collapsed input memref and the
-    // updated indices.
+    // 3. Linearize the gather's base offsets through the source memref. On the
+    // collapsed memref the trailing offset must be scaled by the source's
+    // trailing dim and shifted by the subview's static offset.
+    // Pick new_idxs[k] = idxs[k] * 3 (that's step 2), and solve for new_off:
+    //   new_off = 1 + base_off * 3
+    //           = subview_offset + base_off * stride
+    SmallVector<Value> newOffsets(op.getOffsets());
+    bool trailingOffsetIsZero = isZeroInteger(newOffsets.back());
+    if (!trailingOffsetIsZero) {
+      Value strideVal =
+          arith::ConstantIndexOp::create(rewriter, op.getLoc(), srcTrailingDim);
+      newOffsets.back() = arith::MulIOp::create(rewriter, op.getLoc(),
+                                                newOffsets.back(), strideVal);
+    }
+    if (subviewOffset != 0) {
+      Value subviewOffsetValue =
+          arith::ConstantIndexOp::create(rewriter, op.getLoc(), subviewOffset);
+      if (trailingOffsetIsZero) {
+        newOffsets.back() = subviewOffsetValue;
+      } else {
+        newOffsets.back() =
+            arith::AddIOp::create(rewriter, op.getLoc(), newOffsets.back(),
+                                  subviewOffsetValue)
+                .getResult();
+      }
+    }
----------------
sommerlukas wrote:

Nit/question: Could we simplify the logic here with `createOrFold`? If the trailing offset of the operation is zero, the multiplication with the stride and the addition to the subview offset should fold.

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