Update GPU package extensions

Author: mtfishman

NDTensors/Project.toml

@@ -45,13 +45,13 @@ TBLIS = "48530278-0828-4a49-9772-0f3830dfa1e9"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
-NDTensorsAMDGPUExt = ["AMDGPU", "GPUArraysCore"]
-NDTensorsCUDAExt = ["CUDA", "GPUArraysCore"]
+NDTensorsAMDGPUExt = ["AMDGPU", "GPUArraysCore", "TypeParameterAccessors"]
+NDTensorsCUDAExt = ["CUDA", "GPUArraysCore", "TypeParameterAccessors"]
 NDTensorsGPUArraysCoreExt = "GPUArraysCore"
 NDTensorsHDF5Ext = "HDF5"
 NDTensorsJLArraysExt = ["GPUArraysCore", "JLArrays"]
 NDTensorsMappedArraysExt = ["MappedArrays"]
-NDTensorsMetalExt = ["GPUArraysCore", "Metal"]
+NDTensorsMetalExt = ["GPUArraysCore", "Metal", "TypeParameterAccessors"]
 NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"
 NDTensorscuTENSORExt = "cuTENSOR"

NDTensors/ext/NDTensorsAMDGPUExt/set_types.jl

@@ -1,11 +1,5 @@
-# TypeParameterAccessors definitions
-using NDTensors.TypeParameterAccessors:
-  TypeParameterAccessors, Position, default_type_parameters
+using AMDGPU: ROCArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
-using AMDGPU: AMDGPU, ROCArray
-
-function TypeParameterAccessors.default_type_parameters(::Type{<:ROCArray})
-  return (default_type_parameters(AbstractArray)..., AMDGPU.Mem.HIPBuffer)
-end
+using TypeParameterAccessors: TypeParameterAccessors, Position
 
 TypeParameterAccessors.position(::Type{<:ROCArray}, ::typeof(storagemode)) = Position(3)

NDTensors/ext/NDTensorsCUDAExt/set_types.jl

@@ -1,13 +1,7 @@
-# TypeParameterAccessors definitions
-using CUDA: CUDA, CuArray
-using NDTensors.TypeParameterAccessors:
-  TypeParameterAccessors, Position, default_type_parameters
+using CUDA: CuArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
+using TypeParameterAccessors: TypeParameterAccessors, Position
 
 function TypeParameterAccessors.position(::Type{<:CuArray}, ::typeof(storagemode))
   return Position(3)
 end
-
-function TypeParameterAccessors.default_type_parameters(::Type{<:CuArray})
-  return (default_type_parameters(AbstractArray)..., CUDA.Mem.DeviceBuffer)
-end

NDTensors/ext/NDTensorsMetalExt/set_types.jl

@@ -1,13 +1,7 @@
-using Metal: Metal, MtlArray
-# `TypeParameterAccessors.jl` definitions.
-
-using NDTensors.TypeParameterAccessors: TypeParameterAccessors, Position
+using Metal: MtlArray
 using NDTensors.GPUArraysCoreExtensions: storagemode
+using TypeParameterAccessors: TypeParameterAccessors, Position
 
 function TypeParameterAccessors.position(::Type{<:MtlArray}, ::typeof(storagemode))
   return Position(3)
 end
-
-function TypeParameterAccessors.default_type_parameters(::Type{<:MtlArray})
-  return (Float32, 1, Metal.DefaultStorageMode)
-end