Add power control for Apex

Creates a character device (apex_power) that when open/closed
toggles the PCIe bus as well as the Apex regulator. This can
be used to signficantly reduce idle power when the TPU is
idle.

Note that this is disabled by default in the phanbell device tree.
To enable the apex-power status must be changed to "okay" and for
full power savings apex-regulators should remove the always-on
property.

Change-Id: I906373b5edf481820b692847ebaba57943b85897
diff --git a/arch/arm64/boot/dts/freescale/fsl-imx8mq-som.dtsi b/arch/arm64/boot/dts/freescale/fsl-imx8mq-som.dtsi
index e701eb5..40230ef 100644
--- a/arch/arm64/boot/dts/freescale/fsl-imx8mq-som.dtsi
+++ b/arch/arm64/boot/dts/freescale/fsl-imx8mq-som.dtsi
@@ -49,7 +49,7 @@
 
 	apex_power {
 		compatible = "google,apex-power";
-		status = "okay";
+		status = "disabled";
 		power-supply = <&reg_apex>;
 	};
 
diff --git a/drivers/char/Kconfig b/drivers/char/Kconfig
index e08863b..af87ec5 100644
--- a/drivers/char/Kconfig
+++ b/drivers/char/Kconfig
@@ -170,6 +170,24 @@
 
 	  If unsure, say N.
 
+config GOOGLE_APEX_POWER
+	tristate "Google Apex TPU PCI-based power control"
+	depends on PCI
+	help
+	  Character device for managing the power of the Google Apex
+	  TPU.
+
+	  If you say Y here, a special character device node, /dev/apexp*
+	  will be created which exposes PCI power control operations to
+	  user space. This is used to reduce leakage current by disabling
+	  power to the TPU completely via the PCI-tree and various PMICs.
+
+	  Note: your board's devicetree must include a "compatible =
+	  google,apex-power" stanza somewhere in it for this driver to
+	  function.
+
+	  If you don't know what this is, you should say N here.
+
 source "drivers/tty/hvc/Kconfig"
 
 config VIRTIO_CONSOLE
diff --git a/drivers/char/Makefile b/drivers/char/Makefile
index e021811..b71520f 100644
--- a/drivers/char/Makefile
+++ b/drivers/char/Makefile
@@ -17,6 +17,8 @@
 obj-$(CONFIG_BFIN_OTP)		+= bfin-otp.o
 obj-$(CONFIG_FSL_OTP)		+= fsl_otp.o
 
+obj-$(CONFIG_GOOGLE_APEX_POWER)	+= apex-power.o
+
 obj-$(CONFIG_PRINTER)		+= lp.o
 
 obj-$(CONFIG_APM_EMULATION)	+= apm-emulation.o
