unicorn who dev
Complete FireViewer Android learning SDK, source and verification records
dee7f43 verified | """Native LiteRT lowerings for shape-dependent reductions and shape arithmetic.""" | |
| import operator | |
| import torch | |
| from litert_torch import fx_infra | |
| from litert_torch.backend.experimental import torch_tfl | |
| from litert_torch.backend.experimental.torch_tfl._lowerings import _ir_operation | |
| from litert_torch.backend.lowerings import registry,utils | |
| from litert_converter.mlir import ir | |
| fx_infra.decomp.update_pre_lower_decomp(torch_tfl.decomps) | |
| def dynamic_layer_norm(x,normalized_shape,weight,bias,eps): | |
| axes=list(range(x.ndim-len(normalized_shape),x.ndim)) | |
| mean=x.mean(axes,keepdim=True) | |
| variance=((x-mean)*(x-mean)).mean(axes,keepdim=True) | |
| invstd=torch.rsqrt(variance+eps) | |
| output=(x-mean)*invstd | |
| if weight is not None:output=output*weight | |
| if bias is not None:output=output+bias | |
| return output,mean,invstd | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.native_layer_norm.default:dynamic_layer_norm}) | |
| def lower_floor_div(lctx,lhs,rhs): | |
| if isinstance(lhs,int) and isinstance(rhs,int):return lhs//rhs | |
| if isinstance(lhs,int):lhs=utils.splat(lhs,rhs.type.element_type,rhs.type.shape) | |
| if isinstance(rhs,int):rhs=utils.splat(rhs,lhs.type.element_type,lhs.type.shape) | |
| if lhs.type!=rhs.type:raise TypeError('Shape floor division requires matching scalar integer types') | |
| return _ir_operation('tfl.floor_div',results=[lhs.type],operands=[lhs,rhs]) | |
| def lower_relu(lctx,x): | |
| return _ir_operation('tfl.relu',results=utils.node_meta_to_ir_types(lctx.node),operands=[x]) | |
| def lower_cumsum(lctx,x,dim,dtype=None): | |
| result=utils.node_meta_to_ir_types(lctx.node) | |
| if x.type.element_type!=result[0].element_type: | |
| typ=ir.RankedTensorType.get(x.type.shape,result[0].element_type) | |
| x=_ir_operation('tfl.cast',results=[typ],operands=[x]) | |
| axis=utils.splat(dim,ir.IntegerType.get_signless(32),[]) | |
| return _ir_operation('tfl.cumsum',results=result,operands=[x,axis],attributes={'exclusive':ir.BoolAttr.get(False),'reverse':ir.BoolAttr.get(False)}) | |
| def lower_all_true(lctx,x): | |
| if len(x.type.shape)==0:return x | |
| import numpy as np | |
| axes=utils.numpy_array_constant(np.arange(len(x.type.shape),dtype=np.int32)) | |
| return _ir_operation('tfl.reduce_all',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,axes],attributes={'keep_dims':ir.BoolAttr.get(False)}) | |
| # Preserve symbolic dimensions in the operator schema. The upstream int[] | |
| # schema can specialize reshape dimensions during decomposition. | |
| def dynamic_reshape(x,shape): | |
| return x.reshape(shape).clone() | |
| def dynamic_reshape_fake(x,shape): | |
| return x.reshape(shape).clone() | |
| def lower_reshape(lctx,x,shape): | |
| return _ir_operation('tfl.reshape',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,utils.convert_shape_to_ir_value(shape)]) | |
| def dynamic_view(x,shape):return torch.ops.vds.reshape(x,shape) | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.view.default:dynamic_view,torch.ops.aten._unsafe_view.default:dynamic_view}) | |
| def dynamic_arange(start,end,step=1,*,dtype=None,layout=None,device=None,pin_memory=None): | |
| result=torch.ops.tfl.range(start,end,step) | |
| return result.to(dtype=dtype) if dtype is not None else result | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.arange.start_step:dynamic_arange}) | |
| def dynamic_broadcast(x,shape):return x.expand(shape).clone() | |
| def dynamic_broadcast_fake(x,shape):return x.expand(shape).clone() | |
| def lower_broadcast(lctx,x,shape): | |
| return _ir_operation('tfl.broadcast_to',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,utils.convert_shape_to_ir_value(shape)]) | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.expand.default:lambda x,shape,implicit=False:torch.ops.vds.broadcast_to(x,shape)}) | |
| def dynamic_linear(x:torch.Tensor,weight:torch.Tensor,bias:torch.Tensor|None)->torch.Tensor: | |
| return torch.nn.functional.linear(x,weight,bias) | |
| def dynamic_linear_fake(x,weight,bias): | |
| return x.new_empty(x.shape[:-1] if weight.ndim==1 else (*x.shape[:-1],weight.shape[0])) | |
| def lower_linear(lctx,x,weight,bias): | |
| if len(weight.type.shape)==1: | |
| from litert_converter.mlir.dialects import stablehlo | |
| from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_squeeze_lowering | |
| weight=stablehlo.reshape(ir.RankedTensorType.get([1,weight.type.shape[0]],weight.type.element_type),weight) | |
| if bias is None:bias=utils.splat(0.0,weight.type.element_type,[1]) | |
| typ=utils.node_meta_to_ir_types(lctx.node)[0] | |
| expanded=ir.RankedTensorType.get([*typ.shape,1],typ.element_type) | |
| y=_ir_operation('tfl.fully_connected',results=[expanded],operands=[x,weight,bias],attributes={'fused_activation_function':ir.StringAttr.get('NONE'),'weights_format':ir.StringAttr.get('DEFAULT'),'keep_num_dims':ir.BoolAttr.get(True),'asymmetric_quantize_inputs':ir.BoolAttr.get(False)}) | |
| return _tfl_squeeze_lowering(lctx,y,[-1]) | |
| if bias is None:bias=utils.splat(0.0,weight.type.element_type,[weight.type.shape[0]]) | |
| return _ir_operation('tfl.fully_connected',results=utils.node_meta_to_ir_types(lctx.node),operands=[x,weight,bias],attributes={'fused_activation_function':ir.StringAttr.get('NONE'),'weights_format':ir.StringAttr.get('DEFAULT'),'keep_num_dims':ir.BoolAttr.get(True),'asymmetric_quantize_inputs':ir.BoolAttr.get(False)}) | |
| fx_infra.decomp.update_pre_lower_decomp({torch.ops.aten.linear.default:lambda x,weight,bias=None:torch.ops.vds.linear(x,weight,bias)}) | |
| fx_infra.decomp.update_pre_convert_decomp({ | |
| torch.ops.aten.linear.default:lambda x,weight,bias=None:torch.ops.vds.linear(x,weight,bias), | |
| torch.ops.aten.view.default:dynamic_view, | |
| torch.ops.aten._unsafe_view.default:dynamic_view, | |
| torch.ops.aten.native_layer_norm.default:dynamic_layer_norm, | |
| torch.ops.aten.arange.start_step:dynamic_arange, | |
| torch.ops.aten.expand.default:lambda x,shape,implicit=False:torch.ops.vds.broadcast_to(x,shape), | |
| }) | |
| def dynamic_resize2d(x,size,align_corners,nearest): | |
| if nearest:return torch.nn.functional.interpolate(x,size=size,mode='nearest') | |
| return torch.nn.functional.interpolate(x,size=size,mode='bilinear',align_corners=align_corners) | |
| def dynamic_resize2d_fake(x,size,align_corners,nearest):return x.new_empty((x.shape[0],x.shape[1],*size)) | |
| def lower_resize2d(lctx,x,size,align_corners,nearest): | |
| from litert_converter.mlir.dialects import stablehlo | |
| nhwc=stablehlo.transpose(x,[0,2,3,1]) | |
| result=utils.node_meta_to_ir_types(lctx.node)[0] | |
| typ=ir.RankedTensorType.get([result.shape[i] for i in [0,2,3,1]],result.element_type) | |
| resized=_ir_operation('tfl.resize_nearest_neighbor' if nearest else 'tfl.resize_bilinear',results=[typ],operands=[nhwc,utils.convert_shape_to_ir_value(size)],attributes={'align_corners':ir.BoolAttr.get(align_corners),'half_pixel_centers':ir.BoolAttr.get(not align_corners and not nearest)}) | |
| return stablehlo.transpose(resized,[0,3,1,2]) | |
| def bilinear_vec(x,output_size,align_corners,scale_factors=None): | |
| size=output_size if output_size is not None else [int(x.shape[-2+i]*scale_factors[i]) for i in range(2)] | |
| return torch.ops.vds.resize2d(x,size,align_corners,False) | |
| def bilinear_default(x,output_size,align_corners,scales_h=None,scales_w=None): | |
| return torch.ops.vds.resize2d(x,output_size,align_corners,False) | |
| fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.upsample_bilinear2d.vec:bilinear_vec,torch.ops.aten.upsample_bilinear2d.default:bilinear_default}) | |
| from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_mul_lowering | |
| def lower_bool_mul(lctx,lhs,rhs,fused_activation_function='NONE'): | |
| if isinstance(lhs,ir.Value) and isinstance(rhs,ir.Value) and str(lhs.type.element_type)=='i1' and str(rhs.type.element_type)=='i1': | |
| return _ir_operation('tfl.logical_and',results=utils.node_meta_to_ir_types(lctx.node),operands=[lhs,rhs]) | |
| return _tfl_mul_lowering(lctx,lhs,rhs,fused_activation_function) | |
| def dynamic_strided_slice(x,begin,end,strides): | |
| return x[tuple(slice(b,e,s) for b,e,s in zip(begin,end,strides))].clone() | |
| def dynamic_strided_slice_fake(x,begin,end,strides): | |
| return x[tuple(slice(b,e,s) for b,e,s in zip(begin,end,strides))].clone() | |
| def lower_strided_slice(lctx,x,begin,end,strides): | |
| from litert_torch.backend.experimental.torch_tfl._lowerings import _tfl_strided_slice_lowering | |
| return _tfl_strided_slice_lowering(lctx,x,begin,end,strides) | |
| def dynamic_slice(x,dim=0,start=None,end=None,step=1): | |
| n=x.shape[dim];b=0 if start is None else start;e=n if end is None else end | |
| if b<0:b+=n | |
| if e<0:e+=n | |
| b=max(0,min(b,n));e=max(b,min(e,n)) | |
| begin=[0]*x.ndim;ends=list(x.shape);strides=[1]*x.ndim | |
| begin[dim]=b;ends[dim]=e;strides[dim]=step | |
| return torch.ops.vds.strided_slice(x,begin,ends,strides) | |
| def dynamic_embedding(weight,indices,padding_idx=-1,scale_grad_by_freq=False,sparse=False): | |
| flat=torch.ops.vds.reshape(indices,[-1]).to(torch.int32) | |
| result=torch.ops.tfl.embedding_lookup(flat,weight) | |
| return torch.ops.vds.reshape(result,[*indices.shape,weight.shape[-1]]) | |
| fx_infra.decomp.update_pre_convert_decomp({torch.ops.aten.slice.Tensor:dynamic_slice,torch.ops.aten.embedding.default:dynamic_embedding}) | |