| 132 | } |
| 133 | |
| 134 | SmallVector<TensorPtr> apply_on_physical_tensor( |
| 135 | const OpDef& def, const SmallVector<TensorPtr>& inputs, |
| 136 | SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) { |
| 137 | auto&& matmul = def.cast_final_safe<MatrixMul>(); |
| 138 | auto&& cn = inputs[0]->comp_node(); |
| 139 | |
| 140 | using TensorND = megdnn::TensorND; |
| 141 | SmallVector<TensorND> inp_tensornds(inputs.size()); |
| 142 | TensorLayout layout1 = inputs[0]->layout(), layout2 = inputs[1]->layout(); |
| 143 | |
| 144 | DnnOprCaller<megdnn::MatrixMul> dnn_opr(cn, matmul.param(), matmul.policy()); |
| 145 | |
| 146 | if (matmul.dimA == matmul.dimB && matmul.dimB >= 3) { // only happens in backward |
| 147 | for (size_t i = 1; i + 1 < layout1.ndim; ++i) { |
| 148 | layout1[0] *= layout1[i]; |
| 149 | layout2[0] *= layout2[i]; |
| 150 | } |
| 151 | layout1[1] = layout1[layout1.ndim - 1]; |
| 152 | layout1.ndim = 2; |
| 153 | layout1.init_contiguous_stride(); |
| 154 | layout2[1] = layout2[layout2.ndim - 1]; |
| 155 | layout2.ndim = 2; |
| 156 | layout2.init_contiguous_stride(); |
| 157 | } |
| 158 | |
| 159 | DType dst_dtype; |
| 160 | dnn_opr.op()->deduce_dtype(layout1.dtype, layout2.dtype, dst_dtype); |
| 161 | |
| 162 | // only matters when layout1 has dim 2 |
| 163 | if (matmul.transposeA) |
| 164 | std::swap(layout1.shape[0], layout1.shape[1]); |
| 165 | // only matters when layout2 has dim 2 |
| 166 | if (matmul.transposeB) |
| 167 | std::swap(layout2.shape[0], layout2.shape[1]); |
| 168 | |
| 169 | size_t dim1 = layout1.ndim, dim2 = layout2.ndim; |
| 170 | TensorLayout real_dst_layout(dst_dtype); |
| 171 | if (validated) { |
| 172 | real_dst_layout = output_descs[0].layout; |
| 173 | } else { |
| 174 | size_t ri = 0; |
| 175 | for (size_t i = 0; i < dim1 - 2; i++) |
| 176 | real_dst_layout[ri++] = layout1[i]; |
| 177 | real_dst_layout[ri++] = layout1[dim1 - 2]; |
| 178 | if (dim2 == 2) |
| 179 | real_dst_layout[ri++] = layout2[dim2 - 1]; |
| 180 | real_dst_layout.ndim = ri; |
| 181 | real_dst_layout.init_contiguous_stride(); |
| 182 | } |
| 183 | |
| 184 | if (dim1 == 0 || dim2 == 0 || layout1[layout1.ndim - 1] == 0) { |
| 185 | auto out = Tensor::make(real_dst_layout, cn); |
| 186 | |
| 187 | if (!out->empty()) { |
| 188 | dev_tensor_memset(out->dev_tensor(), 0); |
| 189 | } |
| 190 | return {out}; |
| 191 | } |
nothing calls this directly
no test coverage detected