diff mbox series

[3/5] iommu/tegra-smmu: Use iommu_fwspec in .probe_/.attach_device()

Message ID 20200926080719.6822-4-nicoleotsuka@gmail.com
State New
Headers show
Series iommu/tegra-smmu: Adding PCI support and mappings debugfs node | expand

Commit Message

Nicolin Chen Sept. 26, 2020, 8:07 a.m. UTC
The tegra_smmu_probe_device() function searches in DT for the iommu
phandler to get "smmu" pointer. This works for most of SMMU clients
that exist in the DTB. But a PCI device will not be added to iommu,
since it doesn't have a DT node.

Fortunately, for a client with a DT node, tegra_smmu_probe_device()
calls tegra_smmu_of_xlate() via tegra_smmu_configure(), while for a
PCI device, of_pci_iommu_init() in the IOMMU core calls .of_xlate()
as well, even before running tegra_smmu_probe_device(). And in both
cases, tegra_smmu_of_xlate() prepares a valid iommu_fwspec pointer
that allows us to get the mc->smmu pointer via dev_get_drvdata() by
calling driver_find_device_by_fwnode().

So this patch uses iommu_fwspec in .probe_device() and related code
for a client that does not exist in the DTB, especially a PCI one.

Signed-off-by: Nicolin Chen <nicoleotsuka@gmail.com>
---
 drivers/iommu/tegra-smmu.c | 89 +++++++++++++++++++++++---------------
 drivers/memory/tegra/mc.c  |  2 +-
 include/soc/tegra/mc.h     |  2 +
 3 files changed, 56 insertions(+), 37 deletions(-)

Comments

Dmitry Osipenko Sept. 26, 2020, 2:48 p.m. UTC | #1
26.09.2020 11:07, Nicolin Chen пишет:
...
> +	/* NULL smmu pointer means that SMMU driver is not probed yet */
> +	if (unlikely(!smmu))
> +		return ERR_PTR(-EPROBE_DEFER);

Hello, Nicolin!

Please don't pollute code with likely/unlikely. This is not a
performance-critical code.

