Open191

Babylon.js用のSPZローダを作りたい

にー兄さんにー兄さん

先日、NianticがGaussianSplattingの新しい独自フォーマットであるSPZを発表した

https://scaniverse.com/news/spz-gaussian-splat-open-source-file-format

JPGのように~って感じの謳い文句だけど、JPGエンコーディングほど難しいことはしていなくて
単純に精度による圧縮を書けているらしい(それでも十分な圧縮率になる)

にー兄さんにー兄さん

精度による圧縮は、別にSPZに限らず前からずっとやられていて
例えば.splatでも情報によってはかなり精度を削っているし、
LumaAIがLumaAPIから返すGSデータも、数値のMin,Maxをメタデータに保持しておいて正規化するようなこともやっている

SPZも中身は結構シンプルな気はしているなぁ
sizeが対数関数的に圧縮されているのはなるほどってなった

にー兄さんにー兄さん

spzのここ好きポイントはローダがOSSになっていること!
C++で書かれているけど純粋なロジックだけ抽出してあって、
make地獄にならないところ!!たすかる~

https://github.com/nianticlabs/spz

にー兄さんにー兄さん

libzへの依存だけあるようなので、ここは何とかして解決しなくてはいけない

にー兄さんにー兄さん

今回の目的

今回目指す形は、このSPZフォーマットのGSデータをBabylon.jsにロードすること

サブミッションとしては、そういうライブラリを公開すること

にー兄さんにー兄さん

そもそも、SPZはNianticStudioや8thwallで使われているので
100%同じようなことをNianticでもやっている

しかしそこの部分のコードは公開されていないというだけで、
不可能なことはないはず

にー兄さんにー兄さん

色々作戦はあるけど、
ライブラリとして公開することも視野に入れるのであればある程度メンテナンス性も考慮したい

つまり、まだ発表されて間もないSPZのエンコードデコードをTypeScriptで自前で書くようなことは避けたい(普通に自分にできる気がしない)
パフォーマンス面で課題が出てきそうだし、元のソースをキャッチアップしなくちゃいけないし

にー兄さんにー兄さん

今考えている構成は下記

意図としては

  • SPZのコードを参照しておくことで、SPZに変更があってもSPZをアプデするだけで対応できるようにしたい
    • 例えばgit submodule的にリポジトリに含めるなどする
    • libzの依存解決はどうしようか
  • SPZのcppコードをWASMに変換すればよくない?ってなりそうだけどRustをかましている
    • Rustのwasm-bindgenがいい感じな気がしたため
    • この記事ではwasm-bindgen(wasm-pack)が使えなさそうとも書いてある
    • Emscriptenなんかこわいので別の方法があればそうしたい
にー兄さんにー兄さん

まずは適当に自分で作ったC++のコードを
テキトーに作ったRustコードから呼び出すところをやっていく

にー兄さんにー兄さん

まず困ったのがコイツだった

#include <bzlib.h>

自分の開発環境はWinだけど、このヘッダがそもそもない
なのでこのチュートリアル通りにはできないなぁって感じだった

にー兄さんにー兄さん

フォルダ構成とbuild.rs

build.rs
extern crate bindgen;
extern crate cc;

use std::env;
use std::path::PathBuf;

fn main() {
    //println!("cargo:rustc-link-lib=c++");
    println!("cargo:rerun-if-changed=hello/hello");

    cc::Build::new()
        .file("hello/hello.cpp")
        //.flag("-std=c++17")
        .include("hello")
        .compile("hello");

    let bindings = bindgen::Builder::default()
        .header("hello/hello.hpp")
        .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
        .generate()
        .expect("Unable to generate bindings");

    let out_path = PathBuf::from(env::var("OUT_DIR").unwrap());
    bindings
        .write_to_file(out_path.join("bindings.rs"))
        .expect("Couldn't write bindings!");
}
lib.rs
#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

include!(concat!(env!("OUT_DIR"), "/bindings.rs"));

この状態で

cargo build --lib # libオプションはいらないかも

すると、bindings.rsが生成され、
それをmain.rsから利用できた

にー兄さんにー兄さん

この状態で、wasm-bindgenを導入し、wasm-packでビルドできるか試してみた
結果、Winでは無理そうだった

wasmなのでtargetがwasm32-unknown-unknownとしてビルドされるんだけど
この時にccクレートではccコマンドを探しに行っちゃうみたい
Winでは自動でclでコンパイルするようだったんだけど、ターゲットプラットフォームがあれなので、そうなってしまうみたいだった

これはccのオプションにターゲットを指定した時と同じ挙動だった

にー兄さんにー兄さん

次のムーブとして、
Cmakeを含むC++のプロジェクトをRUst Bindgenで使ってみる(wasm化はしない)
をやってみる

最終的にlibzはvcpkgからインポートしたいので
cmakeでプロジェクトを整えたうえでzlibのインポートをしてみる

にー兄さんにー兄さん

思ったけど、Rustを噛まさない方が簡単なのではないかという気持ちになってきた
いまはEmscriptenなしでclang単体でもwasm出力はできるらしいので、其れでやってみようかなぁ

にー兄さんにー兄さん

WinのPowershellではコマンド通り実行してもなんかエラーになっちゃった

にー兄さんにー兄さん

色々試した結果、下記のように実行すればよいらしい
これはPowershellが悪いな

clang++ -std=c++14 -o build/hello.wasm --target=wasm32 -nostdlib "-Wl,--no-entry" "-Wl,--export-all" hello/hello.cpp
にー兄さんにー兄さん

