@@ -52,6 +52,15 @@ def __call__(self, x, out=None, order="K"):
5252 if not isinstance (x , dpt .usm_ndarray ):
5353 raise TypeError (f"Expected dpctl.tensor.usm_ndarray, got { type (x )} )
5454
55+ if order not in "C" , "F" , "K" , "A" ]:
56+ order = "K"
57+ buf_dt , res_dt = _find_buf_dtype (
58+ x .dtype , self .result_type_resolver_fn_ , x .sycl_device
59+ )
60+ if res_dt is None :
61+ raise RuntimeError
62+
63+ orig_out = out
5564 if out is not None :
5665 if not isinstance (out , dpt .usm_ndarray ):
5766 raise TypeError (
@@ -64,8 +73,21 @@ def __call__(self, x, out=None, order="K"):
6473 f"Expected output shape is { x .shape } { out .shape }
6574 )
6675
67- if ti ._array_overlap (x , out ):
68- raise TypeError ("Input and output arrays have memory overlap" )
76+ if res_dt != out .dtype :
77+ raise TypeError (
78+ f"Output array of type { res_dt }
79+ f" got { out .dtype }
80+ )
81+
82+ if (
83+ buf_dt is None
84+ and ti ._array_overlap (x , out )
85+ and not ti ._same_logical_tensors (x , out )
86+ ):
87+ # Allocate a temporary buffer to avoid memory overlapping.
88+ # Note if `buf_dt` is not None, a temporary copy of `x` will be
89+ # created, so the array overlap check isn't needed.
90+ out = dpt .empty_like (out )
6991
7092 if (
7193 dpctl .utils .get_execution_queue ((x .sycl_queue , out .sycl_queue ))
@@ -75,13 +97,6 @@ def __call__(self, x, out=None, order="K"):
7597 "Input and output allocation queues are not compatible"
7698 )
7799
78- if order not in "C" , "F" , "K" , "A" ]:
79- order = "K"
80- buf_dt , res_dt = _find_buf_dtype (
81- x .dtype , self .result_type_resolver_fn_ , x .sycl_device
82- )
83- if res_dt is None :
84- raise RuntimeError
85100 exec_q = x .sycl_queue
86101 if buf_dt is None :
87102 if out is None :
@@ -91,17 +106,20 @@ def __call__(self, x, out=None, order="K"):
91106 if order == "A" :
92107 order = "F" if x .flags .f_contiguous else "C"
93108 out = dpt .empty_like (x , dtype = res_dt , order = order )
94- else :
95- if res_dt != out .dtype :
96- raise TypeError (
97- f"Output array of type { res_dt }
98- f" got { out .dtype }
99- )
100109
101- ht , _ = self .unary_fn_ (x , out , sycl_queue = exec_q )
102- ht .wait ()
110+ ht_unary_ev , unary_ev = self .unary_fn_ (x , out , sycl_queue = exec_q )
111+
112+ if not (orig_out is None or orig_out is out ):
113+ # Copy the out data from temporary buffer to original memory
114+ ht_copy_ev , _ = ti ._copy_usm_ndarray_into_usm_ndarray (
115+ src = out , dst = orig_out , sycl_queue = exec_q , depends = [unary_ev ]
116+ )
117+ ht_copy_ev .wait ()
118+ out = orig_out
103119
120+ ht_unary_ev .wait ()
104121 return out
122+
105123 if order == "K" :
106124 buf = _empty_like_orderK (x , buf_dt )
107125 else :
@@ -117,11 +135,6 @@ def __call__(self, x, out=None, order="K"):
117135 out = _empty_like_orderK (buf , res_dt )
118136 else :
119137 out = dpt .empty_like (buf , dtype = res_dt , order = order )
120- else :
121- if buf_dt != out .dtype :
122- raise TypeError (
123- f"Output array of type { buf_dt } { out .dtype }
124- )
125138
126139 ht , _ = self .unary_fn_ (buf , out , sycl_queue = exec_q , depends = [copy_ev ])
127140 ht_copy_ev .wait ()
0 commit comments