// 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
}