33 lines · plain
1// RUN: mlir-opt %s -cse -convert-vector-to-llvm="enable-amx" | FileCheck %s2 3// With inclusion of memory side-effects, it is expected CSE not to fold multiple 4// "tileload" and "tilezero".5// CHECK-LABEL: do_not_fold_tiles(6// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"7// CHECK: llvm.call_intrinsic "llvm.x86.tilezero.internal"8// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"9// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"10// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"11// CHECK: llvm.call_intrinsic "llvm.x86.tileloadd64.internal"12func.func @do_not_fold_tiles(%arg0: memref<2x32x32xbf16>, %arg1: memref<2x16x32xbf16>) -> memref<16x32xf32> {13 %c1 = arith.constant 1 : index14 %c0 = arith.constant 0 : index15 %c2 = arith.constant 2 : index16 %c16 = arith.constant 16 : index17 %alloca = memref.alloca() : memref<16x32xf32>18 %0 = amx.tile_zero : !amx.tile<16x16xf32>19 %1 = amx.tile_zero : !amx.tile<16x16xf32>20 %2:2 = scf.for %arg2 = %c0 to %c2 step %c1 iter_args(%arg3 = %0, %arg4 = %1) -> (!amx.tile<16x16xf32>, !amx.tile<16x16xf32>) {21 %3 = amx.tile_load %arg0[%arg2, %c0, %c0] : memref<2x32x32xbf16> into !amx.tile<16x32xbf16>22 %4 = amx.tile_load %arg0[%arg2, %c16, %c0] : memref<2x32x32xbf16> into !amx.tile<16x32xbf16>23 %5 = amx.tile_load %arg1[%arg2, %c0, %c0] : memref<2x16x32xbf16> into !amx.tile<16x32xbf16>24 %6 = amx.tile_load %arg1[%arg2, %c0, %c0] : memref<2x16x32xbf16> into !amx.tile<16x32xbf16>25 %7 = amx.tile_mulf %3, %5, %arg3 : !amx.tile<16x32xbf16>, !amx.tile<16x32xbf16>, !amx.tile<16x16xf32>26 %8 = amx.tile_mulf %4, %6, %arg4 : !amx.tile<16x32xbf16>, !amx.tile<16x32xbf16>, !amx.tile<16x16xf32>27 scf.yield %7, %8 : !amx.tile<16x16xf32>, !amx.tile<16x16xf32>28 }29 amx.tile_store %alloca[%c0, %c0], %2#0 : memref<16x32xf32>, !amx.tile<16x16xf32>30 amx.tile_store %alloca[%c0, %c16], %2#1 : memref<16x32xf32>, !amx.tile<16x16xf32>31 return %alloca : memref<16x32xf32>32}33