move everything into auto-scheduler

external/backends/metal.lisp

@@ -2,9 +2,8 @@
   (:use :cl :caten/air :cffi :caten/codegen/renderer :caten/codegen/helpers
    :caten/codegen/shape-inference :caten/avm :caten/codegen/expr :cl-metal)
   (:import-from
-   :caten/polyhedral
-   #:define-auto-scheduler
-   #:make-schedule-options))
+   :caten/codegen/config
+   #:define-auto-scheduler))
 
 (in-package :caten/metal)
 
@@ -17,7 +16,6 @@
 
 (define-auto-scheduler
     (Metal-Auto-Scheduler ())
-    :cost-functions '(:validity :proximity :coincidence)
     :n-global-loop 3)
 (define-hook-auto-scheduler (Metal-Renderer Metal-Auto-Scheduler))
 (defmethod initialize-instance :after ((metal Metal-Renderer) &rest initargs &key &allow-other-keys)

source/codegen/auto-scheduler/auto-scheduler.lisp

@@ -2,7 +2,15 @@
   (:use :cl))
 
 (in-package :caten/codegen/auto-scheduler)
-
-(defun auto-schedule ()
+;; Scheduling commands
+;; - [ ] apply-tile
+;; - [ ] apply-parallel
+;; - [ ] apply-collapse
+;; - [ ] apply-fuse
+;; - [ ] apply-interchange
+;; - [ ] apply-vectorize
+;; Purpose: get a list of optimal scheduling commands
+;; Note: 積極的にISL ASTとBlueprintを変換しながら変形を施していく
+(defun auto-schedule (item)
 
   )

source/codegen/auto-scheduler/config.lisp

@@ -0,0 +1,36 @@
+(defpackage :caten/codegen/config
+  (:use :cl)
+  (:export
+   #:Auto-Scheduler-Config
+   #:define-auto-scheduler
+   #:auto-scheduler-n-global-loops
+   #:auto-scheduler-tile-size))
+
+(in-package :caten/codegen/config)
+
+(defclass Auto-Scheduler-Config ()
+  ((n-global-loops :type fixnum :accessor auto-scheduler-n-global-loops)
+   (auto-scheduler-tile-size :type fixnum :accessor auto-scheduler-tile-size))
+  (:documentation ""))
+
+(defmethod print-object ((config Auto-Scheduler-Config) stream)
+  (print-unreadable-object (config stream :type t)
+    (format stream "  N-Global-Loops   | ~a |~%" (slot-value config 'n-global-loops))))
+
+(defmacro define-auto-scheduler ((name (&rest args))
+                                 &key
+                                   (n-global-loop 0)
+                                   (tile-size 0)
+                                   (documentation ""))
+  "define-auto-scheduler"
+  (let ((instance (gensym)))
+    `(progn
+       (defclass ,name (Auto-Scheduler-Config)
+         nil
+         (:documentation ,documentation))
+       (defun ,name (,@args)
+         (let ((,instance (make-instance ',name)))
+           (setf
+            (auto-scheduler-n-global-loops ,instance) ,n-global-loop
+            (auto-scheduler-tile-size ,instance) ,tile-size)
+           ,instance)))))

source/codegen/auto-scheduler/polyhedral.lisp

@@ -10,7 +10,8 @@
    #:poly-schedule
    #:poly-domain
    #:poly-dependencies
-   #:map-schedule-nodes))
+   #:map-schedule-nodes
+   #:->ast))
 
 (in-package :caten/codegen/polyhedral)
 
@@ -166,3 +167,24 @@ This function returns a list of the results of applying f to each node. NIL is e
             (when (= (length next-nodes) 0) (return-from map-search))
             (setf node (pop next-nodes)))
     (nreverse outputs)))
+
+(defun gid (n) (intern (format nil "_gid~a" n)))
+
+(defmethod ->ast ((poly Polyhedral-IR) rank)
+  (macrolet ((set-option (name level)
+	       `(foreign-funcall ,(format nil "isl_options_set_~(~a~)" name)
+				 :pointer (isl::context-handle isl::*context*)
+				 :int ,level
+				 :void)))
+    (set-option "ast_build_exploit_nested_bounds" 1)
+    (set-option "ast_build_detect_min_max" 1)
+    (set-option "ast_build_scale_strides" 1)
+    (set-option "ast_build_allow_else" 0)
+    (set-option "ast_build_allow_or" 0))
+  (let* ((schedule (isl:schedule-set-options (isl:copy (poly-schedule poly)) :separate))
+	 (ast-build (isl:ast-build-from-context (isl:set-from-str "{:}")))
+         (rank (* 2 rank)) ;; rank * tile_bands * vectorizing
+         (ast-build (isl:ast-build-set-iterators ast-build (apply #'isl:make-id-list (loop for i upfrom 0 below rank collect (gid i)))))
+         (ast-build (isl:ast-build-set-options ast-build (isl:union-map-from-str "{}")))
+	 (ast-build-node (isl:ast-build-node-from-schedule ast-build schedule)))
+    ast-build-node))

source/codegen/auto-scheduler/unroll.lisp

@@ -0,0 +1,39 @@
+(defpackage :caten/codegen/unroll
+  (:use :cl :caten/aasm :caten/air :caten/codegen/polyhedral)
+  (:export #:apply-packed-funcall))
+
+(in-package :caten/codegen/unroll)
+
+(defun schedule-apply-schedule-option (si idx)
+  (declare (type Polyhedral-IR si))
+  (let ((bands (map-schedule-nodes #'(lambda (type node) (when (eql type :schedule-node-band) node)) si)))
+    (print bands)
+    ))
+
+(defun apply-packed-funcall (schedule-node gid unroll-by)
+  "Groups the iteration into several packed-funcall.
+Packed-Funcall can be also transformed into Unrolling, or Vectorizing.
+For example, the following code:
+```
+for (int i=0; i<a; i++) {
+  T0(c0);
+}
+```
+is mutated into:
+```
+for (int i=0; i<(a-UNROLL_BY); i+=UNROLL_BY) {
+        [packed_funcall]
+                 { T0(c0+0)
+  T0'(c0, 0~4) = { T0(c0+1)
+                 { T0(c0+2)
+                 { T0(c0+3)
+}
+for (int i=a - (mod a UNROLL_BY); i<a; i+=1) {
+ T0(c0) // Loop Reminder (TODO: Optimize Index Computation)
+}
+```
+"
+  (declare (type node schedule-node) (type integer unroll-by) (type symbol gid))
+  (schedule-apply-schedule-option (getattr schedule-node :polyhedral) nil)
+
+  )

source/codegen/auto-scheduler/vectorize.lisp

@@ -0,0 +1,6 @@
+(defpackage :caten/codegen/vectorize
+  (:use :cl))
+
+(in-package :caten/codegen/vectorize)
+
+(defgeneric %mutate-vectorize (op))

source/codegen/backends/clang.lisp

@@ -2,16 +2,13 @@
   (:use :cl :caten/air :cffi :caten/codegen/renderer :caten/codegen/helpers
    :caten/codegen/shape-inference :caten/avm :caten/codegen/expr)
   (:import-from
-   :caten/polyhedral
-   #:define-auto-scheduler
-   #:make-schedule-options))
+   :caten/codegen/config
+   #:define-auto-scheduler))
 
 (in-package :caten/codegen/backends/clang)
 
-(define-auto-scheduler (Clang-Auto-Scheduler (&key (fuse-softmax 0) (n-global-loop (1- (ctx:getenv :OMP)))))
+(define-auto-scheduler (Clang-Auto-Scheduler (&key (n-global-loop (1- (ctx:getenv :OMP)))))
     ;; Use outermost loop parallelism for maximize memory locality (better softmax/layernorm scheduling)
-    :schedule-option (make-schedule-options :schedule-outer-coincidence fuse-softmax)
-    :cost-functions '(:validity :proximity :coincidence)
     :n-global-loop n-global-loop ;; OMP=1 -> The outermost loop is GLOBAL, otherwise everything is a local loop
     :tile-size 32 ;; [TODO] Automatic Parameter Tuning
     )

source/codegen/scop.lisp

@@ -305,16 +305,17 @@ Reference: https://www.researchgate.net/publication/347152973_PET-to-MLIR_A_poly
     (assert (eql (node-type node) :Schedule-Item))
     (assert (getattr node :blueprint) () "Cannot create a domain w/o lowered blueprint")
     (multiple-value-bind (domain read write schedule) (analyze-scop node (map 'list #'car (Getattr node :dynamic-shapes)))
-      (setf (getattr node :polyhedral) (make-polyhedral-ir domain read write schedule))
+      (setf (getattr node :polyhedral) (caten/codegen/polyhedral:make-polyhedral-ir domain read write schedule))
       (when (>= (ctx:getenv :JIT_DEBUG) 2)
         (format t "~a~%" (getattr node :polyhedral)))
       node)))
 
 (defmethod auto-schedule (auto-scheduler (node Node))
   (assert (getattr node :polyhedral))
-  (caten/polyhedral:auto-schedule auto-scheduler (getattr node :polyhedral))
+  ;;(caten/polyhedral:auto-schedule auto-scheduler (getattr node :polyhedral))
+  ;; (caten/codegen/unroll::apply-packed-funcall node nil 4)
   ;; Load blueprint from optimized polyhedral IR
   (setf (getattr node :blueprint)
         (lower-into-bp-from-polyhedral
-         (caten/polyhedral:->ast (getattr node :polyhedral) (getattr node :rank))
+         (caten/codegen/polyhedral:->ast (getattr node :polyhedral) (getattr node :rank))
          node)))

source/isl/package.lisp

@@ -367,6 +367,7 @@
    #:isl-printer-to-str
    ;; from schedule-node.lisp
    #:schedule-get-root
+   #:schedule-node-band-set-ast-build-options
    #:schedule-node-delete
    #:schedule-node-first-child
    #:schedule-node-insert-mark

source/isl/schedule-node.lisp

@@ -119,3 +119,8 @@
 (define-isl-function schedule-node-first-child %isl-schedule-node-first-child
   (:give schedule-node)
   (:take schedule-node))
+
+(define-isl-function schedule-node-band-set-ast-build-options %isl-schedule-node-band-set-ast-build-options
+  (:give schedule-node)
+  (:take schedule-node)
+  (:take union-set))