dense_tensor.mlir 1.1 KB
Newer Older
1
// RUN: infrtexec -i %s | FileCheck %s
2
// CHECK-LABEL: dense_shape0
Y
Yan Chunwei 已提交
3 4
func @dense_shape0() {
  %shape = ts.build_shape [1:i64, 57:i64]
5
  %a = dt.create_uninit_tensor.f32 [12:i64, 23:i64] -> !infrt.dense_tensor<CPU, FP32, NCHW>
Y
Yan Chunwei 已提交
6

7
  Infrt.return
Y
Yan Chunwei 已提交
8 9
}

10 11 12
func @predict(%a: !infrt.dense_tensor<CPU, FP32, NCHW>, %b: !infrt.dense_tensor<CPU, FP32, NCHW>) -> (!infrt.dense_tensor<CPU, FP32, NCHW>, !infrt.dense_tensor<CPU, FP32, NCHW>) {
  %a0 = dt.shallow_copy_tensor %a : !infrt.dense_tensor<CPU, FP32, NCHW> -> !infrt.dense_tensor<CPU, FP32, NCHW>
  %b0 = dt.shallow_copy_tensor %b : !infrt.dense_tensor<CPU, FP32, NCHW> -> !infrt.dense_tensor<CPU, FP32, NCHW>
Y
Yan Chunwei 已提交
13

14
  Infrt.return %a0, %b0: !infrt.dense_tensor<CPU, FP32, NCHW>, !infrt.dense_tensor<CPU, FP32, NCHW>
Y
Yan Chunwei 已提交
15 16 17 18 19
}


func @main() {
  %shape = ts.build_shape [1:i64, 57:i64]
20
  %a = dt.create_uninit_tensor.f32 [12:i64, 23:i64] -> !infrt.dense_tensor<CPU, FP32, NCHW>
Y
Yan Chunwei 已提交
21

22 23
  %b, %c = Infrt.call @predict(%a, %a) : (!infrt.dense_tensor<CPU, FP32, NCHW>, !infrt.dense_tensor<CPU, FP32, NCHW>) -> (!infrt.dense_tensor<CPU, FP32, NCHW>, !infrt.dense_tensor<CPU, FP32, NCHW>)
  Infrt.return
Y
Yan Chunwei 已提交
24
}