[Mlir-commits] [mlir] [MLIR][XeGPU] Extend unrolling support for scatter ops with chunk_size (PR #144447)
Charitha Saumya
llvmlistbot at llvm.org
Mon Jun 16 17:02:08 PDT 2025
================
@@ -402,30 +402,58 @@ struct UnrollCreateDescOp : public UnrollPattern<xegpu::CreateDescOp> {
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
xegpu::TensorDescType tdescTy = op.getType();
+ TypedValue<::mlir::VectorType> indiceVec = op.getOffsets();
+ VectorType indiceVecTy = indiceVec.getType();
- // check if the tensor descriptor type is a 1d vector type
- if (tdescTy.getRank() > 1)
+ if (!tdescTy.isScattered())
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape)
return failure();
- auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
-
- TypedValue<::mlir::VectorType> indiceVec = op.getOffsets();
- VectorType indiceVecTy = indiceVec.getType();
+ SmallVector<int64_t> targetIndiceShape(*targetShape);
+ int64_t originalChunkSize = tdescTy.getChunkSize();
+ // IndiceVec is 1 dim lower than tdescTy when chunkSize is larger than 1.
+ if (originalChunkSize > 1)
+ targetIndiceShape.pop_back();
+ auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
SmallVector<Type> convertedIndiceTypes =
- getUnrolledTypes(indiceVecTy, *targetShape);
+ getUnrolledTypes(indiceVecTy, targetIndiceShape);
SmallVector<Value> convertedIndiceVec =
- pack(indiceVec, convertedIndiceTypes, *targetShape, loc, rewriter);
+ pack(indiceVec, convertedIndiceTypes, targetIndiceShape, loc, rewriter);
SmallVector<Value> newOps;
- for (auto indice : convertedIndiceVec) {
- auto newOp = rewriter.create<xegpu::CreateDescOp>(loc, newTdescTy,
- op.getSource(), indice);
- newOps.push_back(newOp);
+
+ // more indices is need when chunkSize > 1. Since a big load from one
----------------
charithaintc wrote:
Capitalize the fist word in a sentence.
https://github.com/llvm/llvm-project/pull/144447
More information about the Mlir-commits
mailing list