[Backend][Vitis] Fix Vitis input data and expose target to pytorch frontend

allo/backend/hls.py

@@ -388,17 +388,17 @@ def __call__(self, *args, shell=True):
             # prepare data
             func = find_func_in_module(self.module, self.top_func_name)
             inputs, _ = get_func_inputs_outputs(func)
-            for i, ((in_dtype, in_shape), arg) in enumerate(zip(inputs, args)):
+            for i, ((_, in_shape), arg) in enumerate(zip(inputs, args)):
                 write_tensor_to_file(
                     arg,
-                    in_dtype,
                     in_shape,
-                    f"in_data_{i}",
-                    f"{self.project}/input_{i}.h",
+                    f"{self.project}/input{i}.data",
                 )
             # check if the build folder exists
             bitstream_folder = f"{self.project}/build_dir.{self.mode}.{os.environ['XDEVICE'].rsplit('/')[-1].split('.')[0]}"
-            if not os.path.exists(bitstream_folder):
+            if not os.path.exists(
+                os.path.join(bitstream_folder, f"{self.top_func_name}.xclbin")
+            ):
                 cmd = (
                     f"cd {self.project}; make run TARGET={self.mode} PLATFORM=$XDEVICE"
                 )
@@ -414,7 +414,11 @@ def __call__(self, *args, shell=True):
                 print("Build folder exists, skip building")
                 # run the executable
                 prefix = f"XCL_EMULATION_MODE={self.mode}" if self.mode != "hw" else ""
-                cmd = f"cd {self.project}; make host PLATFORM=$XDEVICE; {prefix} ./{self.top_func_name} ../{bitstream_folder}/{self.top_func_name}.xclbin"
+                prefix += f" cd {self.project};"
+                if not os.path.exists(f"{self.project}/{self.top_func_name}"):
+                    prefix += " make host PLATFORM=$XDEVICE;"
+                cmd = f"{prefix} ./{self.top_func_name} ../{bitstream_folder}/{self.top_func_name}.xclbin"
+                print(cmd)
                 process = subprocess.Popen(cmd, shell=True)
                 process.wait()
                 if process.returncode != 0:

allo/backend/vitis.py

@@ -406,18 +406,13 @@ def update_makefile(file_name, ext_libs):
         outfile.write(makefile)
 
 
-def write_tensor_to_file(tensor, dtype, shape, name, file_path):
-    # generate C buffers
+def write_tensor_to_file(tensor, shape, file_path):
     with open(file_path, "w", encoding="utf-8") as f:
         if len(shape) == 0:
             # scalar
-            f.write(f"const {ctype_map[dtype]} {name} = {tensor};\n")
+            f.write(f"{tensor}\n")
         else:
-            f.write(f"const {ctype_map[dtype]} {name}")
-            # pylint: disable=bad-builtin
-            f.write(f"[{', '.join(map(str, shape))}] = {{")
-            f.write(", ".join([str(i) for i in tensor.flatten()]))
-            f.write("};\n")
+            f.write("\n".join([str(i) for i in tensor.flatten()]))
 
 
 def read_tensor_from_file(dtype, shape, file_path):

allo/frontend/pytorch.py

@@ -27,6 +27,9 @@ def from_pytorch(
     leaf_modules=None,
     verbose=False,
     enable_tensor=False,
+    target="llvm",
+    mode="csim",
+    project="top.prj",
 ):
     sig = inspect.signature(model.forward)
     input_names = [
@@ -64,7 +67,7 @@ def from_pytorch(
     s = customize(
         code, verbose=verbose, global_vars=global_vars, enable_tensor=enable_tensor
     )
-    mod = s.build()
+    mod = s.build(target=target, mode=mode, project=project)
     if verbose:
         print(s.module)
     return mod

examples/torch/toy.py

@@ -28,3 +28,8 @@ def forward(self, x, y):
 np_inputs = [x.detach().numpy() for x in example_inputs]
 res = llvm_mod(*np_inputs)
 torch.testing.assert_close(res, golden.detach().numpy())
+print("Passed!")
+
+# Use other backends
+mod = allo.frontend.from_pytorch(model, example_inputs=example_inputs, target="vhls")
+print(mod.hls_code)