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

Code before:

    Value localOffsets =
        computeOffsets(rewriter, scatterOp, meta.first, meta.second);
    Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);

    xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
                                  flatMemref, localOffsets, scatterOp.getMask(),
                                  /*chunk_size=*/IntegerAttr{},
                                  /*l1_hint=*/xegpu::CachePolicyAttr{},
                                  /*l2_hint=*/xegpu::CachePolicyAttr{},
                                  /*l3_hint=*/xegpu::CachePolicyAttr{},
                                  /*layout=*/nullptr);
    rewriter.eraseOp(scatterOp);
    return success();
  }
};

struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
  using Base::Base;

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

    VectorType vecTy = loadOp.getResult().getType();
    auto memTy = cast<MemRefType>(loadOp.getBase().getType());
    if (failed(storeLoadPreconditions(rewriter, loadOp, vecTy, memTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;
    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;

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

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

    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);

Review comment (suggestion):
```suggestion
    MemRefType memTy = loadOp.getBase().getType();
```
I don't think the cast is necessary

Code after:

    Value localOffsets =
        computeOffsets(rewriter, scatterOp, meta.first, meta.second);
    Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);

    xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
                                  flatMemref, localOffsets, scatterOp.getMask(),
                                  /*chunk_size=*/IntegerAttr{},
                                  /*l1_hint=*/xegpu::CachePolicyAttr{},
                                  /*l2_hint=*/xegpu::CachePolicyAttr{},
                                  /*l3_hint=*/xegpu::CachePolicyAttr{},
                                  /*layout=*/nullptr);
    rewriter.eraseOp(scatterOp);
    return success();
  }
};

struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
  using Base::Base;

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

    VectorType vecTy = loadOp.getResult().getType();
    MemRefType memTy = loadOp.getBase().getType();
    if (failed(storeLoadPreconditions(rewriter, loadOp, vecTy, memTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;
    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;

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

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

    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);