[tvm] branch main updated: [CODEGEN][METAL] Fix ramp codegen (#14330)

[junrushao]

This is an automated email from the ASF dual-hosted git repository.

junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new a5ed21d12a [CODEGEN][METAL] Fix ramp codegen (#14330)
a5ed21d12a is described below

commit a5ed21d12aea4a18667915afba6906fedcd85aa5
Author: Tianqi Chen <tqc...@users.noreply.github.com>
AuthorDate: Sun Mar 19 02:58:47 2023 -0400

    [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.
---
 src/target/source/codegen_metal.cc | 11 +++++++++++
 src/target/source/codegen_metal.h  |  1 +
 2 files changed, 12 insertions(+)

diff --git a/src/target/source/codegen_metal.cc 
b/src/target/source/codegen_metal.cc
index 928d961d50..ad9560eef2 100644
--- 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
index 9fb8f80303..4e464c6636 100644
--- 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.