101 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -canonicalize -split-input-file | FileCheck %s2 3func.func @gemm(%a : memref<?x?xf32>, %b : memref<?x?xf32>, %c : memref<?x?xf32>)4{5 linalg.matmul ins(%a, %b: memref<?x?xf32>, memref<?x?xf32>)6 outs(%c: memref<?x?xf32>)7 return8}9 10// CHECK: func @gemm11// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<?x?xf32>12// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: memref<?x?xf32>13// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: memref<?x?xf32>14// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index15// CHECK: scf.for16// CHECK: scf.for17// CHECK: scf.for18// CHECK: %[[svA:.+]] = memref.subview %[[ARG0]]19// CHECK: %[[svB:.+]] = memref.subview %[[ARG1]]20// CHECK: %[[svC:.+]] = memref.subview %[[ARG2]]21 22// CHECK: %[[tmpA:.*]] = memref.alloc() : memref<1024xi8>23// CHECK: %[[VA:.*]] = memref.view %[[tmpA]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32>24// CHECK: %[[svAA:.+]] = memref.subview %[[VA]]25 26// CHECK: %[[tmpC:.*]] = memref.alloc() : memref<1024xi8>27// CHECK: %[[VC:.*]] = memref.view %[[tmpC]][%[[C0]]][] : memref<1024xi8> to memref<16x16xf32>28// CHECK: %[[svCC:.+]] = memref.subview %[[VC]]29 30// CHECK: linalg.copy ins(%[[svA]] : memref<?x?xf32, strided<[?, 1], offset: ?>>) outs(%[[svAA]] : memref<?x?xf32, strided<[16, 1]>>)31// CHECK: linalg.copy ins(%[[svC]] : memref<?x?xf32, strided<[?, 1], offset: ?>>) outs(%[[svCC]] : memref<?x?xf32, strided<[16, 1]>>)32// CHECK: linalg.matmul ins(%[[VA]], %[[svB]]{{.*}} outs(%[[VC]]33// CHECK: linalg.copy ins(%[[svCC]] : memref<?x?xf32, strided<[16, 1]>>) outs(%[[svC]] : memref<?x?xf32, strided<[?, 1], offset: ?>>)34// CHECK: memref.dealloc %[[tmpA]]35// CHECK: memref.dealloc %[[tmpC]]36 37module attributes {transform.with_named_sequence} {38 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {39 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op40 %1, %loops:3 = transform.structured.tile_using_for %0 tile_sizes [16, 16, 16] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)41 %2 = transform.structured.promote %1 { operands_to_promote = [0, 2], force_full_tiles = [false, false], use_full_tiles_by_default } : (!transform.any_op) -> !transform.any_op42 transform.yield43 }44}45 46// -----47 48func.func @matmul_f32(%A: memref<512x256xf32>, %B: memref<256x512xf32>, %C: memref<256x256xf32>, %s0: index, %s1: index, %s2: index) {49 %c0 = arith.constant 0 : index50 %c256 = arith.constant 256 : index51 %c512 = arith.constant 512 : index52 scf.for %arg4 = %c0 to %c512 step %s0 {53 scf.for %arg5 = %c0 to %c512 step %s1 {54 scf.for %arg6 = %c0 to %c256 step %s2 {55 %i0 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg4)[%s0]56 %i1 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg5)[%s1]57 %i2 = affine.min affine_map<(d0)[s0] -> (-d0 + 256, s0)>(%arg6)[%s2]58 %0 = memref.subview %A[%arg4, %arg6][%i0, %i2][1, 1] : memref<512x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>59 %1 = memref.subview %B[%arg6, %arg5][%i2, %i1][1, 1] : memref<256x512xf32> to memref<?x?xf32, strided<[512, 1], offset: ?>>60 %2 = memref.subview %C[%arg4, %arg5][%i0, %i1][1, 1] : memref<256x256xf32> to memref<?x?xf32, strided<[256, 1], offset: ?>>61 linalg.matmul62 ins(%0, %1: memref<?x?xf32, strided<[256, 1], offset: ?>>,63 memref<?x?xf32, strided<[512, 1], offset: ?>>)64 outs(%2: memref<?x?xf32, strided<[256, 1], offset: ?>>)65 }66 }67 }68 return69}70 71// CHECK-LABEL: func.func @matmul_f32(72// CHECK-SAME: %[[ARG0:.*]]: memref<512x256xf32>73// CHECK-SAME: %[[ARG1:.*]]: memref<256x512xf32>74// CHECK-SAME: %[[ARG2:.*]]: memref<256x256xf32>75// CHECK-SAME: %[[ARG3:.*]]: index, %[[ARG4:.*]]: index, %[[ARG5:.*]]: index76// CHECK: %[[C4:.*]] = arith.constant 4 : index77 78// CHECK: %[[i0:.*]] = affine.min79// CHECK: %[[i1:.*]] = affine.min80// CHECK: %[[i2:.*]] = affine.min81 82// CHECK: %[[VAL_13:.*]] = arith.muli %[[i0]], %[[i2]] : index83// CHECK: %[[VAL_14:.*]] = arith.muli %[[VAL_13]], %[[C4]] : index84// CHECK: %[[VAL_15:.*]] = memref.alloc(%[[VAL_14]]) : memref<?xi8>85 86// CHECK: %[[VAL_18:.*]] = arith.muli %[[i2]], %[[i1]] : index87// CHECK: %[[VAL_19:.*]] = arith.muli %[[VAL_18]], %[[C4]] : index88// CHECK: %[[VAL_20:.*]] = memref.alloc(%[[VAL_19]]) : memref<?xi8>89 90// CHECK: %[[VAL_23:.*]] = arith.muli %[[i0]], %[[i1]] : index91// CHECK: %[[VAL_24:.*]] = arith.muli %[[VAL_23]], %[[C4]] : index92// CHECK: %[[VAL_25:.*]] = memref.alloc(%[[VAL_24]]) : memref<?xi8>93 94module attributes {transform.with_named_sequence} {95 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {96 %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op97 %1 = transform.structured.promote %0 { use_original_subview_size } : (!transform.any_op) -> !transform.any_op98 transform.yield99 }100}101