derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
509 lines
39 KiB
MLIR
509 lines
39 KiB
MLIR
// RUN: mlir-opt -xevm-attach-target='chip=cri' -test-xegpu-propagate-layouts="layout-kind=subgroup" -split-input-file %s | FileCheck %s
|
|
|
|
gpu.module @test {
|
|
// CHECK-LABEL: store_nd
|
|
// CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
|
|
func.func @store_nd(%src: memref<256x128xf32>) {
|
|
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32>
|
|
// CHECK-SAME: -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
|
|
// CHECK-SAME: : !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
|
|
// CHECK-SAME: -> vector<256x128xf32>
|
|
// CHECK: xegpu.store_nd %[[LOAD]], %[[TDESC]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
|
|
// CHECK-SAME: : vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>>
|
|
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
|
|
%load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
|
|
xegpu.store_nd %load, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}>
|
|
: vector<256x128xf32>, !xegpu.tensor_desc<256x128xf32>
|
|
return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_transpose
|
|
// CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
|
|
// CHECK-SAME: %[[ARG_1:.*]]: memref<128x256xf32>
|
|
func.func @vector_transpose(%src: memref<256x128xf32>, %src1: memref<128x256xf32>) {
|
|
// CHECK: %[[TDESC_LD:.*]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 16], order = [0, 1]>>
|
|
// CHECK: %[[TDESC_ST:.*]] = xegpu.create_nd_tdesc %[[ARG_1]] : memref<128x256xf32> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64], order = [1, 0]>>
|
|
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC_LD]][0, 0] <{layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 16], order = [0, 1]>}>
|
|
// CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [64, 16], order = [0, 1]>> -> vector<256x128xf32>
|
|
|
|
// CHECK: %[[TRANSPOSED:.*]] = vector.transpose %2, [1, 0]
|
|
// CHECK-SAME {layout_result_0 = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64], order = [1, 0]>} : vector<256x128xf32> to vector<128x256xf32>
|
|
|
|
// CHECK: xegpu.store_nd %[[TRANSPOSED]], %[[TDESC_ST]][0, 0]
|
|
// CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64], order = [1, 0]>}> : vector<128x256xf32>,
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64], order = [1, 0]>>
|
|
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
|
|
%tdesc1 = xegpu.create_nd_tdesc %src1 : memref<128x256xf32> -> !xegpu.tensor_desc<128x256xf32>
|
|
%load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
|
|
%trans = vector.transpose %load, [1, 0] : vector<256x128xf32> to vector<128x256xf32>
|
|
xegpu.store_nd %trans, %tdesc1[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [16, 64], order = [1, 0]>}>
|
|
: vector<128x256xf32>, !xegpu.tensor_desc<128x256xf32>
|
|
return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_transpose
|
|
// CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
|
|
// CHECK-SAME: %[[ARG_1:.*]]: memref<128x256xf32>
|
|
gpu.func @vector_transpose(%src: memref<256x128xf32>, %src1: memref<128x256xf32>) kernel attributes
|
|
{known_block_size = array<i32: 1, 32, 16>} {
|
|
// CHECK: %[[TDESC_LD:.*]] = xegpu.create_nd_tdesc %[[ARG_0]] : memref<256x128xf32> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], order = [0, 1]>>
|
|
// CHECK: %[[TDESC_ST:.*]] = xegpu.create_nd_tdesc %[[ARG_1]] : memref<128x256xf32> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32]>>
|
|
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC_LD]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], order = [0, 1]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], order = [0, 1]>> -> vector<256x128xf32>
|
|
|
|
// CHECK: %[[TRANSPOSED:.*]] = vector.transpose %2, [1, 0]
|
|
// CHECK-SAME {layout_result_0 = #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32]>} : vector<256x128xf32> to vector<128x256xf32>
|
|
|
|
// CHECK: xegpu.store_nd %[[TRANSPOSED]], %[[TDESC_ST]][0, 0]
|
|
// CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32]>}> : vector<128x256xf32>,
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x256xf32, #xegpu.layout<sg_layout = [4, 8], sg_data = [32, 32]>>
|
|
%tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32> -> !xegpu.tensor_desc<256x128xf32>
|
|
%tdesc1 = xegpu.create_nd_tdesc %src1 : memref<128x256xf32> -> !xegpu.tensor_desc<128x256xf32>
|
|
%load = xegpu.load_nd %tdesc[0, 0] : !xegpu.tensor_desc<256x128xf32> -> vector<256x128xf32>
|
|
%trans = vector.transpose %load, [1, 0] : vector<256x128xf32> to vector<128x256xf32>
|
|
xegpu.store_nd %trans, %tdesc1[0, 0] : vector<128x256xf32>, !xegpu.tensor_desc<128x256xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: dpas
|
|
// CHECK-SAME: %[[A_MEMREF:.*]]: memref<128x128xf16>, %[[B_MEMREF:.*]]: memref<128x128xf16>
|
|
// CHECK-SAME: %[[CD_MEMREF:.*]]: memref<128x128xf32>
|
|
gpu.func @dpas(%a: memref<128x128xf16>, %b: memref<128x128xf16>, %d: memref<128x128xf32>) kernel attributes
|
|
{known_block_size = array<i32: 1, 64, 16>} {
|
|
// CHECK: %[[TDESC_A:.*]] = xegpu.create_nd_tdesc %[[A_MEMREF]] : memref<128x128xf16> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>>
|
|
|
|
// CHECK: %[[A_LOADED:.*]] = xegpu.load_nd %[[TDESC_A]][0, 0]
|
|
// CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>}>
|
|
// CHECK-SAME: : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>> -> vector<128x128xf16>
|
|
|
|
// CHECK: %[[TDESC_B:.*]] = xegpu.create_nd_tdesc %[[B_MEMREF]] : memref<128x128xf16> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>>
|
|
|
|
// CHECK: %[[B_LOADED:.*]] = xegpu.load_nd %[[TDESC_B]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>}>
|
|
// CHECK-SAME: : !xegpu.tensor_desc<128x128xf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>> -> vector<128x128xf16>
|
|
|
|
// CHECK: %[[DPAS_RES:.*]] = xegpu.dpas %[[A_LOADED]], %[[B_LOADED]]
|
|
// CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128]>,
|
|
// CHECK-SAME: layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16]>,
|
|
// CHECK-SAME: layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} :
|
|
// CHECK-SAME: vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>
|
|
|
|
// CHECK: %[[TDESC_ST:.*]] = xegpu.create_nd_tdesc %[[CD_MEMREF]] : memref<128x128xf32> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
|
|
// CHECK: xegpu.store_nd %[[DPAS_RES]], %[[TDESC_ST]][0, 0]
|
|
// CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
|
|
// CHECK-SAME: vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf16> -> !xegpu.tensor_desc<128x128xf16>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf16> -> !xegpu.tensor_desc<128x128xf16>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] : !xegpu.tensor_desc<128x128xf16> -> vector<128x128xf16>
|
|
%dpas = xegpu.dpas %load_a, %load_b : vector<128x128xf16>, vector<128x128xf16> -> vector<128x128xf32>
|
|
%tdesc_cd = xegpu.create_nd_tdesc %d : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32>
|
|
xegpu.store_nd %dpas, %tdesc_cd[0, 0] : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_row_reduction
|
|
// CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 64]>, dims = [1]>}
|
|
gpu.func @vector_row_reduction(%src: memref<32x64xf32>, %dst: memref<32xf32>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<32xf32>
|
|
%tdesc_src = xegpu.create_nd_tdesc %src : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32>
|
|
%load = xegpu.load_nd %tdesc_src[0, 0] : !xegpu.tensor_desc<32x64xf32> -> vector<32x64xf32>
|
|
%reduce = vector.multi_reduction <add>, %load, %cst [1] : vector<32x64xf32> to vector<32xf32>
|
|
%tdesc_dst = xegpu.create_nd_tdesc %dst : memref<32xf32> -> !xegpu.tensor_desc<32xf32>
|
|
xegpu.store_nd %reduce, %tdesc_dst[0] <{layout = #xegpu.layout<sg_layout = [32], sg_data = [1]>}>
|
|
: vector<32xf32>, !xegpu.tensor_desc<32xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_row_reduction_scalar
|
|
// CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %{{.*}}, %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [32, 1], sg_data = [1, 64]>, dims = [0, 1]>}
|
|
gpu.func @vector_row_reduction_scalar(%src: memref<32x64xf32>, %dst: memref<32xf32>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
|
|
%cst = arith.constant 0.000000e+00 : f32
|
|
%tdesc_src = xegpu.create_nd_tdesc %src : memref<32x64xf32> -> !xegpu.tensor_desc<32x64xf32>
|
|
%load = xegpu.load_nd %tdesc_src[0, 0] : !xegpu.tensor_desc<32x64xf32> -> vector<32x64xf32>
|
|
%reduce = vector.multi_reduction <add>, %load, %cst [0, 1] : vector<32x64xf32> to f32
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_nest_reduction
|
|
gpu.func @vector_nest_reduction(%src: memref<32x128xf32>, %dst: memref<32xf32>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<32xf32>
|
|
%cst1 = arith.constant dense<0.000000e+00> : vector<32x128xf32>
|
|
%tdesc_src = xegpu.create_nd_tdesc %src : memref<32x128xf32> -> !xegpu.tensor_desc<32x128xf32>
|
|
%load = xegpu.load_nd %tdesc_src[0, 0] : !xegpu.tensor_desc<32x128xf32> -> vector<32x128xf32>
|
|
%bcast1 = vector.broadcast %load: vector<32x128xf32> to vector<4x32x128xf32>
|
|
|
|
// CHECK: %[[BCAST1:.*]] = vector.broadcast %{{.*}} {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>} : vector<32x128xf32> to vector<4x32x128xf32>
|
|
// CHECK: %[[BCAST:.*]] = vector.multi_reduction <add>, %[[BCAST1]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>} [0] : vector<4x32x128xf32> to vector<32x128xf32>
|
|
// CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %[[BCAST]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>, dims = [1]>} [1] : vector<32x128xf32> to vector<32xf32>
|
|
|
|
%bcast = vector.multi_reduction <add>, %bcast1, %cst1 [0]: vector<4x32x128xf32> to vector<32x128xf32>
|
|
%reduce = vector.multi_reduction <add>, %bcast, %cst [1] : vector<32x128xf32> to vector<32xf32>
|
|
%mask = arith.constant dense<1>: vector<32xi1>
|
|
%offset = vector.step : vector<32xindex>
|
|
xegpu.store %reduce, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>, dims = [1]>} : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: vector_nest_reduction_with_nest_slice_layout
|
|
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>, dims = [1]>} dense<0.000000e+00> : vector<32xf32>
|
|
// CHECK: %[[CST0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>} dense<0.000000e+00> : vector<32x128xf32>
|
|
// CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<32x128xf32> -> !xegpu.tensor_desc<32x128xf32, #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>>
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]][0, 0] <{layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>}>
|
|
// CHECK-SAME: -> vector<32x128xf32>
|
|
// CHECK: %[[BCAST1:.*]] = vector.broadcast %[[LOAD]] {layout_result_0 = #xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>} : vector<32x128xf32> to vector<4x32x128xf32>
|
|
// CHECK: %[[REDUCE1:.*]] = vector.multi_reduction <add>, %[[BCAST1]], %[[CST0]]
|
|
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>} [0] : vector<4x32x128xf32> to vector<32x128xf32>
|
|
// CHECK: %[[REDUCE2:.*]] = vector.multi_reduction <add>, %[[REDUCE1]], %[[CST]]
|
|
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 16]>, dims = [0]>, dims = [1]>} [1] : vector<32x128xf32> to vector<32xf32>
|
|
// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>} dense<true> : vector<32xi1>
|
|
// CHECK: %[[OFFSET:.*]] = vector.step {layout_result_0 = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>} : vector<32xindex>
|
|
// CHECK: xegpu.store %[[REDUCE2]], %{{.*}}[%[[OFFSET]]], %[[MASK]]
|
|
// CHECK-SAME: <{layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>}>
|
|
// CHECK-SAME: : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
|
|
gpu.func @vector_nest_reduction_with_nest_slice_layout(%src: memref<32x128xf32>, %dst: memref<32xf32>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<32xf32>
|
|
%cst1 = arith.constant dense<0.000000e+00> : vector<32x128xf32>
|
|
%tdesc_src = xegpu.create_nd_tdesc %src : memref<32x128xf32> -> !xegpu.tensor_desc<32x128xf32>
|
|
%load = xegpu.load_nd %tdesc_src[0, 0] : !xegpu.tensor_desc<32x128xf32> -> vector<32x128xf32>
|
|
%bcast1 = vector.broadcast %load: vector<32x128xf32> to vector<4x32x128xf32>
|
|
%bcast = vector.multi_reduction <add>, %bcast1, %cst1 [0]: vector<4x32x128xf32> to vector<32x128xf32>
|
|
%reduce = vector.multi_reduction <add>, %bcast, %cst [1] : vector<32x128xf32> to vector<32xf32>
|
|
%mask = arith.constant dense<1>: vector<32xi1>
|
|
%offset = vector.step : vector<32xindex>
|
|
xegpu.store %reduce, %dst[%offset], %mask {layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<sg_layout = [1, 4, 8], sg_data = [4, 8, 32]>, dims = [0]>, dims = [1]>} : vector<32xf32>, memref<32xf32>, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#data_layout = #xegpu.layout<sg_layout = [1, 16], sg_data = [4, 16]>
|
|
gpu.module @test {
|
|
// CHECK-LABEL: gpu.func @vector_reduction_preserve_round_robin_layout
|
|
gpu.func @vector_reduction_preserve_round_robin_layout(%arg0: memref<2048xf32, 1>) kernel {
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0{{.*}} <{layout = #xegpu.layout<sg_layout = [1, 16], sg_data = [4, 16]>}> {{.*}} -> vector<8x256xf32>
|
|
// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[LOAD]], {{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [1, 16], sg_data = [4, 16]>, dims = [1]>} [1] : vector<8x256xf32> to vector<8xf32>
|
|
%offset = arith.constant dense<0> : vector<8x256xindex>
|
|
%offset2 = arith.constant dense<1024> : vector<8xindex>
|
|
%acc = arith.constant dense<0.000000e+00> : vector<8xf32>
|
|
%mask = arith.constant dense<true> : vector<8x256xi1>
|
|
%mask2 = arith.constant dense<true> : vector<8xi1>
|
|
%val = xegpu.load %arg0[%offset], %mask : memref<2048xf32, 1>, vector<8x256xindex>, vector<8x256xi1> -> vector<8x256xf32>
|
|
%reduce = vector.multi_reduction <add>, %val, %acc [1] : vector<8x256xf32> to vector<8xf32>
|
|
xegpu.store %reduce, %arg0[%offset2], %mask2 { layout = #xegpu.slice<#data_layout, dims = [1]> } : vector<8xf32>, memref<2048xf32, 1>, vector<8xindex>, vector<8xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: for_loop_dpas
|
|
gpu.func @for_loop_dpas(%arg0: memref<2048x8192xf16>, %arg1: memref<8192x4096xf16>, %arg2: memref<2048x4096xf32>) kernel attributes {known_block_size = array<i32: 4, 1, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<128x128xf32>
|
|
%c128 = arith.constant 128 : index
|
|
%c8192 = arith.constant 8192 : index
|
|
%c0 = arith.constant 0 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%0 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%block_id_x]
|
|
%1 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%block_id_y]
|
|
// CHECK: %2 = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%{{.*}} = %{{.*}}) -> (vector<128x128xf32>) {
|
|
// CHECK-NEXT: xegpu.create_nd_tdesc %{{.*}} : memref<2048x8192xf16> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>,
|
|
// CHECK-SAME: #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 128]>>
|
|
|
|
// CHECK-NEXT: xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 128]>}>
|
|
|
|
// CHECK-NEXT: xegpu.create_nd_tdesc %{{.*}} : memref<8192x4096xf16> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>,
|
|
// CHECK-SAME: #xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>>
|
|
|
|
// CHECK-NEXT: xegpu.load_nd %6[%arg3, %block_id_y] <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>}>
|
|
|
|
// CHECK-NEXT: xegpu.dpas %{{.*}} {
|
|
// CHECK-SAME: layout_a = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 128]>,
|
|
// CHECK-SAME: layout_b = #xegpu.layout<sg_layout = [2, 2], sg_data = [128, 64]>,
|
|
// CHECK-SAME: layout_cd = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}
|
|
// CHECK-SAME: : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
|
|
|
|
// CHECK-NEXT: scf.yield %{{.*}} : vector<128x128xf32>
|
|
// CHECK-NEXT: } {layout_operand_3 = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>, layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}
|
|
// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [64, 64]>}>
|
|
|
|
%2 = scf.for %arg3 = %c0 to %c8192 step %c128 iter_args(%arg4 = %cst) -> (vector<128x128xf32>) {
|
|
%4 = xegpu.create_nd_tdesc %arg0 : memref<2048x8192xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%5 = xegpu.load_nd %4[%block_id_x, %arg3] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
|
|
%6 = xegpu.create_nd_tdesc %arg1 : memref<8192x4096xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%7 = xegpu.load_nd %6[%arg3, %block_id_y] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
|
|
%8 = xegpu.dpas %5, %7, %arg4 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
|
|
scf.yield %8 : vector<128x128xf32>
|
|
}
|
|
%3 = xegpu.create_nd_tdesc %arg2 : memref<2048x4096xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
xegpu.store_nd %2, %3[%0, %1] : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: for_loop_misaligned_dpas_fail
|
|
gpu.func @for_loop_misaligned_dpas_fail(%arg0: memref<2048x8192xf16>, %arg1: memref<8192x4096xf16>, %arg2: memref<2048x4096xf32>) kernel attributes {known_block_size = array<i32: 8, 1, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<128x128xf32>
|
|
%c128 = arith.constant 128 : index
|
|
%c8192 = arith.constant 8192 : index
|
|
%c0 = arith.constant 0 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%0 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%block_id_x]
|
|
%1 = affine.apply affine_map<()[s0] -> (s0 * 128)>()[%block_id_y]
|
|
%2 = scf.for %arg3 = %c0 to %c8192 step %c128 iter_args(%arg4 = %cst) -> (vector<128x128xf32>) {
|
|
%4 = xegpu.create_nd_tdesc %arg0 : memref<2048x8192xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%5 = xegpu.load_nd %4[%block_id_x, %arg3] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
|
|
%6 = xegpu.create_nd_tdesc %arg1 : memref<8192x4096xf16> -> !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%7 = xegpu.load_nd %6[%arg3, %block_id_y] : !xegpu.tensor_desc<128x128xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<128x128xf16>
|
|
// Couldn not find a layout whose sg_data would be aligned on the reduction dimension.
|
|
// CHECK: xegpu.dpas %{{.*}} {layout_cd = #xegpu.layout<sg_layout = [2, 4], sg_data = [64, 32]>} :
|
|
%8 = xegpu.dpas %5, %7, %arg4 : vector<128x128xf16>, vector<128x128xf16>, vector<128x128xf32> -> vector<128x128xf32>
|
|
scf.yield %8 : vector<128x128xf32>
|
|
}
|
|
%3 = xegpu.create_nd_tdesc %arg2 : memref<2048x4096xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
xegpu.store_nd %2, %3[%0, %1] : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: dpas_fails
|
|
gpu.func @dpas_fails(%arg0: memref<2048x8192xf16>, %arg1: memref<8192x4096xf16>, %arg2: memref<2048x4096xf32>) kernel attributes {known_block_size = array<i32: 8, 1, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<32x64xf32>
|
|
%c16 = arith.constant 16 : index
|
|
%c8192 = arith.constant 8192 : index
|
|
%c0 = arith.constant 0 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%4 = xegpu.create_nd_tdesc %arg0 : memref<2048x8192xf16> -> !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%5 = xegpu.load_nd %4[%block_id_x, %c0] : !xegpu.tensor_desc<32x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<32x16xf16>
|
|
%6 = xegpu.create_nd_tdesc %arg1 : memref<8192x4096xf16> -> !xegpu.tensor_desc<16x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
// CHECK: xegpu.load_nd %{{.*}}[%{{.*}}, %{{.*}}] :
|
|
%7 = xegpu.load_nd %6[%c0, %block_id_y] : !xegpu.tensor_desc<16x64xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<16x64xf16>
|
|
// We have 8 SGs, but currently attempt to use only the largest inst size, so the 32x16 A tile is too small -> fail propagation.
|
|
// CHECK: xegpu.dpas %{{.*}} {layout_cd = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>} :
|
|
%8 = xegpu.dpas %5, %7, %cst : vector<32x16xf16>, vector<16x64xf16>, vector<32x64xf32> -> vector<32x64xf32>
|
|
%3 = xegpu.create_nd_tdesc %arg2 : memref<2048x4096xf32> -> !xegpu.tensor_desc<32x64xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
// CHECK: xegpu.store_nd %{{.*}} <{layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16]>}> :
|
|
xegpu.store_nd %8, %3[%block_id_x, %block_id_y] : vector<32x64xf32>, !xegpu.tensor_desc<32x64xf32, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: store_fails
|
|
gpu.func @store_fails(%arg0: memref<2048x8192xf16>) kernel attributes {known_block_size = array<i32: 8, 1, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<8x16xf16>
|
|
%c0 = arith.constant 0 : index
|
|
%tdesc = xegpu.create_nd_tdesc %arg0 : memref<2048x8192xf16> -> !xegpu.tensor_desc<8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
%loaded = xegpu.load_nd %tdesc[%c0, %c0] : !xegpu.tensor_desc<8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<8x16xf16>
|
|
%loaded_add = arith.addf %loaded, %cst : vector<8x16xf16>
|
|
// 8 subgroups could load 1x16 each, but the current infra only considers the largest inst size (larger than 1x16) -> fail propagation.
|
|
// CHECK: xegpu.store_nd %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] :
|
|
xegpu.store_nd %loaded_add, %tdesc[%c0, %c0] : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @xevm_module{
|
|
// CHECK-LABEL: load_store_matrix
|
|
gpu.func @load_store_matrix(%arg0: !xegpu.mem_desc<1x64x128xf32>, %sg_id_lt_2: i1) {
|
|
%c0 = arith.constant 0 : index
|
|
scf.if %sg_id_lt_2 {
|
|
// CHECK: xegpu.load_matrix %{{.*}} <{layout = #xegpu.layout<sg_layout = [1, 4, 2], sg_data = [1, 8, 16]>}>
|
|
%1 = xegpu.load_matrix %arg0[%c0, %c0, %c0] : !xegpu.mem_desc<1x64x128xf32>, index, index, index -> vector<1x32x32xf32>
|
|
// CHECK: xegpu.store_matrix %{{.*}} <{layout = #xegpu.layout<sg_layout = [1, 4, 2], sg_data = [1, 8, 16]>}>
|
|
xegpu.store_matrix %1, %arg0[%c0, %c0, %c0] <{layout = #xegpu.layout<sg_layout = [1, 4, 2], sg_data = [1, 8, 16]>}> : vector<1x32x32xf32>, !xegpu.mem_desc<1x64x128xf32>, index, index, index
|
|
}
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: gpu.func @convert_layout
|
|
// CHECK-SAME: %[[ARG0:.*]]: !xegpu.mem_desc<32x128xf32>
|
|
gpu.func @convert_layout(%arg0: !xegpu.mem_desc<32x128xf32>) {
|
|
// CHECK: %[[C0:.*]] = arith.constant 0 : index
|
|
%c0 = arith.constant 0 : index
|
|
// CHECK: %[[V0:.*]] = xegpu.load_matrix %[[ARG0]][%[[C0]], %[[C0]]]
|
|
// CHECK-SAME: layout = #xegpu.layout<sg_layout = [4, 8], sg_data = [8, 16]>
|
|
%1 = xegpu.load_matrix %arg0[%c0, %c0] : !xegpu.mem_desc<32x128xf32>, index, index -> vector<32x128xf32>
|
|
%2 = xegpu.convert_layout %1
|
|
<{input_layout = #xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>,
|
|
target_layout = #xegpu.layout<sg_layout=[4, 8], sg_data=[8, 16]>}>
|
|
: vector<32x128xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
gpu.module @test {
|
|
// CHECK-LABEL: dpas_mx
|
|
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>} dense<0.000000e+00> : vector<128x256xbf16>
|
|
// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x512xf8E5M2> -> !xegpu.tensor_desc<128x512xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>>
|
|
// CHECK: %[[T1:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x512xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>> -> vector<128x512xf8E5M2>
|
|
// CHECK: %[[T2:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<512x256xf8E5M2> -> !xegpu.tensor_desc<512x256xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>>
|
|
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T2]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<512x256xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>> -> vector<512x256xf8E5M2>
|
|
// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x16xf8E8M0FNU> -> !xegpu.tensor_desc<128x16xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x16xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<128x16xf8E8M0FNU>
|
|
// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<16x256xf8E8M0FNU> -> !xegpu.tensor_desc<16x256xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
|
|
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<16x256xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>> -> vector<16x256xf8E8M0FNU>
|
|
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
|
|
// CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>} :
|
|
// CHECK-SAME: (vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU>) -> vector<128x256xbf16>
|
|
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x256xbf16> -> !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
|
|
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> : vector<128x256xbf16>, !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
|
|
gpu.func @dpas_mx(%arg0: memref<128x512xf8E5M2>, %arg1: memref<512x256xf8E5M2>, %arg2: memref<128x256xbf16>,
|
|
%arg3: memref<128x16xf8E8M0FNU>, %arg4: memref<16x256xf8E8M0FNU>) kernel attributes
|
|
{known_block_size = array<i32: 1, 64, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<128x256xbf16>
|
|
%tdesc_a = xegpu.create_nd_tdesc %arg0 : memref<128x512xf8E5M2> -> !xegpu.tensor_desc<128x512xf8E5M2>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] : !xegpu.tensor_desc<128x512xf8E5M2> -> vector<128x512xf8E5M2>
|
|
%tdesc_b = xegpu.create_nd_tdesc %arg1 : memref<512x256xf8E5M2> -> !xegpu.tensor_desc<512x256xf8E5M2>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] : !xegpu.tensor_desc<512x256xf8E5M2> -> vector<512x256xf8E5M2>
|
|
%tdesc_a_scale = xegpu.create_nd_tdesc %arg3 : memref<128x16xf8E8M0FNU> -> !xegpu.tensor_desc<128x16xf8E8M0FNU>
|
|
%load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] : !xegpu.tensor_desc<128x16xf8E8M0FNU> -> vector<128x16xf8E8M0FNU>
|
|
%tdesc_b_scale = xegpu.create_nd_tdesc %arg4 : memref<16x256xf8E8M0FNU> -> !xegpu.tensor_desc<16x256xf8E8M0FNU>
|
|
%load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] : !xegpu.tensor_desc<16x256xf8E8M0FNU> -> vector<16x256xf8E8M0FNU>
|
|
%dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : (vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU>) -> vector<128x256xbf16>
|
|
%tdesc_cd = xegpu.create_nd_tdesc %arg2 : memref<128x256xbf16> -> !xegpu.tensor_desc<128x256xbf16>
|
|
xegpu.store_nd %dpas_mx, %tdesc_cd[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> : vector<128x256xbf16>, !xegpu.tensor_desc<128x256xbf16>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
gpu.module @test {
|
|
// CHECK-LABEL: dpas_mx_fp4
|
|
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} dense<0.000000e+00> : vector<128x128xf32>
|
|
// CHECK: %[[T0:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x512xf4E2M1FN> -> !xegpu.tensor_desc<128x512xf4E2M1FN, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>>
|
|
// CHECK: %[[T1:.*]] = xegpu.load_nd %[[T0]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x512xf4E2M1FN, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>> -> vector<128x512xf4E2M1FN>
|
|
// CHECK: %[[T2:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<512x128xf4E2M1FN> -> !xegpu.tensor_desc<512x128xf4E2M1FN, #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>>
|
|
// CHECK: %[[T3:.*]] = xegpu.load_nd %[[T2]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<512x128xf4E2M1FN, #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>> -> vector<512x128xf4E2M1FN>
|
|
// CHECK: %[[T4:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x16xf8E8M0FNU> -> !xegpu.tensor_desc<128x16xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
// CHECK: %[[T5:.*]] = xegpu.load_nd %[[T4]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<128x16xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<128x16xf8E8M0FNU>
|
|
// CHECK: %[[T6:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<16x128xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
// CHECK: %[[T7:.*]] = xegpu.load_nd %[[T6]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<16x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<16x128xf8E8M0FNU>
|
|
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
|
|
// CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} :
|
|
// CHECK-SAME: (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>) -> vector<128x128xf32>
|
|
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
|
|
gpu.func @dpas_mx_fp4(%arg0: memref<128x512xf4E2M1FN>, %arg1: memref<512x128xf4E2M1FN>, %arg2: memref<128x128xf32>,
|
|
%arg3: memref<128x16xf8E8M0FNU>, %arg4: memref<16x128xf8E8M0FNU>) kernel attributes
|
|
{known_block_size = array<i32: 1, 64, 16>} {
|
|
%cst = arith.constant dense<0.000000e+00> : vector<128x128xf32>
|
|
%tdesc_a = xegpu.create_nd_tdesc %arg0 : memref<128x512xf4E2M1FN> -> !xegpu.tensor_desc<128x512xf4E2M1FN>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] : !xegpu.tensor_desc<128x512xf4E2M1FN> -> vector<128x512xf4E2M1FN>
|
|
%tdesc_b = xegpu.create_nd_tdesc %arg1 : memref<512x128xf4E2M1FN> -> !xegpu.tensor_desc<512x128xf4E2M1FN>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] : !xegpu.tensor_desc<512x128xf4E2M1FN> -> vector<512x128xf4E2M1FN>
|
|
%tdesc_a_scale = xegpu.create_nd_tdesc %arg3 : memref<128x16xf8E8M0FNU> -> !xegpu.tensor_desc<128x16xf8E8M0FNU>
|
|
%load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] : !xegpu.tensor_desc<128x16xf8E8M0FNU> -> vector<128x16xf8E8M0FNU>
|
|
%tdesc_b_scale = xegpu.create_nd_tdesc %arg4 : memref<16x128xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU>
|
|
%load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] : !xegpu.tensor_desc<16x128xf8E8M0FNU> -> vector<16x128xf8E8M0FNU>
|
|
%dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>) -> vector<128x128xf32>
|
|
%tdesc_cd = xegpu.create_nd_tdesc %arg2 : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32>
|
|
xegpu.store_nd %dpas_mx, %tdesc_cd[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test {
|
|
// CHECK-LABEL: sum_reduction_broadcast_transpose
|
|
// CHECK-SAME: %[[ARG0:.*]]: memref<64x256xf16>
|
|
// CHECK-SAME: %[[ARG1:.*]]: memref<64x32xf16>
|
|
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>}
|
|
// CHECK-SAME: dense<0.000000e+00> : vector<64xf16>
|
|
// CHECK: %[[TDESC0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<64x256xf16> -> !xegpu.tensor_desc<64x32xf16,
|
|
// CHECK-SAME: #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>>
|
|
// CHECK: %[[FOR:.*]] = scf.for %[[ARG2:.*]] = %{{.*}} iter_args(%[[ARG3:.*]] = %[[CST]]) -> (vector<64xf16>) {
|
|
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC0]][0, %[[ARG2]]]
|
|
// CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>}> :
|
|
// CHECK-SAME: !xegpu.tensor_desc<64x32xf16, #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>> -> vector<64x32xf16>
|
|
// CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[ARG3]]
|
|
// CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>} [1]
|
|
// CHECK-SAME: : vector<64x32xf16> to vector<64xf16>
|
|
// CHECK: scf.yield %[[REDUCE]] : vector<64xf16>
|
|
// CHECK: } {layout_operand_3 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>, layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>}
|
|
// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FOR]] {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], order = [0, 1]>}
|
|
// CHECK-SAME: : vector<64xf16> to vector<32x64xf16>
|
|
// CHECK: %[[TRANS:.*]] = vector.transpose %[[BCAST]], [1, 0] {layout_result_0 = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>}
|
|
// CHECK-SAME: : vector<32x64xf16> to vector<64x32xf16>
|
|
// CHECK: %[[TDESC1:.*]] = xegpu.create_nd_tdesc %[[ARG1]] : memref<64x32xf16> ->
|
|
// CHECK-SAME: !xegpu.tensor_desc<64x32xf16, #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>>
|
|
// CHECK: xegpu.store_nd %[[TRANS]], %[[TDESC1]][0, 0] <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>}>
|
|
// CHECK-SAME: : vector<64x32xf16>, !xegpu.tensor_desc<64x32xf16, #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>>
|
|
gpu.func @sum_reduction_broadcast_transpose(%input: memref<64x256xf16>, %output: memref<64x32xf16>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
|
|
%c0 = arith.constant 0 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c256 = arith.constant 256 : index
|
|
%init = arith.constant dense<0.000000e+00> : vector<64xf16>
|
|
%tdesc_input = xegpu.create_nd_tdesc %input : memref<64x256xf16> -> !xegpu.tensor_desc<64x32xf16>
|
|
%result = scf.for %i = %c0 to %c256 step %c32 iter_args(%acc = %init) -> (vector<64xf16>) {
|
|
%loaded = xegpu.load_nd %tdesc_input[0, %i] : !xegpu.tensor_desc<64x32xf16> -> vector<64x32xf16>
|
|
%reduced = vector.multi_reduction <add>, %loaded, %acc [1] : vector<64x32xf16> to vector<64xf16>
|
|
scf.yield %reduced : vector<64xf16>
|
|
}
|
|
%broadcasted = vector.broadcast %result : vector<64xf16> to vector<32x64xf16>
|
|
%transposed = vector.transpose %broadcasted, [1, 0] : vector<32x64xf16> to vector<64x32xf16>
|
|
%tdesc_output = xegpu.create_nd_tdesc %output : memref<64x32xf16> -> !xegpu.tensor_desc<64x32xf16>
|
|
xegpu.store_nd %transposed, %tdesc_output[0, 0] { layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>} : vector<64x32xf16>, !xegpu.tensor_desc<64x32xf16>
|
|
gpu.return
|
|
}
|
|
}
|