Merge remote-tracking branch 'origin/master' into field-offset

.github/workflows/scripts/ti_build/alter.py

@@ -33,11 +33,15 @@ def add_aot_env():
 
 def _write_ti_bashrc():
     path = get_cache_home() / "ti.bashrc"
+    envs = get_cache_home() / "ti-env.sh"
+    _write_env(envs)
     with open(path, "w") as f:
         f.write(
             "[ -f /etc/bashrc ] && source /etc/bashrc\n"
             "[ -f ~/.bashrc ] && source ~/.bashrc\n"
             r'export PS1="\[\e]0;[Taichi Build Environment]\a\]\[\033[01;31m\][Taichi Build] \[\033[00m\]$PS1"'
+            "\n"
+            f"source {envs}\n"
         )
 
     return path
@@ -47,11 +51,15 @@ def _write_ti_zshrc():
     dotdir = get_cache_home() / "zdotdir"
     dotdir.mkdir(parents=True, exist_ok=True)
     path = dotdir / ".zshrc"
+    envs = get_cache_home() / "ti-env.sh"
+    _write_env(envs)
     with open(path, "w") as f:
         f.write(
             "[ -f /etc/zsh/zshrc ] && source /etc/zsh/zshrc\n"
             "[ -f $HOME/.zshrc ] && source $HOME/.zshrc\n"
             r"export PROMPT='%{$fg_bold[red]%}[Taichi Build] %{$reset_color%}'$PROMPT"
+            "\n"
+            f"source {envs}\n"
         )
     return dotdir
 
@@ -138,10 +146,13 @@ def enter_shell():
             os.execl(shell.exe, shell.exe)
 
 
-def write_env(path):
-    cmake_args.writeback()
+def _write_env(path):
     envs = os.environ.get_changed_envs()
     envstr = ""
+
+    if isinstance(path, Path):
+        path = str(path)
+
     if path.endswith(".ps1"):
         envstr = "\n".join([f'$env:{k}="{v}"' for k, v in envs.items()])
     elif path.endswith(".sh"):
@@ -156,6 +167,10 @@ def write_env(path):
     with open(path, "w") as f:
         f.write(envstr)
 
+
+def write_env(path):
+    cmake_args.writeback()
+    _write_env(path)
     misc.info(f"Environment variables written to {path}")
 
 

.github/workflows/scripts/ti_build/bootstrap.py

@@ -2,6 +2,7 @@
 
 # -- stdlib --
 from pathlib import Path
+from types import ModuleType
 from typing import Optional
 import importlib
 import os
@@ -55,21 +56,32 @@ def restart():
         os.execl(sys.executable, sys.executable, "-S", *sys.argv)
 
 
+def _try_import(name: str) -> Optional[ModuleType]:
+    try:
+        return importlib.import_module(name)
+    except ModuleNotFoundError:
+        return None
+
+
 def ensure_dependencies(*deps: str):
     """
     Automatically install dependencies if they are not installed.
     """
 
+    pip = _try_import("pip")
+    ensurepip = _try_import("ensurepip")
+
     if not sys.flags.no_site:
-        # First run, do pip checks
-        try:
-            import pip
-        except ModuleNotFoundError:
-            print("!! pip not found, build.py needs at least a functional pip to work.", flush=True)
-            exit(1)
+        # First run, restart with no_site
+        if not pip and not ensurepip:
+            print(
+                "!! pip or ensurepip not found, build.py needs at least a functional pip/ensurepip to work.", flush=True
+            )
+            sys.exit(1)
 
         restart()
 
+    # Second run
     v = sys.version_info
     bootstrap_root = get_cache_home() / "bootstrap" / f"{v.major}.{v.minor}"
     bootstrap_root.mkdir(parents=True, exist_ok=True)
@@ -85,12 +97,18 @@ def ensure_dependencies(*deps: str):
 
     print("Installing dependencies...", flush=True)
     py = sys.executable
