#pragma once

#include "logging.h"
#include "device_functions-generated.h"
#include "model_interface.h"
#include "raii_wrapper.h"
#include "model.h"
#include "macros.h"
#include "jagged.h"
#include <algorithm>
#include <deque>
#include <fstream>
#include <iostream>
#include <string>
#include <unordered_map>
#include <math.h>
#include <iomanip>


void expand_490_constant_folding (
  const int64_t* const src, // input tensor
  int64_t* const dst, // output tensor
  const int64_t head_size, // how many times to repeat the first part of the tensor.
  cudaStream_t stream);

void gemm_rcr_bias_permute_5_constant_folding(
  void*,
  void*,
  void*,
  void*,
  uint8_t*,

    int,


  int64_t*,

  int64_t*,


  int64_t*,

  int64_t*,


  int64_t*,

  int64_t*,

  cudaStream_t
);

void split_7_constant_folding(
    void *[] /*outputs*/,
    int64_t **[] /*output_shapes*/,
    const bool [] /*output_masks*/,
    const void * /*input*/,
    const int64_t * /*input_shape*/,
    int64_t /*real_num_splits*/,
    int64_t /*all_num_splits*/,
    int64_t [] /*split_sizes*/,
    int64_t /*split_dim*/,
    int64_t /*rank*/,
    cudaStream_t stream
);

void ait_reshape_8_constant_folding(
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*,
  int64_t*
);

    void mem_eff_attention_11_constant_folding(void* output,
                   void* query,
                   void* key,
                   void* value,
                   int64_t* batch_size,
                   int64_t* seq_len_kv,
                   int64_t* seq_len_q,
                   int num_heads,
                   int head_size,
                   int head_size_v,
                   float p_dropout,
                   float softmax_scale,
                   bool is_causal,
                   bool fixed_seq_length_kv,
                   int32_t* lengths_kv,
                   bool fixed_seq_length_q,
                   int32_t* lengths_q,
                   void* workspace,
                   cudaStream_t stream);
    

namespace ait {

// Model is the class that actually performs inference. It owns memory for
// intermediate tensors and dynamic dimensions. Constants are owned by
// the model's owning container object, and input/output memory is owned
// by the user.
// Once an inference run has started, it is not safe to re-use the Model
// until the run has finished!
class ConstantFolder : public ModelBase<ConstantFolder> {
  
  

  public:
    ConstantFolder(
        size_t blob_size,
        size_t workspace_size,
        size_t unique_workspace_size,
        size_t num_inputs,
        size_t num_outputs,
        size_t num_unbound_constants,
        uint8_t* constants,
        AITemplateAllocator& allocator)
        : ModelBase(
            blob_size,
            workspace_size,
            unique_workspace_size,
            num_inputs,
            num_outputs,
            num_unbound_constants,
            constants,
            allocator) {
         constant_name_to_ptr_["pool_probe"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_probe));
     constant_name_to_ptr_["pool_mha_qkv_weight"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_mha_qkv_weight));
     constant_name_to_ptr_["pool_mha_qkv_bias"] = const_cast<const void**>(reinterpret_cast<void**>(&pool_mha_qkv_bias));
    auto* blob_ptr = static_cast<uint8_t*>(blob_.get());
        expand_490_0 = reinterpret_cast<decltype(expand_490_0)>(blob_ptr + 322486272);
    reshape_493_0 = reinterpret_cast<decltype(reshape_493_0)>(blob_ptr + 0);
    split_494_0 = reinterpret_cast<decltype(split_494_0)>(blob_ptr + 322486272);
    split_494_1 = reinterpret_cast<decltype(split_494_1)>(blob_ptr + 429981696);
    split_494_2 = reinterpret_cast<decltype(split_494_2)>(blob_ptr + 537477120);
    
