Mark temporary conversion outputs as linear for eager storage recycling

cunumeric/_ufunc/ufunc.py

@@ -218,7 +218,7 @@ def _maybe_cast_input(
                 f"{arr.dtype} to {to_dtype} with casting rule '{casting}'"
             )
 
-        return arr.astype(to_dtype)
+        return arr._astype(to_dtype, temporary=True)
 
     def _maybe_create_result(
         self,
@@ -366,7 +366,7 @@ def _resolve_dtype(
         if not precision_fixed:
             if arr.dtype in self._resolution_cache:
                 to_dtype = self._resolution_cache[arr.dtype]
-                arr = arr.astype(to_dtype)
+                arr = arr._astype(to_dtype, temporary=True)
                 return arr, np.dtype(self._types[to_dtype.char])
 
         chosen = None
@@ -385,7 +385,9 @@ def _resolve_dtype(
         to_dtype = np.dtype(chosen)
         self._resolution_cache[arr.dtype] = to_dtype
 
-        return arr.astype(to_dtype), np.dtype(self._types[to_dtype.char])
+        return arr._astype(to_dtype, temporary=True), np.dtype(
+            self._types[to_dtype.char]
+        )
 
     def __call__(
         self,
@@ -449,7 +451,7 @@ def _resolve_dtype(
         if not precision_fixed:
             if arr.dtype in self._resolution_cache:
                 to_dtype = self._resolution_cache[arr.dtype]
-                arr = arr.astype(to_dtype)
+                arr = arr._astype(to_dtype, temporary=True)
                 return arr, to_dtypes(self._types[to_dtype.char])
 
         chosen = None
@@ -468,7 +470,9 @@ def _resolve_dtype(
         to_dtype = np.dtype(chosen)
         self._resolution_cache[arr.dtype] = to_dtype
 
-        return arr.astype(to_dtype), to_dtypes(self._types[to_dtype.char])
+        return arr._astype(to_dtype, temporary=True), to_dtypes(
+            self._types[to_dtype.char]
+        )
 
     def __call__(
         self,
@@ -597,15 +601,16 @@ def _resolve_dtype(
 
         if key in self._types:
             arrs = [
-                arr.astype(to_dtype) for arr, to_dtype in zip(arrs, to_dtypes)
+                arr._astype(to_dtype, temporary=True)
+                for arr, to_dtype in zip(arrs, to_dtypes)
             ]
             return arrs, np.dtype(self._types[key])
 
         if not precision_fixed:
             if key in self._resolution_cache:
                 to_dtypes = self._resolution_cache[key]
                 arrs = [
-                    arr.astype(to_dtype)
+                    arr._astype(to_dtype, temporary=True)
                     for arr, to_dtype in zip(arrs, to_dtypes)
                 ]
                 return arrs, np.dtype(self._types[to_dtypes])
@@ -638,7 +643,10 @@ def _resolve_dtype(
             )
 
         self._resolution_cache[key] = chosen
-        arrs = [arr.astype(to_dtype) for arr, to_dtype in zip(arrs, chosen)]
+        arrs = [
+            arr._astype(to_dtype, temporary=True)
+            for arr, to_dtype in zip(arrs, chosen)
+        ]
 
         return arrs, np.dtype(self._types[chosen])
 

cunumeric/array.py

@@ -1860,6 +1860,17 @@ def astype(
         Multiple GPUs, Multiple CPUs
 
         """
+        return self._astype(dtype, order, casting, subok, copy, False)
+
+    def _astype(
+        self,
+        dtype: npt.DTypeLike,
+        order: OrderType = "C",
+        casting: CastingKind = "unsafe",
+        subok: bool = True,
+        copy: bool = True,
+        temporary: bool = False,
+    ) -> ndarray:
         dtype = np.dtype(dtype)
         if self.dtype == dtype:
             return self
@@ -1879,7 +1890,7 @@ def astype(
                 f"to the rule '{casting}'"
             )
         result = ndarray(self.shape, dtype=dtype, inputs=(self,))
-        result._thunk.convert(self._thunk, warn=False)
+        result._thunk.convert(self._thunk, warn=False, temporary=temporary)
         return result
 
     @add_boilerplate()

cunumeric/deferred.py

@@ -1184,6 +1184,7 @@ def convert(
         rhs: Any,
         warn: bool = True,
         nan_op: ConvertCode = ConvertCode.NOOP,
+        temporary: bool = False,
     ) -> None:
         lhs_array = self
         rhs_array = rhs
@@ -1210,6 +1211,9 @@ def convert(
 
         task.execute()
 
+        if temporary:
+            lhs.set_linear()
+
     @auto_convert([1, 2])
     def convolve(self, v: Any, lhs: Any, mode: ConvolveMode) -> None:
         input = self.base

cunumeric/eager.py

@@ -490,6 +490,7 @@ def convert(
         rhs: Any,
         warn: bool = True,
         nan_op: ConvertCode = ConvertCode.NOOP,
+        temporary: bool = False,
     ) -> None:
         self.check_eager_args(rhs)
         if self.deferred is not None:

cunumeric/thunk.py

@@ -157,6 +157,7 @@ def convert(
         rhs: Any,
         warn: bool = True,
         nan_op: ConvertCode = ConvertCode.NOOP,
+        temporary: bool = False,
     ) -> None:
         ...