forked from OSchip/llvm-project
29 lines
1.8 KiB
MLIR
29 lines
1.8 KiB
MLIR
// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" -mlir-disable-threading=true | FileCheck %s
|
|
|
|
// CHECK-LABEL: func @matmul_tensors(
|
|
// CHECK-SAME: %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
|
|
// CHECK-SAME: %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
|
|
func @matmul_tensors(
|
|
%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
|
|
-> tensor<?x?xf32> {
|
|
// CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) {
|
|
// CHECK: %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) {
|
|
// CHECK: %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) {
|
|
// CHECK: %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
|
|
// CHECK: %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
|
|
// CHECK: %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
|
|
// CHECK: %[[sTD:.*]] = linalg.matmul ins(%[[sTA]], %[[sTB]] : tensor<?x?xf32>, tensor<?x?xf32>)
|
|
// CHECK-SAME: init(%[[sTC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
|
|
// CHECK: %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}] : tensor<?x?xf32> into tensor<?x?xf32>
|
|
// CHECK: scf.yield %[[TD]] : tensor<?x?xf32>
|
|
// CHECK: scf.yield %[[TD2]] : tensor<?x?xf32>
|
|
// CHECK: scf.yield %[[TD1]] : tensor<?x?xf32>
|
|
%0 = linalg.matmul ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
|
|
init(%arg2: tensor<?x?xf32>)
|
|
-> tensor<?x?xf32>
|
|
|
|
// CHECK: return %[[TD0]] : tensor<?x?xf32>
|
|
return %0 : tensor<?x?xf32>
|
|
}
|