mlir/test/Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir

   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>
  38 func.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
  48 func.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
  56 module 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 }