[x86] Separate vector instruction selection and CodeGen passes

Makefile

@@ -481,6 +481,7 @@ SOURCE_FILES = \
   InjectHostDevBufferCopies.cpp \
   Inline.cpp \
   InlineReductions.cpp \
+  InstructionSelector.cpp \
   IntegerDivisionTable.cpp \
   Interval.cpp \
   Introspection.cpp \
@@ -579,7 +580,8 @@ SOURCE_FILES = \
   Var.cpp \
   VectorizeLoops.cpp \
   WasmExecutor.cpp \
-  WrapCalls.cpp
+  WrapCalls.cpp \
+  X86Optimize.cpp
 
 # The externally-visible header files that go into making Halide.h.
 # Don't include anything here that includes llvm headers.
@@ -662,6 +664,7 @@ HEADER_FILES = \
   InjectHostDevBufferCopies.h \
   Inline.h \
   InlineReductions.h \
+  InstructionSelector.h \
   IntegerDivisionTable.h \
   Interval.h \
   Introspection.h \
@@ -745,7 +748,8 @@ HEADER_FILES = \
   Util.h \
   Var.h \
   VectorizeLoops.h \
-  WrapCalls.h
+  WrapCalls.h \
+  X86Optimize.h
 
 OBJECTS = $(SOURCE_FILES:%.cpp=$(BUILD_DIR)/%.o)
 HEADERS = $(HEADER_FILES:%.h=$(SRC_DIR)/%.h)

src/Bounds.cpp

@@ -1111,6 +1111,11 @@ class Bounds : public IRVisitor {
         op->value.accept(this);
     }
 
+    void visit(const VectorInstruction *op) override {
+        // TODO(rootjalex): we may need to implement bounds queries.
+        internal_error << "Unexpected VectorInstruction in bounds query: " << Expr(op) << "\n";
+    }
+
     void visit(const Call *op) override {
         TRACK_BOUNDS_INTERVAL;
         TRACK_BOUNDS_INFO("name:", op->name);

src/CMakeLists.txt

@@ -82,6 +82,7 @@ set(HEADER_FILES
     InjectHostDevBufferCopies.h
     Inline.h
     InlineReductions.h
+    InstructionSelector.h
     IntegerDivisionTable.h
     Interval.h
     Introspection.h
@@ -166,6 +167,7 @@ set(HEADER_FILES
     VectorizeLoops.h
     WasmExecutor.h
     WrapCalls.h
+    X86Optimize.h
     )
 
 set(SOURCE_FILES
@@ -245,6 +247,7 @@ set(SOURCE_FILES
     InjectHostDevBufferCopies.cpp
     Inline.cpp
     InlineReductions.cpp
+    InstructionSelector.cpp
     IntegerDivisionTable.cpp
     Interval.cpp
     Introspection.cpp
@@ -344,6 +347,7 @@ set(SOURCE_FILES
     VectorizeLoops.cpp
     WasmExecutor.cpp
     WrapCalls.cpp
+    X86Optimize.cpp
     )
 
 ##

src/CodeGen_C.cpp

@@ -2829,6 +2829,11 @@ Expr CodeGen_C::scalarize_vector_reduce(const VectorReduce *op) {
     return Shuffle::make_concat(lanes);
 }
 
+void CodeGen_C::visit(const VectorInstruction *op) {
+    internal_error << "CodeGen_C should never receive a VectorInstruction, received:\n"
+                   << Expr(op) << "\n";
+}
+
 void CodeGen_C::visit(const VectorReduce *op) {
     stream << get_indent() << "// Vector reduce: " << op->op << "\n";
 

src/CodeGen_C.h

@@ -235,6 +235,7 @@ class CodeGen_C : public IRPrinter {
     void visit(const Fork *) override;
     void visit(const Acquire *) override;
     void visit(const Atomic *) override;
+    void visit(const VectorInstruction *) override;
     void visit(const VectorReduce *) override;
 
     void visit_binop(Type t, const Expr &a, const Expr &b, const char *op);

src/CodeGen_LLVM.cpp

@@ -4022,11 +4022,16 @@ void CodeGen_LLVM::visit(const Shuffle *op) {
     }
 }
 
+void CodeGen_LLVM::visit(const VectorInstruction *op) {
+    internal_error << "CodeGen_LLVM received VectorInstruction node, should be handled by architecture-specific CodeGen class:\n"
+                   << Expr(op) << "\n";
+}
+
 void CodeGen_LLVM::visit(const VectorReduce *op) {
     codegen_vector_reduce(op, Expr());
 }
 
-void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &init) {
+Expr CodeGen_LLVM::split_vector_reduce(const VectorReduce *op, const Expr &init) const {
     Expr val = op->value;
     const int output_lanes = op->type.lanes();
     const int native_lanes = native_vector_bits() / op->type.bits();
@@ -4066,8 +4071,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
             equiv = max(equiv, init);
         }
         equiv = cast(op->type, equiv);
-        equiv.accept(this);
-        return;
+        return equiv;
     }
 
     if (op->type.is_bool() && op->op == VectorReduce::And) {
@@ -4078,8 +4082,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
         if (init.defined()) {
             equiv = min(equiv, init);
         }
-        equiv.accept(this);
-        return;
+        return equiv;
     }
 
     if (elt == Float(16) && upgrade_type_for_arithmetic(elt) != elt) {
@@ -4089,8 +4092,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
             equiv = binop(equiv, init);
         }
         equiv = cast(op->type, equiv);
-        equiv.accept(this);
-        return;
+        return equiv;
     }
 
     if (output_lanes == 1) {
@@ -4189,8 +4191,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
             if (initial_value.defined()) {
                 equiv = binop(initial_value, equiv);
             }
-            equiv.accept(this);
-            return;
+            return equiv;
         }
     }
 
@@ -4213,8 +4214,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
             equiv = binop(equiv, init);
         }
         equiv = common_subexpression_elimination(equiv);
-        equiv.accept(this);
-        return;
+        return equiv;
     }
 
     if (factor > 2 && ((factor & 1) == 0)) {
@@ -4246,8 +4246,7 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
             equiv = binop(equiv, init);
         }
         equiv = common_subexpression_elimination(equiv);
