mlir/test/Integration/Dialect/Vector/CPU/test-transpose.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:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
   4 // RUN: FileCheck %s
   5
   6 func @entry() {
   7   %f0 = arith.constant 0.0: f32
   8   %f1 = arith.constant 1.0: f32
   9   %f2 = arith.constant 2.0: f32
  10   %f3 = arith.constant 3.0: f32
  11   %f4 = arith.constant 4.0: f32
  12   %f5 = arith.constant 5.0: f32
  13   %f6 = arith.constant 6.0: f32
  14   %f7 = arith.constant 7.0: f32
  15   %f8 = arith.constant 8.0: f32
  16
  17   // Construct test vectors and matrices.
  18   %0 = vector.broadcast %f1 : f32 to vector<2xf32>
  19   %a = vector.insert %f2, %0[1] : f32 into vector<2xf32>
  20   %1 = vector.broadcast %f3 : f32 to vector<2xf32>
  21   %b = vector.insert %f4, %1[1] : f32 into vector<2xf32>
  22   %2 = vector.broadcast %f5 : f32 to vector<2xf32>
  23   %c = vector.insert %f6, %2[1] : f32 into vector<2xf32>
  24   %3 = vector.broadcast %f7 : f32 to vector<2xf32>
  25   %d = vector.insert %f8, %3[1] : f32 into vector<2xf32>
  26   %4 = vector.broadcast %f0 : f32 to vector<2x2xf32>
  27   %5 = vector.insert %a, %4[0] : vector<2xf32> into vector<2x2xf32>
  28   %A = vector.insert %b, %5[1] : vector<2xf32> into vector<2x2xf32>
  29   %6 = vector.broadcast %f0 : f32 to vector<2x2xf32>
  30   %7 = vector.insert %c, %6[0] : vector<2xf32> into vector<2x2xf32>
  31   %B = vector.insert %d, %7[1] : vector<2xf32> into vector<2x2xf32>
  32   %8 = vector.broadcast %f0 : f32 to vector<3x2xf32>
  33   %9 = vector.insert %a, %8[0] : vector<2xf32> into vector<3x2xf32>
  34   %10 = vector.insert %b, %9[1] : vector<2xf32> into vector<3x2xf32>
  35   %C = vector.insert %c, %10[2] : vector<2xf32> into vector<3x2xf32>
  36   %cst = arith.constant dense<0.000000e+00> : vector<2x4xf32>
  37   %11 = vector.insert_strided_slice %A, %cst {offsets = [0, 0], strides = [1, 1]} : vector<2x2xf32> into vector<2x4xf32>
  38   %D = vector.insert_strided_slice %B, %11 {offsets = [0, 2], strides = [1, 1]} : vector<2x2xf32> into vector<2x4xf32>
  39
  40   vector.print %A : vector<2x2xf32>
  41   vector.print %B : vector<2x2xf32>
  42   vector.print %C : vector<3x2xf32>
  43   vector.print %D : vector<2x4xf32>
  44   //
  45   // test matrices:
  46   //
  47   // CHECK: ( ( 1, 2 ), ( 3, 4 ) )
  48   // CHECK: ( ( 5, 6 ), ( 7, 8 ) )
  49   // CHECK: ( ( 1, 2 ), ( 3, 4 ), ( 5, 6 ) )
  50   // CHECK: ( ( 1, 2, 5, 6 ), ( 3, 4, 7, 8 ) )
  51
  52   %tA = vector.transpose %A, [1, 0] : vector<2x2xf32> to vector<2x2xf32>
  53   %tB = vector.transpose %B, [1, 0] : vector<2x2xf32> to vector<2x2xf32>
  54   %tC = vector.transpose %C, [1, 0] : vector<3x2xf32> to vector<2x3xf32>
  55   %tD = vector.transpose %D, [1, 0] : vector<2x4xf32> to vector<4x2xf32>
  56
  57   vector.print %tA : vector<2x2xf32>
  58   vector.print %tB : vector<2x2xf32>
  59   vector.print %tC : vector<2x3xf32>
  60   vector.print %tD : vector<4x2xf32>
  61   //
  62   // transposed matrices:
  63   //
  64   // CHECK: ( ( 1, 3 ), ( 2, 4 ) )
  65   // CHECK: ( ( 5, 7 ), ( 6, 8 ) )
  66   // CHECK: ( ( 1, 3, 5 ), ( 2, 4, 6 ) )
  67   // CHECK: ( ( 1, 3 ), ( 2, 4 ), ( 5, 7 ), ( 6, 8 ) )
  68
  69   %idD = vector.transpose %D, [0, 1] : vector<2x4xf32> to vector<2x4xf32>
  70   %ttD = vector.transpose %tD, [1, 0] : vector<4x2xf32> to vector<2x4xf32>
  71
  72   vector.print %idD : vector<2x4xf32>
  73   vector.print %ttD : vector<2x4xf32>
  74   //
  75   // back to original after transpose matrices:
  76   //
  77   // CHECK: ( ( 1, 2, 5, 6 ), ( 3, 4, 7, 8 ) )
  78   // CHECK: ( ( 1, 2, 5, 6 ), ( 3, 4, 7, 8 ) )
  79
  80   // Construct test tensor.
  81   %p = vector.broadcast %f1 : f32 to vector<2x2x2xf32>
  82   %q = vector.insert %f2, %p[0, 0, 1] : f32 into vector<2x2x2xf32>
  83   %r = vector.insert %f3, %q[0, 1, 0] : f32 into vector<2x2x2xf32>
  84   %s = vector.insert %f4, %r[0, 1, 1] : f32 into vector<2x2x2xf32>
  85   %t = vector.insert %f5, %s[1, 0, 0] : f32 into vector<2x2x2xf32>
  86   %u = vector.insert %f6, %t[1, 0, 1] : f32 into vector<2x2x2xf32>
  87   %v = vector.insert %f7, %u[1, 1, 0] : f32 into vector<2x2x2xf32>
  88   %w = vector.insert %f8, %v[1, 1, 1] : f32 into vector<2x2x2xf32>
  89
  90   vector.print %w : vector<2x2x2xf32>
  91   //
  92   // test tensors:
  93   //
  94   // CHECK: ( ( ( 1, 2 ), ( 3, 4 ) ), ( ( 5, 6 ), ( 7, 8 ) ) )
  95
  96   %tP = vector.transpose %w, [0, 1, 2] : vector<2x2x2xf32> to vector<2x2x2xf32>
  97   %tQ = vector.transpose %w, [0, 2, 1] : vector<2x2x2xf32> to vector<2x2x2xf32>
  98   %tR = vector.transpose %w, [1, 0, 2] : vector<2x2x2xf32> to vector<2x2x2xf32>
  99   %tS = vector.transpose %w, [2, 0, 1] : vector<2x2x2xf32> to vector<2x2x2xf32>
 100   %tT = vector.transpose %w, [1, 2, 0] : vector<2x2x2xf32> to vector<2x2x2xf32>
 101   %tU = vector.transpose %w, [2, 1, 0] : vector<2x2x2xf32> to vector<2x2x2xf32>
 102
 103   vector.print %tP : vector<2x2x2xf32>
 104   vector.print %tQ : vector<2x2x2xf32>
 105   vector.print %tR : vector<2x2x2xf32>
 106   vector.print %tS : vector<2x2x2xf32>
 107   vector.print %tT : vector<2x2x2xf32>
 108   vector.print %tU : vector<2x2x2xf32>
 109   //
 110   // transposed tensors:
 111   //
 112   // CHECK: ( ( ( 1, 2 ), ( 3, 4 ) ), ( ( 5, 6 ), ( 7, 8 ) ) )
 113   // CHECK: ( ( ( 1, 3 ), ( 2, 4 ) ), ( ( 5, 7 ), ( 6, 8 ) ) )
 114   // CHECK: ( ( ( 1, 2 ), ( 5, 6 ) ), ( ( 3, 4 ), ( 7, 8 ) ) )
 115   // CHECK: ( ( ( 1, 3 ), ( 5, 7 ) ), ( ( 2, 4 ), ( 6, 8 ) ) )
 116   // CHECK: ( ( ( 1, 5 ), ( 2, 6 ) ), ( ( 3, 7 ), ( 4, 8 ) ) )
 117   // CHECK: ( ( ( 1, 5 ), ( 3, 7 ) ), ( ( 2, 6 ), ( 4, 8 ) ) )
 118
 119   return
 120 }