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
158 lines
7.6 KiB
MLIR
158 lines
7.6 KiB
MLIR
// RUN: mlir-opt %s --convert-vector-to-llvm="enable-x86" --convert-to-llvm -reconcile-unrealized-casts \
|
|
// RUN: | mlir-translate --mlir-to-llvmir \
|
|
// RUN: | FileCheck %s
|
|
|
|
// CHECK-LABEL: define void @amx_tile_zero
|
|
func.func @amx_tile_zero(%out: memref<?x?xf32>, %idx: index)
|
|
{
|
|
// CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
%zero = x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %out[%idx, %idx], %zero : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_load_store
|
|
func.func @amx_tile_load_store(%base: memref<?x?xi8>, %out: memref<?x?xi8>,
|
|
%idx: index)
|
|
{
|
|
// CHECK: call x86_amx @llvm.x86.tileloadd64.internal
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
%val = x86.amx.tile_load %base[%idx, %idx] : memref<?x?xi8> into !x86.amx.tile<16x64xi8>
|
|
x86.amx.tile_store %out[%idx, %idx], %val : memref<?x?xi8>, !x86.amx.tile<16x64xi8>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_load_store_strided
|
|
func.func @amx_tile_load_store_strided(%base: memref<?xi8>, %out: memref<?xi8>,
|
|
%idx: index, %stride: index)
|
|
{
|
|
// CHECK: call x86_amx @llvm.x86.tileloadd64.internal
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
%val = x86.amx.tile_load %base[%idx], %stride
|
|
: memref<?xi8> into !x86.amx.tile<16x64xi8>
|
|
x86.amx.tile_store %out[%idx], %val, %stride
|
|
: memref<?xi8>, !x86.amx.tile<16x64xi8>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_mulf_bf16
|
|
func.func @amx_tile_mulf_bf16(
|
|
%matA: memref<?x?xbf16>, %matB: memref<?x?xbf16>, %idx: index,
|
|
%out: memref<?x?xf32>)
|
|
{
|
|
// CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)
|
|
%acc = x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-COUNT-2: call x86_amx @llvm.x86.tileloadd64.internal
|
|
%tA = x86.amx.tile_load %matA[%idx, %idx] : memref<?x?xbf16> into !x86.amx.tile<16x32xbf16>
|
|
%tB = x86.amx.tile_load %matB[%idx, %idx] : memref<?x?xbf16> into !x86.amx.tile<16x32xbf16>
|
|
// CHECK: call x86_amx @llvm.x86.tdpbf16ps.internal
|
|
%tRes = x86.amx.tile_mulf %tA, %tB, %acc
|
|
: !x86.amx.tile<16x32xbf16>, !x86.amx.tile<16x32xbf16>, !x86.amx.tile<16x16xf32>
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
x86.amx.tile_store %out[%idx, %idx], %tRes : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_mulf_f16
|
|
func.func @amx_tile_mulf_f16(
|
|
%matA: memref<?x?xf16>, %matB: memref<?x?xf16>, %idx: index,
|
|
%out: memref<?x?xf32>)
|
|
{
|
|
// CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)
|
|
%acc = x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-COUNT-2: call x86_amx @llvm.x86.tileloadd64.internal
|
|
%tA = x86.amx.tile_load %matA[%idx, %idx] : memref<?x?xf16> into !x86.amx.tile<16x32xf16>
|
|
%tB = x86.amx.tile_load %matB[%idx, %idx] : memref<?x?xf16> into !x86.amx.tile<16x32xf16>
|
|
// CHECK: call x86_amx @llvm.x86.tdpfp16ps.internal
|
|
%tRes = x86.amx.tile_mulf %tA, %tB, %acc
|
|
: !x86.amx.tile<16x32xf16>, !x86.amx.tile<16x32xf16>, !x86.amx.tile<16x16xf32>
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
x86.amx.tile_store %out[%idx, %idx], %tRes : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_mulf_f8
|
|
func.func @amx_tile_mulf_f8(
|
|
%matA: memref<?x?xf8E4M3FN>, %matB: memref<?x?xf8E5M2>, %idx: index,
|
|
%out: memref<?x?xf32>)
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
// CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)
|
|
%acc = x86.amx.tile_zero : !x86.amx.tile<16x16xf32>
|
|
// CHECK-COUNT-4: call x86_amx @llvm.x86.tileloadd64.internal
|
|
%tA = x86.amx.tile_load %matA[%idx, %idx] : memref<?x?xf8E4M3FN> into !x86.amx.tile<16x64xf8E4M3FN>
|
|
%tB = x86.amx.tile_load %matB[%idx, %idx] : memref<?x?xf8E5M2> into !x86.amx.tile<16x64xf8E5M2>
|
|
%tA1 = x86.amx.tile_load %matA[%c0, %c0] : memref<?x?xf8E4M3FN> into !x86.amx.tile<16x64xf8E4M3FN>
|
|
%tB1 = x86.amx.tile_load %matB[%c0, %c0] : memref<?x?xf8E5M2> into !x86.amx.tile<16x64xf8E5M2>
|
|
// CHECK: call x86_amx @llvm.x86.tdpbf8ps.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdpbhf8ps.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdphbf8ps.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdphf8ps.internal
|
|
%tRes = x86.amx.tile_mulf %tA, %tA1, %acc
|
|
: !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x16xf32>
|
|
%tRes1 = x86.amx.tile_mulf %tA, %tB, %acc
|
|
: !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x16xf32>
|
|
%tRes2 = x86.amx.tile_mulf %tB, %tA, %acc
|
|
: !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x16xf32>
|
|
%tRes3 = x86.amx.tile_mulf %tB, %tB1, %acc
|
|
: !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x16xf32>
|
|
// CHECK-COUNT-4: call void @llvm.x86.tilestored64.internal
|
|
x86.amx.tile_store %out[%idx, %c0], %tRes : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %out[%idx, %c16], %tRes1 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %out[%c0, %idx], %tRes2 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %out[%c16, %idx], %tRes3 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_muli
|
|
func.func @amx_tile_muli(%matA: memref<?x?xi8>, %matB: memref<?x?xi8>,
|
|
%matC: memref<?x?xi32>, %idx: index, %out: memref<?x?xi8>)
|
|
{
|
|
%c0 = arith.constant 0 : index
|
|
%c16 = arith.constant 16 : index
|
|
// CHECK-COUNT-3: call x86_amx @llvm.x86.tileloadd64.internal
|
|
%tA = x86.amx.tile_load %matA[%idx, %idx] : memref<?x?xi8> into !x86.amx.tile<16x64xi8>
|
|
%tB = x86.amx.tile_load %matB[%idx, %idx] : memref<?x?xi8> into !x86.amx.tile<16x64xi8>
|
|
%acc = x86.amx.tile_load %matC[%idx, %idx] : memref<?x?xi32> into !x86.amx.tile<16x16xi32>
|
|
// CHECK: call x86_amx @llvm.x86.tdpbuud.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdpbssd.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdpbusd.internal
|
|
// CHECK: call x86_amx @llvm.x86.tdpbsud.internal
|
|
%res = x86.amx.tile_muli %tA zext, %tB zext, %acc
|
|
: !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
%res1 = x86.amx.tile_muli %tA, %tB, %acc
|
|
: !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
%res2 = x86.amx.tile_muli %tA zext, %tB, %acc
|
|
: !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
%res3 = x86.amx.tile_muli %tA, %tB zext, %acc
|
|
: !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
// CHECK-COUNT-4: call void @llvm.x86.tilestored64.internal
|
|
x86.amx.tile_store %out[%c0, %c0], %res : memref<?x?xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %out[%c0, %c16], %res1 : memref<?x?xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %out[%c16, %c0], %res2 : memref<?x?xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %out[%c16, %c16], %res3 : memref<?x?xi8>, !x86.amx.tile<16x16xi32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: define void @amx_tile_type_through_cf
|
|
func.func @amx_tile_type_through_cf(%src: memref<?x?xi8>, %out: memref<?x?xi8>,
|
|
%idx: index, %cond: i1) {
|
|
cf.cond_br %cond, ^bb1, ^bb2
|
|
^bb1: // pred: ^bb0
|
|
// CHECK: call x86_amx @llvm.x86.tileloadd64.internal
|
|
%0 = x86.amx.tile_load %src[%idx, %idx] : memref<?x?xi8> into !x86.amx.tile<16x64xi8>
|
|
cf.br ^bb3(%0 : !x86.amx.tile<16x64xi8>)
|
|
^bb2: // pred: ^bb0
|
|
// CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xi8>
|
|
cf.br ^bb3(%1 : !x86.amx.tile<16x64xi8>)
|
|
^bb3(%2: !x86.amx.tile<16x64xi8>): // 2 preds: ^bb1, ^bb2
|
|
cf.br ^bb4
|
|
^bb4: // pred: ^bb3
|
|
// CHECK: call void @llvm.x86.tilestored64.internal
|
|
x86.amx.tile_store %out[%idx, %idx], %2 : memref<?x?xi8>, !x86.amx.tile<16x64xi8>
|
|
return
|
|
}
|