cppファイルの中身はこんな感じ
attributeかextern"C"をすると、wasmをNode.jsで呼び出したときに関数名が保たれる

hello/hello.cpp
#include "hello.hpp"

__attribute__((export_name("hello")))
int hello()
{
  return 42;
}

// or

extern "C"
{
  int hello()
  {
    return 42;
  }
}
hello/hello.hpp
extern "C"
{
  int hello();
}
にー兄さんにー兄さん

build/hello.wasmが出力されるので、それをNode.jsで実行するにはこうなる

main.mjs
import fs from "node:fs/promises";

const main = async () => {
  const wasmCode = await fs.readFile("./build/hello.wasm");

  const { instance } = await WebAssembly.instantiate(wasmCode, {});
  const helloFunc = instance.exports.hello;
  const result = helloFunc();
  console.log(result);
};

main();
node ./main.mjs
にー兄さんにー兄さん

いろいろ試したものの、ここからclang単体でspzとzlibを使ったC++のネイティブアプリをビルドすることができなかった・・・

spzはちょっと変更を加えてまぁまだしも、やはりzlibをいい感じに食わせる方法が全くわからないといった感じ

にー兄さんにー兄さん

ここにきて、もうclangで消耗するくらいなら、
いっそのことcmakeでEmscriptenを使ってwasmを出力するようにしたほうがいいのでは?tって考えになった。いや確かに、今回の目的はspzをwasmで使うことなので

にー兄さんにー兄さん

Rustを使いたい要因にwasm-bindgenというかwasm-packがあったけど
それでやってくれるようなTS型定義ファイル生成をやってくれるツールが
Emscriptenの場合でも使えるらしい
https://github.com/ted537/tsembind

にー兄さんにー兄さん

最終的に、これは使わないかな…ってなったけど
Node.jsからwasmビルドできるツールチェーンが公開されていて面白かった

この記事ではまさに、zlibをwasm化しているし、サンプルプロジェクトもある
https://qiita.com/ukyo/items/ebc30ebf365718e5d16d
https://github.com/ukyo/zlib-wasm-without-emscripten-sample

そしてそのツールはこちら
https://github.com/dcodeIO/webassembly

ただ、記事も5年以上経過、ライブラリも7年前更新なので、厳しそう

にー兄さんにー兄さん

とりま、web.devの内容を、公式dockerイメージでやってみたがいい感じに動かせた
https://web.dev/articles/emscripting-a-c-library?hl=ja
https://hub.docker.com/r/emscripten/emsdk

こんな感じのbatファイルで実行

run.bat
docker run^
 --rm^
 --privileged^
 -v /path/to/dir:/src^
 emscripten/emsdk^
 em++ main.cpp --std=c++17 -o ./build/main.js -s USE_ZLIB=1

priviledgeは、つけるとメモリのエラーを回避できそう(docker toolboxだからかな)
https://github.com/nodejs/help/issues/1754
https://qiita.com/muddydixon/items/d2982ab0846002bf3ea8

にー兄さんにー兄さん

最終的に、spzのspz::normalizeはwasm化して実行できた
しかしなぜかGaussianCloidとかが実行できな

にー兄さんにー兄さん

せっかくdockerでやっているので、devcontainerも整えた
と言っても、.devcontainer/devcontainer.jsonに以下を追記しただけ

.devcontainer/devcontainer.json
{
  "image": "emscripten/emsdk",
  "customizations": {
    "vscode": {
      "extensions": [
        "ms-vscode.cpptools-extension-pack"
      ]
    }
  }
}

なんならimageしか自分で加筆していない。これだけでいいんだたらもっと早くやっておけばよかったな

にー兄さんにー兄さん

色々いじってみた結果、やっぱり出力したwasmを自分でインポートして使うと
なんかエラーになっちゃうけど、

出力されたjsをnode.jsで実行するのであれば、GaussianCloudの圧縮などが普通に動いた、すごい

にー兄さんにー兄さん

ESModuleで出力したら動いた!!!
このコードが

#include <emscripten.h>

#include "spz/src/cc/splat-types.h"
#include "spz/src/cc/splat-types.cc"

#include "spz/src/cc/load-spz.h"
#include "spz/src/cc/load-spz.cc"

using namespace std;
// using namespace spz;

extern "C"
{
  EMSCRIPTEN_KEEPALIVE
  float hoge()
  {
    spz::GaussianCloud gs = {
        1, 0, false, {0.0, 0.0, 0.0}, {1.0, 1.0, 1.0}, {1.0, 0.0, 0.0, 0.0}, {1.0}, {1.0, 1.0, 1.0}, {}};

    vector<uint8_t> binGS;

    spz::saveSpz(gs, &binGS);

    spz::PackedGaussians unpackedGS = spz::loadSpzPacked(binGS);

    return unpackedGS.numPoints;
  }
}
にー兄さんにー兄さん

BIND系を何もやらないで関数の引数で渡すと、
js側ではundefinedとして受け取られるらしいな

にー兄さんにー兄さん

ということでembindに向き合う

たとえば、今回登場するのはJSにexportする関数と、
vectorと、GaussianCloudである

なのでこんな感じにbind

