1// RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s 2 3#matmat_accesses = [ 4 affine_map<(i, j, k) -> (i, k)>, 5 affine_map<(i, j, k) -> (k, j)>, 6 affine_map<(i, j, k) -> (i, j)> 7] 8#matmat_trait = { 9 indexing_maps = #matmat_accesses, 10 iterator_types = ["parallel", "parallel", "reduction"] 11} 12 13// CHECK-LABEL: func @matmul 14// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<2x4xf32>, 15// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: vector<4x3xf32>, 16// CHECK-SAME: %[[C:[a-zA-Z0-9]*]]: vector<2x3xf32> 17// CHECK-DAG: %[[vcst:.*]] = arith.constant dense<0.000000e+00> : vector<8xf32> 18// CHECK-DAG: %[[vcst_0:.*]] = arith.constant dense<0.000000e+00> : vector<12xf32> 19// CHECK-DAG: %[[vcst_1:.*]] = arith.constant dense<0.000000e+00> : vector<2x3xf32> 20// CHECK: %[[a0:.*]] = vector.extract %[[A]][0] : vector<4xf32> from vector<2x4xf32> 21// CHECK: %[[a1:.*]] = vector.insert_strided_slice %[[a0]], %[[vcst]] {offsets = [0], strides = [1]} : vector<4xf32> into vector<8xf32> 22// CHECK: %[[a2:.*]] = vector.extract %[[A]][1] : vector<4xf32> from vector<2x4xf32> 23// CHECK: %[[a3:.*]] = vector.insert_strided_slice %[[a2]], %[[a1]] {offsets = [4], strides = [1]} : vector<4xf32> into vector<8xf32> 24// CHECK: %[[b0:.*]] = vector.extract %[[B]][0] : vector<3xf32> from vector<4x3xf32> 25// CHECK: %[[b1:.*]] = vector.insert_strided_slice %[[b0]], %[[vcst_0]] {offsets = [0], strides = [1]} : vector<3xf32> into vector<12xf32> 26// CHECK: %[[b2:.*]] = vector.extract %[[B]][1] : vector<3xf32> from vector<4x3xf32> 27// CHECK: %[[b3:.*]] = vector.insert_strided_slice %[[b2]], %[[b1]] {offsets = [3], strides = [1]} : vector<3xf32> into vector<12xf32> 28// CHECK: %[[b4:.*]] = vector.extract %[[B]][2] : vector<3xf32> from vector<4x3xf32> 29// CHECK: %[[b5:.*]] = vector.insert_strided_slice %[[b4]], %[[b3]] {offsets = [6], strides = [1]} : vector<3xf32> into vector<12xf32> 30// CHECK: %[[b6:.*]] = vector.extract %[[B]][3] : vector<3xf32> from vector<4x3xf32> 31// CHECK: %[[b7:.*]] = vector.insert_strided_slice %[[b6]], %[[b5]] {offsets = [9], strides = [1]} : vector<3xf32> into vector<12xf32> 32// CHECK: %[[mm1:.*]] = vector.matrix_multiply %[[a3]], %[[b7]] {lhs_columns = 4 : i32, lhs_rows = 2 : i32, rhs_columns = 3 : i32} : (vector<8xf32>, vector<12xf32>) -> vector<6xf32> 33// CHECK: %[[mm2:.*]] = vector.extract_strided_slice %[[mm1]] {offsets = [0], sizes = [3], strides = [1]} : vector<6xf32> to vector<3xf32> 34// CHECK: %[[mm3:.*]] = vector.insert %[[mm2]], %[[vcst_1]] [0] : vector<3xf32> into vector<2x3xf32> 35// CHECK: %[[mm4:.*]] = vector.extract_strided_slice %[[mm1]] {offsets = [3], sizes = [3], strides = [1]} : vector<6xf32> to vector<3xf32> 36// CHECK: %[[mm5:.*]] = vector.insert %[[mm4]], %[[mm3]] [1] : vector<3xf32> into vector<2x3xf32> 37// CHECK: %[[mm6:.*]] = arith.addf %[[C]], %[[mm5]] : vector<2x3xf32> 38func.func @matmul(%arg0: vector<2x4xf32>, 39 %arg1: vector<4x3xf32>, 40 %arg2: vector<2x3xf32>) -> vector<2x3xf32> { 41 %0 = vector.contract #matmat_trait %arg0, %arg1, %arg2 42 : vector<2x4xf32>, vector<4x3xf32> into vector<2x3xf32> 43 return %0 : vector<2x3xf32> 44} 45 46// CHECK-LABEL: func @matmul_scalable 47// CHECK-NOT: vector.matrix_multiply 48func.func @matmul_scalable(%arg0: vector<2x4xf32>, 49 %arg1: vector<4x[3]xf32>, 50 %arg2: vector<2x[3]xf32>) -> vector<2x[3]xf32> { 51 %0 = vector.contract #matmat_trait %arg0, %arg1, %arg2 52 : vector<2x4xf32>, vector<4x[3]xf32> into vector<2x[3]xf32> 53 return %0 : vector<2x[3]xf32> 54} 55 56module attributes {transform.with_named_sequence} { 57 transform.named_sequence @__transform_main(%module_op: !transform.any_op {transform.readonly}) { 58 %f = transform.structured.match ops{["func.func"]} in %module_op 59 : (!transform.any_op) -> !transform.any_op 60 61 transform.apply_patterns to %f { 62 transform.apply_patterns.vector.lower_contraction lowering_strategy = "matmulintrinsics" 63 } : !transform.any_op 64 65 transform.apply_patterns to %f { 66 transform.apply_patterns.vector.lower_shape_cast 67 } : !transform.any_op 68 transform.yield 69 } 70} 71