Sparse Halide

[soufianekhiat]

I would to open a discussion about feasability and a possible new feature to change index for being able to remap sparse data to dense memory layout.

Here few of my prototype:

Version 0 (working but limited)

I'm exploring the feasability to extend the usage of Halide for Sparse data.

• Read Only (first)
• Read/Write Sparse

Based on:
"Halide is a programming language designed to make it easier to write high-performance image and array processing code on modern machines."
That can fit on this definition.

A simple way to tackle the problem is to keep everything in dense memory, and considering sparse as a simple reindexing to a dense memory layout.

Here my first prototype:

SparseTile2D sparse2DIn;
sparse2DIn.Resize(Vec2u(8, 8), 3);
sparse2DIn.fullBoundMin = { 0, 0 };
sparse2DIn.fullBoundMax = { 32, 32 };
sparse2DIn.AddTile({ 2, 2 });
sparse2DIn.AddTile({ 10, 16 });
sparse2DIn.AddTile({ 0, 20 });
sparse2DIn.Randomize01();

Func fSparse2DIn{ "FuncSparse2DIn" };
fSparse2DIn = FuncFromSparse2D(sparse2DIn, "_In"); // <= Here everything happen

Func fDenseOut{ "FuncDense2DOut" };
fDenseOut(x, y) =
	(
		fSparse2DIn(x - 1, y - 1) + fSparse2DIn(x, y - 1) + fSparse2DIn(x + 1, y - 1) +
		fSparse2DIn(x - 1, y + 0) + fSparse2DIn(x, y + 0) + fSparse2DIn(x + 1, y + 0) +
		fSparse2DIn(x - 1, y + 1) + fSparse2DIn(x, y + 1) + fSparse2DIn(x + 1, y + 1)
	) / 9.0f;

fDenseOut.compile_jit();

std::vector<f32> denseOut;
denseOut.resize(32 * 32);
std::fill(denseOut.begin(), denseOut.end(), 0.0f);
Halide::Buffer<f32> bufferOut(&denseOut[0], { 32, 32 }, "Dense2DBufferOut");

fDenseOut.realize(bufferOut);

std::vector<UInt8> byteData(32 * 32 * 4);
for (int j = 0; j < 32; ++j)
{
	for (int i = 0; i < 32; ++i)
	{
		byteData[4*(j*32 + i) + 0] = static_cast<UInt8>(std::round(255.0f * std::clamp(denseOut[j*32 + i], 0.0f, 1.0f)));
		byteData[4*(j*32 + i) + 1] = static_cast<UInt8>(std::round(255.0f * std::clamp(denseOut[j*32 + i], 0.0f, 1.0f)));
		byteData[4*(j*32 + i) + 2] = static_cast<UInt8>(std::round(255.0f * std::clamp(denseOut[j*32 + i], 0.0f, 1.0f)));
		byteData[4*(j*32 + i) + 3] = 0xFF;
	}
}
stbi_write_png("../test.png", 32, 32, 4, &byteData[0], 0);

The reindexing job is done with:

template < typename Type >
Expr select(std::vector<Expr> cs, std::vector<Expr> vs, Type const defaultValue)
{
	Expr c = cs.back();
	Expr v = vs.back();

	cs.pop_back();
	vs.pop_back();

	if (cs.size() > 1)
	{
		return select(c, v, select< Type >(cs, vs, defaultValue));
	}
	else
	{
		return select(c, v, select(cs[0], vs[0], Expr(defaultValue)));
	}
}

Func	FuncFromSparse2D(SparseTile2D& data, std::string const& sDeco)
{
	Func newFunc{ ("Sparse2D_" + std::to_string(*reinterpret_cast<UInt64*>(&data))) + sDeco };

	std::vector<Expr>	conds;
	std::vector<Expr>	trues;

	Buffer<f32> pagesBuffer(data.pPages, { (int)(data.memSize) }, ("BufferSparseDense2D_Pages" + sDeco).c_str());

	// Generate bunch of if/else to properly reindex the proper memory pointer.
	UInt32 pageStart = 0u;
	for (TileBound2D const& b : data.tileBounds)
	{
		conds.push_back(((int)b.origin.x) <= x && x < ((int)b.origin.x + (int)data.tileSize.x) &&
						((int)b.origin.y) <= y && y < ((int)b.origin.y + (int)data.tileSize.y));
		Expr xx = clamp(x - (int)b.origin.x, 0, (int)(data.tileSize.x - 1));
		Expr yy = clamp(y - (int)b.origin.y, 0, (int)(data.tileSize.y - 1));
		Expr index = yy * Expr((int)data.tileSize.x) + xx;
		trues.push_back(pagesBuffer(Halide::clamp(index + (int)pageStart, 0, ((int)(data.memSize - 1)))));

		pageStart += data.tileSize.x * data.tileSize.y;
	}

	std::reverse(conds.begin(), conds.end());
	std::reverse(trues.begin(), trues.end());
	newFunc(x, y) = select< f32 >(conds, trues, 0.0f);

	return newFunc;
}

