#include <linux/mfd/syscon.h>
 #include <linux/module.h>
 #include <linux/of.h>
+#include <linux/of_address.h>
 #include <linux/of_device.h>
 #include <linux/platform_device.h>
 #include <linux/slab.h>
 
 #include "clk-mtk.h"
 #include "clk-gate.h"
+#include "clk-mux.h"
 
 const struct mtk_gate_regs cg_regs_dummy = { 0, 0, 0 };
 EXPORT_SYMBOL_GPL(cg_regs_dummy);
        const struct mtk_clk_desc *mcd;
        struct clk_hw_onecell_data *clk_data;
        struct device_node *node = pdev->dev.of_node;
-       int r;
+       void __iomem *base;
+       int num_clks, r;
 
        mcd = of_device_get_match_data(&pdev->dev);
        if (!mcd)
                return -EINVAL;
 
-       clk_data = mtk_alloc_clk_data(mcd->num_clks);
+       /* Composite clocks needs us to pass iomem pointer */
+       if (mcd->composite_clks) {
+               if (!mcd->shared_io)
+                       base = devm_platform_ioremap_resource(pdev, 0);
+               else
+                       base = of_iomap(node, 0);
+
+               if (IS_ERR_OR_NULL(base))
+                       return IS_ERR(base) ? PTR_ERR(base) : -ENOMEM;
+       }
+
+       /* Calculate how many clk_hw_onecell_data entries to allocate */
+       num_clks = mcd->num_clks + mcd->num_composite_clks;
+       num_clks += mcd->num_fixed_clks + mcd->num_factor_clks;
+       num_clks += mcd->num_mux_clks;
+
+       clk_data = mtk_alloc_clk_data(num_clks);
        if (!clk_data)
                return -ENOMEM;
 
-       r = mtk_clk_register_gates(&pdev->dev, node, mcd->clks, mcd->num_clks,
-                                  clk_data);
-       if (r)
-               goto free_data;
+       if (mcd->fixed_clks) {
+               r = mtk_clk_register_fixed_clks(mcd->fixed_clks,
+                                               mcd->num_fixed_clks, clk_data);
+               if (r)
+                       goto free_data;
+       }
+
+       if (mcd->factor_clks) {
+               r = mtk_clk_register_factors(mcd->factor_clks,
+                                            mcd->num_factor_clks, clk_data);
+               if (r)
+                       goto unregister_fixed_clks;
+       }
+
+       if (mcd->mux_clks) {
+               r = mtk_clk_register_muxes(&pdev->dev, mcd->mux_clks,
+                                          mcd->num_mux_clks, node,
+                                          mcd->clk_lock, clk_data);
+               if (r)
+                       goto unregister_factors;
+       };
+
+       if (mcd->composite_clks) {
+               /* We don't check composite_lock because it's optional */
+               r = mtk_clk_register_composites(&pdev->dev,
+                                               mcd->composite_clks,
+                                               mcd->num_composite_clks,
+                                               base, mcd->clk_lock, clk_data);
+               if (r)
+                       goto unregister_muxes;
+       }
+
+       if (mcd->clks) {
+               r = mtk_clk_register_gates(&pdev->dev, node, mcd->clks,
+                                          mcd->num_clks, clk_data);
+               if (r)
+                       goto unregister_composites;
+       }
 
        r = of_clk_add_hw_provider(node, of_clk_hw_onecell_get, clk_data);
        if (r)
        return r;
 
 unregister_clks:
-       mtk_clk_unregister_gates(mcd->clks, mcd->num_clks, clk_data);
+       if (mcd->clks)
+               mtk_clk_unregister_gates(mcd->clks, mcd->num_clks, clk_data);
+unregister_composites:
+       if (mcd->composite_clks)
+               mtk_clk_unregister_composites(mcd->composite_clks,
+                                             mcd->num_composite_clks, clk_data);
+unregister_muxes:
+       if (mcd->mux_clks)
+               mtk_clk_unregister_muxes(mcd->mux_clks,
+                                        mcd->num_mux_clks, clk_data);
+unregister_factors:
+       if (mcd->factor_clks)
+               mtk_clk_unregister_factors(mcd->factor_clks,
+                                          mcd->num_factor_clks, clk_data);
+unregister_fixed_clks:
+       if (mcd->fixed_clks)
+               mtk_clk_unregister_fixed_clks(mcd->fixed_clks,
+                                             mcd->num_fixed_clks, clk_data);
 free_data:
        mtk_free_clk_data(clk_data);
+       if (mcd->shared_io && base)
+               iounmap(base);
        return r;
 }
 EXPORT_SYMBOL_GPL(mtk_clk_simple_probe);
        struct device_node *node = pdev->dev.of_node;
 
        of_clk_del_provider(node);
-       mtk_clk_unregister_gates(mcd->clks, mcd->num_clks, clk_data);
+       if (mcd->clks)
+               mtk_clk_unregister_gates(mcd->clks, mcd->num_clks, clk_data);
+       if (mcd->composite_clks)
+               mtk_clk_unregister_composites(mcd->composite_clks,
+                                             mcd->num_composite_clks, clk_data);
+       if (mcd->mux_clks)
+               mtk_clk_unregister_muxes(mcd->mux_clks,
+                                        mcd->num_mux_clks, clk_data);
+       if (mcd->factor_clks)
+               mtk_clk_unregister_factors(mcd->factor_clks,
+                                          mcd->num_factor_clks, clk_data);
+       if (mcd->fixed_clks)
+               mtk_clk_unregister_fixed_clks(mcd->fixed_clks,
+                                             mcd->num_fixed_clks, clk_data);
        mtk_free_clk_data(clk_data);
 
        return 0;