         params_[1].shape_ptrs = {ParamDim(1, 1, &pool_probe_dim_0), ParamDim(1, 1, &pool_probe_dim_1), ParamDim(1152, 1152, &pool_probe_dim_2)};
     params_[2].shape_ptrs = {ParamDim(3456, 3456, &pool_mha_qkv_weight_dim_0), ParamDim(1152, 1152, &pool_mha_qkv_weight_dim_1)};
     params_[3].shape_ptrs = {ParamDim(3456, 3456, &pool_mha_qkv_bias_dim_0)};
     params_[0].shape_ptrs = {ParamDim(64, 64, &reshape_495_0_dim_0), ParamDim(729, 729, &reshape_495_0_dim_2), ParamDim(16, 16, &reshape_495_0_dim_1), ParamDim(72, 72, &reshape_497_0_dim_3)};

      
      
    }

    ~ConstantFolder() {
      
      
    }

    void SetUpInputsOutputs() {
        
if (pool_probe == nullptr) {
    throw std::runtime_error("Constant pool_probe was not set! Set the value with set_constant.");
}
    

if (pool_mha_qkv_weight == nullptr) {
    throw std::runtime_error("Constant pool_mha_qkv_weight was not set! Set the value with set_constant.");
}
    

if (pool_mha_qkv_bias == nullptr) {
    throw std::runtime_error("Constant pool_mha_qkv_bias was not set! Set the value with set_constant.");
}
    
     reshape_495_0 = split_494_0;
     reshape_496_0 = split_494_1;
     reshape_497_0 = split_494_2;
     mem_eff_attention_498_0 = static_cast<decltype(mem_eff_attention_498_0)>(params_[0].ptr);

if (mem_eff_attention_498_0 == nullptr) {
    throw std::runtime_error("Constant mem_eff_attention_498_0 was not set! Set the value with set_constant.");
}
    
    }

    void ResetConstants(uint8_t* constants) {
        /*
         * This can be called if we want to use a different piece of memory
         * for the constants to be consumed.
         */
        
    }

    void DeviceToDeviceCopies(StreamType stream) {
  
    }


    ///////////////////////////////////////////////////////////////////////////
    // default RunImpl implemenation
    void RunImpl(StreamType stream) {
        
  
  
    {
        expand_490_constant_folding(
            static_cast<int64_t*>(pool_probe),
            static_cast<int64_t*>(expand_490_0),
            46656l,
            stream);
    }
    
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    {
    

    gemm_rcr_bias_permute_5_constant_folding(

        expand_490_0,
        pool_mha_qkv_weight,

        pool_mha_qkv_bias,

        reshape_493_0,
        global_workspace_,
        1,

        &reshape_491_0_dim_0,

        &reshape_491_0_dim_1,


        &pool_mha_qkv_weight_dim_0,

        &pool_mha_qkv_weight_dim_1,


        &reshape_491_0_dim_0,

        &pool_mha_qkv_weight_dim_0,

        stream
    );
    }
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
  {

    void *outputs[] = {
      split_494_0,
      split_494_1,
      split_494_2
    };


      int64_t *split_494_0_shape[] = {
        &split_494_0_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };
      int64_t *split_494_1_shape[] = {
        &split_494_1_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };
      int64_t *split_494_2_shape[] = {
        &split_494_2_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };

    int64_t **output_shapes[] = {
      split_494_0_shape, split_494_1_shape, split_494_2_shape
    };

    const int64_t reshape_493_0_shape[] = {
      reshape_493_0_dim_0, reshape_493_0_dim_1, reshape_493_0_dim_2, reshape_493_0_dim_3, reshape_493_0_dim_4
    };

    int64_t split_sizes[] = {
      1, 1, 1
    };

    bool output_masks[] = {
      true, true, true
    };

    split_7_constant_folding(
        outputs,
        output_shapes,
        output_masks,
        reshape_493_0,
        reshape_493_0_shape,
        3/*real_num_splits*/,
        3/*all_num_splits*/,
        split_sizes,
        0/*split_dim*/,
        5/*rank*/,
        stream
    );
  }
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    ait_reshape_8_constant_folding(
        &split_494_0_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_495_0_dim_0,
        &reshape_495_0_dim_1,
        &reshape_495_0_dim_2,
        &reshape_495_0_dim_3
    );
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    ait_reshape_8_constant_folding(
        &split_494_1_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_496_0_dim_0,
        &reshape_496_0_dim_1,
        &reshape_496_0_dim_2,
        &reshape_496_0_dim_3
    );
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    ait_reshape_8_constant_folding(
        &split_494_2_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_497_0_dim_0,
        &reshape_497_0_dim_1,
        &reshape_497_0_dim_2,
        &reshape_497_0_dim_3
    );
      DeviceCheckLastError(__FILE__, __LINE__);
  
  
    mem_eff_attention_11_constant_folding(
        mem_eff_attention_498_0,
        reshape_495_0, reshape_496_0, reshape_497_0,
        &reshape_495_0_dim_0,
        &reshape_497_0_dim_2,
        &reshape_495_0_dim_2,
        16,
        72,
        72,
        0,
        0.11785113019775792,
        false,
        true,
        static_cast<int32_t*>(nullptr),
        true,
        static_cast<int32_t*>(nullptr),
        global_workspace_,
        stream
    );
    
      DeviceCheckLastError(__FILE__, __LINE__);
  
    }




    void ProfileImpl(StreamType stream, size_t iters, const std::string& filename) {
#ifdef OPTIMIZE_FOR_COMPILATION_TIME
      throw std::runtime_error("Profile is disabled, please recompile without OPTIMIZE_FOR_COMPILE_TIME flag");
#else
      std::ofstream ss(filename);
      if (!ss) {
        throw std::runtime_error(std::string("Could not open file ") + filename);
      }

      int deviceId;
      char* L2CacheSlab = nullptr;
      DevicePropertyType deviceProperties;
      GetDevice(&deviceId);
      GetDeviceProperties(&deviceProperties, deviceId);
      const size_t L2SizeInBytes = deviceProperties.l2CacheSize;
      DeviceMalloc((void**) &L2CacheSlab, L2SizeInBytes);

      ss << "{\n";
      
      {
        std::cout << "Profiling: " << "expand_490" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    {
        expand_490_constant_folding(
            static_cast<int64_t*>(pool_probe),
            static_cast<int64_t*>(expand_490_0),
            46656l,
            stream);
    }
    
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "expand_490" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"1\", \"1\", \"1152\"]]"
           << ", \"output_sizes\": " << "[[\"64\", \"729\", \"1152\"]]"
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "gemm_rcr_bias_permute_492" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    {
    

    gemm_rcr_bias_permute_5_constant_folding(

        expand_490_0,
        pool_mha_qkv_weight,

        pool_mha_qkv_bias,

        reshape_493_0,
        global_workspace_,
        1,

        &reshape_491_0_dim_0,

        &reshape_491_0_dim_1,


        &pool_mha_qkv_weight_dim_0,

        &pool_mha_qkv_weight_dim_1,


        &reshape_491_0_dim_0,

        &pool_mha_qkv_weight_dim_0,

        stream
    );
    }
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "gemm_rcr_bias_permute_492" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"46656\", \"1152\"], [\"3456\", \"1152\"], [\"3456\"]]"
           << ", \"output_sizes\": " << "[[\"46656\", \"3456\"]]"
        
          << ", \"layout\": " << "\"20314\""
        
          << ", \"shape\": " << "\"(729, 3, 16)\""
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "split_494" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
  {

    void *outputs[] = {
      split_494_0,
      split_494_1,
      split_494_2
    };


      int64_t *split_494_0_shape[] = {
        &split_494_0_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };
      int64_t *split_494_1_shape[] = {
        &split_494_1_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };
      int64_t *split_494_2_shape[] = {
        &split_494_2_dim_0, &reshape_493_0_dim_1, &reshape_493_0_dim_2, &reshape_493_0_dim_3, &reshape_493_0_dim_4
      };

    int64_t **output_shapes[] = {
      split_494_0_shape, split_494_1_shape, split_494_2_shape
    };

    const int64_t reshape_493_0_shape[] = {
      reshape_493_0_dim_0, reshape_493_0_dim_1, reshape_493_0_dim_2, reshape_493_0_dim_3, reshape_493_0_dim_4
    };

    int64_t split_sizes[] = {
      1, 1, 1
    };

    bool output_masks[] = {
      true, true, true
    };

    split_7_constant_folding(
        outputs,
        output_shapes,
        output_masks,
        reshape_493_0,
        reshape_493_0_shape,
        3/*real_num_splits*/,
        3/*all_num_splits*/,
        split_sizes,
        0/*split_dim*/,
        5/*rank*/,
        stream
    );
  }
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "split_494" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"3\", \"64\", \"16\", \"729\", \"72\"]]"
           << ", \"output_sizes\": " << "[[\"1\", \"64\", \"16\", \"729\", \"72\"], [\"1\", \"64\", \"16\", \"729\", \"72\"], [\"1\", \"64\", \"16\", \"729\", \"72\"]]"
        
          << ", \"split_sizes\": " << "\"[1, 1, 1]]\""
        
          << ", \"dim\": " << "\"0]\""
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "reshape_495" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    ait_reshape_8_constant_folding(
        &split_494_0_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_495_0_dim_0,
        &reshape_495_0_dim_1,
        &reshape_495_0_dim_2,
        &reshape_495_0_dim_3
    );
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "reshape_495" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"1\", \"64\", \"16\", \"729\", \"72\"]]"
           << ", \"output_sizes\": " << "[[\"64\", \"16\", \"729\", \"72\"]]"
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "reshape_496" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    ait_reshape_8_constant_folding(
        &split_494_1_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_496_0_dim_0,
        &reshape_496_0_dim_1,
        &reshape_496_0_dim_2,
        &reshape_496_0_dim_3
    );
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "reshape_496" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"1\", \"64\", \"16\", \"729\", \"72\"]]"
           << ", \"output_sizes\": " << "[[\"64\", \"16\", \"729\", \"72\"]]"
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "reshape_497" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    ait_reshape_8_constant_folding(
        &split_494_2_dim_0,
        &reshape_493_0_dim_1,
        &reshape_493_0_dim_2,
        &reshape_493_0_dim_3,
        &reshape_493_0_dim_4,
        &reshape_497_0_dim_0,
        &reshape_497_0_dim_1,
        &reshape_497_0_dim_2,
        &reshape_497_0_dim_3
    );
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "reshape_497" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"1\", \"64\", \"16\", \"729\", \"72\"]]"
           << ", \"output_sizes\": " << "[[\"64\", \"16\", \"729\", \"72\"]]"
        
           << " } ";
        
          ss << ",\n";
        
      }
      
      {
        std::cout << "Profiling: " << "mem_eff_attention_498" << " (" << iters << " iterations)" << std::endl;
        std::vector<std::pair<EventType, EventType>> call_events(iters);
        for (auto& [call_start, call_end] : call_events) {
          CreateEvent(&call_start);
          CreateEvent(&call_end);
        }
        for (auto& [call_start, call_end]: call_events) {
          DeviceMemset(L2CacheSlab, 0x73, L2SizeInBytes);
          EventRecord(call_start, stream);
            
    mem_eff_attention_11_constant_folding(
        mem_eff_attention_498_0,
        reshape_495_0, reshape_496_0, reshape_497_0,
        &reshape_495_0_dim_0,
        &reshape_497_0_dim_2,
        &reshape_495_0_dim_2,
        16,
        72,
        72,
        0,
        0.11785113019775792,
        false,
        true,
        static_cast<int32_t*>(nullptr),
        true,
        static_cast<int32_t*>(nullptr),
        global_workspace_,
        stream
    );
    
          EventRecord(call_end, stream);
          DeviceCheckLastError(__FILE__, __LINE__);
        }
        EventSynchronize(std::get<1>(call_events.back()));
        float milliseconds = 0.0;
        for (auto& [call_start, call_end] : call_events) {
          float call_milliseconds = 0.0;
          EventElapsedTime(&call_milliseconds, call_start, call_end);
          DestroyEvent(call_start);
          DestroyEvent(call_end);
          milliseconds += call_milliseconds;
        }
        ss << "\"" << "mem_eff_attention_498" << "\": { \"ms_per_iter\": "
           << std::setprecision(4) << (milliseconds/iters)
           << ", \"qps\": " << 1000 * iters / milliseconds
           << ", \"input_sizes\": " << "[[\"64\", \"16\", \"729\", \"72\"], [\"64\", \"16\", \"729\", \"72\"], [\"64\", \"16\", \"729\", \"72\"]]"
           << ", \"output_sizes\": " << "[[\"64\", \"729\", \"16\", \"72\"]]"
        
          << ", \"causal\": " << "\"False\""
        
           << " } ";
        
          ss << "\n";
        
      }
      
      ss << "}\n";

      DeviceToDeviceCopies(stream);
      std::cout << "AIT per op profiling finished." << std::endl;
      FreeDeviceMemory(L2CacheSlab);
#endif
    }

    static std::unique_ptr<ConstantFolder> Create(
      AITemplateAllocator& allocator,
      uint8_t* constants
    ) {
      return std::make_unique<ConstantFolder>(
          644972544,
          214990848 * (1 + 0),
          0 * (1 + 0),
          0,
          1,
          3,
          constants,
          allocator
      );
    }

  private:
   void* pool_probe {nullptr};
   void* pool_mha_qkv_weight {nullptr};
   void* pool_mha_qkv_bias {nullptr};
   void* expand_490_0 {nullptr};
   void* reshape_493_0 {nullptr};
   void* split_494_0 {nullptr};
   void* split_494_1 {nullptr};
   void* split_494_2 {nullptr};
   void* reshape_495_0 {nullptr};
   void* reshape_496_0 {nullptr};
   void* reshape_497_0 {nullptr};
   void* mem_eff_attention_498_0 {nullptr};
   int64_t pool_probe_dim_0 { 1 };
   int64_t pool_probe_dim_1 { 1 };
   int64_t pool_probe_dim_2 { 1152 };
   int64_t pool_mha_qkv_weight_dim_0 { 3456 };
   int64_t pool_mha_qkv_weight_dim_1 { 1152 };
   int64_t pool_mha_qkv_bias_dim_0 { 3456 };
   int64_t expand_490_0_dim_0 { 64 };
   int64_t expand_490_0_dim_1 { 729 };
   int64_t reshape_493_0_dim_0 { 3 };
   int64_t reshape_493_0_dim_1 { 64 };
   int64_t reshape_493_0_dim_2 { 16 };
   int64_t reshape_493_0_dim_3 { 729 };
   int64_t reshape_493_0_dim_4 { 72 };
   int64_t reshape_491_0_dim_0 { 46656 };
   int64_t reshape_491_0_dim_1 { 1152 };
   int64_t split_494_0_dim_0 { 1 };
   int64_t split_494_1_dim_0 { 1 };
   int64_t split_494_2_dim_0 { 1 };
   int64_t reshape_495_0_dim_0 { 64 };
   int64_t reshape_495_0_dim_1 { 16 };
   int64_t reshape_495_0_dim_2 { 729 };
   int64_t reshape_495_0_dim_3 { 72 };
   int64_t reshape_496_0_dim_0 { 64 };
   int64_t reshape_496_0_dim_1 { 16 };
   int64_t reshape_496_0_dim_2 { 729 };
   int64_t reshape_496_0_dim_3 { 72 };
   int64_t reshape_497_0_dim_0 { 64 };
   int64_t reshape_497_0_dim_1 { 16 };
   int64_t reshape_497_0_dim_2 { 729 };
   int64_t reshape_497_0_dim_3 { 72 };


};
} // namespace ait