flang/test/HLFIR/matmul_transpose.fir

   1 // Test hlfir.matmul_transpose operation parse, verify (no errors), and unparse
   2
   3 // RUN: fir-opt %s | fir-opt | FileCheck %s
   4
   5 // arguments are expressions of known shape
   6 func.func @matmul_transpose0(%arg0: !hlfir.expr<2x2xi32>, %arg1: !hlfir.expr<2x2xi32>) {
   7   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<2x2xi32>, !hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
   8   return
   9 }
  10 // CHECK-LABEL: func.func @matmul_transpose0
  11 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x2xi32>,
  12 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<2x2xi32>) {
  13 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<2x2xi32>, !hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
  14 // CHECK-NEXT:    return
  15 // CHECK-NEXT:  }
  16
  17 // arguments are expressions of assumed shape
  18 func.func @matmul_transpose1(%arg0: !hlfir.expr<?x?xi32>, %arg1: !hlfir.expr<?x?xi32>) {
  19   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<?x?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
  20   return
  21 }
  22 // CHECK-LABEL: func.func @matmul_transpose1
  23 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?xi32>,
  24 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x?xi32>) {
  25 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?xi32>, !hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
  26 // CHECK-NEXT:    return
  27 // CHECK-NEXT:  }
  28
  29 // arguments are expressions where only some dimensions are known #1
  30 func.func @matmul_transpose2(%arg0: !hlfir.expr<?x2xi32>, %arg1: !hlfir.expr<?x2xi32>) {
  31   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<?x2xi32>, !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x2xi32>
  32   return
  33 }
  34 // CHECK-LABEL: func.func @matmul_transpose2
  35 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x2xi32>,
  36 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x2xi32>) {
  37 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x2xi32>, !hlfir.expr<?x2xi32>) -> !hlfir.expr<2x2xi32>
  38 // CHECK-NEXT:    return
  39 // CHECK-NEXT:  }
  40
  41 // arguments are expressions where only some dimensions are known #2
  42 func.func @matmul_transpose3(%arg0: !hlfir.expr<2x?xi32>, %arg1: !hlfir.expr<2x?xi32>) {
  43   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<2x?xi32>, !hlfir.expr<2x?xi32>) -> !hlfir.expr<?x?xi32>
  44   return
  45 }
  46 // CHECK-LABEL: func.func @matmul_transpose3
  47 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x?xi32>,
  48 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<2x?xi32>) {
  49 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<2x?xi32>, !hlfir.expr<2x?xi32>) -> !hlfir.expr<?x?xi32>
  50 // CHECK-NEXT:    return
  51 // CHECK-NEXT:  }
  52
  53 // arguments are logicals
  54 func.func @matmul_transpose4(%arg0: !hlfir.expr<?x?x!fir.logical<4>>, %arg1: !hlfir.expr<?x?x!fir.logical<4>>) {
  55   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<?x?x!fir.logical<4>>, !hlfir.expr<?x?x!fir.logical<4>>) -> !hlfir.expr<?x?x!fir.logical<4>>
  56   return
  57 }
  58 // CHECK-LABEL: func.func @matmul_transpose4
  59 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?x!fir.logical<4>>,
  60 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<?x?x!fir.logical<4>>) {
  61 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?x!fir.logical<4>>, !hlfir.expr<?x?x!fir.logical<4>>) -> !hlfir.expr<?x?x!fir.logical<4>>
  62 // CHECK-NEXT:    return
  63 // CHECK-NEXT:  }
  64
  65 // rhs is rank 1
  66 func.func @matmul_transpose6(%arg0: !hlfir.expr<?x?xi32>, %arg1: !hlfir.expr<?xi32>) {
  67   %res = hlfir.matmul_transpose %arg0 %arg1 : (!hlfir.expr<?x?xi32>, !hlfir.expr<?xi32>) -> !hlfir.expr<?xi32>
  68   return
  69 }
  70 // CHECK-LABEL: func.func @matmul_transpose6
  71 // CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?xi32>,
  72 // CHECK:           %[[ARG1:.*]]: !hlfir.expr<?xi32>) {
  73 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!hlfir.expr<?x?xi32>, !hlfir.expr<?xi32>) -> !hlfir.expr<?xi32>
  74 // CHECK-NEXT:    return
  75 // CHECK-NEXT:  }
  76
  77 // arguments are boxed arrays
  78 func.func @matmul_transpose7(%arg0: !fir.box<!fir.array<2x2xf32>>, %arg1: !fir.box<!fir.array<2x2xf32>>) {
  79   %res = hlfir.matmul_transpose %arg0 %arg1 : (!fir.box<!fir.array<2x2xf32>>, !fir.box<!fir.array<2x2xf32>>) -> !hlfir.expr<2x2xf32>
  80   return
  81 }
  82 // CHECK-LABEL: func.func @matmul_transpose7
  83 // CHECK:           %[[ARG0:.*]]: !fir.box<!fir.array<2x2xf32>>,
  84 // CHECK:           %[[ARG1:.*]]: !fir.box<!fir.array<2x2xf32>>) {
  85 // CHECK-NEXT:    %[[RES:.*]] = hlfir.matmul_transpose %[[ARG0]] %[[ARG1]] : (!fir.box<!fir.array<2x2xf32>>, !fir.box<!fir.array<2x2xf32>>) -> !hlfir.expr<2x2xf32>
  86 // CHECK-NEXT:    return
  87 // CHECK-NEXT:  }