已合并
Support broadcast scenarios in the cdist backward operator. #5279
HandsoemLemon创建于 6月26日
Support broadcast scenarios in the cdist backward operator. #5279
已合并
共 1 个文件变更+45-5
| @@ -44,15 +44,55 @@ at::Tensor _cdist_backward( | |||
| 44 | if (p_in_range && (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950)) { | 44 | if (p_in_range && (c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend950)) { |
| 45 | return acl_op::_cdist_backward(grad, x1, x2, p, cdist); | 45 | return acl_op::_cdist_backward(grad, x1, x2, p, cdist); |
| 46 | } | 46 | } |
| 47 | - auto output_size = x1.sizes(); | ||
| 48 | - auto output_dtype = grad.scalar_type(); | ||
| 49 | 47 | ||
| 48 | + int64_t c1 = x1.size(-1); | ||
| 49 | + int64_t c2 = x2.size(-1); | ||
| 50 | + int64_t r1 = x1.size(-2); | ||
| 51 | + int64_t r2 = x2.size(-2); | ||
| 52 | + int64_t dim1 = static_cast<int64_t>(x1.dim()); | ||
| 53 | + int64_t dim2 = static_cast<int64_t>(x2.dim()); | ||
| 54 | + TORCH_CHECK(c1 == c2, "X1 and X2 must have the same number of columns. X1: ", c1, " X2: ", c2, | ||
| 55 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 56 | + | ||
| 57 | + at::IntArrayRef batch_tensor1(x1.sizes().data(), dim1 - 2); | ||
| 58 | + at::IntArrayRef batch_tensor2(x2.sizes().data(), dim2 - 2); | ||
| 59 | + std::vector<int64_t> expand_batch_portion = at::infer_size(batch_tensor1, batch_tensor2); | ||
| 60 | + std::vector<int64_t> tensor1_expand_size(expand_batch_portion); | ||
| 61 | + tensor1_expand_size.insert(tensor1_expand_size.end(), {r1, c1}); | ||
| 62 | + std::vector<int64_t> tensor2_expand_size(expand_batch_portion); | ||
| 63 | + tensor2_expand_size.insert(tensor2_expand_size.end(), {r2, c2}); | ||
| 64 | + | ||
| 65 | + bool empty_batch = false; | ||
| 66 | + for (const auto size : expand_batch_portion) { | ||
| 67 | + empty_batch = empty_batch || (size == 0); | ||
| 68 | + } | ||
| 69 | + if (r1 == 0 || r2 == 0 || c1 == 0 || empty_batch) { | ||
| 70 | + return at::zeros_like(x1, x1.options()); | ||
| 71 | + } | ||
| 72 | + | ||
| 73 | + at::Tensor x1_broadcast = x1; | ||
| 74 | + if (x1.sizes().vec() != tensor1_expand_size) { | ||
| 75 | + x1_broadcast = x1.expand(tensor1_expand_size); | ||
| 76 | + } | ||
| 77 | + x1_broadcast = x1_broadcast.contiguous(); | ||
| 78 | + | ||
| 79 | + at::Tensor x2_broadcast = x2; | ||
| 80 | + if (x2.sizes().vec() != tensor2_expand_size) { | ||
| 81 | + x2_broadcast = x2.expand(tensor2_expand_size); | ||
| 82 | + } | ||
| 83 | + x2_broadcast = x2_broadcast.contiguous(); | ||
| 84 | + | ||
| 85 | + auto grad_contiguous = grad.contiguous(); | ||
C | |||
| 86 | + auto cdist_contiguous = cdist.contiguous(); | ||
| 87 | + auto output_size = x1_broadcast.sizes(); | ||
| 88 | + auto output_dtype = x1_broadcast.scalar_type(); | ||
| 50 | at::Tensor out = at_npu::native::OpPreparation::apply_tensor_without_format( | 89 | at::Tensor out = at_npu::native::OpPreparation::apply_tensor_without_format( |
| 51 | output_size, | 90 | output_size, |
| 52 | - grad.options().dtype(output_dtype)); | 91 | + x1_broadcast.options().dtype(output_dtype)); |
| 53 | 92 | ||
| 54 | - EXEC_NPU_CMD(aclnnCdistBackward, grad, x1, x2, cdist, p_cast, out); | 93 | + |
| 94 | + EXEC_NPU_CMD(aclnnCdistBackward, grad_contiguous, x1_broadcast, x2_broadcast, cdist_contiguous, p_cast, out); | ||
| 55 | 95 | ||
| 56 | return out; | 96 | return out; |
| 57 | } | 97 | } |
| 58 | -} | 98 | +} |
必须要做转连续操作吗