-
-    if run(py, "-m", "pip", "install", "pip", "--no-user", f"--target={bootstrap_root}"):
-        raise Exception("Unable to install pip!")
-
-    pipcmd = [py, "-S", "-m", "pip", "install", "--no-user", f"--target={bootstrap_root}", "-U"]
-    if run(*pipcmd, *deps, env={"PYTHONPATH": str(bootstrap_root)}):
+    pip_install = ["-m", "pip", "install", "--no-user", f"--target={bootstrap_root}", "-U"]
+
+    if ensurepip:
+        wheels = Path(ensurepip.__path__[0]).glob("**/*.whl")
+        wheels = os.pathsep.join(map(str, wheels))
+        if run(py, "-S", *pip_install, "pip", env={"PYTHONPATH": wheels}):
+            raise Exception("Unable to install pip! (ensurepip method)")
+    else:  # pip must exist
+        if run(py, *pip_install, "pip"):
+            raise Exception("Unable to install pip! (pip method)")
+
+    if run(py, "-S", *pip_install, *deps, env={"PYTHONPATH": str(bootstrap_root)}):
         raise Exception("Unable to install dependencies!")
 
     restart()

c_api/docs/taichi/taichi_core.h.md

@@ -305,6 +305,7 @@ Types of kernel and compute graph argument.
 - `enumeration.argument_type.ndarray`: ND-array wrapped around a `handle.memory`.
 - `enumeration.argument_type.texture`: Texture wrapped around a `handle.image`.
 - `enumeration.argument_type.scalar`: Typed scalar.
+- `enumeration.argument_type.tensor`: Typed tensor.
 
 
 `bit_field.memory_usage`
@@ -450,6 +451,23 @@ Scalar value represented by a power-of-two number of bits.
 
 A typed scalar value.
 
+`union.tensor_value`
+
+Tensor value represented by a power-of-two number of bits.
+
+- `union.tensor_value.x8`: Tensor value that fits into 8 bits.
+- `union.tensor_value.x16`: Tensor value that fits into 16 bits.
+- `union.tensor_value.x32`: Tensor value that fits into 32 bits.
+- `union.tensor_value.x64`: Tensor value that fits into 64 bits.
+
+`structure.tensor_value_with_length`
+
+A tensor value with a length.
+
+`structure.tensor`
+
+A typed tensor value.
+
 `union.argument_value`
 
 A scalar or structured argument value.
@@ -459,6 +477,7 @@ A scalar or structured argument value.
 - `union.argument_value.ndarray`: An ND-array to be bound.
 - `union.argument_value.texture`: A texture to be bound.
 - `union.argument_value.scalar`: An scalar to be bound.
+- `union.argument_value.tensor`: A tensor to be bound.
 
 `structure.argument`
 

c_api/include/taichi/cpp/taichi.hpp

@@ -837,6 +837,25 @@ class ComputeGraph {
     return compute_graph_;
   }
 };
