// If we're in the first block, then the argument to that block is the
    // one we're looking for.
    if (auto func = cast<func::FuncOp>(op)) {
      if (blockArg.getOwner()->isEntryBlock()) {
        *out = blockArg;
        return success();
      }
    }

    // If we're in an inner block, the we have to find all ops that branch
    // to that block, and trace through them.