cleanup

Author: AntonOresten

src/experimental/autotune.jl

@@ -15,11 +15,6 @@ const TUNING_PRESETS = (
     thorough = (warmup=2, reps=7, refine_topk=4, refine_reps=6),
 )
 
-const TUNING_KEYS = (
-    :warmup, :reps, :refine_topk, :refine_reps,
-    :seed, :force, :precompile_workers,
-)
-
 struct TuningOptions
     warmup::Int
     reps::Int
@@ -30,65 +25,38 @@ struct TuningOptions
     precompile_workers::Int
 end
 
-struct TuningSession
-    id::UInt
-end
-
-const NEXT_TUNING_SESSION_ID = Threads.Atomic{UInt}(0)
-
-TuningSession() =
-    TuningSession(Threads.atomic_add!(NEXT_TUNING_SESSION_ID, UInt(1)) + UInt(1))
-
 _tuning_defaults() = (
     seed=nothing,
     force=false,
     precompile_workers=Threads.nthreads(),
 )
 
-function _check_int(name::Symbol, value; min::Int)
-    (value isa Integer && !(value isa Bool)) ||
-        throw(ArgumentError("tuning.$name must be an integer, got $(typeof(value))"))
-    value >= min ||
-        throw(ArgumentError("tuning.$name must be >= $min, got $value"))
-    return Int(value)
-end
-
-function _check_seed(seed)
-    seed === nothing && return nothing
-    (seed isa Integer && !(seed isa Bool)) ||
-        throw(ArgumentError("tuning.seed must be an integer or nothing, got $(typeof(seed))"))
-    return Int(seed)
-end
-
-function _check_bool(name::Symbol, value)
-    value isa Bool ||
-        throw(ArgumentError("tuning.$name must be a Bool, got $(typeof(value))"))
-    return value
-end
+# Lower bound for each count field; the others have no minimum.
+const _TUNING_MINIMA = (warmup=0, reps=1, refine_topk=0, refine_reps=1,
+                        precompile_workers=0)
 
 function normalize_tuning(tuning::NamedTuple)
-    valid_keys = (:preset, TUNING_KEYS...)
-    unknown = setdiff(collect(keys(tuning)), collect(valid_keys))
+    valid_keys = (:preset, fieldnames(TuningOptions)...)
+    unknown = setdiff(keys(tuning), valid_keys)
     isempty(unknown) ||
         throw(ArgumentError("Unknown tuning option(s): $(join(unknown, ", "))"))
 
     preset = get(tuning, :preset, :default)
-    preset isa Symbol || throw(ArgumentError("tuning.preset must be a Symbol"))
-    hasproperty(TUNING_PRESETS, preset) ||
+    preset isa Symbol && hasproperty(TUNING_PRESETS, preset) ||
         throw(ArgumentError("Unknown tuning preset `$preset`; use :fast, :default, or :thorough"))
 
     overrides = NamedTuple(k => v for (k, v) in pairs(tuning) if k !== :preset)
     values = merge(_tuning_defaults(), getproperty(TUNING_PRESETS, preset), overrides)
 
-    return TuningOptions(
-        _check_int(:warmup, values.warmup; min=0),
-        _check_int(:reps, values.reps; min=1),
-        _check_int(:refine_topk, values.refine_topk; min=0),
-        _check_int(:refine_reps, values.refine_reps; min=1),
-        _check_seed(values.seed),
-        _check_bool(:force, values.force),
-        _check_int(:precompile_workers, values.precompile_workers; min=0),
-    )
+    # The struct's field types coerce/reject bad value types; we only enforce
+    # the lower bounds that the types can't. Pull fields by name since `values`
+    # is in merge order, not struct-field order.
+    opts = TuningOptions((getproperty(values, f) for f in fieldnames(TuningOptions))...)
+    for (name, lo) in pairs(_TUNING_MINIMA)
+        getfield(opts, name) >= lo ||
+            throw(ArgumentError("tuning.$name must be >= $lo, got $(getfield(opts, name))"))
+    end
+    return opts
 end
 
 @inline _hint_from_cfg(cfg, name::Symbol, fallback) =
@@ -101,22 +69,14 @@ function hints_from_cfg(cfg; static_num_ctas=nothing, static_occupancy=nothing)
     )
 end
 
-function _check_static_hint_conflict(configs, hint::Symbol, static_value)
-    static_value === nothing && return nothing
-    for cfg in configs
-        if hasproperty(cfg, hint)
-            throw(ArgumentError(
-                "`$hint` is both a static kwarg and an axis in the search space. " *
-                "Pick one."))
-        end
-    end
-    return nothing
-end
-
 function _check_static_hint_conflicts(configs; static_num_ctas=nothing,
                                       static_occupancy=nothing)
-    _check_static_hint_conflict(configs, :num_ctas, static_num_ctas)
-    _check_static_hint_conflict(configs, :occupancy, static_occupancy)
+    statics = (num_ctas=static_num_ctas, occupancy=static_occupancy)
+    for (hint, static_value) in pairs(statics)
+        static_value === nothing && continue
+        any(cfg -> hasproperty(cfg, hint), configs) && throw(ArgumentError(
+            "`$hint` is both a static kwarg and an axis in the search space. Pick one."))
+    end
     return nothing
 end
 
