Merge pull request #289 from dingobye/concat

Fix a type inconsistency issue in codegen_llvm.

Author: nicolasvasilache

tc/core/polyhedral/codegen_llvm.cc

@@ -243,8 +243,26 @@ class CodeGen_TC : public Halide::Internal::CodeGen_X86 {
       CodeGen_X86::visit(call);
     }
   }
+
   void visit(const Halide::Internal::Variable* op) override {
     value = getValue(iteratorMap_->at(op->name));
+
+    // Generate code for type casting if necessary.
+    llvm::Type* ty = llvm_type_of(op->type);
+    if (value->getType() != ty) {
+      if (op->type.is_int()) {
+        value = builder->CreateIntCast(value, ty, true);
+      } else if (op->type.is_uint()) {
+        value = builder->CreateIntCast(value, ty, false);
+      } else if (op->type.is_float()) {
+        value = builder->CreateFPCast(value, ty);
+      } else {
+        CHECK(false) << "Type inconsistency not handled. "
+                     << "Variable " << op->name << " is " << op->type
+                     << ", but its corresponding llvm::Value is "
+                     << toString(value->getType()) << ".";
+      }
+    }
   }
 
  public:

tc/core/polyhedral/codegen_llvm.h

@@ -17,6 +17,7 @@
 
 #include <memory>
 #include <string>
+#include <type_traits>
 
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
@@ -27,7 +28,12 @@
 
 namespace tc {
 
-static inline std::string toString(llvm::Value* llvmObject) {
+template <
+    typename T,
+    typename std::enable_if<
+        std::is_base_of<llvm::Value, T>::value ||
+        std::is_base_of<llvm::Type, T>::value>::type* = nullptr>
+static inline std::string toString(T* llvmObject) {
   std::string output;
   llvm::raw_string_ostream rso(output);
   llvmObject->print(rso);

test/test_mapper_llvm.cc

@@ -216,6 +216,40 @@ def convolution(float(N,C,H,W) I, float(O,C,KH,KW) W1, float(O) B) -> (tmp, O1)
   checkRtol(output - expected, {I, W1, B}, C * KH * KW, 1e-6);
 }
 
+TEST(LLVMCodegen, Concat) {
+  string tc = R"TC(
+def concat(float(M, N) A, float(M, N) B) -> (O1) {
+    O1(n, i, m) = i == 0 ? A(m, n) : B(m, n) where i in 0:2
+}
+)TC";
+  auto N = 16;
+  auto M = 24;
+
+  auto ctx = isl::with_exceptions::globalIslCtx();
+  auto scop = polyhedral::Scop::makeScop(ctx, tc);
+  auto context = scop->makeContext(
+      std::unordered_map<std::string, int>{{"N", N}, {"M", M}});
+  scop = Scop::makeSpecializedScop(*scop, context);
+
+  Jit jit;
+  jit.codegenScop("concat", *scop);
+  auto fptr = (void (*)(float*, float*, float*))jit.getSymbolAddress("concat");
+
+  at::Tensor A = at::CPU(at::kFloat).rand({M, N});
+  at::Tensor B = at::CPU(at::kFloat).rand({M, N});
+  at::Tensor O1 = at::CPU(at::kFloat).rand({N, 2, M});
+  at::Tensor O1c = at::CPU(at::kFloat).rand({N, 2, M});
+
+  for (int n = 0; n < N; ++n) {
+    for (int m = 0; m < M; ++m) {
+      O1c[n][0][m] = A[m][n];
+      O1c[n][1][m] = B[m][n];
+    }
+  }
+  fptr(A.data<float>(), B.data<float>(), O1.data<float>());
+  checkRtol(O1c - O1, {A, B}, N * M);
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);