...
> -static struct platform_driver tegra_mc_driver = {
> +struct platform_driver tegra_mc_driver = {
>  	.driver = {
>  		.name = "tegra-mc",
>  		.of_match_table = tegra_mc_of_match,
> diff --git a/include/soc/tegra/mc.h b/include/soc/tegra/mc.h
> index 1238e35653d1..49a4cf64c4b9 100644
> --- a/include/soc/tegra/mc.h
> +++ b/include/soc/tegra/mc.h
> @@ -184,4 +184,6 @@ struct tegra_mc {
>  int tegra_mc_write_emem_configuration(struct tegra_mc *mc, unsigned long rate);
>  unsigned int tegra_mc_get_emem_device_count(struct tegra_mc *mc);
>  
> +extern struct platform_driver tegra_mc_driver;

No global variables, please. See for the example:

https://elixir.bootlin.com/linux/v5.9-rc6/source/drivers/devfreq/tegra20-devfreq.c#L100

The tegra_get_memory_controller() is now needed by multiple Tegra
drivers, I think it should be good to have it added into the MC driver
and then make it globally available for all drivers by making use of
of_find_matching_node_and_match().

diff --git a/drivers/memory/tegra/mc.c b/drivers/memory/tegra/mc.c
index e1db209fd2ea..ed1bd6d00aaf 100644
--- a/drivers/memory/tegra/mc.c
+++ b/drivers/memory/tegra/mc.c
@@ -43,6 +43,29 @@ static const struct of_device_id tegra_mc_of_match[] = {
 };
 MODULE_DEVICE_TABLE(of, tegra_mc_of_match);

+struct tegra_mc *tegra_get_memory_controller(void)
+{
+	struct platform_device *pdev;
+	struct device_node *np;
+	struct tegra_mc *mc;
+
+	np = of_find_matching_node_and_match(NULL, tegra_mc_of_match, NULL);
+	if (!np)
+		return ERR_PTR(-ENOENT);
+
+	pdev = of_find_device_by_node(np);
+	of_node_put(np);
+	if (!pdev)
+		return ERR_PTR(-ENODEV);
+
+	mc = platform_get_drvdata(pdev);
+	if (!mc)
+		return ERR_PTR(-EPROBE_DEFER);
+
+	return mc;
+}
+EXPORT_SYMBOL_GPL(tegra_get_memory_controller);
Nicolin Chen Sept. 26, 2020, 8:42 p.m. UTC | #2
Hi Dmitry,

Thank you for the review and comments!

On Sat, Sep 26, 2020 at 05:48:17PM +0300, Dmitry Osipenko wrote:
> 26.09.2020 11:07, Nicolin Chen пишет:
> ...
> > +	/* NULL smmu pointer means that SMMU driver is not probed yet */
> > +	if (unlikely(!smmu))
> > +		return ERR_PTR(-EPROBE_DEFER);
> 
> Hello, Nicolin!
> 
> Please don't pollute code with likely/unlikely. This is not a
> performance-critical code.

Will drop it. Thanks.

> ...
> > -static struct platform_driver tegra_mc_driver = {
> > +struct platform_driver tegra_mc_driver = {
> >  	.driver = {
> >  		.name = "tegra-mc",
> >  		.of_match_table = tegra_mc_of_match,
> > diff --git a/include/soc/tegra/mc.h b/include/soc/tegra/mc.h
> > index 1238e35653d1..49a4cf64c4b9 100644
> > --- a/include/soc/tegra/mc.h
> > +++ b/include/soc/tegra/mc.h
> > @@ -184,4 +184,6 @@ struct tegra_mc {
> >  int tegra_mc_write_emem_configuration(struct tegra_mc *mc, unsigned long rate);
> >  unsigned int tegra_mc_get_emem_device_count(struct tegra_mc *mc);
> >  
> > +extern struct platform_driver tegra_mc_driver;
> 
> No global variables, please. See for the example:
> 
> https://elixir.bootlin.com/linux/v5.9-rc6/source/drivers/devfreq/tegra20-devfreq.c#L100

Will fix it. Thanks for the example.

> 
> The tegra_get_memory_controller() is now needed by multiple Tegra
> drivers, I think it should be good to have it added into the MC driver
> and then make it globally available for all drivers by making use of
> of_find_matching_node_and_match().
> 
> diff --git a/drivers/memory/tegra/mc.c b/drivers/memory/tegra/mc.c
> index e1db209fd2ea..ed1bd6d00aaf 100644
> --- a/drivers/memory/tegra/mc.c
> +++ b/drivers/memory/tegra/mc.c
> @@ -43,6 +43,29 @@ static const struct of_device_id tegra_mc_of_match[] = {
>  };
>  MODULE_DEVICE_TABLE(of, tegra_mc_of_match);
> 
> +struct tegra_mc *tegra_get_memory_controller(void)
> +{
> +	struct platform_device *pdev;
> +	struct device_node *np;
> +	struct tegra_mc *mc;
> +
> +	np = of_find_matching_node_and_match(NULL, tegra_mc_of_match, NULL);
> +	if (!np)
> +		return ERR_PTR(-ENOENT);
> +
> +	pdev = of_find_device_by_node(np);
> +	of_node_put(np);
> +	if (!pdev)
> +		return ERR_PTR(-ENODEV);
> +
> +	mc = platform_get_drvdata(pdev);
> +	if (!mc)
> +		return ERR_PTR(-EPROBE_DEFER);
> +
> +	return mc;
> +}
> +EXPORT_SYMBOL_GPL(tegra_get_memory_controller);

Will try this one and integrate into my next version.

Thanks
Nic
Thierry Reding Sept. 28, 2020, 7:29 a.m. UTC | #3
On Sat, Sep 26, 2020 at 05:48:17PM +0300, Dmitry Osipenko wrote:
> 26.09.2020 11:07, Nicolin Chen пишет:
> ...
> > +	/* NULL smmu pointer means that SMMU driver is not probed yet */
> > +	if (unlikely(!smmu))
> > +		return ERR_PTR(-EPROBE_DEFER);
> 
> Hello, Nicolin!
> 
> Please don't pollute code with likely/unlikely. This is not a
> performance-critical code.
> 
> ...
> > -static struct platform_driver tegra_mc_driver = {
> > +struct platform_driver tegra_mc_driver = {
> >  	.driver = {
> >  		.name = "tegra-mc",
> >  		.of_match_table = tegra_mc_of_match,
> > diff --git a/include/soc/tegra/mc.h b/include/soc/tegra/mc.h
> > index 1238e35653d1..49a4cf64c4b9 100644
> > --- a/include/soc/tegra/mc.h
> > +++ b/include/soc/tegra/mc.h
> > @@ -184,4 +184,6 @@ struct tegra_mc {
> >  int tegra_mc_write_emem_configuration(struct tegra_mc *mc, unsigned long rate);
> >  unsigned int tegra_mc_get_emem_device_count(struct tegra_mc *mc);
> >  
> > +extern struct platform_driver tegra_mc_driver;
> 
> No global variables, please. See for the example:
> 
> https://elixir.bootlin.com/linux/v5.9-rc6/source/drivers/devfreq/tegra20-devfreq.c#L100
> 
> The tegra_get_memory_controller() is now needed by multiple Tegra
> drivers, I think it should be good to have it added into the MC driver
> and then make it globally available for all drivers by making use of
> of_find_matching_node_and_match().
> 
> diff --git a/drivers/memory/tegra/mc.c b/drivers/memory/tegra/mc.c
> index e1db209fd2ea..ed1bd6d00aaf 100644
> --- a/drivers/memory/tegra/mc.c
> +++ b/drivers/memory/tegra/mc.c
> @@ -43,6 +43,29 @@ static const struct of_device_id tegra_mc_of_match[] = {
>  };
>  MODULE_DEVICE_TABLE(of, tegra_mc_of_match);
> 
> +struct tegra_mc *tegra_get_memory_controller(void)
> +{
> +	struct platform_device *pdev;
> +	struct device_node *np;
> +	struct tegra_mc *mc;
> +
> +	np = of_find_matching_node_and_match(NULL, tegra_mc_of_match, NULL);
> +	if (!np)
> +		return ERR_PTR(-ENOENT);
> +
> +	pdev = of_find_device_by_node(np);
> +	of_node_put(np);
> +	if (!pdev)
> +		return ERR_PTR(-ENODEV);
> +
> +	mc = platform_get_drvdata(pdev);
> +	if (!mc)
> +		return ERR_PTR(-EPROBE_DEFER);
> +
> +	return mc;
> +}
> +EXPORT_SYMBOL_GPL(tegra_get_memory_controller);

We already have tegra_smmu_find(), which should be enough for this
particular use-case.

Thierry
Thierry Reding Sept. 28, 2020, 7:52 a.m. UTC | #4
On Sat, Sep 26, 2020 at 01:07:17AM -0700, Nicolin Chen wrote:
> The tegra_smmu_probe_device() function searches in DT for the iommu
> phandler to get "smmu" pointer. This works for most of SMMU clients
> that exist in the DTB. But a PCI device will not be added to iommu,
> since it doesn't have a DT node.
> 
> Fortunately, for a client with a DT node, tegra_smmu_probe_device()
> calls tegra_smmu_of_xlate() via tegra_smmu_configure(), while for a
> PCI device, of_pci_iommu_init() in the IOMMU core calls .of_xlate()
> as well, even before running tegra_smmu_probe_device(). And in both
> cases, tegra_smmu_of_xlate() prepares a valid iommu_fwspec pointer
> that allows us to get the mc->smmu pointer via dev_get_drvdata() by
> calling driver_find_device_by_fwnode().
> 
> So this patch uses iommu_fwspec in .probe_device() and related code
> for a client that does not exist in the DTB, especially a PCI one.
> 
> Signed-off-by: Nicolin Chen <nicoleotsuka@gmail.com>
> ---
>  drivers/iommu/tegra-smmu.c | 89 +++++++++++++++++++++++---------------
>  drivers/memory/tegra/mc.c  |  2 +-
>  include/soc/tegra/mc.h     |  2 +
>  3 files changed, 56 insertions(+), 37 deletions(-)
> 
> diff --git a/drivers/iommu/tegra-smmu.c b/drivers/iommu/tegra-smmu.c
> index b10e02073610..97a7185b4578 100644
> --- a/drivers/iommu/tegra-smmu.c
> +++ b/drivers/iommu/tegra-smmu.c
> @@ -13,6 +13,7 @@
>  #include <linux/platform_device.h>
>  #include <linux/slab.h>
>  #include <linux/spinlock.h>
> +#include <linux/dma-iommu.h>

Why is this needed? I don't see any of the symbols declared in that file
used here.

>  #include <linux/dma-mapping.h>
>  
>  #include <soc/tegra/ahb.h>
> @@ -61,6 +62,8 @@ struct tegra_smmu_as {
>  	u32 attr;
>  };
>  
> +static const struct iommu_ops tegra_smmu_ops;
> +
>  static struct tegra_smmu_as *to_smmu_as(struct iommu_domain *dom)
>  {
>  	return container_of(dom, struct tegra_smmu_as, domain);
> @@ -484,60 +487,49 @@ static void tegra_smmu_as_unprepare(struct tegra_smmu *smmu,
>  static int tegra_smmu_attach_dev(struct iommu_domain *domain,
>  				 struct device *dev)
>  {
> +	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
>  	struct tegra_smmu *smmu = dev_iommu_priv_get(dev);
>  	struct tegra_smmu_as *as = to_smmu_as(domain);
> -	struct device_node *np = dev->of_node;
> -	struct of_phandle_args args;
> -	unsigned int index = 0;
> -	int err = 0;
> -
> -	while (!of_parse_phandle_with_args(np, "iommus", "#iommu-cells", index,
> -					   &args)) {
> -		unsigned int swgroup = args.args[0];
> -
> -		if (args.np != smmu->dev->of_node) {
> -			of_node_put(args.np);
> -			continue;
> -		}
> +	int index, err = 0;
>  
> -		of_node_put(args.np);
> +	if (!fwspec || fwspec->ops != &tegra_smmu_ops)
> +		return -ENOENT;
>  
> +	for (index = 0; index < fwspec->num_ids; index++) {
>  		err = tegra_smmu_as_prepare(smmu, as);
> -		if (err < 0)
> -			return err;
> +		if (err)
> +			goto err_disable;

I'd personally drop the err_ prefix here because it's pretty obvious
that we're going to do this as a result of an error happening.

>  
> -		tegra_smmu_enable(smmu, swgroup, as->id);
> -		index++;
> +		tegra_smmu_enable(smmu, fwspec->ids[index], as->id);
>  	}
>  
>  	if (index == 0)
>  		return -ENODEV;
>  
>  	return 0;
> +
> +err_disable:
> +	for (index--; index >= 0; index--) {
> +		tegra_smmu_disable(smmu, fwspec->ids[index], as->id);
> +		tegra_smmu_as_unprepare(smmu, as);
> +	}

I think a more idiomatic version of doing this would be:

	while (index--) {
		...
	}

> +
> +	return err;
>  }
>  
>  static void tegra_smmu_detach_dev(struct iommu_domain *domain, struct device *dev)
>  {
> +	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
>  	struct tegra_smmu_as *as = to_smmu_as(domain);
> -	struct device_node *np = dev->of_node;
>  	struct tegra_smmu *smmu = as->smmu;
> -	struct of_phandle_args args;
>  	unsigned int index = 0;
>  
> -	while (!of_parse_phandle_with_args(np, "iommus", "#iommu-cells", index,
> -					   &args)) {
> -		unsigned int swgroup = args.args[0];
> -
> -		if (args.np != smmu->dev->of_node) {
> -			of_node_put(args.np);
> -			continue;
> -		}
> -
> -		of_node_put(args.np);
> +	if (!fwspec || fwspec->ops != &tegra_smmu_ops)
> +		return;
>  
> -		tegra_smmu_disable(smmu, swgroup, as->id);
> +	for (index = 0; index < fwspec->num_ids; index++) {
> +		tegra_smmu_disable(smmu, fwspec->ids[index], as->id);
>  		tegra_smmu_as_unprepare(smmu, as);
> -		index++;
>  	}
>  }
>  
> @@ -845,10 +837,25 @@ static int tegra_smmu_configure(struct tegra_smmu *smmu, struct device *dev,
>  	return 0;
>  }
>  
> +static struct tegra_smmu *tegra_smmu_get_by_fwnode(struct fwnode_handle *fwnode)
> +{
> +	struct device *dev = driver_find_device_by_fwnode(&tegra_mc_driver.driver, fwnode);
> +	struct tegra_mc *mc;
> +
> +	if (!dev)
> +		return NULL;
> +
> +	put_device(dev);
> +	mc = dev_get_drvdata(dev);
> +
> +	return mc->smmu;
> +}
> +

As I mentioned in another reply, I think tegra_smmu_find() should be all
you need in this case.

>  static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
>  {
>  	struct device_node *np = dev->of_node;
>  	struct tegra_smmu *smmu = NULL;
> +	struct iommu_fwspec *fwspec;
>  	struct of_phandle_args args;
>  	unsigned int index = 0;
>  	int err;
> @@ -868,8 +875,6 @@ static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
>  			 * supported by the Linux kernel, so abort after the
>  			 * first match.
>  			 */
> -			dev_iommu_priv_set(dev, smmu);
> -
>  			break;
>  		}
>  
> @@ -877,8 +882,20 @@ static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
>  		index++;
>  	}
>  
> -	if (!smmu)
> -		return ERR_PTR(-ENODEV);
> +	/* Any device should be able to get smmu pointer using fwspec */
> +	if (!smmu) {
> +		fwspec = dev_iommu_fwspec_get(dev);
> +		if (!fwspec || fwspec->ops != &tegra_smmu_ops)
> +			return ERR_PTR(-ENODEV);
> +
> +		smmu = tegra_smmu_get_by_fwnode(fwspec->iommu_fwnode);
> +	}
> +
> +	/* NULL smmu pointer means that SMMU driver is not probed yet */
> +	if (unlikely(!smmu))
> +		return ERR_PTR(-EPROBE_DEFER);
> +
> +	dev_iommu_priv_set(dev, smmu);

Can this not be unified with the OF code above? Basically in all cases
where we use Tegra SMMU, a fwnode_node should correspond 1:1 to a struct
device_node, which makes the code you added here effectively the same as
what's already there.

Thierry
Nicolin Chen Sept. 28, 2020, 10:18 p.m. UTC | #5
On Mon, Sep 28, 2020 at 09:52:12AM +0200, Thierry Reding wrote:
> On Sat, Sep 26, 2020 at 01:07:17AM -0700, Nicolin Chen wrote:
> > @@ -13,6 +13,7 @@
> >  #include <linux/platform_device.h>
> >  #include <linux/slab.h>
> >  #include <linux/spinlock.h>
> > +#include <linux/dma-iommu.h>
> 
> Why is this needed? I don't see any of the symbols declared in that file
> used here.

Hmm..I think I mixed with some other change that needs this header
file. Removing it....

> >  #include <linux/dma-mapping.h>
> >  
> >  #include <soc/tegra/ahb.h>
> > @@ -61,6 +62,8 @@ struct tegra_smmu_as {
> >  	u32 attr;
> >  };
> >  
> > +static const struct iommu_ops tegra_smmu_ops;
> > +
> >  static struct tegra_smmu_as *to_smmu_as(struct iommu_domain *dom)
> >  {
> >  	return container_of(dom, struct tegra_smmu_as, domain);
> > @@ -484,60 +487,49 @@ static void tegra_smmu_as_unprepare(struct tegra_smmu *smmu,
> >  static int tegra_smmu_attach_dev(struct iommu_domain *domain,
> >  				 struct device *dev)
> >  {
> > +	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
> >  	struct tegra_smmu *smmu = dev_iommu_priv_get(dev);
> >  	struct tegra_smmu_as *as = to_smmu_as(domain);
> > -	struct device_node *np = dev->of_node;
> > -	struct of_phandle_args args;
> > -	unsigned int index = 0;
> > -	int err = 0;
> > -
> > -	while (!of_parse_phandle_with_args(np, "iommus", "#iommu-cells", index,
> > -					   &args)) {
> > -		unsigned int swgroup = args.args[0];
> > -
> > -		if (args.np != smmu->dev->of_node) {
> > -			of_node_put(args.np);
> > -			continue;
> > -		}
> > +	int index, err = 0;
> >  
> > -		of_node_put(args.np);
> > +	if (!fwspec || fwspec->ops != &tegra_smmu_ops)
> > +		return -ENOENT;
> >  
> > +	for (index = 0; index < fwspec->num_ids; index++) {
> >  		err = tegra_smmu_as_prepare(smmu, as);
> > -		if (err < 0)
> > -			return err;
> > +		if (err)
> > +			goto err_disable;
> 
> I'd personally drop the err_ prefix here because it's pretty obvious
> that we're going to do this as a result of an error happening.

Changing to "goto disable".

> > @@ -845,10 +837,25 @@ static int tegra_smmu_configure(struct tegra_smmu *smmu, struct device *dev,
> >  	return 0;
> >  }
> >  
> > +static struct tegra_smmu *tegra_smmu_get_by_fwnode(struct fwnode_handle *fwnode)
> > +{
> > +	struct device *dev = driver_find_device_by_fwnode(&tegra_mc_driver.driver, fwnode);
> > +	struct tegra_mc *mc;
> > +
> > +	if (!dev)
> > +		return NULL;
> > +
> > +	put_device(dev);
> > +	mc = dev_get_drvdata(dev);
> > +
> > +	return mc->smmu;
> > +}
> > +
> 
> As I mentioned in another reply, I think tegra_smmu_find() should be all
> you need in this case.

This function is used by .probe_device() where its dev pointer is
an SMMU client. IIUC, tegra_smmu_find() needs np pointer of "mc".
For a PCI device that doesn't have a DT node with iommus property,
not sure how we can use tegra_smmu_find().
Dmitry Osipenko Sept. 29, 2020, 4:06 a.m. UTC | #6
...
>> As I mentioned in another reply, I think tegra_smmu_find() should be all
>> you need in this case.
> 
> This function is used by .probe_device() where its dev pointer is
> an SMMU client. IIUC, tegra_smmu_find() needs np pointer of "mc".
> For a PCI device that doesn't have a DT node with iommus property,
> not sure how we can use tegra_smmu_find().

Perhaps you could get np from struct pci_dev.bus?
Nicolin Chen Sept. 29, 2020, 5:36 a.m. UTC | #7
On Tue, Sep 29, 2020 at 07:06:37AM +0300, Dmitry Osipenko wrote:
> ...
> >> As I mentioned in another reply, I think tegra_smmu_find() should be all
> >> you need in this case.
> > 
> > This function is used by .probe_device() where its dev pointer is
> > an SMMU client. IIUC, tegra_smmu_find() needs np pointer of "mc".
> > For a PCI device that doesn't have a DT node with iommus property,
> > not sure how we can use tegra_smmu_find().
> 
> Perhaps you could get np from struct pci_dev.bus?

Possibly yes...but I was hoping the solution would be potentially
reused by other devices that don't have DT nodes, USB for example,
though I haven't tested with current version.
diff mbox series

Patch

diff --git a/drivers/iommu/tegra-smmu.c b/drivers/iommu/tegra-smmu.c
index b10e02073610..97a7185b4578 100644
--- a/drivers/iommu/tegra-smmu.c
+++ b/drivers/iommu/tegra-smmu.c
@@ -13,6 +13,7 @@ 
 #include <linux/platform_device.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>
+#include <linux/dma-iommu.h>
 #include <linux/dma-mapping.h>
 
 #include <soc/tegra/ahb.h>
@@ -61,6 +62,8 @@  struct tegra_smmu_as {
 	u32 attr;
 };
 
+static const struct iommu_ops tegra_smmu_ops;
+
 static struct tegra_smmu_as *to_smmu_as(struct iommu_domain *dom)
 {
 	return container_of(dom, struct tegra_smmu_as, domain);
@@ -484,60 +487,49 @@  static void tegra_smmu_as_unprepare(struct tegra_smmu *smmu,
 static int tegra_smmu_attach_dev(struct iommu_domain *domain,
 				 struct device *dev)
 {
+	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
 	struct tegra_smmu *smmu = dev_iommu_priv_get(dev);
 	struct tegra_smmu_as *as = to_smmu_as(domain);
-	struct device_node *np = dev->of_node;
-	struct of_phandle_args args;
-	unsigned int index = 0;
-	int err = 0;
-
-	while (!of_parse_phandle_with_args(np, "iommus", "#iommu-cells", index,
-					   &args)) {
-		unsigned int swgroup = args.args[0];
-
-		if (args.np != smmu->dev->of_node) {
-			of_node_put(args.np);
-			continue;
-		}
+	int index, err = 0;
 
-		of_node_put(args.np);
+	if (!fwspec || fwspec->ops != &tegra_smmu_ops)
+		return -ENOENT;
 
+	for (index = 0; index < fwspec->num_ids; index++) {
 		err = tegra_smmu_as_prepare(smmu, as);
-		if (err < 0)
-			return err;
+		if (err)
+			goto err_disable;
 
-		tegra_smmu_enable(smmu, swgroup, as->id);
-		index++;
+		tegra_smmu_enable(smmu, fwspec->ids[index], as->id);
 	}
 
 	if (index == 0)
 		return -ENODEV;
 
 	return 0;
+
+err_disable:
+	for (index--; index >= 0; index--) {
+		tegra_smmu_disable(smmu, fwspec->ids[index], as->id);
+		tegra_smmu_as_unprepare(smmu, as);
+	}
+
+	return err;
 }
 
 static void tegra_smmu_detach_dev(struct iommu_domain *domain, struct device *dev)
 {
+	struct iommu_fwspec *fwspec = dev_iommu_fwspec_get(dev);
 	struct tegra_smmu_as *as = to_smmu_as(domain);
-	struct device_node *np = dev->of_node;
 	struct tegra_smmu *smmu = as->smmu;
-	struct of_phandle_args args;
 	unsigned int index = 0;
 
-	while (!of_parse_phandle_with_args(np, "iommus", "#iommu-cells", index,
-					   &args)) {
-		unsigned int swgroup = args.args[0];
-
-		if (args.np != smmu->dev->of_node) {
-			of_node_put(args.np);
-			continue;
-		}
-
-		of_node_put(args.np);
+	if (!fwspec || fwspec->ops != &tegra_smmu_ops)
+		return;
 
-		tegra_smmu_disable(smmu, swgroup, as->id);
+	for (index = 0; index < fwspec->num_ids; index++) {
+		tegra_smmu_disable(smmu, fwspec->ids[index], as->id);
 		tegra_smmu_as_unprepare(smmu, as);
-		index++;
 	}
 }
 
@@ -845,10 +837,25 @@  static int tegra_smmu_configure(struct tegra_smmu *smmu, struct device *dev,
 	return 0;
 }
 
+static struct tegra_smmu *tegra_smmu_get_by_fwnode(struct fwnode_handle *fwnode)
+{
+	struct device *dev = driver_find_device_by_fwnode(&tegra_mc_driver.driver, fwnode);
+	struct tegra_mc *mc;
+
+	if (!dev)
+		return NULL;
+
+	put_device(dev);
+	mc = dev_get_drvdata(dev);
+
+	return mc->smmu;
+}
+
 static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
 {
 	struct device_node *np = dev->of_node;
 	struct tegra_smmu *smmu = NULL;
+	struct iommu_fwspec *fwspec;
 	struct of_phandle_args args;
 	unsigned int index = 0;
 	int err;
@@ -868,8 +875,6 @@  static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
 			 * supported by the Linux kernel, so abort after the
 			 * first match.
 			 */
-			dev_iommu_priv_set(dev, smmu);
-
 			break;
 		}
 
@@ -877,8 +882,20 @@  static struct iommu_device *tegra_smmu_probe_device(struct device *dev)
 		index++;
 	}
 
-	if (!smmu)
-		return ERR_PTR(-ENODEV);
+	/* Any device should be able to get smmu pointer using fwspec */
+	if (!smmu) {
+		fwspec = dev_iommu_fwspec_get(dev);
+		if (!fwspec || fwspec->ops != &tegra_smmu_ops)
+			return ERR_PTR(-ENODEV);
+
+		smmu = tegra_smmu_get_by_fwnode(fwspec->iommu_fwnode);
+	}
+
+	/* NULL smmu pointer means that SMMU driver is not probed yet */
+	if (unlikely(!smmu))
+		return ERR_PTR(-EPROBE_DEFER);
+
+	dev_iommu_priv_set(dev, smmu);
 
 	return &smmu->iommu;
 }
diff --git a/drivers/memory/tegra/mc.c b/drivers/memory/tegra/mc.c
index ec8403557ed4..4e23fd8d8433 100644
--- a/drivers/memory/tegra/mc.c
+++ b/drivers/memory/tegra/mc.c
@@ -735,7 +735,7 @@  static const struct dev_pm_ops tegra_mc_pm_ops = {
 	.resume = tegra_mc_resume,
 };
 
-static struct platform_driver tegra_mc_driver = {
+struct platform_driver tegra_mc_driver = {
 	.driver = {
 		.name = "tegra-mc",
 		.of_match_table = tegra_mc_of_match,
diff --git a/include/soc/tegra/mc.h b/include/soc/tegra/mc.h
index 1238e35653d1..49a4cf64c4b9 100644
--- a/include/soc/tegra/mc.h
+++ b/include/soc/tegra/mc.h
@@ -184,4 +184,6 @@  struct tegra_mc {
 int tegra_mc_write_emem_configuration(struct tegra_mc *mc, unsigned long rate);
 unsigned int tegra_mc_get_emem_device_count(struct tegra_mc *mc);
 
+extern struct platform_driver tegra_mc_driver;
+
 #endif /* __SOC_TEGRA_MC_H__ */