| 67 | } |
| 68 | |
| 69 | CondTakeImpl::Output CondTakeImpl::exec( |
| 70 | _megdnn_tensor_in data, _megdnn_tensor_in mask, _megdnn_workspace workspace, |
| 71 | DynOutMallocPolicyCall malloc_policy) { |
| 72 | check_exec_get_size(data.layout, mask.layout, workspace.size); |
| 73 | auto handle = concrete_handle(this->handle()); |
| 74 | |
| 75 | auto wk_bundle = make_bundle(data.layout, mask.layout); |
| 76 | wk_bundle.set(workspace.raw_ptr); |
| 77 | |
| 78 | // float_mask = mask.astype(float32) |
| 79 | void* float_mask = wk_bundle.get(0); |
| 80 | TensorLayout float_mask_layout(TensorShape(mask.layout), dtype::Float32()); |
| 81 | megdnn_assert(wk_bundle.get_size(0) >= float_mask_layout.access_bytes()); |
| 82 | auto astype_opr = handle->create_operator<TypeCvt>(); |
| 83 | astype_opr->exec(mask, {float_mask, float_mask_layout}); |
| 84 | |
| 85 | void* one_elem = wk_bundle.get(1); |
| 86 | TensorLayout one_elem_layout(TensorShape({1}), dtype::Float32()); |
| 87 | |
| 88 | auto fill_opr = handle->create_operator<Fill>(); |
| 89 | fill_opr->param().value = param().val; |
| 90 | fill_opr->exec({one_elem, one_elem_layout}, {}); |
| 91 | |
| 92 | CnnlTensorDescriptor mask_desc, val_desc; |
| 93 | ShapeInfo one_elem_shape = {1}; |
| 94 | val_desc.set(1, one_elem_shape, CNNL_DTYPE_FLOAT, CNNL_LAYOUT_ARRAY); |
| 95 | mask_desc.set( |
| 96 | float_mask_layout.ndim, float_mask_layout.shape, CNNL_DTYPE_FLOAT, |
| 97 | CNNL_LAYOUT_ARRAY); |
| 98 | float alpha = 1.f, beta = 1.f; |
| 99 | if (param().mode == Param::Mode::EQ || param().mode == Param::Mode::NEQ) { |
| 100 | cnnl_check(cnnlAssignSub( |
| 101 | handle->cnnl_handle(), &alpha, val_desc.desc(), one_elem, |
| 102 | wk_bundle.get(2), wk_bundle.get_size(2), &beta, mask_desc.desc(), |
| 103 | float_mask)); |
| 104 | cnnl_check( |
| 105 | cnnlAbs(handle->cnnl_handle(), mask_desc.desc(), float_mask, |
| 106 | mask_desc.desc(), float_mask)); |
| 107 | fill_opr->param().value = param().eps; |
| 108 | fill_opr->exec({one_elem, one_elem_layout}, {}); |
| 109 | } |
| 110 | cnnlLogicOp_t op = get_logic_op(param().mode); |
| 111 | if (param().mode == Param::Mode::EQ) { |
| 112 | op = get_logic_op(Param::Mode::LT); |
| 113 | } |
| 114 | if (param().mode == Param::Mode::NEQ) { |
| 115 | op = get_logic_op(Param::Mode::GEQ); |
| 116 | } |
| 117 | |
| 118 | cnnl_check(cnnlLogicOp( |
| 119 | handle->cnnl_handle(), op, mask_desc.desc(), float_mask, val_desc.desc(), |
| 120 | one_elem, wk_bundle.get(3), wk_bundle.get_size(3), mask_desc.desc(), |
| 121 | float_mask)); |
| 122 | |
| 123 | void* num_true_ptr = wk_bundle.get(1); |
| 124 | CnnlTensorDescriptor num_true_desc; |
| 125 | num_true_desc.set(1, one_elem_shape, CNNL_DTYPE_INT32, CNNL_LAYOUT_ARRAY); |
| 126 | cnnl_check(cnnlNumTrue_v3( |
nothing calls this directly
no test coverage detected