Fix test_torch index dfx checks-op-plugin-AtomGit

Mop_plugin/ops/aclops/IndexPutKernelNpu.cpp+25-0

	@@ -21,6 +21,7 @@
21	#include "op_plugin/utils/OpAdapter.h"	21	#include "op_plugin/utils/OpAdapter.h"
22	#include "op_plugin/utils/AdvancedIndex.h"	22	#include "op_plugin/utils/AdvancedIndex.h"
23	#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"	23	#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"
		24	+#include "torch_npu/csrc/aten/mirror/NPUMemoryOverlap.h"
24		25
25	namespace acl_op {	26	namespace acl_op {
26	using npu_preparation = at_npu::native::OpPreparation;	27	using npu_preparation = at_npu::native::OpPreparation;
	@@ -34,6 +35,17 @@ const std::string indexed_sizes_str = "indexed_sizes";
34	const std::string indexed_strides_str = "indexed_strides";	35	const std::string indexed_strides_str = "indexed_strides";
35	const std::string aicore_str = "AiCore";	36	const std::string aicore_str = "AiCore";
36		37
		38	+void check_no_overlap(const at::Tensor& a, const at::Tensor& b)
		39	+{
		40	+ const auto overlap_status = at_npu::native::get_overlap_status(a, b);
		41	+ TORCH_CHECK(overlap_status != at_npu::native::MemOverlapStatus::PARTIAL &&
		42	+ overlap_status != at_npu::native::MemOverlapStatus::FULL,
		43	+ "unsupported operation: some elements of the input tensor and "
		44	+ "the written-to tensor refer to a single memory location. "
		45	+ "Please clone() the tensor before performing the operation.",
		46	+ OPS_ERROR(ErrCode::NOT_SUPPORT));
		47	+}
		48	+
37	bool is_aicpu_valid(const at::Tensor &self, const std::vector<at::Tensor> &all_defined_indices,	49	bool is_aicpu_valid(const at::Tensor &self, const std::vector<at::Tensor> &all_defined_indices,
38	const at::SmallVector<int64_t, N> masks)	50	const at::SmallVector<int64_t, N> masks)
39	{	51	{
	@@ -225,6 +237,19 @@ at::Tensor &_index_put_impl_(at::Tensor &self, const c10::List<c10::optional<at:
225	if (self.device().type() == at::kCPU) {	237	if (self.device().type() == at::kCPU) {
226	return at::native::_index_put_impl_(self, indices, value, accumulate, unsafe);	238	return at::native::_index_put_impl_(self, indices, value, accumulate, unsafe);
227	}	239	}
		240	+ if (at_npu::native::has_internal_overlap(self) == at_npu::native::MemOverlap::YES) {
		241	+ TORCH_WARN(
		242	+ "Use of index_put_ on expanded tensors is deprecated. "
		243	+ "Please clone() the tensor before performing this operation. "
		244	+ "This also applies to advanced indexing e.g. tensor[indices] = tensor");
		245	+ }
		246	+ check_no_overlap(self, value);
		247	+ for (const c10::optional<at::Tensor>& index : indices) {
		248	+ if (index.has_value()) {
		249	+ check_no_overlap(self, *index);
		250	+ }
		251	+ }
		252	+
228	bool needCast = op_plugin::AdvanceIndex::checkIndexTensorTypes(indices);	253	bool needCast = op_plugin::AdvanceIndex::checkIndexTensorTypes(indices);
229	at::SmallVector<int64_t, N> masks;	254	at::SmallVector<int64_t, N> masks;
230	std::vector<at::Tensor> all_defined_indices;	255	std::vector<at::Tensor> all_defined_indices;

Mop_plugin/ops/opapi/IndexAddKernelNpuOpApi.cpp+26-0

	@@ -27,6 +27,19 @@ at::Tensor& index_add_out(
27	const at::Scalar& alpha,	27	const at::Scalar& alpha,
28	at::Tensor& result) {	28	at::Tensor& result) {
29	DO_COMPATIBILITY(aclnnIndexAdd, acl_op::index_add_out(self, dim, index, source, alpha, result));	29	DO_COMPATIBILITY(aclnnIndexAdd, acl_op::index_add_out(self, dim, index, source, alpha, result));
		30	+ auto self_sizes = self.sizes().vec();
		31	+ auto source_sizes = source.sizes().vec();
		32	+ if (source.dim() != 0 && self.dim() != 0) {
		33	+ auto wrapped_dim = at::maybe_wrap_dim(dim, self.dim());
		34	+ self_sizes.erase(self_sizes.begin() + wrapped_dim);
		35	+ source_sizes.erase(source_sizes.begin() + wrapped_dim);
		36	+ }
		37	+ TORCH_CHECK(self_sizes == source_sizes,
		38	+ "source tensor shape must match self tensor shape, excluding the specified dimension. Got self.shape = ",
		39	+ self.sizes(),
		40	+ " source.shape = ",
		41	+ source.sizes(),
		42	+ OPS_ERROR(ErrCode::PARAM));
30	at_npu::native::OpPreparation::check_tensor({self, index, source},	43	at_npu::native::OpPreparation::check_tensor({self, index, source},
31	result,	44	result,
32	result.scalar_type(),	45	result.scalar_type(),
	@@ -45,6 +58,19 @@ at::Tensor index_add(
45	const at::Tensor& source,	58	const at::Tensor& source,
46	const at::Scalar& alpha) {	59	const at::Scalar& alpha) {
47	DO_COMPATIBILITY(aclnnIndexAdd, acl_op::index_add(self, dim, index, source, alpha));	60	DO_COMPATIBILITY(aclnnIndexAdd, acl_op::index_add(self, dim, index, source, alpha));
		61	+ auto self_sizes = self.sizes().vec();
		62	+ auto source_sizes = source.sizes().vec();
		63	+ if (source.dim() != 0 && self.dim() != 0) {
		64	+ auto wrapped_dim = at::maybe_wrap_dim(dim, self.dim());
		65	+ self_sizes.erase(self_sizes.begin() + wrapped_dim);
		66	+ source_sizes.erase(source_sizes.begin() + wrapped_dim);
		67	+ }
		68	+ TORCH_CHECK(self_sizes == source_sizes,
		69	+ "source tensor shape must match self tensor shape, excluding the specified dimension. Got self.shape = ",
		70	+ self.sizes(),
		71	+ " source.shape = ",
		72	+ source.sizes(),
		73	+ OPS_ERROR(ErrCode::PARAM));
48	at::Tensor result = at_npu::native::OpPreparation::apply_tensor_without_format(self.sizes(), self.options());	74	at::Tensor result = at_npu::native::OpPreparation::apply_tensor_without_format(self.sizes(), self.options());
49	EXEC_NPU_CMD(aclnnIndexAdd, result.copy_(self), dim, index, source, alpha, result.copy_(self));	75	EXEC_NPU_CMD(aclnnIndexAdd, result.copy_(self), dim, index, source, alpha, result.copy_(self));
50	return result;	76	return result;

Mop_plugin/ops/opapi/IndexPutKernelNpuOpApi.cpp+28-1

	@@ -18,10 +18,24 @@
18	#include "op_plugin/OpApiInterface.h"	18	#include "op_plugin/OpApiInterface.h"
19	#include "op_plugin/utils/op_api_common.h"	19	#include "op_plugin/utils/op_api_common.h"
20	#include "op_plugin/utils/AdvancedIndex.h"	20	#include "op_plugin/utils/AdvancedIndex.h"
		21	+#include "torch_npu/csrc/aten/mirror/NPUMemoryOverlap.h"
21		22
22	namespace op_api {	23	namespace op_api {
23	using npu_preparation = at_npu::native::OpPreparation;	24	using npu_preparation = at_npu::native::OpPreparation;
24		25
		26	+namespace {
		27	+void check_no_overlap(const at::Tensor& a, const at::Tensor& b)
		28	+{
		29	+ const auto overlap_status = at_npu::native::get_overlap_status(a, b);
		30	+ TORCH_CHECK(overlap_status != at_npu::native::MemOverlapStatus::PARTIAL &&
		31	+ overlap_status != at_npu::native::MemOverlapStatus::FULL,
		32	+ "unsupported operation: some elements of the input tensor and "
		33	+ "the written-to tensor refer to a single memory location. "
		34	+ "Please clone() the tensor before performing the operation.",
		35	+ OPS_ERROR(ErrCode::NOT_SUPPORT));
		36	+}
		37	+} // namespace
		38	+
25	at::Tensor index_put(	39	at::Tensor index_put(
26	const at::Tensor& self,	40	const at::Tensor& self,
27	const c10::List<c10::optional<at::Tensor>>& indices,	41	const c10::List<c10::optional<at::Tensor>>& indices,
	@@ -50,6 +64,19 @@ at::Tensor& _index_put_impl_(
50	if (self.device().type() == at::kCPU) {	64	if (self.device().type() == at::kCPU) {
51	return at::native::_index_put_impl_(self, indices, value, accumulate, unsafe);	65	return at::native::_index_put_impl_(self, indices, value, accumulate, unsafe);
52	}	66	}
		67	+ if (at_npu::native::has_internal_overlap(self) == at_npu::native::MemOverlap::YES) {
		68	+ TORCH_WARN(
		69	+ "Use of index_put_ on expanded tensors is deprecated. "
		70	+ "Please clone() the tensor before performing this operation. "
		71	+ "This also applies to advanced indexing e.g. tensor[indices] = tensor");
		72	+ }
		73	+ check_no_overlap(self, value);
		74	+ for (const c10::optional<at::Tensor>& index : indices) {
		75	+ if (index.has_value()) {
		76	+ check_no_overlap(self, *index);
		77	+ }
		78	+ }
		79	+
53	bool needCast = op_plugin::AdvanceIndex::checkIndexTensorTypes(indices);	80	bool needCast = op_plugin::AdvanceIndex::checkIndexTensorTypes(indices);
54	auto indices_after = op_plugin::AdvanceIndex::npu_expand_tensors(self, indices, needCast, true);	81	auto indices_after = op_plugin::AdvanceIndex::npu_expand_tensors(self, indices, needCast, true);
55	std::vector<at::Tensor> all_defined_indices;	82	std::vector<at::Tensor> all_defined_indices;
	@@ -84,4 +111,4 @@ at::Tensor& _index_put_impl_(
84	return self;	111	return self;
85	}	112	}
86		113
87	-}	114	+}