From 3abda29648fe79ec6053fccee12b8a29de3abf1b Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Mon, 21 Dec 2020 18:49:33 -0500
Subject: [PATCH] add debug to ane lib

---
 ane/2_compile/hwx_parse.py |  5 ++++-
 ane/lib/ane.mm             | 42 ++++++++++++++++++++++++++++++++++++++
 ane/lib/ane.py             | 16 ++++++++++++++-
 ane/lib/build.sh           |  2 +-
 4 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/ane/2_compile/hwx_parse.py b/ane/2_compile/hwx_parse.py
index 716eb07c91..30288941d6 100755
--- a/ane/2_compile/hwx_parse.py
+++ b/ane/2_compile/hwx_parse.py
@@ -44,7 +44,8 @@ for k,v in syms.items():
 
 
 # **** document what we know ***
-from ane import ANE_Struct
+from ane import ANE_Struct, ANE
+ane = ANE()
 
 aneb = set()
 for typ, num, nam in ANE_Struct:
@@ -108,6 +109,8 @@ f1 = g.headers[0].commands[1][2][0].section_data
 f2 = a.headers[0].commands[1][2][0].section_data
 for i in range(0, len(f2), 0x300):
   print("===== op %d =====" % (i//0x300))
+  dbg = ane.debug(f2[i:i+0x300])
+  print(dbg)
   if len(f1) < 0x300:
     print(compare(f1, f2[i:i+0x300]))
   else:
diff --git a/ane/lib/ane.mm b/ane/lib/ane.mm
index 2df6814e24..907eef220c 100644
--- a/ane/lib/ane.mm
+++ b/ane/lib/ane.mm
@@ -1,6 +1,7 @@
 #include <stdio.h>
 #include <unistd.h>
 #include <stdlib.h>
+#include <sstream>
 
 #import <IOSurface/IOSurfaceRef.h>
 
@@ -158,5 +159,46 @@ int ANE_Run(uint64_t program_handle, void *in_surf, void *out_surf) {
   return 0;
 }
 
+int ANECCompile(CFDictionaryRef param_1, CFDictionaryRef param_2, unsigned long param_3);
+int ANE_CompilePlist(char *path, bool debug=false) {
+  CFTypeRef ikeys[2];
+  ikeys[0] = CFSTR("NetworkPlistName");
+  ikeys[1] = CFSTR("NetworkPlistPath");
+
+  CFTypeRef ivalues[2];
+  ivalues[0] = CFStringCreateWithCString(kCFAllocatorDefault, path, kCFStringEncodingUTF8);
+  ivalues[1] = CFSTR("./");
+
+  CFDictionaryRef iDictionary = CFDictionaryCreate(kCFAllocatorDefault, ikeys, ivalues, 2, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
+  CFArrayRef array = CFArrayCreate(kCFAllocatorDefault, (const void**)&iDictionary, 1, &kCFTypeArrayCallBacks);
+
+  CFMutableDictionaryRef optionsDictionary = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
+  CFMutableDictionaryRef flagsDictionary = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);
+
+  // h11 (or anything?) works here too, and creates different outputs that don't run
+  CFDictionaryAddValue(flagsDictionary, CFSTR("TargetArchitecture"), CFSTR("h13"));
+  CFDictionaryAddValue(optionsDictionary, CFSTR("OutputFileName"), CFSTR("model.hwx"));
+
+  if (debug) {
+    CFDictionaryAddValue(flagsDictionary, CFSTR("CompileANEProgramForDebugging"), kCFBooleanTrue);
+    int debug_mask = 0x7fffffff;
+    CFDictionaryAddValue(flagsDictionary, CFSTR("DebugMask"), CFNumberCreate(kCFAllocatorDefault, 3, &debug_mask));
+  }
+
+  return ANECCompile(optionsDictionary, flagsDictionary, 0);
+}
+
+void _Z24ZinIrRegBitPrintOutDebugILj7EE11ZinIrStatusjRN11ZinHWTraitsIXT_EE6HwTypeEiRNSt3__113basic_ostreamIcNS5_11char_traitsIcEEEE(
+  unsigned long param_1, void *param_2,int param_3, std::basic_ostream<char> *param_4);
+char *ANE_RegDebug(int a1, void *dat, int a2) {
+  std::ostringstream ss;
+  _Z24ZinIrRegBitPrintOutDebugILj7EE11ZinIrStatusjRN11ZinHWTraitsIXT_EE6HwTypeEiRNSt3__113basic_ostreamIcNS5_11char_traitsIcEEEE(a1, dat, a2, &ss);
+  std::string cppstr = ss.str();
+  const char *str = cppstr.c_str();
+  char *ret = (char *)malloc(strlen(str)+1);
+  strcpy(ret, str);
+  return ret;
+}
+
 }
 
diff --git a/ane/lib/ane.py b/ane/lib/ane.py
index 7136ed1b11..5453628592 100755
--- a/ane/lib/ane.py
+++ b/ane/lib/ane.py
@@ -10,7 +10,7 @@ libane = None
 def init_libane():
   global libane
   libane = cdll.LoadLibrary(os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), 
+    os.path.dirname(os.path.abspath(os.path.realpath(__file__))), 
     "libane.dylib"))
 
   libane.ANE_Compile.argtypes = [c_char_p, c_int]
@@ -24,6 +24,8 @@ def init_libane():
   libane.ANE_Run.argtypes = [c_void_p]*3
   libane.ANE_Run.restype = c_int
 
+  libane.ANE_RegDebug.restype = c_char_p
+
 ANE_Struct = [
 # aneTD.Header
   ("u32", 0x1C, "NextCommandOffset"),
@@ -134,6 +136,18 @@ class ANE:
   def tensor(self, shape):
     return ANETensor(shape)
 
+  def debug(self, dat, mems=0):
+    add = [0x30, 0x1d4, 0x220, 0x29c, 0x2f0, 0x30c, 0x32c]
+    ptr = 0x2b
+    ddat = dat[0:0x28]
+    for a in add:
+      pm = dat[ptr]
+      ddat += b"\x00" * (a-len(ddat))
+      ddat += dat[ptr+1:ptr+1+pm+4]
+      ptr += pm+8
+    ddat += b"\x00" * 8
+    return libane.ANE_RegDebug(0, create_string_buffer(ddat), mems).decode('utf-8')
+
   def filln(self, dat, nvdict, base=0x4000):
     for n,v in nvdict.items():
       styp, num = ANE_Struct_Dict[n]
diff --git a/ane/lib/build.sh b/ane/lib/build.sh
index 29985f80ed..0b652ceaeb 100755
--- a/ane/lib/build.sh
+++ b/ane/lib/build.sh
@@ -1,3 +1,3 @@
 #!/bin/bash
-clang++ ane.mm --shared -F /System/Library/PrivateFrameworks/ -framework ANEServices -framework IOSurface -framework Foundation -framework IOKit -o libane.dylib
+clang++ ane.mm --shared -F /System/Library/PrivateFrameworks/ -framework ANEServices -framework IOSurface -framework Foundation -framework IOKit -framework ANECompiler -o libane.dylib