[Mlir-commits] [mlir] [mlir][tosa] Lower MATMUL_T to Linalg (PR #222661)

[Jack Frankland]

================
@@ -95,6 +95,328 @@ func.func @matmul_dyn_output(%arg0: tensor<1x1x8xf32>, %arg1: tensor<1x8x1xf32>)
 
 // -----
 
+// CHECK-LABEL: @matmul_t
+func.func @matmul_t(%arg0: tensor<2x5x3xf32>, %arg1: tensor<2x6x3xf32>) -> tensor<2x5x6xf32> {
+  // CHECK: %[[C0:.+]] = arith.constant 0
+  // CHECK: %[[INIT:.+]] = tensor.empty()
+  // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf32>, tensor<2x6x3xf32>) outs(%[[FILLED]] : tensor<2x5x6xf32>) -> tensor<2x5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf32>, tensor<2x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x5x6xf32>
+  return %0 : tensor<2x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_f16_fp8e4m3_to_bf16
+func.func @matmul_t_f16_fp8e4m3_to_bf16(%arg0: tensor<2x5x3xf16>, %arg1: tensor<2x6x3xf8E4M3FN>) -> tensor<2x5x6xbf16> {
+  // CHECK: %[[FILLED:.+]] = linalg.fill {{.*}} -> tensor<2x5x6xbf16>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf16>, tensor<2x6x3xf8E4M3FN>) outs(%[[FILLED]] : tensor<2x5x6xbf16>) -> tensor<2x5x6xbf16>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E4M3FN>}> : () -> tensor<1xf8E4M3FN>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf16>, tensor<2x6x3xf8E4M3FN>, tensor<1xf16>, tensor<1xf8E4M3FN>) -> tensor<2x5x6xbf16>
+  return %0 : tensor<2x5x6xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_fp8e5m2_f16_to_bf16
+func.func @matmul_t_fp8e5m2_f16_to_bf16(%arg0: tensor<2x5x3xf8E5M2>, %arg1: tensor<2x6x3xf16>) -> tensor<2x5x6xbf16> {
+  // CHECK: %[[FILLED:.+]] = linalg.fill {{.*}} -> tensor<2x5x6xbf16>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%arg0, %arg1 : tensor<2x5x3xf8E5M2>, tensor<2x6x3xf16>) outs(%[[FILLED]] : tensor<2x5x6xbf16>) -> tensor<2x5x6xbf16>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf8E5M2>}> : () -> tensor<1xf8E5M2>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x5x3xf8E5M2>, tensor<2x6x3xf16>, tensor<1xf8E5M2>, tensor<1xf16>) -> tensor<2x5x6xbf16>
+  return %0 : tensor<2x5x6xbf16>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_broadcast
+func.func @matmul_t_broadcast(%arg0: tensor<4x5x3xf32>, %arg1: tensor<1x6x3xf32>) -> tensor<4x5x6xf32> {
+  // CHECK: %[[SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0] [1, 6, 3] [1, 1, 1] : tensor<1x6x3xf32> to tensor<6x3xf32>
+  // CHECK: %[[BCAST_INIT:.+]] = tensor.empty() : tensor<4x6x3xf32>
+  // CHECK: %[[BCAST:.+]] = linalg.broadcast ins(%[[SLICE]] : tensor<6x3xf32>) outs(%[[BCAST_INIT]] : tensor<4x6x3xf32>) dimensions = [0]
+  // CHECK: %[[C0:.+]] = arith.constant 0
+  // CHECK: %[[INIT:.+]] = tensor.empty()
+  // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%arg0, %[[BCAST]] : tensor<4x5x3xf32>, tensor<4x6x3xf32>) outs(%[[FILLED]] : tensor<4x5x6xf32>) -> tensor<4x5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x5x3xf32>, tensor<1x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x5x6xf32>
+  return %0 : tensor<4x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_broadcast_both_rank4
+func.func @matmul_t_broadcast_both_rank4(%arg0: tensor<2x1x5x3xf32>, %arg1: tensor<1x3x6x3xf32>) -> tensor<2x3x5x6xf32> {
+  // CHECK: %[[B_SLICE:.+]] = tensor.extract_slice %arg1[0, 0, 0, 0] [1, 3, 6, 3] [1, 1, 1, 1] : tensor<1x3x6x3xf32> to tensor<3x6x3xf32>
+  // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%[[B_SLICE]] : tensor<3x6x3xf32>) {{.*}} dimensions = [0]
+  // CHECK: %[[A_SLICE:.+]] = tensor.extract_slice %arg0[0, 0, 0, 0] [2, 1, 5, 3] [1, 1, 1, 1] : tensor<2x1x5x3xf32> to tensor<2x5x3xf32>
+  // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%[[A_SLICE]] : tensor<2x5x3xf32>) {{.*}} dimensions = [1]
+  // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x5x3xf32> into tensor<6x5x3xf32>
+  // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<2x3x6x3xf32> into tensor<6x6x3xf32>
+  // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<6x5x3xf32>, tensor<6x6x3xf32>)
+  // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [2, 3, 5, 6] : tensor<6x5x6xf32> into tensor<2x3x5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<2x1x5x3xf32>, tensor<1x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<2x3x5x6xf32>
+  return %0 : tensor<2x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_a_missing_batch_dim
+func.func @matmul_t_a_missing_batch_dim(%arg0: tensor<3x5x3xf32>, %arg1: tensor<4x3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+  // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<3x5x3xf32>) {{.*}} dimensions = [0]
+  // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+  // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+  // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+  // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<3x5x3xf32>, tensor<4x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+  return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_b_missing_batch_dim
+func.func @matmul_t_b_missing_batch_dim(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+  // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<3x6x3xf32>) {{.*}} dimensions = [0]
+  // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+  // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+  // CHECK: %[[MATMUL:.+]] = linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+  // CHECK: tensor.expand_shape %[[MATMUL]] {{\[}}[0, 1], [2], [3]] output_shape [4, 3, 5, 6] : tensor<12x5x6xf32> into tensor<4x3x5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x3x5x3xf32>, tensor<3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+  return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2_a
+func.func @matmul_t_rank2_a(%arg0: tensor<5x3xf32>, %arg1: tensor<4x3x6x3xf32>) -> tensor<4x3x5x6xf32> {
+  // CHECK: %[[A_BCAST:.+]] = linalg.broadcast ins(%arg0 : tensor<5x3xf32>) {{.*}} dimensions = [0, 1]
+  // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %[[A_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+  // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %arg1 {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xf32>, tensor<4x3x6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+  return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2_b
+func.func @matmul_t_rank2_b(%arg0: tensor<4x3x5x3xf32>, %arg1: tensor<6x3xf32>) -> tensor<4x3x5x6xf32> {
+  // CHECK: %[[B_BCAST:.+]] = linalg.broadcast ins(%arg1 : tensor<6x3xf32>) {{.*}} dimensions = [0, 1]
+  // CHECK: %[[A_COLLAPSED:.+]] = tensor.collapse_shape %arg0 {{\[}}[0, 1], [2], [3]] : tensor<4x3x5x3xf32> into tensor<12x5x3xf32>
+  // CHECK: %[[B_COLLAPSED:.+]] = tensor.collapse_shape %[[B_BCAST]] {{\[}}[0, 1], [2], [3]] : tensor<4x3x6x3xf32> into tensor<12x6x3xf32>
+  // CHECK: linalg.batch_matmul indexing_maps =
+  // CHECK-SAME: ins(%[[A_COLLAPSED]], %[[B_COLLAPSED]] : tensor<12x5x3xf32>, tensor<12x6x3xf32>)
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<4x3x5x3xf32>, tensor<6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<4x3x5x6xf32>
+  return %0 : tensor<4x3x5x6xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @matmul_t_rank2
+func.func @matmul_t_rank2(%arg0: tensor<5x3xf32>, %arg1: tensor<6x3xf32>) -> tensor<5x6xf32> {
+  // CHECK: %[[C0:.+]] = arith.constant 0
+  // CHECK: %[[INIT:.+]] = tensor.empty()
+  // CHECK: %[[FILLED:.+]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+  // CHECK: linalg.matmul indexing_maps =
+  // CHECK-SAME: ins(%arg0, %arg1 : tensor<5x3xf32>, tensor<6x3xf32>) outs(%[[FILLED]] : tensor<5x6xf32>) -> tensor<5x6xf32>
+  %a_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %b_zp = "tosa.const"() <{values = dense<0.0> : tensor<1xf32>}> : () -> tensor<1xf32>
+  %0 = tosa.matmul_t %arg0, %arg1, %a_zp, %b_zp : (tensor<5x3xf32>, tensor<6x3xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<5x6xf32>
----------------
FranklandJack wrote:

Yeah added ;)

https://github.com/llvm/llvm-project/pull/222661