torch_ops_count.py

import pandas as pd
import re
import numpy as np
from sys import argv

# `xls_file` and `td files` are total_ops and generatedAtenOps.td file
# respectively.
xls_file = argv[1]
td_file = argv[2]

# Given a python list `arr` and `name` the function writes
# the list to the file.
def write_to_file(arr, name):
    textfile = open(name, "w")
    for element in arr:
        textfile.write(element + "\n")
    textfile.close()


xls_dataframe = pd.read_excel("torch.xlsx", index_col=None)

# Extracting total op name form xls as str type.
total_torch_ops = xls_dataframe["Op Name"].to_numpy().astype('str')

torch_mlir_ops = []

with open(td_file, 'r') as file:
    for line in file:
        # `def Torch_` specifies the op definition.
        if ("def Torch_" in line):
            for word in line.split("<"):
                # Everything after `aten.` and before `"` is op_name
                # according to torch converntion.
                if ("aten." in word):
                    torch_mlir_ops.append(
                        word.split(",")[0].split("aten.")[1].split('"')[0])



# Common ops are the core ops that are present in
# `td_file` as well as `xls_file`.
common_ops = []
for op in torch_mlir_ops:
    if (op in total_torch_ops):
        common_ops.append(op)

print(
    f'The total ops present in torch-mlir is {len(torch_mlir_ops)}. The names of the ops have been written to torch-mlir_ops.txt.'
)
print(
    f'The common ops present in torch-mlir and pyTorch core is {len(common_ops)}. The names of the common_ops have been written to common-torch_ops.txt.'
)

write_to_file(torch_mlir_ops, "torch-mlir_ops.txt")
write_to_file(common_ops, "common-torch_ops.txt")