def register_unary_op(op): """ Register a lowering implementation for the unary op `op`. Because the lowering implementations compile the final op separately using a lambda and compile_internal, `op` needs to be tied to each lowering implementation using a closure. This function makes and lowers a closure for one op. """ to_lower_op = make_unary_op(op) cuda_lower(op, MaskedType)(to_lower_op)
def register_const_op(op): to_lower_op = make_const_op(op) cuda_lower(op, MaskedType, types.Number)(to_lower_op) cuda_lower(op, types.Number, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.Boolean)(to_lower_op) cuda_lower(op, types.Boolean, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.NPDatetime)(to_lower_op) cuda_lower(op, types.NPDatetime, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.NPTimedelta)(to_lower_op) cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op)
cuda_lower(op, MaskedType, types.Number)(to_lower_op) cuda_lower(op, types.Number, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.Boolean)(to_lower_op) cuda_lower(op, types.Boolean, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.NPDatetime)(to_lower_op) cuda_lower(op, types.NPDatetime, MaskedType)(to_lower_op) cuda_lower(op, MaskedType, types.NPTimedelta)(to_lower_op) cuda_lower(op, types.NPTimedelta, MaskedType)(to_lower_op) # register all lowering at init for binary_op in arith_ops + bitwise_ops + comparison_ops: register_arithmetic_op(binary_op) register_const_op(binary_op) # null op impl can be shared between all ops cuda_lower(binary_op, MaskedType, NAType)(masked_scalar_null_op_impl) cuda_lower(binary_op, NAType, MaskedType)(masked_scalar_null_op_impl) # register all lowering at init for unary_op in unary_ops: register_unary_op(unary_op) @cuda_lower(operator.is_, MaskedType, NAType) @cuda_lower(operator.is_, NAType, MaskedType) def masked_scalar_is_null_impl(context, builder, sig, args): """ Implement `MaskedType` is `NA` """ if isinstance(sig.args[1], NAType): masked_type, na = sig.args
def register_const_op(op): to_lower_op = make_const_op(op) cuda_lower(op, MaskedType, types.Number)(to_lower_op) cuda_lower(op, types.Number, MaskedType)(to_lower_op)