124 lines · plain
1// RUN: mlir-opt %s --convert-vector-to-llvm="enable-amx" --convert-to-llvm -reconcile-unrealized-casts \2// RUN: | mlir-translate --mlir-to-llvmir \3// RUN: | FileCheck %s4 5// CHECK-LABEL: define void @amx_tile_zero6func.func @amx_tile_zero(%out: memref<?x?xf32>, %idx: index)7{8 // CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)9 // CHECK: call void @llvm.x86.tilestored64.internal10 %zero = amx.tile_zero : !amx.tile<16x16xf32>11 amx.tile_store %out[%idx, %idx], %zero : memref<?x?xf32>, !amx.tile<16x16xf32>12 return13}14 15// CHECK-LABEL: define void @amx_tile_load_store16func.func @amx_tile_load_store(%base: memref<?x?xi8>, %out: memref<?x?xi8>,17 %idx: index)18{19 // CHECK: call x86_amx @llvm.x86.tileloadd64.internal20 // CHECK: call void @llvm.x86.tilestored64.internal21 %val = amx.tile_load %base[%idx, %idx] : memref<?x?xi8> into !amx.tile<16x64xi8>22 amx.tile_store %out[%idx, %idx], %val : memref<?x?xi8>, !amx.tile<16x64xi8>23 return24}25 26// CHECK-LABEL: define void @amx_tile_load_store_strided27func.func @amx_tile_load_store_strided(%base: memref<?xi8>, %out: memref<?xi8>,28 %idx: index, %stride: index)29{30 // CHECK: call x86_amx @llvm.x86.tileloadd64.internal31 // CHECK: call void @llvm.x86.tilestored64.internal32 %val = amx.tile_load %base[%idx], %stride33 : memref<?xi8> into !amx.tile<16x64xi8>34 amx.tile_store %out[%idx], %val, %stride35 : memref<?xi8>, !amx.tile<16x64xi8>36 return37}38 39// CHECK-LABEL: define void @amx_tile_mulf_bf1640func.func @amx_tile_mulf_bf16(41 %matA: memref<?x?xbf16>, %matB: memref<?x?xbf16>, %idx: index,42 %out: memref<?x?xf32>)43{44 // CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)45 %acc = amx.tile_zero : !amx.tile<16x16xf32>46 // CHECK-COUNT-2: call x86_amx @llvm.x86.tileloadd64.internal47 %tA = amx.tile_load %matA[%idx, %idx] : memref<?x?xbf16> into !amx.tile<16x32xbf16>48 %tB = amx.tile_load %matB[%idx, %idx] : memref<?x?xbf16> into !amx.tile<16x32xbf16>49 // CHECK: call x86_amx @llvm.x86.tdpbf16ps.internal50 %tRes = amx.tile_mulf %tA, %tB, %acc51 : !amx.tile<16x32xbf16>, !amx.tile<16x32xbf16>, !amx.tile<16x16xf32>52 // CHECK: call void @llvm.x86.tilestored64.internal53 amx.tile_store %out[%idx, %idx], %tRes : memref<?x?xf32>, !amx.tile<16x16xf32>54 return55}56 57// CHECK-LABEL: define void @amx_tile_mulf_f1658func.func @amx_tile_mulf_f16(59 %matA: memref<?x?xf16>, %matB: memref<?x?xf16>, %idx: index,60 %out: memref<?x?xf32>)61{62 // CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)63 %acc = amx.tile_zero : !amx.tile<16x16xf32>64 // CHECK-COUNT-2: call x86_amx @llvm.x86.tileloadd64.internal65 %tA = amx.tile_load %matA[%idx, %idx] : memref<?x?xf16> into !amx.tile<16x32xf16>66 %tB = amx.tile_load %matB[%idx, %idx] : memref<?x?xf16> into !amx.tile<16x32xf16>67 // CHECK: call x86_amx @llvm.x86.tdpfp16ps.internal68 %tRes = amx.tile_mulf %tA, %tB, %acc69 : !amx.tile<16x32xf16>, !amx.tile<16x32xf16>, !amx.tile<16x16xf32>70 // CHECK: call void @llvm.x86.tilestored64.internal71 amx.tile_store %out[%idx, %idx], %tRes : memref<?x?xf32>, !amx.tile<16x16xf32>72 return73}74 75// CHECK-LABEL: define void @amx_tile_muli76func.func @amx_tile_muli(%matA: memref<?x?xi8>, %matB: memref<?x?xi8>,77 %matC: memref<?x?xi32>, %idx: index, %out: memref<?x?xi8>)78{79 %c0 = arith.constant 0 : index80 %c16 = arith.constant 16 : index81 // CHECK-COUNT-3: call x86_amx @llvm.x86.tileloadd64.internal82 %tA = amx.tile_load %matA[%idx, %idx] : memref<?x?xi8> into !amx.tile<16x64xi8>83 %tB = amx.tile_load %matB[%idx, %idx] : memref<?x?xi8> into !amx.tile<16x64xi8>84 %acc = amx.tile_load %matC[%idx, %idx] : memref<?x?xi32> into !amx.tile<16x16xi32>85 // CHECK: call x86_amx @llvm.x86.tdpbuud.internal86 // CHECK: call x86_amx @llvm.x86.tdpbssd.internal87 // CHECK: call x86_amx @llvm.x86.tdpbusd.internal88 // CHECK: call x86_amx @llvm.x86.tdpbsud.internal89 %res = amx.tile_muli %tA zext, %tB zext, %acc90 : !amx.tile<16x64xi8>, !amx.tile<16x64xi8>, !amx.tile<16x16xi32>91 %res1 = amx.tile_muli %tA, %tB, %acc92 : !amx.tile<16x64xi8>, !amx.tile<16x64xi8>, !amx.tile<16x16xi32>93 %res2 = amx.tile_muli %tA zext, %tB, %acc94 : !amx.tile<16x64xi8>, !amx.tile<16x64xi8>, !amx.tile<16x16xi32>95 %res3 = amx.tile_muli %tA, %tB zext, %acc96 : !amx.tile<16x64xi8>, !amx.tile<16x64xi8>, !amx.tile<16x16xi32>97 // CHECK-COUNT-4: call void @llvm.x86.tilestored64.internal98 amx.tile_store %out[%c0, %c0], %res : memref<?x?xi8>, !amx.tile<16x16xi32>99 amx.tile_store %out[%c0, %c16], %res1 : memref<?x?xi8>, !amx.tile<16x16xi32>100 amx.tile_store %out[%c16, %c0], %res2 : memref<?x?xi8>, !amx.tile<16x16xi32>101 amx.tile_store %out[%c16, %c16], %res3 : memref<?x?xi8>, !amx.tile<16x16xi32>102 return103}104 105// CHECK-LABEL: define void @amx_tile_type_through_cf106func.func @amx_tile_type_through_cf(%src: memref<?x?xi8>, %out: memref<?x?xi8>,107 %idx: index, %cond: i1) {108 cf.cond_br %cond, ^bb1, ^bb2109^bb1: // pred: ^bb0110 // CHECK: call x86_amx @llvm.x86.tileloadd64.internal111 %0 = amx.tile_load %src[%idx, %idx] : memref<?x?xi8> into !amx.tile<16x64xi8>112 cf.br ^bb3(%0 : !amx.tile<16x64xi8>)113^bb2: // pred: ^bb0114 // CHECK: call x86_amx @llvm.x86.tilezero.internal(i16 16, i16 64)115 %1 = amx.tile_zero : !amx.tile<16x64xi8>116 cf.br ^bb3(%1 : !amx.tile<16x64xi8>)117^bb3(%2: !amx.tile<16x64xi8>): // 2 preds: ^bb1, ^bb2118 cf.br ^bb4119^bb4: // pred: ^bb3120 // CHECK: call void @llvm.x86.tilestored64.internal121 amx.tile_store %out[%idx, %idx], %2 : memref<?x?xi8>, !amx.tile<16x64xi8>122 return123}124