Searched refs:enqueue_op_map (Results 1 – 1 of 1) sorted by relevance
66 func::FuncOp func_op, llvm::StringMap<Operation*>* enqueue_op_map, in FindTPUEmbeddingOps() argument81 if (failed(ExtractEmbeddingAttribute(op, enqueue_op_map))) in FindTPUEmbeddingOps()95 const llvm::StringMap<Operation*>& enqueue_op_map, in UpdateEmbeddingEnqueueOpInput() argument99 for (const auto& it : enqueue_op_map) { in UpdateEmbeddingEnqueueOpInput()151 llvm::StringMap<Operation*> enqueue_op_map; in runOnOperation() local154 if (failed(FindTPUEmbeddingOps(func_op, &enqueue_op_map, in runOnOperation()159 if (enqueue_op_map.size() != recv_activation_op_map.size()) { in runOnOperation()167 if (failed(UpdateEmbeddingEnqueueOpInput(enqueue_op_map, in runOnOperation()