colin
/
onnx-mlir
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
onnx-mlir/test/mlir/krnl/krnl_disconnect_dim_from_al...

44 lines
1.7 KiB
MLIR
Raw Normal View History

Add krnl.dim operation and lowering pass (#261) * Reorganize main function. * Follow review comments. * Emit constants are globals in Krnl and LLVM dialects. * Add krnl.dim op. * Add test with alloc with map. * Code clean-up. * Code clean-up. * Add comment for function. * Update DisconnectKrnlDimFromAlloc.cpp
2020-08-15 00:59:33 +08:00
// RUN: onnx-mlir-opt --disconnect-dims %s -split-input-file | FileCheck %s
func @test_krnl_dim_lowering(%arg0: memref<?x?xf32>) -> index {
%c1 = constant 1 : index
%c0 = constant 0 : index
%0 = dim %arg0, %c0 : memref<?x?xf32>
%1 = alloc(%0) : memref<?x10xf32>
%d0 = "krnl.dim"(%1, %c0) : (memref<?x10xf32>, index) -> index
%d1 = "krnl.dim"(%1, %c1) : (memref<?x10xf32>, index) -> index
%e = addi %d0, %d1 : index
return %e : index
// CHECK-LABEL: test_krnl_dim_lowering
// CHECK: [[CONST0:%.+]] = constant 0 : index
// CHECK: [[CONST10:%.+]] = constant 10 : index
// CHECK: [[DIM:%.+]] = dim %arg0, [[CONST0]] : memref<?x?xf32>
// CHECK: [[ALLOC:%.+]] = alloc([[DIM]]) : memref<?x10xf32>
// CHECK: [[SUM:%.+]] = addi [[DIM]], [[CONST10]] : index
// CHECK: return [[SUM]] : index
}
// -----
#map = affine_map<(d0, d1) -> (d1, d0)>
func @test_krnl_dim_lowering_with_map(%arg0: memref<?x?xf32>) -> index {
%c1 = constant 1 : index
%c0 = constant 0 : index
%0 = dim %arg0, %c0 : memref<?x?xf32>
%1 = alloc(%0) : memref<?x10xf32, #map>
%d0 = "krnl.dim"(%1, %c0) : (memref<?x10xf32, #map>, index) -> index
%d1 = "krnl.dim"(%1, %c1) : (memref<?x10xf32, #map>, index) -> index
%e = addi %d0, %d1 : index
return %e : index
// CHECK: [[MAP:#.+]] = affine_map<(d0, d1) -> (d1, d0)>
// CHECK-LABEL: test_krnl_dim_lowering_with_map
// CHECK: [[CONST0:%.+]] = constant 0 : index
// CHECK: [[CONST10:%.+]] = constant 10 : index
// CHECK: [[DIM:%.+]] = dim %arg0, [[CONST0]] : memref<?x?xf32>
// CHECK: [[ALLOC:%.+]] = alloc([[DIM]]) : memref<?x10xf32, [[MAP]]>
// CHECK: [[SUM:%.+]] = addi [[DIM]], [[CONST10]] : index
// CHECK: return [[SUM]] : index
}