compiler/Translation/SPIRV/IndexComputation.cpp - 3p/openxla/iree - Git at Google

 // Copyright 2019 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.

 //===- IndexComputation.cpp ------------------------------------*- C++//-*-===//
 //
 // For an IREE dispatch function, compute the map from workitem ID to index of
 // tensor computed within that workitem.
 //
 //===----------------------------------------------------------------------===//
 #include "compiler/Translation/SPIRV/IndexComputation.h"

 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"

 static llvm::cl::opt<bool> doAffineExprSimplify(
     "simplify-spirv-affine-exprs",
     llvm::cl::desc("Simplify affine expressions during code-generation."),
     llvm::cl::init(true));

 namespace mlir {
 namespace iree_compiler {

 //===----------------------------------------------------------------------===//
 // Reshape Utility Functions
 //===----------------------------------------------------------------------===//

 namespace {
 /// Handles shapes for scalars. Shape of scalars are represented as empty vetor,
 /// i.e. {}. Its easier to do index propogation to handle the scalar as vector
 /// of size 1.
 inline SmallVector<int64_t, 4> handleIfScalar(ArrayRef<int64_t> shape) {
   SmallVector<int64_t, 4> resultShape;
   if (shape.empty()) {
     return {1};
   }
   return SmallVector<int64_t, 4>(shape.begin(), shape.end());
 }

 /// Reshapes are often used to either add a dimension of size 1 or remove a
 /// dimension of size 1. Recognizing such cases can make the code-generation
 /// easier. The AffineMap needs to either add a constant 0 in the range for such
 /// added dimensions or drop those dimensions.
 inline LogicalResult getAffineExprForAddOrRemoveDimension(
     Builder &builder, ArrayRef<AffineExpr> resultExprs,
     ArrayRef<int64_t> resultShape, ArrayRef<int64_t> operandShape,
     SmallVectorImpl<AffineExpr> &operandExprs) {
   auto resultIndex = resultShape.size();
   auto operandIndex = operandShape.size();
   operandExprs.resize(operandShape.size());
   // Try to match up the dimensions of the operand and result by ignoring any
   // dimensions of size of 1 that are introduced.
   while (resultIndex > 0 && operandIndex > 0) {
     if (resultShape[resultIndex - 1] == -1 ||
         operandShape[operandIndex - 1] == -1) {
       return failure();
     }
     if (resultShape[resultIndex - 1] == operandShape[operandIndex - 1]) {
       operandExprs[operandIndex - 1] = resultExprs[resultIndex - 1];
       resultIndex--;
       operandIndex--;
       continue;
     }
     if (resultShape[resultIndex - 1] == 1) {
       // This is a dimension that is added on the operand. This affine
       // expression corresponding to this dimension is dropped.
       resultIndex--;
       continue;
     }
     if (operandShape[operandIndex - 1] == 1) {
       // This is a dimension of size 1 of the operand that is dropped. Add a
       // constant expr 0.
       operandExprs[operandIndex - 1] = builder.getAffineConstantExpr(0);
       operandIndex--;
       continue;
     }
     return failure();
   }
   // Any remaining dimensions should be 1.
   while (resultIndex > 0) {
     if (resultShape[resultIndex - 1] != 1) {
       return failure();
     }
     resultIndex--;
   }
   while (operandIndex > 0) {
     if (operandShape[operandIndex - 1] != 1) {
       return failure();
     }
     // This is a dimension of size 1 that is dropped. Add a constant expression
     // 0.
     operandExprs[operandIndex - 1] = builder.getAffineConstantExpr(0);
     operandIndex--;
   }
   return success();
 }

 /// Constructs the strides of an array assuming a row-major packed layout.
 // TODO(ravishankarm): This assumes the shape are static. When using dynamic
 // shapes, parameters of each dimension can be used to construct AffineExpr for
 // strides along each dimension. Note that multiplying two symbolic constants is
 // technically not affine, but you could use another symbol to represent the
 // product, so it should be still representable as affine exprs.
 inline LogicalResult getRowMajorPackedStrides(
     Builder &builder, ArrayRef<int64_t> shape,
     SmallVectorImpl<AffineExpr> &strides) {
   strides.resize(shape.size());
   int64_t stride = 1;
   for (auto dim : enumerate(reverse(shape))) {
     if (dim.value() < 0) {
       // TODO(ravishankarm) : Better error message.
       return failure();
     }
     strides[shape.size() - 1 - dim.index()] =
         builder.getAffineConstantExpr(stride);
     stride *= dim.value();
   }
   return success();
 }

 /// Linearizes the index of the result position accessed using the shape of the
 /// result tensor and delinearizes it to get the position of the operand.
 inline LogicalResult getAffineExprForReshape(
     Builder &builder, unsigned numDims, unsigned numSymbols,
     ArrayRef<AffineExpr> resultExprs, ArrayRef<int64_t> resultShape,
     ArrayRef<int64_t> operandShape, SmallVectorImpl<AffineExpr> &operandExprs) {
   // To linearize the index, assume that the memory is laid out in
   // packed-row-major layout based on the shape.
   // TODO(ravishankarm) : When there is stride information, use that to map from
   // index to memory location.
   SmallVector<AffineExpr, 4> resultStrides;
   if (failed(getRowMajorPackedStrides(builder, resultShape, resultStrides))) {
     return failure();
   }
   AffineExpr linearizedExpr;
   for (auto index : enumerate(resultExprs)) {
     auto val = getAffineBinaryOpExpr(AffineExprKind::Mul, index.value(),
                                      resultStrides[index.index()]);
     if (doAffineExprSimplify) {
       val = simplifyAffineExpr(val, numDims, numSymbols);
     }
     linearizedExpr = (index.index() ? getAffineBinaryOpExpr(AffineExprKind::Add,
                                                             linearizedExpr, val)
                                     : val);
     if (doAffineExprSimplify) {
       linearizedExpr = simplifyAffineExpr(val, numDims, numSymbols);
     }
   }

   // Unlinearize the index, assuming row-major-packed layout.
   // TODO(ravishankarm) : When there is stride information, use that to map from
   // memory location to index.
   SmallVector<AffineExpr, 4> operandStrides;
   if (failed(getRowMajorPackedStrides(builder, operandShape, operandStrides))) {
     return failure();
   }
   operandExprs.resize(operandStrides.size());
   for (auto stride : enumerate(operandStrides)) {
     if (stride.index() == operandStrides.size() - 1) {
       operandExprs[stride.index()] = linearizedExpr;
       break;
     }
     auto expr = getAffineBinaryOpExpr(AffineExprKind::FloorDiv, linearizedExpr,
                                       stride.value());
     operandExprs[stride.index()] =
         (doAffineExprSimplify ? simplifyAffineExpr(expr, numDims, numSymbols)
                               : expr);

     linearizedExpr = getAffineBinaryOpExpr(AffineExprKind::Mod, linearizedExpr,
                                            stride.value());
     if (doAffineExprSimplify) {
       linearizedExpr = simplifyAffineExpr(linearizedExpr, numDims, numSymbols);
     }
   }
   return success();
 }
 }  // namespace

 LogicalResult getReshapeOperandMap(Builder &builder, AffineMap resultIndexMap,
                                    ArrayRef<int64_t> resultShapeRef,
                                    ArrayRef<int64_t> operandShapeRef,
                                    AffineMap &operandIndexMap) {
   auto resultShape = handleIfScalar(resultShapeRef);
   auto operandShape = handleIfScalar(operandShapeRef);
   auto resultExprs = resultIndexMap.getResults();
   assert(resultShape.size() == resultExprs.size() &&
          "Ranks of the Domain of index map and result must be the same");
   SmallVector<AffineExpr, 4> operandExprs;
   if (failed(getAffineExprForAddOrRemoveDimension(
           builder, resultExprs, resultShape, operandShape, operandExprs)) &&
       failed(getAffineExprForReshape(
           builder, resultIndexMap.getNumDims(), resultIndexMap.getNumSymbols(),
           resultExprs, resultShape, operandShape, operandExprs))) {
     return failure();
   }
   assert(operandExprs.size() == operandShape.size() &&
          "expected as many exprs for the operand as the rank of the operand");
   operandIndexMap =
       AffineMap::get(resultIndexMap.getNumDims(),
                      resultIndexMap.getNumSymbols(), operandExprs);

   return success();
 }

 LogicalResult IndexPropagation::propagateIndexMap(
     Operation *op, IndexComputationCache &indexMap) const {
   if (op->getNumResults() == 0) {
     // Nothing to do for this op.
     return success();
   }
   if (op->getNumResults() != 1) {
     return op->emitError(
         "default index propagation handles case with a single-return value");
   }
   // Initialize the storage for all the operands.
   for (auto arg : op->getOperands()) {
     indexMap[arg];
   }
   for (auto &resultIndexMap : indexMap[op->getResult(0)]) {
     SmallVector<AffineMap, 4> operandIndices;
     if (failed(this->propagateIndexMap(op, resultIndexMap.first,
                                        operandIndices))) {
       return failure();
     }
     assert(operandIndices.size() == op->getNumOperands() &&
            "Expected as many indices as operands");
     for (auto arg : enumerate(op->getOperands())) {
       indexMap[arg.value()][operandIndices[arg.index()]];
       resultIndexMap.second.push_back(operandIndices[arg.index()]);
     }
   }
   return success();
 }

 void dumpIndexCache(IndexComputationCache &indexMap) {
   for (auto &el : indexMap) {
     // llvm::errs() << "Value : " << *(el.first);
     // llvm::errs().flush();
     if (isa<OpResult>(el.first)) {
       llvm::errs() << "Operation : " << el.first->getDefiningOp()->getName();
     } else if (isa<BlockArgument>(el.first)) {
       llvm::errs() << "BlockArgument";
     }
     for (auto &used : el.second) {
       llvm::errs() << "\n\t" << used.first << " : [";
       std::string sep = "";
       for (auto &operand : used.second) {
         llvm::errs() << sep << operand;
         sep = ", ";
       }
       llvm::errs() << "]";
     }
     llvm::errs() << "\n";
   }
   llvm::errs() << "\n";
 }

 }  // namespace iree_compiler
 }  // namespace mlir
	// Copyright 2019 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.

	//===- IndexComputation.cpp ------------------------------------- C++//--===//
	//
	// For an IREE dispatch function, compute the map from workitem ID to index of
	// tensor computed within that workitem.
	//
	//===----------------------------------------------------------------------===//
	#include "compiler/Translation/SPIRV/IndexComputation.h"

	#include "llvm/Support/CommandLine.h"
	#include "llvm/Support/raw_ostream.h"

	static llvm::cl::opt<bool> doAffineExprSimplify(
	"simplify-spirv-affine-exprs",
	llvm::cl::desc("Simplify affine expressions during code-generation."),
	llvm::cl::init(true));

	namespace mlir {
	namespace iree_compiler {

	//===----------------------------------------------------------------------===//
	// Reshape Utility Functions
	//===----------------------------------------------------------------------===//

	namespace {
	/// Handles shapes for scalars. Shape of scalars are represented as empty vetor,
	/// i.e. {}. Its easier to do index propogation to handle the scalar as vector
	/// of size 1.
	inline SmallVector<int64_t, 4> handleIfScalar(ArrayRef<int64_t> shape) {
	SmallVector<int64_t, 4> resultShape;
	if (shape.empty()) {
	return {1};
	}
	return SmallVector<int64_t, 4>(shape.begin(), shape.end());
	}

	/// Reshapes are often used to either add a dimension of size 1 or remove a
	/// dimension of size 1. Recognizing such cases can make the code-generation
	/// easier. The AffineMap needs to either add a constant 0 in the range for such
	/// added dimensions or drop those dimensions.
	inline LogicalResult getAffineExprForAddOrRemoveDimension(
	Builder &builder, ArrayRef<AffineExpr> resultExprs,
	ArrayRef<int64_t> resultShape, ArrayRef<int64_t> operandShape,
	SmallVectorImpl<AffineExpr> &operandExprs) {
	auto resultIndex = resultShape.size();
	auto operandIndex = operandShape.size();
	operandExprs.resize(operandShape.size());
	// Try to match up the dimensions of the operand and result by ignoring any
	// dimensions of size of 1 that are introduced.
	while (resultIndex > 0 && operandIndex > 0) {
	if (resultShape[resultIndex - 1] == -1 \|\|
	operandShape[operandIndex - 1] == -1) {
	return failure();
	}
	if (resultShape[resultIndex - 1] == operandShape[operandIndex - 1]) {
	operandExprs[operandIndex - 1] = resultExprs[resultIndex - 1];
	resultIndex--;
	operandIndex--;
	continue;
	}
	if (resultShape[resultIndex - 1] == 1) {
	// This is a dimension that is added on the operand. This affine
	// expression corresponding to this dimension is dropped.
	resultIndex--;
	continue;
	}
	if (operandShape[operandIndex - 1] == 1) {
	// This is a dimension of size 1 of the operand that is dropped. Add a
	// constant expr 0.
	operandExprs[operandIndex - 1] = builder.getAffineConstantExpr(0);
	operandIndex--;
	continue;
	}
	return failure();
	}
	// Any remaining dimensions should be 1.
	while (resultIndex > 0) {
	if (resultShape[resultIndex - 1] != 1) {
	return failure();
	}
	resultIndex--;
	}
	while (operandIndex > 0) {
	if (operandShape[operandIndex - 1] != 1) {
	return failure();
	}
	// This is a dimension of size 1 that is dropped. Add a constant expression
	// 0.
	operandExprs[operandIndex - 1] = builder.getAffineConstantExpr(0);
	operandIndex--;
	}
	return success();
	}

	/// Constructs the strides of an array assuming a row-major packed layout.
	// TODO(ravishankarm): This assumes the shape are static. When using dynamic
	// shapes, parameters of each dimension can be used to construct AffineExpr for
	// strides along each dimension. Note that multiplying two symbolic constants is
	// technically not affine, but you could use another symbol to represent the
	// product, so it should be still representable as affine exprs.
	inline LogicalResult getRowMajorPackedStrides(
	Builder &builder, ArrayRef<int64_t> shape,
	SmallVectorImpl<AffineExpr> &strides) {
	strides.resize(shape.size());
	int64_t stride = 1;
	for (auto dim : enumerate(reverse(shape))) {
	if (dim.value() < 0) {
	// TODO(ravishankarm) : Better error message.
	return failure();
	}
	strides[shape.size() - 1 - dim.index()] =
	builder.getAffineConstantExpr(stride);
	stride *= dim.value();
	}
	return success();
	}

	/// Linearizes the index of the result position accessed using the shape of the
	/// result tensor and delinearizes it to get the position of the operand.
	inline LogicalResult getAffineExprForReshape(
	Builder &builder, unsigned numDims, unsigned numSymbols,
	ArrayRef<AffineExpr> resultExprs, ArrayRef<int64_t> resultShape,
	ArrayRef<int64_t> operandShape, SmallVectorImpl<AffineExpr> &operandExprs) {
	// To linearize the index, assume that the memory is laid out in
	// packed-row-major layout based on the shape.
	// TODO(ravishankarm) : When there is stride information, use that to map from
	// index to memory location.
	SmallVector<AffineExpr, 4> resultStrides;
	if (failed(getRowMajorPackedStrides(builder, resultShape, resultStrides))) {
	return failure();
	}
	AffineExpr linearizedExpr;
	for (auto index : enumerate(resultExprs)) {
	auto val = getAffineBinaryOpExpr(AffineExprKind::Mul, index.value(),
	resultStrides[index.index()]);
	if (doAffineExprSimplify) {
	val = simplifyAffineExpr(val, numDims, numSymbols);
	}
	linearizedExpr = (index.index() ? getAffineBinaryOpExpr(AffineExprKind::Add,
	linearizedExpr, val)
	: val);
	if (doAffineExprSimplify) {
	linearizedExpr = simplifyAffineExpr(val, numDims, numSymbols);
	}
	}

	// Unlinearize the index, assuming row-major-packed layout.
	// TODO(ravishankarm) : When there is stride information, use that to map from
	// memory location to index.
	SmallVector<AffineExpr, 4> operandStrides;
	if (failed(getRowMajorPackedStrides(builder, operandShape, operandStrides))) {
	return failure();
	}
	operandExprs.resize(operandStrides.size());
	for (auto stride : enumerate(operandStrides)) {
	if (stride.index() == operandStrides.size() - 1) {
	operandExprs[stride.index()] = linearizedExpr;
	break;
	}
	auto expr = getAffineBinaryOpExpr(AffineExprKind::FloorDiv, linearizedExpr,
	stride.value());
	operandExprs[stride.index()] =
	(doAffineExprSimplify ? simplifyAffineExpr(expr, numDims, numSymbols)
	: expr);

	linearizedExpr = getAffineBinaryOpExpr(AffineExprKind::Mod, linearizedExpr,
	stride.value());
	if (doAffineExprSimplify) {
	linearizedExpr = simplifyAffineExpr(linearizedExpr, numDims, numSymbols);
	}
	}
	return success();
	}
	} // namespace

	LogicalResult getReshapeOperandMap(Builder &builder, AffineMap resultIndexMap,
	ArrayRef<int64_t> resultShapeRef,
	ArrayRef<int64_t> operandShapeRef,
	AffineMap &operandIndexMap) {
	auto resultShape = handleIfScalar(resultShapeRef);
	auto operandShape = handleIfScalar(operandShapeRef);
	auto resultExprs = resultIndexMap.getResults();
	assert(resultShape.size() == resultExprs.size() &&
	"Ranks of the Domain of index map and result must be the same");
	SmallVector<AffineExpr, 4> operandExprs;
	if (failed(getAffineExprForAddOrRemoveDimension(
	builder, resultExprs, resultShape, operandShape, operandExprs)) &&
	failed(getAffineExprForReshape(
	builder, resultIndexMap.getNumDims(), resultIndexMap.getNumSymbols(),
	resultExprs, resultShape, operandShape, operandExprs))) {
	return failure();
	}
	assert(operandExprs.size() == operandShape.size() &&
	"expected as many exprs for the operand as the rank of the operand");
	operandIndexMap =
	AffineMap::get(resultIndexMap.getNumDims(),
	resultIndexMap.getNumSymbols(), operandExprs);

	return success();
	}

	LogicalResult IndexPropagation::propagateIndexMap(
	Operation *op, IndexComputationCache &indexMap) const {
	if (op->getNumResults() == 0) {
	// Nothing to do for this op.
	return success();
	}
	if (op->getNumResults() != 1) {
	return op->emitError(
	"default index propagation handles case with a single-return value");
	}
	// Initialize the storage for all the operands.
	for (auto arg : op->getOperands()) {
	indexMap[arg];
	}
	for (auto &resultIndexMap : indexMap[op->getResult(0)]) {
	SmallVector<AffineMap, 4> operandIndices;
	if (failed(this->propagateIndexMap(op, resultIndexMap.first,
	operandIndices))) {
	return failure();
	}
	assert(operandIndices.size() == op->getNumOperands() &&
	"Expected as many indices as operands");
	for (auto arg : enumerate(op->getOperands())) {
	indexMap[arg.value()][operandIndices[arg.index()]];
	resultIndexMap.second.push_back(operandIndices[arg.index()]);
	}
	}
	return success();
	}

	void dumpIndexCache(IndexComputationCache &indexMap) {
	for (auto &el : indexMap) {
	// llvm::errs() << "Value : " << *(el.first);
	// llvm::errs().flush();
	if (isa<OpResult>(el.first)) {
	llvm::errs() << "Operation : " << el.first->getDefiningOp()->getName();
	} else if (isa<BlockArgument>(el.first)) {
	llvm::errs() << "BlockArgument";
	}
	for (auto &used : el.second) {
	llvm::errs() << "\n\t" << used.first << " : [";
	std::string sep = "";
	for (auto &operand : used.second) {
	llvm::errs() << sep << operand;
	sep = ", ";
	}
	llvm::errs() << "]";
	}
	llvm::errs() << "\n";
	}
	llvm::errs() << "\n";
	}

	} // namespace iree_compiler
	} // namespace mlir