Remove MTL from MtlArray printing

JuliaGPU · Jul 22, 2024 · 4ec2c4a · 4ec2c4a
1 parent 78e8fe2
commit 4ec2c4a
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 13 deletions.
diff --git a/docs/src/usage/array.md b/docs/src/usage/array.md
@@ -28,16 +28,16 @@ can construct `MtlArray`s in the same way as regular `Array` objects:
 
 ```jldoctest
 julia> MtlArray{Int}(undef, 2)
-2-element MtlVector{Int64, Metal.MTL.PrivateStorage}:
+2-element MtlVector{Int64, Metal.PrivateStorage}:
  0
  0
 
 julia> MtlArray{Int}(undef, (1,2))
-1×2 MtlMatrix{Int64, Metal.MTL.PrivateStorage}:
+1×2 MtlMatrix{Int64, Metal.PrivateStorage}:
  0  0
 
 julia> similar(ans)
-1×2 MtlMatrix{Int64, Metal.MTL.PrivateStorage}:
+1×2 MtlMatrix{Int64, Metal.PrivateStorage}:
  0  0
 ```
 
@@ -46,7 +46,7 @@ Copying memory to or from the GPU can be expressed using constructors as well, o
 
 ```jldoctest
 julia> a = MtlArray([1,2])
-2-element MtlVector{Int64, Metal.MTL.PrivateStorage}:
+2-element MtlVector{Int64, Metal.PrivateStorage}:
  1
  2
 
@@ -73,11 +73,11 @@ perform simple element-wise operations you can use `map` or `broadcast`:
 julia> a = MtlArray{Float32}(undef, (1,2));
 
 julia> a .= 5
-1×2 MtlMatrix{Float32, Metal.MTL.PrivateStorage}:
+1×2 MtlMatrix{Float32, Metal.PrivateStorage}:
  5.0  5.0
 
 julia> map(sin, a)
-1×2 MtlMatrix{Float32, Metal.MTL.PrivateStorage}:
+1×2 MtlMatrix{Float32, Metal.PrivateStorage}:
  -0.958924  -0.958924
 ```
 
@@ -86,23 +86,23 @@ To reduce the dimensionality of arrays, Metal.jl implements the various flavours
 
 ```jldoctest
 julia> a = Metal.ones(2,3)
-2×3 MtlMatrix{Float32, Metal.MTL.PrivateStorage}:
+2×3 MtlMatrix{Float32, Metal.PrivateStorage}:
  1.0  1.0  1.0
  1.0  1.0  1.0
 
 julia> reduce(+, a)
 6.0f0
 
 julia> mapreduce(sin, *, a; dims=2)
-2×1 MtlMatrix{Float32, Metal.MTL.PrivateStorage}:
+2×1 MtlMatrix{Float32, Metal.PrivateStorage}:
  0.59582335
  0.59582335
 
 julia> b = Metal.zeros(1)
-1-element MtlVector{Float32, Metal.MTL.PrivateStorage}:
+1-element MtlVector{Float32, Metal.PrivateStorage}:
  0.0
 
 julia> Base.mapreducedim!(identity, +, b, a)
-1×1 MtlMatrix{Float32, Metal.MTL.PrivateStorage}:
+1×1 MtlMatrix{Float32, Metal.PrivateStorage}:
  6.0
 ```
diff --git a/lib/mtl/storage_type.jl b/lib/mtl/storage_type.jl
@@ -43,6 +43,11 @@ See also [`Metal.SharedStorage`](@ref) and [`Metal.ManagedStorage`](@ref).
 struct PrivateStorage     <: StorageMode end
 struct Memoryless  <: StorageMode end
 
+# Remove the ".MTL" when printing
+Base.show(io::IO, ::Type{<:PrivateStorage}) = print(io, "Metal.PrivateStorage")
+Base.show(io::IO, ::Type{<:SharedStorage}) = print(io, "Metal.SharedStorage")
+Base.show(io::IO, ::Type{<:ManagedStorage}) = print(io, "Metal.ManagedStorage")
+
 """
     CPUStorage
 

diff --git a/src/array.jl b/src/array.jl
@@ -480,18 +480,18 @@ Uses Adapt.jl to act inside some wrapper structs.
 
 ```jldoctests
 julia> mtl(ones(3)')
-1×3 adjoint(::MtlVector{Float32, Metal.MTL.PrivateStorage}) with eltype Float32:
+1×3 adjoint(::MtlVector{Float32, Metal.PrivateStorage}) with eltype Float32:
  1.0  1.0  1.0
 
 julia> mtl(zeros(1,3); storage=Metal.SharedStorage)
-1×3 MtlMatrix{Float32, Metal.MTL.SharedStorage}:
+1×3 MtlMatrix{Float32, Metal.SharedStorage}:
  0.0  0.0  0.0
 
 julia> mtl(1:3)
 1:3
 
 julia> MtlArray(1:3)
-3-element MtlVector{Int64, Metal.MTL.PrivateStorage}:
+3-element MtlVector{Int64, Metal.PrivateStorage}:
  1
  2
  3