af_err af_sort(af_array *out, const af_array in, const unsigned dim, const bool isAscending)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
if(info.elements() == 0) {
return af_retain_array(out, in);
}
DIM_ASSERT(1, info.elements() > 0);
af_array val;
switch(type) {
case f32: val = sort<float >(in, dim, isAscending); break;
case f64: val = sort<double >(in, dim, isAscending); break;
case s32: val = sort<int >(in, dim, isAscending); break;
case u32: val = sort<uint >(in, dim, isAscending); break;
case s16: val = sort<short >(in, dim, isAscending); break;
case u16: val = sort<ushort >(in, dim, isAscending); break;
case s64: val = sort<intl >(in, dim, isAscending); break;
case u64: val = sort<uintl >(in, dim, isAscending); break;
case u8: val = sort<uchar >(in, dim, isAscending); break;
case b8: val = sort<char >(in, dim, isAscending); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out, val);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
const af_array keys, const af_array values,
const unsigned dim, const bool isAscending)
{
try {
ArrayInfo info = getInfo(keys);
af_dtype type = info.getType();
ArrayInfo vinfo = getInfo(values);
DIM_ASSERT(3, info.elements() > 0);
DIM_ASSERT(4, info.dims() == vinfo.dims());
// Only Dim 0 supported
ARG_ASSERT(5, dim == 0);
af_array oKey;
af_array oVal;
switch(type) {
case f32: sort_by_key_tmplt<float >(&oKey, &oVal, keys, values, dim, isAscending); break;
case f64: sort_by_key_tmplt<double >(&oKey, &oVal, keys, values, dim, isAscending); break;
case s32: sort_by_key_tmplt<int >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u32: sort_by_key_tmplt<uint >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u8: sort_by_key_tmplt<uchar >(&oKey, &oVal, keys, values, dim, isAscending); break;
case b8: sort_by_key_tmplt<char >(&oKey, &oVal, keys, values, dim, isAscending); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out_keys , oKey);
std::swap(*out_values , oVal);
}
CATCHALL;
return AF_SUCCESS;
}
static void print(af_array arr)
{
const ArrayInfo info = getInfo(arr);
T *data = new T[info.elements()];
af_array arrT;
AF_CHECK(af_reorder(&arrT, arr, 1, 0, 2, 3));
//FIXME: Use alternative function to avoid copies if possible
AF_CHECK(af_get_data_ptr(data, arrT));
const ArrayInfo infoT = getInfo(arrT);
AF_CHECK(af_destroy_array(arrT));
std::ios_base::fmtflags backup = std::cout.flags();
std::cout << "[" << info.dims() << "]\n";
#ifndef NDEBUG
std::cout <<" Offsets: ["<<info.offsets()<<"]"<<std::endl;
std::cout <<" Strides: ["<<info.strides()<<"]"<<std::endl;
#endif
printer(std::cout, data, infoT, infoT.ndims() - 1);
delete[] data;
std::cout.flags(backup);
}
af_err af_sort_index(af_array *out, af_array *indices, const af_array in, const unsigned dim, const bool isAscending)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
DIM_ASSERT(2, info.elements() > 0);
// Only Dim 0 supported
ARG_ASSERT(3, dim == 0);
af_array val;
af_array idx;
switch(type) {
case f32: sort_index<float >(&val, &idx, in, dim, isAscending); break;
case f64: sort_index<double >(&val, &idx, in, dim, isAscending); break;
case s32: sort_index<int >(&val, &idx, in, dim, isAscending); break;
case u32: sort_index<uint >(&val, &idx, in, dim, isAscending); break;
case u8: sort_index<uchar >(&val, &idx, in, dim, isAscending); break;
case b8: sort_index<char >(&val, &idx, in, dim, isAscending); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out , val);
std::swap(*indices, idx);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_shift(af_array *out, const af_array in, const af::dim4 &sdims)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
DIM_ASSERT(1, info.elements() > 0);
af_array output;
switch(type) {
case f32: output = shift<float >(in, sdims); break;
case c32: output = shift<cfloat >(in, sdims); break;
case f64: output = shift<double >(in, sdims); break;
case c64: output = shift<cdouble>(in, sdims); break;
case b8: output = shift<char >(in, sdims); break;
case s32: output = shift<int >(in, sdims); break;
case u32: output = shift<uint >(in, sdims); break;
case u8: output = shift<uchar >(in, sdims); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out,output);
}
CATCHALL;
return AF_SUCCESS;
}
af_err convert(af_array* out, const af_array in, const float r, const float g, const float b)
{
try {
ArrayInfo info = getInfo(in);
af_dtype iType = info.getType();
af::dim4 inputDims = info.dims();
// 2D is not required.
if(info.elements() == 0) {
dim_t my_dims[] = {0, 0, 0, 0};
return af_create_handle(out, AF_MAX_DIMS, my_dims, iType);
}
// If RGB is input, then assert 3 channels
// else 1 channel
if (isRGB2GRAY) ARG_ASSERT(1, (inputDims[2]==3));
else ARG_ASSERT(1, (inputDims[2]==1));
af_array output = 0;
switch(iType) {
case f64: output = convert<double, double, isRGB2GRAY>(in, r, g, b); break;
case f32: output = convert<float , float , isRGB2GRAY>(in, r, g, b); break;
case u32: output = convert<uint , float , isRGB2GRAY>(in, r, g, b); break;
case s32: output = convert<int , float , isRGB2GRAY>(in, r, g, b); break;
case u16: output = convert<ushort, float , isRGB2GRAY>(in, r, g, b); break;
case s16: output = convert<short , float , isRGB2GRAY>(in, r, g, b); break;
case u8: output = convert<uchar , float , isRGB2GRAY>(in, r, g, b); break;
default: TYPE_ERROR(1, iType); break;
}
std::swap(*out, output);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_sort(af_array *out, const af_array in, const unsigned dim, const bool dir)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
DIM_ASSERT(1, info.elements() > 0);
// Only Dim 0 supported
ARG_ASSERT(2, dim == 0);
af_array val;
switch(type) {
case f32: val = sort<float >(in, dim, dir); break;
case f64: val = sort<double >(in, dim, dir); break;
case s32: val = sort<int >(in, dim, dir); break;
case u32: val = sort<uint >(in, dim, dir); break;
case u8: val = sort<uchar >(in, dim, dir); break;
// case s8: val = sort<char >(in, dim, dir); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out, val);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_reorder(af_array *out, const af_array in, const af::dim4 &rdims)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
DIM_ASSERT(1, info.elements() > 0);
// Check that dimensions are not repeated
// allDims is to check if all dimensions are there exactly once
// If all dimensions are present, the allDims will be -1, -1, -1, -1
// after the loop
// Example:
// rdims = {2, 0, 3, 1}
// i = 0 => 2 found and cond is true so alldims[2] = -1
// i = 1 => 0 found and cond is true so alldims[0] = -1
// i = 2 => 3 found and cond is true so alldims[3] = -1
// i = 3 => 1 found and cond is true so alldims[1] = -1
// rdims = {2, 0, 3, 2} // Failure case
// i = 3 => 2 found so cond is false (since alldims[2] = -1 when i = 0) so failed.
int allDims[] = {0, 1, 2, 3};
for(int i = 0; i < 4; i++) {
DIM_ASSERT(i + 2, rdims[i] == allDims[rdims[i]]);
allDims[rdims[i]] = -1;
}
// If reorder is a (batched) transpose, then call transpose
if(info.dims()[3] == 1) {
if(rdims[0] == 1 && rdims[1] == 0 &&
rdims[2] == 2 && rdims[3] == 3) {
return af_transpose(out, in, false);
}
}
af_array output;
switch(type) {
case f32: output = reorder<float >(in, rdims); break;
case c32: output = reorder<cfloat >(in, rdims); break;
case f64: output = reorder<double >(in, rdims); break;
case c64: output = reorder<cdouble>(in, rdims); break;
case b8: output = reorder<char >(in, rdims); break;
case s32: output = reorder<int >(in, rdims); break;
case u32: output = reorder<uint >(in, rdims); break;
case u8: output = reorder<uchar >(in, rdims); break;
case s64: output = reorder<intl >(in, rdims); break;
case u64: output = reorder<uintl >(in, rdims); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out,output);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_topk(af_array *values, af_array *indices, const af_array in,
const int k, const int dim, const af_topk_function order) {
try {
af::topkFunction ord = (order == AF_TOPK_DEFAULT ? AF_TOPK_MAX : order);
ArrayInfo inInfo = getInfo(in);
ARG_ASSERT(2, (inInfo.ndims() > 0));
if (inInfo.elements() == 1) {
dim_t dims[1] = {1};
af_err errValue = af_constant(indices, 0, 1, dims, u32);
return errValue == AF_SUCCESS ? af_retain_array(values, in)
: errValue;
}
int rdim = dim;
auto &inDims = inInfo.dims();
if (rdim == -1) {
for (dim_t d = 0; d < 4; d++) {
if (inDims[d] > 1) {
rdim = d;
break;
}
}
}
ARG_ASSERT(2, (inInfo.dims()[rdim] >= k));
ARG_ASSERT(4, (k <= 256)); // TODO(umar): Remove this limitation
if (rdim != 0)
AF_ERROR("topk is supported along dimenion 0 only.",
AF_ERR_NOT_SUPPORTED);
af_dtype type = inInfo.getType();
switch (type) {
// TODO(umar): FIX RETURN VALUES HERE
case f32: topk<float>(values, indices, in, k, rdim, ord); break;
case f64: topk<double>(values, indices, in, k, rdim, ord); break;
case u32: topk<uint>(values, indices, in, k, rdim, ord); break;
case s32: topk<int>(values, indices, in, k, rdim, ord); break;
default: TYPE_ERROR(1, type);
}
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
const af_array keys, const af_array values,
const unsigned dim, const bool isAscending)
{
try {
ArrayInfo kinfo = getInfo(keys);
af_dtype ktype = kinfo.getType();
ArrayInfo vinfo = getInfo(values);
DIM_ASSERT(4, kinfo.dims() == vinfo.dims());
if(kinfo.elements() == 0) {
dim_t my_dims[] = {0, 0, 0, 0};
AF_CHECK(af_create_handle(out_keys, AF_MAX_DIMS, my_dims, ktype));
AF_CHECK(af_create_handle(out_values, AF_MAX_DIMS, my_dims, ktype));
return AF_SUCCESS;
}
TYPE_ASSERT(kinfo.isReal());
af_array oKey;
af_array oVal;
switch(ktype) {
case f32: sort_by_key_tmplt<float >(&oKey, &oVal, keys, values, dim, isAscending); break;
case f64: sort_by_key_tmplt<double >(&oKey, &oVal, keys, values, dim, isAscending); break;
case s32: sort_by_key_tmplt<int >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u32: sort_by_key_tmplt<uint >(&oKey, &oVal, keys, values, dim, isAscending); break;
case s16: sort_by_key_tmplt<short >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u16: sort_by_key_tmplt<ushort >(&oKey, &oVal, keys, values, dim, isAscending); break;
case s64: sort_by_key_tmplt<intl >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u64: sort_by_key_tmplt<uintl >(&oKey, &oVal, keys, values, dim, isAscending); break;
case u8: sort_by_key_tmplt<uchar >(&oKey, &oVal, keys, values, dim, isAscending); break;
case b8: sort_by_key_tmplt<char >(&oKey, &oVal, keys, values, dim, isAscending); break;
default: TYPE_ERROR(1, ktype);
}
std::swap(*out_keys , oKey);
std::swap(*out_values , oVal);
}
CATCHALL;
return AF_SUCCESS;
}
af_err af_reorder(af_array *out, const af_array in, const af::dim4 &rdims)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
DIM_ASSERT(1, info.elements() > 0);
// Check that dimensions are not repeated
int allDims[] = {0, 1, 2, 3};
for(int i = 0; i < 3; i++) {
DIM_ASSERT(i + 2, rdims[i] == allDims[rdims[i]]);
allDims[rdims[i]] = -1;
}
// If reorder is a (batched) transpose, then call transpose
if(info.dims()[3] == 1) {
if(rdims[0] == 1 && rdims[1] == 0 &&
rdims[2] == 2 && rdims[3] == 3) {
return af_transpose(out, in);
}
}
af_array output;
switch(type) {
case f32: output = reorder<float >(in, rdims); break;
case c32: output = reorder<cfloat >(in, rdims); break;
case f64: output = reorder<double >(in, rdims); break;
case c64: output = reorder<cdouble>(in, rdims); break;
case b8: output = reorder<char >(in, rdims); break;
case s32: output = reorder<int >(in, rdims); break;
case u32: output = reorder<uint >(in, rdims); break;
case u8: output = reorder<uchar >(in, rdims); break;
case s8: output = reorder<char >(in, rdims); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out,output);
}
CATCHALL;
return AF_SUCCESS;
}
static void print(const char *exp, af_array arr, const int precision, std::ostream &os = std::cout, bool transpose = true)
{
if(exp == NULL) {
os << "No Name Array" << std::endl;
} else {
os << exp << std::endl;
}
const ArrayInfo info = getInfo(arr);
vector<T> data(info.elements());
af_array arrT;
if(transpose) {
AF_CHECK(af_reorder(&arrT, arr, 1, 0, 2, 3));
} else {
arrT = arr;
}
//FIXME: Use alternative function to avoid copies if possible
AF_CHECK(af_get_data_ptr(&data.front(), arrT));
const ArrayInfo infoT = getInfo(arrT);
if(transpose) {
AF_CHECK(af_release_array(arrT));
}
std::ios_base::fmtflags backup = os.flags();
os << "[" << info.dims() << "]\n";
#ifndef NDEBUG
os <<" Offsets: [" << info.offsets() << "]" << std::endl;
os <<" Strides: [" << info.strides() << "]" << std::endl;
#endif
printer(os, &data.front(), infoT, infoT.ndims() - 1, precision);
os.flags(backup);
}
af_err af_sort_index(af_array *out, af_array *indices, const af_array in, const unsigned dim, const bool isAscending)
{
try {
ArrayInfo info = getInfo(in);
af_dtype type = info.getType();
if(info.elements() <= 0) {
dim_t my_dims[] = {0, 0, 0, 0};
AF_CHECK(af_create_handle(out, AF_MAX_DIMS, my_dims, type));
AF_CHECK(af_create_handle(indices, AF_MAX_DIMS, my_dims, type));
return AF_SUCCESS;
}
af_array val;
af_array idx;
switch(type) {
case f32: sort_index<float >(&val, &idx, in, dim, isAscending); break;
case f64: sort_index<double >(&val, &idx, in, dim, isAscending); break;
case s32: sort_index<int >(&val, &idx, in, dim, isAscending); break;
case u32: sort_index<uint >(&val, &idx, in, dim, isAscending); break;
case s16: sort_index<short >(&val, &idx, in, dim, isAscending); break;
case u16: sort_index<ushort >(&val, &idx, in, dim, isAscending); break;
case s64: sort_index<intl >(&val, &idx, in, dim, isAscending); break;
case u64: sort_index<uintl >(&val, &idx, in, dim, isAscending); break;
case u8: sort_index<uchar >(&val, &idx, in, dim, isAscending); break;
case b8: sort_index<char >(&val, &idx, in, dim, isAscending); break;
default: TYPE_ERROR(1, type);
}
std::swap(*out , val);
std::swap(*indices, idx);
}
CATCHALL;
return AF_SUCCESS;
}
// Save an image to memory.
af_err af_save_image_memory(void **ptr, const af_array in_, const af_image_format format)
{
try {
FI_Init();
// set your own FreeImage error handler
FreeImage_SetOutputMessage(FreeImageErrorHandler);
// try to guess the file format from the file extension
FREE_IMAGE_FORMAT fif = (FREE_IMAGE_FORMAT)format;
if(fif == FIF_UNKNOWN || fif > 34) { // FreeImage FREE_IMAGE_FORMAT has upto 34 enums as of 3.17
AF_ERROR("FreeImage Error: Unknown Filetype", AF_ERR_NOT_SUPPORTED);
}
ArrayInfo info = getInfo(in_);
// check image color type
uint channels = info.dims()[2];
DIM_ASSERT(1, channels <= 4);
DIM_ASSERT(1, channels != 2);
int fi_bpp = channels * 8;
// sizes
uint fi_w = info.dims()[1];
uint fi_h = info.dims()[0];
// create the result image storage using FreeImage
FIBITMAP* pResultBitmap = FreeImage_Allocate(fi_w, fi_h, fi_bpp);
if(pResultBitmap == NULL) {
AF_ERROR("FreeImage Error: Error creating image or file", AF_ERR_RUNTIME);
}
// FI assumes [0-255]
// If array is in 0-1 range, multiply by 255
af_array in;
double max_real, max_imag;
bool free_in = false;
AF_CHECK(af_max_all(&max_real, &max_imag, in_));
if (max_real <= 1) {
af_array c255;
AF_CHECK(af_constant(&c255, 255.0, info.ndims(), info.dims().get(), f32));
AF_CHECK(af_mul(&in, in_, c255, false));
AF_CHECK(af_release_array(c255));
free_in = true;
} else {
in = in_;
}
// FI = row major | AF = column major
uint nDstPitch = FreeImage_GetPitch(pResultBitmap);
uchar* pDstLine = FreeImage_GetBits(pResultBitmap) + nDstPitch * (fi_h - 1);
af_array rr = 0, gg = 0, bb = 0, aa = 0;
AF_CHECK(channel_split(in, info.dims(), &rr, &gg, &bb, &aa)); // convert array to 3 channels if needed
uint step = channels; // force 3 channels saving
uint indx = 0;
af_array rrT = 0, ggT = 0, bbT = 0, aaT = 0;
if(channels == 4) {
AF_CHECK(af_transpose(&rrT, rr, false));
AF_CHECK(af_transpose(&ggT, gg, false));
AF_CHECK(af_transpose(&bbT, bb, false));
AF_CHECK(af_transpose(&aaT, aa, false));
ArrayInfo cinfo = getInfo(rrT);
float* pSrc0 = pinnedAlloc<float>(cinfo.elements());
float* pSrc1 = pinnedAlloc<float>(cinfo.elements());
float* pSrc2 = pinnedAlloc<float>(cinfo.elements());
float* pSrc3 = pinnedAlloc<float>(cinfo.elements());
AF_CHECK(af_get_data_ptr((void*)pSrc0, rrT));
AF_CHECK(af_get_data_ptr((void*)pSrc1, ggT));
AF_CHECK(af_get_data_ptr((void*)pSrc2, bbT));
AF_CHECK(af_get_data_ptr((void*)pSrc3, aaT));
// Copy the array into FreeImage buffer
for (uint y = 0; y < fi_h; ++y) {
for (uint x = 0; x < fi_w; ++x) {
*(pDstLine + x * step + 2) = (uchar) pSrc0[indx]; // b
*(pDstLine + x * step + 1) = (uchar) pSrc1[indx]; // g
*(pDstLine + x * step + 0) = (uchar) pSrc2[indx]; // r
*(pDstLine + x * step + 3) = (uchar) pSrc3[indx]; // a
++indx;
}
pDstLine -= nDstPitch;
}
pinnedFree(pSrc0);
pinnedFree(pSrc1);
pinnedFree(pSrc2);
pinnedFree(pSrc3);
} else if(channels == 3) {
AF_CHECK(af_transpose(&rrT, rr, false));
AF_CHECK(af_transpose(&ggT, gg, false));
AF_CHECK(af_transpose(&bbT, bb, false));
ArrayInfo cinfo = getInfo(rrT);
float* pSrc0 = pinnedAlloc<float>(cinfo.elements());
float* pSrc1 = pinnedAlloc<float>(cinfo.elements());
float* pSrc2 = pinnedAlloc<float>(cinfo.elements());
AF_CHECK(af_get_data_ptr((void*)pSrc0, rrT));
AF_CHECK(af_get_data_ptr((void*)pSrc1, ggT));
AF_CHECK(af_get_data_ptr((void*)pSrc2, bbT));
// Copy the array into FreeImage buffer
for (uint y = 0; y < fi_h; ++y) {
for (uint x = 0; x < fi_w; ++x) {
*(pDstLine + x * step + 2) = (uchar) pSrc0[indx]; // b
*(pDstLine + x * step + 1) = (uchar) pSrc1[indx]; // g
*(pDstLine + x * step + 0) = (uchar) pSrc2[indx]; // r
++indx;
}
pDstLine -= nDstPitch;
}
pinnedFree(pSrc0);
pinnedFree(pSrc1);
pinnedFree(pSrc2);
} else {
AF_CHECK(af_transpose(&rrT, rr, false));
ArrayInfo cinfo = getInfo(rrT);
float* pSrc0 = pinnedAlloc<float>(cinfo.elements());
AF_CHECK(af_get_data_ptr((void*)pSrc0, rrT));
for (uint y = 0; y < fi_h; ++y) {
for (uint x = 0; x < fi_w; ++x) {
*(pDstLine + x * step) = (uchar) pSrc0[indx];
++indx;
}
pDstLine -= nDstPitch;
}
pinnedFree(pSrc0);
}
FIMEMORY *stream = FreeImage_OpenMemory();
// now save the result image
if (!(FreeImage_SaveToMemory(fif, pResultBitmap, stream, 0) == TRUE)) {
AF_ERROR("FreeImage Error: Failed to save image", AF_ERR_RUNTIME);
}
*ptr = stream;
FreeImage_Unload(pResultBitmap);
if(free_in) AF_CHECK(af_release_array(in ));
if(rr != 0) AF_CHECK(af_release_array(rr ));
if(gg != 0) AF_CHECK(af_release_array(gg ));
if(bb != 0) AF_CHECK(af_release_array(bb ));
if(aa != 0) AF_CHECK(af_release_array(aa ));
if(rrT!= 0) AF_CHECK(af_release_array(rrT));
if(ggT!= 0) AF_CHECK(af_release_array(ggT));
if(bbT!= 0) AF_CHECK(af_release_array(bbT));
if(aaT!= 0) AF_CHECK(af_release_array(aaT));
} CATCHALL
return AF_SUCCESS;
}
af_err af_assign_gen(af_array *out,
const af_array lhs,
const dim_t ndims, const af_index_t* indexs,
const af_array rhs_)
{
af_array output = 0;
af_array rhs = rhs_;
// spanner is sequence index used for indexing along the
// dimensions after ndims
af_index_t spanner;
spanner.idx.seq = af_span;
spanner.isSeq = true;
try {
ARG_ASSERT(2, (ndims>0));
ARG_ASSERT(3, (indexs!=NULL));
int track = 0;
vector<af_seq> seqs(4, af_span);
for (dim_t i = 0; i < ndims; i++) {
if (indexs[i].isSeq) {
track++;
seqs[i] = indexs[i].idx.seq;
}
}
if (track==(int)ndims) {
// all indexs are sequences, redirecting to af_assign
return af_assign_seq(out, lhs, ndims, &(seqs.front()), rhs);
}
ARG_ASSERT(1, (lhs!=0));
ARG_ASSERT(4, (rhs!=0));
ArrayInfo lInfo = getInfo(lhs);
ArrayInfo rInfo = getInfo(rhs);
dim4 lhsDims = lInfo.dims();
dim4 rhsDims = rInfo.dims();
af_dtype lhsType= lInfo.getType();
af_dtype rhsType= rInfo.getType();
ARG_ASSERT(2, (ndims == 1) || (ndims == (dim_t)lInfo.ndims()));
if (ndims == 1 && ndims != (dim_t)lInfo.ndims()) {
af_array tmp_in = 0, tmp_out = 0;
AF_CHECK(af_flat(&tmp_in, lhs));
AF_CHECK(af_assign_gen(&tmp_out, tmp_in, ndims, indexs, rhs_));
AF_CHECK(af_moddims(out, tmp_out, lInfo.ndims(), lInfo.dims().get()));
AF_CHECK(af_release_array(tmp_in));
AF_CHECK(af_release_array(tmp_out));
return AF_SUCCESS;
}
ARG_ASSERT(1, (lhsType==rhsType));
ARG_ASSERT(3, (rhsDims.ndims()>0));
ARG_ASSERT(1, (lhsDims.ndims()>=rhsDims.ndims()));
ARG_ASSERT(2, (lhsDims.ndims()>=ndims));
if (*out != lhs) {
int count = 0;
AF_CHECK(af_get_data_ref_count(&count, lhs));
if (count > 1) {
AF_CHECK(af_copy_array(&output, lhs));
} else {
AF_CHECK(af_retain_array(&output, lhs));
}
} else {
output = lhs;
}
dim4 oDims = toDims(seqs, lhsDims);
// if af_array are indexs along any
// particular dimension, set the length of
// that dimension accordingly before any checks
for (dim_t i=0; i<ndims; i++) {
if (!indexs[i].isSeq) {
oDims[i] = getInfo(indexs[i].idx.arr).elements();
}
}
for (dim_t i = ndims; i < (dim_t)lInfo.ndims(); i++) {
oDims[i] = 1;
}
bool is_vector = true;
for (int i = 0; is_vector && i < oDims.ndims() - 1; i++) {
is_vector &= oDims[i] == 1;
}
//TODO: Move logic out of this
is_vector &= rInfo.isVector() || rInfo.isScalar();
if (is_vector) {
if (oDims.elements() != (dim_t)rInfo.elements() &&
rInfo.elements() != 1) {
AF_ERROR("Size mismatch between input and output", AF_ERR_SIZE);
}
if (rInfo.elements() == 1) {
AF_CHECK(af_tile(&rhs, rhs_, oDims[0], oDims[1], oDims[2], oDims[3]));
} else {
// If both out and rhs are vectors of equal elements, reshape rhs to out dims
AF_CHECK(af_moddims(&rhs, rhs_, oDims.ndims(), oDims.get()));
}
} else {
for (int i = 0; i < 4; i++) {
if (oDims[i] != rhsDims[i]) {
AF_ERROR("Size mismatch between input and output", AF_ERR_SIZE);
}
}
}
af_index_t idxrs[4];
// set all dimensions above ndims to spanner index
for (dim_t i=ndims; i<4; ++i) idxrs[i] = spanner;
for (dim_t i=0; i<ndims; ++i) {
if (!indexs[i].isSeq) {
// check if all af_arrays have atleast one value
// to enable indexing along that dimension
ArrayInfo idxInfo = getInfo(indexs[i].idx.arr);
af_dtype idxType = idxInfo.getType();
ARG_ASSERT(3, (idxType!=c32));
ARG_ASSERT(3, (idxType!=c64));
ARG_ASSERT(3, (idxType!=b8 ));
idxrs[i].idx.arr = indexs[i].idx.arr;
idxrs[i].isSeq = indexs[i].isSeq;
} else {
// af_seq is being used for this dimension
// just copy the index to local variable
idxrs[i] = indexs[i];
}
}
try {
switch(rhsType) {
case c64: genAssign<cdouble>(output, idxrs, rhs); break;
case f64: genAssign<double >(output, idxrs, rhs); break;
case c32: genAssign<cfloat >(output, idxrs, rhs); break;
case f32: genAssign<float >(output, idxrs, rhs); break;
case u64: genAssign<uintl >(output, idxrs, rhs); break;
case u32: genAssign<uint >(output, idxrs, rhs); break;
case s64: genAssign<intl >(output, idxrs, rhs); break;
case s32: genAssign<int >(output, idxrs, rhs); break;
case u8: genAssign<uchar >(output, idxrs, rhs); break;
case b8: genAssign<char >(output, idxrs, rhs); break;
default: TYPE_ERROR(1, rhsType);
}
} catch(...) {
if (*out != lhs) {
AF_CHECK(af_release_array(output));
if (is_vector) { AF_CHECK(af_release_array(rhs)); }
}
throw;
}
if (is_vector) { AF_CHECK(af_release_array(rhs)); }
}
CATCHALL;
std::swap(*out, output);
return AF_SUCCESS;
}
static void save_t(T* pDstLine, const af_array in, const dim4 dims, uint nDstPitch)
{
af_array rr = 0, gg = 0, bb = 0, aa = 0;
AF_CHECK(channel_split(in, dims, &rr, &gg, &bb, &aa)); // convert array to 3 channels if needed
af_array rrT = 0, ggT = 0, bbT = 0, aaT = 0;
T *pSrc0 = 0, *pSrc1 = 0, *pSrc2 = 0, *pSrc3 = 0;
uint step = channels; // force 3 channels saving
uint indx = 0;
AF_CHECK(af_transpose(&rrT, rr, false));
if(channels >= 3) AF_CHECK(af_transpose(&ggT, gg, false));
if(channels >= 3) AF_CHECK(af_transpose(&bbT, bb, false));
if(channels >= 4) AF_CHECK(af_transpose(&aaT, aa, false));
ArrayInfo cinfo = getInfo(rrT);
pSrc0 = pinnedAlloc<T>(cinfo.elements());
if(channels >= 3) pSrc1 = pinnedAlloc<T>(cinfo.elements());
if(channels >= 3) pSrc2 = pinnedAlloc<T>(cinfo.elements());
if(channels >= 4) pSrc3 = pinnedAlloc<T>(cinfo.elements());
AF_CHECK(af_get_data_ptr((void*)pSrc0, rrT));
if(channels >= 3) AF_CHECK(af_get_data_ptr((void*)pSrc1, ggT));
if(channels >= 3) AF_CHECK(af_get_data_ptr((void*)pSrc2, bbT));
if(channels >= 4) AF_CHECK(af_get_data_ptr((void*)pSrc3, aaT));
const uint fi_w = dims[1];
const uint fi_h = dims[0];
// Copy the array into FreeImage buffer
for (uint y = 0; y < fi_h; ++y) {
for (uint x = 0; x < fi_w; ++x) {
if(channels == 1) {
*(pDstLine + x * step + FI_RGBA_RED) = (T) pSrc0[indx]; // r -> 0
} else if(channels >=3) {
if((af_dtype) af::dtype_traits<T>::af_type == u8) {
*(pDstLine + x * step + FI_RGBA_BLUE) = (T) pSrc2[indx]; // b -> 0
*(pDstLine + x * step + FI_RGBA_GREEN) = (T) pSrc1[indx]; // g -> 1
*(pDstLine + x * step + FI_RGBA_RED) = (T) pSrc0[indx]; // r -> 2
} else {
// Non 8-bit types do not use ordering
// See Pixel Access Functions Chapter in FreeImage Doc
*(pDstLine + x * step + 0) = (T) pSrc0[indx]; // r -> 0
*(pDstLine + x * step + 1) = (T) pSrc1[indx]; // g -> 1
*(pDstLine + x * step + 2) = (T) pSrc2[indx]; // b -> 2
}
}
if(channels >= 4) *(pDstLine + x * step + FI_RGBA_ALPHA) = (T) pSrc3[indx]; // a
++indx;
}
pDstLine = (T*)(((uchar*)pDstLine) - nDstPitch);
}
pinnedFree(pSrc0);
if(channels >= 3) pinnedFree(pSrc1);
if(channels >= 3) pinnedFree(pSrc2);
if(channels >= 4) pinnedFree(pSrc3);
if(rr != 0) AF_CHECK(af_release_array(rr ));
if(gg != 0) AF_CHECK(af_release_array(gg ));
if(bb != 0) AF_CHECK(af_release_array(bb ));
if(aa != 0) AF_CHECK(af_release_array(aa ));
if(rrT!= 0) AF_CHECK(af_release_array(rrT));
if(ggT!= 0) AF_CHECK(af_release_array(ggT));
if(bbT!= 0) AF_CHECK(af_release_array(bbT));
if(aaT!= 0) AF_CHECK(af_release_array(aaT));
}
