[Mlir-commits] [mlir] [MLIR][XeGPU] Enable strict property assembly format (PR #217281)
Mehdi Amini
llvmlistbot at llvm.org
Wed Aug 19 03:50:11 PDT 2026
https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/217281
Enable strict property assembly format mode for XeGPU. Spell the DPAS and atomic layout properties directly in custom assembly formats so their syntax remains explicit under strict parsing.
Update XeGPU tests to use the direct layout spelling for those operations.
Assisted-by: Codex
>From 0c30408dcac4cc05ff224f187a1ae5295c14367a Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 6 May 2026 16:59:11 -0700
Subject: [PATCH] [MLIR][XeGPU] Enable strict property assembly format
Enable strict property assembly format mode for XeGPU. Spell the DPAS and
atomic layout properties directly in custom assembly formats so their syntax
remains explicit under strict parsing.
Update XeGPU tests to use the direct layout spelling for those operations.
Assisted-by: Codex
---
.../mlir/Dialect/XeGPU/IR/XeGPUDialect.td | 1 +
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 24 ++-
.../XeGPUToXeVM/loadstore_matrix.mlir | 4 +-
mlir/test/Dialect/XeGPU/invalid.mlir | 32 ++--
.../test/Dialect/XeGPU/peephole-optimize.mlir | 16 +-
.../XeGPU/propagate-layout-inst-data.mlir | 26 +--
.../XeGPU/propagate-layout-subgroup.mlir | 35 ++--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 14 +-
.../XeGPU/resolve-layout-conflicts.mlir | 14 +-
.../XeGPU/sg-to-lane-distribute-unit.mlir | 37 ++--
.../Dialect/XeGPU/sg-to-lane-distribute.mlir | 89 ++++------
mlir/test/Dialect/XeGPU/transform-ops.mlir | 6 +-
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 158 ++++++++--------
.../Dialect/XeGPU/xegpu-recover-layout.mlir | 32 ++--
.../Dialect/XeGPU/xegpu-unroll-patterns.mlir | 8 +-
.../XeGPU/xegpu-wg-to-sg-elemwise.mlir | 30 ++--
.../test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir | 49 +++--
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir | 168 ++++++++----------
.../Dialect/XeGPU/WG/simple_gemm.mlir | 26 +--
19 files changed, 369 insertions(+), 400 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
index b1490c7742a26..652cf18ade580 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td
@@ -36,6 +36,7 @@ def XeGPU_Dialect : Dialect {
let useDefaultTypePrinterParser = true;
let useDefaultAttributePrinterParser = true;
+ let useStrictPropertiesInAssemblyFormat = 1;
let extraClassDeclaration = [{
/// Checks if the given memref type represents shared local memory (SLM).
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 49b98922cee4c..a643e21b947e0 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1010,11 +1010,12 @@ def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>,
Example 1 (Workgroup level, 2D):
```mlir
- %d = xegpu.dpas %a, %b, %c <{
- layout_a = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 128]>,
- layout_b = #xegpu.layout<sg_layout = [4, 8], sg_data = [128, 16]>,
- layout_cd = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16]>}
- : vector<64x128xf16>, vector<128x128xf16>, vector<64x128xf32> -> vector<64x128xf32>
+ %d = xegpu.dpas %a, %b, %c
+ layout_a(#xegpu.layout<sg_layout = [4, 8], sg_data = [16, 128]>)
+ layout_b(#xegpu.layout<sg_layout = [4, 8], sg_data = [128, 16]>)
+ layout_cd(#xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16]>)
+ : vector<64x128xf16>, vector<128x128xf16>, vector<64x128xf32>
+ -> vector<64x128xf32>
```
Example 2 (Lane level, 1D):
@@ -1077,7 +1078,12 @@ def XeGPU_DpasOp : XeGPU_Op<"dpas", [Pure, AllElementTypesMatch<["lhs", "rhs"]>,
}];
let assemblyFormat = [{
- $lhs `,` $rhs (`,` $acc^)? attr-dict `:` type($lhs)`,` type($rhs) (`,` type($acc)^)? `->` type($result)
+ $lhs `,` $rhs (`,` $acc^)?
+ (`layout_a` `(` $layout_a^ `)`)?
+ (`layout_b` `(` $layout_b^ `)`)?
+ (`layout_cd` `(` $layout_cd^ `)`)?
+ attr-dict `:`
+ type($lhs)`,` type($rhs) (`,` type($acc)^)? `->` type($result)
}];
let hasVerifier = 1;
@@ -1141,7 +1147,9 @@ def XeGPU_AtomicRMWOp: XeGPU_Op<"atomic_rmw", [Pure,
let results = (outs XeGPU_ValueType:$result);
let assemblyFormat = [{
- $kind $tensorDesc `,` $mask `,` $value attr-dict `:`
+ $kind $tensorDesc `,` $mask `,` $value
+ (`layout` `(` $layout^ `)`)?
+ attr-dict `:`
qualified(type($tensorDesc)) `,` type($mask) `,` type($value) `->` type($result)
}];
}
@@ -1678,7 +1686,7 @@ def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
$a `,` $b (`,` $acc^)?
(`scale_a` `=` $scale_a^)?
(`scale_b` `=` $scale_b^)?
- attr-dict `:` functional-type(operands, results)
+ prop-dict attr-dict `:` functional-type(operands, results)
}];
let hasVerifier = 1;
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
index 07fb09fa2c24b..b7b8ab8c393e5 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_matrix.mlir
@@ -262,12 +262,12 @@ gpu.module @test_kernel [#xevm.target<chip = "pvc">] {
%c16 = arith.constant 16 : index
%c48 = arith.constant 48 : index
- %1 = xegpu.load_matrix %0[%c16, %c48] {subgroup_block_io}: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index -> vector<8xf16>
+ %1 = xegpu.load_matrix %0[%c16, %c48] <{subgroup_block_io}>: !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index -> vector<8xf16>
//CHECK: %[[storeDataI16:.*]] = vector.bitcast %[[loaded]] : vector<8xf16> to vector<8xi16>
//CHECK: xevm.blockstore %[[ptr2:.*]], %[[storeDataI16]] : (!llvm.ptr<3>, vector<8xi16>)
- xegpu.store_matrix %1, %0[%c16, %c48] {subgroup_block_io}: vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index
+ xegpu.store_matrix %1, %0[%c16, %c48] <{subgroup_block_io}>: vector<8xf16>, !xegpu.mem_desc<32x64xf16, #xegpu.mem_layout<block = [16, 16]>>, index, index
gpu.return %1: vector<8xf16>
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index dc68f5136b4a9..93e4ecd1d88f4 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -482,7 +482,7 @@ func.func @layout_rank_mismatch_sg_lane(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected sg_layout and lane_layout to have the same rank}}
- {layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [16, 2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}
+ <{layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [16, 2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -493,7 +493,7 @@ func.func @layout_rank_mismatch_sg_inst(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected sg_layout and inst_data to have the same rank}}
- {layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [16, 2, 1], inst_data = [16, 2]>}
+ <{layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [16, 2, 1], inst_data = [16, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -504,7 +504,7 @@ func.func @layout_rank_mismatch_inst_lane(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected inst_data and lane_layout to have the same rank}}
- {layout = #xegpu.layout<inst_data = [16, 2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}
+ <{layout = #xegpu.layout<inst_data = [16, 2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -515,7 +515,7 @@ func.func @layout_rank_mismatch_lane_data(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected lane_data and lane_layout to have the same rank}}
- {layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2, 1]>}
+ <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [8, 1], lane_data = [1, 2, 1]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -526,7 +526,7 @@ func.func @layout_rank_mismatch_sg_data(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected sg_data and sg_layout to have the same rank}}
- {layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [16, 2, 1], inst_data = [16, 2]>}
+ <{layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [16, 2, 1], inst_data = [16, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -546,7 +546,7 @@ func.func @layout_sg_data_missing(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{sg_layout and sg_data must be used together}}
- {layout = #xegpu.layout<sg_layout = [2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}
+ <{layout = #xegpu.layout<sg_layout = [2, 1], lane_layout = [8, 1], lane_data = [1, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -557,7 +557,7 @@ func.func @layout_lane_data_missing(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{lane_layout and lane_data must be used together}}
- {layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1]>}
+ <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -568,7 +568,7 @@ func.func @layout_order_without_layout(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected sg_layout/lane_layout being used with order}}
- {layout = #xegpu.layout<inst_data = [16, 2], order = [0, 1]>}
+ <{layout = #xegpu.layout<inst_data = [16, 2], order = [0, 1]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -579,7 +579,7 @@ func.func @layout_order_rank_mismatch_sg(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected order and sg_layout to have the same rank}}
- {layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [16, 2], order = [0, 1, 2]>}
+ <{layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [16, 2], order = [0, 1, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -590,7 +590,7 @@ func.func @layout_order_rank_mismatch_lane(%src: memref<?xf32>) {
%mask = arith.constant dense<1>: vector<4xi1>
%2 = xegpu.load %src[%offsets], %mask
// expected-error at below {{expected order and lane_layout to have the same rank}}
- {layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2], order = [0, 1, 2]>}
+ <{layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2], order = [0, 1, 2]>}>
: memref<?xf32>, vector<4xindex>, vector<4xi1> -> vector<4xf32>
return
}
@@ -665,7 +665,7 @@ func.func @store_mem_desc_1d_exceeds_shape(%arg0: !xegpu.mem_desc<16xf16>, %arg1
// -----
func.func @simt_store_matrix_vector_nonlinear(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [32, 1]>>, %arg1: vector<2x16xf32>) {
// expected-error at +1 {{With subgroup_block_io, accessed data must be contiguous and coalesced}}
- xegpu.store_matrix %arg1, %arg0[0, 0] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} :
+ xegpu.store_matrix %arg1, %arg0[0, 0] <{subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
vector<2x16xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [32, 1]>>
return
}
@@ -673,7 +673,7 @@ func.func @simt_store_matrix_vector_nonlinear(%arg0: !xegpu.mem_desc<32x32xf32,
// -----
func.func @simt_store_matrix_vector_noncoalesced(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [1, 16]>>, %arg1: vector<16x2xf32>) {
// expected-error at +1 {{With subgroup_block_io, the distributed dimensions must be contiguous}}
- xegpu.store_matrix %arg1, %arg0[0, 0] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} :
+ xegpu.store_matrix %arg1, %arg0[0, 0] <{subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}> :
vector<16x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [1, 16]>>
return
}
@@ -681,7 +681,7 @@ func.func @simt_store_matrix_vector_noncoalesced(%arg0: !xegpu.mem_desc<32x32xf3
// -----
func.func @simt_store_matrix_vector_noncoalesced(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [32, 1], block = [1, 17]>>, %arg1: vector<16x2xf32>) {
// expected-error at +1 {{With subgroup_block_io, the block shape must match the lane layout}}
- xegpu.store_matrix %arg1, %arg0[0, 0] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+ xegpu.store_matrix %arg1, %arg0[0, 0] <{subgroup_block_io, layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
vector<16x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [32, 1], block = [1, 17]>>
return
}
@@ -777,7 +777,7 @@ func.func @dpas_mx_scale_k_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf
#layout_cd = #xegpu.layout<sg_layout = [1, 1], sg_data = [8, 16]>
func.func @dpas_mx_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{A shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b {layout_a = #layout_a, layout_b = #layout_b, layout_cd = #layout_cd} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b <{layout_a = #layout_a, layout_b = #layout_b, layout_cd = #layout_cd}> : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<8x16xf32>
return
}
@@ -785,7 +785,7 @@ func.func @dpas_mx_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector
#layout_a_scale_invalid = #xegpu.layout<sg_layout = [1, 1], sg_data = [5, 3]>
func.func @dpas_mx_scale_a_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a_val: vector<8x2xf8E8M0FNU>) {
// expected-error at +1 {{ScaleA shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val {layout_a_scale = #layout_a_scale_invalid} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>) -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val <{layout_a_scale = #layout_a_scale_invalid}> : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>) -> vector<8x16xf32>
return
}
@@ -793,7 +793,7 @@ func.func @dpas_mx_scale_a_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b
#layout_b_scale_invalid = #xegpu.layout<sg_layout = [1, 1], sg_data = [3, 11]>
func.func @dpas_mx_scale_b_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a_val: vector<8x2xf8E8M0FNU>, %scale_b_val: vector<2x16xf8E8M0FNU>) {
// expected-error at +1 {{ScaleB shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val {layout_b_scale = #layout_b_scale_invalid} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val <{layout_b_scale = #layout_b_scale_invalid}> : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
return
}
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index fa5ff9ca56d34..cb1430d815e85 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -10,7 +10,7 @@
// CHECK: %[[BDESC:.*]] = xegpu.create_nd_tdesc %[[T0]], shape : [64, %[[C32]]], strides : [%[[C32]], 1] : i64
// CHECK-SAME: -> !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>>
// CHECK-NEXT: %[[B:.*]] = xegpu.load_nd %[[BDESC]][%{{.*}}, %[[C16]]]
-// CHECK-SAME: {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}
+// CHECK-SAME: <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}>
// CHECK-SAME: : !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>>
// CHECK-SAME: -> vector<16x8xi32>
// CHECK: %[[BITCAST:.*]] = vector.bitcast %[[B]] : vector<16x8xi32> to vector<16x16xf16>
@@ -38,7 +38,7 @@ gpu.func @no_scf(%arg0: memref<64x64xf16>, %arg1: vector<8x16xf16>) -> vector<8x
// CHECK: %[[T1:.*]] = xegpu.create_nd_tdesc %[[T0]], shape : [64, %[[C16]]], strides : [%[[C16]], 1] : i64
// CHECK-SAME: -> !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>>
// CHECK: %[[T2:.*]] = xegpu.load_nd %[[T1]][%{{.*}}, %[[C16]]]
-// CHECK-SAME: {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}
+// CHECK-SAME: <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}>
// CHECK-SAME: : !xegpu.tensor_desc<16x8xi32, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>> -> vector<16x8xi32>
// CHECK: %[[T3:.*]] = vector.bitcast %[[T2]] : vector<16x8xi32> to vector<16x32xi8>
#a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>
@@ -284,12 +284,12 @@ gpu.func @array_length(%arg0: vector<8x16xf16>, %arg1: memref<256x256xf16>, %arg
%c32 = arith.constant 32 : index
%c256 = arith.constant 256 : index
%0 = xegpu.create_nd_tdesc %arg2 : memref<256x256xf32> -> !xegpu.tensor_desc<8x16xf32, #a>
- %1 = xegpu.load_nd %0[%c0, %c0] { layout = #a } : !xegpu.tensor_desc<8x16xf32, #a> -> vector<8x16xf32>
+ %1 = xegpu.load_nd %0[%c0, %c0] <{ layout = #a }> : !xegpu.tensor_desc<8x16xf32, #a> -> vector<8x16xf32>
%3 = xegpu.create_nd_tdesc %arg1 : memref<256x256xf16>
-> !xegpu.tensor_desc<32x16xf16, #b, #xegpu.block_tdesc_attr<array_length = 2 : i64>>
%4:4 = scf.for %arg3 = %c0 to %c256 step %c32 iter_args(%arg4 = %1, %arg5 = %1, %arg6 = %1, %arg7 = %1)
-> (vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>, vector<8x16xf32>) {
- %6 = xegpu.load_nd %3[%c0, %arg3] { layout = #b }
+ %6 = xegpu.load_nd %3[%c0, %arg3] <{ layout = #b }>
: !xegpu.tensor_desc<32x16xf16, #b, #xegpu.block_tdesc_attr<array_length = 2 : i64>> -> vector<2x32x16xf16>
%19 = vector.extract %6[0] { layout_result_0 = #b } : vector<32x16xf16> from vector<2x32x16xf16>
%20 = vector.extract %6[1] { layout_result_0 = #b } : vector<32x16xf16> from vector<2x32x16xf16>
@@ -377,7 +377,7 @@ gpu.module @xevm_test {
%offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
%mask = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1> : vector<16xi1>
- xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
@@ -418,7 +418,7 @@ gpu.module @xevm_test {
%offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
%mask = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1> : vector<16xi1>
- xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
@@ -457,7 +457,7 @@ gpu.module @xevm_test {
%offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
%mask = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1> : vector<16xi1>
- xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask <{layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
@@ -500,7 +500,7 @@ gpu.module @xevm_test {
%offset = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<0> : vector<16xindex>
%mask = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<1> : vector<16xi1>
- xegpu.store %reduce_bcast, %dst[%offset], %mask {layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ xegpu.store %reduce_bcast, %dst[%offset], %mask <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 2ab9d7e4f3f66..96c756cc8d66e 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -36,7 +36,7 @@ func.func @load_store_no_array_len(%arg0: memref<8x32xf32>, %arg1: memref<8x32xf
// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
-// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] layout_a(#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>) :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -342,9 +342,9 @@ gpu.module @test_collapse_dims [#xevm.target<O = 3, chip = "pvc">] {
%data = arith.constant dense<0.0> : vector<32x32xf32>
// CHECK: xegpu.store {{.*}} <{{{.*}}layout = #xegpu.layout<inst_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 2]>{{.*}}}> :
- xegpu.store %data, %ptr_i64[%1], %mask {
+ xegpu.store %data, %ptr_i64[%1], %mask <{
layout = #xegpu.layout<inst_data = [32, 32]>
- } : vector<32x32xf32>, i64, vector<32x32xindex>, vector<32x32xi1>
+ }> : vector<32x32xf32>, i64, vector<32x32xindex>, vector<32x32xi1>
gpu.return
}
@@ -413,7 +413,7 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: <{layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -455,7 +455,7 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x1024xf8E5M2>, %arg1: memref<1024x32xf
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<32x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: <{layout_a = #xegpu.layout<inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [64, 16], lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<16x1024xf4E2M1FN>, vector<1024x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -638,13 +638,13 @@ func.func @complete_scatter_load_inst_data(%src: memref<512xf32>) {
// A=[8,16]->[1,16]/[1,1]; B=[16,16]->[1,16]/[2,1] (VNNI); CD=[8,16]->[1,16]/[1,1].
gpu.module @test {
// CHECK-LABEL: func.func @complete_dpas_inst_data(
-// CHECK: xegpu.dpas %{{.*}}, %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
+// CHECK: xegpu.dpas %{{.*}}, %{{.*}}, %{{.*}} layout_a(#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>) : vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
func.func @complete_dpas_inst_data(%arg0: vector<8x16xf16>, %arg1: vector<16x16xf16>) {
%cst = arith.constant dense<0.000000e+00> : vector<8x16xf32>
- %0 = xegpu.dpas %arg0, %arg1, %cst {
- layout_a = #xegpu.layout<inst_data = [8, 16]>,
- layout_b = #xegpu.layout<inst_data = [16, 16]>,
- layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ %0 = xegpu.dpas %arg0, %arg1, %cst
+ layout_a(#xegpu.layout<inst_data = [8, 16]>)
+ layout_b(#xegpu.layout<inst_data = [16, 16]>)
+ layout_cd(#xegpu.layout<inst_data = [8, 16]>)
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
return
}
@@ -657,16 +657,16 @@ func.func @complete_dpas_inst_data(%arg0: vector<8x16xf16>, %arg1: vector<16x16x
gpu.module @test {
// CHECK-LABEL: func.func @complete_dpas_mx_inst_data(
// CHECK: xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{[0-9a-zA-Z]+}} scale_b = %{{[0-9a-zA-Z]+}}
-// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: <{layout_a = #xegpu.layout<inst_data = [8, 32], lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1], lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16], lane_layout = [1, 16], lane_data = [4, 1]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
func.func @complete_dpas_mx_inst_data(%arg0: vector<16x1024xf8E5M2>, %arg1: vector<1024x32xf8E5M2>,
%arg2: vector<16x32xf8E8M0FNU>, %arg3: vector<32x32xf8E8M0FNU>) {
%cst = arith.constant dense<0.000000e+00> : vector<16x32xbf16>
- %0 = xegpu.dpas_mx %arg0, %arg1, %cst scale_a = %arg2 scale_b = %arg3 {
+ %0 = xegpu.dpas_mx %arg0, %arg1, %cst scale_a = %arg2 scale_b = %arg3 <{
layout_a = #xegpu.layout<inst_data = [8, 32]>,
layout_a_scale = #xegpu.layout<inst_data = [8, 1]>,
layout_b = #xegpu.layout<inst_data = [32, 16]>,
layout_b_scale = #xegpu.layout<inst_data = [1, 16]>,
- layout_cd = #xegpu.layout<inst_data = [8, 16]>}
+ layout_cd = #xegpu.layout<inst_data = [8, 16]>}>
: (vector<16x1024xf8E5M2>, vector<1024x32xf8E5M2>, vector<16x32xbf16>, vector<16x32xf8E8M0FNU>, vector<32x32xf8E8M0FNU>) -> vector<16x32xbf16>
return
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index b4f9483cab895..9e24bec1219da 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -101,9 +101,9 @@ gpu.module @test {
// CHECK-SAME: : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>> -> vector<128x128xf16>
// CHECK: %[[DPAS_RES:.*]] = xegpu.dpas %[[A_LOADED]], %[[B_LOADED]]
- // CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>,
- // CHECK-SAME: layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>,
- // CHECK-SAME: layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} :
+ // CHECK-SAME: layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>)
+ // CHECK-SAME: layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>)
+ // CHECK-SAME: layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>) :
// CHECK-SAME: vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>
// CHECK: %[[TDESC_ST:.*]] = xegpu.create_nd_tdesc %[[CD_MEMREF]] : memref<128x128xf32> ->
@@ -171,7 +171,7 @@ gpu.module @test {
%reduce = vector.multi_reduction <add>, %bcast, %cst [1] : vector<32x128xf32> to vector<32xf32>
%mask = arith.constant dense<1>: vector<32xi1>
%offset = vector.step : vector<32xindex>
- xegpu.store %reduce, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>, dims = [1]>} : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
+ xegpu.store %reduce, %dst[%offset], %mask <{layout = #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>, dims = [1]>}> : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -204,7 +204,7 @@ gpu.module @test {
%reduce = vector.multi_reduction <add>, %bcast, %cst [1] : vector<32x128xf32> to vector<32xf32>
%mask = arith.constant dense<1>: vector<32xi1>
%offset = vector.step : vector<32xindex>
- xegpu.store %reduce, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>} : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
+ xegpu.store %reduce, %dst[%offset], %mask <{layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>}> : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -223,7 +223,7 @@ gpu.module @test {
%mask2 = arith.constant dense<true> : vector<8xi1>
%val = xegpu.load %arg0[%offset], %mask : memref<2048xf32, 1>, vector<8x256xindex>, vector<8x256xi1> -> vector<8x256xf32>
%reduce = vector.multi_reduction <add>, %val, %acc [1] : vector<8x256xf32> to vector<8xf32>
- xegpu.store %reduce, %arg0[%offset2], %mask2 { layout = #xegpu.slice<#data_layout, dims = [1]> } : vector<8xf32>, memref<2048xf32, 1>, vector<8xindex>, vector<8xi1>
+ xegpu.store %reduce, %arg0[%offset2], %mask2 <{ layout = #xegpu.slice<#data_layout, dims = [1]> }> : vector<8xf32>, memref<2048xf32, 1>, vector<8xindex>, vector<8xi1>
gpu.return
}
}
@@ -253,10 +253,9 @@ gpu.module @test {
// CHECK-NEXT: xegpu.load_nd %6[%arg3, %block_id_y] <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>}>
- // CHECK-NEXT: xegpu.dpas %{{.*}} {
- // CHECK-SAME: layout_a = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 128]>,
- // CHECK-SAME: layout_b = #xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>,
- // CHECK-SAME: layout_cd = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}
+ // CHECK-NEXT: xegpu.dpas %{{[^ ]+}}, %{{[^ ]+}}, %{{[^ ]+}} layout_a(#xegpu.layout<sg_layout = [2, 2], sg_data = [64, 128]>)
+ // CHECK-SAME: layout_b(#xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>)
+ // CHECK-SAME: layout_cd(#xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>)
// CHECK-SAME: : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
// CHECK-NEXT: scf.yield %{{.*}} : vector<128x128xf32>
@@ -294,7 +293,7 @@ gpu.module @test {
%5 = xegpu.load_nd %4[%block_id_x, %arg3] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
%6 = xegpu.create_nd_tdesc %arg1 : memref<8192x4096xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
%7 = xegpu.load_nd %6[%arg3, %block_id_y] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
- // CHECK: xegpu.dpas %{{.*}} {layout_a = #xegpu.layout<sg_layout = [2, 4], sg_data = [64, 128]>, layout_b = #xegpu.layout<sg_layout = [2, 4], sg_data = [128, 32]>, layout_cd = #xegpu.layout<sg_layout = [2, 4], sg_data = [64, 32]>} :
+ // CHECK: xegpu.dpas %{{.*}} layout_a(#xegpu.layout<sg_layout = [2, 4], sg_data = [64, 128]>) layout_b(#xegpu.layout<sg_layout = [2, 4], sg_data = [128, 32]>) layout_cd(#xegpu.layout<sg_layout = [2, 4], sg_data = [64, 32]>) :
%8 = xegpu.dpas %5, %7, %arg4 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
scf.yield %8 : vector<128x128xf32>
}
@@ -320,7 +319,7 @@ gpu.module @test {
%6 = xegpu.create_nd_tdesc %arg1 : memref<8192x4096xf16> -> !xegpu.tensor_desc<16x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
// CHECK: xegpu.load_nd %{{.*}}[%{{.*}}, %{{.*}}] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>}> :
%7 = xegpu.load_nd %6[%c0, %block_id_y] : !xegpu.tensor_desc<16x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<16x64xf16>
- // CHECK: xegpu.dpas %{{.*}} {layout_a = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>} :
+ // CHECK: xegpu.dpas %{{.*}} layout_a(#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>) layout_b(#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>) layout_cd(#xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>) :
%8 = xegpu.dpas %5, %7, %cst : vector<32x16xf16>, vector<16x64xf16>, vector<32x64xf32> -> vector<32x64xf32>
%3 = xegpu.create_nd_tdesc %arg2 : memref<2048x4096xf32> -> !xegpu.tensor_desc<32x64xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>}> :
@@ -395,7 +394,7 @@ gpu.module @test {
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x256xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>> -> vector<16x256xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
- // CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>} :
+ // CHECK-SAME: <{layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> :
// CHECK-SAME: (vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU>) -> vector<128x256xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x256xbf16> -> !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> : vector<128x256xbf16>, !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
@@ -436,7 +435,7 @@ gpu.module @test {
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<16x128xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
- // CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} :
+ // CHECK-SAME: <{layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
// CHECK-SAME: (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>) -> vector<128x128xf32>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
@@ -499,7 +498,7 @@ gpu.module @test {
%broadcasted = vector.broadcast %result : vector<64xf16> to vector<32x64xf16>
%transposed = vector.transpose %broadcasted, [1, 0] : vector<32x64xf16> to vector<64x32xf16>
%tdesc_output = xegpu.create_nd_tdesc %output : memref<64x32xf16> -> !xegpu.tensor_desc<64x32xf16>
- xegpu.store_nd %transposed, %tdesc_output[0, 0] { layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>} : vector<64x32xf16>, !xegpu.tensor_desc<64x32xf16>
+ xegpu.store_nd %transposed, %tdesc_output[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>}> : vector<64x32xf16>, !xegpu.tensor_desc<64x32xf16>
gpu.return
}
}
@@ -614,7 +613,7 @@ gpu.module @test {
%mask = arith.constant dense<1> : vector<256xi1>
// CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>, layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
// CHECK-SAME: : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}
+ xegpu.store %val, %dest[%offset], %mask <{chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}>
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
@@ -645,7 +644,7 @@ gpu.module @test {
%mask = arith.constant dense<1> : vector<256xi1>
// CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}>
// CHECK-SAME: : vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}
+ xegpu.store %val, %dest[%offset], %mask <{chunk_size = 1, l1_hint = #xegpu.cache_hint<cached>}>
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
@@ -689,7 +688,7 @@ gpu.func @forward_fill_loop_carried(%arg0: memref<128x64xf16>, %arg1: memref<64x
%3 = xegpu.load_nd %2[%c0, %c0] <{layout = #xegpu.layout<sg_layout = [1, 1], sg_data = [64, 64]>}> : !xegpu.tensor_desc<64x64xf16> -> vector<64x64xf16>
// CHECK: scf.for
%4:2 = scf.for %arg4 = %c0 to %c128 step %c64 iter_args(%acc0 = %cst, %acc = %cst_0) -> (vector<128x64xf32>, vector<128xf32>) {
- %8 = xegpu.dpas %1, %3, %acc0 {layout_a = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>, layout_b = #xegpu.layout<sg_layout = [1, 1], sg_data = [64, 64]>, layout_cd = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>} : vector<128x64xf16>, vector<64x64xf16>, vector<128x64xf32> -> vector<128x64xf32>
+ %8 = xegpu.dpas %1, %3, %acc0 layout_a(#xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>) layout_b(#xegpu.layout<sg_layout = [1, 1], sg_data = [64, 64]>) layout_cd(#xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>) : vector<128x64xf16>, vector<64x64xf16>, vector<128x64xf32> -> vector<128x64xf32>
%9 = vector.broadcast %acc : vector<128xf32> to vector<64x128xf32>
%10 = vector.transpose %9, [1, 0] : vector<64x128xf32> to vector<128x64xf32>
%11 = arith.mulf %8, %10 : vector<128x64xf32>
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 5d45ec79a5a8a..d1881dffce863 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -10,7 +10,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
-// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -40,7 +40,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T1]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
-// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK: %[[T4:.*]] = xegpu.dpas %[[T2]], %[[T3]], %[[CST]] layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf16> -> vector<8x16xf16>
// CHECK: %[[T5:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T4]], %[[T5]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -63,7 +63,7 @@ func.func @dpas_f16_result_f16(%arg0: memref<8x16xf16>, %arg1: memref<16x16xf16>
gpu.module @test {
// CHECK-LABEL: func.func @dpas_i8(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<8x32xi8>, %[[ARG1:[0-9a-zA-Z]+]]: vector<32x16xi8>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xi32>) {
-// CHECK: %[[T0:.*]] = xegpu.dpas %[[ARG0]], %[[ARG1]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+// CHECK: %[[T0:.*]] = xegpu.dpas %[[ARG0]], %[[ARG1]] layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
func.func @dpas_i8(%arg0: vector<8x32xi8>, %arg1: vector<32x16xi8>, %arg2: memref<8x16xi32>) {
%c0 = arith.constant 0 : index
@@ -344,7 +344,7 @@ gpu.module @test {
// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>,
// CHECK-SAME: %[[ARG3:[0-9a-zA-Z]+]]: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>) {
// CHECK: %[[T2:.*]] = arith.addf %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>} : vector<16x16xf16>
-// CHECK: %[[T3:.*]] = xegpu.dpas %{{.*}}, %[[T2]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
+// CHECK: %[[T3:.*]] = xegpu.dpas %{{.*}}, %[[T2]] layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) : vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
// CHECK-NEXT: xegpu.store_nd %[[T3]], %[[ARG2]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK-NEXT: xegpu.store_nd %[[T2]], %[[ARG3]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
func.func @binary_op_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor_desc<16x16xf16>, %arg2: !xegpu.tensor_desc<8x16xf32>, %arg3: !xegpu.tensor_desc<16x16xf16>) {
@@ -371,7 +371,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<8x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<8x16xf16>
// CHECK-NEXT: %[[T5:.*]] = xegpu.load_nd %[[T1]][%[[ARG3]], 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
-// CHECK-NEXT: %[[T6:.*]] = xegpu.dpas %[[T4]], %[[T5]], %[[ARG6]] {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-NEXT: %[[T6:.*]] = xegpu.dpas %[[T4]], %[[T5]], %[[ARG6]] layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) :
// CHECK-SAME: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
// CHECK-NEXT: scf.yield %[[T6]] : vector<8x16xf32>
// CHECK-NEXT: } {layout_operand_3 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
@@ -1010,7 +1010,7 @@ gpu.module @test {
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: <{layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
@@ -1052,7 +1052,7 @@ gpu.module @test {
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
-// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+// CHECK-SAME: <{layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
diff --git a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
index 6a625ff515029..658c568bb1731 100644
--- a/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
+++ b/mlir/test/Dialect/XeGPU/resolve-layout-conflicts.mlir
@@ -20,9 +20,9 @@ func.func @load_nd_with_conflicting_tensor_desc(%arg0: memref<64x64xf16>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<64x64xf16>
-> !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
- %1 = xegpu.load_nd %0 [%c0, %c0] {layout = #inst_data_8x16} : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
+ %1 = xegpu.load_nd %0 [%c0, %c0] <{layout = #inst_data_8x16}> : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
-> vector<16x16xf16>
- xegpu.prefetch_nd %0 [%c0, %c0] {layout = #inst_data_16x16} : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
+ xegpu.prefetch_nd %0 [%c0, %c0] <{layout = #inst_data_16x16}> : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
return
}
@@ -44,11 +44,11 @@ func.func @multiple_tensor_desc_conflicts(%arg0: memref<64x64xf16>) {
%c0 = arith.constant 0 : index
%tdesc1 = xegpu.create_nd_tdesc %arg0 : memref<64x64xf16>
-> !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
- %load1 = xegpu.load_nd %tdesc1 [%c0, %c0] {layout = #inst_data_8x16} : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
+ %load1 = xegpu.load_nd %tdesc1 [%c0, %c0] <{layout = #inst_data_8x16}> : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
-> vector<32x16xf16>
- %load2 = xegpu.load_nd %tdesc1 [%c0, %c0] {layout = #inst_data_32x16} : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
+ %load2 = xegpu.load_nd %tdesc1 [%c0, %c0] <{layout = #inst_data_32x16}> : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
-> vector<32x16xf16>
- xegpu.prefetch_nd %tdesc1 [%c0, %c0] {layout = #inst_data_16x16} : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
+ xegpu.prefetch_nd %tdesc1 [%c0, %c0] <{layout = #inst_data_16x16}> : !xegpu.tensor_desc<32x16xf16, #inst_data_8x16>
return
}
@@ -74,12 +74,12 @@ func.func @load_nd_with_conflicting_tensor_desc_in_loop(%arg0: memref<64x64xf16>
-> !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
%1:2 = scf.for %i = %c0 to %c4 step %c1 iter_args(%acc = %cst, %tdesc = %0)
-> (vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>) {
- %2 = xegpu.load_nd %tdesc [%c0, %c0] {layout = #inst_data_8x16} : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
+ %2 = xegpu.load_nd %tdesc [%c0, %c0] <{layout = #inst_data_8x16}> : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
-> vector<16x16xf16>
%3 = arith.addf %acc, %2 {layout_result_0 = #inst_data_8x16} : vector<16x16xf16>
scf.yield %3, %tdesc : vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
} {layout_result_0 = #inst_data_8x16}
- xegpu.prefetch_nd %0 [%c0, %c0] {layout = #inst_data_16x16} : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
+ xegpu.prefetch_nd %0 [%c0, %c0] <{layout = #inst_data_16x16}> : !xegpu.tensor_desc<16x16xf16, #inst_data_16x16>
return
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 8607d45d5828c..b1d52fa0b32eb 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -31,7 +31,7 @@ gpu.func @cerate_nd_tedesc_nonmemref_source(%arg0: ui64) {
gpu.func @load_nd() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
- %1 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %1 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
gpu.return
}
@@ -43,7 +43,7 @@ gpu.func @load_nd() {
gpu.func @load_nd_packed() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
- %1 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ %1 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
gpu.return
}
@@ -55,7 +55,7 @@ gpu.func @load_nd_packed() {
gpu.func @load_nd_transpose() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<16x8xf32>
- %1 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}
+ %1 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x8xf32> -> vector<16x8xf32>
gpu.return
}
@@ -67,7 +67,7 @@ gpu.func @load_nd_transpose() {
gpu.func @load_nd_array_length() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<array_length = 2>>
- %1 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %1 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<array_length = 2>> -> vector<64x16xf16>
gpu.return
}
@@ -82,9 +82,9 @@ gpu.func @store_nd() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
%1 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
- %2 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %2 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
- xegpu.store_nd %2, %1[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.store_nd %2, %1[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<16x16xf16>, !xegpu.tensor_desc<16x16xf16>
gpu.return
}
@@ -107,14 +107,11 @@ gpu.func @dpas() {
%0 = "some_op"() : () -> !xegpu.tensor_desc<8x16xf16>
%1 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
%5 = arith.constant dense<0.0> : vector<8x16xf32>
- %2 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %2 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %3 = xegpu.load_nd %1[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ %3 = xegpu.load_nd %1[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
- %4 = xegpu.dpas %2, %3, %5
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %4 = xegpu.dpas %2, %3, %5 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
%anchor = xegpu.convert_layout %4
<{
@@ -134,7 +131,7 @@ gpu.func @elementwise() {
%c0 = arith.constant 0 : index
%0 = arith.constant dense<1.0> : vector<16x16xf32>
%1 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf32>
- %2 = xegpu.load_nd %1[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %2 = xegpu.load_nd %1[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
%3 = arith.addf %0, %2
: vector<16x16xf32>
@@ -246,7 +243,7 @@ gpu.func @arith_constant_non_splat_2d_vertical_lanedata() {
gpu.func @prefetch_nd() {
%c0 = arith.constant 0 : index
%0 = "some_op"() : () -> !xegpu.tensor_desc<16x16xf16>
- xegpu.prefetch_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.prefetch_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf16>
gpu.return
}
@@ -1497,11 +1494,11 @@ gpu.func @xegpu_dpas_mx(%arg0: !xegpu.mem_desc<8x8xf16>, %arg1: !xegpu.mem_desc<
%scale_B = arith.constant dense<0.5> : vector<2x16xf8E8M0FNU>
%4 = xegpu.dpas_mx %A, %B scale_a = %scale_A scale_b = %scale_B
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ <{layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>,
- layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: (vector<8x16xf8E4M3FN>, vector<16x16xf8E4M3FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
%anchor = xegpu.convert_layout %4
@@ -1562,20 +1559,20 @@ gpu.func @convert_layout_partial_subgroup() {
%scale_a1 = vector.extract_strided_slice %cvt
{offsets = [8, 0], sizes = [8, 2], strides = [1, 1]}
: vector<16x4xf8E8M0FNU> to vector<8x2xf8E8M0FNU>
- %res0 = xegpu.dpas_mx %a, %b scale_a = %scale_a0 scale_b = %scale_b {
+ %res0 = xegpu.dpas_mx %a, %b scale_a = %scale_a0 scale_b = %scale_b <{
layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>,
layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>,
layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
- } : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
- %res1 = xegpu.dpas_mx %a, %b scale_a = %scale_a1 scale_b = %scale_b {
+ }> : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
+ %res1 = xegpu.dpas_mx %a, %b scale_a = %scale_a1 scale_b = %scale_b <{
layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>,
layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>,
layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
- } : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
+ }> : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
index 9eed7160c0baa..06fd5f13dcf22 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute.mlir
@@ -39,7 +39,7 @@ gpu.func @gemm(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x1024xbf16>, %ar
%2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> ->
!xegpu.tensor_desc<8x16xf32>
%3 = xegpu.load_nd %2[%0, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>
%5 = xegpu.create_nd_tdesc %arg0: memref<1024x1024xbf16>
-> !xegpu.tensor_desc<8x16xbf16>
@@ -48,21 +48,18 @@ gpu.func @gemm(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x1024xbf16>, %ar
%4 = scf.for %arg3 = %c0 to %c1024 step %c16 iter_args(%arg4 = %3) -> (vector<8x16xf32>) {
%7 = xegpu.load_nd %5[%0, %arg3]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xbf16> -> vector<8x16xbf16>
%8 = xegpu.load_nd %6[%arg3, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xbf16> -> vector<16x16xbf16>
- %9 = xegpu.dpas %7, %8, %arg4
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %9 = xegpu.dpas %7, %8, %arg4 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xbf16>, vector<16x16xbf16>, vector<8x16xf32> -> vector<8x16xf32>
scf.yield %9 : vector<8x16xf32>
}
- xegpu.store_nd %4, %2[%0, %1] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}: vector<8x16xf32>,
+ xegpu.store_nd %4, %2[%0, %1] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>: vector<8x16xf32>,
!xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -109,7 +106,7 @@ gpu.func @gemm_with_preop(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x1024
%2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> ->
!xegpu.tensor_desc<8x16xf32>
%3 = xegpu.load_nd %2[%0, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>
%5 = xegpu.create_nd_tdesc %arg0: memref<1024x1024xbf16>
-> !xegpu.tensor_desc<8x16xbf16>
@@ -118,22 +115,19 @@ gpu.func @gemm_with_preop(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x1024
%4 = scf.for %arg3 = %c0 to %c1024 step %c16 iter_args(%arg4 = %3) -> (vector<8x16xf32>) {
%7 = xegpu.load_nd %5[%0, %arg3]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xbf16> -> vector<8x16xbf16>
%preop = arith.addf %7, %cst : vector<8x16xbf16>
%8 = xegpu.load_nd %6[%arg3, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xbf16> -> vector<16x16xbf16>
- %9 = xegpu.dpas %preop, %8, %arg4
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %9 = xegpu.dpas %preop, %8, %arg4 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xbf16>, vector<16x16xbf16>, vector<8x16xf32> -> vector<8x16xf32>
scf.yield %9 : vector<8x16xf32>
}
- xegpu.store_nd %4, %2[%0, %1] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}: vector<8x16xf32>,
+ xegpu.store_nd %4, %2[%0, %1] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>: vector<8x16xf32>,
!xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -175,7 +169,7 @@ gpu.func @gemm_with_postop(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x102
%2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> ->
!xegpu.tensor_desc<8x16xf32>
%3 = xegpu.load_nd %2[%0, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>
%5 = xegpu.create_nd_tdesc %arg0: memref<1024x1024xbf16>
-> !xegpu.tensor_desc<8x16xbf16>
@@ -184,22 +178,19 @@ gpu.func @gemm_with_postop(%arg0: memref<1024x1024xbf16>, %arg1: memref<1024x102
%4 = scf.for %arg3 = %c0 to %c1024 step %c16 iter_args(%arg4 = %3) -> (vector<8x16xf32>) {
%7 = xegpu.load_nd %5[%0, %arg3]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xbf16> -> vector<8x16xbf16>
%8 = xegpu.load_nd %6[%arg3, %1]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xbf16> -> vector<16x16xbf16>
- %9 = xegpu.dpas %7, %8, %arg4
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %9 = xegpu.dpas %7, %8, %arg4 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xbf16>, vector<16x16xbf16>, vector<8x16xf32> -> vector<8x16xf32>
scf.yield %9 : vector<8x16xf32>
}
%postop = math.exp %4 : vector<8x16xf32>
- xegpu.store_nd %postop, %2[%0, %1] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}: vector<8x16xf32>,
+ xegpu.store_nd %postop, %2[%0, %1] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>: vector<8x16xf32>,
!xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -227,24 +218,21 @@ gpu.module @xevm_module{
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16>
-> !xegpu.tensor_desc<8x16xf16>
%1 = xegpu.load_nd %0[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> :
!xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
%2 = xegpu.create_nd_tdesc %arg1: memref<16x16xf16>
-> !xegpu.tensor_desc<16x16xf16>
%3 = xegpu.load_nd %2[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16>
-> vector<16x16xf16>
- %4 = xegpu.dpas %1, %3, %cst
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %4 = xegpu.dpas %1, %3, %cst layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
%5 = math.exp %4
: vector<8x16xf32>
%6 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32> ->
!xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %5, %6[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf32>,
+ xegpu.store_nd %5, %6[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>,
!xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -272,15 +260,15 @@ gpu.module @xevm_module{
%1 = arith.constant dense<1>: vector<16xi1>
%offset = arith.constant dense<12> : vector<16xindex>
%loaded = scf.if %pred -> (vector<16x8xf16>) {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
+ %3 = xegpu.load %src[%offset], %1 <{chunk_size=8,
layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
+ }> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
scf.yield %3 : vector<16x8xf16>
} else {
%3 = arith.constant dense<12.> : vector<16x8xf16>
scf.yield %3 : vector<16x8xf16>
}
- xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
+ xegpu.store %loaded, %src[%offset], %1 <{chunk_size=8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
gpu.return
}
}
@@ -302,10 +290,10 @@ gpu.module @xevm_module{
%1 = arith.constant dense<1>: vector<16xi1>
%offset = arith.constant dense<12> : vector<16xindex>
scf.if %pred {
- %3 = xegpu.load %src[%offset], %1 <{chunk_size=8}> {
+ %3 = xegpu.load %src[%offset], %1 <{chunk_size=8,
layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>
- } : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
- xegpu.store %3, %src[%offset], %1 <{chunk_size=8}> {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
+ }> : memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
+ xegpu.store %3, %src[%offset], %1 <{chunk_size=8, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
}
gpu.return
}
@@ -331,25 +319,22 @@ gpu.module @xevm_module{
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16>
-> !xegpu.tensor_desc<8x16xf16>
%1 = xegpu.load_nd %0[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
%2 = xegpu.create_nd_tdesc %arg1 : memref<16x8xi32>
-> !xegpu.tensor_desc<16x8xi32>
%3 = xegpu.load_nd %2[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}>
: !xegpu.tensor_desc<16x8xi32> -> vector<16x8xi32>
%4 = vector.bitcast %3
: vector<16x8xi32> to vector<16x16xf16>
%5 = vector.transpose %4, [1, 0]
: vector<16x16xf16> to vector<16x16xf16>
- %6 = xegpu.dpas %1, %5, %cst
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %6 = xegpu.dpas %1, %5, %cst layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
%7 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32>
-> !xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %6, %7[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<8x16xf32>,
+ xegpu.store_nd %6, %7[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xf32>,
!xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -412,9 +397,9 @@ gpu.module @xevm_module{
gpu.func @load_store_matrix_3(%arg0: !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
- %1 = xegpu.load_matrix %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} :
+ %1 = xegpu.load_matrix %arg0[%c0, %c1] <{subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}> :
!xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>, index, index -> vector<16x2xf32>
- xegpu.store_matrix %1, %arg0[%c0, %c1] {subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} :
+ xegpu.store_matrix %1, %arg0[%c0, %c1] <{subgroup_block_io, layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}> :
vector<16x2xf32>, !xegpu.mem_desc<32x32xf32, #xegpu.mem_layout<stride = [1, 32], block = [16, 1]>>, index, index
gpu.return
}
@@ -520,7 +505,7 @@ gpu.module @xevm_module {
%cst = arith.constant dense<0.0> : vector<16x16xf32>
%td = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
%ld = xegpu.load_nd %td[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
%r = scf.for %i = %c0 to %c16 step %c1 iter_args(%acc = %cst) -> (vector<16x16xf32>) {
%add = arith.addf %acc, %ld
@@ -530,7 +515,7 @@ gpu.module @xevm_module {
} {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
%td2 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
xegpu.store_nd %r, %td2[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
gpu.return
}
@@ -550,7 +535,7 @@ gpu.module @xevm_module {
%c1_i32 = arith.constant 1 : i32
%td0 = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
%ld0 = xegpu.load_nd %td0[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
%td1 = xegpu.create_nd_tdesc %arg1 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
%r:2 = scf.while (%arg2 = %ld0, %arg3 = %c0_i32) : (vector<16x16xf32>, i32) -> (vector<16x16xf32>, i32) {
@@ -559,16 +544,16 @@ gpu.module @xevm_module {
} do {
^bb0(%arg2: vector<16x16xf32>, %arg3: i32):
xegpu.store_nd %arg2, %td1[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
%next = arith.addi %arg3, %c1_i32 : i32
%ld_next = xegpu.load_nd %td0[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
scf.yield %ld_next, %next : vector<16x16xf32>, i32
}
xegpu.store_nd %r#0, %td1[%c0, %c0]
- {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/transform-ops.mlir b/mlir/test/Dialect/XeGPU/transform-ops.mlir
index 3daa74d223946..df33d7086cfcd 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -218,7 +218,7 @@ func.func @set_anchor_layout_dpas_a(%arg0: memref<4096x4096xf16>, %arg1: memref<
%4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
// CHECK: = xegpu.dpas
- // CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>}
+ // CHECK-SAME: layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>)
%6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
return
}
@@ -243,7 +243,7 @@ func.func @set_anchor_layout_dpas_b(%arg0: memref<4096x4096xf16>, %arg1: memref<
%4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
// CHECK: = xegpu.dpas
- // CHECK-SAME: {layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [16, 16]>}
+ // CHECK-SAME: layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [16, 16]>)
%6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
return
}
@@ -268,7 +268,7 @@ func.func @set_anchor_layout_dpas_c(%arg0: memref<4096x4096xf16>, %arg1: memref<
%4 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
// CHECK: = xegpu.dpas
- // CHECK-SAME: {layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>}
+ // CHECK-SAME: layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>)
%6 = xegpu.dpas %1, %3, %5 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index 2dc177cf8aa69..70211d838e7ac 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -15,7 +15,7 @@ gpu.module @test_kernel {
%n = arith.muli %block_id_y, %c32 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #c>
- %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #c}: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] <{layout = #c}>: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #a>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #b>
@@ -23,16 +23,16 @@ gpu.module @test_kernel {
iter_args(%arg2 = %c_init)
-> (vector<16x32xf32>) {
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
- %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #a}: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] <{layout = #a}>: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
- %b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #b}: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
+ %b = xegpu.load_nd %b_tdesc[%k, %c0] <{layout = #b}>: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
//CHECK-COUNT-8: xegpu.dpas {{.*}}
- %c = xegpu.dpas %a, %b, %arg2 {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
+ %c = xegpu.dpas %a, %b, %arg2 layout_a(#a) layout_b(#b) layout_cd(#c): vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
scf.yield %c
: vector<16x32xf32>
}
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
- xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #c}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
+ xegpu.store_nd %out, %c_tdesc[0, 0] <{layout = #c}>: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
gpu.return
}
}
@@ -52,7 +52,7 @@ gpu.module @test_kernel {
%n = arith.muli %block_id_y, %c32 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #l1>
- %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l1}: !xegpu.tensor_desc<16x32xf32, #l1> -> vector<16x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] <{layout = #l1}>: !xegpu.tensor_desc<16x32xf32, #l1> -> vector<16x32xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l1>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #l2>
@@ -60,16 +60,16 @@ gpu.module @test_kernel {
iter_args(%arg2 = %c_init)
-> (vector<16x32xf32>) {
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
- %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l1}: !xegpu.tensor_desc<16x32xf16, #l1> -> vector<16x32xf16>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] <{layout = #l1}>: !xegpu.tensor_desc<16x32xf16, #l1> -> vector<16x32xf16>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
- %b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #l2}: !xegpu.tensor_desc<32x32xf16, #l2> -> vector<32x32xf16>
+ %b = xegpu.load_nd %b_tdesc[%k, %c0] <{layout = #l2}>: !xegpu.tensor_desc<32x32xf16, #l2> -> vector<32x32xf16>
//CHECK-COUNT-8: xegpu.dpas {{.*}}
- %c = xegpu.dpas %a, %b, %arg2 {layout_a=#l1, layout_b = #l2, layout_cd = #l1}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
+ %c = xegpu.dpas %a, %b, %arg2 layout_a(#l1) layout_b(#l2) layout_cd(#l1): vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
scf.yield %c
: vector<16x32xf32>
}
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l1}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l1>
+ xegpu.store_nd %out, %c_tdesc[0, 0] <{layout = #l1}>: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l1>
gpu.return
}
}
@@ -92,7 +92,7 @@ gpu.module @test_kernel {
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<8x32xf32, #l1>
//CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>
- %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l1}: !xegpu.tensor_desc<8x32xf32, #l1> -> vector<8x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] <{layout = #l1}>: !xegpu.tensor_desc<8x32xf32, #l1> -> vector<8x32xf32>
%c2 = arith.constant 2 : index
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<8x16xf16, #l1>
@@ -102,15 +102,15 @@ gpu.module @test_kernel {
-> (vector<8x32xf32>) {
%a_off = arith.muli %k, %c2 : index
//CHECK: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %a = xegpu.load_nd %a_tdesc[%c0, %a_off] {layout = #l1}: !xegpu.tensor_desc<8x16xf16, #l1> -> vector<8x16xf16>
+ %a = xegpu.load_nd %a_tdesc[%c0, %a_off] <{layout = #l1}>: !xegpu.tensor_desc<8x16xf16, #l1> -> vector<8x16xf16>
//CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
- %b = xegpu.load_nd %b_tdesc[%a_off, %c0] {layout = #l2}: !xegpu.tensor_desc<16x32xf16, #l2> -> vector<16x32xf16>
- %c = xegpu.dpas %a, %b, %arg2 {layout_a=#l1, layout_b = #l2, layout_cd = #l1}: vector<8x16xf16>, vector<16x32xf16>, vector<8x32xf32> -> vector<8x32xf32>
+ %b = xegpu.load_nd %b_tdesc[%a_off, %c0] <{layout = #l2}>: !xegpu.tensor_desc<16x32xf16, #l2> -> vector<16x32xf16>
+ %c = xegpu.dpas %a, %b, %arg2 layout_a(#l1) layout_b(#l2) layout_cd(#l1): vector<8x16xf16>, vector<16x32xf16>, vector<8x32xf32> -> vector<8x32xf32>
scf.yield %c
: vector<8x32xf32>
}
//CHECK-COUNT-2: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l1}: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #l1>
+ xegpu.store_nd %out, %c_tdesc[0, 0] <{layout = #l1}>: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #l1>
gpu.return
}
}
@@ -131,7 +131,7 @@ gpu.module @test_kernel {
%n = arith.muli %block_id_y, %c32 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #c>
- %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #c}: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] <{layout = #c}>: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #a>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #b>
@@ -139,18 +139,18 @@ gpu.module @test_kernel {
iter_args(%arg2 = %c_init)
-> (vector<16x32xf32>) {
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
- %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #a}: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] <{layout = #a}>: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
- %b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #b}: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
+ %b = xegpu.load_nd %b_tdesc[%k, %c0] <{layout = #b}>: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
//CHECK-COUNT-4: math.exp {{.*}} : vector<8x16xf16>
%e = math.exp %a : vector<16x32xf16>
//CHECK-COUNT-8: xegpu.dpas {{.*}}
- %c = xegpu.dpas %e, %b, %arg2 {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
+ %c = xegpu.dpas %e, %b, %arg2 layout_a(#a) layout_b(#b) layout_cd(#c): vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
scf.yield %c
: vector<16x32xf32>
}
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
- xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #c}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
+ xegpu.store_nd %out, %c_tdesc[0, 0] <{layout = #c}>: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
gpu.return
}
}
@@ -172,14 +172,14 @@ gpu.module @test_kernel {
scf.for %k = %c0 to %c1024 step %c32 {
//CHECK-COUNT-8: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l}: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
- %b = xegpu.load_nd %b_tdesc[%c0, %k] {layout = #l}: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] <{layout = #l}>: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
+ %b = xegpu.load_nd %b_tdesc[%c0, %k] <{layout = #l}>: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
//CHECK-COUNT-4: arith.addf {{.*}} : vector<8x16xf16>
%c = arith.addf %a, %b : vector<16x32xf16>
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
- xegpu.store_nd %c, %c_tdesc[%c0, %k] {layout = #l}: vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #l>
+ xegpu.store_nd %c, %c_tdesc[%c0, %k] <{layout = #l}>: vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #l>
}
gpu.return
}
@@ -202,14 +202,14 @@ gpu.module @test_kernel {
scf.for %k = %c0 to %c1024 step %c32 {
//CHECK-COUNT-8: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8xf16> -> vector<8xf16>
- %a = xegpu.load_nd %a_tdesc[%k] {layout = #l}: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
- %b = xegpu.load_nd %b_tdesc[%k] {layout = #l}: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
+ %a = xegpu.load_nd %a_tdesc[%k] <{layout = #l}>: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
+ %b = xegpu.load_nd %b_tdesc[%k] <{layout = #l}>: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
//CHECK-COUNT-4: arith.addf {{.*}} : vector<8xf16>
%c = arith.addf %a, %b : vector<32xf16>
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8xf16>, !xegpu.tensor_desc<8xf16>
- xegpu.store_nd %c, %c_tdesc[%k] {layout = #l}: vector<32xf16>, !xegpu.tensor_desc<32xf16, #l>
+ xegpu.store_nd %c, %c_tdesc[%k] <{layout = #l}>: vector<32xf16>, !xegpu.tensor_desc<32xf16, #l>
}
gpu.return
}
@@ -225,12 +225,12 @@ gpu.module @test_kernel {
%block_id_x = gpu.block_id x
%m = arith.muli %block_id_x, %c64 : index
%0 = xegpu.create_nd_tdesc %a : memref<16x512xf32> -> !xegpu.tensor_desc<16x64xf32, #l>
- %1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<16x64xf32, #l> -> vector<16x64xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #l}>: !xegpu.tensor_desc<16x64xf32, #l> -> vector<16x64xf32>
// CHECK: vector.multi_reduction <add>, {{.*}}, [[ACC:%[0-9A-Za-z]+]] [0] : vector<16x16xf32> to vector<16xf32>
// CHECK-COUNT-3: vector.multi_reduction <add>, {{.*}}, [[ACC]] [0] : vector<16x16xf32> to vector<16xf32>
%2 = vector.multi_reduction <add>, %1, %acc [0]: vector<16x64xf32> to vector<64xf32>
%3 = xegpu.create_nd_tdesc %b : memref<512xf32> -> !xegpu.tensor_desc<64xf32, #r>
- xegpu.store_nd %2, %3[0] {layout = #r}: vector<64xf32>, !xegpu.tensor_desc<64xf32, #r>
+ xegpu.store_nd %2, %3[0] <{layout = #r}>: vector<64xf32>, !xegpu.tensor_desc<64xf32, #r>
gpu.return
}
}
@@ -250,7 +250,7 @@ gpu.module @test_kernel {
%m = arith.muli %block_id_x, %c32 : index
%n = arith.muli %block_id_y, %c32 : index
%0 = xegpu.create_nd_tdesc %a : memref<512x32xf32> -> !xegpu.tensor_desc<32x128xf32, #l>
- %1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<32x128xf32, #l> -> vector<32x128xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #l}>: !xegpu.tensor_desc<32x128xf32, #l> -> vector<32x128xf32>
// CHECK-COUNT-7: arith.addf {{.*}} : vector<16x16xf32>
// CHECK: vector.multi_reduction <add>, {{.*}} [1] : vector<16x16xf32> to vector<16xf32>
@@ -259,7 +259,7 @@ gpu.module @test_kernel {
%2 = vector.multi_reduction <add>, %1, %acc [1]: vector<32x128xf32> to vector<32xf32>
%3 = xegpu.create_nd_tdesc %b : memref<512xf32> -> !xegpu.tensor_desc<32xf32, #r>
- xegpu.store_nd %2, %3[0] {layout = #r}: vector<32xf32>, !xegpu.tensor_desc<32xf32, #r>
+ xegpu.store_nd %2, %3[0] <{layout = #r}>: vector<32xf32>, !xegpu.tensor_desc<32xf32, #r>
gpu.return
}
}
@@ -274,11 +274,11 @@ gpu.module @test_kernel {
%block_id_x = gpu.block_id x
%m = arith.muli %block_id_x, %c64 : index
%0 = xegpu.create_nd_tdesc %a : memref<512xf32> -> !xegpu.tensor_desc<64xf32, #r>
- %1 = xegpu.load_nd %0[0] {layout = #r}: !xegpu.tensor_desc<64xf32, #r> -> vector<64xf32>
+ %1 = xegpu.load_nd %0[0] <{layout = #r}>: !xegpu.tensor_desc<64xf32, #r> -> vector<64xf32>
// CHECK-COUNT-4: vector.broadcast {{.*}} : vector<16xf32> to vector<16x16xf32>
%2 = vector.broadcast %1 : vector<64xf32> to vector<16x64xf32>
%3 = xegpu.create_nd_tdesc %b : memref<16x512xf32> -> !xegpu.tensor_desc<16x64xf32, #l>
- xegpu.store_nd %2, %3[0, 0] {layout = #l}: vector<16x64xf32>, !xegpu.tensor_desc<16x64xf32, #l>
+ xegpu.store_nd %2, %3[0, 0] <{layout = #l}>: vector<16x64xf32>, !xegpu.tensor_desc<16x64xf32, #l>
gpu.return
}
}
@@ -293,12 +293,12 @@ gpu.module @test_kernel {
%block_id_x = gpu.block_id x
%m = arith.muli %block_id_x, %c32 : index
%0 = xegpu.create_nd_tdesc %a : memref<512xf32> -> !xegpu.tensor_desc<32xf32, #r>
- %1 = xegpu.load_nd %0[0] {layout = #r}: !xegpu.tensor_desc<32xf32, #r> -> vector<32xf32>
+ %1 = xegpu.load_nd %0[0] <{layout = #r}>: !xegpu.tensor_desc<32xf32, #r> -> vector<32xf32>
%11 = vector.shape_cast %1 : vector<32xf32> to vector<32x1xf32>
// CHECK-COUNT-8: vector.broadcast {{.*}}: vector<16x1xf32> to vector<16x16xf32>
%2 = vector.broadcast %11 : vector<32x1xf32> to vector<32x64xf32>
%3 = xegpu.create_nd_tdesc %b : memref<16x512xf32> -> !xegpu.tensor_desc<32x64xf32, #l>
- xegpu.store_nd %2, %3[0, 0] {layout = #l} : vector<32x64xf32>, !xegpu.tensor_desc<32x64xf32, #l>
+ xegpu.store_nd %2, %3[0, 0] <{layout = #l}> : vector<32x64xf32>, !xegpu.tensor_desc<32x64xf32, #l>
gpu.return
}
}
@@ -313,11 +313,11 @@ gpu.module @test_kernel {
%block_id_x = gpu.block_id x
%m = arith.muli %block_id_x, %c32 : index
%0 = xegpu.create_nd_tdesc %a : memref<512x8xf32> -> !xegpu.tensor_desc<32x8xf32, #l>
- %1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<32x8xf32, #l> -> vector<32x8xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #l}>: !xegpu.tensor_desc<32x8xf32, #l> -> vector<32x8xf32>
// CHECK-COUNT-2: vector.transpose {{.*}} [1, 0] : vector<16x8xf32> to vector<8x16xf32>
%2 = vector.transpose %1, [1, 0] : vector<32x8xf32> to vector<8x32xf32>
%3 = xegpu.create_nd_tdesc %b : memref<8x512xf32> -> !xegpu.tensor_desc<8x32xf32, #t>
- xegpu.store_nd %2, %3[0, 0] {layout = #t}: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #t>
+ xegpu.store_nd %2, %3[0, 0] <{layout = #t}>: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #t>
gpu.return
}
}
@@ -334,8 +334,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -353,8 +353,8 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184,
192, 200, 208, 216, 224, 232, 240, 248
]> : vector<32xindex>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%cst], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -369,8 +369,8 @@ gpu.module @test_kernel {
//CHECK: arith.addi [[step]], [[cst]] : vector<16xindex>
%step = vector.step : vector<32xindex>
%mask = vector.create_mask %c16 : vector<32xi1>
- %ld = xegpu.load %src[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- xegpu.store %ld, %dst[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ %ld = xegpu.load %src[%step], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ xegpu.store %ld, %dst[%step], %mask <{chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>}> : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -402,12 +402,12 @@ gpu.module @test_kernel {
%c0 = arith.constant 0 : index
%a_tdesc = xegpu.create_nd_tdesc %A : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #b>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #b>
- %a = xegpu.load_nd %a_tdesc[0, 0] {layout = #b}: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
- %b = xegpu.load_nd %b_tdesc[0, 0] {layout = #b}: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
+ %a = xegpu.load_nd %a_tdesc[0, 0] <{layout = #b}>: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
+ %b = xegpu.load_nd %b_tdesc[0, 0] <{layout = #b}>: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
%a1 = xegpu.convert_layout %a <{input_layout = #b, target_layout = #a}> : vector<16x16xf16>
- %c = xegpu.dpas %a1, %b {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>
+ %c = xegpu.dpas %a1, %b layout_a(#a) layout_b(#b) layout_cd(#c): vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>
%c_tdesc = xegpu.create_nd_tdesc %C : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32, #c>
- xegpu.store_nd %c, %c_tdesc[0, 0] {layout = #c}: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #c>
+ xegpu.store_nd %c, %c_tdesc[0, 0] <{layout = #c}>: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #c>
gpu.return
}
}
@@ -422,7 +422,7 @@ gpu.module @test_kernel {
%acc = arith.constant 0.000000e+00 : f16
%c0 = arith.constant 0 : index
%a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #a>
- %a = xegpu.load_nd %a_tdesc[0, 0] {layout = #a}: !xegpu.tensor_desc<16x16xf16, #a> -> vector<16x16xf16>
+ %a = xegpu.load_nd %a_tdesc[0, 0] <{layout = #a}>: !xegpu.tensor_desc<16x16xf16, #a> -> vector<16x16xf16>
%a_reduce = vector.multi_reduction <add>, %a, %acc [0, 1] : vector<16x16xf16> to f16
%13 = xegpu.convert_layout %a_reduce <{target_layout = #xegpu.slice<#a, dims = [0, 1]>}> : f16
memref.store %13, %arg1[%c0] : memref<4xf16>
@@ -549,7 +549,7 @@ gpu.module @test_kernel {
gpu.func @unroll_store_matrix(%value: vector<1x32xf32>, %arg0 : memref<32768xi8, 3>) {
%mdesc = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
// CHECK-COUNT-2: xegpu.store_matrix {{.*}} : vector<1x16xf32>, !xegpu.mem_desc<64x128xf32>, index, index
- xegpu.store_matrix %value, %mdesc[0, 0] {layout = #xegpu.layout<inst_data = [1, 16]>} : vector<1x32xf32>, !xegpu.mem_desc<64x128xf32>
+ xegpu.store_matrix %value, %mdesc[0, 0] <{layout = #xegpu.layout<inst_data = [1, 16]>}> : vector<1x32xf32>, !xegpu.mem_desc<64x128xf32>
gpu.return
}
}
@@ -568,7 +568,7 @@ gpu.module @test_kernel {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -590,7 +590,7 @@ gpu.module @test_kernel {
%mask = vector.create_mask %c17 : vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask <{chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>}> : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
@@ -615,7 +615,7 @@ gpu.module @test_kernel {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17 : vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
gpu.return %ld : vector<32x4xf32>
}
}
@@ -641,7 +641,7 @@ gpu.module @test_kernel {
%mask = vector.create_mask %c17 : vector<32xi1>
%st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask <{chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>}> : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
}
@@ -667,7 +667,7 @@ gpu.module @test_kernel {
]]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 1, layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
gpu.return %ld : vector<1x1x32xf32>
}
@@ -698,11 +698,11 @@ gpu.module @test_kernel {
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x32xf32, #l>
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x32xf32, #l>
- %a = xegpu.load_nd %a_tdesc[%c0, %c0] {layout = #l}: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
- %b = xegpu.load_nd %b_tdesc[%c0, %c0] {layout = #l}: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
+ %a = xegpu.load_nd %a_tdesc[%c0, %c0] <{layout = #l}>: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
+ %b = xegpu.load_nd %b_tdesc[%c0, %c0] <{layout = #l}>: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
%result = arith.addf %a, %b : vector<1x32xf32>
- xegpu.store_nd %result, %c_tdesc[%c0, %c0] {layout = #l}: vector<1x32xf32>, !xegpu.tensor_desc<1x32xf32, #l>
+ xegpu.store_nd %result, %c_tdesc[%c0, %c0] <{layout = #l}>: vector<1x32xf32>, !xegpu.tensor_desc<1x32xf32, #l>
gpu.return
}
}
@@ -735,10 +735,10 @@ gpu.module @test_kernel {
128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]]
]> : vector<1x1x32xindex>
%mask = arith.constant dense<true> : vector<1x1x32xi1>
- %a = xegpu.load %A[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
- %b = xegpu.load %B[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %a = xegpu.load %A[%cst], %mask <{chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
+ %b = xegpu.load %B[%cst], %mask <{chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
%addf = arith.addf %a, %b : vector<1x1x32xf32>
- xegpu.store %addf, %C[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
+ xegpu.store %addf, %C[%cst], %mask <{chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>}> : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
gpu.return
}
}
@@ -762,7 +762,7 @@ gpu.module @test_kernel {
%n = arith.muli %block_id_y, %c32 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #l3>
- %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l3}: !xegpu.tensor_desc<16x32xf32, #l3> -> vector<16x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] <{layout = #l3}>: !xegpu.tensor_desc<16x32xf32, #l3> -> vector<16x32xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x64xf4E2M1FN, #l1>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<64x32xf4E2M1FN, #l2>
@@ -773,19 +773,19 @@ gpu.module @test_kernel {
iter_args(%arg2 = %c_init)
-> (vector<16x32xf32>) {
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x32xf4E2M1FN>
- %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l1}: !xegpu.tensor_desc<16x64xf4E2M1FN, #l1> -> vector<16x64xf4E2M1FN>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] <{layout = #l1}>: !xegpu.tensor_desc<16x64xf4E2M1FN, #l1> -> vector<16x64xf4E2M1FN>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<32x16xf4E2M1FN>
- %b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #l2}: !xegpu.tensor_desc<64x32xf4E2M1FN, #l2> -> vector<64x32xf4E2M1FN>
+ %b = xegpu.load_nd %b_tdesc[%k, %c0] <{layout = #l2}>: !xegpu.tensor_desc<64x32xf4E2M1FN, #l2> -> vector<64x32xf4E2M1FN>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x1xf8E8M0FNU>
- %sa = xegpu.load_nd %scale_a_tdesc[%c0, %c0] {layout = #l1_scale}: !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale> -> vector<16x2xf8E8M0FNU>
+ %sa = xegpu.load_nd %scale_a_tdesc[%c0, %c0] <{layout = #l1_scale}>: !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale> -> vector<16x2xf8E8M0FNU>
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x16xf8E8M0FNU>
- %sb = xegpu.load_nd %scale_b_tdesc[%c0, %c0] {layout = #l2_scale}: !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale> -> vector<2x32xf8E8M0FNU>
+ %sb = xegpu.load_nd %scale_b_tdesc[%c0, %c0] <{layout = #l2_scale}>: !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale> -> vector<2x32xf8E8M0FNU>
//CHECK-COUNT-8: xegpu.dpas_mx {{.*}}
- %c = xegpu.dpas_mx %a, %b, %arg2 scale_a = %sa scale_b = %sb {layout_a=#l1, layout_b = #l2, layout_cd = #l3, layout_a_scale = #l1_scale, layout_b_scale = #l2_scale, layout_result_0 = #l3}: (vector<16x64xf4E2M1FN>, vector<64x32xf4E2M1FN>, vector<16x32xf32>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xf32>
+ %c = xegpu.dpas_mx %a, %b, %arg2 scale_a = %sa scale_b = %sb <{layout_a=#l1, layout_b = #l2, layout_cd = #l3, layout_a_scale = #l1_scale, layout_b_scale = #l2_scale}> {layout_result_0 = #l3}: (vector<16x64xf4E2M1FN>, vector<64x32xf4E2M1FN>, vector<16x32xf32>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xf32>
scf.yield %c : vector<16x32xf32>
} {layout_result_0 = #l3}
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l3}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l3>
+ xegpu.store_nd %out, %c_tdesc[0, 0] <{layout = #l3}>: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l3>
gpu.return
}
}
@@ -801,18 +801,18 @@ gpu.module @test_kernel {
%c0 = arith.constant 0 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<4x8x16xf32> -> !xegpu.tensor_desc<4x8x16xf32, #c_3d>
// CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x8x16xf32>
- %c_init = xegpu.load_nd %c_tdesc[0, 0, 0] {layout = #c_3d}: !xegpu.tensor_desc<4x8x16xf32, #c_3d> -> vector<4x8x16xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0, 0] <{layout = #c_3d}>: !xegpu.tensor_desc<4x8x16xf32, #c_3d> -> vector<4x8x16xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<4x8x32xf16> -> !xegpu.tensor_desc<4x8x32xf16, #a_3d>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<4x32x16xf16> -> !xegpu.tensor_desc<4x32x16xf16, #b_3d>
// CHECK-COUNT-8: xegpu.load_nd {{.*}} -> vector<1x8x16xf16>
- %a = xegpu.load_nd %a_tdesc[0, 0, 0] {layout = #a_3d}: !xegpu.tensor_desc<4x8x32xf16, #a_3d> -> vector<4x8x32xf16>
+ %a = xegpu.load_nd %a_tdesc[0, 0, 0] <{layout = #a_3d}>: !xegpu.tensor_desc<4x8x32xf16, #a_3d> -> vector<4x8x32xf16>
// CHECK-COUNT-8: xegpu.load_nd {{.*}} -> vector<1x16x16xf16>
- %b = xegpu.load_nd %b_tdesc[0, 0, 0] {layout = #b_3d}: !xegpu.tensor_desc<4x32x16xf16, #b_3d> -> vector<4x32x16xf16>
+ %b = xegpu.load_nd %b_tdesc[0, 0, 0] <{layout = #b_3d}>: !xegpu.tensor_desc<4x32x16xf16, #b_3d> -> vector<4x32x16xf16>
// CHECK-COUNT-8: xegpu.dpas {{.*}} : vector<1x8x16xf16>, vector<1x16x16xf16>, vector<1x8x16xf32> -> vector<1x8x16xf32>
- %d = xegpu.dpas %a, %b, %c_init {layout_a = #a_3d, layout_b = #b_3d, layout_cd = #c_3d}
+ %d = xegpu.dpas %a, %b, %c_init layout_a(#a_3d) layout_b(#b_3d) layout_cd(#c_3d)
: vector<4x8x32xf16>, vector<4x32x16xf16>, vector<4x8x16xf32> -> vector<4x8x16xf32>
// CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<1x8x16xf32>, !xegpu.tensor_desc<1x8x16xf32>
- xegpu.store_nd %d, %c_tdesc[0, 0, 0] {layout = #c_3d}: vector<4x8x16xf32>, !xegpu.tensor_desc<4x8x16xf32, #c_3d>
+ xegpu.store_nd %d, %c_tdesc[0, 0, 0] <{layout = #c_3d}>: vector<4x8x16xf32>, !xegpu.tensor_desc<4x8x16xf32, #c_3d>
gpu.return
}
}
@@ -830,31 +830,31 @@ gpu.module @test_kernel {
%c0 = arith.constant 0 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<2x16x32xf32> -> !xegpu.tensor_desc<2x16x32xf32, #c_3d_mx>
// CHECK-COUNT-8: xegpu.load_nd {{.*}} -> vector<1x8x16xf32>
- %c_init = xegpu.load_nd %c_tdesc[0, 0, 0] {layout = #c_3d_mx}: !xegpu.tensor_desc<2x16x32xf32, #c_3d_mx> -> vector<2x16x32xf32>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0, 0] <{layout = #c_3d_mx}>: !xegpu.tensor_desc<2x16x32xf32, #c_3d_mx> -> vector<2x16x32xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<2x16x64xf4E2M1FN> -> !xegpu.tensor_desc<2x16x64xf4E2M1FN, #a_3d_mx>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<2x64x32xf4E2M1FN> -> !xegpu.tensor_desc<2x64x32xf4E2M1FN, #b_3d_mx>
%sa_tdesc = xegpu.create_nd_tdesc %SA : memref<2x16x2xf8E8M0FNU> -> !xegpu.tensor_desc<2x16x2xf8E8M0FNU, #sa_3d>
%sb_tdesc = xegpu.create_nd_tdesc %SB : memref<2x2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x2x32xf8E8M0FNU, #sb_3d>
// CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x8x64xf4E2M1FN>
- %a = xegpu.load_nd %a_tdesc[0, 0, 0] {layout = #a_3d_mx}: !xegpu.tensor_desc<2x16x64xf4E2M1FN, #a_3d_mx> -> vector<2x16x64xf4E2M1FN>
+ %a = xegpu.load_nd %a_tdesc[0, 0, 0] <{layout = #a_3d_mx}>: !xegpu.tensor_desc<2x16x64xf4E2M1FN, #a_3d_mx> -> vector<2x16x64xf4E2M1FN>
// CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x64x16xf4E2M1FN>
- %b = xegpu.load_nd %b_tdesc[0, 0, 0] {layout = #b_3d_mx}: !xegpu.tensor_desc<2x64x32xf4E2M1FN, #b_3d_mx> -> vector<2x64x32xf4E2M1FN>
+ %b = xegpu.load_nd %b_tdesc[0, 0, 0] <{layout = #b_3d_mx}>: !xegpu.tensor_desc<2x64x32xf4E2M1FN, #b_3d_mx> -> vector<2x64x32xf4E2M1FN>
// CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x8x2xf8E8M0FNU>
- %sa = xegpu.load_nd %sa_tdesc[0, 0, 0] {layout = #sa_3d}: !xegpu.tensor_desc<2x16x2xf8E8M0FNU, #sa_3d> -> vector<2x16x2xf8E8M0FNU>
+ %sa = xegpu.load_nd %sa_tdesc[0, 0, 0] <{layout = #sa_3d}>: !xegpu.tensor_desc<2x16x2xf8E8M0FNU, #sa_3d> -> vector<2x16x2xf8E8M0FNU>
// CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x2x16xf8E8M0FNU>
- %sb = xegpu.load_nd %sb_tdesc[0, 0, 0] {layout = #sb_3d}: !xegpu.tensor_desc<2x2x32xf8E8M0FNU, #sb_3d> -> vector<2x2x32xf8E8M0FNU>
+ %sb = xegpu.load_nd %sb_tdesc[0, 0, 0] <{layout = #sb_3d}>: !xegpu.tensor_desc<2x2x32xf8E8M0FNU, #sb_3d> -> vector<2x2x32xf8E8M0FNU>
// dpas_mx: [2,16,64] x [2,64,32] -> [2,16,32] with scales [2,16,2] and [2,2,32]
// unrolled: batch=2, M=16/8=2, K=64/64=1, N=32/16=2 -> 2*2*2=8 results (with k-reduction)
// CHECK-COUNT-8: xegpu.dpas_mx {{.*}} : (vector<1x8x64xf4E2M1FN>, vector<1x64x16xf4E2M1FN>, vector<1x8x16xf32>, vector<1x8x2xf8E8M0FNU>, vector<1x2x16xf8E8M0FNU>) -> vector<1x8x16xf32>
%d = xegpu.dpas_mx %a, %b, %c_init scale_a = %sa scale_b = %sb
- {layout_a = #a_3d_mx, layout_b = #b_3d_mx, layout_cd = #c_3d_mx,
- layout_a_scale = #sa_3d, layout_b_scale = #sb_3d}
+ <{layout_a = #a_3d_mx, layout_b = #b_3d_mx, layout_cd = #c_3d_mx,
+ layout_a_scale = #sa_3d, layout_b_scale = #sb_3d}>
: (vector<2x16x64xf4E2M1FN>, vector<2x64x32xf4E2M1FN>,
vector<2x16x32xf32>,
vector<2x16x2xf8E8M0FNU>, vector<2x2x32xf8E8M0FNU>)
-> vector<2x16x32xf32>
// CHECK-COUNT-8: xegpu.store_nd {{.*}} : vector<1x8x16xf32>, !xegpu.tensor_desc<1x8x16xf32>
- xegpu.store_nd %d, %c_tdesc[0, 0, 0] {layout = #c_3d_mx}: vector<2x16x32xf32>, !xegpu.tensor_desc<2x16x32xf32, #c_3d_mx>
+ xegpu.store_nd %d, %c_tdesc[0, 0, 0] <{layout = #c_3d_mx}>: vector<2x16x32xf32>, !xegpu.tensor_desc<2x16x32xf32, #c_3d_mx>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
index e2a4897fac519..4a0c1b56f3227 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-recover-layout.mlir
@@ -27,14 +27,11 @@ gpu.func @for_basic(%arg0: memref<8x128xf16>, %arg1: memref<128x16xf16>, %arg2:
// CHECK: scf.for
%2 = scf.for %arg3 = %c0 to %c128 step %c16
iter_args(%arg6 = %cst) -> (vector<8x16xf32>) {
- %4 = xegpu.load_nd %0[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %4 = xegpu.load_nd %0[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
- %5 = xegpu.load_nd %1[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ %5 = xegpu.load_nd %1[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
- %6 = xegpu.dpas %4, %5, %arg6
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %6 = xegpu.dpas %4, %5, %arg6 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xf16>, vector<16x16xf16>, vector<8x16xf32> -> vector<8x16xf32>
// Recovery propagates layout to scf.yield vector operand.
// CHECK: scf.yield {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
@@ -47,7 +44,7 @@ gpu.func @for_basic(%arg0: memref<8x128xf16>, %arg1: memref<128x16xf16>, %arg2:
// CHECK-SAME: -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%3 = xegpu.create_nd_tdesc %arg2 : memref<8x16xf32>
-> !xegpu.tensor_desc<8x16xf32>
- xegpu.store_nd %2, %3[%c0, %c0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.store_nd %2, %3[%c0, %c0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
gpu.return
}
@@ -69,7 +66,7 @@ gpu.func @while_basic(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
// CHECK-SAME: -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>
%0 = xegpu.create_nd_tdesc %arg0 : memref<1024xf32>
-> !xegpu.tensor_desc<256xf32>
- %1 = xegpu.load_nd %0[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ %1 = xegpu.load_nd %0[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
: !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// CHECK: xegpu.create_nd_tdesc
// CHECK-SAME: -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [16], sg_data = [16]>>
@@ -86,10 +83,10 @@ gpu.func @while_basic(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
scf.condition(%4) %arg2, %arg3 : vector<256xf32>, i32
} do {
^bb0(%arg2: vector<256xf32>, %arg3: i32):
- xegpu.store_nd %arg2, %2[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ xegpu.store_nd %arg2, %2[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
: vector<256xf32>, !xegpu.tensor_desc<256xf32>
%4 = arith.addi %arg3, %c1_i32 : i32
- %6 = xegpu.load_nd %0[256] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ %6 = xegpu.load_nd %0[256] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
: !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// Recovery propagates layout to scf.yield in the "do" region via
// sibling region propagation (from "before" region arg back to "do" yield).
@@ -99,7 +96,7 @@ gpu.func @while_basic(%arg0: memref<1024xf32>, %arg1: memref<1024xf32>) {
// CHECK: } attributes {layout_operand_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>,
// CHECK-SAME: layout_result_0 = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
}
- xegpu.store_nd %3#0, %2[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}
+ xegpu.store_nd %3#0, %2[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}>
: vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -116,18 +113,18 @@ gpu.func @if_basic(
%arg1: !xegpu.tensor_desc<16x16xf16>,
%arg2: i1,
%arg3: !xegpu.tensor_desc<8x16xf32>) {
- %0 = xegpu.load_nd %arg0[0, 0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %0 = xegpu.load_nd %arg0[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
// CHECK: scf.if
%1 = scf.if %arg2 -> (vector<16x16xf16>) {
- %3 = xegpu.load_nd %arg1[0, 0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ %3 = xegpu.load_nd %arg1[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16>
-> vector<16x16xf16>
// Recovery propagates layout to scf.yield operand in "then" region.
// CHECK: scf.yield {layout_operand_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
scf.yield %3 : vector<16x16xf16>
} else {
- %3 = xegpu.load_nd %arg1[0, 0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
+ %3 = xegpu.load_nd %arg1[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<16x16xf16>
-> vector<16x16xf16>
// Recovery propagates layout to scf.yield operand in "else" region.
@@ -136,12 +133,9 @@ gpu.func @if_basic(
// Recovery sets layout_result_0 on the scf.if for the vector result.
// CHECK: } {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}
}
- %2 = xegpu.dpas %0, %1
- {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ %2 = xegpu.dpas %0, %1 layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<8x16xf16>, vector<16x16xf16> -> vector<8x16xf32>
- xegpu.store_nd %2, %arg3[0, 0] {layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.store_nd %2, %arg3[0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
index 49fdf1cbee174..1aad73433e656 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-unroll-patterns.mlir
@@ -139,7 +139,7 @@ gpu.module @test {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
gpu.return %ld : vector<32xf32>
}
@@ -160,7 +160,7 @@ gpu.module @test {
%mask = vector.create_mask %c17: vector<32xi1>
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask <{chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>}> : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
@@ -184,7 +184,7 @@ gpu.module @test {
%c17 = arith.constant 17: index
%mask = vector.create_mask %c17: vector<32xi1>
- %ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
+ %ld = xegpu.load %src[%cst], %mask <{chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
gpu.return %ld : vector<32x4xf32>
}
@@ -209,7 +209,7 @@ gpu.module @test {
%mask = vector.create_mask %c17: vector<32xi1>
%st_vec = arith.constant dense<1023.>: vector<32x4xf32>
- xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
+ xegpu.store %st_vec, %src[%cst], %mask <{chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>}> : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
gpu.return
}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
index 918d587f12db5..662da6014c6e9 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-elemwise.mlir
@@ -5,7 +5,7 @@ gpu.module @test_elementwise_ops {
// CHECK-LABEL: unary_ops_sg_layout_only
gpu.func @unary_ops_sg_layout_only(%a: memref<24x32xf32>) {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
// CHECK: math.exp {{.*}} : vector<12x8xf32>
%exp = math.exp %load_a
@@ -23,7 +23,7 @@ gpu.module @test_elementwise_ops {
// CHECK-LABEL: unary_ops
gpu.func @unary_ops(%a: memref<24x32xf32>) {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
// CHECK: math.exp {{.*}} : vector<12x8xf32>
%exp = math.exp %load_a
@@ -42,9 +42,9 @@ gpu.module @test_elementwise_ops {
gpu.func @binary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
// CHECK: arith.addf {{.*}}, {{.*}} : vector<12x8xf32>
%addf = arith.addf %load_a, %load_b : vector<24x32xf32>
@@ -62,11 +62,11 @@ gpu.module @test_elementwise_ops {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi1> -> !xegpu.tensor_desc<24x32xi1>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_c = xegpu.load_nd %tdesc_c[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xi1> -> vector<24x32xi1>
// CHECK: arith.select {{.*}}, {{.*}}, {{.*}} : vector<12x8xi1>, vector<12x8xf32>
%select = arith.select %load_c, %load_a, %load_b
@@ -85,9 +85,9 @@ gpu.module @test_elementwise_ops {
gpu.func @type_conversion_ops(%a: memref<24x32xf32>, %b: memref<24x32xi32>) {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
// CHECK: arith.truncf {{.*}} : vector<12x8xf32> to vector<12x8xf16>
%truncf = arith.truncf %load_a
@@ -108,13 +108,13 @@ gpu.module @test_elementwise_ops {
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
%tdesc_d = xegpu.create_nd_tdesc %d : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_c = xegpu.load_nd %tdesc_c[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
- %load_d = xegpu.load_nd %tdesc_d[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
+ %load_d = xegpu.load_nd %tdesc_d[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
// CHECK: arith.cmpf ult, {{.*}}, {{.*}} : vector<12x8xf32>
%cmpf = arith.cmpf ult, %load_a, %load_b
@@ -135,9 +135,9 @@ gpu.module @test_elementwise_ops {
gpu.func @elementwise_ops_rr_assignment(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
// CHECK-COUNT-12: arith.negf {{.*}} : vector<2x2xf32>
// CHECK-NOT: arith.negf
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
index 2969e551d93cb..288f072000c0c 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-rr.mlir
@@ -18,7 +18,7 @@ gpu.module @test_distribution {
// CHECK-NOT: xegpu.load_nd
%tdesc = xegpu.create_nd_tdesc %src: memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
gpu.return
@@ -30,10 +30,10 @@ gpu.module @test_distribution {
// CHECK-NOT: xegpu.store_nd
%tdesc = xegpu.create_nd_tdesc %src: memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
- xegpu.store_nd %load, %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.store_nd %load, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32>
gpu.return
}
@@ -44,7 +44,7 @@ gpu.module @test_distribution {
// CHECK-NOT: xegpu.prefetch_nd
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- xegpu.prefetch_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.prefetch_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
gpu.return
}
@@ -56,22 +56,19 @@ gpu.module @test_distribution {
// CHECK-COUNT-4: xegpu.load_nd {{%.*}}[{{%.*}}, {{%.*}}] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<16x16xf16>
// CHECK-COUNT-4: xegpu.create_nd_tdesc %[[ARG_1]] : memref<128x256xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>>
// CHECK-COUNT-4: xegpu.load_nd {{%.*}}[{{%.*}}, {{%.*}}] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<16x16xf16>
- // CHECK-COUNT-16: xegpu.dpas %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>
+ // CHECK-COUNT-16: xegpu.dpas %{{.*}}, %{{.*}} layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) : vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>
// CHECK-NOT: xegpu.dpas
%tdesc_a = xegpu.create_nd_tdesc %a : memref<256x128xf16>
-> !xegpu.tensor_desc<256x128xf16>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf16>
-> vector<256x128xf16>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<128x256xf16>
-> !xegpu.tensor_desc<128x256xf16>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<128x256xf16>
-> vector<128x256xf16>
- %dpas = xegpu.dpas %load_a, %load_b
- {layout_a = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %dpas = xegpu.dpas %load_a, %load_b layout_a(#xegpu.layout<sg_layout = [8, 4], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<sg_layout = [4, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<256x128xf16>, vector<128x256xf16> -> vector<256x256xf32>
gpu.return
}
@@ -82,7 +79,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<1.0> : vector<256xf32>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf32>
-> !xegpu.tensor_desc<256x64xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [16, 64]>}>
: !xegpu.tensor_desc<256x64xf32>
-> vector<256x64xf32>
// CHECK-COUNT-2: vector.multi_reduction <add>, {{.*}}, %[[C0:.*]] [1] : vector<16x64xf32> to vector<16xf32>
@@ -131,7 +128,7 @@ gpu.module @test_distribution {
gpu.func @vector_transpose(%src: memref<256x128xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 16], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 16], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
// CHECK-COUNT-2: vector.transpose {{.*}}, [1, 0] : vector<32x16xf32> to vector<16x32xf32>
@@ -176,7 +173,7 @@ gpu.module @test_distribution {
%cst_0 = arith.constant dense<0xFF800000> : vector<256xf32>
%block_id_x = gpu.block_id x
%0 = xegpu.create_nd_tdesc %arg0 : memref<4096x128xf32> -> !xegpu.tensor_desc<256x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
- %1 = xegpu.load_nd %0[%block_id_x, 0] {layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 128], inst_data = [8, 16]>} : !xegpu.tensor_desc<256x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<256x128xf32>
+ %1 = xegpu.load_nd %0[%block_id_x, 0] <{layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 128], inst_data = [8, 16]>}> : !xegpu.tensor_desc<256x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<256x128xf32>
%2 = vector.multi_reduction <maximumf>, %1, %cst_0 [1] : vector<256x128xf32> to vector<256xf32>
%3 = vector.shape_cast %2 : vector<256xf32> to vector<256x1xf32>
%4 = vector.broadcast %3 : vector<256x1xf32>to vector<256x128xf32>
@@ -289,7 +286,7 @@ gpu.module @test_distribution {
gpu.func @broadcast(%src: memref<128x1xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<128x1xf32>
-> !xegpu.tensor_desc<128x1xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [4, 1], sg_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 1], sg_data = [16, 1], lane_layout = [8, 1], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<128x1xf32>
-> vector<128x1xf32>
// CHECK-COUNT-4: vector.broadcast {{.*}} : vector<16x1xf32> to vector<16x32xf32>
@@ -313,8 +310,8 @@ gpu.module @test_distribution {
%1 = xegpu.create_nd_tdesc %arg1 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
// CHECK-LABEL: scf.for
scf.for %arg2 = %c0 to %c1024 step %c256 {
- %3 = xegpu.load_nd %0[%arg2] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
- xegpu.store_nd %3, %1[%arg2] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ %3 = xegpu.load_nd %0[%arg2] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ xegpu.store_nd %3, %1[%arg2] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
}
gpu.return
}
@@ -325,7 +322,7 @@ gpu.module @test_distribution {
%c0_i32 = arith.constant 0 : i32
%c256 = arith.constant 256 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
- %1 = xegpu.load_nd %0[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %1 = xegpu.load_nd %0[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
%2 = xegpu.create_nd_tdesc %arg1 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
// CHECK: scf.while ({{.*}}) : (vector<16xf32>, vector<16xf32>, i32) -> (vector<16xf32>, vector<16xf32>, i32)
%3:2 = scf.while (%arg2 = %1, %arg3 = %c0_i32) : (vector<256xf32>, i32) -> (vector<256xf32>, i32) {
@@ -335,12 +332,12 @@ gpu.module @test_distribution {
} do {
// CHECK: ([[arg2:%.+]]: vector<16xf32>, [[arg3:%.+]]: vector<16xf32>, [[arg4:%.+]]: i32)
^bb0(%arg2: vector<256xf32>, %arg3: i32):
- xegpu.store_nd %arg2, %2[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %arg2, %2[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
%4 = arith.addi %arg3, %c1_i32 : i32
- %6 = xegpu.load_nd %0[%c256] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %6 = xegpu.load_nd %0[%c256] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
scf.yield %6, %4 : vector<256xf32>, i32
}
- xegpu.store_nd %3#0, %2[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %3#0, %2[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -353,17 +350,17 @@ gpu.module @test_distribution {
// CHECK-LABEL: scf.if
// CHECK-SAME: (vector<16xf32>, vector<16xf32>)
%4 = scf.if %3 -> (vector<256xf32>) {
- %5 = xegpu.load_nd %1[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %5 = xegpu.load_nd %1[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// CHECK-LABEL: scf.yield
// CHECK-SAME: vector<16xf32>, vector<16xf32>
scf.yield %5 : vector<256xf32>
} else {
- %5 = xegpu.load_nd %2[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %5 = xegpu.load_nd %2[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// CHECK-LABEL: scf.yield
// CHECK-SAME: vector<16xf32>, vector<16xf32>
scf.yield %5 : vector<256xf32>
}
- xegpu.store_nd %4, %1[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %4, %1[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -372,7 +369,7 @@ gpu.module @test_distribution {
%id = gpu.subgroup_id : index
%t = xegpu.create_nd_tdesc %arg0 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>>
- %d = xegpu.load_nd %t[0] {layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}: !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>
+ %d = xegpu.load_nd %t[0] <{layout = #xegpu.layout<sg_layout = [8], sg_data = [16]>}>: !xegpu.tensor_desc<256xf32, #xegpu.layout<sg_layout = [8], sg_data = [16]>> -> vector<256xf32>
%0 = arith.cmpi eq, %id, %c10 : index
// CHECK-LABEL: scf.if
@@ -396,7 +393,7 @@ gpu.module @test_distribution {
%0 = xegpu.create_nd_tdesc %arg0 : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32>
// CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<16x16xf32, #xegpu.layout<inst_data = [16, 16]>> -> vector<16x16xf32>
// CHECK-COUNT-2: xegpu.convert_layout {{.*}} <{input_layout = #xegpu.layout<inst_data = [16, 16]>, target_layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x16xf32>
- %1 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>} : !xegpu.tensor_desc<32x64xf32> -> vector<32x64xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>}> : !xegpu.tensor_desc<32x64xf32> -> vector<32x64xf32>
%2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>,
target_layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16]>}> : vector<32x64xf32>
gpu.return
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 098743cb1e713..701ff9baded87 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -40,7 +40,7 @@ gpu.module @test_distribution {
//CHECK-DAG: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][{{%.*}}, {{%.*}}] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<32x32xf32>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
gpu.return
@@ -52,10 +52,10 @@ gpu.module @test_distribution {
//CHECK: xegpu.store_nd %{{.*}}, {{%.*}}[{{%.*}}, {{%.*}}] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<32x32xf32>, !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
%tdesc = xegpu.create_nd_tdesc %src: memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
- xegpu.store_nd %load, %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}
+ xegpu.store_nd %load, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>}>
: vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32>
gpu.return
}
@@ -75,79 +75,73 @@ gpu.module @test_distribution {
// CHECK-LABEL: dpas
gpu.func @dpas(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {
- // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>
+ // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>) : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>
%tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf16>
-> !xegpu.tensor_desc<128x128xf16>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<128x128xf16>
-> vector<128x128xf16>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf16>
-> !xegpu.tensor_desc<128x128xf16>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<128x128xf16>
-> vector<128x128xf16>
- %dpas = xegpu.dpas %load_a, %load_b
- {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>,
- layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>,
- layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %dpas = xegpu.dpas %load_a, %load_b layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1]>) layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>)
: vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>
gpu.return
}
// CHECK-LABEL: dpas_mx
gpu.func @dpas_mx(%a: memref<128x128xf8E5M2>, %b: memref<128x128xf8E5M2>, %a_scale: memref<128x4xf8E8M0FNU>, %b_scale: memref<4x128xf8E8M0FNU>) {
- // CHECK: %[[DPAS_MX:.*]] = xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{.*}} scale_b = %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : (vector<16x128xf8E5M2>, vector<128x16xf8E5M2>, vector<16x16xbf16>, vector<16x4xf8E8M0FNU>, vector<4x16xf8E8M0FNU>) -> vector<16x16xbf16>
+ // CHECK: %[[DPAS_MX:.*]] = xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{.*}} scale_b = %{{.*}} <{layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : (vector<16x128xf8E5M2>, vector<128x16xf8E5M2>, vector<16x16xbf16>, vector<16x4xf8E8M0FNU>, vector<4x16xf8E8M0FNU>) -> vector<16x16xbf16>
%tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf8E5M2>
-> !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>}>
: !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>>
-> vector<128x128xf8E5M2>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf8E5M2>
-> !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>}
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>}>
: !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
-> vector<128x128xf8E5M2>
%tdesc_a_scale = xegpu.create_nd_tdesc %a_scale : memref<128x4xf8E8M0FNU>
-> !xegpu.tensor_desc<128x4xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>>
- %load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>}
+ %load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<128x4xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>>
-> vector<128x4xf8E8M0FNU>
%tdesc_b_scale = xegpu.create_nd_tdesc %b_scale : memref<4x128xf8E8M0FNU>
-> !xegpu.tensor_desc<4x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
- %load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<4x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
-> vector<4x128xf8E8M0FNU>
%cst = arith.constant dense<0.0> : vector<128x128xbf16>
%dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale
- {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>,
+ <{layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>,
layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>,
layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>,
layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>,
- layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
: (vector<128x128xf8E5M2>, vector<128x128xf8E5M2>, vector<128x128xbf16>, vector<128x4xf8E8M0FNU>, vector<4x128xf8E8M0FNU>) -> vector<128x128xbf16>
gpu.return
}
// CHECK-LABEL: dpas_no_sg_data
gpu.func @dpas_no_sg_data(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {
- // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1], order = [1, 0]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>} : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>
+ // CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} layout_a(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>) layout_b(#xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1], order = [1, 0]>) layout_cd(#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>) : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>
%tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf16>
-> !xegpu.tensor_desc<128x128xf16>
- %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1],
- order = [1, 0]>}
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 1],
+ order = [1, 0]>}>
: !xegpu.tensor_desc<128x128xf16>
-> vector<128x128xf16>
%tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf16>
-> !xegpu.tensor_desc<128x128xf16>
- %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1], order = [1, 0]> }
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1], order = [1, 0]> }>
: !xegpu.tensor_desc<128x128xf16>
-> vector<128x128xf16>
- %dpas = xegpu.dpas %load_a, %load_b
- {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1],
- order = [1, 0]>,
- layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1],
- order = [1, 0]>,
- layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>}
+ %dpas = xegpu.dpas %load_a, %load_b layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1],
+ order = [1, 0]>) layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [2, 1],
+ order = [1, 0]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>)
: vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>
gpu.return
}
@@ -157,7 +151,7 @@ gpu.module @test_distribution {
gpu.func @broadcast_dim1(%src: memref<256x1xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x1xf32>
-> !xegpu.tensor_desc<256x1xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 1], sg_data = [32, 1], lane_layout = [8, 1], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x1xf32>
-> vector<256x1xf32>
// CHECK: vector.broadcast {{.*}} : vector<32x1xf32> to vector<32x32xf32>
@@ -174,7 +168,7 @@ gpu.module @test_distribution {
gpu.func @broadcast_dim0(%src: memref<1x128xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<1x128xf32>
-> !xegpu.tensor_desc<1x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [1, 4], sg_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<1x128xf32>
-> vector<1x128xf32>
// CHECK: vector.broadcast {{.*}} : vector<1x32xf32> to vector<32x32xf32>
@@ -206,9 +200,9 @@ gpu.module @test_distribution {
%3 = xegpu.create_nd_tdesc %arg0 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16>
%4 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16>
// load_nd with offset
- %5 = xegpu.load_nd %2[%0, %1] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}: !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
- %6 = xegpu.load_nd %3[%0, %c0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
- %7 = xegpu.load_nd %4[%c0, %1] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %5 = xegpu.load_nd %2[%0, %1] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}>: !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %6 = xegpu.load_nd %3[%0, %c0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}>: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %7 = xegpu.load_nd %4[%c0, %1] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}>: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
// scf.for loop
// CHECK: [[scf:%.+]]:3 = scf.for [[arg3:%.+]] = [[c0]] to [[c1024]] step [[c128]]
// CHECK-SAME: iter_args([[arg4:%.+]] = {{.*}}, [[arg5:%.+]] = {{.*}}, [[arg6:%.+]] = {{.*}}) ->
@@ -220,17 +214,14 @@ gpu.module @test_distribution {
%8:3 = scf.for %arg3 = %c0 to %c1024 step %c128 iter_args(%arg4 = %6, %arg5 = %7, %arg6 = %5)
-> (vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32>) {
// load_nd with offset inside loop
- %9 = xegpu.dpas %arg4, %arg5, %arg6
- {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>,
- layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>,
- layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}
+ %9 = xegpu.dpas %arg4, %arg5, %arg6 layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>) layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>)
: vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
- %10 = xegpu.load_nd %3[%arg3, %c0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
- %11 = xegpu.load_nd %4[%c0, %arg3] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %10 = xegpu.load_nd %3[%arg3, %c0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}>: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %11 = xegpu.load_nd %4[%c0, %arg3] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}>: !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
scf.yield %10, %11, %9 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32>
}
// store_nd with offset
- xegpu.store_nd %8#2, %2[%0, %1] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ xegpu.store_nd %8#2, %2[%0, %1] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
gpu.return
}
@@ -249,7 +240,7 @@ gpu.module @test_distribution {
// CHECK-NOT: index.sub
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
} {sg_id_range = #xegpu.range<[0, 32]>}
@@ -262,7 +253,7 @@ gpu.module @test_distribution {
// CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C2]]
%tdesc = xegpu.create_nd_tdesc %src2 : memref<128x64xf32>
-> !xegpu.tensor_desc<128x64xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<128x64xf32>
-> vector<128x64xf32>
%exp = math.exp %load : vector<128x64xf32>
@@ -282,7 +273,7 @@ gpu.module @test_distribution {
%c32 = arith.constant 32 : index
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
%cond1 = arith.cmpi sge, %sg_id, %c3 : index
@@ -295,7 +286,7 @@ gpu.module @test_distribution {
// CHECK: %[[SUB:.*]] = index.sub %{{.*}}, %[[C3]]
%td = xegpu.create_nd_tdesc %src1 : memref<128x64xf32>
-> !xegpu.tensor_desc<128x64xf32>
- %ld = xegpu.load_nd %td[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>}
+ %ld = xegpu.load_nd %td[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [32, 16], lane_layout = [8, 4], lane_data = [1, 1]>}>
: !xegpu.tensor_desc<128x64xf32>
-> vector<128x64xf32>
%exp = math.exp %ld : vector<128x64xf32>
@@ -317,7 +308,7 @@ gpu.module @test_distribution {
// CHECK-SAME: : memref<?xf16>, vector<32x4xindex>, vector<32x4xi1> -> vector<32x4xf16>
%offset = arith.constant dense<0> : vector<256x16xindex>
%mask = arith.constant dense<1> : vector<256x16xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}
+ %load = xegpu.load %src[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 4]>, l1_hint = #xegpu.cache_hint<cached>}>
: memref<?xf16>, vector<256x16xindex>, vector<256x16xi1> -> vector<256x16xf16>
gpu.return
}
@@ -333,8 +324,8 @@ gpu.module @test_distribution {
%val = arith.constant dense<25.5> : vector<256xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- xegpu.store %val, %dest[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
- l1_hint = #xegpu.cache_hint<cached>}
+ xegpu.store %val, %dest[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [32], sg_data = [8], inst_data = [8]>,
+ l1_hint = #xegpu.cache_hint<cached>}>
: vector<256xf16>, memref<256xf16>, vector<256xindex>, vector<256xi1>
gpu.return
}
@@ -348,7 +339,7 @@ gpu.module @test_distribution {
// CHECK-SAME: : memref<?xf16>, vector<8xindex>, vector<8xi1> -> vector<8x4xf16>
%offset = arith.constant dense<0> : vector<256xindex>
%mask = arith.constant dense<1> : vector<256xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 4, layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 4]>, l1_hint = #xegpu.cache_hint<cached>}
+ %load = xegpu.load %src[%offset], %mask <{chunk_size = 4, layout = #xegpu.layout<sg_layout = [32, 1], sg_data = [8, 4]>, l1_hint = #xegpu.cache_hint<cached>}>
: memref<?xf16>, vector<256xindex>, vector<256xi1> -> vector<256x4xf16>
gpu.return
}
@@ -399,7 +390,7 @@ gpu.module @test_distribution {
//CHECK: xegpu.store_matrix [[cst]], [[mdesc]][[[off_y]], [[off_x]]] : vector<32x32xf32>, !xegpu.mem_desc<64x128xf32>, index, index
%cst = arith.constant dense<1.0> : vector<64x128xf32>
%mdesc = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
- xegpu.store_matrix %cst, %mdesc[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>} : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
+ xegpu.store_matrix %cst, %mdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [32, 32]>}> : vector<64x128xf32>, !xegpu.mem_desc<64x128xf32>
gpu.return
}
@@ -408,7 +399,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<1.0> : vector<128xf32>
%tdesc = xegpu.create_nd_tdesc %src : memref<4x128xf32>
-> !xegpu.tensor_desc<4x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>}>
: !xegpu.tensor_desc<4x128xf32>
-> vector<4x128xf32>
// CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [0] : vector<4x4xf32> to vector<4xf32>
@@ -425,7 +416,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<1.0> : vector<256xf32>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x64xf32>
-> !xegpu.tensor_desc<256x64xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>}>
: !xegpu.tensor_desc<256x64xf32>
-> vector<256x64xf32>
// CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [1] : vector<16x64xf32> to vector<16xf32>
@@ -442,7 +433,7 @@ gpu.module @test_distribution {
%cst_acc = arith.constant dense<0.0> : vector<4x2x6xf16>
%offset = arith.constant dense<0> : vector<4x2x6x32xindex>
%mask = arith.constant dense<true> : vector<4x2x6x32xi1>
- %load = xegpu.load %src[%offset], %mask {layout = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>} : ui64, vector<4x2x6x32xindex>, vector<4x2x6x32xi1> -> vector<4x2x6x32xf16>
+ %load = xegpu.load %src[%offset], %mask <{layout = #xegpu.layout<sg_layout = [4, 2, 6, 1], sg_data = [1, 1, 1, 32]>}> : ui64, vector<4x2x6x32xindex>, vector<4x2x6x32xi1> -> vector<4x2x6x32xf16>
// CHECK: vector.multi_reduction <add>, {{.*}}, {{.*}} [3] : vector<1x1x1x32xf16> to vector<1x1x1xf16>
%reduce = vector.multi_reduction <add>, %load, %cst_acc [3]
: vector<4x2x6x32xf16> to vector<4x2x6xf16>
@@ -469,7 +460,7 @@ gpu.module @test_distribution {
%cst = arith.constant 0.0 : f32
%tdesc = xegpu.create_nd_tdesc %src : memref<32x32xf32>
-> !xegpu.tensor_desc<32x32xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [8, 8]>}>
: !xegpu.tensor_desc<32x32xf32>
-> vector<32x32xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [0, 1]
@@ -564,7 +555,7 @@ gpu.module @test_distribution {
gpu.func @vector_transpose(%src: memref<256x32xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x32xf32>
-> !xegpu.tensor_desc<256x32xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 32], lane_layout = [16, 1], lane_data = [1, 1], order =[0, 1]>}>
: !xegpu.tensor_desc<256x32xf32>
-> vector<256x32xf32>
//CHECK: vector.transpose {{.*}}, [1, 0] : vector<64x32xf32> to vector<32x64xf32>
@@ -762,7 +753,7 @@ gpu.module @test_distribution {
// CHECK: %[[LOAD:.*]] = xegpu.load {{.*}} <{chunk_size = 1 : i64, layout = #xegpu.slice<#xegpu.layout<inst_data = [8, 16]>, dims = [0]>}>
// CHECK-SAME: memref<4096xf32>, vector<32xindex>, vector<32xi1> -> vector<32xf32>
- %3 = xegpu.load %2[%offset], %mask {chunk_size = 1, layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> } : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
+ %3 = xegpu.load %2[%offset], %mask <{chunk_size = 1, layout = #xegpu.slice<#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, dims = [0]> }> : memref<4096xf32>, vector<256xindex>, vector<256xi1> -> vector<256xf32>
// CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[LOAD]] : vector<32xf32> to vector<32x32xf32>
%4 = vector.broadcast %3 : vector<256xf32> to vector<256x256xf32>
@@ -796,7 +787,7 @@ gpu.module @test_distribution {
%cst_3 = arith.constant dense<1.0> : vector<1x32xf32>
%offset = arith.constant dense<0> : vector<1x32x32xindex>
%mask = arith.constant dense<true> : vector<1x32x32xi1>
- %14 = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>} : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
+ %14 = xegpu.load %src[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [1, 32, 1], sg_data = [1, 1, 32]>}> : memref<?xf32>, vector<1x32x32xindex>, vector<1x32x32xi1> -> vector<1x32x32xf32>
%15 = vector.multi_reduction <add>, %14, %cst_3 [1] : vector<1x32x32xf32> to vector<1x32xf32>
%anchor = xegpu.convert_layout %15
<{
@@ -834,7 +825,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<128xf32>
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [0]
@@ -869,7 +860,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<2x2xf32>
%offset = arith.constant dense<0> : vector<2x2x128x128xindex>
%mask = arith.constant dense<true> : vector<2x2x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>} : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [1, 1, 32, 32]>}> : memref<?xf32>, vector<2x2x128x128xindex>, vector<2x2x128x128xi1> -> vector<2x2x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<2x2x128x128xf32> to vector<2x2xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -901,7 +892,7 @@ gpu.module @test_distribution {
%cst = arith.constant dense<0.0> : vector<32x32xf32>
%offset = arith.constant dense<0> : vector<32x32x128x128xindex>
%mask = arith.constant dense<true> : vector<32x32x128x128xi1>
- %load = xegpu.load %src[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>} : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
+ %load = xegpu.load %src[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [2, 2, 4, 4], sg_data = [16, 16, 32, 32]>}> : memref<?xf32>, vector<32x32x128x128xindex>, vector<32x32x128x128xi1> -> vector<32x32x128x128xf32>
%reduce = vector.multi_reduction <add>, %load, %cst [2, 3] : vector<32x32x128x128xf32> to vector<32x32xf32>
%anchor = xegpu.convert_layout %reduce
<{
@@ -943,7 +934,7 @@ gpu.module @test_distribution {
// CHECK: %[[CONVERT_B:.*]] = xegpu.convert_layout %{{.*}} <{input_layout = #xegpu.layout<inst_data = [32, 16]>, target_layout = #xegpu.layout<inst_data = [16, 16]>}> : vector<32x32xf16>
%9 = xegpu.convert_layout %7 <{input_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [32, 16]>, target_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>}> : vector<256x32xf16>
%10 = xegpu.convert_layout %8 <{input_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [32, 16]>, target_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [16, 16]>}> : vector<32x256xf16>
- %11 = xegpu.dpas %9, %10, %arg4 {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>} : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf32> -> vector<256x256xf32>
+ %11 = xegpu.dpas %9, %10, %arg4 layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>) layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [16, 16]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>) : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf32> -> vector<256x256xf32>
scf.yield %11 : vector<256x256xf32>
}
xegpu.store_nd %6, %2[%0, %1] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>}> : vector<256x256xf32>, !xegpu.tensor_desc<256x256xf32, #xegpu.block_tdesc_attr<boundary_check = false>, #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32], inst_data = [8, 16]>>
@@ -984,7 +975,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[LOAD_OFF_X:.*]] = arith.remui %[[LOAD_MUL_X]], %[[C256:.*]] : index
// CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MDESC]][%[[LOAD_OFF_Y]], %[[LOAD_OFF_X]]] <{layout = #xegpu.layout<inst_data = [16, 16]>}>: !xegpu.mem_desc<128x256xf32>, index, index -> vector<16x32xf32>
%0 = xegpu.create_nd_tdesc %arg0 : memref<128x256xf32> -> !xegpu.tensor_desc<128x256xf32>
- %1 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [4, 16], sg_data = [32, 16], inst_data = [16, 16]>} : !xegpu.tensor_desc<128x256xf32> -> vector<128x256xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<sg_layout = [4, 16], sg_data = [32, 16], inst_data = [16, 16]>}> : !xegpu.tensor_desc<128x256xf32> -> vector<128x256xf32>
%2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [4, 16], sg_data = [32, 16], inst_data = [16, 16]>,
target_layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32], inst_data = [16, 16]>}> : vector<128x256xf32>
%anchor = xegpu.convert_layout %2
@@ -1026,7 +1017,7 @@ gpu.module @test_distribution {
// CHECK-DAG: %[[LOAD_SLM:.*]] = xegpu.load_matrix %[[MDESC]][%[[LOAD_OFF_Z]], %[[LOAD_OFF_Y]], %[[LOAD_OFF_X]]] <{layout = #xegpu.layout<inst_data = [1, 16, 16]>}>: !xegpu.mem_desc<8x128x256xf32>, index, index, index -> vector<1x16x32xf32>
%offset = arith.constant dense<0> : vector<8x128x256xindex>
%mask = arith.constant dense<true> : vector<8x128x256xi1>
- %1 = xegpu.load %arg0[%offset], %mask {chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>} : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
+ %1 = xegpu.load %arg0[%offset], %mask <{chunk_size = 1, layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>}> : memref<?xf32>, vector<8x128x256xindex>, vector<8x128x256xi1> -> vector<8x128x256xf32>
%2 = xegpu.convert_layout %1 <{input_layout = #xegpu.layout<sg_layout = [8, 4, 16], sg_data = [1, 32, 16], inst_data = [1, 16, 16]>,
target_layout = #xegpu.layout<sg_layout = [8, 8, 8], sg_data = [1, 16, 32], inst_data = [1, 16, 16]>}> : vector<8x128x256xf32>
%anchor = xegpu.convert_layout %2
@@ -1062,7 +1053,7 @@ gpu.module @test_distribution {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x256xf32>
-> !xegpu.tensor_desc<256x256xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [32, 32]>}>
: !xegpu.tensor_desc<256x256xf32>
-> vector<256x256xf32>
@@ -1119,7 +1110,7 @@ gpu.module @test_distribution {
%0 = arith.muli %block_id_x, %c128 : index
%1 = arith.muli %block_id_y, %c128 : index
%2 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32>
- %3 = xegpu.load_nd %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %3 = xegpu.load_nd %2[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
%4 = xegpu.create_nd_tdesc %arg0 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16>
%5 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xf16> -> !xegpu.tensor_desc<128x128xf16>
@@ -1132,18 +1123,15 @@ gpu.module @test_distribution {
// CHECK: scf.yield %[[C]] : vector<16x16xf32>
%6 = scf.for %arg3 = %c0 to %c1024 step %c128 iter_args(%arg6 = %3)
-> (vector<128x128xf32>) {
- %8 = xegpu.load_nd %4[0, %arg3] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>} : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
- %9 = xegpu.load_nd %5[%arg3, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>} : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
- %10 = xegpu.dpas %8, %9, %arg6
- {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>,
- layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>,
- layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}
+ %8 = xegpu.load_nd %4[0, %arg3] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}> : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %9 = xegpu.load_nd %5[%arg3, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}> : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
+ %10 = xegpu.dpas %8, %9, %arg6 layout_a(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>) layout_b(#xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>) layout_cd(#xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>)
: vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
scf.yield %10 : vector<128x128xf32>
}
%7 = xegpu.create_nd_tdesc %arg2 : memref<1024x1024xf32>
-> !xegpu.tensor_desc<128x128xf32>
- xegpu.store_nd %6, %7[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]> } : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ xegpu.store_nd %6, %7[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]> }> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
gpu.return
}
@@ -1153,7 +1141,7 @@ gpu.module @test_distribution {
%c0_i32 = arith.constant 0 : i32
%c256 = arith.constant 256 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
- %1 = xegpu.load_nd %0[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %1 = xegpu.load_nd %0[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
%2 = xegpu.create_nd_tdesc %arg1 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
// CHECK: scf.while {{.*}} : (vector<16xf32>, i32) -> (vector<16xf32>, i32)
@@ -1164,12 +1152,12 @@ gpu.module @test_distribution {
} do {
// CHECK: (%[[ARG2:.*]]: vector<16xf32>, %[[ARG3:.*]]: i32)
^bb0(%arg2: vector<256xf32>, %arg3: i32):
- xegpu.store_nd %arg2, %2[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %arg2, %2[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
%4 = arith.addi %arg3, %c1_i32 : i32
- %6 = xegpu.load_nd %0[%c256] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %6 = xegpu.load_nd %0[%c256] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
scf.yield %6, %4 : vector<256xf32>, i32
}
- xegpu.store_nd %3#0, %2[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %3#0, %2[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -1179,9 +1167,9 @@ gpu.module @test_distribution {
%c10_i32 = arith.constant 10 : i32
%c0_i32 = arith.constant 0 : i32
%0 = xegpu.create_nd_tdesc %arg0 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32>
- %1 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %1 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
%2 = xegpu.create_nd_tdesc %arg1 : memref<1024x1024xf32> -> !xegpu.tensor_desc<128x128xf32>
- %3 = xegpu.load_nd %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %3 = xegpu.load_nd %2[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
// CHECK: scf.while {{.*}} : (vector<16x16xf32>, vector<16x16xf32>, i32) -> (vector<16x16xf32>, vector<16x16xf32>, i32)
%4:3 = scf.while (%arg2 = %1, %arg3 = %3, %arg4 = %c0_i32) : (vector<128x128xf32>, vector<128x128xf32>, i32) -> (vector<128x128xf32>, vector<128x128xf32>, i32) {
@@ -1192,12 +1180,12 @@ gpu.module @test_distribution {
// CHECK: (%{{.*}}: vector<16x16xf32>, %{{.*}}: vector<16x16xf32>, %{{.*}}: i32)
^bb0(%arg2: vector<128x128xf32>, %arg3: vector<128x128xf32>, %arg4: i32):
%nx = arith.addi %arg4, %c1_i32 : i32
- %ld0 = xegpu.load_nd %0[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
- %ld1 = xegpu.load_nd %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %ld0 = xegpu.load_nd %0[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
+ %ld1 = xegpu.load_nd %2[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : !xegpu.tensor_desc<128x128xf32> -> vector<128x128xf32>
scf.yield %ld0, %ld1, %nx : vector<128x128xf32>, vector<128x128xf32>, i32
}
- xegpu.store_nd %4#0, %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
- xegpu.store_nd %4#1, %2[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ xegpu.store_nd %4#0, %2[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
+ xegpu.store_nd %4#1, %2[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
gpu.return
}
@@ -1214,19 +1202,19 @@ gpu.module @test_distribution {
%5 = scf.if %4 -> (vector<256xf32>) {
// CHECK-LABEL: xegpu.load_nd
// CHECK-SAME: !xegpu.tensor_desc<16xf32> -> vector<16xf32>
- %2 = xegpu.load_nd %0[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %2 = xegpu.load_nd %0[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// CHECK-LABEL: scf.yield
// CHECK-SAME: vector<16xf32>
scf.yield %2 : vector<256xf32>
} else {
// CHECK-LABEL: xegpu.load_nd
// CHECK-SAME: !xegpu.tensor_desc<16xf32> -> vector<16xf32>
- %3 = xegpu.load_nd %1[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %3 = xegpu.load_nd %1[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
// CHECK-LABEL: scf.yield
// CHECK-SAME: vector<16xf32>
scf.yield %3 : vector<256xf32>
}
- xegpu.store_nd %5, %0[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %5, %0[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -1235,7 +1223,7 @@ gpu.module @test_distribution {
%id = gpu.subgroup_id : index
%t = xegpu.create_nd_tdesc %arg0 : memref<1024xf32> -> !xegpu.tensor_desc<256xf32>
- %d = xegpu.load_nd %t[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
+ %d = xegpu.load_nd %t[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : !xegpu.tensor_desc<256xf32> -> vector<256xf32>
%0 = arith.cmpi eq, %id, %c10 : index
// CHECK-LABEL: scf.if
@@ -1255,7 +1243,7 @@ gpu.module @test_distribution {
// CHECK-SAME: !xegpu.tensor_desc<16xf32>
scf.yield %3 : !xegpu.tensor_desc<256xf32>
}
- xegpu.store_nd %d, %1[0] {layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>} : vector<256xf32>, !xegpu.tensor_desc<256xf32>
+ xegpu.store_nd %d, %1[0] <{layout = #xegpu.layout<sg_layout = [16], sg_data = [16]>}> : vector<256xf32>, !xegpu.tensor_desc<256xf32>
gpu.return
}
@@ -1290,7 +1278,7 @@ gpu.module @test_distribution {
%ptr = arith.index_cast %intptr : index to i64
%mask = arith.constant dense<true> : vector<1x1x16xi1>
- xegpu.store %cst_val, %ptr[%add], %mask {layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [1, 1, 16]>} : vector<1x1x16xf32>, i64, vector<1x1x16xindex>, vector<1x1x16xi1>
+ xegpu.store %cst_val, %ptr[%add], %mask <{layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [1, 1, 16]>}> : vector<1x1x16xf32>, i64, vector<1x1x16xindex>, vector<1x1x16xi1>
gpu.return
}
@@ -1298,7 +1286,7 @@ gpu.module @test_distribution {
gpu.func @bitcast_distribution(%src: memref<256x128xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
// CHECK: vector.bitcast {{.*}} : vector<32x32xf32> to vector<32x64xi16>
@@ -1317,10 +1305,10 @@ gpu.module @test_distribution {
gpu.func @interleave_distribution(%src: memref<256x128xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
-> !xegpu.tensor_desc<256x128xf32>
- %load1 = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ %load1 = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
- %load2 = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ %load2 = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
: !xegpu.tensor_desc<256x128xf32>
-> vector<256x128xf32>
// CHECK: vector.interleave {{.*}}, {{.*}} : vector<32x32xf32> -> vector<32x64xf32>
@@ -1336,7 +1324,7 @@ gpu.module @test_distribution {
// CHECK-LABEL: @deinterleave_distribution
gpu.func @deinterleave_distribution(%src: memref<256x256xf32>) {
%tdesc = xegpu.create_nd_tdesc %src : memref<256x256xf32> -> !xegpu.tensor_desc<256x256xf32>
- %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>} : !xegpu.tensor_desc<256x256xf32> -> vector<256x256xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>}> : !xegpu.tensor_desc<256x256xf32> -> vector<256x256xf32>
// CHECK: {{.*}} = vector.deinterleave {{.*}} : vector<32x64xf32> -> vector<32x32xf32>
%deinterleave:2 = vector.deinterleave %load : vector<256x256xf32> -> vector<256x128xf32>
%anchor = xegpu.convert_layout %deinterleave#0
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_gemm.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_gemm.mlir
index 37f6d33e8ac30..05e6683497236 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_gemm.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_gemm.mlir
@@ -53,34 +53,34 @@ module @gemm attributes {gpu.container_module} {
%m = arith.muli %block_id_x, %c256 : index
%n = arith.muli %block_id_y, %c256 : index
%c_tdesc = xegpu.create_nd_tdesc %C : memref<256x256xf32> -> !xegpu.tensor_desc<256x256xf32, #c>
- %c_init_value = xegpu.load_nd %c_tdesc[%m, %n] {layout = #c}: !xegpu.tensor_desc<256x256xf32, #c> -> vector<256x256xf32>
+ %c_init_value = xegpu.load_nd %c_tdesc[%m, %n] <{layout = #c}>: !xegpu.tensor_desc<256x256xf32, #c> -> vector<256x256xf32>
%a_tdesc = xegpu.create_nd_tdesc %A : memref<256x256xf16> -> !xegpu.tensor_desc<256x32xf16, #a>
%b_tdesc = xegpu.create_nd_tdesc %B : memref<256x256xf16> -> !xegpu.tensor_desc<32x256xf16, #b>
// Prefetch A 3 times.
%a_prefetch_tdesc = xegpu.create_nd_tdesc %A : memref<256x256xf16> -> !xegpu.tensor_desc<256x32xf16, #a_prefetch>
- xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c0] {layout = #a_prefetch}: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
- xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c32] {layout = #a_prefetch}: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
- xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c64] {layout = #a_prefetch}: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
+ xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c0] <{layout = #a_prefetch}>: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
+ xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c32] <{layout = #a_prefetch}>: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
+ xegpu.prefetch_nd %a_prefetch_tdesc[%m, %c64] <{layout = #a_prefetch}>: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
// Prefetch B 3 times.
%b_prefetch_tdesc = xegpu.create_nd_tdesc %B : memref<256x256xf16> -> !xegpu.tensor_desc<32x256xf16, #b_prefetch>
- xegpu.prefetch_nd %b_prefetch_tdesc[%c0, %n] {layout = #b_prefetch}: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
- xegpu.prefetch_nd %b_prefetch_tdesc[%c32, %n] {layout = #b_prefetch}: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
- xegpu.prefetch_nd %b_prefetch_tdesc[%c64, %n] {layout = #b_prefetch}: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
+ xegpu.prefetch_nd %b_prefetch_tdesc[%c0, %n] <{layout = #b_prefetch}>: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
+ xegpu.prefetch_nd %b_prefetch_tdesc[%c32, %n] <{layout = #b_prefetch}>: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
+ xegpu.prefetch_nd %b_prefetch_tdesc[%c64, %n] <{layout = #b_prefetch}>: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
%out = scf.for %k = %c0 to %c256 step %c32
iter_args(%c_value = %c_init_value)
-> (vector<256x256xf32>) {
- %a_value = xegpu.load_nd %a_tdesc[%m, %k] {layout = #a}: !xegpu.tensor_desc<256x32xf16, #a> -> vector<256x32xf16>
- %b_value = xegpu.load_nd %b_tdesc[%k, %n] {layout = #b}: !xegpu.tensor_desc<32x256xf16, #b> -> vector<32x256xf16>
+ %a_value = xegpu.load_nd %a_tdesc[%m, %k] <{layout = #a}>: !xegpu.tensor_desc<256x32xf16, #a> -> vector<256x32xf16>
+ %b_value = xegpu.load_nd %b_tdesc[%k, %n] <{layout = #b}>: !xegpu.tensor_desc<32x256xf16, #b> -> vector<32x256xf16>
// Prefetch next tiles.
%prefetch_offset = arith.addi %k, %c96 : index
- xegpu.prefetch_nd %a_prefetch_tdesc[%m, %prefetch_offset] {layout = #a_prefetch}: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
- xegpu.prefetch_nd %b_prefetch_tdesc[%prefetch_offset, %n] {layout = #b_prefetch}: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
- %c_new_value = xegpu.dpas %a_value, %b_value, %c_value {layout_a = #a, layout_b = #b, layout_cd = #c}
+ xegpu.prefetch_nd %a_prefetch_tdesc[%m, %prefetch_offset] <{layout = #a_prefetch}>: !xegpu.tensor_desc<256x32xf16, #a_prefetch>
+ xegpu.prefetch_nd %b_prefetch_tdesc[%prefetch_offset, %n] <{layout = #b_prefetch}>: !xegpu.tensor_desc<32x256xf16, #b_prefetch>
+ %c_new_value = xegpu.dpas %a_value, %b_value, %c_value layout_a(#a) layout_b(#b) layout_cd(#c)
: vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf32> -> vector<256x256xf32>
scf.yield %c_new_value : vector<256x256xf32>
}
- xegpu.store_nd %out, %c_tdesc[%m, %n] {layout = #c}: vector<256x256xf32>, !xegpu.tensor_desc<256x256xf32, #c>
+ xegpu.store_nd %out, %c_tdesc[%m, %n] <{layout = #c}>: vector<256x256xf32>, !xegpu.tensor_desc<256x256xf32, #c>
gpu.return
}
}
More information about the Mlir-commits
mailing list