[tvm] 10/11: Add explici 64-bit support in metal codegen

[wuwei]

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

wuwei pushed a commit to branch vk-i64
in repository https://gitbox.apache.org/repos/asf/tvm.git

commit 47d45b28b19aba7f6e4eb201b1d9f0bec220f2c5
Author: Wuwei Lin <wu...@apache.org>
AuthorDate: Thu Mar 4 17:01:42 2021 -0500

    Add explici 64-bit support in metal codegen
---
 src/target/source/codegen_metal.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/target/source/codegen_metal.cc 
b/src/target/source/codegen_metal.cc
index f7219cb..c95d578 100644
--- a/src/target/source/codegen_metal.cc
+++ b/src/target/source/codegen_metal.cc
@@ -47,7 +47,7 @@ CodeGenMetal::CodeGenMetal() {
   decl_stream << "#include <metal_stdlib>\n";
   decl_stream << "using namespace metal;\n\n";
   decl_stream << "union __TVMArgUnion {\n"
-              << " int v_long;\n"
+              << " int v_int[2];\n"
               << "};\n\n";
 }
 
@@ -106,6 +106,11 @@ void CodeGenMetal::AddFunction(const PrimFunc& f) {
         PrintType(v.dtype(), decl_stream);
         decl_stream << " " << vid << "[2];\n";
         vref << varg << "." << vid << "[0]";
+      } else if (v.dtype().bits() == 64) {
+        decl_stream << "  ";
+        PrintType(v.dtype(), decl_stream);
+        decl_stream << " " << vid << ";\n";
+        vref << varg << "." << vid;
       } else {
         // For non 32bit type, ref through arg union.
         decl_stream << "  __TVMArgUnion " << vid << ";\n";