megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

/**
 * \file src/jit/impl/mlir/ir/common.h
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
 */

#pragma once

#include "megbrain_build_config.h"
#if MGB_JIT && MGB_JIT_MLIR

#include "megbrain/tensor.h"

#include <mlir/Dialect/StandardOps/IR/Ops.h>
#include <mlir/IR/OperationSupport.h>
#include <mlir/IR/Value.h>

namespace mgb {
namespace jit {

/**
 * \brief Helper function for common value builder
 */
class ValueBuilderHelper {
public:
    ValueBuilderHelper(mlir::OpBuilder& b, mlir::Location location)
            : m_builder{b}, m_location{location} {};

#define cb(name)                                                              \
    mlir::Value name(mlir::ValueRange operands) { return name(operands[0]); } \
    mlir::Value name(mlir::Value lhs)

    // unary functions
    cb(abs);
    cb(ceil);
    cb(cos);
    cb(exp);
    cb(exp2);
    cb(floor);
    cb(log);
    cb(log10);
    cb(log2);
    cb(neg);
    cb(rsqrt);
    cb(sin);
    cb(sqrt);
    cb(tanh);

#undef cb

#define cb(name)                                  \
    mlir::Value name(mlir::ValueRange operands) { \
        return name(operands[0], operands[1]);    \
    }                                             \
    mlir::Value name(mlir::Value lhs, mlir::Value rhs)

    // binary functions
    cb(add);
    cb(bit_and);
    cb(bit_or);
    cb(div);
    cb(divI);
    cb(eq);
    cb(ge);
    cb(gt);
    cb(le);
    cb(lt);
    cb(max);
    cb(min);
    cb(mod);
    cb(modI);
    cb(mul);
    cb(sub);

#undef cb

    // constant functions
    mlir::Value const_f32(float val);
    mlir::Value const_i32(int32_t val);

    // select function
    mlir::Value select(mlir::Value cond, mlir::Value true_val, mlir::Value false_val);

private:
    mlir::OpBuilder& m_builder;
    mlir::Location m_location;
};

template <typename Op>
mlir::Value get_operand(
        mlir::OpBuilder& builder, const mlir::Location& loc, const mlir::Value& val,
        const mlir::ValueRange& index) {
    if (val.getType().isa<mlir::MemRefType>()) {
        return builder.create<Op>(loc, val, index);
    } else {
        return val;
    }
}

mlir::AffineMap get_affinemap(
        mlir::OpBuilder& builder, const mlir::Value& val, const TensorLayout& layout);

mlir::Value get_affine_load_op(
        mlir::OpBuilder& builder, const mlir::Location& loc, const mlir::Value& val,
        const mlir::ValueRange& index, const TensorLayout& dst);

}  // namespace jit
}  // namespace mgb

#endif  // MGB_JIT && MGB_JIT_MLIR

// vim: syntax=cpp.doxygen