Add 1D interpolation and Bilinear interpolation

.dev/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"

.dev/climaformat.jl

@@ -0,0 +1,83 @@
+#!/usr/bin/env julia
+#
+# This is an adapted version of format.jl from JuliaFormatter with the
+# following license:
+#
+#    MIT License
+#    Copyright (c) 2019 Dominique Luna
+#
+#    Permission is hereby granted, free of charge, to any person obtaining a
+#    copy of this software and associated documentation files (the
+#    "Software"), to deal in the Software without restriction, including
+#    without limitation the rights to use, copy, modify, merge, publish,
+#    distribute, sublicense, and/or sell copies of the Software, and to permit
+#    persons to whom the Software is furnished to do so, subject to the
+#    following conditions:
+#
+#    The above copyright notice and this permission notice shall be included
+#    in all copies or substantial portions of the Software.
+#
+#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+#    OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+#    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN
+#    NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+#    DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+#    OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
+#    USE OR OTHER DEALINGS IN THE SOFTWARE.
+#
+using Pkg
+Pkg.activate(@__DIR__)
+Pkg.instantiate()
+
+using JuliaFormatter
+
+help = """
+Usage: climaformat.jl [flags] [FILE/PATH]...
+
+Formats the given julia files using the CLIMA formatting options.  If paths
+are given it will format the julia files in the paths. Otherwise, it will
+format all changed julia files.
+
+    -v, --verbose
+        Print the name of the files being formatted with relevant details.
+
+    -h, --help
+        Print this message.
+"""
+
+function parse_opts!(args::Vector{String})
+    i = 1
+    opts = Dict{Symbol,Union{Int,Bool}}()
+    while i ≤ length(args)
+        arg = args[i]
+        if arg[1] != '-'
+            i += 1
+            continue
+        end
+        if arg == "-v" || arg == "--verbose"
+            opt = :verbose
+        elseif arg == "-h" || arg == "--help"
+            opt = :help
+        else
+            error("invalid option $arg")
+        end
+        if opt in (:verbose, :help)
+            opts[opt] = true
+            deleteat!(args, i)
+        end
+    end
+    return opts
+end
+
+opts = parse_opts!(ARGS)
+if haskey(opts, :help)
+    write(stdout, help)
+    exit(0)
+end
+if isempty(ARGS)
+    filenames = readlines(`git ls-files "*.jl"`)
+else
+    filenames = ARGS
+end
+
+format(filenames; opts...)

Project.toml

