;
; gpu-mi.scm -- GPU-accelerated MI computation using OpenCL
;
; Uses atomese-simd to compute MI for all pairs in parallel on GPU.
;
; MI = log2[count(x,y) * N / (count(x,*) * count(*,y))]
;
; For 1.24M pairs, GPU should be 10-100x faster than CPU.
;
(use-modules (ice-9 format))
(use-modules (srfi srfi-1))
(use-modules (opencog) (opencog exec))
(use-modules (opencog persist))
(use-modules (opencog sensory) (opencog opencl))
(use-modules (opencog matrix))
; ═══════════════════════════════════════════════════════════════════
; CONFIGURATION
; ═══════════════════════════════════════════════════════════════════
(define *kernel-path* "/storage/projects/code/hypercore-kb/opencog-kb/language-learning/mvp/mi-kernel.cl")
(define *mi-key* (Predicate "*-FMI Key-*"))
(define *count-key* (Predicate "*-TruthValueKey-*"))
; ═══════════════════════════════════════════════════════════════════
; GPU SETUP
; ═══════════════════════════════════════════════════════════════════
(define *gpu-node* #f)
(define (init-gpu!)
"Initialize GPU/OpenCL device."
(display "Initializing GPU...\n")
;; Copy kernel to /tmp (required by OpenCL loader)
(copy-file *kernel-path* "/tmp/mi-kernel.cl")
;; Create OpenCL node - empty platform/device = use first available
(set! *gpu-node* (OpenclNode "opencl://:/tmp/mi-kernel.cl"))
;; Open with FloatValue result type
(cog-execute!
(SetValue *gpu-node* (Predicate "*-open-*") (Type 'FloatValue)))
;; Print device info
(let ((desc (cog-value *gpu-node* (Predicate "*-description-*"))))
(format #t "GPU initialized: ~a\n" desc))
*gpu-node*)
(define (close-gpu!)
"Close GPU connection."
(when *gpu-node*
(cog-execute!
(SetValue *gpu-node* (Predicate "*-close-*") (Number 0)))
(set! *gpu-node* #f)))
; ═══════════════════════════════════════════════════════════════════
; DATA EXTRACTION
; ═══════════════════════════════════════════════════════════════════
(define (get-pair-count pair)
"Get count from pair atom."
(let ((val (cog-value pair *count-key*)))
(if val (cog-value-ref val 2) 0.0)))
(define (extract-pair-data pairs llobj)
"Extract counts, left marginals, right marginals from pairs.
Returns (list count-vector left-marg-vector right-marg-vector)."
(define n (length pairs))
(define counts (make-vector n 0.0))
(define left-margs (make-vector n 0.0))
(define right-margs (make-vector n 0.0))
(let ((i 0))
(for-each
(lambda (pair)
(vector-set! counts i (get-pair-count pair))
(vector-set! left-margs i (or (llobj 'left-wild-count pair) 1.0))
(vector-set! right-margs i (or (llobj 'right-wild-count pair) 1.0))
(set! i (+ i 1)))
pairs))
(list counts left-margs right-margs))
; ═══════════════════════════════════════════════════════════════════
; GPU MI COMPUTATION
; ═══════════════════════════════════════════════════════════════════
(define (vector->float-value vec)
"Convert Scheme vector to FloatValue."
(apply FloatValue (vector->list vec)))
(define (compute-mi-gpu! pairs llobj n-total)
"Compute MI for all pairs using GPU.
Returns list of MI values in same order as pairs."
(when (not *gpu-node*)
(init-gpu!))
(define n-pairs (length pairs))
(format #t "Computing MI for ~a pairs on GPU...\n" n-pairs)
(define start-time (current-time))
;; Extract data from AtomSpace
(display " Extracting pair data...\n")
(define data (extract-pair-data pairs llobj))
(define counts (car data))
(define left-margs (cadr data))
(define right-margs (caddr data))
(format #t " Data extracted: ~a pairs\n" n-pairs)
;; Create output buffer
(define result-anchor (Anchor "gpu-mi-result"))
(cog-set-value! result-anchor (Predicate "mi-output")
(OpenclFloatValue (make-list n-pairs 0.0)))
;; Set up input anchors
;; Note: All arrays must be same length for GPU kernel
;; N is replicated for each pair (kernel reads n_value[0] but atomese-simd
;; uses array length to determine work size)
(cog-set-value! (Anchor "gpu-mi-counts") (Predicate "data")
(vector->float-value counts))
(cog-set-value! (Anchor "gpu-mi-left") (Predicate "data")
(vector->float-value left-margs))
(cog-set-value! (Anchor "gpu-mi-right") (Predicate "data")
(vector->float-value right-margs))
(cog-set-value! (Anchor "gpu-mi-n") (Predicate "data")
(apply FloatValue (make-list n-pairs n-total)))
;; Run GPU kernel
(display " Running GPU kernel...\n")
(cog-execute!
(SetValue *gpu-node* (Predicate "*-write-*")
(Section
(Item "compute_mi")
(ConnectorSeq
(ValueOf result-anchor (Predicate "mi-output"))
(ValueOf (Anchor "gpu-mi-counts") (Predicate "data"))
(ValueOf (Anchor "gpu-mi-left") (Predicate "data"))
(ValueOf (Anchor "gpu-mi-right") (Predicate "data"))
(ValueOf (Anchor "gpu-mi-n") (Predicate "data"))))))
;; Read results
(cog-execute! (ValueOf *gpu-node* (Predicate "*-read-*")))
;; Get computed MI values
(define mi-result (cog-execute!
(ValueOf result-anchor (Predicate "mi-output"))))
(define elapsed (- (current-time) start-time))
(format #t " GPU computation done in ~a seconds\n" elapsed)
;; Convert to list
(cog-value->list mi-result))
; ═══════════════════════════════════════════════════════════════════
; BATCH MI COMPUTATION
; ═══════════════════════════════════════════════════════════════════
(define* (batch-mi-gpu! llobj #:key (storage #f) (chunk-size 100000))
"Compute MI for all pairs using GPU, store results on atoms.
llobj: Matrix API object (e.g., from add-pair-stars)
storage: Optional StorageNode to persist results
chunk-size: Process this many pairs at a time (GPU memory limit)
"
(let* ((start-time (current-time))
(_ (display "\n═══ GPU-Accelerated MI Computation ═══\n\n"))
;; Initialize GPU
(_ (init-gpu!))
;; Get all pairs
(_ (display "Fetching all pairs...\n"))
(all-pairs (llobj 'get-all-elts))
(n-pairs (length all-pairs))
(n-total (llobj 'wild-wild-count))
(_ (format #t "Total pairs: ~a\n" n-pairs))
(_ (format #t "Total N: ~a\n" n-total))
;; Process in chunks to avoid GPU memory exhaustion
(n-chunks (ceiling (/ n-pairs chunk-size)))
(_ (format #t "Processing in ~a chunks of ~a\n\n" n-chunks chunk-size)))
;; Process chunks using fold to track state
(let loop ((remaining all-pairs) (processed 0) (chunk-num 0))
(if (null? remaining)
;; Done - cleanup and report
(let* ((_ (close-gpu!))
(elapsed (- (current-time) start-time)))
(format #t "═══ GPU MI Complete ═══\n")
(format #t "Total time: ~,1f seconds (~,2f minutes)\n"
elapsed (/ elapsed 60.0))
(format #t "Throughput: ~,0f pairs/second\n" (/ n-pairs elapsed)))
;; Process next chunk
(let* ((current-chunk-num (+ chunk-num 1))
(chunk-len (min chunk-size (length remaining)))
(chunk (take remaining chunk-len))
(rest (drop remaining chunk-len))
(_ (format #t "Chunk ~a/~a: ~a pairs\n"
current-chunk-num n-chunks chunk-len))
;; Compute MI on GPU
(mi-values (compute-mi-gpu! chunk llobj n-total))
;; Store results on atoms
(_ (display " Storing results...\n"))
(_ (for-each
(lambda (pair mi)
(cog-set-value! pair *mi-key* (FloatValue mi))
(when storage (store-atom pair)))
chunk mi-values))
(new-processed (+ processed chunk-len))
(_ (format #t " Progress: ~a/~a (~,1f%)\n\n"
new-processed n-pairs
(* 100.0 (/ new-processed n-pairs)))))
(loop rest new-processed current-chunk-num))))))
; ═══════════════════════════════════════════════════════════════════
; CONVENIENCE FUNCTIONS
; ═══════════════════════════════════════════════════════════════════
(define (test-gpu-mi)
"Quick test of GPU MI computation."
(init-gpu!)
;; Create test data
(define test-pairs
(list
(Edge (Bond "ANY") (List (Word "test1") (Word "test2")))
(Edge (Bond "ANY") (List (Word "test3") (Word "test4")))))
;; Set some counts
(for-each
(lambda (p)
(cog-set-value! p *count-key* (FloatValue 0 0 10)))
test-pairs)
(display "Test pairs created.\n")
;; Note: This won't work without proper marginals from llobj
;; Just testing GPU connection
(close-gpu!)
(display "GPU test complete.\n"))