mlir/test/Integration/Dialect/Vector/CPU/test-matrix-multiply-col.mlir

   1 // RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm -reconcile-unrealized-casts | \
   2 // RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
   3 // RUN:   -lower-matrix-intrinsics -matrix-allow-contract -matrix-default-layout=column-major \
   4 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
   5 // RUN: FileCheck %s
   6
   7 func @entry() {
   8   %f0 = arith.constant 0.0: f64
   9   %f1 = arith.constant 1.0: f64
  10   %f2 = arith.constant 2.0: f64
  11   %f3 = arith.constant 3.0: f64
  12   %f4 = arith.constant 4.0: f64
  13   %f5 = arith.constant 5.0: f64
  14   %f6 = arith.constant 6.0: f64
  15   %f7 = arith.constant 7.0: f64
  16
  17   // Construct test vectors.
  18   %0 = vector.broadcast %f0 : f64 to vector<4xf64>
  19   %1 = vector.insert %f1, %0[1] : f64 into vector<4xf64>
  20   %2 = vector.insert %f2, %1[2] : f64 into vector<4xf64>
  21   %a = vector.insert %f3, %2[3] : f64 into vector<4xf64>
  22   %3 = vector.broadcast %f4 : f64 to vector<4xf64>
  23   %4 = vector.insert %f5, %3[1] : f64 into vector<4xf64>
  24   %5 = vector.insert %f6, %4[2] : f64 into vector<4xf64>
  25   %b = vector.insert %f7, %5[3] : f64 into vector<4xf64>
  26
  27   vector.print %a : vector<4xf64>
  28   vector.print %b : vector<4xf64>
  29   //
  30   // test vectors:
  31   //
  32   // CHECK: ( 0, 1, 2, 3 )
  33   // CHECK: ( 4, 5, 6, 7 )
  34
  35   // Performs matrix x matrix, interpreting the vectors as
  36   // flattened column-major 2-D matrices.
  37   //
  38   // ( 0, 2 )     (4, 6)     ( 10, 14 )      | /|
  39   //           x          =                  |/ | column-major!
  40   // ( 1, 3 )     (5, 7)     ( 19, 27 )
  41   //
  42   %c = vector.matrix_multiply %a, %b
  43       { lhs_rows = 2: i32, lhs_columns = 2: i32 , rhs_columns = 2: i32 }
  44       : (vector<4xf64>, vector<4xf64>) -> vector<4xf64>
  45
  46   vector.print %c : vector<4xf64>
  47   //
  48   // matrix x matrix:
  49   //
  50   // CHECK: ( 10, 19, 14, 27 )
  51
  52   return
  53 }