experimental/ModelBuilder/MemRefUtils.h - 3p/openxla/iree - Git at Google

 // Copyright 2020 Google LLC
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //      https://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.

 // MemRefUtils.h
 // -----------------------------------------------------------------------------
 //
 // Utils for MLIR ABI interfacing with frameworks
 //
 // The templated free functions below make it possible to allocate dense
 // contiguous buffers with shapes that interoperate properly with the MLIR
 // codegen ABI.
 //
 // ```
 //  // 1. Compile and build a model, prepare the runner.
 //  ModelRunner runner = ...;
 //
 //  // 2. Allocate managed input and outputs with proper shapes and init value.
 //  auto inputLinearInit = [](unsigned idx, float *ptr) { *ptr = 0.032460f; };
 //  auto inputBuffer = makeInitializedStridedMemRefDescriptor<float, 2>(
 //     {B, W0}, inputLinearInit);
 //  auto outputLinearInit = [](unsigned idx, float *ptr) { *ptr = 0.0f; };
 //  auto outputBuffer = makeInitializedStridedMemRefDescriptor<float, 2>(
 //     {B, W3}, outputLinearInit);
 //
 //  // 3. Pack pointers to MLIR ABI compliant buffers and call the named func.
 //  void *packedArgs[2] = {&inputBuffer->descriptor, &outputBuffer->descriptor};
 //  runner.engine->invoke(funcName, llvm::MutableArrayRef<void *>{packedArgs});
 // ```

 #include <algorithm>
 #include <array>
 #include <cassert>
 #include <functional>
 #include <memory>

 #include "mlir/ExecutionEngine/CRunnerUtils.h"

 #ifndef IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_
 #define IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_

 namespace mlir {
 using AllocFunType = std::function<void *(size_t)>;

 namespace detail {

 // Given a shape with sizes greater than 0 along all dimensions,
 // returns the distance, in number of elements, between a slice in a dimension
 // and the next slice in the same dimension.
 //   e.g. shape[3, 4, 5] -> strides[20, 5, 1]
 template <size_t N>
 inline std::array<int64_t, N> makeStrides(const std::array<int64_t, N> &shape) {
   if (N == 0) return shape;
   std::array<int64_t, N> res;
   int64_t running = 1;
   for (int64_t idx = N - 1; idx >= 0; --idx) {
     assert(shape[idx] && "size must be nonnegatice for all shape dimensions");
     res[idx] = running;
     running *= shape[idx];
   }
   return res;
 }

 // Mallocs a StridedMemRefDescriptor<T, N>* that matches the MLIR ABI.
 // This is an implementation detail that is kept in sync with MLIR codegen
 // conventions.
 template <typename T, int N>
 typename std::enable_if<(N >= 1), StridedMemRefType<T, N> *>::type
 makeStridedMemRefDescriptor(void *ptr, void *alignedPtr,
                             const std::array<int64_t, N> &shape,
                             AllocFunType alloc = &::malloc) {
   StridedMemRefType<T, N> *descriptor = static_cast<StridedMemRefType<T, N> *>(
       alloc(sizeof(StridedMemRefType<T, N>)));
   descriptor->basePtr = static_cast<T *>(ptr);
   descriptor->data = static_cast<T *>(alignedPtr);
   descriptor->offset = 0;
   std::copy(shape.begin(), shape.end(), descriptor->sizes);
   auto strides = makeStrides<N>(shape);
   std::copy(strides.begin(), strides.end(), descriptor->strides);
   return descriptor;
 }

 // Mallocs a StridedMemRefDescriptor<T, 0>* (i.e. a pointer to scalar) that
 // matches the MLIR ABI. This is an implementation detail that is kept in sync
 // with MLIR codegen conventions.
 template <typename T, int N>
 typename std::enable_if<(N == 0), StridedMemRefType<T, 0> *>::type
 makeStridedMemRefDescriptor(void *ptr, void *alignedPtr,
                             const std::array<int64_t, N> &shape = {},
                             AllocFunType alloc = &::malloc) {
   StridedMemRefType<T, 0> *descriptor = static_cast<StridedMemRefType<T, 0> *>(
       alloc(sizeof(StridedMemRefType<T, 0>)));
   descriptor->basePtr = static_cast<T *>(ptr);
   descriptor->data = static_cast<T *>(alignedPtr);
   descriptor->offset = 0;
   return descriptor;
 }

 // Mallocs an UnrankedMemRefType<T>* that contains a ranked
 // StridedMemRefDescriptor<T, Rank>* and matches the MLIR ABI. This is an
 // implementation detail that is kept in sync with MLIR codegen conventions.
 template <typename T, int N>
 ::UnrankedMemRefType<T> *allocUnrankedDescriptor(
     void *data, void *alignedData, const std::array<int64_t, N> &shape,
     AllocFunType alloc = &::malloc) {
   ::UnrankedMemRefType<T> *res = static_cast<::UnrankedMemRefType<T> *>(
       alloc(sizeof(::UnrankedMemRefType<T>)));
   res->rank = N;
   res->descriptor = makeStridedMemRefDescriptor<T, N>(data, alignedData, shape);
   return res;
 }

 // Frees an UnrankedMemRefType<T>*
 template <typename T>
 void freeUnrankedDescriptor(::UnrankedMemRefType<T> *desc) {
   free(desc->descriptor);
   free(desc);
 }

 }  // namespace detail

 //===----------------------------------------------------------------------===//
 // Public API
 //===----------------------------------------------------------------------===//

 // Inefficient initializer called on each element during
 // `makeInitializedUnrankedDescriptor`. Takes the linear index and the shape so
 // that it can work in a generic fashion. The user can capture the shape and
 // delinearize if appropriate.
 template <typename T>
 using LinearInitializer = std::function<void(unsigned idx, T *ptr)>;

 inline uint32_t pow2msb(uint32_t val) {
   assert(val > 0);
   val--;
   val |= val >> 1;
   val |= val >> 2;
   val |= val >> 4;
   val |= val >> 8;
   val |= val >> 16;
   return val + 1;
 }

 // No such thing as a portable posize_memalign, roll our own.
 template <typename T>
 std::pair<void *, void *> allocAligned(size_t nElements,
                                        AllocFunType alloc = &::malloc) {
   assert(sizeof(T) < (1ul << 32) && "Elemental type overflows");
   auto size = nElements * sizeof(T);
   auto desiredAlignment = pow2msb(sizeof(T));
   assert((desiredAlignment & (desiredAlignment - 1)) == 0);
   assert(desiredAlignment >= sizeof(T));
   void *data = alloc(size + desiredAlignment);
   uintptr_t addr = reinterpret_cast<uintptr_t>(data);
   uintptr_t rem = addr % desiredAlignment;
   void *alignedData =
       (rem == 0) ? data
                  : reinterpret_cast<void *>(addr + (desiredAlignment - rem));
   assert(reinterpret_cast<uintptr_t>(alignedData) % desiredAlignment == 0);
   return std::pair<void *, void *>(data, alignedData);
 }

 // Entry point to allocate a dense buffer with a given `shape` and initializer
 // of type PointwiseInitializer. Can optionally take specific `alloc` and `free`
 // functions.
 template <typename T, int N, typename FreeFunType = decltype(&::free)>
 std::unique_ptr<::UnrankedMemRefType<float>, FreeFunType>
 makeInitializedUnrankedDescriptor(const std::array<int64_t, N> &shape,
                                   LinearInitializer<T> init,
                                   AllocFunType alloc = &::malloc,
                                   FreeFunType freeFun = &::free) {
   int64_t nElements = 1;
   for (int64_t s : shape) nElements *= s;
   auto allocated = allocAligned<T>(nElements, alloc);
   auto *data = static_cast<T *>(allocated.first);
   auto *alignedData = static_cast<T *>(allocated.second);
   for (unsigned i = 0; i < nElements; ++i) init(i, alignedData);
   return std::unique_ptr<::UnrankedMemRefType<float>, FreeFunType>(
       detail::allocUnrankedDescriptor<T, N>(data, alignedData, shape), freeFun);
 }

 // Entry point to allocate a dense buffer with a given `shape` and initializer
 // of type PointwiseInitializer. Can optionally take specific `alloc` and `free`
 // functions.
 template <typename T, int N, typename FreeFunType = decltype(&::free)>
 std::unique_ptr<StridedMemRefType<T, N>, FreeFunType>
 makeInitializedStridedMemRefDescriptor(const std::array<int64_t, N> &shape,
                                        LinearInitializer<T> init,
                                        AllocFunType alloc = &::malloc,
                                        FreeFunType freeFun = &::free) {
   int64_t nElements = 1;
   for (int64_t s : shape) nElements *= s;
   auto allocated = allocAligned<T>(nElements, alloc);
   auto *data = static_cast<T *>(allocated.first);
   auto *alignedData = static_cast<T *>(allocated.second);
   for (unsigned i = 0; i < nElements; ++i) init(i, alignedData);
   return std::unique_ptr<StridedMemRefType<T, N>, FreeFunType>(
       detail::makeStridedMemRefDescriptor<T, N>(data, alignedData, shape,
                                                 alloc),
       freeFun);
 }

 }  // namespace mlir

 #endif  // IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_
	// Copyright 2020 Google LLC
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// https://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.

	// MemRefUtils.h
	// -----------------------------------------------------------------------------
	//
	// Utils for MLIR ABI interfacing with frameworks
	//
	// The templated free functions below make it possible to allocate dense
	// contiguous buffers with shapes that interoperate properly with the MLIR
	// codegen ABI.
	//
	// ```
	// // 1. Compile and build a model, prepare the runner.
	// ModelRunner runner = ...;
	//
	// // 2. Allocate managed input and outputs with proper shapes and init value.
	// auto inputLinearInit = [](unsigned idx, float ptr) { ptr = 0.032460f; };
	// auto inputBuffer = makeInitializedStridedMemRefDescriptor<float, 2>(
	// {B, W0}, inputLinearInit);
	// auto outputLinearInit = [](unsigned idx, float ptr) { ptr = 0.0f; };
	// auto outputBuffer = makeInitializedStridedMemRefDescriptor<float, 2>(
	// {B, W3}, outputLinearInit);
	//
	// // 3. Pack pointers to MLIR ABI compliant buffers and call the named func.
	// void *packedArgs[2] = {&inputBuffer->descriptor, &outputBuffer->descriptor};
	// runner.engine->invoke(funcName, llvm::MutableArrayRef<void *>{packedArgs});
	// ```

	#include <algorithm>
	#include <array>
	#include <cassert>
	#include <functional>
	#include <memory>

	#include "mlir/ExecutionEngine/CRunnerUtils.h"

	#ifndef IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_
	#define IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_

	namespace mlir {
	using AllocFunType = std::function<void *(size_t)>;

	namespace detail {

	// Given a shape with sizes greater than 0 along all dimensions,
	// returns the distance, in number of elements, between a slice in a dimension
	// and the next slice in the same dimension.
	// e.g. shape[3, 4, 5] -> strides[20, 5, 1]
	template <size_t N>
	inline std::array<int64_t, N> makeStrides(const std::array<int64_t, N> &shape) {
	if (N == 0) return shape;
	std::array<int64_t, N> res;
	int64_t running = 1;
	for (int64_t idx = N - 1; idx >= 0; --idx) {
	assert(shape[idx] && "size must be nonnegatice for all shape dimensions");
	res[idx] = running;
	running *= shape[idx];
	}
	return res;
	}

	// Mallocs a StridedMemRefDescriptor<T, N>* that matches the MLIR ABI.
	// This is an implementation detail that is kept in sync with MLIR codegen
	// conventions.
	template <typename T, int N>
	typename std::enable_if<(N >= 1), StridedMemRefType<T, N> *>::type
	makeStridedMemRefDescriptor(void ptr, void alignedPtr,
	const std::array<int64_t, N> &shape,
	AllocFunType alloc = &::malloc) {
	StridedMemRefType<T, N> descriptor = static_cast<StridedMemRefType<T, N> >(
	alloc(sizeof(StridedMemRefType<T, N>)));
	descriptor->basePtr = static_cast<T *>(ptr);
	descriptor->data = static_cast<T *>(alignedPtr);
	descriptor->offset = 0;
	std::copy(shape.begin(), shape.end(), descriptor->sizes);
	auto strides = makeStrides<N>(shape);
	std::copy(strides.begin(), strides.end(), descriptor->strides);
	return descriptor;
	}

	// Mallocs a StridedMemRefDescriptor<T, 0>* (i.e. a pointer to scalar) that
	// matches the MLIR ABI. This is an implementation detail that is kept in sync
	// with MLIR codegen conventions.
	template <typename T, int N>
	typename std::enable_if<(N == 0), StridedMemRefType<T, 0> *>::type
	makeStridedMemRefDescriptor(void ptr, void alignedPtr,
	const std::array<int64_t, N> &shape = {},
	AllocFunType alloc = &::malloc) {
	StridedMemRefType<T, 0> descriptor = static_cast<StridedMemRefType<T, 0> >(
	alloc(sizeof(StridedMemRefType<T, 0>)));
	descriptor->basePtr = static_cast<T *>(ptr);
	descriptor->data = static_cast<T *>(alignedPtr);
	descriptor->offset = 0;
	return descriptor;
	}

	// Mallocs an UnrankedMemRefType<T>* that contains a ranked
	// StridedMemRefDescriptor<T, Rank>* and matches the MLIR ABI. This is an
	// implementation detail that is kept in sync with MLIR codegen conventions.
	template <typename T, int N>
	::UnrankedMemRefType<T> *allocUnrankedDescriptor(
	void data, void alignedData, const std::array<int64_t, N> &shape,
	AllocFunType alloc = &::malloc) {
	::UnrankedMemRefType<T> res = static_cast<::UnrankedMemRefType<T> >(
	alloc(sizeof(::UnrankedMemRefType<T>)));
	res->rank = N;
	res->descriptor = makeStridedMemRefDescriptor<T, N>(data, alignedData, shape);
	return res;
	}

	// Frees an UnrankedMemRefType<T>*
	template <typename T>
	void freeUnrankedDescriptor(::UnrankedMemRefType<T> *desc) {
	free(desc->descriptor);
	free(desc);
	}

	} // namespace detail

	//===----------------------------------------------------------------------===//
	// Public API
	//===----------------------------------------------------------------------===//

	// Inefficient initializer called on each element during
	// `makeInitializedUnrankedDescriptor`. Takes the linear index and the shape so
	// that it can work in a generic fashion. The user can capture the shape and
	// delinearize if appropriate.
	template <typename T>
	using LinearInitializer = std::function<void(unsigned idx, T *ptr)>;

	inline uint32_t pow2msb(uint32_t val) {
	assert(val > 0);
	val--;
	val \|= val >> 1;
	val \|= val >> 2;
	val \|= val >> 4;
	val \|= val >> 8;
	val \|= val >> 16;
	return val + 1;
	}

	// No such thing as a portable posize_memalign, roll our own.
	template <typename T>
	std::pair<void , void > allocAligned(size_t nElements,
	AllocFunType alloc = &::malloc) {
	assert(sizeof(T) < (1ul << 32) && "Elemental type overflows");
	auto size = nElements * sizeof(T);
	auto desiredAlignment = pow2msb(sizeof(T));
	assert((desiredAlignment & (desiredAlignment - 1)) == 0);
	assert(desiredAlignment >= sizeof(T));
	void *data = alloc(size + desiredAlignment);
	uintptr_t addr = reinterpret_cast<uintptr_t>(data);
	uintptr_t rem = addr % desiredAlignment;
	void *alignedData =
	(rem == 0) ? data
	: reinterpret_cast<void *>(addr + (desiredAlignment - rem));
	assert(reinterpret_cast<uintptr_t>(alignedData) % desiredAlignment == 0);
	return std::pair<void , void >(data, alignedData);
	}

	// Entry point to allocate a dense buffer with a given `shape` and initializer
	// of type PointwiseInitializer. Can optionally take specific `alloc` and `free`
	// functions.
	template <typename T, int N, typename FreeFunType = decltype(&::free)>
	std::unique_ptr<::UnrankedMemRefType<float>, FreeFunType>
	makeInitializedUnrankedDescriptor(const std::array<int64_t, N> &shape,
	LinearInitializer<T> init,
	AllocFunType alloc = &::malloc,
	FreeFunType freeFun = &::free) {
	int64_t nElements = 1;
	for (int64_t s : shape) nElements *= s;
	auto allocated = allocAligned<T>(nElements, alloc);
	auto data = static_cast<T >(allocated.first);
	auto alignedData = static_cast<T >(allocated.second);
	for (unsigned i = 0; i < nElements; ++i) init(i, alignedData);
	return std::unique_ptr<::UnrankedMemRefType<float>, FreeFunType>(
	detail::allocUnrankedDescriptor<T, N>(data, alignedData, shape), freeFun);
	}

	// Entry point to allocate a dense buffer with a given `shape` and initializer
	// of type PointwiseInitializer. Can optionally take specific `alloc` and `free`
	// functions.
	template <typename T, int N, typename FreeFunType = decltype(&::free)>
	std::unique_ptr<StridedMemRefType<T, N>, FreeFunType>
	makeInitializedStridedMemRefDescriptor(const std::array<int64_t, N> &shape,
	LinearInitializer<T> init,
	AllocFunType alloc = &::malloc,
	FreeFunType freeFun = &::free) {
	int64_t nElements = 1;
	for (int64_t s : shape) nElements *= s;
	auto allocated = allocAligned<T>(nElements, alloc);
	auto data = static_cast<T >(allocated.first);
	auto alignedData = static_cast<T >(allocated.second);
	for (unsigned i = 0; i < nElements; ++i) init(i, alignedData);
	return std::unique_ptr<StridedMemRefType<T, N>, FreeFunType>(
	detail::makeStridedMemRefDescriptor<T, N>(data, alignedData, shape,
	alloc),
	freeFun);
	}

	} // namespace mlir

	#endif // IREE_EXPERIMENTAL_MODELBUILDER_MEMREFUTILS_H_