mlir/test/Integration/Dialect/Vector/CPU/test-flat-transpose-row.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=row-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   %6 = vector.broadcast %f0 : f64 to vector<6xf64>
  27   %7 = vector.insert %f1, %6[1] : f64 into vector<6xf64>
  28   %8 = vector.insert %f2, %7[2] : f64 into vector<6xf64>
  29   %9 = vector.insert %f3, %8[3] : f64 into vector<6xf64>
  30   %10 = vector.insert %f4, %9[4] : f64 into vector<6xf64>
  31   %c = vector.insert %f5, %10[5] : f64 into vector<6xf64>
  32
  33   vector.print %a : vector<4xf64>
  34   vector.print %b : vector<4xf64>
  35   vector.print %c : vector<6xf64>
  36   //
  37   // Test vectors:
  38   //
  39   // CHECK: ( 0, 1, 2, 3 )
  40   // CHECK: ( 4, 5, 6, 7 )
  41   // CHECK: ( 0, 1, 2, 3, 4, 5 )
  42
  43   // Performs matrix transpositions interpreting the vectors as
  44   // flattened row-major 2-D matrices.
  45   //
  46   // ( 0, 1 )       ( 0, 2 )
  47   // ( 2, 3 )    -> ( 1, 3 )
  48   //
  49   // ( 4, 5 )       ( 4, 6 )
  50   // ( 6, 7 )    -> ( 5, 7 )
  51   //
  52   // ( 0, 1, 2 )    ( 0, 3 )
  53   // ( 3, 4, 5 ) -> ( 1, 4 )
  54   //                ( 2, 5 )
  55   //
  56   // ( 0, 1 )        ( 0, 2, 4 )
  57   // ( 2, 3 )    ->  ( 1, 3, 5 )
  58   // ( 4, 5 )
  59   //
  60   %d = vector.flat_transpose %a { rows = 2: i32, columns = 2: i32 } : vector<4xf64> -> vector<4xf64>
  61   %e = vector.flat_transpose %b { rows = 2: i32, columns = 2: i32 } : vector<4xf64> -> vector<4xf64>
  62   %f = vector.flat_transpose %c { rows = 2: i32, columns = 3: i32 } : vector<6xf64> -> vector<6xf64>
  63   %g = vector.flat_transpose %c { rows = 3: i32, columns = 2: i32 } : vector<6xf64> -> vector<6xf64>
  64
  65   vector.print %d : vector<4xf64>
  66   vector.print %e : vector<4xf64>
  67   vector.print %f : vector<6xf64>
  68   vector.print %g : vector<6xf64>
  69   //
  70   // Transposed results:
  71   //
  72   // CHECK: ( 0, 2, 1, 3 )
  73   // CHECK: ( 4, 6, 5, 7 )
  74   // CHECK: ( 0, 3, 1, 4, 2, 5 )
  75   // CHECK: ( 0, 2, 4, 1, 3, 5 )
  76
  77   return
  78 }