diff --git a/src/clorch/amp.clj b/src/clorch/amp.clj index 10130c2..c87e45c 100644 --- a/src/clorch/amp.clj +++ b/src/clorch/amp.clj @@ -4,13 +4,15 @@ (:import [org.bytedeco.pytorch Optimizer Scalar Tensor TensorVector] [org.bytedeco.pytorch.global torch torch$DeviceType torch$ScalarType])) -(defn- device-type [device] +(set! *warn-on-reflection* true) + +(defn- device-type ^torch$DeviceType [device] (case device :cuda torch$DeviceType/CUDA :cpu torch$DeviceType/CPU (throw (ex-info "Autocast device must be :cuda or :cpu" {:device device})))) -(defn- scalar-type [dtype] +(defn- scalar-type ^torch$ScalarType [dtype] (case dtype :float16 torch$ScalarType/Half :bfloat16 torch$ScalarType/BFloat16 diff --git a/src/clorch/autograd.clj b/src/clorch/autograd.clj index badd7af..b0ef9af 100644 --- a/src/clorch/autograd.clj +++ b/src/clorch/autograd.clj @@ -2,6 +2,8 @@ (:require [clorch.torch :as torch]) (:import [org.bytedeco.pytorch NoGradGuard])) +(set! *warn-on-reflection* true) + (defn set-requires-grad [t bool] (.set_requires_grad (torch/->tensor t) (boolean bool)) t) diff --git a/src/clorch/cuda.clj b/src/clorch/cuda.clj index 80357d2..6e0c6bb 100644 --- a/src/clorch/cuda.clj +++ b/src/clorch/cuda.clj @@ -4,6 +4,8 @@ [org.bytedeco.pytorch.global torch] [org.bytedeco.javacpp Loader])) +(set! *warn-on-reflection* true) + ;; Force loading with GPU extension if possible (let [old-ext (System/getProperty "org.bytedeco.javacpp.platform.extension")] (try diff --git a/src/clorch/data.clj b/src/clorch/data.clj index bafe104..19dbac3 100644 --- a/src/clorch/data.clj +++ b/src/clorch/data.clj @@ -6,6 +6,8 @@ [java.util.concurrent Callable Future LinkedBlockingQueue ThreadFactory ThreadPoolExecutor TimeUnit TimeoutException] [java.util.concurrent.atomic AtomicInteger])) +(set! *warn-on-reflection* true) + (def ^:private torch-api (delay {:stack (requiring-resolve 'clorch.torch/stack) @@ -38,8 +40,8 @@ (let [cls (Class/forName "org.bytedeco.pytorch.JavaDataset")] (extend cls IDataset - {:get-size (fn [this] (.get (.size this))) - :get-item (fn [this idx] + {:get-size (fn [^org.bytedeco.pytorch.JavaDataset this] (.get (.size this))) + :get-item (fn [^org.bytedeco.pytorch.JavaDataset this idx] (let [example (.get this (clojure.core/long idx))] {:data (.data example) :target (.target example)}))})))) @@ -90,7 +92,7 @@ (instance? @tensor-class first-item) (let [v (org.bytedeco.pytorch.TensorVector.)] (doseq [item items] (.push_back v item)) - (let [stacked ((torch-fn :stack) v 0)] (.retainReference stacked) stacked)) + (let [stacked ((torch-fn :stack) v 0)] (.retainReference ^org.bytedeco.javacpp.Pointer stacked) stacked)) :else (vec items)))) (defn- build-batch [ds idxs collate-fn] @@ -101,7 +103,7 @@ (partition batch-size idxs) (partition-all batch-size idxs))) -(defn- submit-batch [executor ds idxs collate-fn] +(defn- submit-batch [^java.util.concurrent.ExecutorService executor ds idxs collate-fn] (.submit executor ^Callable (reify Callable @@ -125,7 +127,7 @@ (defn- parallel-batches [ds chunks collate-fn num-workers prefetch-factor] - (let [executor (make-thread-worker-executor num-workers) + (let [^ThreadPoolExecutor executor (make-thread-worker-executor num-workers) max-inflight (clojure.core/max 1 (* (clojure.core/int num-workers) (clojure.core/int prefetch-factor)))] (letfn [(fill-queue [remaining q] (loop [remaining remaining @@ -150,7 +152,7 @@ (.getAbsolutePath f))) (defn- process-start-worker [spec-file worker-id] - (let [proc-builder (ProcessBuilder. ["clojure" "-M" "-m" "clorch.data-worker" "--server" spec-file]) + (let [proc-builder (ProcessBuilder. ^java.util.List (vector "clojure" "-M" "-m" "clorch.data-worker" "--server" spec-file)) _ (.redirectError proc-builder java.lang.ProcessBuilder$Redirect/INHERIT) proc (.start proc-builder) rdr (BufferedReader. (InputStreamReader. (.getInputStream proc))) @@ -280,7 +282,7 @@ select (torch-fn :select) X-t (->tensor X) y-t (->tensor y) - n (.size X-t (clojure.core/long 0))] + n (.size ^org.bytedeco.pytorch.Tensor X-t (clojure.core/long 0))] (dataset :size (fn [] n) :get-item (fn [idx] {:data (select X-t 0 idx) :target (select y-t 0 idx)})))) diff --git a/src/clorch/data_worker.clj b/src/clorch/data_worker.clj index 3ca02cb..af7ce1c 100644 --- a/src/clorch/data_worker.clj +++ b/src/clorch/data_worker.clj @@ -3,6 +3,8 @@ [clojure.java.io :as io] [clorch.data :as data])) +(set! *warn-on-reflection* true) + (defn- load-dataset [spec-file] (let [{:keys [factory args worker-init-fn]} (edn/read-string (slurp spec-file)) ctor (requiring-resolve factory) diff --git a/src/clorch/dev.clj b/src/clorch/dev.clj index a9e989d..4384ea7 100644 --- a/src/clorch/dev.clj +++ b/src/clorch/dev.clj @@ -5,6 +5,8 @@ (:import [java.net BindException SocketException] [java.lang System])) +(set! *warn-on-reflection* true) + (defn- configure-javacpp-platform-extension! [] (let [extension (platform/configure-platform-extension!)] (println "clorch.dev: org.bytedeco.javacpp.platform.extension=" (pr-str extension)))) diff --git a/src/clorch/distributed.clj b/src/clorch/distributed.clj index a22fb43..0b43ada 100644 --- a/src/clorch/distributed.clj +++ b/src/clorch/distributed.clj @@ -12,6 +12,8 @@ [org.bytedeco.javacpp.chrono Milliseconds] [org.bytedeco.pytorch Store StringVector TCPStore TCPStoreOptions Tensor])) +(set! *warn-on-reflection* true) + (platform/configure-platform-extension!) (declare destroy-process-group!) @@ -85,8 +87,8 @@ resolved)) (defn- tcp-store - [{:keys [rank world-size timeout-ms master-address master-port]}] - (let [timeout (Milliseconds. timeout-ms) + ^TCPStore [{:keys [rank world-size timeout-ms master-address master-port]}] + (let [timeout (Milliseconds. (long timeout-ms)) workers (SizeTPointer. (long-array [world-size])) options (doto (TCPStoreOptions.) (.port (unchecked-short master-port)) @@ -96,7 +98,7 @@ (.multiTenant false) (.timeout timeout) (.useLibUV true))] - (TCPStore. master-address options))) + (TCPStore. ^String master-address options))) (defn init-process-group! "Initializes this JVM's default NCCL process group. @@ -313,7 +315,7 @@ (.set store rank-key "ready") (dotimes [peer (:world-size context)] (.push_back store-keys (str prefix peer))) - (._wait store store-keys (Milliseconds. timeout-ms)))) + (._wait store store-keys (Milliseconds. (long timeout-ms))))) (defn barrier! "Blocks until every rank enters the barrier." diff --git a/src/clorch/distributed/checkpoint.clj b/src/clorch/distributed/checkpoint.clj index 63e1b3b..f302559 100644 --- a/src/clorch/distributed/checkpoint.clj +++ b/src/clorch/distributed/checkpoint.clj @@ -10,6 +10,8 @@ [org.bytedeco.pytorch Generator] [org.bytedeco.pytorch.global torch])) +(set! *warn-on-reflection* true) + (def ^:private checkpoint-version 1) (defn- metadata-path [path] diff --git a/src/clorch/distributed/launch.clj b/src/clorch/distributed/launch.clj index cff5685..023b0ef 100644 --- a/src/clorch/distributed/launch.clj +++ b/src/clorch/distributed/launch.clj @@ -10,6 +10,8 @@ [java.util UUID] [java.util.concurrent LinkedBlockingQueue TimeUnit TimeoutException])) +(set! *warn-on-reflection* true) + (declare stop!) (defrecord WorkerProcess [rank ^Process process ^File stdout ^File stderr]) @@ -82,7 +84,7 @@ (Files/createDirectories path (make-array java.nio.file.attribute.FileAttribute 0)) (.toFile path))) -(defn- write-worker-config! [config log-directory] +(defn- write-worker-config! ^File [config log-directory] (let [path (.toPath (io/file log-directory "worker-config.edn")) content (ensure-edn! (select-keys config [:main :args]))] (Files/writeString path content StandardCharsets/UTF_8 diff --git a/src/clorch/distributed/nccl.clj b/src/clorch/distributed/nccl.clj index c7b86af..83abbbd 100644 --- a/src/clorch/distributed/nccl.clj +++ b/src/clorch/distributed/nccl.clj @@ -10,6 +10,8 @@ [org.bytedeco.javacpp BytePointer IntPointer Pointer] [org.bytedeco.pytorch ByteVector Store TensorBase TensorVector])) +(set! *warn-on-reflection* true) + (defrecord NcclWork [^CUevent_st event ^ncclComm communicator result released?]) (defn- nccl-message [status] @@ -69,7 +71,7 @@ (throw (ex-info "Unsupported NCCL reduction operation" {:operation operation})))) -(defn- unique-id-bytes [^ncclUniqueId unique-id] +(defn- unique-id-bytes ^bytes [^ncclUniqueId unique-id] (byte-array (map #(.internal unique-id (int %)) (range nccl/NCCL_UNIQUE_ID_BYTES)))) @@ -166,7 +168,7 @@ (let [async-status (IntPointer. 1)] (check-nccl! "ncclCommGetAsyncError" (nccl/ncclCommGetAsyncError - (:communicator work) async-status)) + ^ncclComm (:communicator work) async-status)) (check-nccl! "NCCL asynchronous operation" (.get async-status 0)) (cudart/cudaEventDestroy (:event work)) (reset! (:released? work) true)) diff --git a/src/clorch/distributed/worker.clj b/src/clorch/distributed/worker.clj index a11759b..9dfca5f 100644 --- a/src/clorch/distributed/worker.clj +++ b/src/clorch/distributed/worker.clj @@ -4,6 +4,8 @@ [clorch.cuda :as cuda] [clorch.distributed :as dist])) +(set! *warn-on-reflection* true) + (defn- ensure-java-25! [] (let [feature (.feature (Runtime/version))] (when (< feature 25) diff --git a/src/clorch/distributions.clj b/src/clorch/distributions.clj index a284a9c..3e9a5b7 100644 --- a/src/clorch/distributions.clj +++ b/src/clorch/distributions.clj @@ -8,8 +8,11 @@ (:import [org.apache.commons.math3.distribution BetaDistribution ChiSquaredDistribution FDistribution GammaDistribution TDistribution WeibullDistribution ParetoDistribution] [org.apache.commons.math3.special Gamma] + [org.bytedeco.pytorch Tensor] [org.bytedeco.pytorch.global torch])) +(set! *warn-on-reflection* true) + (def ^:private eps 1e-8) (def ^:private log-2pi 1.8378770664093453) (def ^:private log-pi 1.1447298858494002) @@ -43,7 +46,7 @@ (defn- maybe-shape [sample-shape] (when (seq sample-shape) (vec sample-shape))) -(defn- as-tensor [x] +(defn- as-tensor ^Tensor [x] (if (number? x) (t/tensor x) (t/->tensor x))) (defn- const-like [value c] diff --git a/src/clorch/einsum.clj b/src/clorch/einsum.clj index 1acd033..58514a5 100644 --- a/src/clorch/einsum.clj +++ b/src/clorch/einsum.clj @@ -3,6 +3,8 @@ (:require [clojure.string :as str] [clorch.torch :as torch])) +(set! *warn-on-reflection* true) + (def ^:private index-symbols "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ") diff --git a/src/clorch/linalg.clj b/src/clorch/linalg.clj index e00d675..ddd6749 100644 --- a/src/clorch/linalg.clj +++ b/src/clorch/linalg.clj @@ -6,6 +6,8 @@ (:refer-clojure :exclude [cond]) (:require [clorch.torch :as t])) +(set! *warn-on-reflection* true) + (defn norm [x] (t/linalg-norm x)) (defn vector-norm [x] (t/linalg-vector-norm x)) (defn matrix-norm [x] (t/linalg-matrix-norm x)) diff --git a/src/clorch/nn.clj b/src/clorch/nn.clj index 097bd67..9f422b5 100644 --- a/src/clorch/nn.clj +++ b/src/clorch/nn.clj @@ -5,7 +5,10 @@ [clorch.nn.functional :as F] [clorch.autograd :as autograd]) (:import [org.bytedeco.pytorch LinearImpl LinearOptions ReLUImpl SigmoidImpl SoftmaxImpl DropoutImpl DropoutOptions FlattenImpl Conv1dImpl Conv1dOptions Conv2dImpl Conv2dOptions Conv3dImpl Conv3dOptions ConvTranspose1dImpl ConvTranspose1dOptions ConvTranspose2dImpl ConvTranspose2dOptions ConvTranspose3dImpl ConvTranspose3dOptions BatchNorm1dImpl BatchNorm2dImpl BatchNorm3dImpl BatchNormOptions GroupNormImpl GroupNormOptions InstanceNorm1dImpl InstanceNorm2dImpl InstanceNorm3dImpl InstanceNormOptions RNNImpl RNNOptions LSTMImpl LSTMOptions GRUImpl GRUOptions EmbeddingImpl EmbeddingOptions Module Tensor TensorVector TanhImpl ELUImpl SELUImpl CELUImpl CELUOptions ELUOptions LeakyReLUImpl LeakyReLUOptions ReLU6Impl SoftplusImpl SoftplusOptions SoftsignImpl SoftminImpl MishImpl HardtanhImpl HardtanhOptions LogSigmoidImpl LogSoftmaxImpl HardshrinkImpl HardshrinkOptions SoftshrinkImpl SoftshrinkOptions TanhshrinkImpl ThresholdImpl ThresholdOptions GLUImpl GLUOptions RReLUImpl RReLUOptions AlphaDropoutImpl FeatureAlphaDropoutImpl BilinearImpl BilinearOptions] - [org.bytedeco.javacpp LongPointer])) + [org.bytedeco.pytorch.global torch$ScalarType] + [org.bytedeco.javacpp LongPointer Pointer])) + +(set! *warn-on-reflection* true) ;; --- Core Protocols --- @@ -62,9 +65,10 @@ (satisfies? IModule thing) (-to thing dtype-or-device) (instance? Tensor thing) (torch/to thing dtype-or-device) (instance? TensorVector thing) - (let [out (TensorVector.)] + (let [^TensorVector thing thing + out (TensorVector.)] (dotimes [i (.size thing)] - (.push_back out (torch/to (.get thing (clojure.core/long i)) dtype-or-device))) + (.push_back out ^Tensor (torch/to (.get thing (clojure.core/long i)) dtype-or-device))) out) (instance? Parameter thing) (assoc thing :tensor (torch/to (:tensor thing) dtype-or-device)) (map? thing) (into (empty thing) (for [[k v] thing] [k (to v dtype-or-device)])) @@ -96,8 +100,8 @@ (swap! params-collector conj node)) (instance? TensorVector node) - (dotimes [i (.size node)] - (let [p (.get node (clojure.core/long i))] + (dotimes [i (.size ^TensorVector node)] + (let [p (.get ^TensorVector node (clojure.core/long i))] (when (or (not trainable-only) (.requires_grad ^Tensor p)) (swap! params-collector conj p)))) @@ -154,7 +158,12 @@ (when (string? k) (get m (keyword k)))) nil))] (cond - (instance? Module model) (do (.load_state_dict ^Module model dict) model) + (instance? Module model) + (let [pairs (.pairs (.named_parameters ^Module model))] + (dotimes [i (.size pairs)] + (when-some [v (lookup dict (.getString (.first pairs (clojure.core/long i))))] + (torch/copy- (.second pairs (clojure.core/long i)) v))) + model) (instance? Parameter model) (do (torch/copy- (:tensor model) dict) model) (instance? Tensor model) (do (torch/copy- model dict) model) (map? model) (do (doseq [[k v] model] (load-state-dict v (lookup dict k))) model) @@ -182,15 +191,51 @@ ;; --- Protocol Implementations --- +(def ^:private ^java.util.concurrent.ConcurrentHashMap forward-methods + (java.util.concurrent.ConcurrentHashMap.)) + +(defn- find-forward + "The single public `forward` of class `c` that accepts `arg-classes`, or nil + (none, or more than one)." + ^java.lang.reflect.Method [^Class c arg-classes] + (let [n (count arg-classes) + ms (filterv (fn [^java.lang.reflect.Method m] + (and (= "forward" (.getName m)) + (= n (.getParameterCount m)) + (every? true? (map (fn [^Class p ^Class a] (.isAssignableFrom p a)) + (.getParameterTypes m) arg-classes)))) + (.getMethods c))] + (when (= 1 (count ms)) (first ms)))) + +(defn- invoke-forward + "Calls a native module's `forward`. JavaCPP's Module declares no `forward`; + each *Impl class declares its own, so it can't be resolved at compile time. + The method is looked up once per (module class, argument classes) and + cached; anything ambiguous falls back to Clojure's reflector." + [module & args] + (let [arg-array (object-array args) + k (into [(class module)] (map class) args) + ^java.lang.reflect.Method m + (or (.get forward-methods k) + (when-not (some nil? args) + (when-let [m (find-forward (class module) (rest k))] + (.put forward-methods k m) + m)))] + (if m + (try (.invoke m module arg-array) + (catch java.lang.reflect.InvocationTargetException e + (throw (.getCause e)))) + (clojure.lang.Reflector/invokeInstanceMethod module "forward" arg-array)))) + (extend-type org.bytedeco.pytorch.Module IModule - (-forward [this x] (.forward this x)) - (-train [this bool] (.train this bool)) + (-forward [this x] (invoke-forward this x)) + (-train [this bool] (.train this (clojure.core/boolean bool))) (-to [this dtype-or-device] (cond (keyword? dtype-or-device) (if-let [stype (get torch/dtype-map dtype-or-device)] - (.to this stype false) + (.to this ^torch$ScalarType stype false) (.to this (org.bytedeco.pytorch.Device. (name dtype-or-device)) false)) (instance? org.bytedeco.pytorch.Device dtype-or-device) @@ -226,7 +271,7 @@ (defn zero-grad "Clears the gradients of all parameters in the model or collection." [thing] - (let [params (parameters thing)] + (let [^TensorVector params (parameters thing)] (dotimes [i (.size params)] (let [p (.get params (clojure.core/long i)) g (.grad p)] @@ -262,7 +307,7 @@ (defn- get-params-count ([m] (get-params-count m false)) ([m trainable-only] - (let [p-vec (parameters m :trainable-only trainable-only) + (let [^TensorVector p-vec (parameters m :trainable-only trainable-only) sz (.size p-vec)] (reduce + (for [i (range sz)] (.numel (.get p-vec (clojure.core/long i)))))))) @@ -288,7 +333,7 @@ params-str (str/join ", " (for [[k s] params] (str k "=" s)))] (str module-label "(" params-str ")")) (map? m) - (str (if (record? m) (.getSimpleName (type m)) "Module") " {\n" + (str (if (record? m) (.getSimpleName ^Class (type m)) "Module") " {\n" (str/join "\n" (for [[k v] m] (str indent-str " " k ": " (module-string v (inc indent))))) "\n" indent-str "}") (vector? m) @@ -299,7 +344,7 @@ :else (str m))))) (defmethod print-method org.bytedeco.pytorch.Module [m ^java.io.Writer w] - (.write w (module-string m))) + (.write w ^String (module-string m))) (defn summary "Prints a summary of the model architecture, inspired by torchsummary. @@ -341,9 +386,9 @@ (instance? Module node) (let [named-p (.named_parameters ^Module node)] (dotimes [i (.size named-p)] - (let [pair (.get (.pairs named-p) (clojure.core/long i)) - param-name (.getString (.first pair)) - tensor (.second pair) + (let [pairs (.pairs named-p) + param-name (.getString (.first pairs (clojure.core/long i))) + tensor (.second pairs (clojure.core/long i)) full-name (if (empty? path) param-name (str path "." param-name))] (swap! res assoc full-name tensor)))) @@ -415,7 +460,7 @@ :else (throw (IllegalArgumentException. "Params must be a number or a sequence")))] (LongPointer. (long-array v)))) -(defn- set-lp! [lp values dims] +(defn- set-lp! [^Pointer lp values dims] (.put lp (->lp values dims)) lp) @@ -443,7 +488,7 @@ (when sparse (.put (.sparse opts) (clojure.core/boolean sparse))) (let [emb (EmbeddingImpl. opts)] (when _weight (torch/copy- (.weight emb) _weight)) - (let [p-vec (parameters emb)] + (let [^TensorVector p-vec (parameters emb)] (dotimes [i (.size p-vec)] (autograd/set-requires-grad (.get p-vec (clojure.core/long i)) (not _freeze)))) (when _freeze (.eval emb)) @@ -531,28 +576,36 @@ (defn instancenorm2d [features] (InstanceNorm2dImpl. (InstanceNormOptions. (clojure.core/long features)))) (defn instancenorm3d [features] (InstanceNorm3dImpl. (InstanceNormOptions. (clojure.core/long features)))) -(defn- setup-rnn-opts [opts hidden-size num-layers bias batch-first dropout bidirectional] - (.put (.hidden_size opts) (clojure.core/long hidden-size)) - (when num-layers (.put (.num_layers opts) (clojure.core/long num-layers))) - (when (some? bias) (.put (.bias opts) (clojure.core/boolean bias))) - (when (some? batch-first) (.put (.batch_first opts) (clojure.core/boolean batch-first))) - (when dropout (.put (.dropout opts) (clojure.core/double dropout))) - (when (some? bidirectional) (.put (.bidirectional opts) (clojure.core/boolean bidirectional))) - opts) +(defmacro ^:private setup-rnn-opts + "Sets the options RNN, LSTM and GRU share. A macro so each call site's + options class (RNNOptions, LSTMOptions, GRUOptions: no common type) + is known at compile time." + [opts-class opts hidden-size num-layers bias batch-first dropout bidirectional] + (let [o (with-meta (gensym "opts") {:tag opts-class})] + `(let [~o ~opts + hidden-size# ~hidden-size num-layers# ~num-layers bias# ~bias + batch-first# ~batch-first dropout# ~dropout bidirectional# ~bidirectional] + (.put (.hidden_size ~o) (clojure.core/long hidden-size#)) + (when num-layers# (.put (.num_layers ~o) (clojure.core/long num-layers#))) + (when (some? bias#) (.put (.bias ~o) (clojure.core/boolean bias#))) + (when (some? batch-first#) (.put (.batch_first ~o) (clojure.core/boolean batch-first#))) + (when dropout# (.put (.dropout ~o) (clojure.core/double dropout#))) + (when (some? bidirectional#) (.put (.bidirectional ~o) (clojure.core/boolean bidirectional#))) + ~o))) (defn rnn [input-size hidden-size & {:keys [num-layers bias batch-first dropout bidirectional]}] (let [opts (RNNOptions. (clojure.core/long input-size) (clojure.core/long hidden-size))] - (setup-rnn-opts opts hidden-size num-layers bias batch-first dropout bidirectional) + (setup-rnn-opts RNNOptions opts hidden-size num-layers bias batch-first dropout bidirectional) (RNNImpl. opts))) (defn lstm [input-size hidden-size & {:keys [num-layers bias batch-first dropout bidirectional]}] (let [opts (LSTMOptions. (clojure.core/long input-size) (clojure.core/long hidden-size))] - (setup-rnn-opts opts hidden-size num-layers bias batch-first dropout bidirectional) + (setup-rnn-opts LSTMOptions opts hidden-size num-layers bias batch-first dropout bidirectional) (LSTMImpl. opts))) (defn gru [input-size hidden-size & {:keys [num-layers bias batch-first dropout bidirectional]}] (let [opts (GRUOptions. (clojure.core/long input-size) (clojure.core/long hidden-size))] - (setup-rnn-opts opts hidden-size num-layers bias batch-first dropout bidirectional) + (setup-rnn-opts GRUOptions opts hidden-size num-layers bias batch-first dropout bidirectional) (GRUImpl. opts))) (defn max-pool1d [kernel & {:keys [stride padding dilation ceil-mode] :or {padding 0 dilation 1 ceil-mode false}}] @@ -605,7 +658,7 @@ (defrecord PairwiseModuleRecord [impl] IModule - (-forward [_ x] (.forward impl (torch/->tensor (first x)) (torch/->tensor (second x)))) + (-forward [_ x] (invoke-forward impl (torch/->tensor (first x)) (torch/->tensor (second x)))) (-train [this _] this) (-to [this _] this)) diff --git a/src/clorch/nn/functional.clj b/src/clorch/nn/functional.clj index e6dd1cf..0be72d8 100644 --- a/src/clorch/nn/functional.clj +++ b/src/clorch/nn/functional.clj @@ -3,6 +3,8 @@ (:import [org.bytedeco.pytorch DoubleOptional TensorOptional] [org.bytedeco.pytorch.global torch])) +(set! *warn-on-reflection* true) + (defn relu [tensor] (torch/relu (t/->tensor tensor))) @@ -211,7 +213,7 @@ (defn adaptive-avg-pool3d [tensor output-size] (torch/adaptive_avg_pool3d (t/->tensor tensor) (long-array (if (number? output-size) [output-size output-size output-size] output-size)))) -(defn- ->optional [t] +(defn- ->optional ^TensorOptional [t] (if t (org.bytedeco.pytorch.TensorOptional. (t/->tensor t)) (org.bytedeco.pytorch.TensorOptional.))) @@ -219,7 +221,7 @@ (defn linear [input weight & [bias]] (torch/linear (t/->tensor input) (t/->tensor weight) (->optional bias))) -(defn- ->lp [params dims] +(defn- ->lp ^longs [params dims] (let [v (cond (number? params) (vec (repeatedly dims #(clojure.core/long params))) (sequential? params) (mapv clojure.core/long params) @@ -268,7 +270,7 @@ (long-array v))) (defn pad [input padding & {:keys [mode value] :or {mode :constant value 0.0}}] - (let [p-arr (if (number? padding) + (let [^longs p-arr (if (number? padding) (->pad-array padding (- (count (t/size input)) 2)) (long-array (mapv clojure.core/long padding)))] (torch/pad (t/->tensor input) p-arr (clojure.core/name mode) (org.bytedeco.pytorch.DoubleOptional. (double value))))) diff --git a/src/clorch/nn/init.clj b/src/clorch/nn/init.clj index 9899749..50a49bf 100644 --- a/src/clorch/nn/init.clj +++ b/src/clorch/nn/init.clj @@ -3,6 +3,8 @@ (:import [org.bytedeco.pytorch.global torch] [org.bytedeco.pytorch NoGradGuard Scalar kFanIn kFanOut kReLU kLeakyReLU kTanh kSigmoid kLinear FanModeType Nonlinearity])) +(set! *warn-on-reflection* true) + (defmacro ^:private no-grad [& body] `(with-open [_# (NoGradGuard.)] ~@body)) diff --git a/src/clorch/nn/parallel.clj b/src/clorch/nn/parallel.clj index 06e2e42..6d1e3ac 100644 --- a/src/clorch/nn/parallel.clj +++ b/src/clorch/nn/parallel.clj @@ -2,16 +2,18 @@ "Distributed neural-network wrappers backed by direct NCCL collectives." (:require [clorch.distributed :as dist] [clorch.nn :as nn]) - (:import [org.bytedeco.pytorch Module Scalar SizeTVector Tensor + (:import [org.bytedeco.pytorch Module Optimizer Scalar SizeTVector Store Tensor TensorBase TensorTensorHook TensorVector] [org.bytedeco.pytorch.global torch])) +(set! *warn-on-reflection* true) + (def ^:dynamic *synchronize-gradients* "False inside no-sync so DDP accumulates local gradients without collectives." true) (defn- parameter-vector [model] - (let [parameters (nn/parameters model :trainable-only true)] + (let [^TensorVector parameters (nn/parameters model :trainable-only true)] (when (zero? (.size parameters)) (throw (ex-info "DistributedDataParallel requires trainable parameters" {}))) parameters)) @@ -60,13 +62,13 @@ prefix (str "clorch/ddp/signature/" operation-id "/") rank-key (str prefix (:rank context)) signature (pr-str (parameter-signature parameters)) - store (:store context) + ^Store store (:store context) store-keys (org.bytedeco.pytorch.StringVector.)] (.set store rank-key signature) (dotimes [peer (:world-size context)] (.push_back store-keys (str prefix peer))) (._wait store store-keys - (org.bytedeco.javacpp.chrono.Milliseconds. (:timeout-ms context))) + (org.bytedeco.javacpp.chrono.Milliseconds. (long (:timeout-ms context)))) (let [signatures (mapv #(.get_to_str store (str prefix %)) (range (:world-size context)))] @@ -135,7 +137,7 @@ (reset! (:remaining bucket) (count (:indices bucket))))) (defn- gradient-hook - [context bucket slot ^Tensor parameter synchronize? pending classloader] + ^TensorTensorHook [context bucket slot ^Tensor parameter synchronize? pending classloader] (proxy [TensorTensorHook] [] (call [gradient-base] ;; Native autograd workers do not inherit the application classloader. @@ -276,7 +278,7 @@ (if scaler ((requiring-resolve 'clorch.amp/step!) scaler optimizer) (do - (.step optimizer) + (.step ^Optimizer optimizer) true)))) (defn reducer-statistics diff --git a/src/clorch/optim.clj b/src/clorch/optim.clj index 8d2c905..a694265 100644 --- a/src/clorch/optim.clj +++ b/src/clorch/optim.clj @@ -1,7 +1,10 @@ (ns clorch.optim - (:import [org.bytedeco.pytorch SGD SGDOptions Adam AdamOptions AdamW AdamWOptions RMSprop RMSpropOptions Adagrad AdagradOptions TensorVector])) + (:import [org.bytedeco.pytorch SGD SGDOptions Adam AdamOptions AdamW AdamWOptions RMSprop RMSpropOptions Adagrad AdagradOptions Optimizer TensorVector] + [org.bytedeco.javacpp DoublePointer])) -(defn- ->tensor-vector [params] +(set! *warn-on-reflection* true) + +(defn- ->tensor-vector ^TensorVector [params] (if (instance? TensorVector params) params (let [v (TensorVector.)] @@ -9,7 +12,7 @@ (.push_back v p)) v))) -(defn- put-betas! [betas-pointer [beta1 beta2]] +(defn- put-betas! [^DoublePointer betas-pointer [beta1 beta2]] (let [first-default (.get betas-pointer 0) values (cond (= first-default 0.9) [beta1 beta2] @@ -99,8 +102,8 @@ (.put (.eps opts) (double eps)) (Adagrad. (->tensor-vector params) opts))) -(defn zero-grad [opt] +(defn zero-grad [^Optimizer opt] (.zero_grad opt)) -(defn step [opt] +(defn step [^Optimizer opt] (.step opt)) diff --git a/src/clorch/parity.clj b/src/clorch/parity.clj index f5759d2..c63083d 100644 --- a/src/clorch/parity.clj +++ b/src/clorch/parity.clj @@ -7,6 +7,8 @@ [clojure.string :as str]) (:import [org.bytedeco.pytorch.global torch])) +(set! *warn-on-reflection* true) + (def ^:private noisy-prefixes #{"Get" "Set" "Log" "Init" "Is" "Has" "DeviceType" "LinalgBackend" "report" "show_" "dispatch_"}) diff --git a/src/clorch/platform.clj b/src/clorch/platform.clj index 7a1b063..210c664 100644 --- a/src/clorch/platform.clj +++ b/src/clorch/platform.clj @@ -3,6 +3,8 @@ [clojure.string :as str]) (:import [org.bytedeco.javacpp Loader])) +(set! *warn-on-reflection* true) + (defn- env [k] (System/getenv k)) diff --git a/src/clorch/release_check.clj b/src/clorch/release_check.clj index 4385905..ed8422a 100644 --- a/src/clorch/release_check.clj +++ b/src/clorch/release_check.clj @@ -3,6 +3,8 @@ [clojure.string :as str] [clojure.test :as test])) +(set! *warn-on-reflection* true) + (def ^:private excluded-test-namespaces '#{clorch.memory-test-suite}) @@ -29,6 +31,35 @@ sort vec))) +(defn discover-src-namespaces [] + (let [src-dir (io/file "src")] + (->> (file-seq src-dir) + (filter #(and (.isFile ^java.io.File %) + (str/ends-with? (.getName ^java.io.File %) ".clj"))) + (map #(file->ns-sym src-dir %)) + (remove #{'clorch.release-check}) + sort + vec))) + +(defn check-reflection! + "Loads every src namespace (each sets *warn-on-reflection*) and fails if + the compiler reports reflection in any of them. Must run before anything + else loads them: warnings are only printed when a namespace compiles." + [] + (println "Checking for reflection warnings...") + (let [err (java.io.StringWriter.)] + (binding [*err* (java.io.PrintWriter. err)] + (doseq [ns-sym (discover-src-namespaces)] + (require ns-sym))) + (let [warnings (->> (str/split-lines (str err)) + (filter #(and (str/starts-with? % "Reflection warning") + (str/includes? % "clorch/"))))] + (when (seq warnings) + (run! println warnings) + (throw (ex-info "Reflection warnings in src." + {:count (count warnings)}))) + (println "No reflection warnings.")))) + (defn- require-test-namespaces! [ns-syms] (doseq [ns-sym ns-syms] (require ns-sym :reload))) @@ -66,5 +97,6 @@ (defn -main [& args] (let [mode (parse-mode args)] (println "Release check mode:" (name mode)) + (check-reflection!) (run-clojure-tests!) (maybe-run-cuda-smoke! mode))) diff --git a/src/clorch/torch.clj b/src/clorch/torch.clj index 0eb4752..3ce93e9 100644 --- a/src/clorch/torch.clj +++ b/src/clorch/torch.clj @@ -22,6 +22,8 @@ [org.bytedeco.pytorch.global torch torch$ScalarType] [org.bytedeco.javacpp PointerScope Pointer DoublePointer])) +(set! *warn-on-reflection* true) + (platform/configure-platform-extension!) ;; --- Memory Management --- @@ -211,20 +213,27 @@ (if opts (let [{:keys [dtype device requires-grad]} opts t-dtype (if dtype - (let [stype (get dtype-map dtype (torch/kFloat))] + (let [^torch$ScalarType stype (get dtype-map dtype (torch/kFloat))] (if (not= (.scalar_type ^Tensor t) stype) (.to ^Tensor t stype false false (MemoryFormatOptional.)) t)) t) t-dev (if device (let [dev (Device. (clojure.core/str (name device)))] - (.to ^Tensor t-dtype dev (.scalar_type t-dtype) false false (MemoryFormatOptional.))) + (.to ^Tensor t-dtype dev (.scalar_type ^Tensor t-dtype) false false (MemoryFormatOptional.))) t-dtype) t-grad (if (some? requires-grad) (.set_requires_grad ^Tensor t-dev (clojure.core/boolean requires-grad)) t-dev)] t-grad) t))) -(defn ->tensor [x] +(defn ->tensor + "Coerces `x` to a Tensor: Tensors (and Modules) pass through, tensor + containers unwrap, anything else goes through `tensor`. + + The ^Tensor return hint lets every `(torch/op (->tensor a) ...)` call + resolve its JavaCPP overload at compile time instead of by reflection + on each call (~30-40x per op on small tensors)." + ^Tensor [x] (cond (instance? org.bytedeco.pytorch.Tensor x) x (instance? org.bytedeco.pytorch.Module x) x @@ -608,7 +617,7 @@ ```" [a b] (if (instance? Tensor b) - (.eq (->tensor a) b) + (.eq (->tensor a) ^Tensor b) (.eq (->tensor a) (Scalar. (clojure.core/double b))))) (defn gt @@ -626,7 +635,7 @@ ```" [a b] (if (instance? Tensor b) - (.gt (->tensor a) b) + (.gt (->tensor a) ^Tensor b) (.gt (->tensor a) (Scalar. (clojure.core/double b))))) (defn lt @@ -644,7 +653,7 @@ ```" [a b] (if (instance? Tensor b) - (.lt (->tensor a) b) + (.lt (->tensor a) ^Tensor b) (.lt (->tensor a) (Scalar. (clojure.core/double b))))) (defn ge @@ -662,7 +671,7 @@ ```" [a b] (if (instance? Tensor b) - (.ge (->tensor a) b) + (.ge (->tensor a) ^Tensor b) (.ge (->tensor a) (Scalar. (clojure.core/double b))))) (defn le @@ -680,7 +689,7 @@ ```" [a b] (if (instance? Tensor b) - (.le (->tensor a) b) + (.le (->tensor a) ^Tensor b) (.le (->tensor a) (Scalar. (clojure.core/double b))))) (defn ne @@ -698,7 +707,7 @@ ```" [a b] (if (instance? Tensor b) - (.ne (->tensor a) b) + (.ne (->tensor a) ^Tensor b) (.ne (->tensor a) (Scalar. (clojure.core/double b))))) (defn logical-and @@ -975,7 +984,7 @@ [t dtype-or-device] (if (instance? org.bytedeco.pytorch.Module t) t - (let [t-obj (if (instance? org.bytedeco.pytorch.Tensor t) t (->tensor t))] + (let [^Tensor t-obj (if (instance? org.bytedeco.pytorch.Tensor t) t (->tensor t))] (cond (keyword? dtype-or-device) (if-let [stype (get dtype-map dtype-or-device)] @@ -1494,7 +1503,7 @@ (defn ->tensor-vector "Converts a Clojure sequence of tensors into a native TensorVector." - [tensors] + ^org.bytedeco.pytorch.TensorVector [tensors] (if (instance? org.bytedeco.pytorch.TensorVector tensors) tensors (let [tv (org.bytedeco.pytorch.TensorVector.)] @@ -1699,7 +1708,7 @@ ([a] (.max (->tensor a))) ([a dim & {:keys [keepdim] :or {keepdim false}}] (let [res (.max (->tensor a) (clojure.core/long dim) (clojure.core/boolean keepdim))] - (.get res (clojure.core/long 0))))) + (.get0 res)))) (defn min "Returns the minimum value of all elements in the input tensor. @@ -1721,7 +1730,7 @@ ([a] (.min (->tensor a))) ([a dim & {:keys [keepdim] :or {keepdim false}}] (let [res (.min (->tensor a) (clojure.core/long dim) (clojure.core/boolean keepdim))] - (.get res (clojure.core/long 0))))) + (.get0 res)))) (defn view "Returns a new tensor with the same data as the self tensor but of a different shape. @@ -1909,7 +1918,7 @@ "Convert negative step slicing to positive step using flip. Returns [transformed-tensor new-indexers]." [t-obj indexers] - (loop [t-obj t-obj + (loop [t-obj ^Tensor t-obj idx indexers dim 0 new-indexers []] @@ -2001,7 +2010,7 @@ (let [->sio (fn [v] (if (nil? v) (SymIntOptional.) (SymIntOptional. (SymInt. (clojure.core/long v)))))] (TensorIndex. (Slice. (->sio start) (->sio end) (->sio step))))) :else (throw (IllegalArgumentException. (str "Invalid indexer: " spec)))))) - (.index t-final idx-vec)))) + (.index ^Tensor t-final idx-vec)))) (defn precompute-rope-freqs [dim seq-len & {:keys [theta] :or {theta 10000.0}}] (let [freqs (div (arange 0 dim 2) (clojure.core/double dim)) @@ -2264,7 +2273,9 @@ (defn clip ([t min-val max-val] - (torch/clip (->tensor t) (clojure.core/double min-val) (clojure.core/double max-val)))) + (torch/clip (->tensor t) + (ScalarOptional. (Scalar. (clojure.core/double min-val))) + (ScalarOptional. (Scalar. (clojure.core/double max-val)))))) (defn count-nonzero "Counts the number of non-zero values in the tensor `t`. @@ -2327,7 +2338,7 @@ ([t] (torch/argsort (->tensor t))) ([t dim] - (torch/argsort (->tensor t) (clojure.core/long dim))) + (torch/argsort (->tensor t) (clojure.core/long dim) false)) ([t dim descending] (torch/argsort (->tensor t) (clojure.core/long dim) (clojure.core/boolean descending))) ([t dim descending stable] @@ -2433,12 +2444,12 @@ (->tensor chol) (clojure.core/boolean upper)))) -(defn- ->long-optional [n] +(defn- ->long-optional ^LongOptional [n] (if (some? n) (LongOptional. (clojure.core/long n)) (LongOptional.))) -(defn- ->string-view-optional [s] +(defn- ->string-view-optional ^StringViewOptional [s] (if (some? s) (StringViewOptional. (clojure.core/str s)) (StringViewOptional.))) @@ -2809,19 +2820,19 @@ ```" [t] (torch/linalg_det (->tensor t))) -(defn- tuple2->vector [pair] +(defn- tuple2->vector [^org.bytedeco.pytorch.T_TensorTensor_T pair] [(.get0 pair) (.get1 pair)]) -(defn- tuple3->vector [triple] +(defn- tuple3->vector [^org.bytedeco.pytorch.T_TensorTensorTensor_T triple] [(.get0 triple) (.get1 triple) (.get2 triple)]) -(defn- tuple4->map [quad] +(defn- tuple4->map [^org.bytedeco.pytorch.T_TensorTensorTensorTensor_T quad] {:solution (.get0 quad) :residuals (.get1 quad) :rank (.get2 quad) :singular-values (.get3 quad)}) -(defn- ->double-optional [x] +(defn- ->double-optional ^DoubleOptional [x] (if (some? x) (DoubleOptional. (clojure.core/double x)) (DoubleOptional.))) @@ -2835,7 +2846,7 @@ (defn- ->driver-optional [driver] (if (some? driver) - (StringViewOptional. (->driver-string driver)) + (StringViewOptional. ^String (->driver-string driver)) (StringViewOptional.))) (defn- order->cond-arg [ord] @@ -3310,11 +3321,11 @@ (->double-optional rtol) (clojure.core/boolean hermitian))) (if (instance? Tensor tol-or-opts) - (torch/linalg_matrix_rank (->tensor a) tol-or-opts) + (torch/linalg_matrix_rank (->tensor a) ^Tensor tol-or-opts) (torch/linalg_matrix_rank (->tensor a) (clojure.core/double tol-or-opts))))) ([a tol {:keys [hermitian] :or {hermitian false}}] (if (instance? Tensor tol) - (torch/linalg_matrix_rank (->tensor a) tol (clojure.core/boolean hermitian)) + (torch/linalg_matrix_rank (->tensor a) ^Tensor tol (clojure.core/boolean hermitian)) (torch/linalg_matrix_rank (->tensor a) (clojure.core/double tol) (clojure.core/boolean hermitian))))) (defn linalg-cond @@ -3333,7 +3344,10 @@ ([a] (torch/linalg_cond (->tensor a))) ([a ord] - (torch/linalg_cond (->tensor a) (order->cond-arg ord)))) + (let [arg (order->cond-arg ord)] + (if (instance? ScalarOptional arg) + (torch/linalg_cond (->tensor a) ^ScalarOptional arg) + (torch/linalg_cond (->tensor a) ^String arg))))) (defn linalg-lstsq "Computes a solution to the least squares problem of a system of linear equations. @@ -3606,12 +3620,12 @@ (defn clip-grad-norm-raw "Thin wrapper around native clip_grad_norm_ for TensorVector." - [params max-norm] + [^org.bytedeco.pytorch.TensorVector params max-norm] (torch/clip_grad_norm_ params (clojure.core/double max-norm) 2.0 false)) ;; --- Utilities --- -(defn- walk-and-serialize [archive prefix thing] +(defn- walk-and-serialize [^org.bytedeco.pytorch.OutputArchive archive prefix thing] (cond (instance? Tensor thing) (.write archive (clojure.core/str prefix) ^Tensor thing) @@ -3639,7 +3653,7 @@ (doseq [[i v] (map-indexed vector thing)] (walk-and-serialize archive (if (empty? prefix) (str i) (clojure.core/str prefix "." i)) v)))) -(defn- walk-and-deserialize [archive prefix thing] +(defn- walk-and-deserialize [^org.bytedeco.pytorch.InputArchive archive prefix thing] (cond (instance? Tensor thing) (do (.read archive (clojure.core/str prefix) ^Tensor thing) thing) @@ -3882,7 +3896,7 @@ nil) (defmethod print-method org.bytedeco.pytorch.Tensor [t ^java.io.Writer w] - (.write w (tensor-string t))) + (.write w ^String (tensor-string t))) (defn top-p "Nucleus (Top-p) sampling. diff --git a/test/clorch/nn_test.clj b/test/clorch/nn_test.clj index 847703b..3354c09 100644 --- a/test/clorch/nn_test.clj +++ b/test/clorch/nn_test.clj @@ -65,3 +65,23 @@ (nn/load-state-dict model {:weight new-weights :bias (torch/zeros [3])}) (is (= 0.0 (torch/item-float (torch/sum (torch/->tensor w))))) (is (= 0.0 (torch/item-float (torch/sum (torch/->tensor b))))))))) + +(deftest load-state-dict-native-module-test + (testing "load-state-dict copies a state-dict into a native module by name" + (torch/with-torch + (let [src (nn/linear 4 3) + dst (nn/linear 4 3) + x (torch/randn [2 4])] + (is (identical? dst (nn/load-state-dict dst (nn/state-dict src)))) + (is (torch/allclose (nn/forward src x) (nn/forward dst x))))))) + +(deftest native-module-forward-test + (testing "forward dispatches to each native module's own forward" + (torch/with-torch + (is (= [2 3] (torch/size (nn/forward (nn/linear 4 3) (torch/randn [2 4]))))) + (is (= [2 5] (torch/size (nn/forward (nn/bilinear 3 4 5) [(torch/randn [2 3]) (torch/randn [2 4])])))) + (is (= [2 7 8] (torch/size (nn/forward (nn/embedding 10 8) (torch/tensor [[1 2 3 4 5 6 7] [0 0 0 0 0 0 0]] {:dtype :int64}))))))) + (testing "native errors propagate unwrapped" + (torch/with-torch + (is (thrown? RuntimeException + (nn/forward (nn/linear 4 3) (torch/randn [2 5]))))))) diff --git a/test/clorch/torch_ops_test.clj b/test/clorch/torch_ops_test.clj index 7eeaf60..32a48a6 100644 --- a/test/clorch/torch_ops_test.clj +++ b/test/clorch/torch_ops_test.clj @@ -118,6 +118,17 @@ (is (= [1.0 4.0] (mapv t/item-float (t/tseq (t/minimum x y))))) (is (= [2.0 5.0] (mapv t/item-float (t/tseq (t/fmax x y))))) (is (= [1.0 4.0] (mapv t/item-float (t/tseq (t/fmin x y))))))) + (testing "max/min along a dim" + (let [m (t/tensor [[1.0 5.0] [7.0 2.0]])] + (is (= [5.0 7.0] (mapv t/item-float (t/tseq (t/max m 1))))) + (is (= [1.0 2.0] (mapv t/item-float (t/tseq (t/min m 1))))) + (is (= [1 2] (t/size (t/max m 0 :keepdim true)))))) + (testing "clip" + (is (= [1.5 2.0 4.0] (mapv t/item-float (t/tseq (t/clip (t/tensor [1.0 2.0 9.0]) 1.5 4.0)))))) + (testing "argsort along a dim" + (let [m (t/tensor [[3.0 1.0 2.0] [0.0 5.0 4.0]])] + (is (= [[1.0 2.0 0.0] [0.0 2.0 1.0]] + (mapv #(mapv t/item-float (t/tseq %)) (t/tseq (t/argsort m 1))))))) (testing "matmul" (let [m1 (t/tensor [[1 2] [3 4]]) m2 (t/tensor [[5 6] [7 8]])