+template <typename T>
+struct DataTypeToEnum {
+  static constexpr TiDataType value = TI_DATA_TYPE_UNKNOWN;
+};
+#define DEFINE_DATA_TYPE_ENUM(type, enumv)                    \
+  template <>                                                 \
+  struct DataTypeToEnum<type> {                               \
+    static constexpr TiDataType value = TI_DATA_TYPE_##enumv; \
+  };
+
+DEFINE_DATA_TYPE_ENUM(int32_t, I32);
+DEFINE_DATA_TYPE_ENUM(float, F32);
+DEFINE_DATA_TYPE_ENUM(uint16_t, U16);
+DEFINE_DATA_TYPE_ENUM(int16_t, I16);
+DEFINE_DATA_TYPE_ENUM(uint8_t, U8);
+DEFINE_DATA_TYPE_ENUM(int8_t, I8);
+DEFINE_DATA_TYPE_ENUM(uint64_t, U64);
+DEFINE_DATA_TYPE_ENUM(int64_t, I64);
+#undef DEFINE_DATA_TYPE_ENUM
 
 class Kernel {
  protected:
@@ -884,11 +903,12 @@ class Kernel {
   template <typename T>
   void push_arg(const std::vector<T> &v) {
     int idx = args_.size();
-    // Temporary workaround for setting vec/matrix arguments in a flattened way.
-    args_.resize(args_.size() + v.size());
-    for (int j = 0; j < v.size(); ++j) {
-      at(idx + j) = v[j];
-    }
+    args_.resize(idx + 1);
+    args_[idx].type = TI_ARGUMENT_TYPE_TENSOR;
+    std::memcpy(args_[idx].value.tensor.contents.data.x32, v.data(),
+                v.size() * sizeof(T));
+    args_[idx].value.tensor.contents.length = v.size();
+    args_[idx].value.tensor.type = DataTypeToEnum<T>::value;
   }
 
   template <typename T>

c_api/include/taichi/taichi_core.h

@@ -227,7 +227,7 @@
 #pragma once
 
 #ifndef TI_C_API_VERSION
-#define TI_C_API_VERSION 1005000
+#define TI_C_API_VERSION 1007000
 #endif  // TI_C_API_VERSION
 
 #ifndef TAICHI_H
@@ -463,6 +463,8 @@ typedef enum TiArgumentType {
   TI_ARGUMENT_TYPE_TEXTURE = 3,
   // Typed scalar.
   TI_ARGUMENT_TYPE_SCALAR = 4,
+  // Typed tensor.
+  TI_ARGUMENT_TYPE_TENSOR = 5,
   TI_ARGUMENT_TYPE_MAX_ENUM = 0xffffffff,
 } TiArgumentType;
 
@@ -802,6 +804,36 @@ typedef struct TiScalar {
   TiScalarValue value;
 } TiScalar;
 
+// Union `TiTensorValue`
+//
+// Tensor value represented by a power-of-two number of bits.
+typedef union TiTensorValue {
+  // Tensor value that fits into 8 bits.
+  uint8_t x8[128];
+  // Tensor value that fits into 16 bits.
+  uint16_t x16[64];
+  // Tensor value that fits into 32 bits.
+  uint32_t x32[32];
+  // Tensor value that fits into 64 bits.
+  uint64_t x64[16];
+} TiTensorValue;
+
+// Structure `TiTensorValueWithLength`
+//
+// A tensor value with a length.
+typedef struct TiTensorValueWithLength {
+  uint32_t length;
+  TiTensorValue data;
+} TiTensorValueWithLength;
+
+// Structure `TiTensor`
+//
+// A typed tensor value.
+typedef struct TiTensor {
+  TiDataType type;
+  TiTensorValueWithLength contents;
+} TiTensor;
+
 // Union `TiArgumentValue` (1.4.0)
 //
 // A scalar or structured argument value.
@@ -818,6 +850,8 @@ typedef union TiArgumentValue {
   TiTexture texture;
   // An scalar to be bound.
   TiScalar scalar;
+  // A tensor to be bound.
+  TiTensor tensor;
 } TiArgumentValue;
 
 // Structure `TiArgument` (1.4.0)

c_api/src/taichi_core_impl.cpp

@@ -792,6 +792,28 @@ void ti_launch_kernel(TiRuntime runtime,
         devallocs.emplace_back(std::move(devalloc));
         break;
       }
+      case TI_ARGUMENT_TYPE_TENSOR: {
+        auto &tensor = arg.value.tensor;
+        if (tensor.type == TI_DATA_TYPE_I16 ||
+            tensor.type == TI_DATA_TYPE_U16 ||
+            tensor.type == TI_DATA_TYPE_F16) {
+          for (int j = 0; j < tensor.contents.length; j++) {
+            builder.set_struct_arg_impl({(int)i, j},
+                                        tensor.contents.data.x16[j]);
+          }
+        } else if (tensor.type == TI_DATA_TYPE_I32 ||
+                   tensor.type == TI_DATA_TYPE_U32 ||
+                   tensor.type == TI_DATA_TYPE_F32) {
+          for (int j = 0; j < tensor.contents.length; j++) {
+            builder.set_struct_arg_impl({(int)i, j},
+                                        tensor.contents.data.x32[j]);
+          }
+        } else {
+          ti_set_last_error(TI_ERROR_NOT_SUPPORTED,
+                            ("args[" + std::to_string(i) + "].type").c_str());
+        }
+        break;
+      }
       default: {
         ti_set_last_error(TI_ERROR_ARGUMENT_OUT_OF_RANGE,
                           ("args[" + std::to_string(i) + "].type").c_str());

c_api/src/taichi_llvm_impl.cpp

@@ -52,14 +52,13 @@ taichi::lang::Device &LlvmRuntime::get() {
 
 TiMemory LlvmRuntime::allocate_memory(
     const taichi::lang::Device::AllocParams &params) {
-  const taichi::lang::CompileConfig &config = executor_->get_config();
   taichi::lang::LLVMRuntime *llvm_runtime = executor_->get_llvm_runtime();
   taichi::lang::LlvmDevice *llvm_device = executor_->llvm_device();
 
   taichi::lang::DeviceAllocation devalloc =
-      llvm_device->allocate_memory_runtime(
-          {params, config.ndarray_use_cached_allocator,
-           executor_->get_runtime_jit_module(), llvm_runtime, result_buffer});
+      llvm_device->allocate_memory_runtime({params,
+                                            executor_->get_runtime_jit_module(),
+                                            llvm_runtime, result_buffer});
   return devalloc2devmem(*this, devalloc);
 }