test_mte_customWorkload_new.cpp

/*
 * To be compiled with -march=armv8.5-a+memtag
 */
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <cstdlib>
#include <unistd.h>
#include <sys/auxv.h>
#include <sys/mman.h>
#include <sys/prctl.h>
// #include <fstream>
// #include <vector>
// #include <string>

/*
 * From arch/arm64/include/uapi/asm/hwcap.h
 */
#define HWCAP2_MTE              (1 << 18)

/*
 * From arch/arm64/include/uapi/asm/mman.h
 */
#define PROT_MTE                 0x20

/*
 * From include/uapi/linux/prctl.h
 */
#define PR_SET_TAGGED_ADDR_CTRL 55
#define PR_GET_TAGGED_ADDR_CTRL 56
# define PR_TAGGED_ADDR_ENABLE  (1UL << 0)
# define PR_MTE_TCF_SHIFT       1
# define PR_MTE_TCF_NONE        (0UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_SYNC        (1UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_ASYNC       (2UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_MASK        (3UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TAG_SHIFT       3
# define PR_MTE_TAG_MASK        (0xffffUL << PR_MTE_TAG_SHIFT)

#define GRANULE_SZ 16
#define MALLOC_HDR_SZ 64
#define MALLOC_ARR_LEN 175000 // this has to be atleast the #rows in the workload file!

/*
 * Insert a random logical tag into the given pointer.
 */
//         printf("IN insert_random_tag: tagged ptr: %llu \n", __val);               
#define insert_random_tag(ptr) ({                       \
        uint64_t __val;                                 \
        asm("irg %0, %1" : "=r" (__val) : "r" (ptr));   \
        __val;                                          \
})

/*
 * Set the allocation tag of 1granule on the destination address.
 */
        // printf("IN set_tag: for addr %p \n", tagged_addr);             
        // printf("Done with set_tag!\n");                                
#define set_tag(tagged_addr) do {                                      \
        asm volatile("stg %0, [%0]" : : "r" (tagged_addr) : "memory"); \
} while (0)

void printIntArr(int* arr, int len)
{
        for (int x = 0; x < len; x++)
                printf("arr[%i] : %i, ", x, arr[x]);
        printf("\n");
}

int main(int argc, char const *argv[])
{
        // Make malloc_sizes a LARGE array!
        long int total_mallocs = 0;
        long long int total_mallocs_sz = 0;
        int malloc_sizes[MALLOC_ARR_LEN]; // +ve elems: new malloc, -ve elems: del (-1*elem)th malloc, i.e. tag=0
        int malloc_ids[MALLOC_ARR_LEN];
        long int malloc_orig_offsets[MALLOC_ARR_LEN];

        int useless_arr[1000];

        // long long int first_malloc_addr = 0;
        // // open file:
        FILE *fp;
        fp = fopen("mysqld_v2_workload_20k_rwRe2New2_21175.txt", "r");
        // printf("1\n");
        // // read file: put malloc sizes & freed mallocs in malloc_sizes arr.
        if (fp)
        {
                // printf("2\n");
                char buf[256];
                while (true)
                {
                        if (fscanf(fp, "%s", buf) == EOF)
                                break;
                        if (buf[0] == 'm')
                        {
                                int msz; // char msz[256];
                                long moffset; // char moffset[256];
                                int mid; // char mid[256];
                                
                                if ((fscanf(fp, "%i", &msz) == 0) || (fscanf(fp, "%li", &moffset) == 0) || (fscanf(fp, "%i", &mid) == 0))
                                        printf("Parse: ERRORRRRRR!!!! malloc: sz / offset / id unavailable!!! \n");
                                malloc_sizes[total_mallocs] = msz;
                                malloc_orig_offsets[total_mallocs] = moffset;
                                malloc_ids[total_mallocs] = mid;

                                total_mallocs += 1;
                                total_mallocs_sz = (total_mallocs_sz > (moffset+msz)) ? (total_mallocs_sz) : (moffset+msz);

                                // if (total_mallocs%500 == 0)
                                //         printf("Parse: GOT malloc: sz %i, offset: %li, id: %i, total_mallocs: %i, total_mallocs_sz: %li \n", (msz), (moffset), (mid), total_mallocs, total_mallocs_sz);
                        }
                        else if (buf[0] == 'f')
                        {
                                long foffset; // [256];
                                int fid; // [256];

                                if ( (fscanf(fp, "%li", &foffset) == 0) || (fscanf(fp, "%i", &fid) == 0) )
                                        printf("Parse: ERRORRRRRR!!!! free: offset / id unavailable!!!! \n");
                                
                                malloc_sizes[total_mallocs] = -1 * fid;
                                total_mallocs += 1;

                                if (malloc_orig_offsets[fid] != foffset)
                                        printf("Parse: ERRORRRR!!!! free command fid: %i, foffset: %li, malloc_orig_offsets[fid]: %li offset not matching!!! \n", fid, foffset, malloc_orig_offsets[fid]);

                                // if (total_mallocs %  == 0)
                                //         printf("Parse: GOT free: id: %i, offset: %li, malloc_orig_offsets: %li \n", (fid), (foffset), malloc_orig_offsets[fid]);
                        }
                }
        }

        char *page_ptr;
        unsigned long page_sz = sysconf(_SC_PAGESIZE);
        unsigned long hwcap2 = getauxval(AT_HWCAP2);

        /* check if MTE is present */
        if (!(hwcap2 & HWCAP2_MTE))
                return EXIT_FAILURE;

        /*
         * Enable the tagged address ABI, synchronous or asynchronous MTE
         * tag check faults (based on per-CPU preference) and allow all
         * non-zero tags in the randomly generated set.
         */
        if (prctl(PR_SET_TAGGED_ADDR_CTRL,
                  PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | PR_MTE_TCF_ASYNC |
                  (0xfffe << PR_MTE_TAG_SHIFT),
                  0, 0, 0)) {
                perror("prctl() failed");
                return EXIT_FAILURE;
        }

        // need to explicitly cast since c++ doesn't allow casting void* to unsigned char* [that'd be okay in C]
        // addr (arg0), offset (arg4) need to be a multiple of page sz, but not length.
        long long int total_mmap_sz = page_sz * (total_mallocs_sz/page_sz + 1);
        
        printf("TOTAL #Events: %li sz of mallocs: %lli, did mmap sz %lli, page_ptr: %p, Starting to allocate now! \n", total_mallocs, total_mallocs_sz, total_mmap_sz, page_ptr);
        
        page_ptr = static_cast<char*>(mmap(0, total_mmap_sz, PROT_READ | PROT_WRITE,
                 MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
        if (page_ptr == MAP_FAILED) {
                perror("mmap() failed");
                return EXIT_FAILURE;
        }

        // need to enable MTE protection for the mem starting from page_ptr, of page_sz bytes.
        if (mprotect(page_ptr, total_mmap_sz, PROT_READ | PROT_WRITE | PROT_MTE)) {
                perror("mprotect() failed");
                return EXIT_FAILURE;
        }   

        unsigned long malloc_tags[total_mallocs];
        
        // In this simulation, we always use orig_offset as the start of the mem allocation.
        // This is accurate as far as tagging behavior is concerned, since all metadata is tagged with 0.
        for (int i = 0; i < total_mallocs; i++)
        {
                if (malloc_sizes[i] > 0)
                {
                        // new alloc! generate random tag and tag all granules with it:
                        malloc_tags[i] = (((char*)insert_random_tag(page_ptr + malloc_orig_offsets[i])) - page_ptr) >> 56;
                        for (int gi = 0; gi < (malloc_sizes[i]/GRANULE_SZ); gi++)
                        {
                                // if (gi % 1000 == 0)
                                if ((gi == 0) && (i%5000 == 7))
                                        printf("MALLOC=%i ABOUT to set tag to %lu for gi = %i (offset: %li, sz: %i)", i, malloc_tags[i], gi, malloc_orig_offsets[i], malloc_sizes[i]);
                                set_tag( ((page_ptr + malloc_orig_offsets[i] + (malloc_tags[i] << 56))) + gi*GRANULE_SZ);
                        }

                        // Checking access:
                        for (int bi = 0; bi < malloc_sizes[i]; bi++)
                                useless_arr[bi%1000] = *((page_ptr + malloc_orig_offsets[i] + (malloc_tags[i] << 56)) + bi) ;  
                }
                else
                {
                        // free: Just 0 all granules.
                        int free_index = -1*malloc_sizes[i];
                        if (i % 5000 == 7)
                                printf("FREE-ing i=%i index %i at offset: %li (tag: %lu, sz: %i) \n !!!", i, free_index, malloc_orig_offsets[free_index], malloc_tags[free_index], malloc_sizes[free_index]);

                        // tag mem with 0
                        malloc_tags[free_index] = 0;

                        for (int gi = 0; gi < (malloc_sizes[free_index])/GRANULE_SZ; gi++)
                                set_tag( (page_ptr + malloc_orig_offsets[free_index]) + gi*GRANULE_SZ);
                }
        }

        return 0;
}