diff --git a/drivers/char/apex-power.c b/drivers/char/apex-power.c
new file mode 100644
index 0000000..b501233
--- /dev/null
+++ b/drivers/char/apex-power.c
@@ -0,0 +1,340 @@
+/*
+ * Driver to expose PCI-express hot-removal of devices via file
+ * handles.
+ *
+ * Copyright (C) 2018 Google, Inc.
+ * Author: June Tate-Gans <jtgans@google.com>
+ *
+ * Shamelessly the general structure of the character driver has been
+ * borrowed from bsr.c, the POWER architecture's method for accessing
+ * BSR registers.
+ *
+ * This software is licensed under the terms of the GNU General Public
+ * License version 2, as published by the Free Software Foundation, and
+ * may be copied, distributed, and modified under those terms.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ */
+
+#include <linux/types.h>
+#include <linux/module.h>
+
+#include <asm/cmpxchg.h>
+#include <linux/atomic.h>
+#include <linux/cdev.h>
+#include <linux/device.h>
+#include <linux/delay.h>
+#include <linux/fs.h>
+#include <linux/kernel.h>
+#include <linux/mutex.h>
+#include <linux/of.h>
+#include <linux/of_platform.h>
+#include <linux/pci.h>
+#include <linux/regulator/consumer.h>
+#include <linux/workqueue.h>
+
+#if defined(CONFIG_IMX8MQ_PHANBELL_POWERSAVE)
+#include <linux/busfreq-imx.h>
+#endif
+
+#define APEX_PCI_VENDOR_ID 0x1ac1
+#define APEX_PCI_DEVICE_ID 0x089a
+
+#define APEX_ALLOWED_RETRIES 5
+#define APEX_RETRY_DELAY_MS 100
+
+static atomic_t apex_power_minor = ATOMIC_INIT(0);
+static struct class *apex_power_class;
+static int apex_power_major;
+
+struct apex_power_priv {
+	struct platform_device *pdev;
+	struct device *device;
+	struct cdev cdev;
+	bool owned;
+	struct mutex owned_lock;
+	struct delayed_work delayed_init;
+	struct regulator *supply;
+};
+
+static struct pci_dev *get_apex_pci_device(struct apex_power_priv *apex_power_data, bool rescan) {
+	int retries = 0;
+	struct pci_bus *pci_bus = NULL;
+	struct pci_dev *apex_dev = NULL;
+
+	for (retries = 0; retries < APEX_ALLOWED_RETRIES; retries++) {
+		/* For powering up, rescan the pci bus each retry */
+		if (rescan) {
+			pci_lock_rescan_remove();
+			while ((pci_bus = pci_find_next_bus(pci_bus)) != NULL) {
+				pci_rescan_bus(pci_bus);
+			}
+			pci_unlock_rescan_remove();
+		}
+
+		apex_dev = pci_get_device(APEX_PCI_VENDOR_ID, APEX_PCI_DEVICE_ID, NULL);
+		if (apex_dev) {
+			break;
+		}
+		dev_err(&apex_power_data->pdev->dev, "Unable to find apex device, retrying");
+		msleep(APEX_RETRY_DELAY_MS);
+	}
+	return apex_dev;
+}
+
+static int apex_power_down(struct apex_power_priv *apex_power_data)
+{
+	struct pci_dev *apex_dev = NULL;
+	struct pci_dev *apex_connected_bus = NULL;
+	int ret = 0;
+
+#if defined(CONFIG_IMX8MQ_PHANBELL_POWERSAVE)
+	if (apex_power_data->owned) {
+		release_bus_freq(BUS_FREQ_HIGH);
+	}
+#endif
+
+	apex_dev = get_apex_pci_device(apex_power_data, /*rescan=*/ false);
+	if (!apex_dev) {
+		dev_err(&apex_power_data->pdev->dev, "can't find Apex on PCI bus?!");
+		return -EIO;
+	}
+	apex_connected_bus = apex_dev->bus->self;
+	pci_dev_put(apex_dev);
+
+	pci_stop_and_remove_bus_device_locked(apex_connected_bus);
+
+	ret = regulator_disable(apex_power_data->supply);
+	if (ret) {
+		dev_err(&apex_power_data->pdev->dev, "Unable to disable regulator.");
+	}
+
+	return 0;
+}
+
+static int apex_power_up(struct apex_power_priv *apex_power_data)
+{
+	struct pci_dev *apex_dev = NULL;
+	int ret = 0;
+
+	ret = regulator_enable(apex_power_data->supply);
+	if (ret) {
+		dev_err(&apex_power_data->pdev->dev, "Unable to enable regulator.");
+	}
+
+#if defined(CONFIG_IMX8MQ_PHANBELL_POWERSAVE)
+	request_bus_freq(BUS_FREQ_HIGH);
+#endif
+
+	apex_dev = get_apex_pci_device(apex_power_data, /*rescan=*/ true);
+	if (!apex_dev) {
+		dev_err(&apex_power_data->pdev->dev, "can't find Apex on PCI bus?!");
+
+#if defined(CONFIG_IMX8MQ_PHANBELL_POWERSAVE)
+		release_bus_freq(BUS_FREQ_HIGH);
+#endif
+
+		return -EIO;
+	}
+	pci_dev_put(apex_dev);
+
+	return 0;
+}
+
+static int apex_power_release(struct inode* inode, struct file* filep)
+{
+	int ret = 0;
+	struct apex_power_priv *apex_power_data = 
+		container_of(inode->i_cdev, struct apex_power_priv, cdev);
+	if (mutex_lock_interruptible(&apex_power_data->owned_lock)) {
+		return -ENOLCK;
+	}
+
+	if (apex_power_data->owned != true) {
+		ret = -EPERM;
+		goto out;
+	}
+
+	ret = apex_power_down(apex_power_data);
+	apex_power_data->owned = false;
+
+out:
+	mutex_unlock(&apex_power_data->owned_lock);
+	return ret;
+}
+
+static int apex_power_open(struct inode* inode, struct file* filep)
+{
+	int ret = 0;
+	struct apex_power_priv *apex_power_data = 
+		container_of(inode->i_cdev, struct apex_power_priv, cdev);
+
+	if (mutex_lock_interruptible(&apex_power_data->owned_lock)) {
+		return -ENOLCK;
+	}
+
+	if (apex_power_data->owned != false) {
+		ret = -EPERM;
+		goto out;
+	}
+	apex_power_data->owned = true;
+
+	/* Ensure that we don't accidentally run afoul of the initial delayed
+	   power down. */
+	cancel_delayed_work_sync(&apex_power_data->delayed_init);
+
+	ret = apex_power_up(apex_power_data);
+
+out:
+	mutex_unlock(&apex_power_data->owned_lock);
+	return ret;
+}
+
+static const struct file_operations apex_power_fops = {
+	.owner = THIS_MODULE,
+	.open = apex_power_open,
+	.release = apex_power_release,
+};
+
+static void apex_power_delayed_init_callback(struct work_struct *work)
+{
+	struct pci_dev *apex_dev = NULL;
+	struct apex_power_priv *apex_power_data =
+		container_of(
+			container_of(work, struct delayed_work, work),
+			struct apex_power_priv,
+			delayed_init);
+	apex_dev = pci_get_device(APEX_PCI_VENDOR_ID, APEX_PCI_DEVICE_ID, NULL);
+
+	if (!apex_dev) {
+		dev_info(&apex_power_data->pdev->dev, "rescheduling late init power down");
+		schedule_delayed_work(&apex_power_data->delayed_init,
+				      msecs_to_jiffies(1000));
+		return;
+	}
+
+	pci_dev_put(apex_dev);
+
+	dev_info(&apex_power_data->pdev->dev, "init routines powering down apex");
+
+	if (mutex_lock_interruptible(&apex_power_data->owned_lock)) {
+		dev_err(&apex_power_data->pdev->dev, "Unable to lock mutex.");
+		return;
+	}
+
+	apex_power_down(apex_power_data);
+
+	mutex_unlock(&apex_power_data->owned_lock);
+}
+
+static int apex_power_probe(struct platform_device *pdev)
+{
+	dev_t apex_power_dev;
+	int ret = 0;
+	struct apex_power_priv *apex_power_data = 
+		devm_kzalloc(&pdev->dev, sizeof(struct apex_power_priv), GFP_KERNEL);
+	if (!apex_power_data) {
+		return -ENOMEM;
+	}
+	platform_set_drvdata(pdev, apex_power_data);
+	apex_power_data->pdev = pdev;
+	mutex_init(&apex_power_data->owned_lock);
+
+	// Enable the supply in regulator framework to ensure power isn't removed
+	// until init is complete.
+	apex_power_data->supply = regulator_get_exclusive(&pdev->dev, "power");
+	if (IS_ERR(apex_power_data->supply)) {
+		dev_err(&pdev->dev, "Unable to find regulator.");
+		return -ENODEV;
+	}
+	
+	ret = regulator_enable(apex_power_data->supply);
+	if (ret) {
+		dev_err(&pdev->dev, "Unable to enable regulator.");
+		return -ENODEV;
+	}
+
+	apex_power_dev = MKDEV(apex_power_major, atomic_inc_return(&apex_power_minor));
+	cdev_init(&apex_power_data->cdev, &apex_power_fops);
+	ret = cdev_add(&apex_power_data->cdev, apex_power_dev, 1);
+
+	if (ret) {
+		dev_err(&pdev->dev, "Unable to add cdev!");
+	}
+
+	apex_power_data->device = device_create(apex_power_class, NULL, apex_power_dev,
+					  NULL, "%s", dev_name(&pdev->dev));
+	if (IS_ERR(apex_power_data->device)) {
+		dev_err(&pdev->dev, "device_create failed.");
+		cdev_del(&apex_power_data->cdev);
+	}
+
+	INIT_DELAYED_WORK(&apex_power_data->delayed_init,
+			  apex_power_delayed_init_callback);
+	schedule_delayed_work(&apex_power_data->delayed_init, msecs_to_jiffies(7000));
+
+	dev_info(&pdev->dev, "initialized.");
+
+	return 0;
+}
+
+static int apex_power_remove(struct platform_device *pdev)
+{
+	struct apex_power_priv *apex_power_data = platform_get_drvdata(pdev); 
+	device_del(apex_power_data->device);
+	cdev_del(&apex_power_data->cdev);
+	regulator_disable(apex_power_data->supply);
+	regulator_put(apex_power_data->supply);
+	return 0;
+}
+
+static const struct of_device_id apex_power_dt_ids[] = {
+	{ .compatible = "google,apex-power", },
+	{}
+};
+MODULE_DEVICE_TABLE(of, apex_power_dt_ids);
+
+static struct platform_driver apex_power_driver = {
+	.driver = {
+		.name = "apex-power",
+		.owner = THIS_MODULE,
+		.of_match_table = of_match_ptr(apex_power_dt_ids),
+	},
+	.probe = apex_power_probe,
+	.remove = apex_power_remove,
+};
+
+static int __init apex_power_init(void) {
+	int ret;
+	dev_t apex_power_dev;
+	apex_power_class = class_create(THIS_MODULE, "apex_power");
+	if (IS_ERR(apex_power_class)) {
+		printk(KERN_ERR "apex_power: could not allocate device class\n");
+		return -ENODEV;
+	}
+	ret = alloc_chrdev_region(&apex_power_dev, 0, 255, "apex_power");
+	apex_power_major = MAJOR(apex_power_dev);
+	if (ret < 0) {
+		printk(KERN_ERR "alloc_chrdev_region() failed\n");
+		class_destroy(apex_power_class);
+		return -ENODEV;
+	}
+
+	return platform_driver_register(&apex_power_driver);
+}
+module_init(apex_power_init);
+
+static void __exit apex_power_exit(void) {
+	platform_driver_unregister(&apex_power_driver);
+	unregister_chrdev_region(MKDEV(apex_power_major, 0), 1);
+	class_destroy(apex_power_class);
+}
+module_exit(apex_power_exit);
+
+MODULE_DESCRIPTION("Google Apex power management driver");
+MODULE_AUTHOR("June Tate-Gans <jtgans@google.com>");
+MODULE_ALIAS("platform:apex-power");
+MODULE_LICENSE("GPL");