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
165 lines
9.4 KiB
MLIR
165 lines
9.4 KiB
MLIR
// RUN: mlir-opt --xegpu-wg-to-sg-distribute -split-input-file %s | FileCheck %s
|
|
|
|
gpu.module @test_elementwise_ops {
|
|
|
|
// CHECK-LABEL: unary_ops_sg_layout_only
|
|
gpu.func @unary_ops_sg_layout_only(%a: memref<24x32xf32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
// CHECK: math.exp {{.*}} : vector<12x8xf32>
|
|
%exp = math.exp %load_a
|
|
: vector<24x32xf32>
|
|
// CHECK: arith.negf {{.*}} : vector<12x8xf32>
|
|
%negf = arith.negf %exp
|
|
: vector<24x32xf32>
|
|
%anchor = xegpu.convert_layout %negf
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8]>
|
|
}> : vector<24x32xf32>
|
|
gpu.return
|
|
}
|
|
|
|
// CHECK-LABEL: unary_ops
|
|
gpu.func @unary_ops(%a: memref<24x32xf32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
// CHECK: math.exp {{.*}} : vector<12x8xf32>
|
|
%exp = math.exp %load_a
|
|
: vector<24x32xf32>
|
|
// CHECK: arith.negf {{.*}} : vector<12x8xf32>
|
|
%negf = arith.negf %exp
|
|
: vector<24x32xf32>
|
|
%anchor = xegpu.convert_layout %negf
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
}> : vector<24x32xf32>
|
|
gpu.return
|
|
}
|
|
|
|
// CHECK-LABEL: binary_ops
|
|
gpu.func @binary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
// CHECK: arith.addf {{.*}}, {{.*}} : vector<12x8xf32>
|
|
%addf = arith.addf %load_a, %load_b : vector<24x32xf32>
|
|
// CHECK: math.powf {{.*}}, {{.*}} : vector<12x8xf32>
|
|
%powf = math.powf %addf, %load_b : vector<24x32xf32>
|
|
%anchor = xegpu.convert_layout %powf
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
}> : vector<24x32xf32>
|
|
gpu.return
|
|
}
|
|
|
|
// CHECK-LABEL: ternary_ops
|
|
gpu.func @ternary_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi1>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi1> -> !xegpu.tensor_desc<24x32xi1>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xi1> -> vector<24x32xi1>
|
|
// CHECK: arith.select {{.*}}, {{.*}}, {{.*}} : vector<12x8xi1>, vector<12x8xf32>
|
|
%select = arith.select %load_c, %load_a, %load_b
|
|
: vector<24x32xi1>, vector<24x32xf32>
|
|
// CHECK: math.fma {{.*}}, {{.*}}, {{.*}} : vector<12x8xf32>
|
|
%fma = math.fma %load_a, %load_b, %select
|
|
: vector<24x32xf32>
|
|
%anchor = xegpu.convert_layout %fma
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
}> : vector<24x32xf32>
|
|
gpu.return
|
|
}
|
|
|
|
// CHECK-LABEL: type_conversion_ops
|
|
gpu.func @type_conversion_ops(%a: memref<24x32xf32>, %b: memref<24x32xi32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
|
|
// CHECK: arith.truncf {{.*}} : vector<12x8xf32> to vector<12x8xf16>
|
|
%truncf = arith.truncf %load_a
|
|
: vector<24x32xf32> to vector<24x32xf16>
|
|
// CHECK: arith.bitcast {{.*}} : vector<12x8xf16> to vector<12x8xi16>
|
|
%bitcast = arith.bitcast %truncf
|
|
: vector<24x32xf16> to vector<24x32xi16>
|
|
%anchor = xegpu.convert_layout %bitcast
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
}> : vector<24x32xi16>
|
|
gpu.return
|
|
}
|
|
|
|
// CHECK-LABEL: comparison_ops
|
|
gpu.func @comparison_ops(%a: memref<24x32xf32>, %b: memref<24x32xf32>, %c: memref<24x32xi32>, %d: memref<24x32xi32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_c = xegpu.create_nd_tdesc %c : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
|
|
%tdesc_d = xegpu.create_nd_tdesc %d : memref<24x32xi32> -> !xegpu.tensor_desc<24x32xi32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_c = xegpu.load_nd %tdesc_c[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
|
|
%load_d = xegpu.load_nd %tdesc_d[0, 0] {layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xi32> -> vector<24x32xi32>
|
|
// CHECK: arith.cmpf ult, {{.*}}, {{.*}} : vector<12x8xf32>
|
|
%cmpf = arith.cmpf ult, %load_a, %load_b
|
|
: vector<24x32xf32>
|
|
// CHECK: arith.cmpi eq, {{.*}}, {{.*}} : vector<12x8xi32>
|
|
%cmpi = arith.cmpi eq, %load_c, %load_d
|
|
: vector<24x32xi32>
|
|
%res = arith.select %cmpi, %cmpi, %cmpf : vector<24x32xi1>, vector<24x32xi1>
|
|
%anchor = xegpu.convert_layout %res
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [2, 4], sg_data = [12, 8], lane_layout = [2, 8], lane_data = [1, 1]>
|
|
}> : vector<24x32xi1>
|
|
gpu.return
|
|
}
|
|
|
|
// 1 to N decomposition of elementwise operations
|
|
// CHECK-LABEL: elementwise_ops_rr_assignment
|
|
gpu.func @elementwise_ops_rr_assignment(%a: memref<24x32xf32>, %b: memref<24x32xf32>) {
|
|
%tdesc_a = xegpu.create_nd_tdesc %a : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%tdesc_b = xegpu.create_nd_tdesc %b : memref<24x32xf32> -> !xegpu.tensor_desc<24x32xf32>
|
|
%load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
%load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>}
|
|
: !xegpu.tensor_desc<24x32xf32> -> vector<24x32xf32>
|
|
// CHECK-COUNT-12: arith.negf {{.*}} : vector<2x2xf32>
|
|
// CHECK-NOT: arith.negf
|
|
%negf = arith.negf %load_a
|
|
: vector<24x32xf32>
|
|
// CHECK-COUNT-12: math.powf {{.*}}, {{.*}} : vector<2x2xf32>
|
|
// CHECK-NOT: math.powf
|
|
%powf = math.powf %negf, %load_b
|
|
: vector<24x32xf32>
|
|
%anchor = xegpu.convert_layout %powf
|
|
<{
|
|
input_layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>,
|
|
target_layout = #xegpu.layout<sg_layout = [4, 4], sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>
|
|
}> : vector<24x32xf32>
|
|
gpu.return
|
|
}
|
|
}
|
|
|