Skip to content

Instantly share code, notes, and snippets.

@stellaraccident
Last active April 24, 2021 00:45
Show Gist options
  • Save stellaraccident/e58ae480ba713fb70a20e38a487373f4 to your computer and use it in GitHub Desktop.
Save stellaraccident/e58ae480ba713fb70a20e38a487373f4 to your computer and use it in GitHub Desktop.
TF import with dyn shapes
module {
func @binary_func(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>) -> (tensor<16xf32>, tensor<16xf32>) attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]],\22r\22:[[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]],\22v\22:1}", iree.module.export} {
return %arg0, %arg1 : tensor<16xf32>, tensor<16xf32>
}
func @dict_nest(%arg0: !iree.list<?>, %arg1: tensor<f32>) -> !iree.list<?> attributes {iree.abi = "{\22a\22:[[\22sdict\22,[\22dict\22,[\22sdict\22,[\22a\22,[\22ndarray\22,\22f32\22,1,16]],[\22b\22,[\22ndarray\22,\22f32\22,1,16]]]],[\22list\22,[\22slist\22,[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]]]],[\22ndarray\22,\22f32\22,0]],\22r\22:[[\22sdict\22,[\22dict\22,[\22sdict\22,[\22a\22,[\22ndarray\22,\22f32\22,1,16]],[\22b\22,[\22ndarray\22,\22f32\22,1,16]]]],[\22list\22,[\22slist\22,[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]]]]],\22v\22:1}", iree.module.export} {
%c2 = constant 2 : index
%c0 = constant 0 : index
%c1 = constant 1 : index
%0 = iree.list.get %arg0[%c0] : !iree.list<?> -> !iree.list<?>
%1 = iree.list.get %0[%c0] : !iree.list<?> -> tensor<16xf32>
%2 = iree.list.get %0[%c1] : !iree.list<?> -> tensor<16xf32>
%3 = iree.list.get %arg0[%c1] : !iree.list<?> -> !iree.list<?>
%4 = iree.list.get %3[%c0] : !iree.list<?> -> tensor<16xf32>
%5 = iree.list.get %3[%c1] : !iree.list<?> -> tensor<16xf32>
%6 = iree.list.create %c2 : !iree.list<?>
%7 = iree.list.create %c2 : !iree.list<?>
iree.list.set %7[%c0], %1 : tensor<16xf32> -> !iree.list<?>
iree.list.set %7[%c1], %2 : tensor<16xf32> -> !iree.list<?>
iree.list.set %6[%c0], %7 : !iree.list<?> -> !iree.list<?>
%8 = iree.list.create %c2 : !iree.list<?>
iree.list.set %8[%c0], %4 : tensor<16xf32> -> !iree.list<?>
iree.list.set %8[%c1], %5 : tensor<16xf32> -> !iree.list<?>
iree.list.set %6[%c1], %8 : !iree.list<?> -> !iree.list<?>
return %6 : !iree.list<?>
}
func @dynamic_addition(%arg0: tensor<?x16xf32>, %arg1: tensor<?x16xf32>) -> tensor<?x16xf32> attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,2,null,16],[\22ndarray\22,\22f32\22,2,null,16]],\22r\22:[[\22ndarray\22,\22f32\22,2,null,1
6]],\22v\22:1}", iree.module.export} {
%0 = shapex.get_ranked_shape %arg0 : tensor<?x16xf32> -> !shapex.ranked_shape<[?,16]>
%1 = shapex.get_ranked_shape %arg1 : tensor<?x16xf32> -> !shapex.ranked_shape<[?,16]>
%2 = "shapex.ranked_broadcast_shape"(%0, %1) {lhs_broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>, rhs_broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (!shapex.ranked_shape<[?,16]>, !shapex.ranked_shape<[?,16]>) -> !shapex.ranked_shape<[?,16]>
%3 = "shapex.ranked_broadcast_in_dim"(%arg0, %2) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x16xf32>, !shapex.ranked_shape<[?,16]>) -> tensor<?x16xf32>
%4 = "shapex.ranked_broadcast_in_dim"(%arg1, %2) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x16xf32>, !shapex.ranked_shape<[?,16]>) -> tensor<?x16xf32>
%5 = mhlo.add %3, %4 : tensor<?x16xf32>
return %5 : tensor<?x16xf32>
}
func @return_list(%arg0: tensor<16xf32>, %arg1: tensor<16xf32>) -> (tensor<16xf32>, tensor<16xf32>) attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]],\22r\22:[[\22ndarray\22,\22f32\22,1,16],[\22ndarray\22,\22f32\22,1,16]],\22v\22:1}", iree.module.export} {
return %arg0, %arg1 : tensor<16xf32>, tensor<16xf32>
}
func @unary_func(%arg0: tensor<16xf32>) -> tensor<16xf32> attributes {iree.abi = "{\22a\22:[[\22ndarray\22,\22f32\22,1,16]],\22r\22:[[\22ndarray\22,\22f32\22,1,16]],\22v\22:1}", iree.module.export} {
return %arg0 : tensor<16xf32>
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment