1// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-flat-transpose=1 | FileCheck %s 2 3// Tests for lowering 2-D vector.transpose into vector.flat_transpose. 4// 5// TODO: having ShapeCastOp2DDownCastRewritePattern and 6// ShapeCastOp2DUpCastRewritePattern too early in the greedy rewriting 7// patterns misses opportunities to fold shape casts! 8 9// No shape cast folding expected. 10// 11// CHECK-LABEL: func @transpose44_44( 12// CHECK-SAME: %[[A:.*]]: vector<4x4xf32> 13// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<4x4xf32> 14// CHECK: %[[T8:.*]] = vector.flat_transpose %{{.*}} {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32> 15// CHECK: %[[T9:.*]] = vector.extract_strided_slice %[[T8]] {offsets = [0], sizes = [4], strides = [1]} : vector<16xf32> to vector<4xf32> 16// 17func @transpose44_44(%arg0: vector<4x4xf32>) -> vector<4x4xf32> { 18 %0 = vector.transpose %arg0, [1, 0] : vector<4x4xf32> to vector<4x4xf32> 19 return %0 : vector<4x4xf32> 20} 21 22// Folds preceding shape cast as expected, 23// no following shape cast folding expected. 24// 25// CHECK-LABEL: func @transpose16_44( 26// CHECK-SAME: %[[A:.*]]: vector<16xf32> 27// CHECK: %[[T0:.*]] = vector.flat_transpose %[[A]] {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32> 28// CHECK: %[[T1:.*]] = vector.extract_strided_slice %[[T0]] {offsets = [0], sizes = [4], strides = [1]} : vector<16xf32> to vector<4xf32> 29// 30func @transpose16_44(%arg0: vector<16xf32>) -> vector<4x4xf32> { 31 %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<4x4xf32> 32 %1 = vector.transpose %0, [1, 0] : vector<4x4xf32> to vector<4x4xf32> 33 return %1 : vector<4x4xf32> 34} 35 36// No preceding shape cast folding expected, 37// but FAILS to fold following cast. 38// 39// CHECK-LABEL: func @transpose44_16( 40// CHECK-SAME: %[[A:.*]]: vector<4x4xf32> 41// CHECK: %[[T0:.*]] = vector.extract %[[A]][0] : vector<4x4xf32> 42// CHECK: %[[T8:.*]] = vector.flat_transpose %{{.*}} {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32> 43func @transpose44_16(%arg0: vector<4x4xf32>) -> vector<16xf32> { 44 %0 = vector.transpose %arg0, [1, 0] : vector<4x4xf32> to vector<4x4xf32> 45 %1 = vector.shape_cast %0 : vector<4x4xf32> to vector<16xf32> 46 return %1 : vector<16xf32> 47} 48 49// Folds preceding shape cast as expected, 50// but FAILS to fold following cast. 51// 52// CHECK-LABEL: func @transpose16_16( 53// CHECK-SAME: %[[A:.*]]: vector<16xf32> 54// CHECK: %[[T0:.*]] = vector.flat_transpose %[[A]] {columns = 4 : i32, rows = 4 : i32} : vector<16xf32> -> vector<16xf32> 55// 56func @transpose16_16(%arg0: vector<16xf32>) -> vector<16xf32> { 57 %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<4x4xf32> 58 %1 = vector.transpose %0, [1, 0] : vector<4x4xf32> to vector<4x4xf32> 59 %2 = vector.shape_cast %1 : vector<4x4xf32> to vector<16xf32> 60 return %2 : vector<16xf32> 61} 62