[CODEGEN][METAL] Fix unaligned vector load

This PR fixes the implementation of unaligned vector load.
Previously vector construction was printed as (float2)(v0, v1).
This will cause problem as C have comma expression, and (v0, v1) will be evaluated as v1.
The final result will become float2(v1, v1). The bug affects all codegen that uses
the default implementation, such as metal. We added a testcase on metal to cover this case.

Author: tqchen

src/target/source/codegen_c.cc

@@ -695,7 +695,7 @@ void CodeGenC::VisitExpr_(const BufferLoadNode* op, std::ostream& os) {  // NOLI
         can_vector_load = true;
       }
     }
-
+    LOG(INFO) << "Something wrong";
     if (can_vector_load) {
       std::string ref = GetVecLoad(op->dtype, op->buffer.get(), base.Eval());
       HandleVolatileLoads(ref, op, os);
@@ -722,8 +722,10 @@ void CodeGenC::VisitExpr_(const BufferLoadNode* op, std::ostream& os) {  // NOLI
         value_temp << '[';
         PrintVecElemLoad(sindex, index.dtype(), i, value_temp);
         value_temp << ']';
+        LOG(INFO) << "PrinVecElemLoad" << value_temp.str();
         PrintVecElemLoadExpr(op->dtype, i, value_temp.str(), svalue_expr);
       }
+      LOG(INFO) << svalue_expr.str();
       os << svalue_expr.str();
     }
   }
@@ -802,15 +804,16 @@ void CodeGenC::VisitExpr_(const LetNode* op, std::ostream& os) {  // NOLINT(*)
 }
 
 void CodeGenC::VisitExpr_(const RampNode* op, std::ostream& os) {  // NOLINT(*)
-  // constraint of current logic
-  ICHECK_EQ(op->base.dtype(), DataType::Int(32));
-  os << "((int" << op->lanes << ")(";
+  // NOTE: C have comma expression so cannot use (int2)(v0, v1)
+  // instead should use int2(v0, v1)
+  PrintType(op->dtype, os);
+  os << "(";
   for (int i = 0; i < op->lanes; i++) {
     os << "(" << PrintExpr(op->base) << ")"
        << "+(" << PrintExpr(op->stride) << "*" << i << ")";
     if (i != op->lanes - 1) os << ", ";
   }
-  os << "))";
+  os << ")";
 }
 
 void CodeGenC::VisitExpr_(const ShuffleNode* op, std::ostream& os) {
@@ -999,9 +1002,11 @@ void CodeGenC::PrintVecElemLoadExpr(DataType t, int i, const std::string& value,
   }
 
   if (i == 0) {
-    os << "((";
+    // NOTE: C have comma expression so cannot use (float2)(v0, v1)
+    // instead should use float2(v0, v1)
+    os << "(";
     PrintType(t, os);
-    os << ")(";
+    os << "(";
   }
   os << value;
   if (i != t.lanes() - 1) {

src/target/source/codegen_metal.cc

@@ -299,17 +299,6 @@ void CodeGenMetal::VisitExpr_(const BroadcastNode* op, std::ostream& os) {  // N
   os << ')';
 }
 
-void CodeGenMetal::VisitExpr_(const RampNode* op, std::ostream& os) {  // NOLINT(*)
-  PrintType(op->dtype, os);
-  os << "(";
-  for (int i = 0; i < op->lanes; ++i) {
-    if (i != 0) os << ", ";
-    os << "(" << PrintExpr(op->base) << ")"
-       << "+(" << PrintExpr(op->stride) << "*" << i << ")";
-  }
-  os << ')';
-}
-
 void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) {  // NOLINT(*)
   if (op->op.same_as(builtin::reinterpret())) {
     // generate as_type<TYPE>(ARG)
@@ -369,7 +358,11 @@ runtime::Module BuildMetal(IRModule mod, Target target) {
     code << fsource;
   }
 
-  return MetalModuleCreate(code.str(), fmt, ExtractFuncInfo(mod), source.str());
+  std::string code_str = code.str();
+  if (const auto* f = Registry::Get("tvm_callback_metal_postproc")) {
+    code_str = (*f)(code_str).operator std::string();
+  }
+  return MetalModuleCreate(code_str, fmt, ExtractFuncInfo(mod), source.str());
 }
 
 TVM_REGISTER_GLOBAL("target.build.metal").set_body_typed(BuildMetal);

src/target/source/codegen_metal.h

@@ -51,9 +51,9 @@ class CodeGenMetal final : public CodeGenC {
   void PrintVecElemStore(const std::string& vec, DataType t, int i, const std::string& value) final;
   // overload visitor
   void VisitExpr_(const BroadcastNode* op, std::ostream& os) final;  // NOLINT(*)
-  void VisitExpr_(const RampNode* op, std::ostream& os) final;       // NOLINT(*)
   void VisitExpr_(const CallNode* op, std::ostream& os) final;       // NOLINT(*)
-  void VisitExpr_(const FloatImmNode* op, std::ostream& os) final;
+  void VisitExpr_(const FloatImmNode* op, std::ostream& os) final;   // NOLINT(*)
+
   // reuse parent's function.
   using CodeGenC::PrintType;
 

tests/python/unittest/test_target_codegen_metal.py

@@ -22,6 +22,8 @@
 from tvm.contrib.nvcc import have_fp16, have_int8, have_bf16
 from tvm.contrib import nvcc
 import tvm.testing
+import tvm.script
+from tvm.script import tir as T
 
 tx = te.thread_axis("threadIdx.x")
 bx = te.thread_axis("blockIdx.x")
@@ -54,6 +56,31 @@ def check_inf_nan(dev, n, value, dtype):
     check_inf_nan(dev, 1, float("nan"), "float16")
 
 
+@tvm.testing.requires_gpu
+@tvm.testing.requires_metal
+def test_unaligned_vectorize():
+    @tvm.script.ir_module
+    class IRModule:
+        @T.prim_func
+        def main(A: T.Buffer((2, 3), "float32"), B: T.Buffer((6,), "float32")):
+            T.func_attr({"global_symbol": "main"})
+            for i0_1 in T.thread_binding(3, thread="threadIdx.x"):
+                for i0_0 in T.vectorized(2):
+                    with T.block("block"):
+                        vi0 = T.axis.spatial(6, i0_0 * 3 + i0_1)
+                        B[vi0] = A[vi0 // 3, vi0 % 3]
+
+    target = "metal"
+    dev = tvm.metal()
+
+    a = (np.arange(6).reshape(2, 3)).astype("float32")
+    a_nd = tvm.nd.array(a, dev)
+    b_nd = tvm.nd.empty((6,), "float32", dev)
+    f = tvm.build(IRModule, target=target)
+    f(a_nd, b_nd)
+    np.testing.assert_allclose(b_nd.numpy(), a.reshape(6), atol=1e-5, rtol=1e-5)
+
+
 @tvm.testing.requires_gpu
 @tvm.testing.requires_metal
 def test_metal_erf():