MCPcopy Create free account
hub / github.com/MegEngine/MegCC / CreateCompute

Method CreateCompute

compiler/lib/CodeGen/Elemwise.cpp:51–87  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

49}
50
51void ElemwiseKernel::CreateCompute(
52 Block* entryBlock, mlir::OpBuilder& op_builder, mlir::MLIRContext* ctx,
53 TContext* context) const {
54 Value input_val = entryBlock->getArgument(0);
55 Value output_val = entryBlock->getArgument(1);
56 MemRefType out_memref = output_val.getType().dyn_cast_or_null<MemRefType>();
57
58 SmallVector<AffineMap> indexing_maps;
59 SmallVector<StringRef> iter_type = getNParallelLoopsAttrs(out_memref.getRank());
60
61 indexing_maps.push_back(
62 get_affine_map(input_val.getType().dyn_cast_or_null<MemRefType>(), ctx));
63
64 indexing_maps.push_back(
65 get_affine_map(output_val.getType().dyn_cast_or_null<MemRefType>(), ctx));
66 SmallVector<Value> inputs_val;
67 SmallVector<Value> outputs_val;
68 inputs_val.push_back(input_val);
69 outputs_val.push_back(output_val);
70 op_builder.create<linalg::GenericOp>(
71 op_builder.getUnknownLoc(), inputs_val, outputs_val, indexing_maps,
72 iter_type,
73 [&](OpBuilder& nestedBuilder, Location nestedLoc, ValueRange blockArgs) {
74 Type act_type = blockArgs[0].getType();
75 Value input_val = blockArgs[0];
76 // ConstantOp
77 Value const_zero_op = op_builder.create<arith::ConstantOp>(
78 op_builder.getUnknownLoc(), act_type,
79 FloatAttr::get(act_type, llvm::APFloat(0.f)));
80 Value opResult = op_builder.create<arith::MaxFOp>(
81 op_builder.getUnknownLoc(), act_type, input_val, const_zero_op);
82 op_builder.create<linalg::YieldOp>(
83 op_builder.getUnknownLoc(), opResult);
84 });
85 std::vector<Value> results;
86 op_builder.create<ReturnOp>(op_builder.getUnknownLoc(), results);
87}
88
89void ElemwiseKernel::CreatePass(
90 mlir::PassManager& pm, mlir::MLIRContext* ctx, TContext* context) const {

Callers

nothing calls this directly

Calls 2

getNParallelLoopsAttrsFunction · 0.85
get_affine_mapFunction · 0.85

Tested by

no test coverage detected