@@ -0,0 +1,15 @@
+name = "ClimaInterpolations"
+uuid = "dd0f122e-fa3b-47f3-bcf0-93bbc60d885e"
+authors = ["sriharshakandala <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+
+[extensions]
+ClimaInterpolationsCUDAExt = "CUDA"
+ClimaInterpolationsMetalExt = "Metal"

ext/ClimaInterpolationsCUDAExt.jl

@@ -0,0 +1,19 @@
+module ClimaInterpolationsCUDAExt
+
+using CUDA
+
+import ClimaInterpolations
+import ClimaInterpolations.Interpolation1D:
+    interpolate,
+    interpolate1d!,
+    Order1D,
+    Extrapolate1D,
+    Linear,
+    Flat,
+    LinearExtrapolation,
+    get_stencil
+
+include(joinpath("cuda", "interpolation1d.jl"))
+include(joinpath("cuda", "bilinearinterpolation.jl"))
+
+end

ext/ClimaInterpolationsMetalExt.jl

@@ -0,0 +1,18 @@
+module ClimaInterpolationsMetalExt
+
+using Metal
+
+import ClimaInterpolations
+import ClimaInterpolations.Interpolation1D:
+    interpolate,
+    interpolate1d!,
+    Order1D,
+    Extrapolate1D,
+    Linear,
+    Flat,
+    LinearExtrapolation,
+    get_stencil
+
+include(joinpath("metal", "interpolation1d.jl"))
+include(joinpath("metal", "bilinearinterpolation.jl"))
+end

ext/cuda/bilinearinterpolation.jl

@@ -0,0 +1,153 @@
+import ..ClimaInterpolations.BilinearInterpolation:
+    Bilinear, set_source_range!, interpolatebilinear!, get_stencil_bilinear1d, get_dims
+
+function interpolatebilinear!(
+    ftarget::CuArray{FT,N},
+    bilinear::B,
+    fsource::CuArray{FT,N},
+) where {FT,N,B}
+    @assert N ≥ 2
+    (leveldimssource..., nxs, nys) = size(fsource)
+    (leveldimstarget..., nxt, nyt) = size(ftarget)
+
+    nsourcex, nsourcey, ntargetx, ntargety = get_dims(bilinear)
+
+    @assert leveldimssource == leveldimstarget
+    @assert nxs == nsourcex && nys == nsourcey
+    @assert nxt == ntargetx && nyt == ntargety
+    # max threadgroups per grid is not officially specified by Metal
+    # It is possible this number is unlimited on newer Apple models (>m3)
+    maxxblocks = 1024
+    maxyblocks = 1024
+    # target n loops per dimension per threadgroup
+    targetloopsperblockdimx = 4
+    targetloopsperblockdimy = 4
+    nxblocks = min(cld(ntargetx, targetloopsperblockdimx), maxxblocks)
+    nyblocks = min(cld(ntargety, targetloopsperblockdimy), maxyblocks)
+
+    nxloops = cld(ntargetx, nxblocks)
+    nyloops = cld(ntargety, nyblocks)
+
+    levelcidxs = CartesianIndices(leveldimssource)
+    nlevels = length(levelcidxs)
+
+    kernel = @cuda launch = false interpolatebilinear_kernel!(
+        ftarget,
+        bilinear,
+        fsource,
+        levelcidxs,
+        (nxloops, ntargetx),
+        (nyloops, ntargety),
+        (1, 1),
+    )
+    kernel_config = CUDA.launch_configuration(kernel.fun)
+    nthreads = min(kernel_config.threads, nlevels)
+    nzloops = cld(nlevels, nthreads)
+
+    @cuda threads = nthreads blocks = (nxblocks, nyblocks) interpolatebilinear_kernel!(
+        ftarget,
+        bilinear,
+        fsource,
+        levelcidxs,
+        (nxloops, ntargetx),
+        (nyloops, ntargety),
+        (nzloops, nlevels),
+    )
+    return nothing
+end
+
+function interpolatebilinear_kernel!(
+    ftarget::AbstractArray{FT,N},
+    bilinear::B,
+    fsource::AbstractArray{FT,N},
+    levelcidxs,
+    (nxloops, ntargetx),
+    (nyloops, ntargety),
+    (nzloops, nlevels),
+) where {FT,N,B}
+    (; sourcex, sourcey, targetx, targety, startx, starty) = bilinear
+
+    ixtg, iytg = blockIdx().x, blockIdx().y
+    izt = threadIdx().x
+    nzthreads = blockDim().x
+    (nxblocks, nyblocks) = gridDim().x, gridDim().y
+
+    @inbounds begin
+        for yl = 1:nyloops
+            iy = iytg + (yl - 1) * nyblocks
+            if iy ≤ ntargety
+                y, sty = targety[iy], starty[iy]
+                y1, y2 = sourcey[sty], sourcey[sty+1]
+                dy1 = y - y1
+                dy2 = y2 - y
+                for xl = 1:nxloops
+                    ix = ixtg + (xl - 1) * nxblocks
+                    if ix ≤ ntargetx
+                        x, stx = targetx[ix], startx[ix]
+                        x1, x2 = sourcex[stx], sourcex[stx+1]
+
+                        dx1 = x - x1
+                        dx2 = x2 - x
+
+                        fac = FT(1) / ((x2 - x1) * (y2 - y1))
+
+                        for zl = 1:nzloops
+                            iz = izt + (zl - 1) * nzthreads
+                            if iz ≤ nlevels
+                                levelidx = Tuple(levelcidxs[iz])
+
+                                ftarget[levelidx..., ix, iy] =
+                                    (
+                                        dx1 * (
+                                            dy2 * fsource[levelidx..., stx+1, sty] +
+                                            dy1 * fsource[levelidx..., stx+1, sty+1]
+                                        ) +
+                                        dx2 * (
+                                            dy1 * fsource[levelidx..., stx, sty+1] +
+                                            dy2 * fsource[levelidx..., stx, sty]
+                                        )
+                                    ) * fac
+                            end
+                        end
+                    end
+                end
+            end
+        end
+
+    end
+    return nothing
+end
+
+CUDA.Adapt.@adapt_structure Bilinear
+
+function Bilinear(sourcex::V, sourcey::V, targetx::V, targety::V) where {V<:CuVector}
+    startx = similar(targetx, Int)
+    starty = similar(targety, Int)
+    # Call set_source_range_kernel! with one thread per group
+    # and two threadgroups
+    @cuda threads = 1 blocks = 2 set_source_range_kernel!(
+        startx,
+        sourcex,
+        targetx,
+        starty,
+        sourcey,
+        targety,
+    )
+    return Bilinear(sourcex, sourcey, targetx, targety, startx, starty)
+end
+
+function set_source_range_kernel!(
+    startx::AbstractVector{I},
+    sourcex::AbstractVector{FT},
+    targetx::AbstractVector{FT},
+    starty::AbstractVector{I},
+    sourcey::AbstractVector{FT},
+    targety::AbstractVector{FT},
+) where {I,FT}
+    order = Linear()
+    bid = blockIdx().x
+    tid = threadIdx().x
+    (bid == 1 && tid == 1) && set_source_range!(startx, sourcex, targetx)
+    (bid == 2 && tid == 1) && set_source_range!(starty, sourcey, targety)
+    return nothing
+end

ext/cuda/interpolation1d.jl

@@ -0,0 +1,79 @@
+
+function interpolate1d!(
+    ftarget::CuArray{FT,N},
+    xsource::CuArray{FT,N},
+    xtarget::CuArray{FT,N},
+    fsource::CuArray{FT,N},
+    order,
+    extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+) where {FT,N}
+    @assert N ≥ 1
+    (nsource, coldims_source...) = size(xsource)
+    (ntarget, coldims_target...) = size(xtarget)
+    @assert coldims_source == coldims_target
+    @assert size(ftarget) == size(xtarget)
+    @assert size(fsource) == size(xsource)
+    coldims = coldims_source
+    colcidxs = CartesianIndices(coldims)
+    nitems = length(colcidxs)
+    kernel = @cuda launch = false interpolate1d_kernel!(
+        ftarget,
+        xsource,
+        xtarget,
+        fsource,
+        order,
+        extrapolate,
+        colcidxs,
+    )
+    kernel_config = CUDA.launch_configuration(kernel.fun)
+    nthreads = min(kernel_config.threads, nitems)
+    nblocks = convert(Int, cld(nitems, nthreads))
+
+    @cuda threads = nthreads blocks = nblocks interpolate1d_kernel!(
+        ftarget,
+        xsource,
+        xtarget,
+        fsource,
+        order,
+        extrapolate,
+        colcidxs,
+    )
+    return nothing
+end
+
+function interpolate1d_kernel!(
+    ftarget::AbstractArray{FT,N},
+    xsource::AbstractArray{FT,N},
+    xtarget::AbstractArray{FT,N},
+    fsource::AbstractArray{FT,N},
+    order,
+    extrapolate,
+    colcidxs,
+) where {FT,N}
+    nitems = length(colcidxs)
+    gid = threadIdx().x + (blockIdx().x - 1) * blockDim().x
+
+    if gid ≤ nitems
+        colidx = Tuple(colcidxs[gid])
+        nsource = size(xsource, 1)
+        ntarget = size(xtarget, 1)
+        first = 1
+
+        for i1 = 1:ntarget
+            (st, en) = get_stencil(
+                order,
+                view(xsource, 1:nsource, colidx...),
+                xtarget[i1, colidx...],
+                first = first,
+                extrapolate = extrapolate,
+            )
+            first = st
+            ftarget[i1, colidx...] = interpolate(
+                view(xsource, st:en, colidx...),
+                view(fsource, st:en, colidx...),
+                xtarget[i1, colidx...],
+            )
+        end
+    end
+    return nothing
+end