static int ext2_iomap_begin(struct inode *inode, loff_t offset, loff_t length,
                unsigned flags, struct iomap *iomap)
 {
+       struct block_device *bdev;
        unsigned int blkbits = inode->i_blkbits;
        unsigned long first_block = offset >> blkbits;
        unsigned long max_blocks = (length + (1 << blkbits) - 1) >> blkbits;
                return ret;
 
        iomap->flags = 0;
-       iomap->bdev = inode->i_sb->s_bdev;
+       bdev = inode->i_sb->s_bdev;
+       iomap->bdev = bdev;
        iomap->offset = (u64)first_block << blkbits;
+       if (blk_queue_dax(bdev->bd_queue))
+               iomap->dax_dev = dax_get_by_host(bdev->bd_disk->disk_name);
+       else
+               iomap->dax_dev = NULL;
 
        if (ret == 0) {
                iomap->type = IOMAP_HOLE;
 ext2_iomap_end(struct inode *inode, loff_t offset, loff_t length,
                ssize_t written, unsigned flags, struct iomap *iomap)
 {
+       put_dax(iomap->dax_dev);
        if (iomap->type == IOMAP_MAPPED &&
            written < length &&
            (flags & IOMAP_WRITE))
 
 static int ext4_iomap_begin(struct inode *inode, loff_t offset, loff_t length,
                            unsigned flags, struct iomap *iomap)
 {
+       struct block_device *bdev;
        unsigned int blkbits = inode->i_blkbits;
        unsigned long first_block = offset >> blkbits;
        unsigned long last_block = (offset + length - 1) >> blkbits;
        }
 
        iomap->flags = 0;
-       iomap->bdev = inode->i_sb->s_bdev;
+       bdev = inode->i_sb->s_bdev;
+       iomap->bdev = bdev;
+       if (blk_queue_dax(bdev->bd_queue))
+               iomap->dax_dev = dax_get_by_host(bdev->bd_disk->disk_name);
+       else
+               iomap->dax_dev = NULL;
        iomap->offset = first_block << blkbits;
 
        if (ret == 0) {
        int blkbits = inode->i_blkbits;
        bool truncate = false;
 
+       put_dax(iomap->dax_dev);
        if (!(flags & IOMAP_WRITE) || (flags & IOMAP_FAULT))
                return 0;
 
 
        int                     nimaps = 1, error = 0;
        bool                    shared = false, trimmed = false;
        unsigned                lockmode;
+       struct block_device     *bdev;
 
        if (XFS_FORCED_SHUTDOWN(mp))
                return -EIO;
        }
 
        xfs_bmbt_to_iomap(ip, iomap, &imap);
+
+       /* optionally associate a dax device with the iomap bdev */
+       bdev = iomap->bdev;
+       if (blk_queue_dax(bdev->bd_queue))
+               iomap->dax_dev = dax_get_by_host(bdev->bd_disk->disk_name);
+       else
+               iomap->dax_dev = NULL;
+
        if (shared)
                iomap->flags |= IOMAP_F_SHARED;
        return 0;
        unsigned                flags,
        struct iomap            *iomap)
 {
+       put_dax(iomap->dax_dev);
        if ((flags & IOMAP_WRITE) && iomap->type == IOMAP_DELALLOC)
                return xfs_file_iomap_end_delalloc(XFS_I(inode), offset,
                                length, written, iomap);
 
        u16                     type;   /* type of mapping */
        u16                     flags;  /* flags for mapping */
        struct block_device     *bdev;  /* block device for I/O */
+       struct dax_device       *dax_dev; /* dax_dev for dax operations */
 };
 
 /*