EMSCRIPTEN_BINDINGS(my_module)
{
  emscripten::function("create_test", &create_test);
  emscripten::function("hoge", &hoge);
  emscripten::function("save_spz", &save_spz);
  emscripten::function("load_spz", &load_spz);

  register_vector<float>("vector<float>");
  register_vector<uint8_t>("vector<uint8_t>");

  value_object<GaussianCloud>("GaussianCloud")
      .field("numPoints", &GaussianCloud::numPoints)
      .field("shDegree", &GaussianCloud::shDegree)
      .field("antialiased", &GaussianCloud::antialiased)
      .field("positions", &GaussianCloud::positions)
      .field("scales", &GaussianCloud::scales)
      .field("rotations", &GaussianCloud::rotations)
      .field("alphas", &GaussianCloud::alphas)
      .field("colors", &GaussianCloud::colors)
      .field("sh", &GaussianCloud::sh);
}
にー兄さんにー兄さん

embindをちゃんとすると、extern"C"しなくても関数名が維持されるし、アンダーバーがつかない

しかもstructもちゃんと対応しているので、ちゃんと向こう側で解釈される

にー兄さんにー兄さん

vectorの問題

直近で最終的にやりたいのは、JSから.spzのデータバッファを投げて
wasmでGaussianCloudに変換するというものなので、
巨大な配列をJS->C++でやり取りする必要がある

Emscriptenで配列をちゃんとやりとりするには、それ相応の知識が必要
それこそwasmでポインタ操作するわけだ

色々文献を漁ってみて、参考になりそうなリンクがいくつか見つかった

https://stackoverflow.com/questions/53602955/using-emscripten-how-to-get-c-uint8-t-array-to-js-blob-or-uint8array
https://emscripten.org/docs/api_reference/val.h.html
https://emscripten.org/docs/porting/connecting_cpp_and_javascript/Interacting-with-code.html#interacting-with-code-access-memory
これが本命かな
https://stackoverflow.com/questions/74755250/how-pass-a-large-array-from-js-to-c-using-emscripten

にー兄さんにー兄さん

いや……でもなんか最後きもいな……?

Module.sort([{num: 2}, {num: 1}, {num: 3}]);
    => [{num: 1}, {num: 2}, {num: 3}]
にー兄さんにー兄さん

上の記事めっちゃ分かりやすかった、この方法で行けるなら
C++->JSはこれでやろう

JS側でmallocしてメモリ領域を確保したうえでポインタを渡せば
それが配列として受け取れるってことか

    pointer = Module._malloc(Float64Array.BYTES_PER_ELEMENT * float64Array.length);
    if (pointer == null) throw new Error("Memory allocation failed.");
    Module.HEAPF64.set(float64Array, pointer / Float64Array.BYTES_PER_ELEMENT);
    const sumDoubleArray = Module.cwrap("sumDoubleArray", "number", ["number", "number"]);
    return sumDoubleArray(pointer, float64Array.length);

そして最後にModule._free(pointer)

にー兄さんにー兄さん

シンプルにC++は型やユーティリティで融通が利かないからイラっと来る時があるけど、頑張るぞい!

にー兄さんにー兄さん

動かないのは具体的には、
JS側でUint8Arrayを引数として渡して、C++側でuint8_t arr[]で受け取ると
なんかUnboudedTypesなるエラーが出るって感じ

