*/
 void __iomem *fwnode_iomap(struct fwnode_handle *fwnode, int index)
 {
-       if (IS_ENABLED(CONFIG_OF_ADDRESS) && is_of_node(fwnode))
-               return of_iomap(to_of_node(fwnode), index);
-
-       return NULL;
+       return fwnode_call_ptr_op(fwnode, iomap, index);
 }
 EXPORT_SYMBOL(fwnode_iomap);
 
 
        return 0;
 }
 
+static void __iomem *of_fwnode_iomap(struct fwnode_handle *fwnode, int index)
+{
+#ifdef CONFIG_OF_ADDRESS
+       return of_iomap(to_of_node(fwnode), index);
+#else
+       return NULL;
+#endif
+}
+
 static int of_fwnode_add_links(struct fwnode_handle *fwnode)
 {
        struct property *p;
        .graph_get_remote_endpoint = of_fwnode_graph_get_remote_endpoint,
        .graph_get_port_parent = of_fwnode_graph_get_port_parent,
        .graph_parse_endpoint = of_fwnode_graph_parse_endpoint,
+       .iomap = of_fwnode_iomap,
        .add_links = of_fwnode_add_links,
 };
 EXPORT_SYMBOL_GPL(of_fwnode_ops);
 
        (*graph_get_port_parent)(struct fwnode_handle *fwnode);
        int (*graph_parse_endpoint)(const struct fwnode_handle *fwnode,
                                    struct fwnode_endpoint *endpoint);
+       void __iomem *(*iomap)(struct fwnode_handle *fwnode, int index);
        int (*add_links)(struct fwnode_handle *fwnode);
 };