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
198 lines
12 KiB
MLIR
198 lines
12 KiB
MLIR
// RUN: mlir-opt %s -convert-vector-to-llvm="enable-x86" | mlir-opt | FileCheck %s
|
|
|
|
// CHECK-LABEL: muli(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbuud.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbssd.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbusd.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbsud.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @muli(%arg0: memref<?x?xi8>, %arg1: memref<?x?xi32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xi8>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xi8> into !x86.amx.tile<16x64xi8>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xi32> into !x86.amx.tile<16x16xi32>
|
|
%4 = x86.amx.tile_muli %1 zext, %2 zext, %3 : !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xi32>, !x86.amx.tile<16x16xi32>
|
|
%5 = x86.amx.tile_muli %1, %2, %3 : !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %arg1[%0, %0], %5 : memref<?x?xi32>, !x86.amx.tile<16x16xi32>
|
|
%6 = x86.amx.tile_muli %1 zext, %2, %3 : !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %arg1[%0, %0], %6 : memref<?x?xi32>, !x86.amx.tile<16x16xi32>
|
|
%7 = x86.amx.tile_muli %1, %2 zext, %3 : !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x64xi8>, !x86.amx.tile<16x16xi32>
|
|
x86.amx.tile_store %arg1[%0, %0], %7 : memref<?x?xi32>, !x86.amx.tile<16x16xi32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulbf16(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbf16ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulbf16(%arg0: memref<?x?xbf16>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x32xbf16>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xbf16> into !x86.amx.tile<16x32xbf16>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x32xbf16>, !x86.amx.tile<16x32xbf16>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulfp16(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpfp16ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulfp16(%arg0: memref<?x?xf16>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x32xf16>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xf16> into !x86.amx.tile<16x32xf16>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x32xf16>, !x86.amx.tile<16x32xf16>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulf8E4M3FN(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbf8ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulf8E4M3FN(%arg0: memref<?x?xf8E4M3FN>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xf8E4M3FN>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xf8E4M3FN> into !x86.amx.tile<16x64xf8E4M3FN>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulf8E4M3FNxf8E5M2(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdpbhf8ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulf8E4M3FNxf8E5M2(%arg0: memref<?x?xf8E5M2>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xf8E4M3FN>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xf8E5M2> into !x86.amx.tile<16x64xf8E5M2>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulf8E5M2xf8E4M3FN(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdphbf8ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulf8E5M2xf8E4M3FN(%arg0: memref<?x?xf8E4M3FN>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xf8E5M2>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xf8E4M3FN> into !x86.amx.tile<16x64xf8E4M3FN>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x64xf8E4M3FN>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
// CHECK-LABEL: mulf8E5M2(
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tdphf8ps.internal"
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"
|
|
func.func @mulf8E5M2(%arg0: memref<?x?xf8E5M2>, %arg1: memref<?x?xf32>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_zero : !x86.amx.tile<16x64xf8E5M2>
|
|
%2 = x86.amx.tile_load %arg0[%0, %0] : memref<?x?xf8E5M2> into !x86.amx.tile<16x64xf8E5M2>
|
|
%3 = x86.amx.tile_load %arg1[%0, %0] : memref<?x?xf32> into !x86.amx.tile<16x16xf32>
|
|
%4 = x86.amx.tile_mulf %1, %2, %3 : !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x64xf8E5M2>, !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg1[%0, %0], %4 : memref<?x?xf32>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
/// Intrinsics require stride in number of bytes.
|
|
// CHECK-LABEL: strides_implicit(
|
|
// CHECK: %[[LOAD_STRIDE_1:.+]] = llvm.mlir.constant(32 : i64) : i64
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_1]]
|
|
// CHECK: %[[LOAD_STRIDE_2:.+]] = llvm.mlir.constant(128 : i64) : i64
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_2]]
|
|
// CHECK: llvm.extractvalue %{{.+}}[4, 0]
|
|
// CHECK: %[[LOAD_BUF_STRIDE:.+]] = llvm.extractvalue %{{.+}}[4, 0]
|
|
// CHECK: %[[LOAD_STRIDE_SCALE:.+]] = llvm.mlir.constant(4 : i64) : i64
|
|
// CHECK: %[[LOAD_STRIDE_3:.+]] = llvm.mul %[[LOAD_STRIDE_SCALE]], %[[LOAD_BUF_STRIDE]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_3]]
|
|
// CHECK: %[[STORE_STRIDE_1:.+]] = llvm.mlir.constant(32 : i64) : i64
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_1]]
|
|
// CHECK: %[[STORE_STRIDE_2:.+]] = llvm.mlir.constant(128 : i64) : i64
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_2]]
|
|
// CHECK: llvm.extractvalue %{{.+}}[4, 0]
|
|
// CHECK: %[[STORE_BUF_STRIDE:.+]] = llvm.extractvalue %{{.+}}[4, 0]
|
|
// CHECK: %[[STORE_STRIDE_SCALE:.+]] = llvm.mlir.constant(4 : i64) : i64
|
|
// CHECK: %[[STORE_STRIDE_3:.+]] = llvm.mul %[[STORE_STRIDE_SCALE]], %[[STORE_BUF_STRIDE]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_3]]
|
|
func.func @strides_implicit(%arg0: memref<16x32xi8>,
|
|
%arg1: memref<32x32xbf16, strided<[64, 1]>>,
|
|
%arg2: memref<16x32xf32, strided<[?, 1]>>) {
|
|
%0 = arith.constant 0 : index
|
|
%1 = x86.amx.tile_load %arg0[%0, %0] : memref<16x32xi8> into !x86.amx.tile<16x32xi8>
|
|
%2 = x86.amx.tile_load %arg1[%0, %0] : memref<32x32xbf16, strided<[64, 1]>> into !x86.amx.tile<16x32xbf16>
|
|
%3 = x86.amx.tile_load %arg2[%0, %0] : memref<16x32xf32, strided<[?, 1]>> into !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg0[%0, %0], %1 : memref<16x32xi8>, !x86.amx.tile<16x32xi8>
|
|
x86.amx.tile_store %arg1[%0, %0], %2 : memref<32x32xbf16, strided<[64, 1]>>, !x86.amx.tile<16x32xbf16>
|
|
x86.amx.tile_store %arg2[%0, %0], %3 : memref<16x32xf32, strided<[?, 1]>>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|
|
|
|
/// Intrinsics require stride in number of bytes.
|
|
// CHECK-LABEL: strides_explicit(
|
|
// CHECK-SAME: %[[STRIDE:.+]]: index
|
|
// CHECK-DAG: %[[STRIDE_I64:.+]] = builtin.unrealized_conversion_cast %[[STRIDE]] : index to i64
|
|
// CHECK-DAG: %[[C64:.+]] = arith.constant 64 : index
|
|
// CHECK-DAG: %[[C64_I64:.+]] = builtin.unrealized_conversion_cast %[[C64]] : index to i64
|
|
// CHECK: %[[LOAD_STRIDE_SCALE_1:.+]] = llvm.mlir.constant(1 : i64) : i64
|
|
// CHECK: %[[LOAD_STRIDE_1:.+]] = llvm.mul %[[LOAD_STRIDE_SCALE_1]], %[[STRIDE_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_1]]
|
|
// CHECK: %[[LOAD_STRIDE_SCALE_2:.+]] = llvm.mlir.constant(2 : i64) : i64
|
|
// CHECK: %[[LOAD_STRIDE_2:.+]] = llvm.mul %[[LOAD_STRIDE_SCALE_2]], %[[STRIDE_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_2]]
|
|
// CHECK: %[[LOAD_STRIDE_SCALE_3:.+]] = llvm.mlir.constant(4 : i64) : i64
|
|
// CHECK: %[[LOAD_STRIDE_3:.+]] = llvm.mul %[[LOAD_STRIDE_SCALE_3]], %[[C64_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[LOAD_STRIDE_3]]
|
|
// CHECK: %[[STORE_STRIDE_SCALE_1:.+]] = llvm.mlir.constant(1 : i64) : i64
|
|
// CHECK: %[[STORE_STRIDE_1:.+]] = llvm.mul %[[STORE_STRIDE_SCALE_1]], %[[STRIDE_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_1]]
|
|
// CHECK: %[[STORE_STRIDE_SCALE_2:.+]] = llvm.mlir.constant(2 : i64) : i64
|
|
// CHECK: %[[STORE_STRIDE_2:.+]] = llvm.mul %[[STORE_STRIDE_SCALE_2]], %[[STRIDE_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_2]]
|
|
// CHECK: %[[STORE_STRIDE_SCALE_3:.+]] = llvm.mlir.constant(4 : i64) : i64
|
|
// CHECK: %[[STORE_STRIDE_3:.+]] = llvm.mul %[[STORE_STRIDE_SCALE_3]], %[[C64_I64]]
|
|
// CHECK: llvm.call_intrinsic "llvm.x86.tilestored64.internal"(%{{.+}}, %{{.+}}, %{{.+}}, %[[STORE_STRIDE_3]]
|
|
func.func @strides_explicit(%stride: index,
|
|
%arg0: memref<?xi8>,
|
|
%arg1: memref<16x32xbf16>,
|
|
%arg2: memref<32x32xf32, strided<[64, 1]>>) {
|
|
%0 = arith.constant 0 : index
|
|
%c64 = arith.constant 64 : index
|
|
%1 = x86.amx.tile_load %arg0[%0], %stride : memref<?xi8> into !x86.amx.tile<16x32xi8>
|
|
%2 = x86.amx.tile_load %arg1[%0, %0], %stride : memref<16x32xbf16> into !x86.amx.tile<16x32xbf16>
|
|
%3 = x86.amx.tile_load %arg2[%0, %0], %c64 : memref<32x32xf32, strided<[64, 1]>> into !x86.amx.tile<16x16xf32>
|
|
x86.amx.tile_store %arg0[%0], %1, %stride : memref<?xi8>, !x86.amx.tile<16x32xi8>
|
|
x86.amx.tile_store %arg1[%0, %0], %2, %stride : memref<16x32xbf16>, !x86.amx.tile<16x32xbf16>
|
|
x86.amx.tile_store %arg2[%0, %0], %3, %c64 : memref<32x32xf32, strided<[64, 1]>>, !x86.amx.tile<16x16xf32>
|
|
return
|
|
}
|