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
187 lines
8.2 KiB
MLIR
187 lines
8.2 KiB
MLIR
// RUN: mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
|
|
|
|
// CHECK-LABEL: func.func @blockload2d(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,
|
|
// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32)
|
|
func.func @blockload2d(%a: !llvm.ptr<1>, %base_width_a: i32, %base_height_a: i32,
|
|
%base_pitch_a: i32, %x: i32, %y: i32) -> vector<8xi16> {
|
|
// CHECK: %[[VAR0:.*]] = xevm.blockload2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]]
|
|
// CHECK-DAG: elem_size_in_bits = 16 : i32
|
|
// CHECK-DAG: tile_width = 16 : i32
|
|
// CHECK-DAG: tile_height = 8 : i32
|
|
// CHECK-DAG: v_blocks = 1 : i32
|
|
// CHECK-DAG: transpose = false
|
|
// CHECK-DAG: pack_register = false
|
|
// CHECK-DAG: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>
|
|
// CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi16>
|
|
%loaded_a = xevm.blockload2d %a, %base_width_a, %base_height_a, %base_pitch_a, %x, %y
|
|
<{elem_size_in_bits=16 : i32, tile_width=16 : i32, tile_height=8 : i32, v_blocks=1 : i32,
|
|
transpose=false, pack_register=false, cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}>
|
|
: (!llvm.ptr<1>, i32, i32, i32, i32, i32) -> vector<8xi16>
|
|
return %loaded_a : vector<8xi16>
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @blockstore2d(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,
|
|
// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32,
|
|
// CHECK-SAME: %[[ARG6:.*]]: vector<8xi32>)
|
|
func.func @blockstore2d(%c: !llvm.ptr<1>, %base_width_c: i32, %base_height_c: i32,
|
|
%base_pitch_c: i32, %x: i32, %y: i32, %c_result_casted: vector<8xi32>) {
|
|
// CHECK: xevm.blockstore2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG6]]
|
|
// CHECK-DAG: elem_size_in_bits = 32 : i32
|
|
// CHECK-DAG: tile_width = 16 : i32
|
|
// CHECK-DAG: tile_height = 8 : i32
|
|
// CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
|
|
xevm.blockstore2d %c, %base_width_c, %base_height_c, %base_pitch_c, %x, %y, %c_result_casted
|
|
<{elem_size_in_bits=32 : i32, tile_width=16 : i32, tile_height=8 : i32}>
|
|
: (!llvm.ptr<1>, i32, i32, i32, i32, i32, vector<8xi32>)
|
|
return
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @blockprefetch2d(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,
|
|
// CHECK-SAME: %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32)
|
|
func.func @blockprefetch2d(%ptr: !llvm.ptr<1>, %base_width: i32, %base_height: i32,
|
|
%base_pitch: i32, %x: i32, %y: i32) {
|
|
// CHECK: xevm.blockprefetch2d %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG5]]
|
|
// CHECK-DAG: elem_size_in_bits = 8 : i32
|
|
// CHECK-DAG: tile_width = 32 : i32
|
|
// CHECK-DAG: tile_height = 8 : i32
|
|
// CHECK-DAG: v_blocks = 1 : i32
|
|
// CHECK-DAG: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>
|
|
// CHECK: (!llvm.ptr<1>, i32, i32, i32, i32, i32)
|
|
xevm.blockprefetch2d %ptr, %base_width, %base_height, %base_pitch, %x, %y <{elem_size_in_bits=8 : i32,
|
|
tile_width=32 : i32, tile_height=8 : i32, v_blocks=1 : i32,
|
|
cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>, i32, i32, i32, i32, i32)
|
|
return
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @blockload(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>)
|
|
func.func @blockload(%ptr: !llvm.ptr<1>) -> vector<4xi16> {
|
|
// CHECK: %[[VAR0:.*]] = xevm.blockload %[[ARG0]]
|
|
// CHECK-SAME: cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>
|
|
// CHECK-SAME: (!llvm.ptr<1>) -> vector<4xi16>
|
|
%loaded = xevm.blockload %ptr <{cache_control=#xevm.load_cache_control<L1uc_L2uc_L3uc>}>
|
|
: (!llvm.ptr<1>) -> vector<4xi16>
|
|
return %loaded : vector<4xi16>
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @blockstore(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>,
|
|
// CHECK-SAME: %[[ARG1:.*]]: vector<4xi32>)
|
|
func.func @blockstore(%ptr: !llvm.ptr<1>, %value: vector<4xi32>) {
|
|
// CHECK: xevm.blockstore %[[ARG0]], %[[ARG1]]
|
|
// CHECK-SAME: (!llvm.ptr<1>, vector<4xi32>)
|
|
xevm.blockstore %ptr, %value : (!llvm.ptr<1>, vector<4xi32>)
|
|
return
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @mma(
|
|
// CHECK-SAME: %[[ARG0:.*]]: vector<8xf32>, %[[ARG1:.*]]: vector<8xi16>, %[[ARG2:.*]]: vector<8xi32>)
|
|
func.func @mma(%loaded_c_casted: vector<8xf32>, %loaded_a: vector<8xi16>, %loaded_b_casted: vector<8xi32>) -> vector<8xf32> {
|
|
// CHECK: %0 = xevm.mma %[[ARG1]], %[[ARG2]], %[[ARG0]] {shape = <m = 8, n = 16, k = 16>,
|
|
// CHECK-SAME: types = <d = f32, a = f16, b = f16, c = f32>} : (vector<8xi16>, vector<8xi32>, vector<8xf32>) -> vector<8xf32>
|
|
%c_result = xevm.mma %loaded_a, %loaded_b_casted, %loaded_c_casted { shape=<m=8, n=16, k=16>,
|
|
types=<d=f32, a=f16, b=f16, c=f32> } : (vector<8xi16>, vector<8xi32>, vector<8xf32>) -> vector<8xf32>
|
|
return %c_result : vector<8xf32>
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @mma_mx(
|
|
// CHECK-SAME: %[[ARG0:.*]]: vector<8xf32>, %[[ARG1:.*]]: vector<8xi16>, %[[ARG2:.*]]: vector<8xi32>, %[[ARG3:.*]]: vector<2xi8>, %[[ARG4:.*]]: vector<2xi8>)
|
|
func.func @mma_mx(%loaded_c_casted: vector<8xf32>, %loaded_a: vector<8xi16>, %loaded_b_casted: vector<8xi32>, %scale_a: vector<2xi8>, %scale_b: vector<2xi8>) -> vector<8xf32> {
|
|
// CHECK: %[[VAR0:.*]] = xevm.mma_mx %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]], %[[ARG0]]
|
|
// CHECK-SAME: {shape = <m = 8, n = 16, k = 64>, types = <d = f32, a = e2m1, b = e2m1, c = f32>}
|
|
// CHECK-SAME:: (vector<8xi16>, vector<8xi32>, vector<2xi8>, vector<2xi8>, vector<8xf32>) -> vector<8xf32>
|
|
%c_result = xevm.mma_mx %loaded_a, %loaded_b_casted, %scale_a, %scale_b, %loaded_c_casted { shape=<m=8, n=16, k=64>,
|
|
types=<d=f32, a=e2m1, b=e2m1, c=f32> } : (vector<8xi16>, vector<8xi32>, vector<2xi8>, vector<2xi8>, vector<8xf32>) -> vector<8xf32>
|
|
return %c_result : vector<8xf32>
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @truncf_scalar
|
|
func.func @truncf_scalar() -> i8 {
|
|
// CHECK: %[[VAR0:.*]] = arith.constant 1.0
|
|
%0 = arith.constant 1.0 : bf16
|
|
// CHECK: xevm.truncf %[[VAR0]] {src_etype = bf16, dst_etype = f8} : (bf16) -> i8
|
|
%2 = xevm.truncf %0 { src_etype=bf16, dst_etype=f8 } : (bf16) -> i8
|
|
return %2 : i8
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @truncf_vector
|
|
func.func @truncf_vector() -> vector<8xi4> {
|
|
// CHECK: %[[VAR0:.*]] = arith.constant
|
|
%0 = arith.constant dense<1.0> : vector<8xbf16>
|
|
// CHECK: xevm.truncf %[[VAR0]] {src_etype = bf16, dst_etype = e2m1} : (vector<8xbf16>) -> vector<8xi4>
|
|
%2 = xevm.truncf %0 { src_etype=bf16, dst_etype=e2m1 } : (vector<8xbf16>) -> vector<8xi4>
|
|
return %2 : vector<8xi4>
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @memfence()
|
|
func.func @memfence() {
|
|
// CHECK: xevm.memfence
|
|
// CHECK-DAG: addrspace = #xevm.addr_space<global>
|
|
// CHECK-DAG: scope = #xevm.mem_scope<workgroup>
|
|
xevm.memfence <{addrspace=#xevm.addr_space<global>, scope=#xevm.mem_scope<workgroup>}>
|
|
return
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: func.func @prefetch(
|
|
// CHECK-SAME: %[[ARG0:.*]]: !llvm.ptr<1>)
|
|
func.func @prefetch(%ptr: !llvm.ptr<1>) {
|
|
// CHECK: xevm.prefetch %[[ARG0]]
|
|
// CHECK-SAME: <{cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>)
|
|
xevm.prefetch %ptr <{cache_control = #xevm.load_cache_control<L1uc_L2uc_L3uc>}> : (!llvm.ptr<1>)
|
|
return
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: @xevm_module [#xevm.target<O = 3, chip = "pvc">] {
|
|
gpu.module @xevm_module [#xevm.target<O = 3, chip = "pvc">]{
|
|
}
|
|
|
|
// -----
|
|
// CHECK-LABEL: @xevm_special_ids
|
|
llvm.func @xevm_special_ids() -> i32 {
|
|
// CHECK: xevm.local_id.x : i32
|
|
%1 = xevm.local_id.x : i32
|
|
// CHECK: xevm.local_id.y : i32
|
|
%2 = xevm.local_id.y : i32
|
|
// CHECK: xevm.local_id.z : i32
|
|
%3 = xevm.local_id.z : i32
|
|
// CHECK: xevm.local_size.x : i32
|
|
%4 = xevm.local_size.x : i32
|
|
// CHECK: xevm.local_size.y : i32
|
|
%5 = xevm.local_size.y : i32
|
|
// CHECK: xevm.local_size.z : i32
|
|
%6 = xevm.local_size.z : i32
|
|
// CHECK: xevm.group_id.x : i32
|
|
%7 = xevm.group_id.x : i32
|
|
// CHECK: xevm.group_id.y : i32
|
|
%8 = xevm.group_id.y : i32
|
|
// CHECK: xevm.group_id.z : i32
|
|
%9 = xevm.group_id.z : i32
|
|
// CHECK: xevm.group_count.x : i32
|
|
%10 = xevm.group_count.x : i32
|
|
// CHECK: xevm.group_count.y : i32
|
|
%11 = xevm.group_count.y : i32
|
|
// CHECK: xevm.group_count.z : i32
|
|
%12 = xevm.group_count.z : i32
|
|
// CHECK: xevm.lane_id : i32
|
|
%14 = xevm.lane_id : i32
|
|
// CHECK: xevm.subgroup_size : i32
|
|
%39 = xevm.subgroup_size : i32
|
|
// CHECK: xevm.subgroup_id : i32
|
|
%40 = xevm.subgroup_id : i32
|
|
llvm.return %1 : i32
|
|
}
|