@@ -132,11 +92,11 @@ end
 @inline _converted_args(args_fn, cfg) = map(cuTileconvert, args_fn(cfg))
 @inline _argtypes(args) = Tuple{map(Core.Typeof, args)...}
 
-function _compile_cfg(@nospecialize(f), cfg, args_fn, session::TuningSession;
+function _compile_cfg(@nospecialize(f), cfg, args_fn;
                       sm_arch::VersionNumber, opt_level::Int,
                       static_num_ctas=nothing, static_occupancy=nothing)
     converted = _converted_args(args_fn, cfg)
-    return temporary_cufunction(f, _argtypes(converted), session.id;
+    return temporary_cufunction(f, _argtypes(converted);
         sm_arch, opt_level,
         hints_from_cfg(cfg; static_num_ctas, static_occupancy)...)
 end
@@ -166,34 +126,33 @@ function _time_ms(run_once, get_args;
     return best_ms
 end
 
-function eval_cfg(@nospecialize(f), cfg, grid_fn, args_fn, session::TuningSession;
+function eval_cfg(@nospecialize(f), cfg, grid_fn, args_fn;
                   sm_arch::VersionNumber, opt_level::Int, warmup::Int, reps::Int,
                   static_num_ctas=nothing, static_occupancy=nothing,
                   verify=nothing, reset=nothing)
     grid = _grid_dims(grid_fn(cfg))
-    kernel = _compile_cfg(f, cfg, args_fn, session; sm_arch, opt_level,
+    kernel = _compile_cfg(f, cfg, args_fn; sm_arch, opt_level,
                           static_num_ctas, static_occupancy)
 
     run_once = converted -> kernel(converted...; blocks=grid)
     get_args = () -> _converted_args(args_fn, cfg)
     return _time_ms(run_once, get_args; warmup, reps, verify, reset)
 end
 
-function precompile_cfg(@nospecialize(f), cfg, args_fn, session::TuningSession;
+function precompile_cfg(@nospecialize(f), cfg, args_fn;
                         sm_arch::VersionNumber, opt_level::Int,
                         static_num_ctas=nothing, static_occupancy=nothing)
-    _compile_cfg(f, cfg, args_fn, session; sm_arch, opt_level,
+    _compile_cfg(f, cfg, args_fn; sm_arch, opt_level,
                  static_num_ctas, static_occupancy)
     return nothing
 end
 
 const TimingRecord = Tuple{Any, Float32}
 
 function _measure_cfg!(record::Vector{TimingRecord}, first_error::Base.RefValue,
-                       @nospecialize(f), cfg, grid_fn, args_fn,
-                       session::TuningSession; kwargs...)
+                       @nospecialize(f), cfg, grid_fn, args_fn; kwargs...)
     ms = try
-        eval_cfg(f, cfg, grid_fn, args_fn, session; kwargs...)
+        eval_cfg(f, cfg, grid_fn, args_fn; kwargs...)
     catch err
         err isa InterruptException && rethrow()
         if err isa VerificationError
@@ -209,16 +168,15 @@ function _measure_cfg!(record::Vector{TimingRecord}, first_error::Base.RefValue,
     return nothing
 end
 
-function measure_candidates(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn,
-                            session::TuningSession;
+function measure_candidates(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn;
                             sm_arch::VersionNumber, opt_level::Int,
                             warmup::Int, reps::Int,
                             static_num_ctas=nothing, static_occupancy=nothing,
                             verify=nothing, reset=nothing)
     record = TimingRecord[]
     first_error = Ref{Any}(nothing)
     for cfg in configs
-        _measure_cfg!(record, first_error, f, cfg, grid_fn, args_fn, session;
+        _measure_cfg!(record, first_error, f, cfg, grid_fn, args_fn;
                       sm_arch, opt_level, warmup, reps,
                       static_num_ctas, static_occupancy, verify, reset)
     end
@@ -232,16 +190,15 @@ Compile candidate configurations on worker tasks while the caller task measures
 completed candidates. Measurement stays on the caller task because CUDA state is
 task-local; workers only run the untimed temporary compile path.
 """
-function pipelined_tune(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn,
-                        session::TuningSession;
+function pipelined_tune(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn;
                         sm_arch::VersionNumber, opt_level::Int,
                         warmup::Int, reps::Int, workers::Int,
                         static_num_ctas=nothing, static_occupancy=nothing,
                         verify=nothing, reset=nothing)
     isempty(configs) && return TimingRecord[], nothing, nothing
 
     if iszero(workers) || length(configs) == 1
-        record, first_error = measure_candidates(f, configs, grid_fn, args_fn, session;
+        record, first_error = measure_candidates(f, configs, grid_fn, args_fn;
             sm_arch, opt_level, warmup, reps,
             static_num_ctas, static_occupancy, verify, reset)
         return record, nothing, first_error
@@ -265,7 +222,7 @@ function pipelined_tune(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn
                 for cfg in jobs
                     cancelled[] && break
                     try
-                        precompile_cfg(f, cfg, args_fn, session; sm_arch, opt_level,
+                        precompile_cfg(f, cfg, args_fn; sm_arch, opt_level,
                                         static_num_ctas, static_occupancy)
                         cancelled[] || put!(ready, cfg)
                     catch err
@@ -289,7 +246,7 @@ function pipelined_tune(@nospecialize(f), configs::Vector{Any}, grid_fn, args_fn
     first_error = Ref{Any}(nothing)
     try
         for cfg in ready
-            _measure_cfg!(record, first_error, f, cfg, grid_fn, args_fn, session;
+            _measure_cfg!(record, first_error, f, cfg, grid_fn, args_fn;
                           sm_arch, opt_level, warmup, reps,
                           static_num_ctas, static_occupancy, verify, reset)
         end
@@ -348,15 +305,15 @@ function _no_valid_config_error(first_error, precompile_error)
 end
 
 function _refine_record(@nospecialize(f), record::Vector{TimingRecord}, tuning::TuningOptions,
-                        grid_fn, args_fn, session::TuningSession;
+                        grid_fn, args_fn;
                         sm_arch::VersionNumber, opt_level::Int,
                         static_num_ctas=nothing, static_occupancy=nothing,
                         verify=nothing, reset=nothing)
     (tuning.refine_topk > 0 && length(record) > 1) || return record
 
     sort!(record, by=last)
     top = Any[first(r) for r in record[1:min(tuning.refine_topk, length(record))]]
-    refined, _ = measure_candidates(f, top, grid_fn, args_fn, session;
+    refined, _ = measure_candidates(f, top, grid_fn, args_fn;
         sm_arch, opt_level, warmup=tuning.warmup, reps=tuning.refine_reps,
         static_num_ctas, static_occupancy, verify, reset)
     return isempty(refined) ? record : refined
@@ -383,11 +340,10 @@ function find_or_tune(@nospecialize(f), space::AbstractSearchSpace,
 
     checker = verify !== nothing ? verify() : nothing
     reset = setup !== nothing ? setup() : nothing
-    session = TuningSession()
 
     record, precompile_error, first_error =
         with(_SCOPED_INF_CACHE => _fresh_inf_cache()) do
-            pipelined_tune(f, trials, grid_fn, args_fn, session;
+            pipelined_tune(f, trials, grid_fn, args_fn;
                 sm_arch, opt_level,
                 warmup=tuning.warmup, reps=tuning.reps,
                 workers=tuning.precompile_workers,
@@ -397,7 +353,7 @@ function find_or_tune(@nospecialize(f), space::AbstractSearchSpace,
 
     isempty(record) && _no_valid_config_error(first_error, precompile_error)
 
-    record = _refine_record(f, record, tuning, grid_fn, args_fn, session;
+    record = _refine_record(f, record, tuning, grid_fn, args_fn;
         sm_arch, opt_level,
         static_num_ctas, static_occupancy,
         verify=checker, reset)

src/launch.jl

@@ -125,12 +125,13 @@ TileCacheKey(sm_arch::VersionNumber, bytecode_version::VersionNumber,
 
 struct TemporaryTileCacheKey
     key::TileCacheKey
-    session_id::UInt
 end
 
-# Autotune candidates use this owner so they can share work within one tuning
-# pass without becoming visible to the normal `cufunction` cache keyed by
-# `TileCacheKey`. The winning config is promoted by the final normal launch.
+# Autotune candidates compile under this owner so they share inference/codegen
+# work without becoming visible to (or polluting) the normal `cufunction` cache
+# keyed by `TileCacheKey`. It's a plain marker wrapper: candidate CIs are keyed
+# by content, so the cached set is bounded by the search space, not by how many
+# times you tune. The winning config is promoted by the final normal launch.
 @inline tile_cache_key(key::TileCacheKey) = key
 @inline tile_cache_key(key::TemporaryTileCacheKey) = key.key
 
@@ -583,7 +584,7 @@ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{};
                   key, true)::TileKernel{Core.Typeof(f), tt}
 end
 
-function temporary_cufunction(@nospecialize(f), tt::Type{<:Tuple}, session_id::UInt;
+function temporary_cufunction(@nospecialize(f), tt::Type{<:Tuple};
                               sm_arch::Union{VersionNumber, Nothing}=nothing,
                               opt_level::Union{Int, Nothing}=nothing,
                               num_ctas::Union{Int, Nothing}=nothing,
@@ -593,7 +594,7 @@ function temporary_cufunction(@nospecialize(f), tt::Type{<:Tuple}, session_id::U
     resolved_sm_arch = sm_arch !== nothing ? sm_arch : default_sm_arch()
     key = TileCacheKey(resolved_sm_arch, bytecode_version, opt_level, num_ctas, occupancy,
                        num_worker_warps)
-    owner = TemporaryTileCacheKey(key, session_id)
+    owner = TemporaryTileCacheKey(key)
     argtypes, const_argtypes = unwrap_argtypes(f, tt)
     return invoke_frozen(cufunction_compile, f, tt, argtypes, const_argtypes,
                          owner, false)::TileKernel{Core.Typeof(f), tt}