code: #include <torch/csrc/autograd/profiler_utils.h>
  
  namespace torch { namespace autograd { namespace profiler {
  
  static constexpr auto kConv2dStride = 3;
  static constexpr auto kConv2dPadding = 4;
  static constexpr auto kConv2dDilation = 5;
  static constexpr auto kConv2dGroups = 6;
  
  // List of supported operators
  static constexpr auto kConv2dOp = "aten::conv2d";
  static constexpr auto kGemmOp = "aten::mm";
  static constexpr auto kMulOp = "aten::mul.Tensor";
  static constexpr auto kAddOp = "aten::add.Tensor";
  
  static constexpr auto kInputSize = "input_size";
  static constexpr auto kWeightSize = "weight_size";
  static constexpr auto kStride = "stride";
  static constexpr auto kPadding = "padding";
  static constexpr auto kDilation = "dilation";
  static constexpr auto kGroups = "groups";
  static constexpr auto kMatSize = "mat_size";
  static constexpr auto kMat1Size = "mat1_size";
  static constexpr auto kMat2Size = "mat2_size";
  
  static bool validateInput(const std::string &op_name, size_t min_size,
                         const std::vector<c10::IValue>& inputs,
                         const std::vector<int>& should_be_tensor) {
    std::stringstream ss;
    if (inputs.size() < min_size) {
        ss << "Failed to save extra arguments for flops compuation of op "
           << op_name
           << ", min size: " << min_size
           << ", actual size: " << inputs.size();
        TORCH_WARN(ss.str());
        return false;
    }
    for (auto index : should_be_tensor) {
      if (!inputs[index].isTensor()) {
        ss << "Failed to save extra arguments for flops compuation of op "
           << op_name
           << ", input[" << index
           << "] must be a tensor.";
        TORCH_WARN(ss.str());
        return false;
      }
    }
    return true;
  }
  
  std::unordered_map<std::string, c10::IValue> saveExtraArgs(const at::RecordFunction& fn) {
    // for specific types of fn, return the saved extra args for computing flops
    std::unordered_map<std::string, c10::IValue> map;
    std::vector<c10::IValue> inputs = fn.inputs();
    std::string fname(fn.name().str());
  
    if (inputs.empty()) {
      // Input shape is unavailable, return empty map
      return map;
    }
  
    if (fname == kConv2dOp) {
      std::vector<int> tensors{0, 1};
      bool check = validateInput(fname, kConv2dGroups + 1, inputs, tensors);
      if (!check) {
        return map;
      }
  
      at::Tensor input = inputs[0].toTensor();
      at::Tensor weight = inputs[1].toTensor();
      if (weight.sizes().size() != 4) {
        TORCH_WARN("Failed to compute flops for op aten::conv2d because it requires a 4D kernel tensor.");
        return map;
      }
      map[kInputSize] = at::IValue(input.sizes());
      map[kWeightSize] = at::IValue(weight.sizes());
      map[kStride] = inputs[kConv2dStride];
      map[kPadding] = inputs[kConv2dPadding];
      map[kDilation] = inputs[kConv2dDilation];
      map[kGroups] = inputs[kConv2dGroups];
    } else if (fname == kGemmOp) {
      std::vector<int> tensors{0, 1};
      bool check = validateInput(fname, 2, inputs, tensors);
      if (!check) {
        return map;
      }
  
      at::Tensor left = inputs[0].toTensor();
      at::Tensor right = inputs[1].toTensor();
      map[kMat1Size] = at::IValue(left.sizes());
      map[kMat2Size] = at::IValue(right.sizes());
    } else if (fname == kMulOp) {
      std::vector<int> tensors{0};
      bool check = validateInput(fname, 1, inputs, tensors);
      if (!check) {
        return map;
      }
  
      at::Tensor mat = inputs[0].toTensor();
      map[kMatSize] = at::IValue(mat.sizes());
    } else if (fname == kAddOp) {
      std::vector<int> tensors{0};
      bool check = validateInput(fname, 1, inputs, tensors);
      if (!check) {
        return map;
      }
  
      at::Tensor mat = inputs[0].toTensor();
      map[kMatSize] = at::IValue(mat.sizes());
    }
  
    return map;
  }
  
  uint64_t computeFlops(const std::string &op_name, const std::unordered_map<std::string, c10::IValue> &extra_args) {
    if (op_name == kConv2dOp) {
      if (extra_args.find(kInputSize) == extra_args.end()
          || extra_args.find(kWeightSize) == extra_args.end()) {
        TORCH_WARN("Calculating flops for aten::conv2d requires input_size and weight_size in saved arguments.");
        return 0;
      }
      auto input_sizes_ref = extra_args.at(kInputSize);
      auto kernel_sizes_ref = extra_args.at(kWeightSize);
      if (!input_sizes_ref.isIntList() || !kernel_sizes_ref.isIntList()) {
        TORCH_WARN("Failed to compute flops for op aten::conv2d because it requires input and weight tensor sizes.");
        return 0;
      }
  
      const std::vector<int64_t> input_sizes = input_sizes_ref.toIntVector();
      const std::vector<int64_t> kernel_sizes = kernel_sizes_ref.toIntVector();
      if (input_sizes.size() != 4 || kernel_sizes.size() != 4) {
        TORCH_WARN("Failed to compute flops for op aten::conv2d because both input and weight must be size 4.");
        return 0;
      }
      // format of the input is defined in torch.nn.quantized.functional.conv2d()
      uint64_t minibatch = 0, in_channels = 0, input_h = 0, input_w = 0;
      uint64_t out_channels = 0, kernel_h = 0, kernel_w = 0;
      const uint64_t conv2d_multiply_factor = 2;
      std::tie(minibatch, in_channels, input_h, input_w) = std::make_tuple(input_sizes[0], input_sizes[1],
                                                                           input_sizes[2], input_sizes[3]);
      std::tie(out_channels, std::ignore, kernel_h, kernel_w) = std::make_tuple(kernel_sizes[0], kernel_sizes[1],
                                                                                kernel_sizes[2], kernel_sizes[3]);
  
      // grouping is NOT properly handled yet
      return conv2d_multiply_factor * minibatch * input_h * input_w * kernel_h * kernel_w * in_channels * out_channels;
    } else if (op_name == kGemmOp) {
      if (extra_args.find(kMat1Size) == extra_args.end()
          || extra_args.find(kMat2Size) == extra_args.end()) {
        TORCH_WARN("Calculating flops for aten::mm requires mat1_size and mat2_size in saved arguments.");
        return 0;
      }
      auto mat1_sizes_ref = extra_args.at(kMat1Size);
      auto mat2_sizes_ref = extra_args.at(kMat2Size);
      if (!mat1_sizes_ref.isIntList() || !mat2_sizes_ref.isIntList()) {
        TORCH_WARN("Failed to compute flops for op aten::mm because it requires mat1_size and mat2_size to be IntList.");
        return 0;
      }
  
      std::vector<int64_t> mat1_size = mat1_sizes_ref.toIntVector();
      std::vector<int64_t> mat2_size = mat2_sizes_ref.toIntVector();
      if (mat1_size.size() == 0) {
        return 0;
      } else {
        int64_t overlap_dim = mat1_size.back();
        uint64_t flops = 1;
        for(int64_t dim : mat1_size) {
          flops *= dim;
        }
        flops /= overlap_dim;
        for(int64_t dim : mat2_size) {
          flops *= dim;
        }
        return flops;
      }
    } else if (op_name == kMulOp) {
      if (extra_args.find(kMatSize) == extra_args.end()) {
        TORCH_WARN("Calculating flops for aten::mul.Tensor requires mat_size in saved arguments.");
        return 0;
      }
      auto mat_sizes = extra_args.at(kMatSize);
      if (!mat_sizes.isIntList()) {
        TORCH_WARN("Failed to compute flops for op aten::mul because it requires mat_size to be IntList.");
        return 0;
      }
  
      std::vector<int64_t> mat_size = mat_sizes.toIntVector();
      uint64_t flops = 1;
      for(int64_t dim : mat_size) {
        flops *= dim;
      }
      return flops;
    } else if (op_name == kAddOp) {
      if (extra_args.find(kMatSize) == extra_args.end()) {
        TORCH_WARN("Calculating flops for aten::add.Tensor requires mat_size in saved arguments.");
        return 0;
      }
      auto mat_sizes = extra_args.at(kMatSize);
      if (!mat_sizes.isIntList()) {
        TORCH_WARN("Failed to compute flops for op aten::add because it requires mat_size to be IntList.");
        return 0;
      }
  
      std::vector<int64_t> mat_size = mat_sizes.toIntVector();
      uint64_t flops = 1;
      for(int64_t dim : mat_size) {
        flops *= dim;
      }
      return flops;
    }
    return 0;
  }
  
  } // namespace profiler
  } // namespace autograd
  } // namespace torch
  