Repository: llvm/llvm-project
File: mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (C++)

Code before:

    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
    auto loadNdOp =
        xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc, indices,
                                /*packed=*/nullptr, /*transpose=*/nullptr,
                                /*l1_hint=*/hint,
                                /*l2_hint=*/hint, /*l3_hint=*/hint,
                                /*layout=*/nullptr);
    rewriter.replaceOp(loadOp, loadNdOp);

    return success();
  }
};

struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(vector::StoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = storeOp.getLoc();

    TypedValue<VectorType> vector = storeOp.getValueToStore();
    VectorType vecTy = vector.getType();
    auto storeMemTy = cast<MemRefType>(storeOp.getBase().getType());
    if (failed(storeLoadPreconditions(rewriter, storeOp, vecTy, storeMemTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;

    auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
        rewriter, loc, storeOp.getBase(),
        getAsOpFoldResult(storeOp.getIndices()), vecTy.getRank());

    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(),
        /*array_length=*/1, boundaryCheck, xegpu::MemorySpace::Global);

    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));

    auto storeNdOp =
        xegpu::StoreNdOp::create(rewriter, loc, vector, ndDesc, indices,
                                 /*l1_hint=*/hint,
                                 /*l2_hint=*/hint, /*l3_hint=*/hint,
                                 /*layout=*/nullptr);

    rewriter.replaceOp(storeOp, storeNdOp);

Review comment (bug):
```suggestion
    MemRefType memTy = storeOp.getBase().getType();
    if (failed(storeLoadPreconditions(rewriter, storeOp, vecTy, memTy)))
```

same here

Code after:

    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
    auto loadNdOp =
        xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc, indices,
                                /*packed=*/nullptr, /*transpose=*/nullptr,
                                /*l1_hint=*/hint,
                                /*l2_hint=*/hint, /*l3_hint=*/hint,
                                /*layout=*/nullptr);
    rewriter.replaceOp(loadOp, loadNdOp);

    return success();
  }
};

struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(vector::StoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = storeOp.getLoc();

    TypedValue<VectorType> vector = storeOp.getValueToStore();
    VectorType vecTy = vector.getType();
    MemRefType memTy = storeOp.getBase().getType();
    if (failed(storeLoadPreconditions(rewriter, storeOp, vecTy, memTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;

    auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
        rewriter, loc, storeOp.getBase(),
        getAsOpFoldResult(storeOp.getIndices()), vecTy.getRank());

    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(),
        /*array_length=*/1, boundaryCheck, xegpu::MemorySpace::Global);

    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));

    auto storeNdOp =
        xegpu::StoreNdOp::create(rewriter, loc, vector, ndDesc, indices,
                                 /*l1_hint=*/hint,
                                 /*l2_hint=*/hint, /*l3_hint=*/hint,
                                 /*layout=*/nullptr);

    rewriter.replaceOp(storeOp, storeNdOp);