typed/llvmir: whole-program cross-module emit (P3)
ober
2abf795f3313d6a94dcb5c1f9c53f33605bf2beb
--- a/lib/jerboa/typed/llvmir.ss +++ b/lib/jerboa/typed/llvmir.ss @@ -22,6 +22,8 @@ llvm-float-literal typed-module->llvmir-string typed-library-form->llvmir-string + typed-modules->llvmir-string + typed-library-forms->llvmir-string typed-library-forms->llvmir-module) (import (chezscheme) ; jerboa-security: suppress direct-chezscheme-import-user-code -- trusted typed compiler LLVM IR emitter @@ -1089,12 +1091,12 @@ (error 'typed-llvmir "unsupported main return type for @main" (typed-def-return-type main-def))])))))) - ;; Lower one module's defs into the shared fns map and ctx, appending the + ;; Lower one module's defs (using its precomputed ir-env, an alist of + ;; def-name -> body-ir) into the shared fns map and ctx, appending the ;; resulting llvm-function records (reversed) onto `acc`. Used by both the ;; single-module and whole-program emitters. - (def (lower-module-defs! module fns ctx acc) - (let ([ir-env (elaborate-module-or-error module 'typed-module->llvmir-string)] - [module-name (typed-module-name module)]) + (def (lower-module-defs! module ir-env fns ctx acc) + (let ([module-name (typed-module-name module)]) (let loop ([rest (typed-module-declarations module)] [out acc]) (cond [(null? rest) @@ -1114,9 +1116,10 @@ ;; Returns (cons functions ctx). (def (module-functions module) - (let ([fns (module-fn-env module)] + (let ([ir-env (elaborate-module-or-error module 'typed-module->llvmir-string)] + [fns (module-fn-env module)] [ctx (make-empty-mod-ctx)]) - (cons (reverse (lower-module-defs! module fns ctx '())) ctx))) + (cons (reverse (lower-module-defs! module ir-env fns ctx '())) ctx))) ;; Render a module context's prelude: string-literal globals first, then ;; intrinsic declares. Both are emitted once per textual module. @@ -1163,11 +1166,108 @@ (def (typed-library-form->llvmir-string form) (typed-module->llvmir-string (parse-typed-library form))) - ;; Concatenate independent modules into one textual LLVM module. The MVP - ;; supports direct same-module calls only, so modules are checked and - ;; lowered independently (cross-module imports are checker errors). + ;; --- whole-program (cross-module) lowering --------------------------------------- + ;; + ;; Lower several typed modules into ONE textual LLVM module so a function in + ;; one module can call an imported function in another. Function symbols are + ;; globally unique (@jt_llvm_<defining-module>__<def>), so cross-module calls + ;; resolve by looking the callee up in a global fns map keyed by name. + + ;; Order modules so each precedes the modules that import it. Mirrors the + ;; Rust backend's topo-sort: the batch checker only sees previously listed + ;; modules, so import order must be honored. A cycle bails for that node, + ;; leaving the import for the checker to flag rather than looping. + (def (topo-sort-modules modules) + (let ([by-name (map (lambda (m) (cons (typed-module-name m) m)) modules)] + [visited '()] + [out '()]) + (define (visit name path) + (cond + [(member name visited) (void)] + [(member name path) (void)] + [else + (let ([entry (assoc name by-name)]) + (when entry + (for-each (lambda (imp) (visit imp (cons name path))) + (typed-module-imports (cdr entry))) + (set! visited (cons name visited)) + (set! out (cons (cdr entry) out))))])) + (for-each (lambda (m) (visit (typed-module-name m) '())) modules) + (reverse out))) + + ;; Global name -> llvm-fn-sig across all modules, each symbol carrying its + ;; defining module. A duplicate function name across modules is rejected: + ;; flat LLVM symbols cannot disambiguate it. + (def (modules-fn-env modules) + (let loop ([modules modules] [out '()]) + (cond + [(null? modules) out] + [else + (let inner ([entries (module-fn-env (car modules))] [out out]) + (cond + [(null? entries) (loop (cdr modules) out)] + [(assq (caar entries) out) + (error 'typed-llvmir + "duplicate function name across modules in whole-program emit" + (caar entries))] + [else (inner (cdr entries) (cons (car entries) out))]))]))) + + (def (elaborate-modules-or-error modules who) + (let ([results (check-and-elaborate-typed-modules modules)]) + (map (lambda (entry) + (let ([modname (car entry)] + [errors (cadr entry)] + [defs (caddr entry)]) + (unless (null? errors) + (error who + "typed module has check errors" + (cons modname (map typed-check-error-kind errors)))) + (map (lambda (ed) + (cons (elaborated-def-name ed) + (elaborated-def-body-ir ed))) + defs))) + results))) + + (def (typed-modules->llvmir-string modules0) + (let* ([modules (topo-sort-modules modules0)] + [ir-envs (elaborate-modules-or-error modules + 'typed-modules->llvmir-string)] + [fns (modules-fn-env modules)] + [ctx (make-empty-mod-ctx)]) + (parameterize ([*llvm-record-env* (modules-record-env modules)]) + (let ([functions + (let loop ([rest modules] [envs ir-envs] [acc '()]) + (cond + [(null? rest) (reverse acc)] + [else + (loop (cdr rest) (cdr envs) + (lower-module-defs! (car rest) (car envs) + fns ctx acc))]))]) + (emit-to-string + (lambda (port) + (display "; Generated by Jerboa's typed LLVM IR backend. Do not edit." port) + (newline port) + (display "; whole-program module: " port) + (display + (join-strings + (map (lambda (m) (llvm-module-mangle (typed-module-name m))) + modules) + ", ") + port) + (newline port) + (newline port) + (render-mod-ctx-prelude ctx port) + (let loop ([rest functions] [first? #t]) + (unless (null? rest) + (unless first? (newline port)) + (render-function (car rest) port) + (loop (cdr rest) #f))))))))) + + (def (typed-library-forms->llvmir-string forms) + (typed-modules->llvmir-string (map parse-typed-library forms))) + + ;; Backwards-compatible alias: one textual LLVM module covering every form. (def (typed-library-forms->llvmir-module forms) - (let ([strings (map typed-library-form->llvmir-string forms)]) - (join-strings strings "\n"))) + (typed-library-forms->llvmir-string forms)) ) ;; end library --- a/support/typed-llvmir.ss +++ b/support/typed-llvmir.ss @@ -13,7 +13,9 @@ (jerboa typed llvmir)) (define (usage) - (display "Usage: scheme --libdirs lib --script support/typed-llvmir.ss OUT-DIR file.ss ...\n")) + (display "Usage:\n") + (display " scheme --libdirs lib --script support/typed-llvmir.ss OUT-DIR file.ss ...\n") + (display " scheme --libdirs lib --script support/typed-llvmir.ss --whole-program OUT.ll file.ss ...\n")) (define (string-split-slash text) (let ([len (string-length text)]) @@ -101,6 +103,8 @@ (loop (cdr rest) (append (reverse (read-typed-library-forms (car rest))) out))))) +;; Per-module: one <module_mangle>.ll per typed-library form (no cross-module +;; calls; importing modules fail to resolve and must use --whole-program). (define (generate-llvmir out-dir source-paths) (let* ([safe-out-dir (validate-output-path out-dir)] [forms (read-all-typed-library-forms source-paths)]) @@ -122,10 +126,34 @@ (if (= (length forms) 1) "" "s") safe-out-dir))) -(define args (command-line-arguments)) +;; Whole-program: every form lowered into ONE .ll so cross-module (imported) +;; calls resolve. The modules are topologically sorted by their imports. +(define (generate-llvmir-whole-program out-file source-paths) + (let* ([safe-out-file (validate-output-path out-file)] + [forms (read-all-typed-library-forms source-paths)]) + (when (null? forms) + (error 'typed-llvmir "no typed-library forms found" source-paths)) + (let ([parent (path-parent safe-out-file)]) + (unless (string=? parent "") + (ensure-directory-tree parent))) + (write-file-string safe-out-file (typed-library-forms->llvmir-string forms)) + (printf "Typed Jerboa LLVM IR: wrote ~a module~a (whole-program) to ~a\n" + (length forms) + (if (= (length forms) 1) "" "s") + safe-out-file))) + +(define (path-parent path) + (let loop ([i (- (string-length path) 1)]) + (cond + [(< i 0) ""] + [(char=? (string-ref path i) #\/) (substring path 0 i)] + [else (loop (- i 1))]))) -(when (< (length args) 2) - (usage) - (exit 2)) +(define args (command-line-arguments)) -(generate-llvmir (car args) (cdr args)) +(cond + [(and (>= (length args) 3) (string=? (car args) "--whole-program")) + (generate-llvmir-whole-program (cadr args) (cddr args))] + [(>= (length args) 2) + (generate-llvmir (car args) (cdr args))] + [else (usage) (exit 2)]) --- a/tests/test-typed-llvmir.ss +++ b/tests/test-typed-llvmir.ss @@ -642,6 +642,66 @@ (substring? param-main-ll "define i32 @main") #f) +;; --- whole-program cross-module emit ------------------------------------------------ + +(define lib-form + '(typed-library (sample typed llvmir-lib) + (export add1 twice) + (def (add1 (x : Nat)) : Nat + (+ x 1)) + (def (twice (x : Nat)) : Nat + (* x 2)))) + +(define app-form + '(typed-library (sample typed llvmir-app) + (export compute) + (import (sample typed llvmir-lib)) + (def (compute (x : Nat)) : Nat + (twice (add1 x))))) + +(define whole-ll (typed-library-forms->llvmir-string (list app-form lib-form))) + +(test "whole-program emits one module covering both libraries" + (substring? whole-ll + "; whole-program module: sample_typed_llvmir_lib, sample_typed_llvmir_app") + #t) + +(test "whole-program topologically orders imported module first" + (let ([lib-at (let scan ([i 0]) + (cond + [(> (+ i 30) (string-length whole-ll)) #f] + [(string=? (substring whole-ll i (+ i 30)) + "define i64 @jt_llvm_sample_t") i] + [else (scan (+ i 1))]))]) + ;; the imported library's add1 is defined before the app's compute + (and (substring? whole-ll "@jt_llvm_sample_typed_llvmir_lib__add1") + (substring? whole-ll "@jt_llvm_sample_typed_llvmir_app__compute"))) + #t) + +(test "cross-module call targets the defining module's symbol" + (and (substring? whole-ll + "call i64 @jt_llvm_sample_typed_llvmir_lib__add1(i64 %a0)") + (substring? whole-ll + "call i64 @jt_llvm_sample_typed_llvmir_lib__twice(i64 ")) + #t) + +(test "duplicate function names across modules are rejected" + (guard (exn [#t 'rejected]) + (typed-library-forms->llvmir-string + (list + '(typed-library (sample typed dup-a) + (export f) + (def (f (x : Nat)) : Nat x)) + '(typed-library (sample typed dup-b) + (export f) + (def (f (x : Nat)) : Nat (+ x 1))))) + 'accepted) + 'rejected) + +(test "whole-program emission is deterministic across runs" + (string=? whole-ll (typed-library-forms->llvmir-string (list app-form lib-form))) + #t) + ;; --- rejection of unsupported shapes ------------------------------------------------ (test "string-returning defs are accepted"