[Mlir-commits] [mlir] [MLIR][Conversion][XeGPU][XeVM] Add XeGPUToXeVM conversion pass and tests. (PR #154556)
Sang Ik Lee
llvmlistbot at llvm.org
Wed Aug 20 11:00:03 PDT 2025
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/154556
>From b860103f876016074f35dd6f6401c6b58d31f6ef Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 14 Jul 2025 18:54:41 +0000
Subject: [PATCH 1/4] Add XeGPUToXeVM conversion pass and tests.
---
mlir/include/mlir/Conversion/Passes.h | 1 +
mlir/include/mlir/Conversion/Passes.td | 12 +
.../mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h | 27 +
mlir/lib/Conversion/CMakeLists.txt | 1 +
.../lib/Conversion/XeGPUToXeVM/CMakeLists.txt | 25 +
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 932 ++++++++++++++++++
.../XeGPUToXeVM/create_nd_tdesc.mlir | 48 +
mlir/test/Conversion/XeGPUToXeVM/dpas.mlir | 17 +
mlir/test/Conversion/XeGPUToXeVM/fence.mlir | 15 +
.../Conversion/XeGPUToXeVM/loadstore_nd.mlir | 71 ++
.../XeGPUToXeVM/loadstoreprefetch.mlir | 357 +++++++
.../Conversion/XeGPUToXeVM/prefetch_nd.mlir | 40 +
.../Conversion/XeGPUToXeVM/update_offset.mlir | 25 +
13 files changed, 1571 insertions(+)
create mode 100644 mlir/include/mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h
create mode 100644 mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
create mode 100644 mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/dpas.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/fence.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/prefetch_nd.mlir
create mode 100644 mlir/test/Conversion/XeGPUToXeVM/update_offset.mlir
diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index 91b2ecf8922a3..da061b269daf7 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -82,6 +82,7 @@
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
#include "mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h"
#include "mlir/Conversion/VectorToXeGPU/VectorToXeGPU.h"
+#include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
#include "mlir/Conversion/XeVMToLLVM/XeVMToLLVM.h"
namespace mlir {
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 2058aba7f9e37..323af3e97e2d4 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1555,4 +1555,16 @@ def ConvertXeVMToLLVMPass : Pass<"convert-xevm-to-llvm"> {
let dependentDialects = ["LLVM::LLVMDialect"];
}
+//===----------------------------------------------------------------------===//
+// XeGPUToXeVM
+//===----------------------------------------------------------------------===//
+
+def ConvertXeGPUToXeVMPass : Pass<"convert-xegpu-to-xevm"> {
+ let summary = "Convert XeGPU to XeVM dialect";
+ let dependentDialects = ["xevm::XeVMDialect", "vector::VectorDialect",
+ "memref::MemRefDialect", "arith::ArithDialect",
+ "LLVM::LLVMDialect", "index::IndexDialect",
+ "gpu::GPUDialect", "scf::SCFDialect"];
+}
+
#endif // MLIR_CONVERSION_PASSES
diff --git a/mlir/include/mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h b/mlir/include/mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h
new file mode 100644
index 0000000000000..fb23d24b0161b
--- /dev/null
+++ b/mlir/include/mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h
@@ -0,0 +1,27 @@
+//===-- XeGPUToXeVM.h - Convert XeGPU to XeVM dialect ---------_--*- C++-*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_CONVERSION_XEGPUTOXEVM_XEGPUTOXEVMPASS_H_
+#define MLIR_CONVERSION_XEGPUTOXEVM_XEGPUTOXEVMPASS_H_
+
+#include <memory>
+
+namespace mlir {
+class DialectRegistry;
+class LLVMTypeConverter;
+class RewritePatternSet;
+class Pass;
+
+#define GEN_PASS_DECL_CONVERTXEGPUTOXEVMPASS
+#include "mlir/Conversion/Passes.h.inc"
+
+void populateXeGPUToXeVMConversionPatterns(
+ mlir::RewritePatternSet &patterns, mlir::LLVMTypeConverter &typeConverter);
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_XEGPUTOXEVM_XEGPUTOXEVMPASS_H_
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
index 171f7169fd41d..134fe8e14ca38 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -76,3 +76,4 @@ add_subdirectory(VectorToSCF)
add_subdirectory(VectorToSPIRV)
add_subdirectory(VectorToXeGPU)
add_subdirectory(XeVMToLLVM)
+add_subdirectory(XeGPUToXeVM)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
new file mode 100644
index 0000000000000..ed54b0bb5ee81
--- /dev/null
+++ b/mlir/lib/Conversion/XeGPUToXeVM/CMakeLists.txt
@@ -0,0 +1,25 @@
+add_mlir_conversion_library(MLIRXeGPUToXeVM
+ XeGPUToXeVM.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/XeGPUToXeVM
+
+ DEPENDS
+ MLIRConversionPassIncGen
+
+ LINK_COMPONENTS
+ Core
+
+ LINK_LIBS PUBLIC
+ MLIRFuncDialect
+ MLIRGPUDialect
+ MLIRLLVMCommonConversion
+ MLIRLLVMDialect
+ MLIRXeVMDialect
+ MLIRVectorDialect
+ MLIRArithDialect
+ MLIRIndexDialect
+ MLIRXeGPUDialect
+ MLIRPass
+ MLIRTransforms
+)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
new file mode 100644
index 0000000000000..380409afbc62e
--- /dev/null
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -0,0 +1,932 @@
+//===-- XeVMToLLVM.cpp - XeVM to LLVM dialect conversion --------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
+#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
+
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/Index/IR/IndexDialect.h"
+#include "mlir/Dialect/Index/IR/IndexOps.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Transforms/Patterns.h"
+#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/Support/FormatVariadic.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Types.h"
+
+#include "llvm/ADT/TypeSwitch.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTXEGPUTOXEVMPASS
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+enum class NdDescI32Layout : uint32_t {
+ BasePtr = 0,
+ BaseShapeW = 2,
+ BaseShapeH = 3,
+ TensorOffsetW = 4,
+ TensorOffsetH = 5
+};
+
+static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
+ switch (xeGpuMemspace) {
+ case xegpu::MemorySpace::Global:
+ return static_cast<int>(xevm::AddrSpace::GLOBAL);
+ case xegpu::MemorySpace::SLM:
+ return static_cast<int>(xevm::AddrSpace::SHARED);
+ }
+ llvm_unreachable("Unknown XeGPU memory space.");
+}
+
+template <typename T>
+std::tuple<bool, int32_t, int32_t> checkAllLinear(SmallVector<T> denseAttr) {
+ assert(!denseAttr.empty());
+ const int32_t intercept{static_cast<int32_t>(denseAttr[0])};
+ if (denseAttr.size() < 2)
+ return {true, 0, intercept};
+ const T slope{denseAttr[1] - denseAttr[0]};
+ for (size_t i = 1; i < denseAttr.size(); ++i)
+ if (denseAttr[i] - denseAttr[i - 1] != slope)
+ return {false, 0, 0};
+ return {true, static_cast<int32_t>(slope), intercept};
+}
+
+VectorType encodeVectorTypeTo(VectorType currentVecType, Type toElemType) {
+ auto elemType = currentVecType.getElementType();
+ auto currentBitWidth = elemType.getIntOrFloatBitWidth();
+ auto newBitWidth = toElemType.getIntOrFloatBitWidth();
+ const int size =
+ currentVecType.getNumElements() * currentBitWidth / newBitWidth;
+ return VectorType::get(size, toElemType);
+}
+
+xevm::LoadCacheControl
+translateLoadXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
+ std::optional<xegpu::CachePolicy> L3hint) {
+ auto L1hintVal =
+ L1hint.has_value() ? L1hint.value() : xegpu::CachePolicy::UNCACHED;
+ auto L3hintVal =
+ L3hint.has_value() ? L3hint.value() : xegpu::CachePolicy::UNCACHED;
+ switch (L1hintVal) {
+ case xegpu::CachePolicy::CACHED:
+ if (L3hintVal == xegpu::CachePolicy::CACHED)
+ return xevm::LoadCacheControl::L1C_L2UC_L3C;
+ else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::LoadCacheControl::L1C_L2UC_L3UC;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::UNCACHED:
+ if (L3hintVal == xegpu::CachePolicy::CACHED)
+ return xevm::LoadCacheControl::L1UC_L2UC_L3C;
+ else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::LoadCacheControl::L1UC_L2UC_L3UC;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::STREAMING:
+ if (L3hintVal == xegpu::CachePolicy::CACHED)
+ return xevm::LoadCacheControl::L1S_L2UC_L3C;
+ else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::LoadCacheControl::L1S_L2UC_L3UC;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::READ_INVALIDATE:
+ return xevm::LoadCacheControl::INVALIDATE_READ;
+ default:
+ llvm_unreachable("Unsupported cache control.");
+ }
+}
+
+xevm::StoreCacheControl
+translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
+ std::optional<xegpu::CachePolicy> L3hint) {
+ auto L1hintVal =
+ L1hint.has_value() ? L1hint.value() : xegpu::CachePolicy::UNCACHED;
+ auto L3hintVal =
+ L3hint.has_value() ? L3hint.value() : xegpu::CachePolicy::UNCACHED;
+ switch (L1hintVal) {
+ case xegpu::CachePolicy::UNCACHED:
+ if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::StoreCacheControl::L1UC_L2UC_L3UC;
+ else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
+ return xevm::StoreCacheControl::L1UC_L2UC_L3WB;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::STREAMING:
+ if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::StoreCacheControl::L1S_L2UC_L3UC;
+ else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
+ return xevm::StoreCacheControl::L1S_L2UC_L3WB;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::WRITE_BACK:
+ if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::StoreCacheControl::L1WB_L2UC_L3UC;
+ else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
+ return xevm::StoreCacheControl::L1WB_L2UC_L3WB;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ case xegpu::CachePolicy::WRITE_THROUGH:
+ if (L3hintVal == xegpu::CachePolicy::UNCACHED)
+ return xevm::StoreCacheControl::L1WT_L2UC_L3UC;
+ else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
+ return xevm::StoreCacheControl::L1WT_L2UC_L3WB;
+ else
+ llvm_unreachable("Unsupported cache control.");
+ default:
+ llvm_unreachable("Unsupported cache control.");
+ }
+}
+
+class CreateNdDescToXeVMPattern
+ : public OpConversionPattern<xegpu::CreateNdDescOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::CreateNdDescOp op,
+ xegpu::CreateNdDescOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto source = op.getSource();
+ Type payloadElemTy = rewriter.getI32Type();
+ Type i64Ty = rewriter.getI64Type();
+ VectorType payloadTy = VectorType::get(8, payloadElemTy);
+ VectorType payloadI64Ty = VectorType::get(4, i64Ty);
+ Value payload = arith::ConstantOp::create(
+ rewriter, loc,
+ DenseElementsAttr::get(payloadTy, IntegerAttr::get(payloadElemTy, 0)));
+
+ Value baseAddr;
+ Value baseShapeW;
+ Value baseShapeH;
+ Value offsetW;
+ Value offsetH;
+
+ bool sourceIsMemref = false;
+ auto sourceTy = source.getType();
+ int64_t rank;
+ if (isa<MemRefType>(sourceTy)) {
+ sourceIsMemref = true;
+ baseAddr =
+ memref::ExtractAlignedPointerAsIndexOp::create(rewriter, loc, source);
+ auto sourceMemrefTy = cast<MemRefType>(sourceTy);
+ if (!sourceMemrefTy.hasStaticShape()) {
+ op.emitError() << "Expected static memref shape.";
+ return failure();
+ }
+ rank = sourceMemrefTy.getRank();
+ if (rank != 2) {
+ op.emitError() << "Expected a 2D memref.";
+ return failure();
+ }
+ } else if (sourceTy == rewriter.getIntegerType(64, false)) {
+ rank = op.getMixedSizes().size();
+ } else {
+ op.emitError() << "Expected source to be a 2D memref or ui64.";
+ return failure();
+ }
+ auto createOffset = [&](unsigned idx) -> Value {
+ Value val;
+ OpFoldResult ofr = op.getMixedOffsets()[idx];
+ if (auto v = llvm::dyn_cast_if_present<Value>(ofr)) {
+ val = arith::IndexCastOp::create(rewriter, loc, i64Ty, v);
+ val = arith::TruncIOp::create(rewriter, loc, payloadElemTy, val);
+ } else {
+ int32_t off = llvm::cast<IntegerAttr>(cast<Attribute>(ofr)).getInt();
+ val = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, off);
+ }
+ return val;
+ };
+ auto offsets = op.getMixedOffsets();
+ if (offsets.size() == 2) {
+ offsetW = createOffset(rank - 1);
+ offsetH = createOffset(rank - 2);
+ } else {
+ offsetW = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, 0);
+ offsetH = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, 0);
+ }
+ auto createShape = [&](unsigned idx) -> Value {
+ Value val;
+ OpFoldResult ofr = op.getMixedSizes()[idx];
+ if (auto v = llvm::dyn_cast_if_present<Value>(ofr)) {
+ val = arith::IndexCastOp::create(rewriter, loc, i64Ty, v);
+ val = arith::TruncIOp::create(rewriter, loc, payloadElemTy, val);
+ } else {
+ int32_t off = llvm::cast<IntegerAttr>(cast<Attribute>(ofr)).getInt();
+ val = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, off);
+ }
+ return val;
+ };
+ if (sourceIsMemref) {
+ auto sourceMemrefTy = cast<MemRefType>(sourceTy);
+ baseShapeW = arith::ConstantIntOp::create(
+ rewriter, loc, payloadElemTy, sourceMemrefTy.getDimSize(rank - 1));
+ baseShapeH = arith::ConstantIntOp::create(
+ rewriter, loc, payloadElemTy, sourceMemrefTy.getDimSize(rank - 2));
+ baseAddr = arith::IndexCastUIOp::create(rewriter, loc, i64Ty, baseAddr);
+ } else {
+ baseShapeW = createShape(rank - 1);
+ baseShapeH = createShape(rank - 2);
+ baseAddr = adaptor.getSource();
+ }
+ Value payLoadAsI64 =
+ vector::BitCastOp::create(rewriter, loc, payloadI64Ty, payload);
+ payLoadAsI64 =
+ vector::InsertOp::create(rewriter, loc, baseAddr, payLoadAsI64,
+ static_cast<int>(NdDescI32Layout::BasePtr));
+ payload = vector::BitCastOp::create(rewriter, loc, payloadTy, payLoadAsI64);
+ payload =
+ vector::InsertOp::create(rewriter, loc, baseShapeW, payload,
+ static_cast<int>(NdDescI32Layout::BaseShapeW));
+ payload =
+ vector::InsertOp::create(rewriter, loc, baseShapeH, payload,
+ static_cast<int>(NdDescI32Layout::BaseShapeH));
+ payload = vector::InsertOp::create(
+ rewriter, loc, offsetW, payload,
+ static_cast<int>(NdDescI32Layout::TensorOffsetW));
+ payload = vector::InsertOp::create(
+ rewriter, loc, offsetH, payload,
+ static_cast<int>(NdDescI32Layout::TensorOffsetH));
+ rewriter.replaceOp(op, payload);
+ return success();
+ }
+};
+
+class UpdateNdOffsetToXeVMPattern
+ : public OpConversionPattern<xegpu::UpdateNdOffsetOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::UpdateNdOffsetOp op,
+ xegpu::UpdateNdOffsetOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto offsets = op.getOffsets();
+ auto tdesc = adaptor.getTensorDesc();
+ for (size_t offsetDim = 0; offsetDim < offsets.size(); offsetDim++) {
+ auto offset = offsets[offsetDim];
+ if (auto cst =
+ dyn_cast_if_present<arith::ConstantOp>(offset.getDefiningOp()))
+ if (auto attr = dyn_cast_if_present<IntegerAttr>(cst.getValue());
+ attr && !attr.getInt())
+ continue;
+ const int offsetPos =
+ static_cast<int>(offsetDim ? NdDescI32Layout::TensorOffsetW
+ : NdDescI32Layout::TensorOffsetH);
+ auto oldOffset =
+ vector::ExtractOp::create(rewriter, loc, tdesc, offsetPos);
+ offset = arith::IndexCastUIOp::create(rewriter, loc,
+ rewriter.getI32Type(), offset);
+ auto newOffset = arith::AddIOp::create(rewriter, loc, oldOffset, offset);
+ tdesc =
+ vector::InsertOp::create(rewriter, loc, newOffset, tdesc, offsetPos);
+ }
+ rewriter.replaceOp(op, tdesc);
+ return success();
+ }
+};
+
+template <
+ typename OpType,
+ typename = std::enable_if_t<llvm::is_one_of<
+ OpType, xegpu::LoadNdOp, xegpu::StoreNdOp, xegpu::PrefetchNdOp>::value>>
+class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
+ using OpConversionPattern<OpType>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ctxt = rewriter.getContext();
+
+ auto tdesc = adaptor.getTensorDesc();
+ auto tdescTy = op.getTensorDescType();
+ if (tdescTy.getRank() != 2) {
+ return rewriter.notifyMatchFailure(op, "Expected 2D tensor descriptor.");
+ }
+
+ VectorType payloadI64Ty = VectorType::get(4, rewriter.getI64Type());
+ Value payLoadAsI64 =
+ vector::BitCastOp::create(rewriter, loc, payloadI64Ty, tdesc);
+ Value basePtr =
+ vector::ExtractOp::create(rewriter, loc, payLoadAsI64,
+ static_cast<int>(NdDescI32Layout::BasePtr));
+ Value baseShapeW = vector::ExtractOp::create(
+ rewriter, loc, tdesc, static_cast<int>(NdDescI32Layout::BaseShapeW));
+ Value baseShapeH = vector::ExtractOp::create(
+ rewriter, loc, tdesc, static_cast<int>(NdDescI32Layout::BaseShapeH));
+ // Offsets can come from three sources:
+ // 1. Constant offsets, which are provided by the op.
+ // 2. Offsets as operands, which are provided by the op.
+ // 3. Offsets extracted from the tensor descriptor.
+ Value offsetW;
+ Value offsetH;
+ auto cOffsets = op.getConstOffsets();
+ auto offsets = op.getOffsets();
+ if (cOffsets) {
+ offsetW = arith::ConstantIntOp::create(
+ rewriter, loc, rewriter.getI32Type(), (*cOffsets)[0]);
+ offsetH = arith::ConstantIntOp::create(
+ rewriter, loc, rewriter.getI32Type(), (*cOffsets)[1]);
+ } else if (offsets.size() != 0) {
+ // offsets are provided as operands
+ if (offsets[0].getType() != rewriter.getI32Type()) {
+ if (offsets[0].getType() != rewriter.getIndexType()) {
+ return rewriter.notifyMatchFailure(
+ op, "Expected offsets to be of type i32 or index.");
+ }
+ offsetW = arith::IndexCastUIOp::create(
+ rewriter, loc, rewriter.getI32Type(), offsets[0]);
+ } else {
+ offsetW = offsets[0];
+ }
+ if (offsets[1].getType() != rewriter.getI32Type()) {
+ if (offsets[1].getType() != rewriter.getIndexType()) {
+ return rewriter.notifyMatchFailure(
+ op, "Expected offsets to be of type i32 or index.");
+ }
+ offsetH = arith::IndexCastUIOp::create(
+ rewriter, loc, rewriter.getI32Type(), offsets[1]);
+ } else {
+ offsetH = offsets[1];
+ }
+ } else {
+ // If offsets are not available, we need to extract them from the tensor
+ // descriptor.
+ offsetW = vector::ExtractOp::create(
+ rewriter, loc, tdesc,
+ static_cast<int>(NdDescI32Layout::TensorOffsetW));
+ offsetH = vector::ExtractOp::create(
+ rewriter, loc, tdesc,
+ static_cast<int>(NdDescI32Layout::TensorOffsetH));
+ }
+ auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
+ ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
+ Value basePtrLLVM =
+ LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtr);
+ auto elemType = tdescTy.getElementType();
+ auto elemBitSize = elemType.getIntOrFloatBitWidth();
+ // auto elemBitSizeAttr = rewriter.getIntegerAttr(rewriter.getI32Type(),
+ // elemBitSize);
+ Value elemByteSize = arith::ConstantIntOp::create(
+ rewriter, loc, rewriter.getI32Type(), elemBitSize / 8);
+ Value surfaceW =
+ arith::MulIOp::create(rewriter, loc, baseShapeW, elemByteSize);
+
+ auto tileW = tdescTy.getDimSize(1);
+ auto tileH = tdescTy.getDimSize(0);
+ int32_t vblocks = tdescTy.getArrayLength();
+ if constexpr (std::is_same_v<OpType, xegpu::StoreNdOp>) {
+ VectorType srcVecTy = cast<VectorType>(op.getValue().getType());
+ auto storeCacheControl =
+ translateStoreXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
+ VectorType srcFlatVecTy =
+ VectorType::get(srcVecTy.getNumElements(), srcVecTy.getElementType());
+ Value srcFlatVec = op.getValue();
+ srcFlatVecTy = encodeVectorTypeTo(srcFlatVecTy,
+ rewriter.getIntegerType(elemBitSize));
+ srcFlatVec =
+ vector::BitCastOp::create(rewriter, loc, srcFlatVecTy, srcFlatVec);
+ xevm::BlockStore2dOp::create(
+ rewriter, loc, basePtrLLVM, surfaceW, baseShapeH, surfaceW, offsetW,
+ offsetH, elemBitSize, tileW, tileH, srcFlatVec,
+ xevm::StoreCacheControlAttr::get(ctxt, storeCacheControl));
+ rewriter.eraseOp(op);
+ } else {
+ auto loadCacheControl =
+ translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
+ if constexpr (std::is_same_v<OpType, xegpu::PrefetchNdOp>) {
+ xevm::BlockPrefetch2dOp::create(
+ rewriter, loc, basePtrLLVM, surfaceW, baseShapeH, surfaceW, offsetW,
+ offsetH, elemBitSize, tileW, tileH, vblocks,
+ xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
+ rewriter.eraseOp(op);
+ } else {
+ VectorType dstVecTy = cast<VectorType>(op.getValue().getType());
+ const bool vnni = op.getPacked().value_or(false);
+ auto transposeValue = op.getTranspose();
+ bool transpose =
+ transposeValue.has_value() && transposeValue.value()[0] == 1;
+ VectorType loadedTy = encodeVectorTypeTo(
+ dstVecTy, vnni ? rewriter.getI32Type()
+ : rewriter.getIntegerType(elemBitSize));
+
+ Value resultFlatVec = xevm::BlockLoad2dOp::create(
+ rewriter, loc, loadedTy, basePtrLLVM, surfaceW, baseShapeH,
+ surfaceW, offsetW, offsetH, elemBitSize, tileW, tileH, vblocks,
+ transpose, vnni,
+ xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
+ resultFlatVec = vector::BitCastOp::create(
+ rewriter, loc,
+ encodeVectorTypeTo(loadedTy, dstVecTy.getElementType()),
+ resultFlatVec);
+ rewriter.replaceOp(op, resultFlatVec);
+ }
+ }
+ return success();
+ }
+};
+
+template <
+ typename OpType,
+ typename = std::enable_if_t<llvm::is_one_of<
+ OpType, xegpu::LoadGatherOp, xegpu::StoreScatterOp, xegpu::CreateDescOp,
+ xegpu::UpdateOffsetOp, xegpu::PrefetchOp>::value>>
+int64_t getElemByteSize(OpType op) {
+ // Get the element byte size from the tensor descriptor.
+ auto elemBitWidth =
+ op.getTensorDesc().getType().getElementType().getIntOrFloatBitWidth();
+ return elemBitWidth / 8;
+}
+
+// Add a builder that creates
+// offset * elemByteSize + baseAddr
+auto addOffset = [](ConversionPatternRewriter &rewriter, Location loc,
+ Value baseAddr, Value offset,
+ int64_t elemByteSize) -> Value {
+ Value byteSize = arith::ConstantIntOp::create(
+ rewriter, loc, rewriter.getI64Type(), elemByteSize);
+ Value byteOffset = arith::MulIOp::create(rewriter, loc, offset, byteSize);
+ Value newAddr = arith::AddIOp::create(rewriter, loc, baseAddr, byteOffset);
+ return newAddr;
+};
+
+class CreateDescToXeVMPattern
+ : public OpConversionPattern<xegpu::CreateDescOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::CreateDescOp op, xegpu::CreateDescOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto offsets = adaptor.getOffsets();
+ // Source type can be a 1D memref or ui64
+ // Using "op" instead of "adaptor" since we want to access memref type
+ // instead of LLVM struct type.
+ auto memrefTy = dyn_cast<MemRefType>(op.getSource().getType());
+ Value subGroupAddr;
+ if (memrefTy) {
+ subGroupAddr = memref::ExtractAlignedPointerAsIndexOp::create(
+ rewriter, loc, op.getSource());
+ subGroupAddr = arith::IndexCastUIOp::create(
+ rewriter, loc, rewriter.getI64Type(), subGroupAddr);
+ } else {
+ subGroupAddr = adaptor.getSource();
+ }
+ auto laneAddr =
+ addOffset(rewriter, loc, subGroupAddr, offsets, getElemByteSize(op));
+ rewriter.replaceOp(op, laneAddr);
+ return success();
+ }
+};
+
+class UpdateOffsetToXeVMPattern
+ : public OpConversionPattern<xegpu::UpdateOffsetOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::UpdateOffsetOp op,
+ xegpu::UpdateOffsetOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ Value newOffsetForLane =
+ addOffset(rewriter, loc, adaptor.getTensorDesc(), adaptor.getOffsets(),
+ getElemByteSize(op));
+ rewriter.replaceOp(op, newOffsetForLane);
+ return success();
+ }
+};
+
+template <typename OpType,
+ typename = std::enable_if_t<llvm::is_one_of<
+ OpType, xegpu::LoadGatherOp, xegpu::StoreScatterOp>::value>>
+class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
+ using OpConversionPattern<OpType>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ctxt = rewriter.getContext();
+ auto tdescTy = op.getTensorDescType();
+ auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
+ ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
+ Value basePtrI64;
+ if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
+ basePtrI64 = adaptor.getSource();
+ } else {
+ basePtrI64 = adaptor.getDest();
+ }
+ Value offsets = adaptor.getOffsets();
+ Value mask = adaptor.getMask();
+ if (offsets) {
+ VectorType offsetsVecTy = dyn_cast<VectorType>(offsets.getType());
+ if (offsetsVecTy) {
+ // Offset needs be scalar.
+ return rewriter.notifyMatchFailure(op,
+ "Expected offsets to be a scalar.");
+ } else {
+ basePtrI64 =
+ addOffset(rewriter, loc, basePtrI64, offsets, getElemByteSize(op));
+ }
+ }
+ Value basePtrLLVM =
+ LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
+ VectorType srcOrDstVecTy = op.getValueType();
+ VectorType srcOrDstFlatVecTy = VectorType::get(
+ srcOrDstVecTy.getNumElements(), srcOrDstVecTy.getElementType());
+ Value maskForLane;
+ VectorType maskVecTy = dyn_cast<VectorType>(mask.getType());
+ if (maskVecTy) {
+ return rewriter.notifyMatchFailure(op, "Expected mask to be a scalar.");
+ } else
+ maskForLane = mask;
+ if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
+ scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, {srcOrDstVecTy},
+ maskForLane, true, true);
+ rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
+ Value loaded =
+ LLVM::LoadOp::create(rewriter, loc, srcOrDstFlatVecTy, basePtrLLVM);
+ loaded.getDefiningOp()->setAttr("cache_control",
+ xevm::LoadCacheControlAttr::get(
+ ctxt, translateLoadXeGPUCacheHint(
+ op.getL1Hint(), op.getL3Hint())));
+ if (srcOrDstVecTy != srcOrDstFlatVecTy) {
+ loaded =
+ vector::ShapeCastOp::create(rewriter, loc, srcOrDstVecTy, loaded);
+ }
+ scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
+ rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+ // If mask is false, we yield a vector of zeros.
+ auto eTy = srcOrDstVecTy.getElementType();
+ loaded = arith::ConstantOp::create(
+ rewriter, loc,
+ eTy.isFloat()
+ ? DenseElementsAttr::get(srcOrDstVecTy, FloatAttr::get(eTy, 0.0))
+ : DenseElementsAttr::get(srcOrDstVecTy,
+ IntegerAttr::get(eTy, 0)));
+ scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
+ rewriter.replaceOp(op, ifOp.getResult(0));
+ } else {
+ scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, maskForLane, false);
+ auto body = ifOp.getBody();
+ rewriter.setInsertionPointToStart(body);
+ VectorType valTy = op.getValue().getType();
+ Value srcFlatVec = op.getValue();
+ if (valTy != srcOrDstFlatVecTy) {
+ srcFlatVec = vector::ShapeCastOp::create(rewriter, loc,
+ srcOrDstFlatVecTy, srcFlatVec);
+ }
+ auto storeOp = LLVM::StoreOp::create(rewriter, loc, srcFlatVec, basePtrLLVM);
+ storeOp.getOperation()->setAttr(
+ "cache_control",
+ xevm::StoreCacheControlAttr::get(ctxt,
+ translateStoreXeGPUCacheHint(
+ op.getL1Hint(), op.getL3Hint())));
+ rewriter.eraseOp(op);
+ }
+ return success();
+ }
+};
+
+class PrefetchToXeVMPattern : public OpConversionPattern<xegpu::PrefetchOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::PrefetchOp op, xegpu::PrefetchOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ctxt = rewriter.getContext();
+ auto tdescTy = op.getTensorDescType();
+ auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
+ ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
+ Value basePtrI64 = adaptor.getSource();
+ Value offsets = adaptor.getOffsets();
+ if (offsets) {
+ VectorType offsetsVecTy = dyn_cast<VectorType>(offsets.getType());
+ if (offsetsVecTy) {
+ // Offset needs be scalar.
+ return rewriter.notifyMatchFailure(op,
+ "Expected offsets to be a scalar.");
+ } else {
+ basePtrI64 =
+ addOffset(rewriter, loc, basePtrI64, offsets, getElemByteSize(op));
+ }
+ }
+ Value ptrLLVM =
+ LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
+ xevm::PrefetchOp::create(
+ rewriter, loc, ptrLLVM,
+ xevm::LoadCacheControlAttr::get(
+ ctxt, translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint())));
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+class FenceToXeVMPattern : public OpConversionPattern<xegpu::FenceOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::FenceOp op, xegpu::FenceOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ xevm::MemScope memScope{xevm::MemScope::WORKGROUP};
+ switch (op.getFenceScope()) {
+ case xegpu::FenceScope::Workgroup:
+ memScope = xevm::MemScope::WORKGROUP;
+ break;
+ case xegpu::FenceScope::GPU:
+ memScope = xevm::MemScope::DEVICE;
+ break;
+ llvm_unreachable("Unknown XeGPU fence scope.");
+ }
+ xevm::AddrSpace addrSpace{xevm::AddrSpace::GLOBAL};
+ switch (op.getMemoryKind()) {
+ case xegpu::MemorySpace::Global:
+ addrSpace = xevm::AddrSpace::GLOBAL;
+ break;
+ case xegpu::MemorySpace::SLM:
+ addrSpace = xevm::AddrSpace::SHARED;
+ break;
+ llvm_unreachable("Unknown XeGPU fence scope.");
+ }
+ xevm::MemfenceOp::create(rewriter, loc, memScope, addrSpace);
+ rewriter.eraseOp(op);
+ return success();
+ }
+};
+
+class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::DpasOp op, xegpu::DpasOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ctxt = rewriter.getContext();
+ auto aTy = cast<VectorType>(op.getLhs().getType());
+ auto bTy = cast<VectorType>(op.getRhs().getType());
+ auto resultType = cast<VectorType>(op.getResultType());
+
+ auto encodePrecision = [&](Type type) -> xevm::ElemType {
+ if (type == rewriter.getBF16Type())
+ return xevm::ElemType::BF16;
+ else if (type == rewriter.getF16Type())
+ return xevm::ElemType::F16;
+ else if (type == rewriter.getTF32Type())
+ return xevm::ElemType::TF32;
+ else if (type.isInteger(8)) {
+ if (type.isUnsignedInteger())
+ return xevm::ElemType::U8;
+ return xevm::ElemType::S8;
+ } else if (type == rewriter.getF32Type())
+ return xevm::ElemType::F32;
+ else if (type.isInteger(32))
+ return xevm::ElemType::S32;
+ llvm_unreachable("add more support for ElemType");
+ };
+ xevm::ElemType precATy = encodePrecision(aTy.getElementType());
+ xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
+ // auto precA = xevm::ElemTypeAttr::get(ctxt, precATy);
+ // auto precB = xevm::ElemTypeAttr::get(ctxt, precBTy);
+ Value c = op.getAcc();
+ if (!c) {
+ auto elementTy = resultType.getElementType();
+ Attribute initValueAttr;
+ if (isa<FloatType>(elementTy))
+ initValueAttr = FloatAttr::get(elementTy, 0.0);
+ else
+ initValueAttr = IntegerAttr::get(elementTy, 0);
+ c = arith::ConstantOp::create(
+ rewriter, loc, DenseElementsAttr::get(resultType, initValueAttr));
+ }
+
+ Value aVec = op.getLhs();
+ Value bVec = op.getRhs();
+ auto cvecty = cast<VectorType>(c.getType());
+ xevm::ElemType precCTy = encodePrecision(cvecty.getElementType());
+ xevm::ElemType precDTy = encodePrecision(resultType.getElementType());
+ // auto precC = xevm::ElemTypeAttr::get(ctxt, precCTy);
+ // auto precD = xevm::ElemTypeAttr::get(ctxt, precDTy);
+ VectorType cNty =
+ VectorType::get(cvecty.getNumElements(), cvecty.getElementType());
+ if (cvecty != cNty)
+ c = vector::ShapeCastOp::create(rewriter, loc, cNty, c);
+ // below are uArch dependent values, should move away from hardcoding
+ constexpr int32_t systolicDepth{8};
+ constexpr int32_t executionSize{16};
+ Value dpasRes = xevm::MMAOp::create(
+ rewriter, loc, cNty, aVec, bVec, c,
+ xevm::MMAShapeAttr::get(ctxt, cvecty.getNumElements(), executionSize,
+ systolicDepth *
+ getNumOperandsPerDword(precATy)),
+ xevm::MMATypesAttr::get(ctxt, precDTy, precATy, precBTy, precCTy));
+ if (cvecty != cNty)
+ dpasRes = vector::ShapeCastOp::create(rewriter, loc, resultType, dpasRes);
+ rewriter.replaceOp(op, dpasRes);
+ return success();
+ }
+
+private:
+ static unsigned getNumOperandsPerDword(xevm::ElemType pTy) {
+ switch (pTy) {
+ case xevm::ElemType::TF32:
+ return 1;
+ case xevm::ElemType::BF16:
+ case xevm::ElemType::F16:
+ return 2;
+ case xevm::ElemType::U8:
+ case xevm::ElemType::S8:
+ return 4;
+ default:
+ llvm_unreachable("unsupported xevm::ElemType");
+ }
+ }
+};
+
+static std::optional<LLVM::AtomicBinOp>
+matchSimpleAtomicOp(arith::AtomicRMWKind arithKind) {
+ switch (arithKind) {
+ case arith::AtomicRMWKind::addf:
+ return LLVM::AtomicBinOp::fadd;
+ case arith::AtomicRMWKind::addi:
+ return LLVM::AtomicBinOp::add;
+ case arith::AtomicRMWKind::assign:
+ return LLVM::AtomicBinOp::xchg;
+ case arith::AtomicRMWKind::maximumf:
+ return LLVM::AtomicBinOp::fmax;
+ case arith::AtomicRMWKind::maxs:
+ return LLVM::AtomicBinOp::max;
+ case arith::AtomicRMWKind::maxu:
+ return LLVM::AtomicBinOp::umax;
+ case arith::AtomicRMWKind::minimumf:
+ return LLVM::AtomicBinOp::fmin;
+ case arith::AtomicRMWKind::mins:
+ return LLVM::AtomicBinOp::min;
+ case arith::AtomicRMWKind::minu:
+ return LLVM::AtomicBinOp::umin;
+ case arith::AtomicRMWKind::ori:
+ return LLVM::AtomicBinOp::_or;
+ case arith::AtomicRMWKind::andi:
+ return LLVM::AtomicBinOp::_and;
+ default:
+ return std::nullopt;
+ }
+ llvm_unreachable("Invalid AtomicRMWKind");
+}
+
+class AtomicRMWToXeVMPattern : public OpConversionPattern<xegpu::AtomicRMWOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::AtomicRMWOp op, xegpu::AtomicRMWOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto ctxt = rewriter.getContext();
+ auto tdesc = op.getTensorDesc().getType();
+ auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
+ ctxt, getNumericXeVMAddrSpace(tdesc.getMemorySpace()));
+ Value basePtrI64 = arith::IndexCastOp::create(
+ rewriter, loc, rewriter.getI64Type(), adaptor.getTensorDesc());
+ Value basePtrLLVM =
+ LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
+ VectorType srcOrDstVecTy = cast<VectorType>(op.getValue().getType());
+ VectorType srcOrDstFlatVecTy = VectorType::get(
+ srcOrDstVecTy.getNumElements(), srcOrDstVecTy.getElementType());
+ Value srcFlatVec = vector::ShapeCastOp::create(
+ rewriter, loc, srcOrDstFlatVecTy, op.getValue());
+ auto atomicKind = matchSimpleAtomicOp(op.getKind());
+ assert(atomicKind.has_value());
+ Value resVec = srcFlatVec;
+ for (int i = 0; i < srcOrDstVecTy.getNumElements(); i++) {
+ auto val = vector::ExtractOp::create(rewriter, loc, resVec, i);
+ Value idx = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
+ rewriter.getIndexAttr(i));
+ Value currPtr =
+ LLVM::GEPOp::create(rewriter, loc, ptrTypeLLVM,
+ srcOrDstVecTy.getElementType(), basePtrLLVM, idx);
+ Value newVal =
+ LLVM::AtomicRMWOp::create(rewriter, loc, atomicKind.value(), currPtr,
+ val, LLVM::AtomicOrdering::seq_cst);
+ resVec = vector::InsertOp::create(rewriter, loc, newVal, resVec, i);
+ }
+ rewriter.replaceOp(op, resVec);
+ return success();
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Pass Definition
+//===----------------------------------------------------------------------===//
+
+struct ConvertXeGPUToXeVMPass
+ : public impl::ConvertXeGPUToXeVMPassBase<ConvertXeGPUToXeVMPass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ LLVMTypeConverter typeConverter(&getContext());
+ typeConverter.addConversion([&](VectorType type) -> Type {
+ unsigned rank = type.getRank();
+ auto elemType = type.getElementType();
+ // If the element type is index, convert it to i64.
+ if (llvm::isa<IndexType>(elemType))
+ elemType = IntegerType::get(&getContext(), 64);
+ // If the vector is a scalar or has a single element, return the element
+ if (rank < 1 || type.getNumElements() == 1)
+ return elemType;
+ // Otherwise, convert the vector to a flat vector type.
+ unsigned sum = 1;
+ for (unsigned i = 0; i < rank; i++) {
+ sum *= type.getShape()[i];
+ }
+ return VectorType::get(sum, elemType);
+ });
+ typeConverter.addConversion([&](xegpu::TensorDescType type) -> Type {
+ if (type.isScattered()) {
+ return IntegerType::get(&getContext(), 64);
+ }
+ auto i32Type = IntegerType::get(&getContext(), 32);
+ return VectorType::get(8, i32Type);
+ });
+
+ auto ui64MaterializationCast = [](OpBuilder &builder, Type type,
+ ValueRange inputs,
+ Location loc) -> Value {
+ if (inputs.size() != 1)
+ return {};
+ auto input = inputs.front();
+ if (input.getType() == builder.getIntegerType(64, false)) {
+ Value cast =
+ index::CastUOp::create(builder, loc, builder.getIndexType(), input)
+ .getResult();
+ return arith::IndexCastOp::create(builder, loc, type, cast).getResult();
+ }
+ return {};
+ };
+
+ auto vector1DMaterializationCast = [](OpBuilder &builder, Type type,
+ ValueRange inputs,
+ Location loc) -> Value {
+ if (inputs.size() != 1)
+ return {};
+ auto input = inputs.front();
+ if (auto vecTy = dyn_cast<VectorType>(input.getType())) {
+ if (vecTy.getNumElements() == 1) {
+ // If the vector has a single element, return the element type.
+ Value cast =
+ vector::ExtractOp::create(builder, loc, input, 0).getResult();
+ if (vecTy.getElementType() == builder.getIndexType())
+ cast = arith::IndexCastOp::create(builder, loc, type, cast)
+ .getResult();
+ return cast;
+ }
+ }
+ return {};
+ };
+ typeConverter.addSourceMaterialization(ui64MaterializationCast);
+ typeConverter.addSourceMaterialization(vector1DMaterializationCast);
+ typeConverter.addTargetMaterialization(ui64MaterializationCast);
+ typeConverter.addTargetMaterialization(vector1DMaterializationCast);
+ ConversionTarget target(getContext());
+ target.addLegalDialect<xevm::XeVMDialect, LLVM::LLVMDialect,
+ vector::VectorDialect, arith::ArithDialect,
+ memref::MemRefDialect, gpu::GPUDialect,
+ index::IndexDialect>();
+ target.addIllegalDialect<xegpu::XeGPUDialect>();
+
+ RewritePatternSet patterns(&getContext());
+ populateXeGPUToXeVMConversionPatterns(patterns, typeConverter);
+ scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
+ patterns, target);
+ if (failed(applyPartialConversion(getOperation(), target,
+ std::move(patterns))))
+ signalPassFailure();
+ }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Pattern Population
+//===----------------------------------------------------------------------===//
+void mlir::populateXeGPUToXeVMConversionPatterns(
+ RewritePatternSet &patterns, LLVMTypeConverter &typeConverter) {
+ patterns.add<CreateNdDescToXeVMPattern, UpdateNdOffsetToXeVMPattern,
+ LoadStorePrefetchNdToXeVMPattern<xegpu::LoadNdOp>,
+ LoadStorePrefetchNdToXeVMPattern<xegpu::StoreNdOp>,
+ LoadStorePrefetchNdToXeVMPattern<xegpu::PrefetchNdOp>>(
+ typeConverter, patterns.getContext());
+ patterns.add<CreateDescToXeVMPattern, UpdateOffsetToXeVMPattern,
+ AtomicRMWToXeVMPattern, PrefetchToXeVMPattern,
+ LoadStoreToXeVMPattern<xegpu::LoadGatherOp>,
+ LoadStoreToXeVMPattern<xegpu::StoreScatterOp>>(
+ typeConverter, patterns.getContext());
+ patterns.add<FenceToXeVMPattern, DpasToXeVMPattern>(typeConverter,
+ patterns.getContext());
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
new file mode 100644
index 0000000000000..4fba920f023c4
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -0,0 +1,48 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+gpu.module @create_nd_tdesc {
+ // CHECK-LABEL: gpu.func @create_nd_tdesc
+ // CHECK-SAME: %[[ARG0:.*]]: memref<8x16xf32, 1>, %[[ARG1:.*]]: ui64
+ // CHECK-SAME: %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index
+ gpu.func @create_nd_tdesc(%src: memref<8x16xf32, 1>, %ptr: ui64, %shape1: index, %shape2: index,
+ %stride1: index, %stride2: index) kernel {
+ // CHECK: %[[VAR0:.*]] = index.castu %[[ARG1]] : ui64 to index
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xi32>
+ // CHECK: %[[C0_I32:.*]] = arith.constant 0 : i32
+ // CHECK: %[[C0_I32_0:.*]] = arith.constant 0 : i32
+ // CHECK: %[[VAR2:.*]] = arith.index_cast %[[ARG3]] : index to i64
+ // CHECK: %[[VAR3:.*]] = arith.trunci %[[VAR2]] : i64 to i32
+ // CHECK: %[[VAR4:.*]] = arith.index_cast %[[ARG2]] : index to i64
+ // CHECK: %[[VAR5:.*]] = arith.trunci %[[VAR4]] : i64 to i32
+ // CHECK: %[[VAR6:.*]] = vector.bitcast %[[CST]] : vector<8xi32> to vector<4xi64>
+ // CHECK: %[[VAR7:.*]] = vector.insert %[[VAR1]], %[[VAR6]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[VAR8:.*]] = vector.bitcast %[[VAR7]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[VAR9:.*]] = vector.insert %[[VAR3]], %[[VAR8]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[VAR10:.*]] = vector.insert %[[VAR5]], %[[VAR9]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[VAR11:.*]] = vector.insert %[[C0_I32]], %[[VAR10]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[VAR12:.*]] = vector.insert %[[C0_I32_0]], %[[VAR11]] [5] : i32 into vector<8xi32>
+ %ptr_tdesc = xegpu.create_nd_tdesc %ptr, shape:[%shape1, %shape2], strides:[%stride1, %stride2]
+ : ui64 -> !xegpu.tensor_desc<8x16xf32>
+
+ // CHECK: %[[MEMSPACECAST:.*]] = memref.memory_space_cast %[[ARG0]] : memref<8x16xf32, 1> to memref<8x16xf32>
+ %srcce = memref.memory_space_cast %src : memref<8x16xf32, 1> to memref<8x16xf32>
+
+ // CHECK: %[[CST_1:.*]] = arith.constant dense<0> : vector<8xi32>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[MEMSPACECAST]] : memref<8x16xf32> -> index
+ // CHECK: %[[C0_I32_2:.*]] = arith.constant 0 : i32
+ // CHECK: %[[C0_I32_3:.*]] = arith.constant 0 : i32
+ // CHECK: %[[C16_I32:.*]] = arith.constant 16 : i32
+ // CHECK: %[[C8_I32:.*]] = arith.constant 8 : i32
+ // CHECK: %[[VAR13:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[VAR14:.*]] = vector.bitcast %[[CST_1]] : vector<8xi32> to vector<4xi64>
+ // CHECK: %[[VAR15:.*]] = vector.insert %[[VAR13]], %[[VAR14]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[VAR16:.*]] = vector.bitcast %[[VAR15]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[VAR17:.*]] = vector.insert %[[C16_I32]], %[[VAR16]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[VAR18:.*]] = vector.insert %[[C8_I32]], %[[VAR17]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[VAR19:.*]] = vector.insert %[[C0_I32_2]], %[[VAR18]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[VAR20:.*]] = vector.insert %[[C0_I32_3]], %[[VAR19]] [5] : i32 into vector<8xi32>
+ %src_tdesc = xegpu.create_nd_tdesc %srcce : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/dpas.mlir b/mlir/test/Conversion/XeGPUToXeVM/dpas.mlir
new file mode 100644
index 0000000000000..15940fc4aca26
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/dpas.mlir
@@ -0,0 +1,17 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+#sg_map_a_f16 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+#sg_map_b_f16 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>
+#sg_map_c_f32 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+
+gpu.module @load_store_check {
+ //CHECK: func.func @dpas(%[[ARG0:.*]]: vector<8xf16>, %[[ARG1:.*]]: vector<16xf16>, %[[ARG2:.*]]: vector<8xf32>) -> vector<8xf32>
+ func.func @dpas(%a_loaded: vector<8xf16>, %b_loaded: vector<16xf16>, %c_loaded: vector<8xf32>) -> vector<8xf32> {
+ // Loads are checked in a separate test.
+ // CHECK: %[[D:.*]] = xevm.mma %[[ARG0]], %[[ARG1]], %[[ARG2]] {shape = <m = 8, n = 16, k = 16>, types = <d = f32, a = f16, b = f16, c = f32>}
+ // CHECK-SAME: : (vector<8xf16>, vector<16xf16>, vector<8xf32>) -> vector<8xf32>
+ %d = xegpu.dpas %a_loaded, %b_loaded, %c_loaded {a_layout = #sg_map_a_f16, b_layout = #sg_map_b_f16, c_layout = #sg_map_c_f32}
+ : vector<8xf16>, vector<16xf16>, vector<8xf32> -> vector<8xf32>
+ return %d : vector<8xf32>
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/fence.mlir b/mlir/test/Conversion/XeGPUToXeVM/fence.mlir
new file mode 100644
index 0000000000000..cedfcace398a6
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/fence.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+gpu.module @fence_check {
+ gpu.func @fence(%dst: memref<8x16xf32, 1>) kernel {
+ %tid_x = gpu.thread_id x
+ %tid_x_i32 = arith.index_cast %tid_x : index to i32
+ %tid_x_f32 = arith.sitofp %tid_x_i32 : i32 to f32
+
+ // CHECK: xevm.memfence <{addrspace = #xevm.addr_space<global>, scope = #xevm.mem_scope<workgroup>}>
+ xegpu.fence memory_kind = global, fence_scope = workgroup
+ %c0 = arith.constant 0 : index
+ memref.store %tid_x_f32, %dst[%c0, %c0] : memref<8x16xf32, 1>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
new file mode 100644
index 0000000000000..c692da632d458
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
@@ -0,0 +1,71 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+gpu.module @load_store_check {
+ gpu.func @load_store(%src: memref<8x16xf32, 1>, %dst: memref<8x16xf32, 1>) kernel {
+ %srcce = memref.memory_space_cast %src : memref<8x16xf32, 1> to memref<8x16xf32>
+ %dstte = memref.memory_space_cast %dst : memref<8x16xf32, 1> to memref<8x16xf32>
+
+ // CHECK: %[[LD_PTR_AS_I64:.*]] = arith.index_castui {{.*}} : index to i64
+ // CHECK: %[[LD_CREATE_DESC_I64:.*]] = vector.bitcast {{.*}} : vector<8xi32> to vector<4xi64>
+ // CHECK: %[[LD_DESC_0:.*]] = vector.insert %[[LD_PTR_AS_I64]], %[[LD_CREATE_DESC_I64]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[LD_DESC_1:.*]] = vector.bitcast %[[LD_DESC_0]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[LD_DESC_2:.*]] = vector.insert {{.*}}, %[[LD_DESC_1]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC_3:.*]] = vector.insert {{.*}}, %[[LD_DESC_2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC_4:.*]] = vector.insert {{.*}}, %[[LD_DESC_3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC:.*]] = vector.insert {{.*}}, %[[LD_DESC_4]] [5] : i32 into vector<8xi32>
+ %src_tdesc = xegpu.create_nd_tdesc %srcce : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32>
+
+
+ //CHECK: %[[LD_DESC_I64:.*]] = vector.bitcast %[[LD_DESC]] : vector<8xi32> to vector<4xi64>
+ //CHECK: %[[LD_INTPTR:.*]] = vector.extract %[[LD_DESC_I64]][0] : i64 from vector<4xi64>
+ //CHECK: %[[LD_BASE_W:.*]] = vector.extract %[[LD_DESC]][2] : i32 from vector<8xi32>
+ //CHECK: %[[LD_BASE_H:.*]] = vector.extract %[[LD_DESC]][3] : i32 from vector<8xi32>
+ //CHECK: %[[LD_TILE_W:.*]] = arith.constant 0 : i32
+ //CHECK: %[[LD_TILE_H:.*]] = arith.constant 0 : i32
+ //CHECK: %[[LD_LLVMPTR:.*]] = llvm.inttoptr %[[LD_INTPTR]] : i64 to !llvm.ptr<1>
+ //CHECK: %[[LD_SIZEOF_F32:.*]] = arith.constant 4 : i32
+ //CHECK: %[[LD_BASE_ROW_IN_BYTES:.*]] = arith.muli %[[LD_BASE_W]], %[[LD_SIZEOF_F32]] : i32
+ //CHECK: %[[LD_LOADED_I32:.*]] = xevm.blockload2d %[[LD_LLVMPTR]], %[[LD_BASE_ROW_IN_BYTES]],
+ //CHECK-SAME: %[[LD_BASE_H]], %[[LD_BASE_ROW_IN_BYTES]], %[[LD_TILE_W]], %[[LD_TILE_H]]
+ //CHECK-SAME: <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>, elem_size_in_bits = 32 : i32,
+ //CHECK-SAME: pack_register = false, tile_height = 8 : i32, tile_width = 16 : i32, transpose = false,
+ //CHECK-SAME: v_blocks = 1 : i32}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi32>
+ %loaded = xegpu.load_nd %src_tdesc[0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<8x16xf32> -> vector<8xf32>
+ //CHECK: %[[LD_LOADED_F32:.*]] = vector.bitcast %[[LD_LOADED_I32]] : vector<8xi32> to vector<8xf32>
+
+ %tid_x = gpu.thread_id x
+ %tid_x_i32 = arith.index_cast %tid_x : index to i32
+ %tid_x_f32 = arith.sitofp %tid_x_i32 : i32 to f32
+ //CHECK: %[[LOADED_F32_MODIFIED:.*]] = vector.insert %{{.*}}, %[[LD_LOADED_F32]] [0] : f32 into vector<8xf32>
+ %loaded_modified = vector.insert %tid_x_f32, %loaded[0] : f32 into vector<8xf32>
+
+ // CHECK: %[[PTR_AS_I64:.*]] = arith.index_castui {{.*}} : index to i64
+ // CHECK: %[[CREATE_DESC_I64:.*]] = vector.bitcast {{.*}} : vector<8xi32> to vector<4xi64>
+ // CHECK: %[[DESC_0:.*]] = vector.insert %[[PTR_AS_I64]], %[[CREATE_DESC_I64]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[DESC_1:.*]] = vector.bitcast %[[DESC_0]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[DESC_2:.*]] = vector.insert {{.*}}, %[[DESC_1]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[DESC_3:.*]] = vector.insert {{.*}}, %[[DESC_2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[DESC_4:.*]] = vector.insert {{.*}}, %[[DESC_3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[DESC:.*]] = vector.insert {{.*}}, %[[DESC_4]] [5] : i32 into vector<8xi32>
+ %dst_tdesc = xegpu.create_nd_tdesc %dstte : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>>
+
+ //CHECK: %[[DESC_I64:.*]] = vector.bitcast %[[DESC]] : vector<8xi32> to vector<4xi64>
+ //CHECK: %[[INTPTR:.*]] = vector.extract %[[DESC_I64]][0] : i64 from vector<4xi64>
+ //CHECK: %[[BASE_W:.*]] = vector.extract %[[DESC]][2] : i32 from vector<8xi32>
+ //CHECK: %[[BASE_H:.*]] = vector.extract %[[DESC]][3] : i32 from vector<8xi32>
+ //CHECK: %[[TILE_W:.*]] = arith.constant 0 : i32
+ //CHECK: %[[TILE_H:.*]] = arith.constant 0 : i32
+ //CHECK: %[[LLVMPTR:.*]] = llvm.inttoptr %[[INTPTR]] : i64 to !llvm.ptr<1>
+ //CHECK: %[[SIZEOF_F32:.*]] = arith.constant 4 : i32
+ //CHECK: %[[BASE_ROW_IN_BYTES:.*]] = arith.muli %[[BASE_W]], %[[SIZEOF_F32]] : i32
+ //CHECK: %[[FLAT_VALUE_I32:.*]] = vector.bitcast %[[LOADED_F32_MODIFIED]] : vector<8xf32> to vector<8xi32>
+ //CHECK: xevm.blockstore2d %[[LLVMPTR]], %[[BASE_ROW_IN_BYTES]], %[[BASE_H]], %[[BASE_ROW_IN_BYTES]],
+ //CHECK-SAME: %[[TILE_W]], %[[TILE_H]], %[[FLAT_VALUE_I32]]
+ //CHECK-SAME: <{cache_control = #xevm.store_cache_control<L1wb_L2uc_L3uc>, elem_size_in_bits = 32 : i32,
+ //CHECK-SAME: tile_height = 8 : i32, tile_width = 16 : i32}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
+ xegpu.store_nd %loaded_modified, %dst_tdesc[0, 0] <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<8xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
new file mode 100644
index 0000000000000..f6d023307313a
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstoreprefetch.mlir
@@ -0,0 +1,357 @@
+// RUN: mlir-opt %s --split-input-file -convert-xegpu-to-xevm | FileCheck %s
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_ui64_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: ui64
+gpu.func @load_gather_ui64_src_constant_offset(%src: ui64) {
+ // CHECK: %[[VAR0:.*]] = index.castu %[[ARG0]] : ui64 to index
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR4:.*]] = vector.extract %[[CST_0]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR5:.*]] = arith.muli %[[VAR3]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR6:.*]] = arith.addi %[[VAR1]], %[[VAR5]] : i64
+ %2 = xegpu.create_tdesc %src, %0 : ui64, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR7:.*]] = llvm.inttoptr %[[VAR6]] : i64 to !llvm.ptr<1>
+ // CHECK: %[[VAR8:.*]] = scf.if %[[VAR4]] -> (vector<2xf32>) {
+ // CHECK: %[[VAR9:.*]] = llvm.load %[[VAR7]] {cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}
+ // CHECK-SAME: : !llvm.ptr<1> -> vector<2xf32>
+ // CHECK: scf.yield %[[VAR9]] : vector<2xf32>
+ // CHECK: } else {
+ // CHECK: %[[CST_1:.*]] = arith.constant dense<0.000000e+00> : vector<2xf32>
+ // CHECK: scf.yield %[[CST_1]] : vector<2xf32>
+ %3 = xegpu.load %2, %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<1xi1> -> vector<2xf32>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_memref_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>
+gpu.func @load_gather_memref_src_constant_offset(%src: memref<256xf32>) {
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST_0]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR3:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR4:.*]] = arith.muli %[[VAR1]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR5:.*]] = arith.addi %[[VAR3]], %[[VAR4]] : i64
+ %2 = xegpu.create_tdesc %src, %0 : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>
+ // CHECK: %[[VAR6:.*]] = llvm.inttoptr %[[VAR5]] : i64 to !llvm.ptr<1>
+ // CHECK: %[[VAR7:.*]] = scf.if %[[VAR2]] -> (f32) {
+ // CHECK: %[[VAR8:.*]] = llvm.load %[[VAR6]] {cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}
+ // CHECK-SAME: : !llvm.ptr<1> -> vector<1xf32>
+ // CHECK: %[[VAR9:.*]] = vector.extract %[[VAR8]][0] : f32 from vector<1xf32>
+ // CHECK: scf.yield %[[VAR9]] : f32
+ // CHECK: } else {
+ // CHECK: %[[CST_1:.*]] = arith.constant dense<0.000000e+00> : vector<1xf32>
+ // CHECK: %[[VAR8:.*]] = vector.extract %[[CST_1:.*]][0] : f32 from vector<1xf32>
+ // CHECK: scf.yield %[[VAR8]] : f32
+ %3 = xegpu.load %2, %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>, vector<1xi1> -> vector<1xf32>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_memref_src_value_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf16>, %[[ARG1:.*]]: vector<1xindex>
+gpu.func @load_gather_memref_src_value_offset(%src: memref<256xf16>, %offset: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf16> -> index
+ // CHECK: %[[VAR3:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C2_I64:.*]] = arith.constant 2 : i64
+ // CHECK: %[[VAR4:.*]] = arith.muli %[[VAR1]], %[[C2_I64]] : i64
+ // CHECK: %[[VAR5:.*]] = arith.addi %[[VAR3]], %[[VAR4]] : i64
+ %2 = xegpu.create_tdesc %src, %offset : memref<256xf16>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x8xf16, #xegpu.scatter_tdesc_attr<chunk_size = 8>>
+ // CHECK: %[[VAR6:.*]] = llvm.inttoptr %[[VAR5]] : i64 to !llvm.ptr<1>
+ // CHECK: %[[VAR7:.*]] = scf.if %[[VAR2]] -> (vector<8xf16>) {
+ // CHECK: %[[VAR8:.*]] = llvm.load %[[VAR6]] {cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}
+ // CHECK-SAME: : !llvm.ptr<1> -> vector<8xf16>
+ // CHECK: scf.yield %[[VAR8]] : vector<8xf16>
+ // CHECK: } else {
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<0.000000e+00> : vector<8xf16>
+ // CHECK: scf.yield %[[CST_0]] : vector<8xf16>
+ %3 = xegpu.load %2, %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x8xf16, #xegpu.scatter_tdesc_attr<chunk_size = 8>>, vector<1xi1> -> vector<8xf16>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @load_gather_memref_src_load_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf16>, %[[ARG1:.*]]: vector<1xindex>, %[[ARG2:.*]]: vector<1xindex>
+gpu.func @load_gather_memref_src_load_offset(%src: memref<256xf16>, %offset1: vector<1xindex>, %offset2: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG2]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR4:.*]] = vector.extract %[[CST]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf16> -> index
+ // CHECK: %[[VAR5:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C2_I64:.*]] = arith.constant 2 : i64
+ // CHECK: %[[VAR6:.*]] = arith.muli %[[VAR3]], %[[C2_I64]] : i64
+ // CHECK: %[[VAR7:.*]] = arith.addi %[[VAR5]], %[[VAR6]] : i64
+ %2 = xegpu.create_tdesc %src, %offset1 : memref<256xf16>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x8xf16, #xegpu.scatter_tdesc_attr<chunk_size = 8>>
+ // CHECK: %[[C2_I64_0:.*]] = arith.constant 2 : i64
+ // CHECK: %[[VAR8:.*]] = arith.muli %[[VAR1]], %[[C2_I64_0]] : i64
+ // CHECK: %[[VAR9:.*]] = arith.addi %[[VAR7]], %[[VAR8]] : i64
+ // CHECK: %[[VAR10:.*]] = llvm.inttoptr %[[VAR9]] : i64 to !llvm.ptr<1>
+ // CHECK: %[[VAR11:.*]] = scf.if %[[VAR4]] -> (vector<8xf16>) {
+ // CHECK: %[[VAR12:.*]] = llvm.load %[[VAR10]] {cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}
+ // CHECK-SAME: : !llvm.ptr<1> -> vector<8xf16>
+ // CHECK: scf.yield %[[VAR12]] : vector<8xf16>
+ // CHECK: } else {
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<0.000000e+00> : vector<8xf16>
+ // CHECK: scf.yield %[[CST_0]] : vector<8xf16>
+ %3 = xegpu.load %2[%offset2], %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x8xf16, #xegpu.scatter_tdesc_attr<chunk_size = 8>>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_ui64_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: ui64
+gpu.func @store_scatter_ui64_src_constant_offset(%src: ui64) {
+ // CHECK: %[[VAR0:.*]] = index.castu %[[ARG0]] : ui64 to index
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR4:.*]] = vector.extract %[[CST_0]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[CST_1:.*]] = arith.constant dense<2.900000e+00> : vector<2xf32>
+ %2 = arith.constant dense<2.9>: vector<2xf32>
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR5:.*]] = arith.muli %[[VAR3]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR6:.*]] = arith.addi %[[VAR1]], %[[VAR5]] : i64
+ %3 = xegpu.create_tdesc %src, %0 : ui64, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR7:.*]] = llvm.inttoptr %[[VAR6]] : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[VAR4]] {
+ // CHECK: llvm.store %[[CST_1]], %[[VAR7]] {cache_control = #xevm.store_cache_control<L1wb_L2uc_L3uc>}
+ // CHECK-SAME: : vector<2xf32>, !llvm.ptr<1>
+ xegpu.store %2, %3, %1 <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<2xf32>, !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<1xi1>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_memref_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>
+gpu.func @store_scatter_memref_src_constant_offset(%src: memref<256xf32>) {
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST_0]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[CST_1:.*]] = arith.constant dense<2.900390e+00> : vector<2xf16>
+ %2 = arith.constant dense<2.9>: vector<2xf16>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR3:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C2_I64:.*]] = arith.constant 2 : i64
+ // CHECK: %[[VAR4:.*]] = arith.muli %[[VAR1]], %[[C2_I64]] : i64
+ // CHECK: %[[VAR5:.*]] = arith.addi %[[VAR3]], %[[VAR4]] : i64
+ %3 = xegpu.create_tdesc %src, %0 : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf16, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR6:.*]] = llvm.inttoptr %[[VAR5]] : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[VAR2]] {
+ // CHECK: llvm.store %[[CST_1]], %[[VAR6]] {cache_control = #xevm.store_cache_control<L1wb_L2uc_L3uc>}
+ // CHECK-SAME: : vector<2xf16>, !llvm.ptr<1>
+ xegpu.store %2, %3, %1 <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<2xf16>, !xegpu.tensor_desc<1x2xf16, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<1xi1>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_memref_src_value_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>, %[[ARG1:.*]]: vector<1xindex>
+gpu.func @store_scatter_memref_src_value_offset(%src: memref<256xf32>, %offset: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<2.900000e+00> : vector<1xf32>
+ %2 = arith.constant dense<2.9>: vector<1xf32>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR3:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR4:.*]] = arith.muli %[[VAR1]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR5:.*]] = arith.addi %[[VAR3]], %[[VAR4]] : i64
+ %3 = xegpu.create_tdesc %src, %offset : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>
+ // CHECK: %[[VAR6:.*]] = llvm.inttoptr %[[VAR5]] : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[VAR2]] {
+ // CHECK: llvm.store %[[CST_0]], %[[VAR6]] {cache_control = #xevm.store_cache_control<L1wb_L2uc_L3uc>}
+ // CHECK-SAME: : vector<1xf32>, !llvm.ptr<1>
+ xegpu.store %2, %3, %1 <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<1xf32>, !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>, vector<1xi1>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @store_scatter_memref_src_store_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>, %[[ARG1:.*]]: vector<1xindex>, %[[ARG2:.*]]: vector<1xindex>
+gpu.func @store_scatter_memref_src_store_offset(%src: memref<256xf32>, %offset: vector<1xindex>, %offset2: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG2]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<true> : vector<1xi1>
+ // CHECK: %[[VAR4:.*]] = vector.extract %[[CST]][0] : i1 from vector<1xi1>
+ %1 = arith.constant dense<1>: vector<1xi1>
+ // CHECK: %[[CST_0:.*]] = arith.constant dense<2.900000e+00> : vector<1xf32>
+ %2 = arith.constant dense<2.9>: vector<1xf32>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR5:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR6:.*]] = arith.muli %[[VAR3]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR7:.*]] = arith.addi %[[VAR5]], %[[VAR6]] : i64
+ %3 = xegpu.create_tdesc %src, %offset : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>
+ // CHECK: %[[C4_I64_1:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR8:.*]] = arith.muli %[[VAR1]], %[[C4_I64_1]] : i64
+ // CHECK: %[[VAR9:.*]] = arith.addi %[[VAR7]], %[[VAR8]] : i64
+ // CHECK: %[[VAR10:.*]] = llvm.inttoptr %[[VAR9]] : i64 to !llvm.ptr<1>
+ // CHECK: scf.if %[[VAR4]] {
+ // CHECK: llvm.store %[[CST_0]], %[[VAR10]] {cache_control = #xevm.store_cache_control<L1wb_L2uc_L3uc>}
+ // CHECK-SAME: : vector<1xf32>, !llvm.ptr<1>
+ xegpu.store %2, %3[%offset2], %1 <{l1_hint = #xegpu.cache_hint<write_back>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : vector<1xf32>, !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>, vector<1xindex>, vector<1xi1>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @prefetch_ui64_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: ui64
+gpu.func @prefetch_ui64_src_constant_offset(%src: ui64) {
+ // CHECK: %[[VAR0:.*]] = index.castu %[[ARG0]] : ui64 to index
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR4:.*]] = arith.muli %[[VAR3]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR5:.*]] = arith.addi %[[VAR1]], %[[VAR4]] : i64
+ %1 = xegpu.create_tdesc %src, %0 : ui64, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR6:.*]] = llvm.inttoptr %[[VAR5]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.prefetch %[[VAR6]] <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}> : (!llvm.ptr<1>)
+ xegpu.prefetch %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @prefetch_memref_src_constant_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>
+gpu.func @prefetch_memref_src_constant_offset(%src: memref<256xf32>) {
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ %0 = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR2:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR3:.*]] = arith.muli %[[VAR1]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR4:.*]] = arith.addi %[[VAR2]], %[[VAR3]] : i64
+ %1 = xegpu.create_tdesc %src, %0 : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR5:.*]] = llvm.inttoptr %[[VAR4]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.prefetch %[[VAR5]] <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}> : (!llvm.ptr<1>)
+ xegpu.prefetch %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @prefetch_memref_src_value_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>, %[[ARG1:.*]]: vector<1xindex>
+gpu.func @prefetch_memref_src_value_offset(%src: memref<256xf32>, %offset: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR2:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR3:.*]] = arith.muli %[[VAR1]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR4:.*]] = arith.addi %[[VAR2]], %[[VAR3]] : i64
+ %1 = xegpu.create_tdesc %src, %offset : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[VAR5:.*]] = llvm.inttoptr %[[VAR4]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.prefetch %[[VAR5]] <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}> : (!llvm.ptr<1>)
+ xegpu.prefetch %1 <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ gpu.return
+}
+}
+// -----
+
+gpu.module @test {
+// CHECK-LABEL: @prefetch_memref_src_prefetch_offset
+// CHECK-SAME: %[[ARG0:.*]]: memref<256xf32>, %[[ARG1:.*]]: vector<1xindex>, %[[ARG2:.*]]: vector<1xindex>
+gpu.func @prefetch_memref_src_prefetch_offset(%src: memref<256xf32>, %offset: vector<1xindex>, %offset2: vector<1xindex>) {
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[ARG2]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[VAR2:.*]] = vector.extract %[[ARG1]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR3:.*]] = arith.index_cast %[[VAR2]] : index to i64
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<256xf32> -> index
+ // CHECK: %[[VAR4:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR5:.*]] = arith.muli %[[VAR3]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR6:.*]] = arith.addi %[[VAR4]], %[[VAR5]] : i64
+ %1 = xegpu.create_tdesc %src, %offset : memref<256xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>
+ // CHECK: %[[C4_I64_0:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR7:.*]] = arith.muli %[[VAR1]], %[[C4_I64_0]] : i64
+ // CHECK: %[[VAR8:.*]] = arith.addi %[[VAR6]], %[[VAR7]] : i64
+ // CHECK: %[[VAR9:.*]] = llvm.inttoptr %[[VAR8]] : i64 to !llvm.ptr<1>
+ // CHECK: xevm.prefetch %[[VAR9]] <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>}> : (!llvm.ptr<1>)
+ xegpu.prefetch %1[%offset2] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<1x2xf32, #xegpu.scatter_tdesc_attr<chunk_size = 2>>, vector<1xindex>
+ gpu.return
+}
+}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/prefetch_nd.mlir b/mlir/test/Conversion/XeGPUToXeVM/prefetch_nd.mlir
new file mode 100644
index 0000000000000..8513b4f9857fb
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/prefetch_nd.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm -split-input-file %s | FileCheck %s
+
+gpu.module @fence_check {
+ gpu.func @fence(%src: memref<8x16xf32, 1>, %dst: memref<8x16xf32, 1>) kernel {
+ %srcce = memref.memory_space_cast %src : memref<8x16xf32, 1> to memref<8x16xf32>
+ %dstte = memref.memory_space_cast %dst : memref<8x16xf32, 1> to memref<8x16xf32>
+
+ // CHECK: %[[LD_PTR_AS_I64:.*]] = arith.index_castui {{.*}} : index to i64
+ // CHECK: %[[LD_CREATE_DESC_I64:.*]] = vector.bitcast {{.*}} : vector<8xi32> to vector<4xi64>
+ // CHECK: %[[LD_DESC_0:.*]] = vector.insert %[[LD_PTR_AS_I64]], %[[LD_CREATE_DESC_I64]] [0] : i64 into vector<4xi64>
+ // CHECK: %[[LD_DESC_1:.*]] = vector.bitcast %[[LD_DESC_0]] : vector<4xi64> to vector<8xi32>
+ // CHECK: %[[LD_DESC_2:.*]] = vector.insert {{.*}}, %[[LD_DESC_1]] [2] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC_3:.*]] = vector.insert {{.*}}, %[[LD_DESC_2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC_4:.*]] = vector.insert {{.*}}, %[[LD_DESC_3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[LD_DESC:.*]] = vector.insert {{.*}}, %[[LD_DESC_4]] [5] : i32 into vector<8xi32>
+ %src_tdesc = xegpu.create_nd_tdesc %srcce : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32,
+ #xegpu.block_tdesc_attr<memory_space = global>, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+
+ //CHECK: %[[LD_DESC_I64:.*]] = vector.bitcast %[[LD_DESC]] : vector<8xi32> to vector<4xi64>
+ //CHECK: %[[PREF_INTPTR:.*]] = vector.extract %[[LD_DESC_I64]][0] : i64 from vector<4xi64>
+ //CHECK: %[[PREF_BASE_W:.*]] = vector.extract %[[LD_DESC]][2] : i32 from vector<8xi32>
+ //CHECK: %[[PREF_BASE_H:.*]] = vector.extract %[[LD_DESC]][3] : i32 from vector<8xi32>
+ //CHECK: %[[PREF_TILE_W:.*]] = arith.constant 0 : i32
+ //CHECK: %[[PREF_TILE_H:.*]] = arith.constant 0 : i32
+ //CHECK: %[[PREF_LLVMPTR:.*]] = llvm.inttoptr %[[PREF_INTPTR]] : i64 to !llvm.ptr<1>
+ //CHECK: %[[PREF_SIZEOF_F32:.*]] = arith.constant 4 : i32
+ //CHECK: %[[PREF_BASE_ROW_IN_BYTES:.*]] = arith.muli %[[PREF_BASE_W]], %[[PREF_SIZEOF_F32]] : i32
+ //CHECK: xevm.blockprefetch2d %[[PREF_LLVMPTR]], %[[PREF_BASE_ROW_IN_BYTES]], %[[PREF_BASE_H]],
+ //CHECK-SAME: %[[PREF_BASE_ROW_IN_BYTES]], %[[PREF_TILE_W]], %[[PREF_TILE_H]]
+ //CHECK-SAME: <{cache_control = #xevm.load_cache_control<L1c_L2uc_L3uc>, elem_size_in_bits = 32 : i32,
+ //CHECK-SAME: tile_height = 8 : i32, tile_width = 16 : i32, v_blocks = 1 : i32}>
+ //CHECK-SAME: : (!llvm.ptr<1>, i32, i32, i32, i32, i32)
+ xegpu.prefetch_nd %src_tdesc[0, 0] <{l1_hint = #xegpu.cache_hint<cached>, l2_hint = #xegpu.cache_hint<uncached>}>
+ : !xegpu.tensor_desc<8x16xf32, #xegpu.block_tdesc_attr<memory_space = global>,
+ #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+
+ gpu.return
+ }
+}
+
diff --git a/mlir/test/Conversion/XeGPUToXeVM/update_offset.mlir b/mlir/test/Conversion/XeGPUToXeVM/update_offset.mlir
new file mode 100644
index 0000000000000..e9d7fd4cf40a6
--- /dev/null
+++ b/mlir/test/Conversion/XeGPUToXeVM/update_offset.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt -convert-xegpu-to-xevm %s | FileCheck %s
+
+gpu.module @update_offset {
+ // CHECK-LABEL: gpu.func @update_offset
+ // CHECK-SAME: %[[ARG0:.*]]: memref<128xf32>
+ gpu.func @update_offset(%src: memref<128xf32>) kernel {
+ // CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<1xindex>
+ %offset = arith.constant dense<0> : vector<1xindex>
+ // CHECK: %[[VAR0:.*]] = vector.extract %[[CST]][0] : index from vector<1xindex>
+ // CHECK: %[[VAR1:.*]] = arith.index_cast %[[VAR0]] : index to i64
+ // CHECK: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[ARG0]] : memref<128xf32> -> index
+ // CHECK: %[[VAR2:.*]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[C4_I64:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR3:.*]] = arith.muli %[[VAR1]], %[[C4_I64]] : i64
+ // CHECK: %[[VAR4:.*]] = arith.addi %[[VAR2]], %[[VAR3]] : i64
+ %src_tdesc = xegpu.create_tdesc %src, %offset : memref<128xf32>, vector<1xindex>
+ -> !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>
+ // CHECK: %[[C4_I64_0:.*]] = arith.constant 4 : i64
+ // CHECK: %[[VAR5:.*]] = arith.muli %[[VAR1]], %[[C4_I64_0]] : i64
+ // CHECK: %[[VAR6:.*]] = arith.addi %[[VAR4]], %[[VAR5]] : i64
+ %new_tdesc = xegpu.update_offset %src_tdesc, %offset : !xegpu.tensor_desc<1xf32, #xegpu.scatter_tdesc_attr<>>
+ , vector<1xindex>
+ gpu.return
+ }
+}
>From 97d5eceb4c15fbc508cc081d88cce3f540271caf Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 20 Aug 2025 15:25:41 +0000
Subject: [PATCH 2/4] Apply clang format.
---
.../lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 18 +++++++++---------
1 file changed, 9 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 380409afbc62e..32983152ef5bd 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -558,10 +558,10 @@ class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
Value loaded =
LLVM::LoadOp::create(rewriter, loc, srcOrDstFlatVecTy, basePtrLLVM);
- loaded.getDefiningOp()->setAttr("cache_control",
- xevm::LoadCacheControlAttr::get(
- ctxt, translateLoadXeGPUCacheHint(
- op.getL1Hint(), op.getL3Hint())));
+ loaded.getDefiningOp()->setAttr(
+ "cache_control", xevm::LoadCacheControlAttr::get(
+ ctxt, translateLoadXeGPUCacheHint(
+ op.getL1Hint(), op.getL3Hint())));
if (srcOrDstVecTy != srcOrDstFlatVecTy) {
loaded =
vector::ShapeCastOp::create(rewriter, loc, srcOrDstVecTy, loaded);
@@ -588,12 +588,12 @@ class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
srcFlatVec = vector::ShapeCastOp::create(rewriter, loc,
srcOrDstFlatVecTy, srcFlatVec);
}
- auto storeOp = LLVM::StoreOp::create(rewriter, loc, srcFlatVec, basePtrLLVM);
+ auto storeOp =
+ LLVM::StoreOp::create(rewriter, loc, srcFlatVec, basePtrLLVM);
storeOp.getOperation()->setAttr(
- "cache_control",
- xevm::StoreCacheControlAttr::get(ctxt,
- translateStoreXeGPUCacheHint(
- op.getL1Hint(), op.getL3Hint())));
+ "cache_control", xevm::StoreCacheControlAttr::get(
+ ctxt, translateStoreXeGPUCacheHint(
+ op.getL1Hint(), op.getL3Hint())));
rewriter.eraseOp(op);
}
return success();
>From 148df78d506d26a26e5974848dce42df1c12ba8b Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 20 Aug 2025 15:29:01 +0000
Subject: [PATCH 3/4] Remove commented out code.
---
mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 6 ------
1 file changed, 6 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 32983152ef5bd..89f40c22e7a68 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -380,8 +380,6 @@ class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtr);
auto elemType = tdescTy.getElementType();
auto elemBitSize = elemType.getIntOrFloatBitWidth();
- // auto elemBitSizeAttr = rewriter.getIntegerAttr(rewriter.getI32Type(),
- // elemBitSize);
Value elemByteSize = arith::ConstantIntOp::create(
rewriter, loc, rewriter.getI32Type(), elemBitSize / 8);
Value surfaceW =
@@ -695,8 +693,6 @@ class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
};
xevm::ElemType precATy = encodePrecision(aTy.getElementType());
xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
- // auto precA = xevm::ElemTypeAttr::get(ctxt, precATy);
- // auto precB = xevm::ElemTypeAttr::get(ctxt, precBTy);
Value c = op.getAcc();
if (!c) {
auto elementTy = resultType.getElementType();
@@ -714,8 +710,6 @@ class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
auto cvecty = cast<VectorType>(c.getType());
xevm::ElemType precCTy = encodePrecision(cvecty.getElementType());
xevm::ElemType precDTy = encodePrecision(resultType.getElementType());
- // auto precC = xevm::ElemTypeAttr::get(ctxt, precCTy);
- // auto precD = xevm::ElemTypeAttr::get(ctxt, precDTy);
VectorType cNty =
VectorType::get(cvecty.getNumElements(), cvecty.getElementType());
if (cvecty != cNty)
>From edd191c9939255c1639e6fbeea0ed7e6786634a8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 20 Aug 2025 17:59:49 +0000
Subject: [PATCH 4/4] Remove dead code.
---
mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 13 -------------
1 file changed, 13 deletions(-)
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index 89f40c22e7a68..776380974c549 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -56,19 +56,6 @@ static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
llvm_unreachable("Unknown XeGPU memory space.");
}
-template <typename T>
-std::tuple<bool, int32_t, int32_t> checkAllLinear(SmallVector<T> denseAttr) {
- assert(!denseAttr.empty());
- const int32_t intercept{static_cast<int32_t>(denseAttr[0])};
- if (denseAttr.size() < 2)
- return {true, 0, intercept};
- const T slope{denseAttr[1] - denseAttr[0]};
- for (size_t i = 1; i < denseAttr.size(); ++i)
- if (denseAttr[i] - denseAttr[i - 1] != slope)
- return {false, 0, 0};
- return {true, static_cast<int32_t>(slope), intercept};
-}
-
VectorType encodeVectorTypeTo(VectorType currentVecType, Type toElemType) {
auto elemType = currentVecType.getElementType();
auto currentBitWidth = elemType.getIntOrFloatBitWidth();
More information about the Mlir-commits
mailing list