From 1b6dc9c8c995b65267e8dbb1aea39cbf5f3179ec Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Mon, 7 Jul 2025 23:43:12 -0700 Subject: [PATCH 01/15] Add matrix_desc and operations --- docs/rfcs/XeGPU.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 16e52b48f..b9bdd3ff1 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -329,6 +329,20 @@ Attribute `Memory_kind` describes the memory kind. "global" means the global mem `nbarrier` and `fence` operations lower to uniform instructions, so there is no need to specify the `sg_map`. +## XeGPU operations to access share local memory +Users must create `matrix_desc` to hold a matrix in the share local memory. The matrix must be row-major. The matrix can attach a attribute for its memory layout, for example, a blocked layout or just original non-blocked row-major layout (aka. linear layout). +User can get a subview of an existing `matrix_desc` to get a new `matrix_desc`, potentially having a stride. Then user can use load_matrix and store_matrix to move the matrix data between share local memory and vectors (registers). The matrix is typically 2d and but can be multi-dimension. XeGPU's load_matrix and store_matrix works at workgroup level only. It uses xegpu.layout to describe how the matrix is decomposed to data fragments and maps to work items. The workgroup level operation loads the entire matrix to vector. + +The motivation of `matrix_desc` data type and related operations is to simplify the programming model. Rather than trying to reuse `tensor_desc` to describe the matrix/tile in the share local memory, it is straightforward to use a dedicate data type to describe it. The use of share local memory is usually very local not exposed to workgroup level user, for example, supporting the lowering of transpose, reduction, and convert layout operations. So the createion of matrix_desc doesn't take a memref as input and implictly allocate share local memory. The share local memory may be blocked to facilitate the optimized lowering to load chunk or 1d block load. + + +| Ops | Syntax | Example | +| :--- | :---- | :--- | +|create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16, @mem_layout=block[8, 16] > | +|matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @layout_type=1> -> matrix_desc<128x128xbf16, @row_stride=256, @mem_layout=block[8, 16]> | +|load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]> -> vector<128x256xbf16> | +|store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]>, vector<128x256xbf16> | + ## XeGPU Attributes to support Work Item Level semantics **Attribute xegpu.sg_map** From 39198a30b84fb3f462131b48e389f200b7102d71 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Tue, 8 Jul 2025 09:04:15 -0700 Subject: [PATCH 02/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index b9bdd3ff1..0a54c6508 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -343,6 +343,26 @@ The motivation of `matrix_desc` data type and related operations is to simplify |load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]> -> vector<128x256xbf16> | |store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]>, vector<128x256xbf16> | +User creates `matrix_desc` to hold a matrix in the share local memory. The operation allocates a share local memory for the matrix, assuming the matrix is row-major and contiguous. The block attribute indicates the matrix has a blocked layout. +```mlir +%mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16, @mem_layout=block[8, 16]> +``` +User creates a subview of matrix. +```mlir +%mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16> +%mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, row_stride=128> +``` + +Users load a matrix from share local memory to vector. +```mlir +vec_a = load_matrix matrix_desc_a: matrix_desc<256x128xbf16, @mem_layout=block[8, 16]> -> vector<256x128xbf6> +``` + +Users store a matrix to share local memory from vector. +```mlir +store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @mem_layout=block[8, 16]>, vector<256x128xbf6> +``` + ## XeGPU Attributes to support Work Item Level semantics **Attribute xegpu.sg_map** From c1ab2984fad72d309e3ddca20b9d57153ffdf87b Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Tue, 8 Jul 2025 20:03:53 -0700 Subject: [PATCH 03/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 86 ++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 75 insertions(+), 11 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 0a54c6508..a4d576f5e 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -338,31 +338,95 @@ The motivation of `matrix_desc` data type and related operations is to simplify | Ops | Syntax | Example | | :--- | :---- | :--- | -|create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16, @mem_layout=block[8, 16] > | -|matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @layout_type=1> -> matrix_desc<128x128xbf16, @row_stride=256, @mem_layout=block[8, 16]> | -|load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]> -> vector<128x256xbf16> | -|store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @mem_layout=block[8, 16]>, vector<128x256xbf16> | +|create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16> | +|matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @layout_type=1> -> matrix_desc<128x128xbf16, @stride=[256,1], @block=[8, 16]> | +|load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @block=[8, 16]> -> vector<128x256xbf16> | +|store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @block=[8, 16]>, vector<128x256xbf16> | -User creates `matrix_desc` to hold a matrix in the share local memory. The operation allocates a share local memory for the matrix, assuming the matrix is row-major and contiguous. The block attribute indicates the matrix has a blocked layout. +User creates `matrix_desc` to hold a matrix in the share local memory. The operation allocates a share local memory for the matrix, assuming the matrix is row-major and contiguous. ```mlir -%mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16, @mem_layout=block[8, 16]> +%mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> ``` -User creates a subview of matrix. +User creates a subview of matrix. The new matrix maybe associated with `block` and `strides` atttribute to describe the memory layout. The `strides` attributes allows matrix_desc being further decomposed to subgroup and work item level. The `block` attribute indicates the matrix has a blocked layout. ```mlir -%mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16> -%mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, row_stride=128> +%mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> +%mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> ``` Users load a matrix from share local memory to vector. ```mlir -vec_a = load_matrix matrix_desc_a: matrix_desc<256x128xbf16, @mem_layout=block[8, 16]> -> vector<256x128xbf6> +vec_a = load_matrix matrix_desc_a: matrix_desc<256x128xbf16, @block=[8, 16]> -> vector<256x128xbf6> ``` Users store a matrix to share local memory from vector. ```mlir -store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @mem_layout=block[8, 16]>, vector<256x128xbf6> +store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @block=[8, 16]>, vector<256x128xbf6> ``` +**Cooperative Transpose Example** +Suppose we have wg-level user input code +```mlir +#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[1, 0] } +#Coop_wg = {sg_layout = [8, 4] , sg_data= [32, 8], order=[1, 0] } +#dpas_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } + +%at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > +%a = vector.transpose %1 #Coop_wg :vector<32x256xf16> -> vector<256x32xf16> +%a_dpas = Conv_layout %2 #Coop_wg #dpas_wg +``` + +After an optimization pass which optimize the transpose-A pattern, the transformed code uses store_matrix and load_matrix. Note the load_nd and store_matrix has smaller sg_data so the subgroups perform cooperative transpose. +```mlir +#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[1, 0] } +#dpas_t_wg = {sg_layout = [4, 8], sg_data= [32, 32], order=[1, 0] } + +%at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > +%m1 = create_matrix_desc : matrix_desc<32x256xf16> +%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, strides=[1, 32], #coop_wg> +store_matrix %m1t[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, strides=[1, 32], #Coop_t_wg> +barrier +%m4 = matrix_desc_subview : matrix_desc<256x32xf16, #dpas_t_wg > +%a_dpas = load_matrix %m2 [0, 0] #dpas_t_wg : vector<256x32xf16>, matrix_desc<256x32xf16, #dpas_t_wg> +``` + +After wg->sg level distribution and blocking, this lowers to the following sg-level code. +```mlir +%at1 = load_nd %tdesc, sg_coords1: tensor_desc<32x256xf16> -> vector<8x16xf16> +%at2 = load_nd %tdesc, sg_coords2: tensor_desc<32x256xf16> -> vector<8x16xf16> +%m1 = create_matrix_desc : matrix_desc<32x256xf16> + +%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> +store_matrix %m1t, sg_slm_coord1, %at1: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> +store_matrix %m1t, sg_slm_coord2, %at2: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> +barrier +%m4 = matrix_desc_subview : matrix_desc<256x32xf16, block=[16, 16] > +%a_dpas1 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> +%a_dpas2 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> +%a_dpas3 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> +%a_dpas4 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> +``` + +After assigned with lane_layout +```mlir +#Coop_t_lane ={lane_layout = [1, 16] , lane_data= [1, 1]} +#dpas_t_lane = {lane_layout = [2, 8], lane_data= [1, 2]} + +%at1 = load_nd %tdesc, sg_coords1: tensor_desc<32x256xf16, #Coop_t_lane> -> vector<8x16xf16> +%at2 = load_nd %tdesc, sg_coords2: tensor_desc<32x256xf16, #Coop_t_lane> -> vector<8x16xf16> +%m1 = create_matrix_desc : matrix_desc<32x256xf16> + +%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> +store_matrix %m1t, sg_slm_coord1, %at1: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> +store_matrix %m1t, sg_slm_coord2, %at2: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> +barrier +%m4 = matrix_desc_subview : matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane > +%a_dpas1 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> +%a_dpas2 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> +%a_dpas3 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> +%a_dpas4 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> +``` + + ## XeGPU Attributes to support Work Item Level semantics **Attribute xegpu.sg_map** From 1bbcdaa6237c4fa57e86609d23f62fa09098cd51 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Tue, 8 Jul 2025 20:05:41 -0700 Subject: [PATCH 04/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index a4d576f5e..c5d12038f 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -408,7 +408,7 @@ barrier After assigned with lane_layout ```mlir -#Coop_t_lane ={lane_layout = [1, 16] , lane_data= [1, 1]} +#Coop_t_lane ={lane_layout = [1, 16] , lane_data= [8, 1]} #dpas_t_lane = {lane_layout = [2, 8], lane_data= [1, 2]} %at1 = load_nd %tdesc, sg_coords1: tensor_desc<32x256xf16, #Coop_t_lane> -> vector<8x16xf16> From 8906c08928001759a8252551de1057034f7f533f Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Wed, 9 Jul 2025 07:51:05 -0700 Subject: [PATCH 05/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index c5d12038f..52bae60a2 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -366,7 +366,7 @@ store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @block=[8, 16]>, ve **Cooperative Transpose Example** Suppose we have wg-level user input code ```mlir -#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[1, 0] } +#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1] } #Coop_wg = {sg_layout = [8, 4] , sg_data= [32, 8], order=[1, 0] } #dpas_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } @@ -377,7 +377,7 @@ Suppose we have wg-level user input code After an optimization pass which optimize the transpose-A pattern, the transformed code uses store_matrix and load_matrix. Note the load_nd and store_matrix has smaller sg_data so the subgroups perform cooperative transpose. ```mlir -#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[1, 0] } +#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1] } #dpas_t_wg = {sg_layout = [4, 8], sg_data= [32, 32], order=[1, 0] } %at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > From 3c0c8ad99cebee22554a33bf099bfd7f8454cdb1 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Wed, 9 Jul 2025 16:20:37 -0700 Subject: [PATCH 06/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 113 +++++++++++++++++++++++++++++++-------------- 1 file changed, 79 insertions(+), 34 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 52bae60a2..227460a78 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -377,55 +377,100 @@ Suppose we have wg-level user input code After an optimization pass which optimize the transpose-A pattern, the transformed code uses store_matrix and load_matrix. Note the load_nd and store_matrix has smaller sg_data so the subgroups perform cooperative transpose. ```mlir -#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1] } -#dpas_t_wg = {sg_layout = [4, 8], sg_data= [32, 32], order=[1, 0] } +#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1 } +#dpas_t_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } %at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > -%m1 = create_matrix_desc : matrix_desc<32x256xf16> -%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, strides=[1, 32], #coop_wg> -store_matrix %m1t[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, strides=[1, 32], #Coop_t_wg> +%m = create_matrix_desc : matrix_desc<32x256xf16> +%mt = matrix_desc_subview %m: matrix_desc<32x256xf16, strides=[1, 32], #coop_t_wg> +store_matrix %mt[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, strides=[1, 32], #coop_t_wg> barrier -%m4 = matrix_desc_subview : matrix_desc<256x32xf16, #dpas_t_wg > -%a_dpas = load_matrix %m2 [0, 0] #dpas_t_wg : vector<256x32xf16>, matrix_desc<256x32xf16, #dpas_t_wg> +%ma = matrix_desc_subview : matrix_desc<256x32xf16, #dpas_t_wg> +%a_dpas = load_matrix %ma [0, 0] #dpas_t_wg : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` -After wg->sg level distribution and blocking, this lowers to the following sg-level code. +After wg->sg level distribution, this lowers to the following sg-level code. ```mlir -%at1 = load_nd %tdesc, sg_coords1: tensor_desc<32x256xf16> -> vector<8x16xf16> -%at2 = load_nd %tdesc, sg_coords2: tensor_desc<32x256xf16> -> vector<8x16xf16> -%m1 = create_matrix_desc : matrix_desc<32x256xf16> +#coop_t_inst ={ inst_data=[8, 16] } +#dpas_t_inst = {inst_data=[16, 16] } +create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16, #coop_t_inst> +%at = load_nd %tdesc: tensor_desc<8x32xf16, #coop_t_inst> -> vector<8x32xf16> -%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -store_matrix %m1t, sg_slm_coord1, %at1: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -store_matrix %m1t, sg_slm_coord2, %at2: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> +%m = create_matrix_desc : matrix_desc<32x256xf16> +%mt_sg = create_matrix_desc %m [sg_idy*8, sg_idx*32]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #coop_t_inst > +store_matrix %mt_sg, %at: vector<8x32xf16>, matrix_desc<8x32xf16, block=[16, 16], strides=[1, 32], #coop_t_inst > barrier -%m4 = matrix_desc_subview : matrix_desc<256x32xf16, block=[16, 16] > -%a_dpas1 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> -%a_dpas2 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> -%a_dpas3 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> -%a_dpas4 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16]> +%ma = matrix_desc_subview %m : matrix_desc<256x32xf16, block=[16, 16], #dpas_t_inst> +%ma_sg = matrix_desc_subview %ma [sg_idy*32, sg_idx*32%32]: matrix_desc<32x32xf16, block=[16, 16] , #dpas_t_inst> +%a_dpas = load_matrix ma_sg: matrix_desc<32x32xf16, block=[16, 16], #dpas_t_inst >-> vector<32x32xf16> + ``` -After assigned with lane_layout +After blocking according to inst_data. ```mlir -#Coop_t_lane ={lane_layout = [1, 16] , lane_data= [8, 1]} -#dpas_t_lane = {lane_layout = [2, 8], lane_data= [1, 2]} +create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16> +%at = load_nd %tdesc, sg_coords1: tensor_desc<8x32xf16> -> vector<8x32xf16> +%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> +%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> +%m = create_matrix_desc : matrix_desc<32x256xf16> +%mt_sg = create_matrix_desc %m1 [0, 0]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -%at1 = load_nd %tdesc, sg_coords1: tensor_desc<32x256xf16, #Coop_t_lane> -> vector<8x16xf16> -%at2 = load_nd %tdesc, sg_coords2: tensor_desc<32x256xf16, #Coop_t_lane> -> vector<8x16xf16> -%m1 = create_matrix_desc : matrix_desc<32x256xf16> +%mt_inst0 = create_matrix_desc % mt_sg [sg_idy*8, sg_idx*32]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -> matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> +%mt_inst1 = create_matrix_desc % mt_sg [sg_idy*8, sg_idx*32+16]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -> matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> -%m1t = matrix_desc_subview %m1: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> -store_matrix %m1t, sg_slm_coord1, %at1: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> -store_matrix %m1t, sg_slm_coord2, %at2: vector<8x16xf16>, matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #Coop_t_lane> +store_matrix %mt_inst0, %at0: vector<8x16xf16>, matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> +store_matrix %mt_inst1, %at1: vector<8x16xf16>, matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> barrier -%m4 = matrix_desc_subview : matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane > -%a_dpas1 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> -%a_dpas2 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> -%a_dpas3 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> -%a_dpas4 = load_matrix %m2 sg_slm_coord3 #dpas_t_wg : vector<16x16xf16>, matrix_desc<256x32xf16, block=[16, 16], #dpas_t_lane> -``` +%ma = matrix_desc_subview %m: matrix_desc<256x32xf16, block=[16, 16] > +%ma_inst0 = matrix_desc_subview % ma [sg_idy*32, sg_idx*32%32]: matrix_desc<16x16xf16, block=[16, 16] > +% ma_inst1 = matrix_desc_subview %ma [sg_idy*32, sg_idx*32%32+16]: matrix_desc<16x16xf16, block=[16, 16] > +% ma_inst2 = matrix_desc_subview %ma [sg_idy*32+16, sg_idx*32%32]: matrix_desc<16x16xf16, block=[16, 16] > +% ma_inst3 = matrix_desc_subview %ma [sg_idy*32+16, sg_idx*32%32+16]: matrix_desc<16x16xf16, block=[16, 16] > +%a_dpas_0 = load_matrix ma_inst0: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> +%a_dpas_1 = load_matrix ma_inst1: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> +%a_dpas_2 = load_matrix ma_inst2: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> +%a_dpas_3 = load_matrix ma_inst3: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> +``` + +MaterializeSLMAccess pass replace matrix_desc and related ops to memref and load and load_1d. Pseudo code used to simplify the code. +```mlir +create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16> +%at = load_nd %tdesc, sg_coords1: tensor_desc<8x32xf16> -> vector<8x32xf16> +%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> +%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> + +%blk_y=sg_idy*8 /16: index +%blk_in_y=sg_idy*8 %16: index +%sg_idx_vec = %sg_idx*32 + [0..15] : vector<16xindex> +%blk_x=%sg_idx_vec /16: vector<16xindex > +%blk_in_x=%sg_idx_vec %16: vector<16xindex > +%sg_start_offset_vec = %blk_y * 16 + %blk_in_y + %blk_x * 512 + %blk_in_x*16 +%tdesc0 = xegpu.create_tdesc %m, %sg_start_offset_vec: memref<8192xf16, 3>, %sg_start_offset_vec ->tdesc<8x16xf16, chunk=8, scope=slm> + +%sg_idx_vec2 = %sg_idx*32 + [16..31] : vector<16xindex> +%blk_x2=%sg_idx_vec /16: vector<16xindex > +%blk_in_x2=%sg_idx_vec %16: vector<16xindex > +%sg_start_offset_vec2 = %blk_y * 16 + %blk_in_y + %blk_x * 512 + %blk_in_x*16 +%tdesc1 = xegpu.create_tdesc %m, %sg_start_offset_vec2: memref<8192xf16, 3>, %sg_start_offset_vec ->tdesc<8x16xf16, chunk=8, scope=slm> + +xegpu.store %tdesc0, %at0: tdesc<8x32xf16, chunk=8, scope=slm>, vector<8x16xf16> +xegpu.store %tdesc1, %at1: tdesc<8x32xf16, chunk=8, scope=slm>, vector<8x16xf16> +barrier +%inst_start_offset0 = sg_idy*2* 512 +%tdesc0 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > +%inst_start_offset0 = sg_idy*2* 512 + 256 +%tdesc1 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > +%inst_start_offset0 = sg_idy*2* 512 + 512 +%tdesc2 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > +%inst_start_offset0 = sg_idy*2* 512 + 512 + 256 +%tdesc3 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > + +a_dpas_0 = Load_nd %tdesc0: tdesc<256xf16 > -> vector<256xf16> +a_dpas_1 = Load_nd %tdesc1: tdesc<256xf16 > -> vector<256xf16> +a_dpas_2 = Load_nd %tdesc2: tdesc<256xf16 > -> vector<256xf16> +a_dpas_3 = Load_nd %tdesc3: tdesc<256xf16 > -> vector<256xf16> +``` ## XeGPU Attributes to support Work Item Level semantics From bfc834b7b9bf387b8cb79c16a652b138aeaa0c0e Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Thu, 10 Jul 2025 19:43:36 -0700 Subject: [PATCH 07/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 227460a78..b8ebaaa8a 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -330,10 +330,9 @@ Attribute `Memory_kind` describes the memory kind. "global" means the global mem `nbarrier` and `fence` operations lower to uniform instructions, so there is no need to specify the `sg_map`. ## XeGPU operations to access share local memory -Users must create `matrix_desc` to hold a matrix in the share local memory. The matrix must be row-major. The matrix can attach a attribute for its memory layout, for example, a blocked layout or just original non-blocked row-major layout (aka. linear layout). -User can get a subview of an existing `matrix_desc` to get a new `matrix_desc`, potentially having a stride. Then user can use load_matrix and store_matrix to move the matrix data between share local memory and vectors (registers). The matrix is typically 2d and but can be multi-dimension. XeGPU's load_matrix and store_matrix works at workgroup level only. It uses xegpu.layout to describe how the matrix is decomposed to data fragments and maps to work items. The workgroup level operation loads the entire matrix to vector. +XeGPU introduced `matrix_desc` data type to simplify programming share local memory (slm). In Xe2, the slm access is not directly programmed by user, but as an implementation to support transpose, reduction, and convert layout of a workgroup level tile. There is a common programming pattern that user allocates a 2d matrix in slm with row-major contiguous layout, then distribute the matrix to each subgroup and lane. The distribution process involves complex address computation, espeically when the memory access try to view the 2d matrix in a transposed view. The address computation becomes even more complicated for Xe2's 1d block load as it requires to block the slm so the innermost block is contiguous in slm. `matrix_desc` data type simplified the distribution by encoding the transposed and blocked layout as attribute, which separates the logical addresses compution in distribution and later physical address computation. The distribution process works on a logical address on top of row-major contiguous view of 2d matrix, and later materialized to physical address using the slm's memory layout attributes as required by 1d block load and regular load. -The motivation of `matrix_desc` data type and related operations is to simplify the programming model. Rather than trying to reuse `tensor_desc` to describe the matrix/tile in the share local memory, it is straightforward to use a dedicate data type to describe it. The use of share local memory is usually very local not exposed to workgroup level user, for example, supporting the lowering of transpose, reduction, and convert layout operations. So the createion of matrix_desc doesn't take a memref as input and implictly allocate share local memory. The share local memory may be blocked to facilitate the optimized lowering to load chunk or 1d block load. +Users must create `matrix_desc` to hold a matrix in the share local memory. The matrix must be row-major. The matrix can attach a attribute for its memory layout, for example, a blocked layout or just original non-blocked row-major layout (aka. linear layout). User can get a subview of an existing `matrix_desc` to get a new `matrix_desc`, potentially having strided and blocked layout attributes. Then user can use load_matrix and store_matrix to move the matrix data between slm and vectors (registers). The matrix is typically 2d and but can be multi-dimension. XeGPU's load_matrix and store_matrix works at workgroup scope only. | Ops | Syntax | Example | @@ -344,10 +343,11 @@ The motivation of `matrix_desc` data type and related operations is to simplify |store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @block=[8, 16]>, vector<128x256xbf16> | User creates `matrix_desc` to hold a matrix in the share local memory. The operation allocates a share local memory for the matrix, assuming the matrix is row-major and contiguous. + ```mlir %mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> ``` -User creates a subview of matrix. The new matrix maybe associated with `block` and `strides` atttribute to describe the memory layout. The `strides` attributes allows matrix_desc being further decomposed to subgroup and work item level. The `block` attribute indicates the matrix has a blocked layout. +User creates a subview of matrix. The new matrix maybe associated with `block` and `strides` atttributes to describe the memory layout. The `strides` attributes allows matrix_desc being further decomposed to subgroup and work item level. The `block` attribute indicates the matrix has a blocked layout. The `block` attribute facilitates the optimized lowering to 1d block load, and `strides` for load with chunk. It can also attach `xegpu.layout` attribute to describe how the matrix is decomposed to data fragments and maps to work items. ```mlir %mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> %mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> From be46fadf0d598b3679a1d648716cdfcd7c742cbd Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Thu, 10 Jul 2025 21:59:27 -0700 Subject: [PATCH 08/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index b8ebaaa8a..bae6f5ddf 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -329,12 +329,40 @@ Attribute `Memory_kind` describes the memory kind. "global" means the global mem `nbarrier` and `fence` operations lower to uniform instructions, so there is no need to specify the `sg_map`. -## XeGPU operations to access share local memory -XeGPU introduced `matrix_desc` data type to simplify programming share local memory (slm). In Xe2, the slm access is not directly programmed by user, but as an implementation to support transpose, reduction, and convert layout of a workgroup level tile. There is a common programming pattern that user allocates a 2d matrix in slm with row-major contiguous layout, then distribute the matrix to each subgroup and lane. The distribution process involves complex address computation, espeically when the memory access try to view the 2d matrix in a transposed view. The address computation becomes even more complicated for Xe2's 1d block load as it requires to block the slm so the innermost block is contiguous in slm. `matrix_desc` data type simplified the distribution by encoding the transposed and blocked layout as attribute, which separates the logical addresses compution in distribution and later physical address computation. The distribution process works on a logical address on top of row-major contiguous view of 2d matrix, and later materialized to physical address using the slm's memory layout attributes as required by 1d block load and regular load. +## matrix_desc Type: Simplified Shared Local Memory (SLM) Abstraction +To streamline programming of shared local memory (SLM) on Intel Xe architecture, the XeGPU dialect introduces a new type: matrix_desc. This abstraction is designed to simplify the management of workgroup-level tiles in SLM, especially in scenarios involving layout transformations such as transpose, reduction, and blocking. +**Background and Motivation** +On Xe2 GPUs, SLM remains accessible for direct use by programmers. However, in tile-based programming — particularly when applying layout transformations such as transpose, re-layout — SLM is more commonly used as a backing store to facilitate structured tile movement across subgroups and lanes. + +Prior to the introduction of matrix_desc, SLM usage was modeled using the nd_tdesc type, which was originally designed for global memory access. As such, it lacked layout-specific attributes like blocking and stride metadata, which are essential for modeling tiled or transposed views in SLM. Developers were responsible for manually computing physical addresses — a process that became particularly complex when applying transformations such as transpose or blocking as required by chunked load or 1D block load. + +This complexity was further compounded by hierarchical distribution, where workgroup-level tiles are subdivided across subgroups, instructions, and individual lanes — each step requiring separate address transformation logic. This made the code error-prone and difficult to optimize. + +**Design and Semantics** + +The matrix_desc type addresses these challenges by: + +-Encoding layout transformations (e.g., transpose, blocking) as static attributes of the descriptor. + +-Separating logical and physical address computation: + + -The distribution and unrolling process operates on a conceptual row-major 2D matrix. + + -The physical address materialization then maps logical coordinates to hardware-compliant SLM addresses, guided by layout attributes in matrix_desc. + +This separation simplifies distribution and unrolling passes and enables systematic, robust transformations during compilation. The descriptor encapsulates all necessary layout metadata to generate correct and efficient SLM access patterns — supporting both regular loads and 1D block loads — without requiring the user to write explicit address arithmetic. Users must create `matrix_desc` to hold a matrix in the share local memory. The matrix must be row-major. The matrix can attach a attribute for its memory layout, for example, a blocked layout or just original non-blocked row-major layout (aka. linear layout). User can get a subview of an existing `matrix_desc` to get a new `matrix_desc`, potentially having strided and blocked layout attributes. Then user can use load_matrix and store_matrix to move the matrix data between slm and vectors (registers). The matrix is typically 2d and but can be multi-dimension. XeGPU's load_matrix and store_matrix works at workgroup scope only. +**Basic Usage** + +To represent a matrix stored in shared local memory (SLM), users must create a matrix_desc object. The underlying memory is assumed to follow a row-major layout, and the base matrix_desc represents a raw, unannotated matrix in this layout. The base matrix may be n-dimensional. + +Only matrix_desc instances created via subview may carry an xegpu.layout attribute, which specifies the mapping of lanes and registers to fragments of the matrix. This attribute guides the tile distribution process based on the assumed row-major view of the original matrix. In addition, subviewed matrix_desc instances may carry layout metadata such as blocking and striding, which are used to control physical address computation when accessing SLM. + +Data movement between SLM and vector registers is performed using load_matrix and store_matrix, which operate at workgroup scope and require the input matrix_desc to be 2D. If the original matrix is higher-dimensional, it must be subviewed to a 2D shape before it can be used with these operations. + | Ops | Syntax | Example | | :--- | :---- | :--- | |create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16> | From c75986c3dfa4d1bdbb609232dcaecb0ccc42b393 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Thu, 10 Jul 2025 22:01:43 -0700 Subject: [PATCH 09/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index bae6f5ddf..a48ff0547 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -340,21 +340,10 @@ This complexity was further compounded by hierarchical distribution, where workg **Design and Semantics** -The matrix_desc type addresses these challenges by: - --Encoding layout transformations (e.g., transpose, blocking) as static attributes of the descriptor. - --Separating logical and physical address computation: - - -The distribution and unrolling process operates on a conceptual row-major 2D matrix. - - -The physical address materialization then maps logical coordinates to hardware-compliant SLM addresses, guided by layout attributes in matrix_desc. +The matrix_desc type addresses these challenges by encoding layout transformations—such as transpose and blocking—as static attributes of the descriptor, and by clearly separating logical and physical address computation. The distribution and unrolling process operates on a conceptual row-major 2D matrix, enabling clean and structured logical access, while the physical address materialization phase maps these logical coordinates to hardware-compliant SLM addresses, guided by the layout attributes attached to the matrix_desc. This separation simplifies distribution and unrolling passes and enables systematic, robust transformations during compilation. The descriptor encapsulates all necessary layout metadata to generate correct and efficient SLM access patterns — supporting both regular loads and 1D block loads — without requiring the user to write explicit address arithmetic. -Users must create `matrix_desc` to hold a matrix in the share local memory. The matrix must be row-major. The matrix can attach a attribute for its memory layout, for example, a blocked layout or just original non-blocked row-major layout (aka. linear layout). User can get a subview of an existing `matrix_desc` to get a new `matrix_desc`, potentially having strided and blocked layout attributes. Then user can use load_matrix and store_matrix to move the matrix data between slm and vectors (registers). The matrix is typically 2d and but can be multi-dimension. XeGPU's load_matrix and store_matrix works at workgroup scope only. - - **Basic Usage** To represent a matrix stored in shared local memory (SLM), users must create a matrix_desc object. The underlying memory is assumed to follow a row-major layout, and the base matrix_desc represents a raw, unannotated matrix in this layout. The base matrix may be n-dimensional. From d5b06840d336b14a13c1e7613f30a49cc9aab3e8 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Thu, 10 Jul 2025 22:42:48 -0700 Subject: [PATCH 10/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 58 +++++++++++++++++++++++++++++++++------------- 1 file changed, 42 insertions(+), 16 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index a48ff0547..5f94858b9 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -330,8 +330,11 @@ Attribute `Memory_kind` describes the memory kind. "global" means the global mem `nbarrier` and `fence` operations lower to uniform instructions, so there is no need to specify the `sg_map`. ## matrix_desc Type: Simplified Shared Local Memory (SLM) Abstraction + To streamline programming of shared local memory (SLM) on Intel Xe architecture, the XeGPU dialect introduces a new type: matrix_desc. This abstraction is designed to simplify the management of workgroup-level tiles in SLM, especially in scenarios involving layout transformations such as transpose, reduction, and blocking. + **Background and Motivation** + On Xe2 GPUs, SLM remains accessible for direct use by programmers. However, in tile-based programming — particularly when applying layout transformations such as transpose, re-layout — SLM is more commonly used as a backing store to facilitate structured tile movement across subgroups and lanes. Prior to the introduction of matrix_desc, SLM usage was modeled using the nd_tdesc type, which was originally designed for global memory access. As such, it lacked layout-specific attributes like blocking and stride metadata, which are essential for modeling tiled or transposed views in SLM. Developers were responsible for manually computing physical addresses — a process that became particularly complex when applying transformations such as transpose or blocking as required by chunked load or 1D block load. @@ -359,29 +362,29 @@ Data movement between SLM and vector registers is performed using load_matrix an |load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @block=[8, 16]> -> vector<128x256xbf16> | |store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @block=[8, 16]>, vector<128x256xbf16> | -User creates `matrix_desc` to hold a matrix in the share local memory. The operation allocates a share local memory for the matrix, assuming the matrix is row-major and contiguous. +Users create a `matrix_desc` to represent a matrix stored in shared local memory (SLM). The operation allocates SLM for the matrix, assuming a row-major contiguous layout. ```mlir %mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> ``` -User creates a subview of matrix. The new matrix maybe associated with `block` and `strides` atttributes to describe the memory layout. The `strides` attributes allows matrix_desc being further decomposed to subgroup and work item level. The `block` attribute indicates the matrix has a blocked layout. The `block` attribute facilitates the optimized lowering to 1d block load, and `strides` for load with chunk. It can also attach `xegpu.layout` attribute to describe how the matrix is decomposed to data fragments and maps to work items. +Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. The resulting matrix_desc may be annotated with layout attributes such as @block and @strides to describe its memory layout more precisely. The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. Additionally, a subview may carry an xegpu.layout attribute that defines how the matrix is logically partitioned into fragments and mapped to work items. ```mlir %mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> %mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> ``` -Users load a matrix from share local memory to vector. +Users can load a matrix from shared local memory into a vector value using the load_matrix operation. The result is a vector type in the IR, representing a tile stored in registers. ```mlir vec_a = load_matrix matrix_desc_a: matrix_desc<256x128xbf16, @block=[8, 16]> -> vector<256x128xbf6> ``` -Users store a matrix to share local memory from vector. +Users can store a matrix from a vector value into shared local memory using the store_matrix operation. ```mlir store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @block=[8, 16]>, vector<256x128xbf6> ``` **Cooperative Transpose Example** -Suppose we have wg-level user input code +This example demonstrates a cooperative transpose pattern in which a matrix tile is loaded by a workgroup and collaboratively transposed across subgroups or threads. The operation is broken into two steps: a local transpose using vector.transpose and a cooperative re-layout using xegpu.convert_layout, where neighboring subgroups within a workgroup exchange data to form the desired transposed tile layout. ```mlir #Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1] } #Coop_wg = {sg_layout = [8, 4] , sg_data= [32, 8], order=[1, 0] } @@ -391,19 +394,42 @@ Suppose we have wg-level user input code %a = vector.transpose %1 #Coop_wg :vector<32x256xf16> -> vector<256x32xf16> %a_dpas = Conv_layout %2 #Coop_wg #dpas_wg ``` +In this flow: -After an optimization pass which optimize the transpose-A pattern, the transformed code uses store_matrix and load_matrix. Note the load_nd and store_matrix has smaller sg_data so the subgroups perform cooperative transpose. -```mlir -#Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1 } -#dpas_t_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } +vector.transpose applies a local transpose within each thread’s register tile. + +xegpu.convert_layout performs a cooperative data exchange among threads/subgroups to assemble a larger tile in the transposed layout. + +The result is a matrix tile conforming to the #dpas_wg layout, ready for compute instructions such as DPAS. + +After an optimization pass that targets the transpose-A pattern, the code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. + +It is generally preferred to fuse transpose and convert_layout earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. + +```mlir +#Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], order = [0, 1] } // original layout +#dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], order = [1, 0] } // target DPAS layout + +%at = xegpu.load_nd %tdesc + : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> + +%m = xegpu.create_matrix_desc + : matrix_desc<32x256xf16> + +%mt = xegpu.matrix_desc_subview %m[0, 0] + : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> + +xegpu.store_matrix %mt[0, 0], %at + : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> + +xegpu.barrier + +%ma = xegpu.matrix_desc_subview %m[0, 0] + : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> + +%a_dpas = xegpu.load_matrix %ma[0, 0] + : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> -%at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > -%m = create_matrix_desc : matrix_desc<32x256xf16> -%mt = matrix_desc_subview %m: matrix_desc<32x256xf16, strides=[1, 32], #coop_t_wg> -store_matrix %mt[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, strides=[1, 32], #coop_t_wg> -barrier -%ma = matrix_desc_subview : matrix_desc<256x32xf16, #dpas_t_wg> -%a_dpas = load_matrix %ma [0, 0] #dpas_t_wg : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` After wg->sg level distribution, this lowers to the following sg-level code. From 3657d362c46109eebdcc9fdc2983d22a101e0d62 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Thu, 10 Jul 2025 23:06:13 -0700 Subject: [PATCH 11/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 360 ++++++++++++++++++++++++++++++++++----------- 1 file changed, 273 insertions(+), 87 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 5f94858b9..b9abf89e3 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -369,8 +369,11 @@ Users create a `matrix_desc` to represent a matrix stored in shared local memory ``` Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. The resulting matrix_desc may be annotated with layout attributes such as @block and @strides to describe its memory layout more precisely. The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. Additionally, a subview may carry an xegpu.layout attribute that defines how the matrix is logically partitioned into fragments and mapped to work items. ```mlir -%mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0]: matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> -%mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster*64]: matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> +%mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0] + : matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> + +%mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster * 64] + : matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> ``` Users can load a matrix from shared local memory into a vector value using the load_matrix operation. The result is a vector type in the IR, representing a tile stored in registers. @@ -384,6 +387,7 @@ store_matrix matrix_desc_b, vec_a :matrix_desc<256x128xbf16, @block=[8, 16]>, ve ``` **Cooperative Transpose Example** + This example demonstrates a cooperative transpose pattern in which a matrix tile is loaded by a workgroup and collaboratively transposed across subgroups or threads. The operation is broken into two steps: a local transpose using vector.transpose and a cooperative re-layout using xegpu.convert_layout, where neighboring subgroups within a workgroup exchange data to form the desired transposed tile layout. ```mlir #Coop_t_wg ={sg_layout = [4, 8], sg_data= [8, 32], order=[0, 1] } @@ -396,13 +400,14 @@ This example demonstrates a cooperative transpose pattern in which a matrix tile ``` In this flow: -vector.transpose applies a local transpose within each thread’s register tile. +1. vector.transpose applies a local transpose within each thread’s register tile. -xegpu.convert_layout performs a cooperative data exchange among threads/subgroups to assemble a larger tile in the transposed layout. +2. xegpu.convert_layout performs a cooperative data exchange among threads/subgroups to assemble a larger tile in the transposed layout. -The result is a matrix tile conforming to the #dpas_wg layout, ready for compute instructions such as DPAS. +3. The result is a matrix tile conforming to the #dpas_wg layout, ready for compute instructions such as DPAS. -After an optimization pass that targets the transpose-A pattern, the code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. +**After an optimization pass that targets the transpose-A pattern** +The code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. It is generally preferred to fuse transpose and convert_layout earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. @@ -432,87 +437,268 @@ xegpu.barrier ``` -After wg->sg level distribution, this lowers to the following sg-level code. -```mlir -#coop_t_inst ={ inst_data=[8, 16] } -#dpas_t_inst = {inst_data=[16, 16] } -create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16, #coop_t_inst> -%at = load_nd %tdesc: tensor_desc<8x32xf16, #coop_t_inst> -> vector<8x32xf16> - -%m = create_matrix_desc : matrix_desc<32x256xf16> -%mt_sg = create_matrix_desc %m [sg_idy*8, sg_idx*32]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32], #coop_t_inst > -store_matrix %mt_sg, %at: vector<8x32xf16>, matrix_desc<8x32xf16, block=[16, 16], strides=[1, 32], #coop_t_inst > -barrier -%ma = matrix_desc_subview %m : matrix_desc<256x32xf16, block=[16, 16], #dpas_t_inst> -%ma_sg = matrix_desc_subview %ma [sg_idy*32, sg_idx*32%32]: matrix_desc<32x32xf16, block=[16, 16] , #dpas_t_inst> -%a_dpas = load_matrix ma_sg: matrix_desc<32x32xf16, block=[16, 16], #dpas_t_inst >-> vector<32x32xf16> - -``` - -After blocking according to inst_data. -```mlir -create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16> -%at = load_nd %tdesc, sg_coords1: tensor_desc<8x32xf16> -> vector<8x32xf16> -%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> -%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> -%m = create_matrix_desc : matrix_desc<32x256xf16> -%mt_sg = create_matrix_desc %m1 [0, 0]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> - -%mt_inst0 = create_matrix_desc % mt_sg [sg_idy*8, sg_idx*32]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -> matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> -%mt_inst1 = create_matrix_desc % mt_sg [sg_idy*8, sg_idx*32+16]: matrix_desc<32x256xf16, block=[16, 16], strides=[1, 32]> -> matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> - -store_matrix %mt_inst0, %at0: vector<8x16xf16>, matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> -store_matrix %mt_inst1, %at1: vector<8x16xf16>, matrix_desc<8x16xf16, block=[16, 16], strides=[1, 32]> -barrier -%ma = matrix_desc_subview %m: matrix_desc<256x32xf16, block=[16, 16] > -%ma_inst0 = matrix_desc_subview % ma [sg_idy*32, sg_idx*32%32]: matrix_desc<16x16xf16, block=[16, 16] > -% ma_inst1 = matrix_desc_subview %ma [sg_idy*32, sg_idx*32%32+16]: matrix_desc<16x16xf16, block=[16, 16] > -% ma_inst2 = matrix_desc_subview %ma [sg_idy*32+16, sg_idx*32%32]: matrix_desc<16x16xf16, block=[16, 16] > -% ma_inst3 = matrix_desc_subview %ma [sg_idy*32+16, sg_idx*32%32+16]: matrix_desc<16x16xf16, block=[16, 16] > -%a_dpas_0 = load_matrix ma_inst0: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> -%a_dpas_1 = load_matrix ma_inst1: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> -%a_dpas_2 = load_matrix ma_inst2: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> -%a_dpas_3 = load_matrix ma_inst3: matrix_desc<16x16xf16, block=[16, 16]>, vector<16x16xf16> -``` - -MaterializeSLMAccess pass replace matrix_desc and related ops to memref and load and load_1d. Pseudo code used to simplify the code. -```mlir -create_nd_tdesc %tdesc_sg [widy*32+sg_idy*8, widx*256+sg_idx*32] : : memref<4096x4096xf16> -> : tensor_desc<8x32xf16> -%at = load_nd %tdesc, sg_coords1: tensor_desc<8x32xf16> -> vector<8x32xf16> -%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> -%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> - -%blk_y=sg_idy*8 /16: index -%blk_in_y=sg_idy*8 %16: index -%sg_idx_vec = %sg_idx*32 + [0..15] : vector<16xindex> -%blk_x=%sg_idx_vec /16: vector<16xindex > -%blk_in_x=%sg_idx_vec %16: vector<16xindex > -%sg_start_offset_vec = %blk_y * 16 + %blk_in_y + %blk_x * 512 + %blk_in_x*16 -%tdesc0 = xegpu.create_tdesc %m, %sg_start_offset_vec: memref<8192xf16, 3>, %sg_start_offset_vec ->tdesc<8x16xf16, chunk=8, scope=slm> - -%sg_idx_vec2 = %sg_idx*32 + [16..31] : vector<16xindex> -%blk_x2=%sg_idx_vec /16: vector<16xindex > -%blk_in_x2=%sg_idx_vec %16: vector<16xindex > -%sg_start_offset_vec2 = %blk_y * 16 + %blk_in_y + %blk_x * 512 + %blk_in_x*16 -%tdesc1 = xegpu.create_tdesc %m, %sg_start_offset_vec2: memref<8192xf16, 3>, %sg_start_offset_vec ->tdesc<8x16xf16, chunk=8, scope=slm> - -xegpu.store %tdesc0, %at0: tdesc<8x32xf16, chunk=8, scope=slm>, vector<8x16xf16> -xegpu.store %tdesc1, %at1: tdesc<8x32xf16, chunk=8, scope=slm>, vector<8x16xf16> - -barrier -%inst_start_offset0 = sg_idy*2* 512 -%tdesc0 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > -%inst_start_offset0 = sg_idy*2* 512 + 256 -%tdesc1 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > -%inst_start_offset0 = sg_idy*2* 512 + 512 -%tdesc2 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > -%inst_start_offset0 = sg_idy*2* 512 + 512 + 256 -%tdesc3 = xegpu.create_nd_tdesc %m1, % inst_start_offset0 : memref<8192xf16, 3>, index->tdesc<256xf16 > - -a_dpas_0 = Load_nd %tdesc0: tdesc<256xf16 > -> vector<256xf16> -a_dpas_1 = Load_nd %tdesc1: tdesc<256xf16 > -> vector<256xf16> -a_dpas_2 = Load_nd %tdesc2: tdesc<256xf16 > -> vector<256xf16> -a_dpas_3 = Load_nd %tdesc3: tdesc<256xf16 > -> vector<256xf16> +**Adding attributes for Instruction-Level Blocking** +In this example, the matrix is distributed according to hardware capability, using instruction-level blocking. This type of blocking does not change the physical memory layout (i.e., there is no memory-level tiling); instead, it affects how data is accessed and lowered to instructions like store_scatter and load_gather. + +Each lane handles 2 f16 elements (32 bits), and the matrix_desc uses a @strides attribute to represent memory layout. During lowering, offsets must be computed by composing: 1) The logical layout (e.g., xegpu.layout), and 2) The memory layout (e.g., @strides), to produce final physical offsets. +```mlir +#Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [1, 32], order = [0, 1] } +#dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [1, 32], order = [1, 0] } +// Load a tile cooperatively using logical layout #Coop_t_wg +%at = xegpu.load_nd %tdesc + : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> + +// Allocate a matrix_desc over shared memory +%m = xegpu.create_matrix_desc + : matrix_desc<32x256xf16> + +// Subview with strides, using the same layout for consistent offset mapping +%mt = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> + +// Store matrix with instruction-level blocking (per-inst offsets computed using strides) +xegpu.store_matrix %mt[0, 0], %at + : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> + +// Synchronize before loading the transposed result +xegpu.barrier + +// Subview for DPAS layout (transposed, larger tile per subgroup) +%ma = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> + +// Load matrix cooperatively with new layout +%a_dpas = xegpu.load_matrix %ma[0, 0] + : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> +``` +**Optimized with Blocking: Lowering to store_chunk and 1D Block Load** +This pattern demonstrates a more optimized strategy using instruction-level blocking, enabling the use of efficient memory instructions such as store_chunk and 1D block load. For correct and efficient lowering, several constraints must be satisfied: + +The inst_data field must specify a meaningful 2D shape that aligns with the capabilities of store_chunk and 1D block load. + +Blocking must be explicitly expressed in the memory layout via the @block attribute. Two related matrix_desc subviews (e.g., producer and consumer) must have consistent block sizes. In some cases, the block shape may be transposed between the two to accommodate different access orders. + +Each instruction must access only within its assigned matrix block boundary — no cross-block accesses are allowed. + +During lowering, store_matrix is lowered to store_chunk if the matrix has strides, and load_matrix is lowered to 1D block load if the matrix has a blocked layout. + +```mlir +#Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [8, 16], order = [0, 1] } +#dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [16, 16], order = [1, 0] } +// Load matrix cooperatively from global memory +%at = xegpu.load_nd %tdesc + : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> + +// Allocate shared local memory buffer +%m = xegpu.create_matrix_desc + : matrix_desc<32x256xf16> + +// Subview with both block and stride attributes +%mt = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> + +// Store cooperatively into SLM using blocking-aware lowering (store_chunk) +xegpu.store_matrix %mt[0, 0], %at + : vector<32x256xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> + +// Synchronize before reuse +xegpu.barrier + +// Subview with matching block shape for 1D block load +%ma = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> + +// Load cooperatively from SLM using 1D block load +%a_dpas = xegpu.load_matrix %ma[0, 0] + : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -> vector<256x32xf16> +``` + +**Workgroup to Subgroup Distribution** +This example illustrates how data is distributed from workgroup to subgroups. It demonstrates how load_matrix and store_matrix cooperate with matrix_desc subviews to enable efficient subgroup distribution. In this step, the sg_layout and sg_data attributes are removed from the layout specification, leaving only the inst_data attribute. + +The matrix is assumed to be stored in row-major contiguous layout, and indexing into it is performed using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the load_matrix/store_matrix abstraction is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. + +```mlir +#coop_t_inst = { inst_data = [8, 16] } +#dpas_t_inst = { inst_data = [16, 16] } + +// Each subgroup loads its portion of the global matrix using inst_data layout +%tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] + : memref<4096x4096xf16> -> tensor_desc<8x32xf16, #coop_t_inst> + +%at = xegpu.load_nd %tdesc_sg + : tensor_desc<8x32xf16, #coop_t_inst> -> vector<8x32xf16> + +// Allocate a matrix in SLM with row-major layout +%m = xegpu.create_matrix_desc + : matrix_desc<32x256xf16> + +// Subgroup subview: logical coords computed assuming row-major layout +%mt_sg = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32] + : matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> + -> matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32]> + +// Store vector into SLM using per-inst tile and logical mapping +xegpu.store_matrix %mt_sg, %at + : vector<8x32xf16>, matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> + +// Barrier to synchronize SLM access +xegpu.barrier + +// Subview for DPAS tile shape +%ma = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> + +// Subgroup-level subview for cooperative DPAS load +%ma_sg = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * (32 % 32)] + : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> + -> matrix_desc<32x32xf16, @block=[16, 16]> + +// Load matrix cooperatively from SLM using 1D block load +%a_dpas = xegpu.load_matrix %ma_sg + : matrix_desc<32x32xf16, @block=[16, 16], #dpas_t_inst> -> vector<32x32xf16> +``` +**Unrolling Guided by Inst_data** +This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This pattern ensures that each store operation writes within its assigned block boundary, respecting the @block and @strides attributes. + +On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. + +```mlir +// Load global matrix fragment cooperatively +%tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] + : memref<4096x4096xf16> -> tensor_desc<8x32xf16> + +%at = xegpu.load_nd %tdesc_sg + : tensor_desc<8x32xf16> -> vector<8x32xf16> + +// Extract 16-column instruction tiles +%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> +%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> + +// Create shared memory backing matrix +%m = xegpu.create_matrix_desc : matrix_desc<32x256xf16> + +// Compute instruction-tile-level subviews +%mt_inst0 = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32] + : matrix_desc<32x256xf16> -> matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> + +%mt_inst1 = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32 + 16] + : matrix_desc<32x256xf16> -> matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> + +// Store unrolled tiles into shared local memory +xegpu.store_matrix %mt_inst0, %at0 + : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> + +xegpu.store_matrix %mt_inst1, %at1 + : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> + +// Synchronize to ensure SLM is ready +xegpu.barrier + +// Create 32×32 transposed matrix view for DPAS-style consumption +%ma = xegpu.matrix_desc_subview %m + : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16]> + +// Compute 16×16 instruction tiles for DPAS load +%ma_inst0 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32] + : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> + +%ma_inst1 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32 + 16] + : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> + +%ma_inst2 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32] + : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> + +%ma_inst3 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32 + 16] + : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> + +// Load unrolled tiles for compute +%a_dpas_0 = xegpu.load_matrix %ma_inst0 + : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> + +%a_dpas_1 = xegpu.load_matrix %ma_inst1 + : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> + +%a_dpas_2 = xegpu.load_matrix %ma_inst2 + : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> + +%a_dpas_3 = xegpu.load_matrix %ma_inst3 + : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> +``` + +**MaterializeSLMAccess: Lowering matrix_desc to Physical Memory Access** +This step lowers high-level matrix_desc abstractions and cooperative memory operations (store_matrix, load_matrix) into explicit memory operations (store_chunk, load_1d) over shared local memory (memref). It performs full address materialization using the matrix's layout attributes (@strides, @block) and logical lane coordinates. + +Key Concepts: +Matrix-to-Memory Conversion: Replace matrix_desc-based tile abstractions with raw memref and compute physical offsets explicitly. + +Chunked Store: Each thread stores a small fragment (e.g., 8×16) using logical offsets composed with layout metadata. Lowered to store_chunk. + +1D Block Load: A transposed layout (e.g., 256×32) is blocked as 16×16 tiles. Contiguous blocks are loaded using load_1d, which requires computing the physical offset of the first element per tile. + +Offset Calculation: Logical per-lane coordinates are transformed into logical block coordinates, then to physical offsets using block size and strides. + +```mlir +// Load global input tile into vector +%tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] + : memref<4096x4096xf16> -> tensor_desc<8x32xf16> +%at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16> -> vector<8x32xf16> + +// Unroll 8x32 into two 8x16 tiles +%at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> +%at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> + +// Shared local memory buffer +%m = memref.alloc : memref<8192xf16, 3> + +// Compute blocked offset vectors for SLM store +%blk_y = divi_signed %sg_idy * 8, 16 : index +%blk_in_y = remi_signed %sg_idy * 8, 16 : index +%sg_idx_vec = addi %sg_idx * 32, dense<[0, ..., 15]> : vector<16xindex> +%blk_x = divi_unsigned %sg_idx_vec, 16 +%blk_in_x = remi_unsigned %sg_idx_vec, 16 +%offset_vec0 = addi (addi (addi (%blk_y * 16, %blk_in_y), %blk_x * 512), %blk_in_x * 16) + +// Create tensor_desc for SLM store +%tdesc0 = xegpu.create_tdesc %m, %offset_vec0 + : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> + +// Repeat for second tile +%sg_idx_vec2 = addi %sg_idx * 32, dense<[16, ..., 31]> : vector<16xindex> +%blk_x2 = divi_unsigned %sg_idx_vec2, 16 +%blk_in_x2 = remi_unsigned %sg_idx_vec2, 16 +%offset_vec1 = addi (addi (addi (%blk_y * 16, %blk_in_y), %blk_x2 * 512), %blk_in_x2 * 16) +%tdesc1 = xegpu.create_tdesc %m, %offset_vec1 + : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> + +// Transpose and store +%at0_t = vector.transpose %at0 : vector<8x16xf16> -> vector<16x8xf16> +%at1_t = vector.transpose %at1 : vector<8x16xf16> -> vector<16x8xf16> +xegpu.store %tdesc0, %at0_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> +xegpu.store %tdesc1, %at1_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> + +// Barrier to ensure SLM visibility +xegpu.barrier + +// ---------------------- Load 1D Block ---------------------- + +// Compute per-tile physical offsets +%inst_start_offset0 = mul %sg_idy, 2 * 512 +%inst_start_offset1 = add %inst_start_offset0, 256 +%inst_start_offset2 = add %inst_start_offset0, 512 +%inst_start_offset3 = add %inst_start_offset0, 768 + +// Create tdesc for 1D block loads +%tdesc0 = xegpu.create_nd_tdesc %m, %inst_start_offset0 : memref<8192xf16, 3>, index -> tdesc<256xf16> +%tdesc1 = xegpu.create_nd_tdesc %m, %inst_start_offset1 : memref<8192xf16, 3>, index -> tdesc<256xf16> +%tdesc2 = xegpu.create_nd_tdesc %m, %inst_start_offset2 : memref<8192xf16, 3>, index -> tdesc<256xf16> +%tdesc3 = xegpu.create_nd_tdesc %m, %inst_start_offset3 : memref<8192xf16, 3>, index -> tdesc<256xf16> + +// Load 1D tiles +%a_dpas_0 = xegpu.load_nd %tdesc0 : tdesc<256xf16> -> vector<256xf16> +%a_dpas_1 = xegpu.load_nd %tdesc1 : tdesc<256xf16> -> vector<256xf16> +%a_dpas_2 = xegpu.load_nd %tdesc2 : tdesc<256xf16> -> vector<256xf16> +%a_dpas_3 = xegpu.load_nd %tdesc3 : tdesc<256xf16> -> vector<256xf16> ``` ## XeGPU Attributes to support Work Item Level semantics From 3b79e51bbeae8e84b9ee01642ff3359fc197d28b Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Fri, 11 Jul 2025 17:43:42 -0700 Subject: [PATCH 12/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 203 ++++++++++++++------------------------------- 1 file changed, 62 insertions(+), 141 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index b9abf89e3..7de5b2dba 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -351,14 +351,15 @@ This separation simplifies distribution and unrolling passes and enables systema To represent a matrix stored in shared local memory (SLM), users must create a matrix_desc object. The underlying memory is assumed to follow a row-major layout, and the base matrix_desc represents a raw, unannotated matrix in this layout. The base matrix may be n-dimensional. -Only matrix_desc instances created via subview may carry an xegpu.layout attribute, which specifies the mapping of lanes and registers to fragments of the matrix. This attribute guides the tile distribution process based on the assumed row-major view of the original matrix. In addition, subviewed matrix_desc instances may carry layout metadata such as blocking and striding, which are used to control physical address computation when accessing SLM. +matrix_desc_view creates a matrix_desc instance with memory layout attributes such as @block and @stride. These attributes define the blocking and striding parameters, which govern physical address computation when accessing shared local memory (SLM). Additionally, an xegpu.layout attribute is added to specify the mapping of lanes and registers to fragments of the matrix, guiding tile distribution based on the assumed row-major view of the matrix. The matrix_desc_subview creates a subview on top of the matrix_desc produced by matrix_desc_view, inheriting all of its layout attributes. The subview is then subject to decomposition and distribution. Data movement between SLM and vector registers is performed using load_matrix and store_matrix, which operate at workgroup scope and require the input matrix_desc to be 2D. If the original matrix is higher-dimensional, it must be subviewed to a 2D shape before it can be used with these operations. | Ops | Syntax | Example | | :--- | :---- | :--- | |create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16> | -|matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @layout_type=1> -> matrix_desc<128x128xbf16, @stride=[256,1], @block=[8, 16]> | +|matrix_desc_view | operation ::= xegpu.matrix_desc_view \$mdesc attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_a_layout = xegpu.matrix_desc_view %mdesc:matrix_desc<256x128xbf16> -> matrix_desc<256x128xbf16, @stride=[1, 256], @block=[8, 16]> | +|matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @stride=[256,1], @block=[8, 16]> -> matrix_desc<128x128xbf16, @stride=[256,1], @block=[8, 16]> | |load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @block=[8, 16]> -> vector<128x256xbf16> | |store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @block=[8, 16]>, vector<128x256xbf16> | @@ -367,10 +368,16 @@ Users create a `matrix_desc` to represent a matrix stored in shared local memory ```mlir %mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> ``` -Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. The resulting matrix_desc may be annotated with layout attributes such as @block and @strides to describe its memory layout more precisely. The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. Additionally, a subview may carry an xegpu.layout attribute that defines how the matrix is logically partitioned into fragments and mapped to work items. +matrix_desc_view annotates matrix_desc with layout attributes such as @block and @strides to describe its memory layout more precisely. The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. + +```mlir +%mdesc_a_layout = xegpu.matrix_desc_view %mdesc:matrix_desc<256x128xbf16> -> matrix_desc<256x128xbf16, @stride=[1, 256], @block=[8, 16]> +``` +Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. Subview inherit memory layout attributes from the base matrix_desc. Additionally, a view may carry an xegpu.layout attribute that defines how the matrix is logically partitioned into fragments and mapped to work items. + ```mlir %mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0] - : matrix_desc<3x256x128xbf16> -> matrix_desc<256x128xbf16, @block=[8, 16]> + : matrix_desc<3x256x128xbf16, @block=[8, 16]> -> matrix_desc<256x128xbf16, @block=[8, 16]> %mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster * 64] : matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> @@ -394,7 +401,7 @@ This example demonstrates a cooperative transpose pattern in which a matrix tile #Coop_wg = {sg_layout = [8, 4] , sg_data= [32, 8], order=[1, 0] } #dpas_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } -%at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16 > +%at = load_nd %tdesc: tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16 > %a = vector.transpose %1 #Coop_wg :vector<32x256xf16> -> vector<256x32xf16> %a_dpas = Conv_layout %2 #Coop_wg #dpas_wg ``` @@ -406,7 +413,7 @@ In this flow: 3. The result is a matrix tile conforming to the #dpas_wg layout, ready for compute instructions such as DPAS. -**After an optimization pass that targets the transpose-A pattern** +**After optimization that targets the transpose-A pattern** The code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. It is generally preferred to fuse transpose and convert_layout earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. @@ -415,68 +422,38 @@ It is generally preferred to fuse transpose and convert_layout earlier in the pi #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], order = [0, 1] } // original layout #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], order = [1, 0] } // target DPAS layout -%at = xegpu.load_nd %tdesc - : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> - -%m = xegpu.create_matrix_desc - : matrix_desc<32x256xf16> - -%mt = xegpu.matrix_desc_subview %m[0, 0] - : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> - -xegpu.store_matrix %mt[0, 0], %at - : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> - -xegpu.barrier - -%ma = xegpu.matrix_desc_subview %m[0, 0] - : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> - -%a_dpas = xegpu.load_matrix %ma[0, 0] - : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> - +%at = xegpu.load_nd %tdesc : tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> +%mt = xegpu.matrix_desc_view %m : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +xegpu.store_matrix %mt[0, 0], %at : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +gpu.barrier +%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> +%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` -**Adding attributes for Instruction-Level Blocking** -In this example, the matrix is distributed according to hardware capability, using instruction-level blocking. This type of blocking does not change the physical memory layout (i.e., there is no memory-level tiling); instead, it affects how data is accessed and lowered to instructions like store_scatter and load_gather. +**Adding attributes for Instruction-Level Blocking: basic blocking** +In this example, the xegpu.layout is extended to support instruction-level blocking. The basic blocking assumes 16 lanes, and each lane handles 2 f16 elements (32 bits). This basic blocking does not change the physical memory layout (i.e., there is no memory-level tiling); instead, it loweres to instructions like store_scatter and load_gather. -Each lane handles 2 f16 elements (32 bits), and the matrix_desc uses a @strides attribute to represent memory layout. During lowering, offsets must be computed by composing: 1) The logical layout (e.g., xegpu.layout), and 2) The memory layout (e.g., @strides), to produce final physical offsets. ```mlir #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [1, 32], order = [0, 1] } #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [1, 32], order = [1, 0] } -// Load a tile cooperatively using logical layout #Coop_t_wg -%at = xegpu.load_nd %tdesc - : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> -// Allocate a matrix_desc over shared memory -%m = xegpu.create_matrix_desc - : matrix_desc<32x256xf16> - -// Subview with strides, using the same layout for consistent offset mapping -%mt = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> - -// Store matrix with instruction-level blocking (per-inst offsets computed using strides) -xegpu.store_matrix %mt[0, 0], %at - : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +%at = xegpu.load_nd %tdesc: tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = xegpu.create_matrix_desc: matrix_desc<8192xf16> +%mt = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +xegpu.store_matrix %mt[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> -// Synchronize before loading the transposed result -xegpu.barrier +gpu.barrier -// Subview for DPAS layout (transposed, larger tile per subgroup) -%ma = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> - -// Load matrix cooperatively with new layout -%a_dpas = xegpu.load_matrix %ma[0, 0] - : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> +%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> +%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` **Optimized with Blocking: Lowering to store_chunk and 1D Block Load** -This pattern demonstrates a more optimized strategy using instruction-level blocking, enabling the use of efficient memory instructions such as store_chunk and 1D block load. For correct and efficient lowering, several constraints must be satisfied: +This pattern demonstrates a more optimized strategy for instruction-level blocking, enabling the use of efficient memory instructions such as store_chunk and 1D block load. For correct and efficient lowering, several constraints must be satisfied: The inst_data field must specify a meaningful 2D shape that aligns with the capabilities of store_chunk and 1D block load. -Blocking must be explicitly expressed in the memory layout via the @block attribute. Two related matrix_desc subviews (e.g., producer and consumer) must have consistent block sizes. In some cases, the block shape may be transposed between the two to accommodate different access orders. +Blocking must be explicitly expressed in the memory layout via the @block attribute. Two related matrix_desc subviews (e.g., producer and consumer) must have consistent block sizes. If one matrix_desc is transposed, the block shape should match the transposed shape of the other one. Each instruction must access only within its assigned matrix block boundary — no cross-block accesses are allowed. @@ -485,38 +462,21 @@ During lowering, store_matrix is lowered to store_chunk if the matrix has stride ```mlir #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [8, 16], order = [0, 1] } #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [16, 16], order = [1, 0] } -// Load matrix cooperatively from global memory -%at = xegpu.load_nd %tdesc - : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> -// Allocate shared local memory buffer -%m = xegpu.create_matrix_desc - : matrix_desc<32x256xf16> - -// Subview with both block and stride attributes -%mt = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> - -// Store cooperatively into SLM using blocking-aware lowering (store_chunk) -xegpu.store_matrix %mt[0, 0], %at - : vector<32x256xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> - -// Synchronize before reuse -xegpu.barrier - -// Subview with matching block shape for 1D block load -%ma = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> +%at = xegpu.load_nd %tdesc : tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> +%mt = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> +xegpu.store_matrix %mt[0, 0], %at : vector<32x256xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> -// Load cooperatively from SLM using 1D block load -%a_dpas = xegpu.load_matrix %ma[0, 0] - : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -> vector<256x32xf16> +gpu.barrier +%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> +%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -> vector<256x32xf16> ``` **Workgroup to Subgroup Distribution** This example illustrates how data is distributed from workgroup to subgroups. It demonstrates how load_matrix and store_matrix cooperate with matrix_desc subviews to enable efficient subgroup distribution. In this step, the sg_layout and sg_data attributes are removed from the layout specification, leaving only the inst_data attribute. -The matrix is assumed to be stored in row-major contiguous layout, and indexing into it is performed using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the load_matrix/store_matrix abstraction is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. +The matrix is assumed to be stored in row-major contiguous layout, and indexing into it is performed using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the matrix_desc data type is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. ```mlir #coop_t_inst = { inst_data = [8, 16] } @@ -525,126 +485,87 @@ The matrix is assumed to be stored in row-major contiguous layout, and indexing // Each subgroup loads its portion of the global matrix using inst_data layout %tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] : memref<4096x4096xf16> -> tensor_desc<8x32xf16, #coop_t_inst> - %at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16, #coop_t_inst> -> vector<8x32xf16> - -// Allocate a matrix in SLM with row-major layout %m = xegpu.create_matrix_desc : matrix_desc<32x256xf16> - -// Subgroup subview: logical coords computed assuming row-major layout -%mt_sg = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32] +%mt = xegpu.matrix_desc_view %m + : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> +%mt_sg = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32] : matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> - -> matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32]> - -// Store vector into SLM using per-inst tile and logical mapping + -> matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> xegpu.store_matrix %mt_sg, %at : vector<8x32xf16>, matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> -// Barrier to synchronize SLM access -xegpu.barrier +gpu.barrier -// Subview for DPAS tile shape -%ma = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> - -// Subgroup-level subview for cooperative DPAS load +%ma = xegpu.matrix_desc_view %m + : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> %ma_sg = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * (32 % 32)] : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> -> matrix_desc<32x32xf16, @block=[16, 16]> - -// Load matrix cooperatively from SLM using 1D block load %a_dpas = xegpu.load_matrix %ma_sg : matrix_desc<32x32xf16, @block=[16, 16], #dpas_t_inst> -> vector<32x32xf16> ``` **Unrolling Guided by Inst_data** -This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This pattern ensures that each store operation writes within its assigned block boundary, respecting the @block and @strides attributes. - -On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. +This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This pattern ensures that each load and store operation writes within its assigned block boundary, respecting the @block and @strides attributes. On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. ```mlir -// Load global matrix fragment cooperatively %tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] : memref<4096x4096xf16> -> tensor_desc<8x32xf16> - -%at = xegpu.load_nd %tdesc_sg - : tensor_desc<8x32xf16> -> vector<8x32xf16> - -// Extract 16-column instruction tiles +%at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16> -> vector<8x32xf16> %at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> %at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> - -// Create shared memory backing matrix -%m = xegpu.create_matrix_desc : matrix_desc<32x256xf16> - -// Compute instruction-tile-level subviews -%mt_inst0 = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32] - : matrix_desc<32x256xf16> -> matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> - -%mt_inst1 = xegpu.matrix_desc_subview %m[%sg_idy * 8, %sg_idx * 32 + 16] - : matrix_desc<32x256xf16> -> matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> - -// Store unrolled tiles into shared local memory +%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> +%mt = xegpu.matrix_desc_view %m + : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> +%mt_inst0 = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32] + : matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> +%mt_inst1 = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32 + 16] + : matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> xegpu.store_matrix %mt_inst0, %at0 : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> - xegpu.store_matrix %mt_inst1, %at1 : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> -// Synchronize to ensure SLM is ready -xegpu.barrier +gpu.barrier -// Create 32×32 transposed matrix view for DPAS-style consumption %ma = xegpu.matrix_desc_subview %m - : matrix_desc<32x256xf16> -> matrix_desc<256x32xf16, @block=[16, 16]> - -// Compute 16×16 instruction tiles for DPAS load + : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16]> %ma_inst0 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32] : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> - %ma_inst1 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32 + 16] : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> - %ma_inst2 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32] : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> - %ma_inst3 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32 + 16] : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> - -// Load unrolled tiles for compute %a_dpas_0 = xegpu.load_matrix %ma_inst0 : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> - %a_dpas_1 = xegpu.load_matrix %ma_inst1 : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> - %a_dpas_2 = xegpu.load_matrix %ma_inst2 : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> - %a_dpas_3 = xegpu.load_matrix %ma_inst3 : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> ``` **MaterializeSLMAccess: Lowering matrix_desc to Physical Memory Access** -This step lowers high-level matrix_desc abstractions and cooperative memory operations (store_matrix, load_matrix) into explicit memory operations (store_chunk, load_1d) over shared local memory (memref). It performs full address materialization using the matrix's layout attributes (@strides, @block) and logical lane coordinates. +This step lowers high-level matrix_desc operations (store_matrix, load_matrix) into low-level memory operations (store_chunk, load_1d) over shared local memory. It performs full address materialization using the matrix's layout attributes (@strides, @block) and logical lane coordinates. Key Concepts: Matrix-to-Memory Conversion: Replace matrix_desc-based tile abstractions with raw memref and compute physical offsets explicitly. -Chunked Store: Each thread stores a small fragment (e.g., 8×16) using logical offsets composed with layout metadata. Lowered to store_chunk. +Chunked Store: Each thread stores a small fragment (e.g., 8×1) using the logical offset composed with layout metadata. Lowered to store_chunk. -1D Block Load: A transposed layout (e.g., 256×32) is blocked as 16×16 tiles. Contiguous blocks are loaded using load_1d, which requires computing the physical offset of the first element per tile. +1D Block Load: A transposed layout (e.g., 256×32) is blocked as 16×16 tiles. Contiguous blocks are loaded using load_1d, which requires computing the physical offset of the first element per 1D block. Offset Calculation: Logical per-lane coordinates are transformed into logical block coordinates, then to physical offsets using block size and strides. ```mlir -// Load global input tile into vector %tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] : memref<4096x4096xf16> -> tensor_desc<8x32xf16> -%at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16> -> vector<8x32xf16> - -// Unroll 8x32 into two 8x16 tiles +%at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16> -> vector<8x32xf16> %at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> %at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> @@ -678,7 +599,7 @@ xegpu.store %tdesc0, %at0_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf xegpu.store %tdesc1, %at1_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> // Barrier to ensure SLM visibility -xegpu.barrier +gpu.barrier // ---------------------- Load 1D Block ---------------------- From 5dd778e4f28b456fd32c7b25abd44e16b7dde64d Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Fri, 11 Jul 2025 18:49:08 -0700 Subject: [PATCH 13/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 34 +++++++++++++++++++++------------- 1 file changed, 21 insertions(+), 13 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 7de5b2dba..9d34da7d3 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -414,6 +414,7 @@ In this flow: 3. The result is a matrix tile conforming to the #dpas_wg layout, ready for compute instructions such as DPAS. **After optimization that targets the transpose-A pattern** + The code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. It is generally preferred to fuse transpose and convert_layout earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. @@ -431,7 +432,8 @@ gpu.barrier %a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` -**Adding attributes for Instruction-Level Blocking: basic blocking** +**Adding Attributes for Basic Instruction-Level Blocking** + In this example, the xegpu.layout is extended to support instruction-level blocking. The basic blocking assumes 16 lanes, and each lane handles 2 f16 elements (32 bits). This basic blocking does not change the physical memory layout (i.e., there is no memory-level tiling); instead, it loweres to instructions like store_scatter and load_gather. ```mlir @@ -448,7 +450,8 @@ gpu.barrier %ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> %a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> ``` -**Optimized with Blocking: Lowering to store_chunk and 1D Block Load** +**Optimized Blocking: Lowering to store_chunk and 1D Block Load** + This pattern demonstrates a more optimized strategy for instruction-level blocking, enabling the use of efficient memory instructions such as store_chunk and 1D block load. For correct and efficient lowering, several constraints must be satisfied: The inst_data field must specify a meaningful 2D shape that aligns with the capabilities of store_chunk and 1D block load. @@ -474,6 +477,7 @@ gpu.barrier ``` **Workgroup to Subgroup Distribution** + This example illustrates how data is distributed from workgroup to subgroups. It demonstrates how load_matrix and store_matrix cooperate with matrix_desc subviews to enable efficient subgroup distribution. In this step, the sg_layout and sg_data attributes are removed from the layout specification, leaving only the inst_data attribute. The matrix is assumed to be stored in row-major contiguous layout, and indexing into it is performed using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the matrix_desc data type is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. @@ -508,6 +512,7 @@ gpu.barrier : matrix_desc<32x32xf16, @block=[16, 16], #dpas_t_inst> -> vector<32x32xf16> ``` **Unrolling Guided by Inst_data** + This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This pattern ensures that each load and store operation writes within its assigned block boundary, respecting the @block and @strides attributes. On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. ```mlir @@ -530,7 +535,7 @@ xegpu.store_matrix %mt_inst1, %at1 gpu.barrier -%ma = xegpu.matrix_desc_subview %m +%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16]> %ma_inst0 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32] : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> @@ -551,6 +556,7 @@ gpu.barrier ``` **MaterializeSLMAccess: Lowering matrix_desc to Physical Memory Access** + This step lowers high-level matrix_desc operations (store_matrix, load_matrix) into low-level memory operations (store_chunk, load_1d) over shared local memory. It performs full address materialization using the matrix's layout attributes (@strides, @block) and logical lane coordinates. Key Concepts: @@ -572,15 +578,16 @@ Offset Calculation: Logical per-lane coordinates are transformed into logical bl // Shared local memory buffer %m = memref.alloc : memref<8192xf16, 3> +// ---------------------- Chunked Load ---------------------- // Compute blocked offset vectors for SLM store %blk_y = divi_signed %sg_idy * 8, 16 : index %blk_in_y = remi_signed %sg_idy * 8, 16 : index %sg_idx_vec = addi %sg_idx * 32, dense<[0, ..., 15]> : vector<16xindex> %blk_x = divi_unsigned %sg_idx_vec, 16 %blk_in_x = remi_unsigned %sg_idx_vec, 16 -%offset_vec0 = addi (addi (addi (%blk_y * 16, %blk_in_y), %blk_x * 512), %blk_in_x * 16) - -// Create tensor_desc for SLM store +// calculate physic addresses with pre-computed strides of the blocked matrix. +// [32x256, strides=1x32] blocked as [2x16x16x16, strides=256x512x1x16] +%offset_vec0 = addi (addi (addi (%blk_in_y, %blk_in_x * 16), %blk_y * 256),%blk_x * 512) %tdesc0 = xegpu.create_tdesc %m, %offset_vec0 : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> @@ -588,34 +595,35 @@ Offset Calculation: Logical per-lane coordinates are transformed into logical bl %sg_idx_vec2 = addi %sg_idx * 32, dense<[16, ..., 31]> : vector<16xindex> %blk_x2 = divi_unsigned %sg_idx_vec2, 16 %blk_in_x2 = remi_unsigned %sg_idx_vec2, 16 -%offset_vec1 = addi (addi (addi (%blk_y * 16, %blk_in_y), %blk_x2 * 512), %blk_in_x2 * 16) +%offset_vec1 = addi (addi (addi (%blk_in_y, %blk_in_x2 * 16), %blk_y * 256),%blk_x2 * 512) %tdesc1 = xegpu.create_tdesc %m, %offset_vec1 : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> -// Transpose and store +// The transpose is added as we remove the transpose attribute out from chunked load/store and expect an explict data transpose. +// it will be no op after lane distribution since each lane owns same data when [8,1] is transpose to [1, 8] %at0_t = vector.transpose %at0 : vector<8x16xf16> -> vector<16x8xf16> %at1_t = vector.transpose %at1 : vector<8x16xf16> -> vector<16x8xf16> xegpu.store %tdesc0, %at0_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> xegpu.store %tdesc1, %at1_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> -// Barrier to ensure SLM visibility gpu.barrier // ---------------------- Load 1D Block ---------------------- - -// Compute per-tile physical offsets +// Compute per-block physical offsets +// pre-computed strides of the blocked matrix: [256x32] blocked as [2x16x16x16, strides=512x256x16x1] +// sg_idx*32 coord to blocked matrix ccord: sg_idx*32%32/16 (0), sg_idx*32%32%16 (0). %32 due matrix shape[1] is 32 +// sg_idy*32 coord to blocked matrix coord: sg_idy*32/16, sg_idy*32%16 (0) +// then map to physical addr using stride [2x16x16x16, strides=512x256x16x1], get sg_idy*32/16 *512 %inst_start_offset0 = mul %sg_idy, 2 * 512 %inst_start_offset1 = add %inst_start_offset0, 256 %inst_start_offset2 = add %inst_start_offset0, 512 %inst_start_offset3 = add %inst_start_offset0, 768 -// Create tdesc for 1D block loads %tdesc0 = xegpu.create_nd_tdesc %m, %inst_start_offset0 : memref<8192xf16, 3>, index -> tdesc<256xf16> %tdesc1 = xegpu.create_nd_tdesc %m, %inst_start_offset1 : memref<8192xf16, 3>, index -> tdesc<256xf16> %tdesc2 = xegpu.create_nd_tdesc %m, %inst_start_offset2 : memref<8192xf16, 3>, index -> tdesc<256xf16> %tdesc3 = xegpu.create_nd_tdesc %m, %inst_start_offset3 : memref<8192xf16, 3>, index -> tdesc<256xf16> -// Load 1D tiles %a_dpas_0 = xegpu.load_nd %tdesc0 : tdesc<256xf16> -> vector<256xf16> %a_dpas_1 = xegpu.load_nd %tdesc1 : tdesc<256xf16> -> vector<256xf16> %a_dpas_2 = xegpu.load_nd %tdesc2 : tdesc<256xf16> -> vector<256xf16> From 63bed6aeca05550bf0373d098a472dfcae70346d Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Mon, 21 Jul 2025 16:01:09 -0700 Subject: [PATCH 14/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 226 ++++++++++++++++++--------------------------- 1 file changed, 92 insertions(+), 134 deletions(-) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 9d34da7d3..0e953058b 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -349,40 +349,32 @@ This separation simplifies distribution and unrolling passes and enables systema **Basic Usage** -To represent a matrix stored in shared local memory (SLM), users must create a matrix_desc object. The underlying memory is assumed to follow a row-major layout, and the base matrix_desc represents a raw, unannotated matrix in this layout. The base matrix may be n-dimensional. - -matrix_desc_view creates a matrix_desc instance with memory layout attributes such as @block and @stride. These attributes define the blocking and striding parameters, which govern physical address computation when accessing shared local memory (SLM). Additionally, an xegpu.layout attribute is added to specify the mapping of lanes and registers to fragments of the matrix, guiding tile distribution based on the assumed row-major view of the matrix. The matrix_desc_subview creates a subview on top of the matrix_desc produced by matrix_desc_view, inheriting all of its layout attributes. The subview is then subject to decomposition and distribution. - -Data movement between SLM and vector registers is performed using load_matrix and store_matrix, which operate at workgroup scope and require the input matrix_desc to be 2D. If the original matrix is higher-dimensional, it must be subviewed to a 2D shape before it can be used with these operations. +To represent a matrix stored in shared local memory (SLM), users must create a matrix_desc object. Create_matrix_desc initializes a matrix_desc instance with memory layout attributes such as @block and @stride. These attributes define the blocking and striding parameters, which govern physical address computation when accessing shared local memory (SLM). The matrix_desc_subview creates a subview on top of the matrix_desc, inheriting all of its layout attributes. Load_matrix and store_matrix performs data movement between SLM and vector registers. xegpu.layout attribute is added to load_matrix and store_matrix to specify the mapping of lanes and registers to fragments of the matrix, guiding tile distribution based on the assumed row-major view of the matrix. | Ops | Syntax | Example | | :--- | :---- | :--- | -|create_matrix_desc | operation ::= xegpu.create_matrix_desc attr-dict : type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc : matrix_desc<256x128xbf16> | -|matrix_desc_view | operation ::= xegpu.matrix_desc_view \$mdesc attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_a_layout = xegpu.matrix_desc_view %mdesc:matrix_desc<256x128xbf16> -> matrix_desc<256x128xbf16, @stride=[1, 256], @block=[8, 16]> | +|create_matrix_desc | operation ::= xegpu.create_matrix_desc $mref attr-dict :type($mref), type(\$mdesc) | %mdesc_a = xegpu.create_matrix_desc %m: memref<65536xi8, 3> -> matrix_desc<256x128xbf16> | |matrix_desc_subview | operation ::= xegpu.matrix_desc_subview \$mdesc, DynamicIndexList<\$coord> attr-dict : type(\$mdesc) -> type(\$mdesc) | %mdesc_coop = xegpu.matrix_desc_subview %mdesc[128, 0]:matrix_desc<256x256xbf16, @stride=[256,1], @block=[8, 16]> -> matrix_desc<128x128xbf16, @stride=[256,1], @block=[8, 16]> | |load_matrix | operation ::= xegpu.load_matrix $mdesc attr-dict : type($mdesc), {type(coords)} -> type($res) | %result = xegpu.load_matrix %mdesc : matrix_desc<128x256xbf16, @block=[8, 16]> -> vector<128x256xbf16> | |store_matrix | operation ::= xegpu.store_matrix $mdesc, $val attr-dict : type($mdesc), {type(coords)}, type($val) | %result = xegpu.store_matrix %mdesc, %val : matrix_desc<128x256xbf16, @block=[8, 16]>, vector<128x256xbf16> | -Users create a `matrix_desc` to represent a matrix stored in shared local memory (SLM). The operation allocates SLM for the matrix, assuming a row-major contiguous layout. +Users create a `matrix_desc` to represent a matrix stored in shared local memory (SLM). The operation takes a memory buffer (1D int8 memref with empty layout) and create a structured representation of the share local memory. The result matrix_desc has proper information including shape, element type, and memory layout attributes (@block and @strides). The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. -```mlir -%mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> -``` -matrix_desc_view annotates matrix_desc with layout attributes such as @block and @strides to describe its memory layout more precisely. The @block attribute indicates that the matrix follows a blocked layout, enabling optimized lowering to 1D block loads. The @strides attribute specifies the logical strides of each dimension and is typically used to support chunked loads. +When there is no input memref operand, it allocates SLM for the matrix, assuming a row-major contiguous layout. ```mlir -%mdesc_a_layout = xegpu.matrix_desc_view %mdesc:matrix_desc<256x128xbf16> -> matrix_desc<256x128xbf16, @stride=[1, 256], @block=[8, 16]> +%mdesc_a = xegpu.create_matrix_desc: matrix_desc<256x128xbf16> +%mdesc_b = xegpu. create_matrix_desc %m : memref<16384xi8, 3>-> matrix_desc<32x256xf16, @strides=[1, 32]> ``` -Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. Subview inherit memory layout attributes from the base matrix_desc. Additionally, a view may carry an xegpu.layout attribute that defines how the matrix is logically partitioned into fragments and mapped to work items. +Users can create a subview of a matrix_desc to represent a sliced or partitioned view of the original matrix. Subviews may reduce the rank of the matrix, allowing users to extract a lower-dimensional matrix from a higher-dimensional one. Subview inherits memory layout attributes from the base matrix_desc. For GEMM use case, matrix operations typically work on 2D matrix_desc. If the original matrix is higher-dimensional, it can be subviewed to a 2D shape before it is used with these operations. ```mlir %mdesc_a = xegpu.matrix_desc_subview %mdescs_a[%mma_cycle_i, 0, 0] : matrix_desc<3x256x128xbf16, @block=[8, 16]> -> matrix_desc<256x128xbf16, @block=[8, 16]> %mdesc_coop_a = xegpu.matrix_desc_subview %mdesc_a[0, %wg_id_x_in_cluster * 64] - : matrix_desc<256x128xbf16> -> matrix_desc<256x64xbf16, @strides=[128, 1]> + : matrix_desc<256x128xbf16, @strides=[128, 1]> -> matrix_desc<256x64xbf16, @strides=[128, 1]> ``` - Users can load a matrix from shared local memory into a vector value using the load_matrix operation. The result is a vector type in the IR, representing a tile stored in registers. ```mlir vec_a = load_matrix matrix_desc_a: matrix_desc<256x128xbf16, @block=[8, 16]> -> vector<256x128xbf6> @@ -401,9 +393,9 @@ This example demonstrates a cooperative transpose pattern in which a matrix tile #Coop_wg = {sg_layout = [8, 4] , sg_data= [32, 8], order=[1, 0] } #dpas_wg = {sg_layout = [8, 4], sg_data= [32, 32], order=[1, 0] } -%at = load_nd %tdesc: tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16 > -%a = vector.transpose %1 #Coop_wg :vector<32x256xf16> -> vector<256x32xf16> -%a_dpas = Conv_layout %2 #Coop_wg #dpas_wg +%at = load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> +%a = vector.transpose %1 {layout_result_0 = #Coop_wg}: vector<32x256xf16> to vector<256x32xf16> +%a_dpas = Conv_layout %2 <{from = #Coop_wg, to = #dpas_wg}>: vector<256x32xf16> ``` In this flow: @@ -415,50 +407,51 @@ In this flow: **After optimization that targets the transpose-A pattern** -The code is transformed to use store_matrix and load_matrix to materialize the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. +The code is transformed to use store_matrix and load_matrix to implement the transpose cooperatively in shared local memory. Note that both load_nd and store_matrix use smaller sg_data values, meaning each subgroup processes a smaller fragment, enabling a cooperative transpose across threads. -It is generally preferred to fuse transpose and convert_layout earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. +It is generally preferred to detect the “transpose + convert_layout” pattern and fuse them earlier in the pipeline, as this affects the blocking strategy for load_matrix and store_matrix (which are the lowered forms of the logical layout conversion and transpose). Early fusion enables better alignment with optimal hardware load instructions. ```mlir #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], order = [0, 1] } // original layout #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], order = [1, 0] } // target DPAS layout -%at = xegpu.load_nd %tdesc : tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> -%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> -%mt = xegpu.matrix_desc_view %m : matrix_desc<32x256xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> -xegpu.store_matrix %mt[0, 0], %at : vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +%at = xegpu.load_nd %tdesc : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> +%mt = xegpu. create_matrix_desc %m : memref<16384xi8, 3>-> matrix_desc<32x256xf16, @strides=[1, 32]> +xegpu.store_matrix %at, %mt[0, 0] #Coop_t_wg: vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32]> gpu.barrier -%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> -%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> +%ma = xegpu.create_matrix_desc %m : memref<16384xi8, 3>-> matrix_desc<256x32xf16> +%a_dpas = xegpu.load_matrix %ma[0, 0] #dpas_t_wg: matrix_desc<256x32xf16> -> vector<256x32xf16> ``` -**Adding Attributes for Basic Instruction-Level Blocking** +**Layout Assignment** +***Basic Blocking: Using regular load and store instruction*** -In this example, the xegpu.layout is extended to support instruction-level blocking. The basic blocking assumes 16 lanes, and each lane handles 2 f16 elements (32 bits). This basic blocking does not change the physical memory layout (i.e., there is no memory-level tiling); instead, it loweres to instructions like store_scatter and load_gather. +In this example, the xegpu.layout is extended to support instruction-level blocking. The basic blocking assumes 16 lanes, and each lane handles 2 f16 elements (32 bits). This basic instruction blocking does not try to block memory layout. It lowers to instructions like chunked store and load_gather. ```mlir #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [1, 32], order = [0, 1] } #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [1, 32], order = [1, 0] } -%at = xegpu.load_nd %tdesc: tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> -%m = xegpu.create_matrix_desc: matrix_desc<8192xf16> -%mt = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> -xegpu.store_matrix %mt[0, 0], %at: vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32], #Coop_t_wg> +%at = xegpu.load_nd %tdesc: tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> +%m = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<32x256xf16, @strides=[1, 32]> +xegpu.store_matrix %at, %mt[0, 0] #Coop_t_wg: vector<32x256xf16>, matrix_desc<32x256xf16, @strides=[1, 32]> gpu.barrier -%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, #dpas_t_wg> -%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, #dpas_t_wg> -> vector<256x32xf16> +%ma = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<256x32xf16> +%a_dpas = xegpu.load_matrix %ma[0, 0] #dpas_t_wg: matrix_desc<256x32xf16> -> vector<256x32xf16> ``` -**Optimized Blocking: Lowering to store_chunk and 1D Block Load** +***Optimized Blocking: Lowering to store_chunk and 1D Block Load*** -This pattern demonstrates a more optimized strategy for instruction-level blocking, enabling the use of efficient memory instructions such as store_chunk and 1D block load. For correct and efficient lowering, several constraints must be satisfied: +This pattern demonstrates a more optimized strategy for instruction-level blocking, enabling the use of efficient memory instructions such as 1D block load. For correct and efficient lowering, several constraints must be satisfied: -The inst_data field must specify a meaningful 2D shape that aligns with the capabilities of store_chunk and 1D block load. +- The inst_data field must specify a meaningful 2D shape that aligns with the capabilities of chunked store and 1D block load. -Blocking must be explicitly expressed in the memory layout via the @block attribute. Two related matrix_desc subviews (e.g., producer and consumer) must have consistent block sizes. If one matrix_desc is transposed, the block shape should match the transposed shape of the other one. +- Blocking must be explicitly expressed in the memory layout via the @block attribute. Two related matrix_desc (e.g., producer and consumer) must have consistent block sizes. If one matrix_desc is transposed, the block shape should match the transposed shape of the other one. -Each instruction must access only within its assigned matrix block boundary — no cross-block accesses are allowed. +- Each instruction must access only within its assigned matrix block boundary — no cross-block accesses are allowed. During lowering, store_matrix is lowered to store_chunk if the matrix has strides, and load_matrix is lowered to 1D block load if the matrix has a blocked layout. @@ -466,21 +459,21 @@ During lowering, store_matrix is lowered to store_chunk if the matrix has stride #Coop_t_wg = { sg_layout = [4, 8], sg_data = [8, 32], inst_data = [8, 16], order = [0, 1] } #dpas_t_wg = { sg_layout = [8, 4], sg_data = [32, 32], inst_data = [16, 16], order = [1, 0] } -%at = xegpu.load_nd %tdesc : tensor_desc<4096x4096xf16, #Coop_t_wg> -> vector<32x256xf16> -%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> -%mt = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> -xegpu.store_matrix %mt[0, 0], %at : vector<32x256xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #Coop_t_wg> +%at = xegpu.load_nd %tdesc : tensor_desc<32x256xf16, #Coop_t_wg> -> vector<32x256xf16> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> +%mt = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> +xegpu.store_matrix %at, %mt[0, 0] #Coop_t_wg : vector<32x256xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> gpu.barrier -%ma = xegpu.matrix_desc_view %m : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -%a_dpas = xegpu.load_matrix %ma[0, 0] : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -> vector<256x32xf16> +%ma = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<256x32xf16, @block=[16, 16]> +%a_dpas = xegpu.load_matrix %ma[0, 0] #dpas_t_wg : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_wg> -> vector<256x32xf16> ``` **Workgroup to Subgroup Distribution** -This example illustrates how data is distributed from workgroup to subgroups. It demonstrates how load_matrix and store_matrix cooperate with matrix_desc subviews to enable efficient subgroup distribution. In this step, the sg_layout and sg_data attributes are removed from the layout specification, leaving only the inst_data attribute. +This example illustrates how load_matrix and store_matrix are distributed from workgroup to subgroups. After distribution, the sg_layout and sg_data attributes are removed from the layout specification, leaving only the inst_data attribute. -The matrix is assumed to be stored in row-major contiguous layout, and indexing into it is performed using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the matrix_desc data type is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. +The distribution process assumes matrix stored in row-major contiguous layout, and performes indexing using logical coordinates. These logical coordinates are used throughout tile distribution and layout transformations. Only at the final lowering stage (e.g., MaterializeSLMAccess) are physical offsets computed using memory layout attributes such as @strides and @block. A key property of the matrix_desc data type is that logical tile decomposition does not alter the block or stride metadata, making logical address computation straightforward. ```mlir #coop_t_inst = { inst_data = [8, 16] } @@ -491,29 +484,19 @@ The matrix is assumed to be stored in row-major contiguous layout, and indexing : memref<4096x4096xf16> -> tensor_desc<8x32xf16, #coop_t_inst> %at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16, #coop_t_inst> -> vector<8x32xf16> -%m = xegpu.create_matrix_desc - : matrix_desc<32x256xf16> -%mt = xegpu.matrix_desc_view %m - : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> -%mt_sg = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32] - : matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> - -> matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> -xegpu.store_matrix %mt_sg, %at - : vector<8x32xf16>, matrix_desc<8x32xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> +%mt = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> +xegpu.store_matrix %at, %mt[%sg_idy * 8, %sg_idx * 32] #coop_t_inst + : vector<8x32xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> gpu.barrier - -%ma = xegpu.matrix_desc_view %m - : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> -%ma_sg = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * (32 % 32)] - : matrix_desc<256x32xf16, @block=[16, 16], #dpas_t_inst> - -> matrix_desc<32x32xf16, @block=[16, 16]> -%a_dpas = xegpu.load_matrix %ma_sg - : matrix_desc<32x32xf16, @block=[16, 16], #dpas_t_inst> -> vector<32x32xf16> +%ma = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<256x32xf16, @block=[16, 16]> +%a_dpas = xegpu.load_matrix %ma[%sg_idy * 32, %sg_idx * (32 % 32)] #dpas_t_inst + : matrix_desc<256x32xf16, @block=[16, 16]> -> vector<32x32xf16> ``` -**Unrolling Guided by Inst_data** -This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This pattern ensures that each load and store operation writes within its assigned block boundary, respecting the @block and @strides attributes. On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. +**Unrolling Guided by Inst_data** +This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This inst_data attributes ensures that each store operation writes within its assigned block boundary, respecting the @block attributes. On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. ```mlir %tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] @@ -521,38 +504,23 @@ This example illustrates how matrix loads and stores can be unrolled into smalle %at = xegpu.load_nd %tdesc_sg : tensor_desc<8x32xf16> -> vector<8x32xf16> %at0 = vector.extract %at[0, 0] : vector<8x32xf16> -> vector<8x16xf16> %at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> -%m = xegpu.create_matrix_desc : matrix_desc<8192xf16> -%mt = xegpu.matrix_desc_view %m - : matrix_desc<8192xf16> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32], #coop_t_inst> -%mt_inst0 = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32] - : matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> -%mt_inst1 = xegpu.matrix_desc_subview %mt[%sg_idy * 8, %sg_idx * 32 + 16] - : matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> -xegpu.store_matrix %mt_inst0, %at0 - : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> -xegpu.store_matrix %mt_inst1, %at1 - : vector<8x16xf16>, matrix_desc<8x16xf16, @block=[16, 16], @strides=[1, 32]> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> +%mt = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> +xegpu.store_matrix %at0, %mt[%sg_idy * 8, %sg_idx * 32] + : vector<8x16xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> +xegpu.store_matrix %at1, %mt[%sg_idy * 8, %sg_idx * 32 + 16] + : vector<8x16xf16>, matrix_desc<32x256xf16, @block=[16, 16], @strides=[1, 32]> gpu.barrier - -%ma = xegpu.matrix_desc_view %m - : matrix_desc<8192xf16> -> matrix_desc<256x32xf16, @block=[16, 16]> -%ma_inst0 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32] - : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> -%ma_inst1 = xegpu.matrix_desc_subview %ma[%sg_idy * 32, %sg_idx * 32 % 32 + 16] - : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> -%ma_inst2 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32] - : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> -%ma_inst3 = xegpu.matrix_desc_subview %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32 + 16] - : matrix_desc<256x32xf16> -> matrix_desc<16x16xf16, @block=[16, 16]> -%a_dpas_0 = xegpu.load_matrix %ma_inst0 - : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> -%a_dpas_1 = xegpu.load_matrix %ma_inst1 - : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> -%a_dpas_2 = xegpu.load_matrix %ma_inst2 - : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> -%a_dpas_3 = xegpu.load_matrix %ma_inst3 - : matrix_desc<16x16xf16, @block=[16, 16]> -> vector<16x16xf16> +%ma = xegpu.create_matrix_desc %m : memref<16384xi8, 3> -> matrix_desc<256x32xf16, @block=[16, 16]> +%a_dpas_0 = xegpu.load_matrix %ma[%sg_idy * 32, %sg_idx * 32 % 32] + : matrix_desc<256x32xf16, @block=[16, 16]> -> vector<16x16xf16> +%a_dpas_1 = xegpu.load_matrix %ma[%sg_idy * 32, %sg_idx * 32 % 32 + 16] + : matrix_desc<256x32xf16, @block=[16, 16]> -> vector<16x16xf16> +%a_dpas_2 = xegpu.load_matrix %ma[%sg_idy * 32 + 16, %sg_idx * 32 % 32] + : matrix_desc<256x32xf16, @block=[16, 16]> -> vector<16x16xf16> +%a_dpas_3 = xegpu.load_matrix %[%sg_idy * 32 + 16, %sg_idx * 32 % 32 + 16] + : matrix_desc<256x32xf16, @block=[16, 16]> -> vector<16x16xf16> ``` **MaterializeSLMAccess: Lowering matrix_desc to Physical Memory Access** @@ -560,13 +528,11 @@ gpu.barrier This step lowers high-level matrix_desc operations (store_matrix, load_matrix) into low-level memory operations (store_chunk, load_1d) over shared local memory. It performs full address materialization using the matrix's layout attributes (@strides, @block) and logical lane coordinates. Key Concepts: -Matrix-to-Memory Conversion: Replace matrix_desc-based tile abstractions with raw memref and compute physical offsets explicitly. +- Chunked Store: Each thread stores a small fragment (e.g., 8×1) using the logical offset composed with layout metadata. Lowered to store_chunk. -Chunked Store: Each thread stores a small fragment (e.g., 8×1) using the logical offset composed with layout metadata. Lowered to store_chunk. +- 1D Block Load: A transposed layout (e.g., 256×32) is blocked as 16×16 tiles. Contiguous blocks are loaded using load_1d, which requires computing the physical offset of the first element per 1D block. -1D Block Load: A transposed layout (e.g., 256×32) is blocked as 16×16 tiles. Contiguous blocks are loaded using load_1d, which requires computing the physical offset of the first element per 1D block. - -Offset Calculation: Logical per-lane coordinates are transformed into logical block coordinates, then to physical offsets using block size and strides. +- Offset Calculation: Logical per-lane coordinates are transformed into logical block coordinates, then to physical offsets using block size and strides. ```mlir %tdesc_sg = xegpu.create_nd_tdesc %base[%widy * 32 + %sg_idy * 8, %widx * 256 + %sg_idx * 32] @@ -576,41 +542,38 @@ Offset Calculation: Logical per-lane coordinates are transformed into logical bl %at1 = vector.extract %at[0, 16] : vector<8x32xf16> -> vector<8x16xf16> // Shared local memory buffer -%m = memref.alloc : memref<8192xf16, 3> +%m = memref.alloca() {alignment = 1024} : memref<16384xi8, 3> + +// ---------------------- Chunked Store ---------------------- +// The transpose is added as we remove the transpose attribute out from chunked load/store and expect an explict data transpose. +// it will be no op after lane distribution since each lane owns same data when [8,1] is transpose to [1, 8] +%at0_t = vector.transpose %at0 : vector<8x16xf16> -> vector<16x8xf16> -// ---------------------- Chunked Load ---------------------- // Compute blocked offset vectors for SLM store -%blk_y = divi_signed %sg_idy * 8, 16 : index -%blk_in_y = remi_signed %sg_idy * 8, 16 : index -%sg_idx_vec = addi %sg_idx * 32, dense<[0, ..., 15]> : vector<16xindex> -%blk_x = divi_unsigned %sg_idx_vec, 16 -%blk_in_x = remi_unsigned %sg_idx_vec, 16 +%blk_y=sg_idy*8 /16: index +%blk_in_y=sg_idy*8 %16: index +%sg_idx_vec = %sg_idx*32 + [0..15] : vector<16xindex> +%blk_x=%sg_idx_vec /16: vector<16xindex > +%blk_in_x=%sg_idx_vec %16: vector<16xindex > + // calculate physic addresses with pre-computed strides of the blocked matrix. // [32x256, strides=1x32] blocked as [2x16x16x16, strides=256x512x1x16] -%offset_vec0 = addi (addi (addi (%blk_in_y, %blk_in_x * 16), %blk_y * 256),%blk_x * 512) -%tdesc0 = xegpu.create_tdesc %m, %offset_vec0 - : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> +%offset_vec0 = %blk_y * 256+ + %blk_x * 512 + %blk_in_y + %blk_in_x*16 +xegpu.store %at0_t, %m, %offset_vec0 @chunk_size=8: vector<16x8xf16>, memref<8192xf16, 3>, vector<16xindex> // Repeat for second tile -%sg_idx_vec2 = addi %sg_idx * 32, dense<[16, ..., 31]> : vector<16xindex> -%blk_x2 = divi_unsigned %sg_idx_vec2, 16 -%blk_in_x2 = remi_unsigned %sg_idx_vec2, 16 -%offset_vec1 = addi (addi (addi (%blk_in_y, %blk_in_x2 * 16), %blk_y * 256),%blk_x2 * 512) -%tdesc1 = xegpu.create_tdesc %m, %offset_vec1 - : memref<8192xf16, 3>, vector<16xindex> -> tdesc<16x8xf16, chunk=8, scope=slm> - -// The transpose is added as we remove the transpose attribute out from chunked load/store and expect an explict data transpose. -// it will be no op after lane distribution since each lane owns same data when [8,1] is transpose to [1, 8] -%at0_t = vector.transpose %at0 : vector<8x16xf16> -> vector<16x8xf16> %at1_t = vector.transpose %at1 : vector<8x16xf16> -> vector<16x8xf16> -xegpu.store %tdesc0, %at0_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> -xegpu.store %tdesc1, %at1_t : tdesc<16x8xf16, chunk=8, scope=slm>, vector<16x8xf16> +%sg_idx_vec2 = %sg_idx*32 + [16..31] : vector<16xindex> +%blk_x2=%sg_idx_vec2 /16: vector<16xindex > +%blk_in_x2=%sg_idx_vec2 %16: vector<16xindex > +%offset_vec1 = %blk_y * 256+ + %blk_x2 * 512 + %blk_in_y+ %blk_in_x2*16 +xegpu.store %at1_t, %m, %offset_vec1: @chunk_size=8: vector<16x8xf16>, memref<8192xf16, 3>, vector<16xindex> gpu.barrier // ---------------------- Load 1D Block ---------------------- // Compute per-block physical offsets -// pre-computed strides of the blocked matrix: [256x32] blocked as [2x16x16x16, strides=512x256x16x1] +// pre-computed strides of the blocked matrix: [256x32] blocked as [16x2x16x16, strides=512x256x16x1] // sg_idx*32 coord to blocked matrix ccord: sg_idx*32%32/16 (0), sg_idx*32%32%16 (0). %32 due matrix shape[1] is 32 // sg_idy*32 coord to blocked matrix coord: sg_idy*32/16, sg_idy*32%16 (0) // then map to physical addr using stride [2x16x16x16, strides=512x256x16x1], get sg_idy*32/16 *512 @@ -619,15 +582,10 @@ gpu.barrier %inst_start_offset2 = add %inst_start_offset0, 512 %inst_start_offset3 = add %inst_start_offset0, 768 -%tdesc0 = xegpu.create_nd_tdesc %m, %inst_start_offset0 : memref<8192xf16, 3>, index -> tdesc<256xf16> -%tdesc1 = xegpu.create_nd_tdesc %m, %inst_start_offset1 : memref<8192xf16, 3>, index -> tdesc<256xf16> -%tdesc2 = xegpu.create_nd_tdesc %m, %inst_start_offset2 : memref<8192xf16, 3>, index -> tdesc<256xf16> -%tdesc3 = xegpu.create_nd_tdesc %m, %inst_start_offset3 : memref<8192xf16, 3>, index -> tdesc<256xf16> - -%a_dpas_0 = xegpu.load_nd %tdesc0 : tdesc<256xf16> -> vector<256xf16> -%a_dpas_1 = xegpu.load_nd %tdesc1 : tdesc<256xf16> -> vector<256xf16> -%a_dpas_2 = xegpu.load_nd %tdesc2 : tdesc<256xf16> -> vector<256xf16> -%a_dpas_3 = xegpu.load_nd %tdesc3 : tdesc<256xf16> -> vector<256xf16> +%a_dpas_0 = xegpu.load_nd %m, %inst_start_offset0 : memref<8192xf16, 3>, index -> vector<256xf16> +%a_dpas_1 = xegpu.load_nd %m, %inst_start_offset1 : memref<8192xf16, 3>, index -> vector<256xf16> +%a_dpas_2 = xegpu.load_nd %m, %inst_start_offset2 : memref<8192xf16, 3>, index -> vector<256xf16> +%a_dpas_3 = xegpu.load_nd %m, %inst_start_offset3 : memref<8192xf16, 3>, index -> vector<256xf16> ``` ## XeGPU Attributes to support Work Item Level semantics From da43f5d3e38ece998bffc77a86b94700918dd170 Mon Sep 17 00:00:00 2001 From: Jianhui Li Date: Mon, 21 Jul 2025 16:02:17 -0700 Subject: [PATCH 15/15] Update XeGPU.md --- docs/rfcs/XeGPU.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/rfcs/XeGPU.md b/docs/rfcs/XeGPU.md index 0e953058b..bea245905 100644 --- a/docs/rfcs/XeGPU.md +++ b/docs/rfcs/XeGPU.md @@ -496,6 +496,7 @@ gpu.barrier ``` **Unrolling Guided by Inst_data** + This example illustrates how matrix loads and stores can be unrolled into smaller instruction tiles for better alignment with hardware capabilities. This inst_data attributes ensures that each store operation writes within its assigned block boundary, respecting the @block attributes. On the load side, the matrix_desc is subviewed into multiple 16×16 instruction tiles, which are then used in separate load_matrix operations. This breakdown enables explicit instruction-level unrolling, allowing each instruction to operate on a fixed tile size that aligns with DPAS or tensor-core instruction requirements. ```mlir