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
749 lines
43 KiB
MLIR
749 lines
43 KiB
MLIR
// RUN: mlir-opt --xegpu-blocking -split-input-file %s | FileCheck %s
|
|
|
|
#a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
#b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [16, 1]>
|
|
#c = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @gemm_with_one_to_n_lowering(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c16 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #c>
|
|
%c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #c}: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #a>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #b>
|
|
%out = scf.for %k = %c0 to %c1024 step %c32
|
|
iter_args(%arg2 = %c_init)
|
|
-> (vector<16x32xf32>) {
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #a}: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #b}: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
|
|
//CHECK-COUNT-8: xegpu.dpas {{.*}}
|
|
%c = xegpu.dpas %a, %b, %arg2 {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
|
|
scf.yield %c
|
|
: vector<16x32xf32>
|
|
}
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
|
|
xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #c}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l1 = #xegpu.layout<inst_data = [8, 16]>
|
|
#l2 = #xegpu.layout<inst_data = [16, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @gemm_with_inst_data_only_attribute(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c16 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #l1>
|
|
%c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l1}: !xegpu.tensor_desc<16x32xf32, #l1> -> vector<16x32xf32>
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l1>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #l2>
|
|
%out = scf.for %k = %c0 to %c1024 step %c32
|
|
iter_args(%arg2 = %c_init)
|
|
-> (vector<16x32xf32>) {
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l1}: !xegpu.tensor_desc<16x32xf16, #l1> -> vector<16x32xf16>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #l2}: !xegpu.tensor_desc<32x32xf16, #l2> -> vector<32x32xf16>
|
|
//CHECK-COUNT-8: xegpu.dpas {{.*}}
|
|
%c = xegpu.dpas %a, %b, %arg2 {layout_a=#l1, layout_b = #l2, layout_cd = #l1}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
|
|
scf.yield %c
|
|
: vector<16x32xf32>
|
|
}
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
|
|
xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l1}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l1 = #xegpu.layout<inst_data = [8, 16]>
|
|
#l2 = #xegpu.layout<inst_data = [16, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @gemm_with_one_to_one_lowering(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c8 = arith.constant 8 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c8 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<8x32xf32, #l1>
|
|
|
|
//CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf32> -> vector<8x16xf32>
|
|
%c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l1}: !xegpu.tensor_desc<8x32xf32, #l1> -> vector<8x32xf32>
|
|
|
|
%c2 = arith.constant 2 : index
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<8x16xf16, #l1>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l2>
|
|
%out = scf.for %k = %c0 to %c1024 step %c16
|
|
iter_args(%arg2 = %c_init)
|
|
-> (vector<8x32xf32>) {
|
|
%a_off = arith.muli %k, %c2 : index
|
|
//CHECK: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %a_off] {layout = #l1}: !xegpu.tensor_desc<8x16xf16, #l1> -> vector<8x16xf16>
|
|
//CHECK-COUNT-2: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%a_off, %c0] {layout = #l2}: !xegpu.tensor_desc<16x32xf16, #l2> -> vector<16x32xf16>
|
|
%c = xegpu.dpas %a, %b, %arg2 {layout_a=#l1, layout_b = #l2, layout_cd = #l1}: vector<8x16xf16>, vector<16x32xf16>, vector<8x32xf32> -> vector<8x32xf32>
|
|
scf.yield %c
|
|
: vector<8x32xf32>
|
|
}
|
|
//CHECK-COUNT-2: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
|
|
xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l1}: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #l1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
#b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [16, 1]>
|
|
#c = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @gemm_with_elemwise_preop(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c16 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #c>
|
|
%c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #c}: !xegpu.tensor_desc<16x32xf32, #c> -> vector<16x32xf32>
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #a>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32x32xf16, #b>
|
|
%out = scf.for %k = %c0 to %c1024 step %c32
|
|
iter_args(%arg2 = %c_init)
|
|
-> (vector<16x32xf32>) {
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x16xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #a}: !xegpu.tensor_desc<16x32xf16, #a> -> vector<16x32xf16>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<16x16xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #b}: !xegpu.tensor_desc<32x32xf16, #b> -> vector<32x32xf16>
|
|
//CHECK-COUNT-4: math.exp {{.*}} : vector<8x16xf16>
|
|
%e = math.exp %a : vector<16x32xf16>
|
|
//CHECK-COUNT-8: xegpu.dpas {{.*}}
|
|
%c = xegpu.dpas %e, %b, %arg2 {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x32xf16>, vector<32x32xf16>, vector<16x32xf32> -> vector<16x32xf32>
|
|
scf.yield %c
|
|
: vector<16x32xf32>
|
|
}
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
|
|
xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #c}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #c>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [8, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @elementwise_with_inst_data_only(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf16>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c32 : index
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l>
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<16x32xf16, #l>
|
|
|
|
scf.for %k = %c0 to %c1024 step %c32 {
|
|
//CHECK-COUNT-8: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8x16xf16> -> vector<8x16xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l}: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%c0, %k] {layout = #l}: !xegpu.tensor_desc<16x32xf16, #l> -> vector<16x32xf16>
|
|
|
|
//CHECK-COUNT-4: arith.addf {{.*}} : vector<8x16xf16>
|
|
%c = arith.addf %a, %b : vector<16x32xf16>
|
|
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf16>, !xegpu.tensor_desc<8x16xf16>
|
|
xegpu.store_nd %c, %c_tdesc[%c0, %k] {layout = #l}: vector<16x32xf16>, !xegpu.tensor_desc<16x32xf16, #l>
|
|
}
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [8]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @elementwise_1D(%A: memref<1024x1024xf16>, %B: memref<1024x1024xf16>, %C: memref<1024x1024xf16>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c32 : index
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf16> -> !xegpu.tensor_desc<32xf16, #l>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf16> -> !xegpu.tensor_desc<32xf16, #l>
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf16> -> !xegpu.tensor_desc<32xf16, #l>
|
|
|
|
scf.for %k = %c0 to %c1024 step %c32 {
|
|
//CHECK-COUNT-8: xegpu.load_nd {{.*}} : !xegpu.tensor_desc<8xf16> -> vector<8xf16>
|
|
%a = xegpu.load_nd %a_tdesc[%k] {layout = #l}: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
|
|
%b = xegpu.load_nd %b_tdesc[%k] {layout = #l}: !xegpu.tensor_desc<32xf16, #l> -> vector<32xf16>
|
|
|
|
//CHECK-COUNT-4: arith.addf {{.*}} : vector<8xf16>
|
|
%c = arith.addf %a, %b : vector<32xf16>
|
|
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8xf16>, !xegpu.tensor_desc<8xf16>
|
|
xegpu.store_nd %c, %c_tdesc[%k] {layout = #l}: vector<32xf16>, !xegpu.tensor_desc<32xf16, #l>
|
|
}
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16, 16]>
|
|
#r = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [0]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @reduce_dim_0(%a: memref<16x512xf32>, %b: memref<512xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
|
|
%acc = arith.constant dense<0.0> : vector<64xf32>
|
|
%c64 = arith.constant 64 : index
|
|
%block_id_x = gpu.block_id x
|
|
%m = arith.muli %block_id_x, %c64 : index
|
|
%0 = xegpu.create_nd_tdesc %a : memref<16x512xf32> -> !xegpu.tensor_desc<16x64xf32, #l>
|
|
%1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<16x64xf32, #l> -> vector<16x64xf32>
|
|
// CHECK: vector.multi_reduction <add>, {{.*}}, [[ACC:%[0-9A-Za-z]+]] [0] : vector<16x16xf32> to vector<16xf32>
|
|
// CHECK-COUNT-3: vector.multi_reduction <add>, {{.*}}, [[ACC]] [0] : vector<16x16xf32> to vector<16xf32>
|
|
%2 = vector.multi_reduction <add>, %1, %acc [0]: vector<16x64xf32> to vector<64xf32>
|
|
%3 = xegpu.create_nd_tdesc %b : memref<512xf32> -> !xegpu.tensor_desc<64xf32, #r>
|
|
xegpu.store_nd %2, %3[0] {layout = #r}: vector<64xf32>, !xegpu.tensor_desc<64xf32, #r>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16, 16]>
|
|
#r = #xegpu.slice<#xegpu.layout<inst_data = [16, 16]>, dims = [1]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @reduce_dim_1(%a: memref<512x32xf32>, %b: memref<512xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
|
|
%c1 = arith.constant 1 : index
|
|
%c32 = arith.constant 32 : index
|
|
%acc = arith.constant dense<0.0> : vector<32xf32>
|
|
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
|
|
%m = arith.muli %block_id_x, %c32 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
%0 = xegpu.create_nd_tdesc %a : memref<512x32xf32> -> !xegpu.tensor_desc<32x128xf32, #l>
|
|
%1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<32x128xf32, #l> -> vector<32x128xf32>
|
|
|
|
// CHECK-COUNT-7: arith.addf {{.*}} : vector<16x16xf32>
|
|
// CHECK: vector.multi_reduction <add>, {{.*}} [1] : vector<16x16xf32> to vector<16xf32>
|
|
// CHECK-COUNT-7: arith.addf {{.*}} : vector<16x16xf32>
|
|
// CHECK: vector.multi_reduction <add>, {{.*}} [1] : vector<16x16xf32> to vector<16xf32>
|
|
|
|
%2 = vector.multi_reduction <add>, %1, %acc [1]: vector<32x128xf32> to vector<32xf32>
|
|
%3 = xegpu.create_nd_tdesc %b : memref<512xf32> -> !xegpu.tensor_desc<32xf32, #r>
|
|
xegpu.store_nd %2, %3[0] {layout = #r}: vector<32xf32>, !xegpu.tensor_desc<32xf32, #r>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#r = #xegpu.layout<inst_data = [16]>
|
|
#l = #xegpu.layout<inst_data = [16, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @broadcast_dim_0(%a: memref<512xf32>, %b: memref<16x512xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
|
|
|
|
%c64 = arith.constant 64 : index
|
|
%block_id_x = gpu.block_id x
|
|
%m = arith.muli %block_id_x, %c64 : index
|
|
%0 = xegpu.create_nd_tdesc %a : memref<512xf32> -> !xegpu.tensor_desc<64xf32, #r>
|
|
%1 = xegpu.load_nd %0[0] {layout = #r}: !xegpu.tensor_desc<64xf32, #r> -> vector<64xf32>
|
|
// CHECK-COUNT-4: vector.broadcast {{.*}} : vector<16xf32> to vector<16x16xf32>
|
|
%2 = vector.broadcast %1 : vector<64xf32> to vector<16x64xf32>
|
|
%3 = xegpu.create_nd_tdesc %b : memref<16x512xf32> -> !xegpu.tensor_desc<16x64xf32, #l>
|
|
xegpu.store_nd %2, %3[0, 0] {layout = #l}: vector<16x64xf32>, !xegpu.tensor_desc<16x64xf32, #l>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#r = #xegpu.layout<inst_data = [16]>
|
|
#l = #xegpu.layout<inst_data = [16, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @broadcast_dim_1(%a: memref<512xf32>, %b: memref<16x512xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
|
|
|
|
%c32 = arith.constant 32 : index
|
|
%block_id_x = gpu.block_id x
|
|
%m = arith.muli %block_id_x, %c32 : index
|
|
%0 = xegpu.create_nd_tdesc %a : memref<512xf32> -> !xegpu.tensor_desc<32xf32, #r>
|
|
%1 = xegpu.load_nd %0[0] {layout = #r}: !xegpu.tensor_desc<32xf32, #r> -> vector<32xf32>
|
|
%11 = vector.shape_cast %1 : vector<32xf32> to vector<32x1xf32>
|
|
// CHECK-COUNT-8: vector.broadcast {{.*}}: vector<16x1xf32> to vector<16x16xf32>
|
|
%2 = vector.broadcast %11 : vector<32x1xf32> to vector<32x64xf32>
|
|
%3 = xegpu.create_nd_tdesc %b : memref<16x512xf32> -> !xegpu.tensor_desc<32x64xf32, #l>
|
|
xegpu.store_nd %2, %3[0, 0] {layout = #l} : vector<32x64xf32>, !xegpu.tensor_desc<32x64xf32, #l>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16, 8]>
|
|
#t = #xegpu.layout<inst_data = [8, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @transpose(%a: memref<512x8xf32>, %b: memref<8x512xf32>) kernel attributes {VectorComputeFunctionINTEL, spirv.entry_point_abi = #spirv.entry_point_abi<>} {
|
|
|
|
%c32 = arith.constant 32 : index
|
|
%block_id_x = gpu.block_id x
|
|
%m = arith.muli %block_id_x, %c32 : index
|
|
%0 = xegpu.create_nd_tdesc %a : memref<512x8xf32> -> !xegpu.tensor_desc<32x8xf32, #l>
|
|
%1 = xegpu.load_nd %0[0, 0] {layout = #l}: !xegpu.tensor_desc<32x8xf32, #l> -> vector<32x8xf32>
|
|
// CHECK-COUNT-2: vector.transpose {{.*}} [1, 0] : vector<16x8xf32> to vector<8x16xf32>
|
|
%2 = vector.transpose %1, [1, 0] : vector<32x8xf32> to vector<8x32xf32>
|
|
%3 = xegpu.create_nd_tdesc %b : memref<8x512xf32> -> !xegpu.tensor_desc<8x32xf32, #t>
|
|
xegpu.store_nd %2, %3[0, 0] {layout = #t}: vector<8x32xf32>, !xegpu.tensor_desc<8x32xf32, #t>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @test_vector_constant_mask(%src: ui64, %dst: ui64) {
|
|
//CHECK: arith.constant dense<true> : vector<16xi1>
|
|
%mask = vector.constant_mask [32] : vector<32xi1>
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
%ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
|
|
xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @test_vector_create_mask(%src: ui64, %dst: ui64) {
|
|
%c16 = arith.constant 16 : index
|
|
//CHECK-COUNT-2: vector.create_mask {{.*}} : vector<16xi1>
|
|
%mask = vector.create_mask %c16 : vector<32xi1>
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
%ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
|
|
xegpu.store %ld, %dst[%cst], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @test_vector_step(%src: ui64, %dst: ui64) {
|
|
%c16 = arith.constant 16 : index
|
|
//CHECK: [[cst:%.+]] = arith.constant dense<16> : vector<16xindex>
|
|
//CHECK: [[step:%.+]] = vector.step : vector<16xindex>
|
|
//CHECK: arith.addi [[step]], [[cst]] : vector<16xindex>
|
|
%step = vector.step : vector<32xindex>
|
|
%mask = vector.create_mask %c16 : vector<32xi1>
|
|
%ld = xegpu.load %src[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
|
|
xegpu.store %ld, %dst[%step], %mask {chunk_size = 1, layout = #l, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#a = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
#b = #xegpu.layout<inst_data = [16, 16], lane_layout = [1, 16], lane_data = [16, 1]>
|
|
#c = #xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>
|
|
|
|
gpu.module @test_kernel {
|
|
//CHECK-LABEL: gpu.func @convert_layout
|
|
//CHECK-SAME: [[arg0:%.+]]: memref<16x16xf16>, [[arg1:%.+]]: memref<16x16xf16>, [[arg2:%.+]]: memref<16x16xf32>
|
|
//CHECK: [[c8:%.+]] = arith.constant 8 : index
|
|
//CHECK: [[c0:%.+]] = arith.constant 0 : index
|
|
//CHECK: [[a:%.+]] = xegpu.create_nd_tdesc [[arg0]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>>
|
|
//CHECK: [[b:%.+]] = xegpu.create_nd_tdesc [[arg1]] : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>>
|
|
//CHECK: [[load_a:%.+]] = xegpu.load_nd [[a]][[[c0]], [[c0]]] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>> -> vector<16x16xf16>
|
|
//CHECK: [[load_b:%.+]] = xegpu.load_nd [[b]][[[c0]], [[c0]]] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>}> : !xegpu.tensor_desc<16x16xf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>> -> vector<16x16xf16>
|
|
//CHECK: [[cvt:%.+]] = xegpu.convert_layout [[load_a]] <{input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>, target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> : vector<16x16xf16>
|
|
//CHECK: [[a0:%.+]] = vector.extract_strided_slice [[cvt]] {offsets = [0, 0], sizes = [8, 16], strides = [1, 1]} : vector<16x16xf16> to vector<8x16xf16>
|
|
//CHECK: [[a1:%.+]] = vector.extract_strided_slice [[cvt]] {offsets = [8, 0], sizes = [8, 16], strides = [1, 1]} : vector<16x16xf16> to vector<8x16xf16>
|
|
//CHECK: [[dpas0:%.+]] = xegpu.dpas [[a0]], [[load_b]]
|
|
//CHECK: [[dpas1:%.+]] = xegpu.dpas [[a1]], [[load_b]]
|
|
//CHECK: [[c_tdesc:%.+]] = xegpu.create_nd_tdesc [[arg2]] : memref<16x16xf32> -> !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
|
|
//CHECK: xegpu.store_nd [[dpas0]], [[c_tdesc]][[[c0]], [[c0]]] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
|
|
//CHECK: xegpu.store_nd [[dpas1]], [[c_tdesc]][[[c8]], [[c0]]] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>}> : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>>
|
|
|
|
gpu.func @convert_layout(%A: memref<16x16xf16>, %B: memref<16x16xf16>, %C: memref<16x16xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #b>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #b>
|
|
%a = xegpu.load_nd %a_tdesc[0, 0] {layout = #b}: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
|
|
%b = xegpu.load_nd %b_tdesc[0, 0] {layout = #b}: !xegpu.tensor_desc<16x16xf16, #b> -> vector<16x16xf16>
|
|
%a1 = xegpu.convert_layout %a <{input_layout = #b, target_layout = #a}> : vector<16x16xf16>
|
|
%c = xegpu.dpas %a1, %b {layout_a=#a, layout_b = #b, layout_cd = #c}: vector<16x16xf16>, vector<16x16xf16> -> vector<16x16xf32>
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32, #c>
|
|
xegpu.store_nd %c, %c_tdesc[0, 0] {layout = #c}: vector<16x16xf32>, !xegpu.tensor_desc<16x16xf32, #c>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
#a = #xegpu.layout<inst_data = [8, 16]>
|
|
gpu.module @test_kernel {
|
|
//CHECK-LABEL: gpu.func @convert_layout_scalar
|
|
// CHECK-NOT: xegpu.convert_layout
|
|
gpu.func @convert_layout_scalar(%arg0: memref<16x16xf16>, %arg1: memref<4xf16>) {
|
|
%acc = arith.constant 0.000000e+00 : f16
|
|
%c0 = arith.constant 0 : index
|
|
%a_tdesc = xegpu.create_nd_tdesc %arg0 : memref<16x16xf16> -> !xegpu.tensor_desc<16x16xf16, #a>
|
|
%a = xegpu.load_nd %a_tdesc[0, 0] {layout = #a}: !xegpu.tensor_desc<16x16xf16, #a> -> vector<16x16xf16>
|
|
%a_reduce = vector.multi_reduction <add>, %a, %acc [0, 1] : vector<16x16xf16> to f16
|
|
%13 = xegpu.convert_layout %a_reduce <{input_layout = #xegpu.slice<#a, dims = [0, 1]>, target_layout = #xegpu.slice<#a, dims = [0, 1]>}> : f16
|
|
memref.store %13, %arg1[%c0] : memref<4xf16>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
#in = #xegpu.slice<#xegpu.layout<inst_data = [1, 16]>, dims = [1]>
|
|
#out = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [1]>
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: gpu.func @convert_layout_drop_inst_data_to_null
|
|
// CHECK-NOT: xegpu.convert_layout
|
|
gpu.func @convert_layout_drop_inst_data_to_null(%arg0: vector<2xf32>) -> vector<2xf32> {
|
|
%0 = xegpu.convert_layout %arg0 <{input_layout = #in, target_layout = #out}> : vector<2xf32>
|
|
gpu.return %0 : vector<2xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: gpu.func @convert_layout_drop_slice_inst_data_to_null
|
|
// CHECK-NOT: xegpu.convert_layout
|
|
gpu.func @convert_layout_drop_slice_inst_data_to_null(%arg0: vector<1xf32>) -> vector<1xf32> {
|
|
%0 = xegpu.convert_layout %arg0 <{input_layout = #xegpu.layout<inst_data = [1]>, target_layout = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16]>, dims = [1, 2]>}> : vector<1xf32>
|
|
gpu.return %0 : vector<1xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
|
|
#lb = #xegpu.layout<inst_data = [4, 32, 2], lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>
|
|
#b = #xegpu.layout<inst_data = [4, 16, 2], lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>
|
|
|
|
gpu.module @test_kernel {
|
|
//CHECK: gpu.func @convert_layout([[arg0:%.+]]: vector<8x32x2xf16>) -> vector<8x32x2xf16> {
|
|
//CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<8x32x2xf16>
|
|
//CHECK: [[e0:%.+]] = vector.extract_strided_slice [[arg0]] {offsets = [0, 0, 0], sizes = [4, 32, 2], strides = [1, 1, 1]} : vector<8x32x2xf16> to vector<4x32x2xf16>
|
|
//CHECK: [[e1:%.+]] = vector.extract_strided_slice [[arg0]] {offsets = [4, 0, 0], sizes = [4, 32, 2], strides = [1, 1, 1]} : vector<8x32x2xf16> to vector<4x32x2xf16>
|
|
//CHECK: [[c0:%.+]] = xegpu.convert_layout [[e0]] <{input_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>, target_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>}> : vector<4x32x2xf16>
|
|
//CHECK: [[c1:%.+]] = xegpu.convert_layout [[e1]] <{input_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 2]>, target_layout = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>}> : vector<4x32x2xf16>
|
|
//CHECK: [[e2:%.+]] = vector.extract_strided_slice [[c0]] {offsets = [0, 0, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
|
|
//CHECK: [[m0:%.+]] = math.exp [[e2]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
|
|
//CHECK: [[i0:%.+]] = vector.insert_strided_slice [[m0]], [[cst]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
|
|
//CHECK: [[e3:%.+]] = vector.extract_strided_slice [[c0]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
|
|
//CHECK: [[m1:%.+]] = math.exp [[e3]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
|
|
//CHECK: [[i1:%.+]] = vector.insert_strided_slice [[m1]], [[i0]] {offsets = [0, 16, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
|
|
//CHECK: [[e4:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 0, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
|
|
//CHECK: [[m2:%.+]] = math.exp [[e4]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
|
|
//CHECK: [[i2:%.+]] = vector.insert_strided_slice [[m2]], [[i1]] {offsets = [4, 0, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
|
|
//CHECK: [[e5:%.+]] = vector.extract_strided_slice [[c1]] {offsets = [0, 16, 0], sizes = [4, 16, 2], strides = [1, 1, 1]} : vector<4x32x2xf16> to vector<4x16x2xf16>
|
|
//CHECK: [[m3:%.+]] = math.exp [[e5]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [4, 1, 1]>} : vector<4x16x2xf16>
|
|
//CHECK: [[i3:%.+]] = vector.insert_strided_slice [[m3]], [[i2]] {offsets = [4, 16, 0], strides = [1, 1, 1]} : vector<4x16x2xf16> into vector<8x32x2xf16>
|
|
//CHECK: gpu.return [[i3]] : vector<8x32x2xf16>
|
|
|
|
gpu.func @convert_layout(%B: vector<8x32x2xf16>) -> vector<8x32x2xf16> {
|
|
%b = xegpu.convert_layout %B <{input_layout = #lb, target_layout = #b}> : vector<8x32x2xf16>
|
|
%e = math.exp %b : vector<8x32x2xf16>
|
|
%anchor = xegpu.convert_layout %e <{input_layout = #b, target_layout = #b}> : vector<8x32x2xf16>
|
|
gpu.return %e : vector<8x32x2xf16>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
//CHECK-LABEL: unroll_load_matrix
|
|
gpu.func @unroll_load_matrix(%arg0: memref<4096xi8, 3>) -> vector<32x32xf32> {
|
|
%0 = xegpu.create_mem_desc %arg0 : memref<4096xi8, 3> -> !xegpu.mem_desc<32x32xf32>
|
|
//CHECK-COUNT-8: xegpu.load_matrix {{.*}} : !xegpu.mem_desc<32x32xf32>, index, index -> vector<8x16xf32>
|
|
//CHECK-COUNT-8: vector.insert_strided_slice {{.*}} : vector<8x16xf32> into vector<32x32xf32>
|
|
%1 = xegpu.load_matrix %0[0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}>: !xegpu.mem_desc<32x32xf32> -> vector<32x32xf32>
|
|
gpu.return %1: vector<32x32xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: unroll_store_matrix
|
|
gpu.func @unroll_store_matrix(%value: vector<1x32xf32>, %arg0 : memref<32768xi8, 3>) {
|
|
%mdesc = xegpu.create_mem_desc %arg0 : memref<32768xi8, 3> -> !xegpu.mem_desc<64x128xf32>
|
|
// CHECK-COUNT-2: xegpu.store_matrix {{.*}} : vector<1x16xf32>, !xegpu.mem_desc<64x128xf32>, index, index
|
|
xegpu.store_matrix %value, %mdesc[0, 0] {layout = #xegpu.layout<inst_data = [1, 16]>} : vector<1x32xf32>, !xegpu.mem_desc<64x128xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: load_with_offsets
|
|
// CHECK-COUNT-2: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16xf32>
|
|
gpu.func @load_with_offsets(%src: ui64) -> vector<32xf32> {
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
|
|
%c17 = arith.constant 17: index
|
|
%mask = vector.create_mask %c17 : vector<32xi1>
|
|
%ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32xf32>
|
|
|
|
gpu.return %ld : vector<32xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: store_with_offsets
|
|
// CHECK-COUNT-2: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16xf32>, ui64, vector<16xindex>, vector<16xi1>
|
|
gpu.func @store_with_offsets(%src: ui64) {
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
|
|
%c17 = arith.constant 17: index
|
|
%mask = vector.create_mask %c17 : vector<32xi1>
|
|
|
|
%st_vec = arith.constant dense<1023.0>: vector<32xf32>
|
|
xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [16]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32xf32>, ui64, vector<32xindex>, vector<32xi1>
|
|
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: load_with_offsets_chunk
|
|
// CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<32x4xf32>
|
|
// CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
|
|
// CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
|
|
// CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
|
|
// CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
|
|
// CHECK-COUNT-4: xegpu.load {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<16xindex>, vector<16xi1> -> vector<16x2xf32>
|
|
gpu.func @load_with_offsets_chunk(%src: ui64) -> vector<32x4xf32> {
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
|
|
%c17 = arith.constant 17: index
|
|
%mask = vector.create_mask %c17 : vector<32xi1>
|
|
%ld = xegpu.load %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<32xindex>, vector<32xi1> -> vector<32x4xf32>
|
|
gpu.return %ld : vector<32x4xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: store_with_offsets_chunk
|
|
// CHECK: [[cst:%.+]] = arith.constant dense<1.023000e+03> : vector<16x2xf32
|
|
// CHECK: [[cst0:%.+]] = arith.constant dense<[130, 138, 146, 154, 162, 170, 178, 186, 194, 202, 210, 218, 226, 234, 242, 250]> : vector<16xindex>
|
|
// CHECK: [[cst1:%.+]] = arith.constant dense<[2, 10, 18, 26, 34, 42, 50, 58, 66, 74, 82, 90, 98, 106, 114, 122]> : vector<16xindex>
|
|
// CHECK: [[cst2:%.+]] = arith.constant dense<[128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]> : vector<16xindex>
|
|
// CHECK: [[cst3:%.+]] = arith.constant dense<[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]> : vector<16xindex>
|
|
// CHECK-COUNT-4: xegpu.store {{.*}}[{{.*}}], {{.*}} <{chunk_size = 2 : i64, l1_hint = #xegpu.cache_hint<cached>}> : vector<16x2xf32>, ui64, vector<16xindex>, vector<16xi1>
|
|
gpu.func @store_with_offsets_chunk(%src: ui64) {
|
|
%cst = arith.constant dense<[
|
|
0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248
|
|
]> : vector<32xindex>
|
|
|
|
%c17 = arith.constant 17: index
|
|
%mask = vector.create_mask %c17 : vector<32xi1>
|
|
|
|
%st_vec = arith.constant dense<1023.>: vector<32x4xf32>
|
|
xegpu.store %st_vec, %src[%cst], %mask {chunk_size = 4, layout = #xegpu.layout<inst_data = [16, 2]>, l1_hint = #xegpu.cache_hint<cached>} : vector<32x4xf32>, ui64, vector<32xindex>, vector<32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: preserve_unit_dim_of_load_inst_data
|
|
// CHECK-SAME: [[arg0:%.+]]: ui64
|
|
// CHECK: [[cst:%.+]] = arith.constant dense<0.000000e+00> : vector<1x1x32xf32>
|
|
// CHECK: [[cst_0:%.+]] = arith.constant dense<true> : vector<1x1x16xi1>
|
|
// CHECK: [[cst_1:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
|
|
// CHECK: [[cst_2:%.+]] = arith.constant dense<{{.*}}> : vector<1x1x16xindex>
|
|
// CHECK: [[ld_0:%.+]] = xegpu.load [[arg0]][[[cst_1]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[ld_1:%.+]] = xegpu.load [[arg0]][[[cst_2]]], [[cst_0]] <{chunk_size = 1 : i64, l1_hint = #xegpu.cache_hint<cached>}> : ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[ins_0:%.+]] = vector.insert_strided_slice [[ld_0]], [[cst]] {offsets = [0, 0, 0], strides = [1, 1, 1]} : vector<1x1x16xf32> into vector<1x1x32xf32>
|
|
// CHECK: [[ins_1:%.+]] = vector.insert_strided_slice [[ld_1]], [[ins_0]] {offsets = [0, 0, 16], strides = [1, 1, 1]} : vector<1x1x16xf32> into vector<1x1x32xf32>
|
|
gpu.func @preserve_unit_dim_of_load_inst_data(%src: ui64) -> vector<1x1x32xf32> {
|
|
%cst = arith.constant dense<[[
|
|
[0, 8, 16, 24, 32, 40, 48, 56,
|
|
64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184,
|
|
192, 200, 208, 216, 224, 232, 240, 248]
|
|
]]> : vector<1x1x32xindex>
|
|
|
|
%mask = arith.constant dense<true> : vector<1x1x32xi1>
|
|
%ld = xegpu.load %src[%cst], %mask {chunk_size = 1, layout = #xegpu.layout<inst_data = [1, 1, 16]>, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
|
|
|
|
gpu.return %ld : vector<1x1x32xf32>
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l = #xegpu.layout<inst_data = [1, 16]>
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: load_store_nd_with_offsets
|
|
// CHECK-SAME: [[arg0:%.+]]: memref<1024x1024xf32>, [[arg1:%.+]]: memref<1024x1024xf32>, [[arg2:%.+]]: memref<1024x1024xf32>
|
|
// CHECK-DAG: [[c16:%.+]] = arith.constant 16 : index
|
|
// CHECK-DAG: [[c0:%.+]] = arith.constant 0 : index
|
|
// CHECK: [[tdesc_a:%.+]] = xegpu.create_nd_tdesc [[arg0]] : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x16xf32>
|
|
// CHECK: [[tdesc_b:%.+]] = xegpu.create_nd_tdesc [[arg1]] : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x16xf32>
|
|
// CHECK: [[tdesc_c:%.+]] = xegpu.create_nd_tdesc [[arg2]] : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x16xf32>
|
|
// CHECK: [[ld_a0:%.+]] = xegpu.load_nd [[tdesc_a]][[[c0]], [[c0]]] : !xegpu.tensor_desc<1x16xf32> -> vector<1x16xf32>
|
|
// CHECK: [[ld_a1:%.+]] = xegpu.load_nd [[tdesc_a]][[[c0]], [[c16]]] : !xegpu.tensor_desc<1x16xf32> -> vector<1x16xf32>
|
|
// CHECK: [[ld_b0:%.+]] = xegpu.load_nd [[tdesc_b]][[[c0]], [[c0]]] : !xegpu.tensor_desc<1x16xf32> -> vector<1x16xf32>
|
|
// CHECK: [[ld_b1:%.+]] = xegpu.load_nd [[tdesc_b]][[[c0]], [[c16]]] : !xegpu.tensor_desc<1x16xf32> -> vector<1x16xf32>
|
|
// CHECK: [[add0:%.+]] = arith.addf [[ld_a0]], [[ld_b0]] : vector<1x16xf32>
|
|
// CHECK: [[add1:%.+]] = arith.addf [[ld_a1]], [[ld_b1]] : vector<1x16xf32>
|
|
// CHECK: xegpu.store_nd [[add0]], [[tdesc_c]][[[c0]], [[c0]]] : vector<1x16xf32>, !xegpu.tensor_desc<1x16xf32>
|
|
// CHECK: xegpu.store_nd [[add1]], [[tdesc_c]][[[c0]], [[c16]]] : vector<1x16xf32>, !xegpu.tensor_desc<1x16xf32>
|
|
gpu.func @load_store_nd_with_offsets(%A: memref<1024x1024xf32>, %B: memref<1024x1024xf32>, %C: memref<1024x1024xf32>) {
|
|
%c0 = arith.constant 0 : index
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x32xf32, #l>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x32xf32, #l>
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<1x32xf32, #l>
|
|
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %c0] {layout = #l}: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
|
|
%b = xegpu.load_nd %b_tdesc[%c0, %c0] {layout = #l}: !xegpu.tensor_desc<1x32xf32, #l> -> vector<1x32xf32>
|
|
|
|
%result = arith.addf %a, %b : vector<1x32xf32>
|
|
xegpu.store_nd %result, %c_tdesc[%c0, %c0] {layout = #l}: vector<1x32xf32>, !xegpu.tensor_desc<1x32xf32, #l>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#inst_data = #xegpu.layout<inst_data = [1, 1, 16]>
|
|
gpu.module @test_kernel {
|
|
// CHECK-LABEL: load_add_store_leading_unit_dims
|
|
// CHECK-SAME: [[arg0:%.+]]: ui64, [[arg1:%.+]]: ui64, [[arg2:%.+]]: ui64
|
|
// CHECK: [[c0:%.+]] = arith.constant dense<true> : vector<1x1x16xi1>
|
|
// CHECK: [[c1:%.+]] = arith.constant dense<[{{\[\[}}0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120]]]> : vector<1x1x16xindex>
|
|
// CHECK: [[c2:%.+]] = arith.constant dense<[{{\[\[}}128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]]]> : vector<1x1x16xindex>
|
|
// CHECK: [[v0:%.+]] = xegpu.load [[arg0]]{{\[}}[[c1]]], [[c0]]
|
|
// CHECK-SAME: ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[v1:%.+]] = xegpu.load [[arg0]]{{\[}}[[c2]]], [[c0]]
|
|
// CHECK-SAME: ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[v2:%.+]] = xegpu.load [[arg1]]{{\[}}[[c1]]], [[c0]]
|
|
// CHECK-SAME: ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[v3:%.+]] = xegpu.load [[arg1]]{{\[}}[[c2]]], [[c0]]
|
|
// CHECK-SAME: ui64, vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf32>
|
|
// CHECK: [[v4:%.+]] = arith.addf [[v0]], [[v2]] : vector<1x1x16xf32>
|
|
// CHECK: [[v5:%.+]] = arith.addf [[v1]], [[v3]] : vector<1x1x16xf32>
|
|
// CHECK: xegpu.store [[v4]], [[arg2]]{{\[}}[[c1]]], [[c0]]
|
|
// CHECK-SAME: vector<1x1x16xf32>, ui64, vector<1x1x16xindex>, vector<1x1x16xi1>
|
|
// CHECK: xegpu.store [[v5]], [[arg2]]{{\[}}[[c2]]], [[c0]]
|
|
// CHECK-SAME: vector<1x1x16xf32>, ui64, vector<1x1x16xindex>, vector<1x1x16xi1>
|
|
gpu.func @load_add_store_leading_unit_dims(%A: ui64, %B: ui64, %C: ui64) {
|
|
%cst = arith.constant dense<[
|
|
[[0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120,
|
|
128, 136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248]]
|
|
]> : vector<1x1x32xindex>
|
|
%mask = arith.constant dense<true> : vector<1x1x32xi1>
|
|
%a = xegpu.load %A[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
|
|
%b = xegpu.load %B[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : ui64, vector<1x1x32xindex>, vector<1x1x32xi1> -> vector<1x1x32xf32>
|
|
%addf = arith.addf %a, %b : vector<1x1x32xf32>
|
|
xegpu.store %addf, %C[%cst], %mask {chunk_size = 1, layout = #inst_data, l1_hint = #xegpu.cache_hint<cached>} : vector<1x1x32xf32>, ui64, vector<1x1x32xindex>, vector<1x1x32xi1>
|
|
gpu.return
|
|
}
|
|
}
|
|
|
|
// -----
|
|
#l1 = #xegpu.layout<inst_data = [8, 32]>
|
|
#l2 = #xegpu.layout<inst_data = [32, 16]>
|
|
#l3 = #xegpu.layout<inst_data = [8, 16]>
|
|
#l1_scale = #xegpu.layout<inst_data = [8, 1]>
|
|
#l2_scale = #xegpu.layout<inst_data = [1, 16]>
|
|
gpu.module @test_kernel {
|
|
gpu.func @dpas_mx(%A: memref<1024x1024xf4E2M1FN>, %B: memref<1024x1024xf4E2M1FN>, %C: memref<1024x1024xf32>, %scale_a: memref<1024x64xf8E8M0FNU>, %scale_b: memref<64x1024xf8E8M0FNU>) {
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
%c32 = arith.constant 32 : index
|
|
%c64 = arith.constant 64 : index
|
|
%c1024 = arith.constant 1024 : index
|
|
%block_id_x = gpu.block_id x
|
|
%block_id_y = gpu.block_id y
|
|
%m = arith.muli %block_id_x, %c16 : index
|
|
%n = arith.muli %block_id_y, %c32 : index
|
|
|
|
%c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #l3>
|
|
%c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l3}: !xegpu.tensor_desc<16x32xf32, #l3> -> vector<16x32xf32>
|
|
|
|
%a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x64xf4E2M1FN, #l1>
|
|
%b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<64x32xf4E2M1FN, #l2>
|
|
%scale_a_tdesc = xegpu.create_nd_tdesc %scale_a : memref<1024x64xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale>
|
|
%scale_b_tdesc = xegpu.create_nd_tdesc %scale_b : memref<64x1024xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale>
|
|
|
|
%out = scf.for %k = %c0 to %c1024 step %c64
|
|
iter_args(%arg2 = %c_init)
|
|
-> (vector<16x32xf32>) {
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x32xf4E2M1FN>
|
|
%a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l1}: !xegpu.tensor_desc<16x64xf4E2M1FN, #l1> -> vector<16x64xf4E2M1FN>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<32x16xf4E2M1FN>
|
|
%b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #l2}: !xegpu.tensor_desc<64x32xf4E2M1FN, #l2> -> vector<64x32xf4E2M1FN>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x1xf8E8M0FNU>
|
|
%sa = xegpu.load_nd %scale_a_tdesc[%c0, %c0] {layout = #l1_scale}: !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale> -> vector<16x2xf8E8M0FNU>
|
|
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x16xf8E8M0FNU>
|
|
%sb = xegpu.load_nd %scale_b_tdesc[%c0, %c0] {layout = #l2_scale}: !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale> -> vector<2x32xf8E8M0FNU>
|
|
//CHECK-COUNT-8: xegpu.dpas_mx {{.*}}
|
|
%c = xegpu.dpas_mx %a, %b, %arg2 scale_a = %sa scale_b = %sb {layout_a=#l1, layout_b = #l2, layout_cd = #l3, layout_a_scale = #l1_scale, layout_b_scale = #l2_scale, layout_result_0 = #l3}: (vector<16x64xf4E2M1FN>, vector<64x32xf4E2M1FN>, vector<16x32xf32>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xf32>
|
|
scf.yield %c : vector<16x32xf32>
|
|
} {layout_result_0 = #l3}
|
|
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
|
|
xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l3}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l3>
|
|
gpu.return
|
|
}
|
|
}
|