maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][Linalg] Add multi-result op cast test. 

https://reviews.llvm.org/D121369 fixed an issue with canonicalizing a
linalg op producer with a cast op consumer. Adding a test to verify
that change.

Reviewed By: hanchung

Differential Revision: https://reviews.llvm.org/D121648
This commit is contained in:
MaheshRavishankar 2022-03-16 04:51:17 +00:00 committed by Mahesh Ravishankar
parent 0bf9aabd09
commit 4b644fca08
1 changed files with 34 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
mlir/test/Dialect/Linalg/canonicalize.mlir
View File

@ -901,3 +901,37 @@ func @fold_conv_op_with_cast_consumer(%arg0 : tensor<?x?x?x?xf32>,
// CHECK-SAME: outs(%[[OUT_CAST]] :
// CHECK: %[[RESULT_CAST:.+]] = tensor.cast %[[CONV]]
// CHECK: return %[[CONV]], %[[RESULT_CAST]]
// -----
func @fold_multi_use_generic_op_with_consumer(%arg0 : tensor<?x?x?xf32>) -> (tensor<?x?x?xf32>, tensor<2x3x4xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%c2 = arith.constant 2 : index
%d0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
%d1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
%d2 = tensor.dim %arg0, %c2 : tensor<?x?x?xf32>
%init1 = linalg.init_tensor [%d1, %d2, %d0] : tensor<?x?x?xf32>
%init2 = linalg.init_tensor [%d2, %d1, %d0] : tensor<?x?x?xf32>
%0:2 = linalg.generic {
iterator_types = ["parallel", "parallel", "parallel"],
indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
affine_map<(d0, d1, d2) -> (d1, d2, d0)>,
affine_map<(d0, d1, d2) -> (d2, d1, d0)>]}
ins(%arg0 : tensor<?x?x?xf32>) outs(%init1, %init2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {
^bb0(%b0 : f32, %b1 : f32, %b2 : f32) :
linalg.yield %b0, %b0 : f32, f32
} -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)
%1 = tensor.cast %0#1 : tensor<?x?x?xf32> to tensor<2x3x4xf32>
return %0#0, %1 : tensor<?x?x?xf32>, tensor<2x3x4xf32>
}
// CHECK: func @fold_multi_use_generic_op_with_consumer
// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?xf32>
// CHECK-DAG: %[[INIT1:.+]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
// CHECK-DAG: %[[CAST:.+]] = tensor.cast %[[ARG0]] : tensor<?x?x?xf32> to tensor<4x3x2xf32>
// CHECK-DAG: %[[INIT2:.+]] = linalg.init_tensor [3, 2, 4] : tensor<3x2x4xf32>
// CHECK: %[[GENERIC:.+]]:2 = linalg.generic
// CHECK-SAME: ins(%[[CAST]] :
// CHECK-SAME: outs(%[[INIT2]], %[[INIT1]] :
// CHECK: %[[RETURN_CAST:.+]] = tensor.cast %[[GENERIC]]#0 : tensor<3x2x4xf32> to tensor<?x?x?xf32>
// CHECK: return %[[RETURN_CAST]], %[[GENERIC]]#1