[TF] TF backend fix and new logic to choose backend (#1393)

* TF backend fix and new logic to choose backend

* fix

* fix

* fix

* fix

* fix backend

* fix

* dlpack alignment

* add flag

* flag

* lint

* lint

* remove unused

* several fixes
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

src/runtime/c_runtime_api.cc

@@ -406,3 +406,4 @@ DGL_REGISTER_GLOBAL("_GetDeviceAttr")
       DeviceAPIManager::Get(ctx)->GetAttr(ctx, kind, ret);
     }
   });
+

src/runtime/ndarray.cc

@@ -350,10 +350,24 @@ int DGLArrayFromDLPack(DLManagedTensor* from,
   API_END();
 }
 
-int DGLArrayToDLPack(DGLArrayHandle from,
-                     DLManagedTensor** out) {
+inline bool is_aligned(const void* ptr, std::uintptr_t alignment) noexcept {
+  auto iptr = reinterpret_cast<std::uintptr_t>(ptr);
+  return !(iptr % alignment);
+}
+
+int DGLArrayToDLPack(DGLArrayHandle from, DLManagedTensor** out,
+                     int alignment) {
   API_BEGIN();
-  *out = NDArray::Internal::ToDLPack(reinterpret_cast<NDArray::Container*>(from));
+  auto* nd_container = reinterpret_cast<NDArray::Container*>(from);
+  DLTensor* nd = &(nd_container->dl_tensor);
+  if (alignment != 0 && !is_aligned(nd->data, alignment)) {
+    std::vector<int64_t> shape_vec(nd->shape, nd->shape + nd->ndim);
+    NDArray copy_ndarray = NDArray::Empty(shape_vec, nd->dtype, nd->ctx);
+    copy_ndarray.CopyFrom(nd);
+    *out = copy_ndarray.ToDLPack();
+  } else {
+    *out = NDArray::Internal::ToDLPack(nd_container);
+  }
   API_END();
 }
 

tests/backend/__init__.py

@@ -6,8 +6,7 @@ import importlib
 import sys
 import numpy as np
 
-mod_name = os.environ.get('DGLBACKEND', 'pytorch').lower()
-mod = importlib.import_module('.%s' % mod_name, __name__)
+mod = importlib.import_module('.%s' % backend_name, __name__)
 thismod = sys.modules[__name__]
 
 for api in backend_unittest.__dict__.keys():
@@ -17,7 +16,6 @@ for api in backend_unittest.__dict__.keys():
         # Tensor APIs used in unit tests MUST be supported across all backends
         globals()[api] = mod.__dict__[api]
 
-
 # Tensor creation with default dtype and context
 
 _zeros = zeros