btensor.h File Reference

#include <optional>
#include <mim/lam.h>
#include <mim/tuple.h>
#include <mim/world.h>
#include <mim/plug/affine/affine.h>
#include <mim/plug/buffer/buffer.h>
#include <mim/plug/core/core.h>
#include <mim/plug/cps/cps.h>
#include <mim/plug/mem/mem.h>
#include "mim/plug/btensor/autogen.h"

Include dependency graph for btensor.h:

This graph shows which files directly or indirectly include this file:

Go to the source code of this file.

Classes
struct	mim::plug::btensor::MapRedPost
	The curried meta arguments of a tensor.map_reduce_post or btensor.map_reduce_post callee (same layout). More...
struct	mim::plug::btensor::MapRedPost::Ranks

Namespaces
namespace	mim
namespace	mim::plug
namespace	mim::plug::btensor
	The btensor Plugin

Functions
const Def *	mim::plug::btensor::op_affine_map (const Def *f, const Def *m, const Def *n, const Def *s_in, const Def *s_out, const Def *idxs, const Def *mem)
	affine.map @(m, n) @(s_in, s_out) f idxs as mem with the address space as of mem ↦ [mem', coords].
void	mim::plug::btensor::app_cps (Lam *caller, const Def *f, DefVec args, const Def *ret)
	Calls the combiner/epilogue f from caller with args and the continuation ret.
std::pair< Lam *, const Def * >	mim::plug::btensor::ds_splice (const Def *arg, const Def *codom, std::string name)
	A fresh fun: Fn arg->type() → codom named name, spliced into direct style as cps.cps2ds_dep fun arg.
DefVec	mim::plug::btensor::to_idxs (World &w, Defs extents, Defs iters)
	Converts the raw i64 loop counters iters into indices Idx extents[d].
const Def *	mim::plug::btensor::folded_idx (const Def *shape, Defs iters)
	The raw i64 loop counters iters as a buffer index into shape, with its literal size-1 axes folded away.
const Def *	mim::plug::btensor::clamp (const Def *x, const Def *hi)
	Clamps the i64 x into [0, hi].
std::pair< DefVec, const Def * >	mim::plug::btensor::pad_coords (Defs iters, const Def *lo, const Def *s_in, bool replicate)
	The read coordinates o#d − lo#d of a pad into Idx s_in#d - clamped to the nearest edge if replicate, or else zeroed out of range - together with the all-axes in-range flag (nullptr if replicate).
std::optional< fe::Vector< u64 > >	mim::plug::btensor::concat_extents (const Def *Sis, u64 nis, u64 r, u64 ax)
	The literal extents of the nis input shapes Sis along the concat axis ax; std::nullopt if one is not literal.
const Def *	mim::plug::btensor::concat_select (const Def *o_ax, const fe::Vector< u64 > &exts, auto read)
	The concat cell at the i64 axis coordinate o_ax: read(i, loc) of the input i owning it - the highest one whose prefix offset is ≤ o_ax - at the axis coordinate loc clamped into that input.
const Def *	mim::plug::btensor::build_pointwise (const Def *result_ty, const Def *mem, const Def *dst, const Def *ins, const Def *s_out, u64 rn, const std::string &name, auto compute)
	A [mem, dst, ins] → [mem, Buf] fun spliced via cps.cps2ds_dep that loops over s_out carrying [mem, dst] and writes compute's element at the loop's (identity) output coordinates.