Produces:

This approch would work for single indirection data. But limit us for richer indirection from "local sparse data". For instance I would like to have NanoVDB version based on Halide.

Version 1 (WIP proposition, open to discussion).

Inspired by Taichi (the C++ front-end is barely usable). I have a new front-end:

enum class MemberType
{
	Int8,
	Int16,
	Int32,
	Int64,
	UInt8,
	UInt16,
	UInt32,
	UInt64,
	f16,
	f32,
	f64
};

enum class MemberSemantic
{
	None,
	Vec2,
	Vec3,
	Vec4,
	Color_sRGB_Rec709,
	Color_sRGBA_Rec709,
	Color_RGB_Rec2020,
	Color_RGBA_Rec2020,
	Mat22,
	Mat33,
	Mat34,
	Mat44,
	Complex,
	Dual,
	Quaternion
};

struct Block;

typedef Block& BlockIt;

struct Block
{
private:
	enum class Structural : Int8
	{
		None = -1,
		Dense,
		Hash, // I don't know if it's feasable with Halide
		Dynamic // Just reserve max memory with an atomic counter like AppendStructuredBuffer on DirectX
	};

	enum class Decoration : UInt8
	{
		None = 0,
		Morton = (1 << 0),
		Bitmasked = (1 << 1), // TODO
		Pointer = (1 << 2) // For Halide just CPU or GPU just an int to offset the 'start' global (or local) memory.
	};

	friend Decoration operator|=(Decoration& eA, Decoration const eB)
	{
		eA = static_cast<Decoration>(static_cast<UInt8>(eA) | static_cast<UInt8>(eB));
		return eA;
	}

public:
	BlockIt Place(UInt32 const uID, std::string const& sName, MemberType const eMemberType, MemberSemantic const eSemantic = MemberSemantic::None);

	BlockIt Dense(std::vector<Halide::Var> const& aIndexingVars, std::vector<UInt32> const& aGridSizes);
	BlockIt Hash(std::vector<Halide::Var> const& aIndexingVars, UInt32 const uHashSize, HashAlgo const eHashAlgo = HashAlgo::A);
	BlockIt Dynamic(std::vector<Halide::Var> const& aIndexingVars, std::vector<UInt32> const& aMaxSizes);

	BlockIt Morton();
	BlockIt Bitmasked();
	BlockIt Pointer();

	Halide::Func	Read(UInt32 const uID, Halide::Func rawByteData, std::vector<Halide::Expr> const& Xs);

	UInt32 ByteSizeMembers() const;
	UInt32 MaxBufferSize() const;

	Structural	eStructural		= Structural::None;
	Decoration	eDecoration		= Decoration::None;
	UInt32		uCurrentByteOffset	= 0u;
	UInt32		uElementCountRaw	= 0u;
	UInt32		uElementCount		= 0u;
	UInt32		uMemberCount		= 0u;

	std::vector<Halide::Var>		aIndexingVars	= {};
	std::unordered_map<UInt32, Member>	aMembers	= {};
	std::vector<SparseHalide::Block>	aChildren	= {};
	std::vector<Halide::Func>		aPlaced		= {};

	Bool		bIsRoot = true;

	std::vector<UInt32> 	aGridSizes	= {};
	UInt32			uHashSize	= 0xFFFFFFFFu;
	HashAlgo		eHashAlgo	= HashAlgo::A; // TBD
	std::vector<UInt32>	aMaxSizes	= {};
};

Here few usecase:

The Read function should produce the offset for the read or write.

RegularGrid AoS

SparseHalide::Block GridAoS;
GridAoS.Dense({ x, y, z }, { 128, 128, 128 })
	.Place(0, "p", SparseHalide::MemberType::f32, SparseHalide::MemberSemantic::Vec3);

//////////////////////////////////////////////////////////////////////////
// High Level equivalent:
struct Root
{
	Vec3 p[128][128][128];
};

// p Reader i = {0, 1, 2}:
f32 value_i = memory(3*(128*(128*x + y) + z) + i);

// p Writer i = {0, 1, 2}:
memory(3*(128*(128*x + y) + z) + i) = value_i;

RegularGrid SoA

SparseHalide::Block GridSoA;
GridSoA.Dense({ x, y, z }, { 128, 128, 128 })
	.Place(0, "x", SparseHalide::MemberType::f32)
	.Place(1, "y", SparseHalide::MemberType::f32)
	.Place(2, "z", SparseHalide::MemberType::f32);
//////////////////////////////////////////////////////////////////////////
// High Level equivalent:
struct Root
{
	Vec3 p[128][128][128];
};

Here the main idea is the operator () or the function Read will produce another function for reading,
Exactly what BoundaryCondition do:

// Reader
f32 value_x = memoryX(3*(128*(128*x + y) + z) + 0);
f32 value_y = memoryY(3*(128*(128*x + y) + z) + 1);
f32 value_z = memoryZ(3*(128*(128*x + y) + z) + 2);

