better MaterializeMergeSchedule.cpp (something still broken downstream)

2026-05-22 06:56:39 +02:00
parent 6aaf1c0870
commit 43ed3914b8
13 changed files with 1433 additions and 1620 deletions
@@ -120,6 +120,15 @@ static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
  if (value == laneArg || isConstantIndexLike(value))
    return true;

+  auto extractOp = value.getDefiningOp<tensor::ExtractOp>();
+  if (extractOp) {
+    auto constantTensor = extractOp.getTensor().getDefiningOp<arith::ConstantOp>();
+    auto denseAttr = constantTensor ? dyn_cast<DenseIntElementsAttr>(constantTensor.getValue()) : nullptr;
+    if (!denseAttr || denseAttr.getType().getRank() != 1 || extractOp.getIndices().size() != 1)
+      return false;
+    return isSupportedLaneOffsetExpr(extractOp.getIndices().front(), laneArg);
+  }
+
  auto addOp = value.getDefiningOp<arith::AddIOp>();
  if (!addOp)
    return false;