にー兄さんにー兄さん
UnboundTypeError: Cannot call load_spz due to unbound types: PKh
    at new UnboundTypeError (file:///workspaces/rust-bindgen-testbed/emcc-sandbox/build/main.mjs:2033:22)
    at throwUnboundTypeError (file:///workspaces/rust-bindgen-testbed/emcc-sandbox/build/main.mjs:2080:13)
    at Object.load_spz (file:///workspaces/rust-bindgen-testbed/emcc-sandbox/build/main.mjs:2633:9)
    at main (file:///workspaces/rust-bindgen-testbed/emcc-sandbox/index.mjs:19:32)
にー兄さんにー兄さん
GaussianCloud load_spz(const int gsPtr, const int length)
{
  auto pointer = (uint8_t*)gsPtr;
  auto spzBuffer = vector<uint8_t>(pointer, pointer+length);

  return spz::loadSpz(spzBuffer);
}
    gsPtr = wasmModule._malloc(Uint8Array.BYTES_PER_ELEMENT * gsBinData.length)
    if (gsPtr == null) { throw new Error("allocation failed") }

    wasmModule.HEAPU8.set(gsBinData, gsPtr / Uint8Array.BYTES_PER_ELEMENT)

    const gsCloud = wasmModule.load_spz(gsPtr, gsBinData.length)
    console.log(gsCloud);

この組み合わせで動いた!
つまりいったんintで受け取って、ポインタへ変換する

にー兄さんにー兄さん

今度はC++の構造体内にあるvectorからJSのTypedArrayを取得する処理
JSでこんな感じにやっても、なんか動かなかった

const gsCloud = wasmModule.load_spz(gsPtr, gsBinData.length)
console.log(gsCloud.positions.data());
// error!

多分data関数がポインタを扱っているため、バインドできなかった的な何かなのかなぁ

にー兄さんにー兄さん

gsCloud.positionsはvector<float>ではあるので、
これのポインタを例のごとくintで受け取り、
HEAP32から受け取る方針にして、いい感じに動くようになった

つまりC++でこんな感じの関数を作成

int vf32_ptr(vector<float> &v)
{
  return (int)(v.data());
}
 // ...

EMSCRIPTEN_BINDINGS(my_module)
{
  emscripten::function("vf32_ptr", &vf32_ptr, allow_raw_pointers());
}
にー兄さんにー兄さん

この関数を使ってポインタを取得すれば、あとはバッファからFloat配列を取得できる

const positionPtr = wasmModule.vf32_ptr(gsCloud.positions)
const positionBuf = new Float32Array(wasmModule.HEAP32.buffer, positionPtr, gsCloud.positions.size())
console.log(positionBuf);
にー兄さんにー兄さん
Float32Array(2797680) [
  -126.175537109375,         204.15625,                 0,  126.175537109375,
          204.15625,                 0, -126.175537109375,        -204.15625,
                  0,  126.175537109375,        -204.15625,                 0,
                  0, -126.175537109375,         204.15625,                 0,
   126.175537109375,         204.15625,                 0, -126.175537109375,
         -204.15625,                 0,  126.175537109375,        -204.15625,
          204.15625,                 0, -126.175537109375,         204.15625,
                  0,  126.175537109375,        -204.15625,                 0,
  -126.175537109375,        -204.15625,                 0,  126.175537109375,
       -194.1640625,               120,        74.1640625,              -120,
         74.1640625,       194.1640625,       -74.1640625,       194.1640625,
                120,        74.1640625,       194.1640625,               120,
                  0,               240,                 0,        74.1640625,
        194.1640625,              -120,       -74.1640625,       194.1640625,
               -120,              -120,        74.1640625,      -194.1640625,
       -194.1640625,               120,       -74.1640625,              -240,
                  0,                 0,      -194.1640625,              -120,
        -74.1640625,      -194.1640625,              -120,        74.1640625,
               -120,       -74.1640625,       194.1640625,                 0,
                  0,               240,               120,        74.1640625,
        194.1640625,       194.1640625,               120,        74.1640625,
        194.1640625,               120,       -74.1640625,               120,
         74.1640625,      -194.1640625,                 0,                 0,
               -240,              -120,       -74.1640625,      -194.1640625,
        194.1640625,              -120,        74.1640625,               120,
  ... 2797580 more items
]

いい感じ!

にー兄さんにー兄さん

ちょっとわからないけど、docsを読んでいたら手軽に使えそう(危険そう)なものがあったので
一旦それを使ってみた

GaussianCloud load_spz(const emscripten::val &data)
{
  vector<uint8_t> dataBuffer = vecFromJSArray<uint8_t>(data);
  return loadSpz(dataBuffer);
}
にー兄さんにー兄さん

こんな感じにサンプルをロードしたら、ロードできた!

import { readFile } from "fs/promises";
import Module from "./build/main.mjs";

const main = async () => {
  const { load_spz } = await Module();

  const gsData = await readFile("./spz/samples/racoonfamily.spz");
  const gaussianCloud = load_spz(new Uint8Array(gsData.buffer));
  console.log(gaussianCloud);
};

main();

にー兄さんにー兄さん

これを実行するにあたり、普通の設定だとOOMが起きてしまったので
メモリをランタイムで確保するような設定にしてみた(大丈夫かな)

em++ main.cpp \
 --std=c++17 \
 -lembind \
 -o ./build/main.mjs \
 -s USE_ZLIB=1 \
 -s WASM=1 \
 -s MODULARIZE=1 \
 -s EXPORT_ES6=1 \
 -s ALLOW_MEMORY_GROWTH # <--
にー兄さんにー兄さん

ひとまずこれで第1関門は突破かなぁ
これがブラウザで動くのかとか、vectorの取り回しを気を付けようとか、そういう部分をやっていこう

にー兄さんにー兄さん

手元の環境では、embindしたC++のコードをwasmビルドして
.spz形式のGSデータをGaussianCloudとしてロードすることに
Node.jsの環境で成功した

さてここからの残件を整理しよう

にー兄さんにー兄さん

個人的には、検証はいろいろもうできたので
そろそろパッケージの本開発に移行したさがある

にー兄さんにー兄さん

残件を整理した

検証開発

  • C++ <-> JSでの配列のやり取りを理解する
    • C++へ.spzのUint8Arrayを渡す方法
    • JSで受け取ったGaussianCloudのpositionなどの配列を参照する方法
  • 出力されたwasmをViteのプロジェクトでロードして動作確認
  • tsembindによる型定義出力の対応

repo

  • monorepo構成のディレクトリ構造を考える
  • CI/CDによるパッケージデプロイ

@spz-wasm/core

  • @spz-wasm/coreのプロジェクト作成
  • .spz -> GaussianCloudへの変換ができるnpmパッケージ @spz-wasm/coreの作成

@spz-wasm/babylonjs

  • @spz-wasm/babylonjsのプロジェクトを作成
  • GausianCloud -> Gaussian Splattingのロジック実装
  • spz-wasm/babylonjsのAPI設計
  • Babylon.js側のnpmパッケージ化
にー兄さんにー兄さん

C++/JSでのvectorの受け渡しは、ここのメモにあるリンクを見て勉強かなぁ
https://zenn.dev/link/comments/a294090fc6af74

今回は、

  • JS側からC++関数呼び出しの引数で巨大なUint8Arrayを渡す
  • JS側からC++の構造体の中にある巨大なvectorを取得する

という用途なので、完全に双方向同じ処理をしたいわけではない

にー兄さんにー兄さん

これを実行してみると

const { HEAP8, _malloc, _free } = await Module();

console.log(HEAP8)
console.log(_malloc);
console.log(_free);
Int8Array(16908288) [
  101, 109, 115, 99, 103, 84, 19, 2, -2, -51, -70, -119,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,
  ... 16908188 more items
]
undefined
undefined

だった、やはりmallocとかはexportしないといけないか

にー兄さんにー兄さん

こうしてみたらいい感じに取得できた

em++ main.cpp \
 --std=c++17 \
 -lembind \
 -o ./build/main.mjs \
 -s USE_ZLIB=1 \
 -s WASM=1 \
 -s MODULARIZE=1 \
 -s EXPORT_ES6=1 \
 -s ALLOW_MEMORY_GROWTH \
 -s EXPORTED_FUNCTIONS="['_malloc', '_free']" # <--
Int8Array(16908288) [
  101, 109, 115, 99, 103, 84, 19, 2, -2, -51, -70, -119,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,   0,  0,  0, 0,  0,   0,   0,    0,
    0,   0,   0,  0,
  ... 16908188 more items
]
[Function (anonymous)]
[Function (anonymous)]
にー兄さんにー兄さん

あと、Viteでwasmが読み込めるかは早めに検証しておきたい
これによってはwasmのコンパイルオプションを変更する必要が出てくる
https://github.com/drumath2237/spz-wasm/issues/2

にー兄さんにー兄さん

monorepoのディレクトリ構成考えるよ~の会

久しぶりにpnpm-workspaceを使ったモノレポ構成を考えてみる
BayuewJSの時はturborepo&lerna-liteを導入していたが果たして

にー兄さんにー兄さん

あんまりturbopackの話を聞かないような気がしたけど、正式版がリリースされていた
turborepoも4日前にリリースされていたし、メンテも続いていそうなので良さそうだな

逆にpnpm-workspaceだけでも開発は事足りるのかな

にー兄さんにー兄さん

ばば~っと構築してみたけど、いい感じな気がする

/
├─ packages/
│    ├─ core/
│    └─ babylonjs/
├─ pnpm-workspace.yaml
└─ package.json

こんな感じの構成で、packages以下にはぞれぞれvite のlibraryテンプレートのプロジェクトを作成

pnpm -F @spz-wasm/babylonjs add --workspace @spz-wasm/core

こうしてやればworkspace内のパッケージ依存も記述可能

にー兄さんにー兄さん

名前をどうしようか問題
spz-wasmだと名前かっこいいんだけど、目的をちゃんととらえていないような感じ
自分はspzの完全なwasmラッパーを作りたいわけではなく、
Babylon.jsで読み込めればいいんだよなぁ

ということは、spz-loaderが妥当なのではと思ってきた

にー兄さんにー兄さん

うんうん、モノレポでlint/formatとビルドを一括でできるようにもしてみた

にー兄さんにー兄さん

名前はキメの問題な気がするので、よりいい名前があるならフランクに変えちゃおう

  • spz-loader
    • @spz-loader/core
    • @spz-loader/babylonjs

にするか

にー兄さんにー兄さん

なんと、Babylon.jsでもspzサポートのドラフトPRが出てしまった😇
https://github.com/BabylonJS/Babylon.js/pull/15849

TSで記述されているっぽい 従来のFileLoaderに拡張されている感じ

にー兄さんにー兄さん

価値は半減しちゃうけど、実装方法違うならまぁいったん作り上げてみるか

  • spzの最新に追従しやすい
  • (wasmによるパフォーマンス:これは要検討。オーバーヘッドがあったらむしろパフォーマンス悪いかも)
にー兄さんにー兄さん

まずはtestbedのプロジェクトをコピーしてきてwasm出力までをやろうかなぁ

凝ったことはせず、それこそC++の実装は変えず(いらないものだけ削除)、
とにかくwasm出力だけを試す

それからtsembindを試す

そしてViteで使ってみる

にー兄さんにー兄さん

C++実装を整えるのも必要だなぁ
ポインタを渡す部分とか

あとspzのGaussianCloudオブジェクトをそのまま使うのではなく、
spz-loader独自のGaussianSplatting用オブジェクトが必要

∵wasm由来のGaussianCloudはdeleteする必要があるため、JS側のメモリにコピーする必要がある

にー兄さんにー兄さん

コピーはコストが高いので、
最終的には生のspz/GaussianCloudを触れるようにしておいた方が良さそうではあるな

にー兄さんにー兄さん

オリジナルのGaussianSplattingインターフェースは、いったんViteとの連携ができてからの話になるな

にー兄さんにー兄さん

お~、pnpm workspaceで-rでビルドすると、依存関係を無視して実行することもあるんだな

にー兄さんにー兄さん

明確に、turborepoを導入すれば解決されるなこれは
そっかぁ、pnpm単体ではこの問題があるのかぁ

にー兄さんにー兄さん

ん~なんか、pnpmの説明によると
-rは依存関係のトポロジを加味した実行をしてくれるらしいけど、なぜか自分の環境だとうまくいかない
もうこれはturborepo導入したほうが早そうなのでそうしよう

にー兄さんにー兄さん

あ、まって、いけたわw(は?)

原因はあれでした、パッケージ名を変えたのに、
依存パッケージの名前を変えてなかったのが原因でした(なるほどね)

つまりこういうこと 凡ミス!!!!

にー兄さんにー兄さん

さて、ようやく問題が片付いたので
Cpp/wasmの移植作業をやろう

その前にディレクトリ構成を考えるかぁ

にー兄さんにー兄さん

spzやC++、dockerなどはpackages/core/lib/の中に集約しよう

packages/core/lib/
├─ spz-wasm/
│    ├─ .devcontainer/devcontainer.json
│    ├─ spz/      <------------------------- git submodules
│    ├─ .gitignore
│    ├─ build.sh
│    ├─ main.cpp
│    ├─ index.ts
│    └─ build/    <-------------------------- ignored
│         ├─ main.wasm.js
│         └─ main.wasm.d.ts
└─ index.ts       <-------------------------- @spz-loader/core entry point
にー兄さんにー兄さん

こんな感じに追加

➜ git submodule add git@github.com:nianticlabs/spz.git packages/core/lib/spz-wasm/spz
Cloning into 'C:/works/misc/spz-loader/packages/core/lib/spz-wasm/spz'...
remote: Enumerating objects: 24, done.
remote: Counting objects: 100% (24/24), done.
remote: Compressing objects: 100% (19/19), done.
remote: Total 24 (delta 3), reused 15 (delta 0), pack-reused 0 (from 0)
Receiving objects: 100% (24/24), 40.40 MiB | 6.19 MiB/s, done.
Resolving deltas: 100% (3/3), done.
warning: in the working copy of '.gitmodules', LF will be replaced by CRLF the next time Git touches it
にー兄さんにー兄さん

WinのPowerShellにおけるdockerのボリュームマウントの時、pwdコマンドのevalがうまくいかない問題は前からちょっと気になっていて、ちゃんとした解決策を思いついていなかった

ただ、pnpmのshell-emulationを使うことで、pwdを正しく評価できるようになり
npm-scriptからdockerによるビルドができるようになった

.npmrc
shell-emulator=true
package.json
{
  "scripts": {
    ...,
    "build:docker": "docker run --rm -v $(pwd)/lib/spz-wasm/:/src emscripten/emsdk ./build.sh",
  },
  ...
}
にー兄さんにー兄さん

そして、ちょっと改変したら
普通にwasmの処理が動いて、spzがブラウザで動いた!!

にー兄さんにー兄さん

特にwasm pluginとかも必要なく動いたな……?
まだtsembindも使っていないけど、いったんts-ignoreを使ってESModuleのコードをimportしている感じ

にー兄さんにー兄さん

wasmのビルドも、特にSINGLE_FILEにするなどもしていない

見てみると、packages/core/dist/index.jsの中に
application/wasmとしてwasmのb64化されたものが追記されている

ん~まじか、これこそwasm-plugin使わないとできないのかと思ったら
Viteが自動でやってくれたんだろうか、すごいな

にー兄さんにー兄さん

ビルド時にlibraryモードを無効化して、previewしてみると
たしかにwasmはwasmでassetsにいることが分かる(spzがなかったのでエラーだったけど)

にー兄さんにー兄さん

どうやらtsembindはesmoduleに対応していないらしい
やってみたらERR_REQUIRE_ESMというエラーコードが出てきた(これは一般的なエラーコードらしい)

にー兄さんにー兄さん

まぁということで、emscriptenのdockerコンテナ上でtypescriptをグローバルインストールして
--emit-tsdオプションを付けたら行けたわ

割といい感じ生成してくれていそうだけどどうだろ

にー兄さんにー兄さん

思った以上にEmscriptenのTS型定義生成が最高すぎるかもしれない件

にー兄さんにー兄さん

TS型定義生成がうまくいったので、ここからはもうcppやtsのロジックをバリバリ書いていく感じになる(主にtsだけど)

にー兄さんにー兄さん

cppではtestbedで使っていた型定義というか、embindを見直した
dtsファイルはこれの通りに生成されるため、名前をVector~~みたいな感じにしたり

あと、C++で扱っているGaussianCloudをRawとし、
TSで扱うGaussianCloudとは別物とした

register_vector<float>("VectorFloat32");
register_vector<uint8_t>("VectorUInt8T");

value_object<GaussianCloud>("RawGaussianCloud")
    .field(...);
にー兄さんにー兄さん

そして、RawGSCloudとGSCloud変換するようなロジックを作成した

まずRawGSCloudのほうにはC++のvectorで扱っているようなポインタが含まれるので
ポインタからFloat32Arrayをコピーして生成するようにした

/**
 * create new Float32Array from cpp FloatVector
 * @param wasmModule emscripten main module
 * @param vec cpp float32 vector
 * @returns copied float32 array
 */
export const floatVectorToFloatArray = (
  wasmModule: MainModule,
  vec: VectorFloat32,
): Float32Array => {
  const pointer = wasmModule.vf32_ptr(vec);
  const size = vec.size();

  const buffer = new Float32Array(wasmModule.HEAPF32.buffer, pointer, size);
  const copiedBuffer = new Float32Array(buffer);

  return copiedBuffer;
};
にー兄さんにー兄さん

TypedArrayはnewするときに、ArrayBufferを渡すとViewが作られ
TypedArrayを渡すとコピーが作られる仕様なので、
TypedArrayを渡してコピーを作成している

これでRawのほうはDisposeできる(リソースの開放)

にー兄さんにー兄さん

GaussianCloudも定義はこんな感じ
とはいえRawのVector部分がFloat32Arrayになっただけではある

export type GaussianCloud = {
  numPoints: number;
  shDegree: number;
  antialiased: boolean;
  positions: Float32Array;
  scales: Float32Array;
  rotations: Float32Array;
  alphas: Float32Array;
  colors: Float32Array;
  sh: Float32Array;
};
にー兄さんにー兄さん

そういえばspzのAPIで、
loadSpzloadSpzPackedで何が違うんだろうって思ったら
まずpackedのほうはPackedGaussianCloudを返す。
これはspzファイルの内容をただ読み込んだだけという感じで、
spzのフィールド値が基本的にvector<uint8_t>になっている

これをunpackGaussians関数でunpackすると、数値的に復元されたGaussianCloudが出てくるという感じ

そしてこんな感じの位置付け

  • loadSpzPackedは、spzをロードしてPackedGaussianCloudを返すまで
  • loadSpzはさらにunpackしてGaussianCloudを返す

という感じなので、基本的にはloadSpzを使えば良さそう

にー兄さんにー兄さん

テスト用にScaniverseでGSデータを取って読み込ませてみた

これはsrc/mainのほうで動いているのでpnpm devで起動していて
assetsフォルダに入ったspzを読み込んでいる様子

うん、点群数は問題なさそう

なんだけど、なんかcolorとかpositionとかが変なデータが入っている
colorに負の値が入ってるのが気になる、なんだこれ
あとなんか全体的にsizeが同じバッファになってる

にー兄さんにー兄さん

バッファがなんか変なのは、HEAPを読み取るときに
HEAPを直接呼んでいたから 正しくは.bufferしなくちゃだった

にー兄さんにー兄さん

試しにcolorだけuint8arrayでやってみるとかもしてみた

alphaも以前、なんか1以上とか負の値が入っていてよくわからんな?

にー兄さんにー兄さん

いったん、issueの内容自体は達成したので、フランクにPRをマージしていく

にー兄さんにー兄さん

Babylon.jsでは、最終的にsplat形式のArrayBufferを作成すればGaussianSplattingオブジェクトを作成できる

CedricさんのSplatFileLoaderの中身を参考にデータ変換を行ってみるか
https://github.com/CedricGuillemet/Babylon.js/blob/065ceadd6f7db15cefc7a688bd5c516614998f46/packages/dev/loaders/src/SPLAT/splatFileLoader.ts#L184-L314

にー兄さんにー兄さん

ParsedPLYを返してるけど、必要なのはその中のdataっていうArrayBufferなので
これさえできればいいかな

まぁ一応ParsedPLYも作っておくか

にー兄さんにー兄さん

さて、GSCloudからSplatのArrayBufferに変換するコードを書いてみた

`packages/babylonjs/lib/index.ts`のコード
import { GaussianSplattingMesh, type Scene } from "@babylonjs/core";
import { loadSpz, type GaussianCloud } from "@spz-loader/core";

export const createGaussianSplattingFromSpz = async (
  data: Uint8Array,
  name?: string,
  scene?: Scene,
  keepInRam?: boolean,
): Promise<GaussianSplattingMesh> => {
  const splat = new GaussianSplattingMesh(
    name ?? "spz splat",
    null,
    scene,
    keepInRam,
  );
  const splatBuffer = await parseSpzToSplat(data);
  await splat.loadDataAsync(splatBuffer);
  return splat;
};

export const parseSpzToSplat = async (
  data: ArrayBuffer,
): Promise<ArrayBuffer> => {
  const gsCloud = await loadSpz(new Uint8Array(data));
  return _convertGaussianCloudToSplatBuffer(gsCloud);
};

const _convertGaussianCloudToSplatBuffer = (
  gsCloud: GaussianCloud,
): ArrayBuffer => {
  // position(3*f32) + scale(3*f32) + color(4*u8) + rotation(4*u8)
  const rowOutputLength = 3 * 4 + 3 * 4 + 4 + 4;
  const splatCount = gsCloud.numPoints;
  const buffer = new ArrayBuffer(rowOutputLength * splatCount);

  const position = new Float32Array(buffer);
  const scale = new Float32Array(buffer);
  const rgba = new Uint8ClampedArray(buffer);
  const rot = new Uint8ClampedArray(buffer);

  // positions
  for (let i = 0; i < splatCount; i++) {
    position[i * 8 + 0] = gsCloud.positions[i * 3 + 0];
    position[i * 8 + 1] = gsCloud.positions[i * 3 + 1];
    position[i * 8 + 2] = gsCloud.positions[i * 3 + 2];
  }

  // colors
  for (let i = 0; i < splatCount; i++) {
    rgba[i * 32 + 24 + 0] = gsCloud.colors[i * 3 + 0];
    rgba[i * 32 + 24 + 1] = gsCloud.colors[i * 3 + 1];
    rgba[i * 32 + 24 + 2] = gsCloud.colors[i * 3 + 2];
    rgba[i * 32 + 24 + 3] = gsCloud.alphas[i];
  }

  // scales
  for (let i = 0; i < splatCount; i++) {
    scale[i * 8 + 3 + 0] = gsCloud.scales[i * 3 + 0];
    scale[i * 8 + 3 + 1] = gsCloud.scales[i * 3 + 1];
    scale[i * 8 + 3 + 2] = gsCloud.scales[i * 3 + 2];
  }

  // rotations
  for (let i = 0; i < splatCount; i++) {
    rot[i * 32 + 28 + 1] = gsCloud.rotations[i * 4 + 0] * 127.5;
    rot[i * 32 + 28 + 2] = gsCloud.rotations[i * 4 + 1] * 127.5;
    rot[i * 32 + 28 + 3] = gsCloud.rotations[i * 4 + 2] * 127.5;
    rot[i * 32 + 28 + 0] = gsCloud.rotations[i * 4 + 3] * 127.5;
  }

  return buffer;
};
にー兄さんにー兄さん

使い勝手としてはこんな感じ

packages/babylonjs/src/main.ts
const spzData = await fetch(spzPath).then((res) => res.arrayBuffer());
await createGaussianSplattingFromSpz(spzData);
にー兄さんにー兄さん

しかし、この状態でdevサーバを開くと
こんな感じのエラーがコンソールに表示された

Server responded with status code 431. See https://vite.dev/guide/troubleshooting.html#_431-request-header-fields-too-large.
にー兄さんにー兄さん

調べてみると、どうやらリクエストヘッダが大きすぎるとのこと

なんかリクエストする時にwasmのb64文字列がそのままURLに入ってるもんで、それはヤバいな
見てみると、spz-wasmのbuild/main.mjsにこんなコードが

bundler-friendlyとはいえ、これが最終的にb64に置き換わっているようで
たしかにこんなに長い文字列のurlだったらクラッシュするだろうなと

にー兄さんにー兄さん

ちなみに、解決策としてsrciptsにこんな感じに記述することも可能
しかし、こうしたところで根本解決はせず、普通にURLが長いエラーでクラッシュした

package.json
  "scripts": {
    "dev": "NODE_OPTIONS=--max_http_header_size=12800000 vite",
にー兄さんにー兄さん

なんかcolorを見ていたら変~~~~な感じのことをしていたので、それをいい感じに補正してやったら
なんかそれっぽいのが出てきた!

だけどなんか回転なのかな、が変、惜しい

にー兄さんにー兄さん

さて、全然目も取れていなかったけど
PRのpushごとにlint/formatとビルドのCIジョブを回したいと思って色々やっていました

当初は、せっかくemscriptenの公式dockerを使っているのでそれを使いたかったけど
なぜか全然GitHub Actions上で動かなかった
おもにPermission系で全然いうことを聞いてくれなくてお手上げになってしまった

なんか、良いところまで来たかなって思ったんだけど
npmでtypescriptをインストしようとしたら/.npmへのアクセス権が無くて失敗したり

にー兄さんにー兄さん

最終的にCIは、emscriptenが使えるactionsがあったため、
それを使いつつem++でコンパイルするという形になった
いや~~環境が違ってくるのでちょっと大丈夫かなぁと思うんだけど、
とりあえずしょうがなし・・・

にー兄さんにー兄さん

scaniverseの表示と比べて、なんかコントラストが弱いのが気になるな?なんだろ

ここに関してちょっと調査をしたところ、どうやらSHの有無や色情報の帯域?によって
色に補正をかけるための係数が違うらしい(↓はload-spz.cc

Scale factor for DC color components. To convert to RGB, we should multiply by 0.282, but it can be useful to represent base colors that are out of range if the higher spherical harmonics bands bring them back into range so we multiply by a smaller value.

にー兄さんにー兄さん

で、コード中では0.15をかけていたけど、0.282をかけて復元するようにしたら
公式のscaniverseと同じような見た目になった

にー兄さんにー兄さん

メモリリーク原因対策の前に、emscriptenのwasmのデバッグについて確認する
このページを参考に
https://developer.chrome.com/blog/wasm-debugging-2020?hl=ja

にー兄さんにー兄さん

基本的に拡張機能を入れれば使えそう

だけどDockerやWSLといった、ブラウザを実行している環境ではなく
その環境をマウントしている別の仮想環境などを使っている場合は注意

wasmの場合はsourcemap的な感じでDWARFというデバッグ情報を埋め込むのだけど
そのパスがビルド環境のものなので、いざブラウザ実行環境でパスを参照しようとしたときにエラーになってしまう

例えば自分は、ブラウザはWin11Homeで実行していて
emscriptenを公式のemscripten/emsdkというdockerコンテナ上でビルドしている
そしてカレントディレクトリは/srcにマウントされる仕組みである

そうなったときに、拡張機能のオプションからパスをoverrideするように設定する

にー兄さんにー兄さん

するとこんな感じに、ChromeのDevToolsから
wasmのデバッグをC++コードを参照しながらできるのである。すごいな、便利やな~~

にー兄さんにー兄さん

メモリリークの件に関して

ここの分野に詳しくないというのもそうなんだけど、
メモリリーク起きてるのかな?という気になってきた

どうやら@spz-loader/coreのほうでは起きていないような挙動になっていて
@spz-loader/babylonjsでは怪しい
具体的にはリロードのたびにメモリが増える挙動になるものの、
ここではJSのメモリ空間を扱っているので
適切にGCが効く場合もある、みたいな

一番考えていたwasmとのメモリ空間での問題がないのであれば
そこまで大きな問題ではないのではと思って
0.1.0リリース後の対応でもいいような優先度になりました

にー兄さんにー兄さん

とはいえ、babylonjsローダが正直普通ではありえんくらいメモリを使っている感じではあるので
早く何とかしたい気持ちではいますね

にー兄さんにー兄さん

やはりやっておこうということで、メモリ使用の最適化を行った
確実にコピーが多いだろうなぁって部分を削ったり、
あとめっちゃ聞いたのはEmscriptenのコンパイラプションだった

にー兄さんにー兄さん

最終的に採用したオプションたち

em++ main.cpp \
  -Oz \
  --closure 1 \
  --llvm-lto 1 \
  -fno-exceptions \
  --std=c++17 \
  -lembind \
  --emit-tsd main.d.ts \
  -o ./build/main.mjs \
  -s USE_ZLIB=1 \
  -s WASM=1 \
  -s SINGLE_FILE=1 \
  -s MODULARIZE=1 \
  -s EXPORT_ES6=1 \
  -s EXPORT_NAME="spzwasm" \
  -s ALLOW_MEMORY_GROWTH \
  -s EXPORTED_FUNCTIONS="['_malloc', '_free']"

これを取り入れたら、もともと肥大しがちだったメモリ使用量が1/2いかになったりめっちゃ良かった
読み込み速度も上がった気がする