Searched refs:s_and_s_op (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/mlir/hlo/lib/Dialect/mhlo/transforms/ |
| D | lhlo_legalize_to_parallel_loops.cc | 488 lmhlo::SelectAndScatterOp s_and_s_op, ArrayRef<Value> /*args*/, in matchAndRewrite() argument
490 auto loc = s_and_s_op.getLoc(); in matchAndRewrite()
491 InitializeOutput(s_and_s_op, &rewriter); in matchAndRewrite()
493 MakeLoopOverShape(loc, s_and_s_op.source(), &rewriter); in matchAndRewrite()
497 auto selected_ivs = SelectIvs(s_and_s_op, loop_over_src, &rewriter); in matchAndRewrite()
500 auto src_elem = rewriter.create<LoadOp>(loc, s_and_s_op.source(), in matchAndRewrite()
504 auto rmw = rewriter.create<GenericAtomicRMWOp>(loc, s_and_s_op.out(), in matchAndRewrite()
509 &s_and_s_op.scatter().front(), &rmw_builder); in matchAndRewrite()
512 rewriter.replaceOp(s_and_s_op, llvm::None); in matchAndRewrite()
517 void InitializeOutput(lmhlo::SelectAndScatterOp s_and_s_op, in InitializeOutput() argument
[all …]
|