[Mlir-commits] [mlir] [MLIR][XeGPU] Fix memref-i4 offsets computation (PR #233474)
Dmitry Chigarev
llvmlistbot at llvm.org
Mon Oct 19 09:05:16 PDT 2026
https://github.com/dchigarev created https://github.com/llvm/llvm-project/pull/233474
None
>From e5991313c9ae03e4ee619b895e4fadb1982c8fa6 Mon Sep 17 00:00:00 2001
From: Dmitry Chigarev <dmitry.chigarev at intel.com>
Date: Mon, 19 Oct 2026 16:03:01 +0000
Subject: [PATCH] [MLIR][XeGPU] Fix memref-i4 offsets computation
Signed-off-by: Dmitry Chigarev <dmitry.chigarev at intel.com>
---
.../mlir/Dialect/XeGPU/Utils/XeGPUUtils.h | 7 ++++
.../Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp | 14 ++------
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 8 ++---
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 33 +++++++++++++++++++
.../XeGPUToXeVM/materializecast.mlir | 20 +++++++++++
.../test/Dialect/XeGPU/peephole-optimize.mlir | 20 +++++++++++
6 files changed, 85 insertions(+), 17 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index ab526793f2aa8..9bf4b7eafd98e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -241,6 +241,13 @@ bool requireTranspose(const DistributeLayoutAttr layout,
/// Returns true if `type` has a static shape and static strides.
bool hasStaticShapeAndStrides(MemRefType type);
+/// Returns `baseAddr + elemOffset * bitWidth / 8`, i.e. advances the byte
+/// address `baseAddr` by `elemOffset` elements of `memrefTy`'s element type.
+/// Sub-byte element types (e.g. i4) are supported. `baseAddr` and `elemOffset`
+/// must have the same index or integer type, which is also the result type.
+Value addElemOffsetInBytes(OpBuilder &builder, Location loc, Value baseAddr,
+ Value elemOffset, MemRefType memrefTy);
+
// Check if dst shape is an expansion of src shape by inserting unit dimensions.
bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
SmallVector<int64_t> &expandedUnitDims);
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index c0b3adf8402e0..a2feb2425b8d9 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -1829,17 +1829,9 @@ struct ConvertXeGPUToXeVMPass
auto offsetCasted =
arith::IndexCastUIOp::create(builder, loc, type, offset);
- // Compute the final address: base address + byte offset
- auto byteSize = arith::ConstantOp::create(
- builder, loc, type,
- builder.getIntegerAttr(type,
- memrefTy.getElementTypeBitWidth() / 8));
- auto byteOffset =
- arith::MulIOp::create(builder, loc, offsetCasted, byteSize);
- auto addrWithOffset =
- arith::AddIOp::create(builder, loc, addrCasted, byteOffset);
-
- return addrWithOffset.getResult();
+ // Compute the final address: base address + byte offset.
+ return xegpu::addElemOffsetInBytes(builder, loc, addrCasted,
+ offsetCasted, memrefTy);
}
return {};
};
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index b430a5430c3c7..8d97441465ccb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -331,12 +331,8 @@ class XeGPUCreateNdDescOpPattern final
// Compute base = aligned pointer + offset (in bytes).
Value alignedPtr = memref::ExtractAlignedPointerAsIndexOp::create(
rewriter, loc, meta.getBaseBuffer());
- Value elemBytes = arith::ConstantIndexOp::create(
- rewriter, loc, memrefType.getElementTypeBitWidth() / 8);
- Value offsetBytes =
- arith::MulIOp::create(rewriter, loc, meta.getOffset(), elemBytes);
- Value baseIdx =
- arith::AddIOp::create(rewriter, loc, alignedPtr, offsetBytes);
+ Value baseIdx = xegpu::addElemOffsetInBytes(
+ rewriter, loc, alignedPtr, meta.getOffset(), memrefType);
source = arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(),
baseIdx);
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 406712b3f5b1b..ca2b5c36d4d9c 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -840,6 +840,39 @@ bool xegpu::hasStaticShapeAndStrides(MemRefType type) {
llvm::none_of(strides, ShapedType::isDynamic);
}
+Value xegpu::addElemOffsetInBytes(OpBuilder &builder, Location loc,
+ Value baseAddr, Value elemOffset,
+ MemRefType memrefTy) {
+ Type type = elemOffset.getType();
+ assert(baseAddr.getType() == type &&
+ "expected base address and offset of the same type");
+ auto makeConst = [&](int64_t value) -> Value {
+ return arith::ConstantOp::create(builder, loc, type,
+ builder.getIntegerAttr(type, value));
+ };
+
+ // finalAddr = baseAddr + elemOffset * bitWidth / 8
+ unsigned elemBits = memrefTy.getElementTypeBitWidth();
+ Value byteOffset;
+ if (elemBits % 8 == 0) {
+ // For byte-sized elements divide by 8 statically so that the division is
+ // folded into the constant:
+ // finalAddr = baseAddr + elemOffset * (bitWidth / 8)
+ byteOffset = arith::MulIOp::create(builder, loc, elemOffset,
+ makeConst(elemBits / 8));
+ } else {
+ // For sub-byte elements multiply first and then divide by 8, otherwise the
+ // byte size would be zero (e.g. 4 / 8 == 0 for i4):
+ // finalAddr = baseAddr + (elemOffset * bitWidth) / 8
+ Value bitSize = makeConst(elemBits);
+ Value eight = makeConst(8);
+ byteOffset = arith::DivUIOp::create(
+ builder, loc, arith::MulIOp::create(builder, loc, elemOffset, bitSize),
+ eight);
+ }
+ return arith::AddIOp::create(builder, loc, baseAddr, byteOffset);
+}
+
// Check if dst shape is an expansion of src shape by inserting unit dimensions.
// Returns true if all dimensions in src match corresponding dimensions in dst
// (after skipping unit dimensions), and populates expandedUnitDims with the
diff --git a/mlir/test/Conversion/XeGPUToXeVM/materializecast.mlir b/mlir/test/Conversion/XeGPUToXeVM/materializecast.mlir
index 969c369cd17e8..cf2647c6b8e6b 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/materializecast.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/materializecast.mlir
@@ -90,3 +90,23 @@ gpu.module @materializecast [#xevm.target<chip = "pvc">]{
gpu.return
}
}
+
+// -----
+gpu.module @materializecast {
+ // The memref offset counts i4 elements, two per byte.
+ // CHECK-LABEL: gpu.func @materialize_subbyte_memref_offset
+ // CHECK: %[[BASE:.*]], %[[OFF:.*]], %{{.*}}:2, %{{.*}}:2 = memref.extract_strided_metadata
+ // CHECK: %[[PTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE]]
+ // CHECK: %[[PTR64:.*]] = arith.index_castui %[[PTR]] : index to i64
+ // CHECK: %[[OFF64:.*]] = arith.index_castui %[[OFF]] : index to i64
+ // CHECK-DAG: %[[C4:.*]] = arith.constant 4 : i64
+ // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : i64
+ // CHECK: %[[BITS:.*]] = arith.muli %[[OFF64]], %[[C4]] : i64
+ // CHECK: %[[BYTES:.*]] = arith.divui %[[BITS]], %[[C8]] : i64
+ // CHECK: arith.addi %[[PTR64]], %[[BYTES]] : i64
+ gpu.func @materialize_subbyte_memref_offset(%src: memref<32x512xi4, strided<[512, 1], offset: ?>>) kernel {
+ %t = xegpu.create_nd_tdesc %src : memref<32x512xi4, strided<[512, 1], offset: ?>> -> !xegpu.tensor_desc<1x64xi4>
+ %v = xegpu.load_nd %t[0, 0] : !xegpu.tensor_desc<1x64xi4> -> vector<4xi4>
+ gpu.return
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 11c58d28377d2..0ca92d0fc972a 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -564,3 +564,23 @@ gpu.func @transpose_dynamic_offset(%arg0: memref<64x64xf16, strided<[64, 1], off
gpu.return %2 : vector<16x16xf16>
}
}
+
+// -----
+// The memref offset of an i4 source counts elements, two per byte.
+#l = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 8], order = [0, 1]>
+// CHECK-LABEL: gpu.func @i4_transpose_memref_offset(
+// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index
+// CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[BASE:.*]], %[[OFFSET:.*]], %{{.+}}:2, %{{.+}}:2 = memref.extract_strided_metadata
+// CHECK: %[[PTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE]]
+// CHECK: %[[BITS:.*]] = arith.muli %[[OFFSET]], %[[C4]] : index
+// CHECK: %[[BYTES:.*]] = arith.divui %[[BITS]], %[[C8]] : index
+// CHECK: arith.addi %[[PTR]], %[[BYTES]] : index
+// CHECK: xegpu.create_nd_tdesc {{.*}} -> !xegpu.tensor_desc<32x8xi32
+gpu.module @xevm_module {
+gpu.func @i4_transpose_memref_offset(%src: memref<512x2048xi4, strided<[2048, 1], offset: ?>>, %i: index, %j: index) -> vector<32x64xi4> {
+ %t = xegpu.create_nd_tdesc %src : memref<512x2048xi4, strided<[2048, 1], offset: ?>> -> !xegpu.tensor_desc<32x64xi4, #l>
+ %v = xegpu.load_nd %t[%i, %j] <{layout = #l}> : !xegpu.tensor_desc<32x64xi4, #l> -> vector<32x64xi4>
+ gpu.return %v : vector<32x64xi4>
+}
+}
More information about the Mlir-commits
mailing list