[CODEGEN][METAL] Fix ramp codegen (#14330)

Fix ramp node codegen for the metal backend. The default C codegen can cause problem in vector indices assignment. Confirmed on apple M2.
apache · Mar 19, 2023 · a5ed21d · a5ed21d
1 parent 2ff41c6
commit a5ed21d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/target/source/codegen_metal.cc b/src/target/source/codegen_metal.cc
@@ -299,6 +299,17 @@ 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)

diff --git a/src/target/source/codegen_metal.h b/src/target/source/codegen_metal.h
@@ -51,6 +51,7 @@ 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;
   // reuse parent's function.