// Writer
memoryX(3*(128*(128*x + y) + z) + 0) = value_x;
memoryY(3*(128*(128*x + y) + z) + 1) = value_y;
memoryZ(3*(128*(128*x + y) + z) + 2) = value_z;

Quadtree MaxDepth 3

SparseHalide::Block QuadTree;
QuadTree
	.Dense({ x, y }, { 4, 4 }).Bitmasked().Pointer()
	.Dense({ x, y }, { 4, 4 }).Bitmasked().Pointer()
	.Dense({ x, y }, { 4, 4 })
		.Place(0, "Vel", SparseHalide::MemberType::f32, SparseHalide::MemberSemantic::Vec3)
		.Place(1, "Density", SparseHalide::MemberType::f32);
//////////////////////////////////////////////////////////////////////////
// High Level equivalent
struct Root
{
	std::bitfield<4*4>	childsMasks;
	Child1*				children1[4][4];
};
struct Child1
{
	std::bitfield<4*4>	childsMasks;
	Child2*				children1[4][4];
};
struct Child2
{
	Vec3	Vel[4][4];
	f32		Density[4][4];
};

// Dense Memory Layout:
// ... memoryBitRoot ...
// --- padding ---
// ... memoryBitChild1...
// --- padding ---
// ... memoryBitChild2...
// --- padding ---
// velocities
// --- padding ---
// densities
// --- padding ---

// Pseudo code, Reader
Vec2i ChildRoot;
ChildRoot.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f/4.0f/4.0f);
ChildRoot.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f/4.0f/4.0f);
Vec2i Child1;
Child1.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f/4.0f);
Child1.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f/4.0f);
Vec2i Child2;
Child2.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f);
Child2.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f);

bool bitRoot = ReadBit(memoryBitRoot, ChildRoot);
bool bitChild1 = ReadBit(memoryBitChild1, Child1);
bool bitChild2 = ReadBit(memoryBitChild2, Child2);

// 64 == 4*4*4

f32 vel_i = select(bitRoot,
					select(bitChild1,
						select(bitChild2,
							ReadFloat(memoryVel, 64, 64, 64, 3, x, y, z, i)
								, 0.0f)
							, 0.0f)
						, 0.0f);

// Pseudo code Writer:
Vec2i ChildRoot;
ChildRoot.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f/4.0f/4.0f);
ChildRoot.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f/4.0f/4.0f);
Vec2i Child1;
Child1.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f/4.0f);
Child1.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f/4.0f);
Vec2i Child2;
Child2.x = RoundToInt(Bound(v.x, minX, maxX)/4.0f);
Child2.y = RoundToInt(Bound(v.y, minY, maxY)/4.0f);

bool bitRoot = SetBit(memoryBitRoot, ChildRoot, true);
bool bitChild1 = SetBit(memoryBitChild1, Child1, true);
bool bitChild2 = SetBit(memoryBitChild2, Child2, true);

// 64 == 4*4*4

SetFloat(memoryVel, 64, 64, 64, 3, x, y, z, i, value_i);
//////////////////////////////////////////////////////////////////////////

Some extra sparse layout:

// VDB
/*
K. Museth, “VDB: High-Resolution Sparse Volumes With Dynamic Topology”. ACM Transactions on Graphics, Volume 32, Issue 3, Pages 27:1-27:22, June 2013, Presented at SIGGRAPH 2013 in Anaheim[.](http://dx.doi.org/10.1145/2487228.2487235)
*/
SparseHalide::Block VDB;
VDB
	.Hash({ x, y, z }, 32)
	.Dense({ x, y, z }, { 16, 16, 16 }).Pointer()
	.Dense({ x, y, z }, { 16, 16, 16 }).Pointer()
	.Dense({ x, y, z }, { 8, 8, 8 })
		.Place(0, "u", SparseHalide::MemberType::f32)
		.Place(1, "v", SparseHalide::MemberType::f32)
		.Place(2, "w", SparseHalide::MemberType::f32);

// SPGrid
/*
SPGrid: a sparse paged grid structure applied to adaptive smoke simulation
Rajsekhar Setaluri, Mridul Aanjaneya, Sean Bauer and Eftychios Sifakis
ACM Transactions on Graphics (SIGGRAPH Asia proceedings), 33, 6, 205:1-12 (2014)
*/
SparseHalide::Block SPGrid;
SPGrid
	.Dense({ x, y, z }, { 512, 512, 512 }).Morton().Bitmasked()
	.Dense({ x, y, z }, { 8, 4, 4 })
	.Place(0, "flag", SparseHalide::MemberType::UInt8)
	.Place(1, "u", SparseHalide::MemberType::f32)
	.Place(2, "v", SparseHalide::MemberType::f32)
	.Place(3, "w", SparseHalide::MemberType::f32);

I'm open to any comment.