-        codegen(equiv);
-        return;
+        return equiv;
     }
 
     // Extract each slice and combine
@@ -4261,8 +4260,13 @@ void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &ini
         }
     }
     equiv = common_subexpression_elimination(equiv);
-    codegen(equiv);
-}  // namespace Internal
+    return equiv;
+}
+
+void CodeGen_LLVM::codegen_vector_reduce(const VectorReduce *op, const Expr &init) {
+    Expr equiv = split_vector_reduce(op, init);
+    equiv.accept(this);
+}
 
 void CodeGen_LLVM::visit(const Atomic *op) {
     if (!op->mutex_name.empty()) {

src/CodeGen_LLVM.h

@@ -48,6 +48,8 @@ struct ExternSignature;
 
 namespace Internal {
 
+class InstructionSelector;
+
 /** A code generator abstract base class. Actual code generators
  * (e.g. CodeGen_X86) inherit from this. This class is responsible
  * for taking a Halide Stmt and producing llvm bitcode, machine
@@ -361,6 +363,7 @@ class CodeGen_LLVM : public IRVisitor {
     void visit(const IfThenElse *) override;
     void visit(const Evaluate *) override;
     void visit(const Shuffle *) override;
+    void visit(const VectorInstruction *) override;
     void visit(const VectorReduce *) override;
     void visit(const Prefetch *) override;
     void visit(const Atomic *) override;
@@ -514,6 +517,11 @@ class CodeGen_LLVM : public IRVisitor {
      * across backends. */
     virtual void codegen_vector_reduce(const VectorReduce *op, const Expr &init);
 
+    /** Split up a VectorReduce node if possible, or generate LLVM
+        intrinsics for full reductions. This is used in
+        `codegen_vector_reduce`. **/
+    virtual Expr split_vector_reduce(const VectorReduce *op, const Expr &init) const;
+
     /** Are we inside an atomic node that uses mutex locks?
         This is used for detecting deadlocks from nested atomics & illegal vectorization. */
     bool inside_atomic_mutex_node;
@@ -621,6 +629,12 @@ class CodeGen_LLVM : public IRVisitor {
      * represents a unique struct type created by a closure or similar.
      */
     std::map<llvm::Value *, llvm::Type *> struct_type_recovery;
+
+    /** Instruction selection uses `split_vector_reduce` and
+     *  `upgrade_type_for_arithmetic`, so needs access to those
+     *  methods.
+     */
+    friend class InstructionSelector;
 };
 
 }  // namespace Internal