Repository: llvm/llvm-project
File: mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp (C++)

Code before:

    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
      return failure();

    xegpu::DistributeLayoutAttr inputLayout = op.getInputLayoutAttr();
    xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
    if (!inputLayout || !targetLayout)
      return rewriter.notifyMatchFailure(op, "missing layout attributes.");

    if (inputLayout.getEffectiveInstDataAsInt().empty() ||
        targetLayout.getEffectiveInstDataAsInt().empty())
      return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");

    inputLayout = inputLayout.dropInstData();
    targetLayout = targetLayout.dropInstData();

    Value newSource = op.getSource();
    SmallVector<Value> newOps;
    if (inputLayout && targetLayout) {
      SmallVector<Type> convertedValTypes =
          getUnrolledTypes(valueTy, *targetShape);
      SmallVector<Value> convertedValues =
          pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
      for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
        auto newOp = xegpu::ConvertLayoutOp::create(rewriter, op.getLoc(), t, v,
                                                    inputLayout, targetLayout);
        newOps.push_back(newOp);
      }
      newSource = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
    }

    rewriter.replaceOp(op, newSource);
    return success();
  }
};

} // namespace

void mlir::xegpu::populateXeGPUUnrollPatterns(
    RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
  patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
               UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp,
               UnrollCreateDescOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
               UnrollPrefetchOp, UnrollUpdateOffsetOp, UnrollLoadMatrixOp,
               UnrollStoreMatrixOp, UnrollLoadGatherOpWithOffset,
               UnrollStoreScatterOpWithOffsets, UnrollConvertLayoutOp>(
      patterns.getContext(), options);
}


Review comment (suggestion):
```suggestion
        auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
```

Code after:
  LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getType());
    assert(valueTy && "the value type must be vector type!");

    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
      return failure();

    xegpu::DistributeLayoutAttr inputLayout = op.getInputLayoutAttr();
    xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
    if (!inputLayout || !targetLayout)
      return rewriter.notifyMatchFailure(op, "missing layout attributes.");

    if (inputLayout.getEffectiveInstDataAsInt().empty() ||
        targetLayout.getEffectiveInstDataAsInt().empty())
      return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");

    inputLayout = inputLayout.dropInstData();
    targetLayout = targetLayout.dropInstData();

    Value newSource = op.getSource();
    SmallVector<Value> newOps;
    if (inputLayout && targetLayout) {
      SmallVector<Type> convertedValTypes =
          getUnrolledTypes(valueTy, *targetShape);
      SmallVector<Value> convertedValues =
          pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
      for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
        auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
                                                    inputLayout, targetLayout);
        newOps.push_back(newOp);
      }
      newSource = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
    }

    rewriter.replaceOp(op, newSource);
    return success();
  }
};

} // namespace

void mlir::xegpu::populateXeGPUUnrollPatterns(
    RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
  patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
               UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp,
               UnrollCreateDescOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
               UnrollPrefetchOp, UnrollUpdateOffsetOp, UnrollLoadMatrixOp,
               UnrollStoreMatrixOp, UnrollLoadGatherOpWithOffset,