1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "mlir-hlo/utils/codegen_utils.h"
17
18 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
19 #include "mlir/Dialect/Func/IR/FuncOps.h"
20 #include "mlir/Dialect/MemRef/IR/MemRef.h"
21 #include "mlir/IR/BuiltinTypes.h"
22 #include "mlir/IR/Location.h"
23 #include "mlir/Pass/Pass.h"
24
25 using llvm::SmallVector;
26
27 namespace mlir {
28 namespace codegen_utils {
29
emitNumElementsComputation(OpBuilder & b,Location loc,Value memref)30 Value emitNumElementsComputation(OpBuilder& b, Location loc, Value memref) {
31 int rank = memref.getType().cast<MemRefType>().getRank();
32 Value numElements;
33 numElements = b.create<mlir::arith::ConstantOp>(
34 loc, b.getIndexType(), b.getIntegerAttr(b.getIndexType(), 1));
35 for (int r = 0; r < rank; ++r) {
36 auto dimSize = b.create<memref::DimOp>(loc, memref, r);
37 numElements = b.create<arith::MulIOp>(loc, numElements, dimSize);
38 }
39 return numElements;
40 }
41
emitNumElementsComputation(OpBuilder & b,Location loc,Operation * op)42 Value emitNumElementsComputation(OpBuilder& b, Location loc, Operation* op) {
43 // only const rank is supported for now
44 assert(op->getDialect()->getNamespace() == "lmhlo");
45 int numOperands = op->getNumOperands();
46 Value resultMemref = op->getOperand(numOperands - 1);
47 return emitNumElementsComputation(b, loc, resultMemref);
48 }
49
calcMultiDimIndex(OpBuilder & b,Location loc,Value linearIndex,ArrayRef<Value> shape)50 SmallVector<Value> calcMultiDimIndex(OpBuilder& b, Location loc,
51 Value linearIndex, ArrayRef<Value> shape) {
52 int rank = shape.size();
53 SmallVector<Value> result;
54 if (rank == 0) return result;
55 if (rank == 1) {
56 result.push_back(linearIndex);
57 return result;
58 }
59
60 // dim_acc_mul_vec = [d, c*d, b*c*d]
61 SmallVector<Value> dimAccMulVec;
62 Value tmpAccMul = shape[rank - 1];
63 dimAccMulVec.emplace_back(tmpAccMul);
64 for (int i = rank - 2; i > 0; --i) {
65 tmpAccMul = b.create<arith::MulIOp>(loc, tmpAccMul, shape[i]);
66 dimAccMulVec.emplace_back(tmpAccMul);
67 }
68 Value blockIndex = linearIndex;
69 for (int i = 0; i < rank; ++i) {
70 Value index;
71 if (i == rank - 1) {
72 index = blockIndex;
73 } else {
74 index = b.create<arith::DivUIOp>(loc, blockIndex, dimAccMulVec.back());
75 blockIndex =
76 b.create<arith::RemUIOp>(loc, blockIndex, dimAccMulVec.back());
77 dimAccMulVec.pop_back();
78 }
79 result.push_back(index);
80 }
81 return result;
82 }
83
calcMultiDimIndex(OpBuilder & b,Location loc,Value linearIndex,Value memref)84 SmallVector<Value> calcMultiDimIndex(OpBuilder& b, Location loc,
85 Value linearIndex, Value memref) {
86 int rank = memref.getType().cast<MemRefType>().getRank();
87 SmallVector<Value> result;
88 if (rank == 0) return result;
89 if (rank == 1) {
90 result.push_back(linearIndex);
91 return result;
92 }
93 // shape = [a, b, c, d]
94 SmallVector<Value, 4> shapeVec;
95 for (int i = 0; i < rank; ++i) {
96 shapeVec.push_back(b.create<memref::DimOp>(loc, memref, i));
97 }
98
99 return calcMultiDimIndex(b, loc, linearIndex, shapeVec);
100 }
101
calcMultiDimIndexForFirstOperand(OpBuilder & b,Location loc,Value linearIndex,Operation * op)102 SmallVector<Value> calcMultiDimIndexForFirstOperand(OpBuilder& b, Location loc,
103 Value linearIndex,
104 Operation* op) {
105 assert(op->getDialect()->getNamespace() == "lmhlo");
106 Value operandMemref = op->getOperand(0);
107 return calcMultiDimIndex(b, loc, linearIndex, operandMemref);
108 }
109
110 } // namespace codegen_utils
111 } // namespace mlir
112