[NDTensors] Optimize permutedims (#1288)

ITensor · Dec 16, 2023 · faa0357 · faa0357
1 parent a00ef70
commit faa0357
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 10 deletions.
diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml
@@ -64,7 +64,7 @@ SparseArrays = "1.6"
 SplitApplyCombine = "1.2.2"
 StaticArrays = "0.12, 1.0"
 Strided = "2"
-StridedViews = "0.2"
+StridedViews = "0.2.2"
 TimerOutputs = "0.5.5"
 TupleTools = "1.2.0"
 VectorInterface = "0.4.2"

diff --git a/NDTensors/src/array/permutedims.jl b/NDTensors/src/array/permutedims.jl
@@ -10,11 +10,15 @@ function permutedims(E::Exposed{<:Array}, perm)
 end
 
 function permutedims!(Edest::Exposed{<:Array}, Esrc::Exposed{<:Array}, perm)
-  @strided unexpose(Edest) .= permutedims(Esrc, perm)
-  return unexpose(Edest)
+  a_dest = unexpose(Edest)
+  a_src = unexpose(Esrc)
+  @strided a_dest .= permutedims(a_src, perm)
+  return a_dest
 end
 
 function permutedims!(Edest::Exposed{<:Array}, Esrc::Exposed{<:Array}, perm, f)
-  @strided unexpose(Edest) .= f.(unexpose(Edest), permutedims(Esrc, perm))
-  return unexpose(Edest)
+  a_dest = unexpose(Edest)
+  a_src = unexpose(Esrc)
+  @strided a_dest .= f.(a_dest, permutedims(a_src, perm))
+  return a_dest
 end
diff --git a/NDTensors/src/dense/densetensor.jl b/NDTensors/src/dense/densetensor.jl
@@ -193,13 +193,20 @@ end
 # Maybe allocate output data.
 # TODO: Remove this in favor of `map!`
 # applied to `PermutedDimsArray`.
-function permutedims!!(R::DenseTensor, T::DenseTensor, perm, f::Function=(r, t) -> t)
+function permutedims!!(R::DenseTensor, T::DenseTensor, perm, f::Function)
   Base.checkdims_perm(R, T, perm)
   RR = convert(promote_type(typeof(R), typeof(T)), R)
   permutedims!(RR, T, perm, f)
   return RR
 end
 
+function permutedims!!(R::DenseTensor, T::DenseTensor, perm)
+  Base.checkdims_perm(R, T, perm)
+  RR = convert(promote_type(typeof(R), typeof(T)), R)
+  permutedims!(RR, T, perm)
+  return RR
+end
+
 # TODO: call permutedims!(R,T,perm,(r,t)->t)?
 function permutedims!(
   R::DenseTensor{<:Number,N,StoreT}, T::DenseTensor{<:Number,N,StoreT}, perm::NTuple{N,Int}
@@ -216,7 +223,7 @@ function permutedims!(
 ) where {N}
   RA = array(R)
   TA = array(T)
-  RA .= permutedims(expose(TA), perm)
+  permutedims!(expose(RA), expose(TA), perm)
   return R
 end
 

diff --git a/NDTensors/src/lib/Unwrap/test/Project.toml b/NDTensors/src/lib/Unwrap/test/Project.toml
@@ -0,0 +1,4 @@
+[deps]
+GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+NDTensors = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/NDTensors/src/tensoroperations/generic_tensor_operations.jl b/NDTensors/src/tensoroperations/generic_tensor_operations.jl
@@ -9,16 +9,31 @@ end
 # and return the result of the permutation.
 # Similar to `BangBang.jl` notation:
 # https://juliafolds.github.io/BangBang.jl/stable/.
-function permutedims!!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)
+function permutedims!!(output_tensor::Tensor, tensor::Tensor, perm, f::Function)
   Base.checkdims_perm(output_tensor, tensor, perm)
   permutedims!(output_tensor, tensor, perm, f)
   return output_tensor
 end
 
-function permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)
+# Equivalent to `permutedims!!(output_tensor, tensor, perm, (r, t) -> t)`
+function permutedims!!(output_tensor::Tensor, tensor::Tensor, perm)
+  Base.checkdims_perm(output_tensor, tensor, perm)
+  permutedims!(output_tensor, tensor, perm)
+  return output_tensor
+end
+
+function permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function)
+  Base.checkdims_perm(output_tensor, tensor, perm)
+  error(
+    "`permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.",
+  )
+  return output_tensor
+end
+
+function permutedims!(output_tensor::Tensor, tensor::Tensor, perm)
   Base.checkdims_perm(output_tensor, tensor, perm)
   error(
-    "`permutedims!(output_tensor::Tensor, tensor::Tensor, perm, f::Function=(r, t) -> t)` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, `perm = $perm`, and `f = $f`.",
+    "`permutedims!(output_tensor::Tensor, tensor::Tensor, perm` not implemented for `typeof(output_tensor) = $(typeof(output_tensor))`, `typeof(tensor) = $(typeof(tensor))`, and `perm = $perm`.",
   )
   return output_tensor
 end