Normalize impl->fpcore outputs in one place

Author: pavpanchekha

src/syntax/platform.rkt

@@ -35,6 +35,7 @@
          prog->spec
          batch-to-spec!
          get-fpcore-impl
+         impl->fpcore
          (struct-out $platform)
          ;; Platform API
          ;; Operator sets
@@ -176,11 +177,16 @@
          (apply cost-proc (map loop args itypes))]))))
 
 ;; Extracts the `fpcore` field of an operator implementation
-;; as a property dictionary and expression.
+;; as a property dictionary and operation.
 (define (impl->fpcore impl)
-  (match (impl-info impl 'fpcore)
-    [(list '! props ... body) (values (props->dict props) body)]
-    [body (values '() body)]))
+  (define-values (props body)
+    (match (impl-info impl 'fpcore)
+      [(list '! props ... body) (values (props->dict props) body)]
+      [body (values '() body)]))
+  (values props
+          (if (symbol? body)
+              (list body)
+              body)))
 
 (define/reset op-hash #f)
 
@@ -192,12 +198,8 @@
     (define h (make-hash))
     (for ([impl (in-list (platform-impls (*active-platform*)))])
       (define-values (_ expr) (impl->fpcore impl))
-      (define expr*
-        (if (symbol? expr)
-            (list expr)
-            expr))
-      (when (list? expr*)
-        (hash-update! h (car expr*) (curry cons impl) '())))
+      (when (list? expr)
+        (hash-update! h (car expr) (curry cons impl) '())))
     (op-hash h))
 
   ; gather all implementations that have the same spec, input representations,
@@ -207,12 +209,8 @@
           (for ([impl (in-list (hash-ref (op-hash) op '()))]
                 #:when (equal? ireprs (impl-info impl 'itype)))
             (define-values (prop-dict* expr) (impl->fpcore impl))
-            (define expr*
-              (if (symbol? expr)
-                  (list expr)
-                  expr)) ; Handle named constants
             (define pattern (cons op (map (lambda (_) (gensym)) ireprs)))
-            (when (and (subset? prop-dict* prop-dict) (pattern-match pattern expr*))
+            (when (and (subset? prop-dict* prop-dict) (pattern-match pattern expr))
               (sow impl)))))
   ; check that we have any matching impls
   (cond

src/syntax/sugar.rkt

@@ -86,6 +86,18 @@
       [(approx impl spec) (approx (loop impl) (loop spec))]
       [(list op args ...) (cons op (map loop args))])))
 
+(define (fpcore-extension-spec expr)
+  (define specs
+    (remove-duplicates (filter identity
+                               (for/list ([impl (in-list (platform-impls (*active-platform*)))])
+                                 (define-values (_ body) (impl->fpcore impl))
+                                 (define subst (pattern-match (expand-expr body) expr))
+                                 (and subst (pattern-substitute (impl-info impl 'spec) subst))))))
+  (match specs
+    ['() #f]
+    [(list spec) spec]
+    [_ (error 'fpcore->spec "ambiguous platform specs for `~a`: ~a" expr specs)]))
+
 ;; Expression pre-processing for normalizing expressions.
 ;; Used for conversion from FPCore to other IRs.
 (define (expand-expr expr)
@@ -156,7 +168,7 @@
       ; applications
       [`(,op ,args ...) `(,op ,@(map (curryr loop env) args))]
       ; constants
-      [(? operator-exists? op) (list expr)]
+      [(? operator-exists? op) (list op)]
       ; variables
       [(? symbol?) (dict-ref env expr expr)]
       ; other
@@ -176,7 +188,13 @@
       [(? symbol?) expr]
       [(list '! _ ... body) (loop body)]
       [(list 'cast arg) (loop arg)]
-      [(list op args ...) `(,op ,@(map loop args))])))
+      [(list op args ...)
+       (define args* (map loop args))
+       (define expr* `(,op ,@args*))
+       (match (fpcore-extension-spec expr*)
+         [#f expr*]
+         [(== expr*) expr*]
+         [spec (loop spec)])])))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;; FPCore -> LImpl
@@ -195,11 +213,8 @@
   (define ireprs (map (lambda (arg) (repr-of arg ctx)) args))
   (define impl (assert-fpcore-impl op prop-dict ireprs))
   (define vars (impl-info impl 'vars))
-  (define pattern
-    (match (impl-info impl 'fpcore)
-      [(list '! _ ... body) body]
-      [body body]))
-  (define subst (pattern-match pattern (cons op args)))
+  (define-values (_ body) (impl->fpcore impl))
+  (define subst (pattern-match body (cons op args)))
   (pattern-substitute (cons impl vars) subst))
 
 ;; Translates from FPCore to an LImpl.
@@ -235,11 +250,8 @@
        (define idx-repr (repr-of idx* ctx))
        (define impl (assert-fpcore-impl 'ref prop-dict (list arr-repr idx-repr)))
        (define vars (impl-info impl 'vars))
-       (define pattern
-         (match (impl-info impl 'fpcore)
-           [(list '! _ ... body) body]
-           [body body]))
-       (define subst (pattern-match pattern (list 'ref arr* idx*)))
+       (define-values (_ body) (impl->fpcore impl))
+       (define subst (pattern-match body (list 'ref arr* idx*)))
        (pattern-substitute (cons impl vars) subst)]
       [(list op args ...)
        (define args* (map (lambda (arg) (loop arg prop-dict)) args))

src/syntax/test-syntax.rkt

@@ -2,6 +2,7 @@
 
 (require "syntax.rkt"
          "platform.rkt"
+         "sugar.rkt"
          "types.rkt")
 
 (module+ test
@@ -24,4 +25,9 @@
                 '+.f64)
   (check-equal? (get-impl 'sin '((:precision . binary64)) (list f64)) 'sin.f64)
 
+  ; fpcore->spec
+  (check-equal? (fpcore->spec '(log1p x)) '(log (+ 1 x)))
+  (check-equal? (fpcore->spec '(hypot x y)) '(sqrt (+ (* x x) (* y y))))
+  (check-equal? (fpcore->spec '(fma x y z)) '(+ (* x y) z))
+
   (void))