Feature request: a handy way of getting what operators are being used in TorchBench models

[zou3519]

Not sure if this feature should go here or elsewhere, but here it is:

As a library implementor, one question I have is, "what is the set of operations I need to add support for for my feature to get good coverage on torchbench models"?

We actually already have the mechanism to do this. Thanks to @ezyang's work on torch_function mode, we're able to, given a model, get out a list of PyTorch python APIs the model uses. @samdow's work on torch_dispatch mode lets us, given a model, get out a list of aten operators that the model uses.

Every time I want to do this, I have to:

1. figure out how to install torchbench (and debug installation issues of the packages)
2. write a script using the above mentioned mechanisms to do this

It would be cool if we autogenerated a list of operations used by the models and shoved them somewhere. Perhaps this is in scope for Operator HUD? https://docs.google.com/document/d/1esIWMJfy7behsupRxGVr4KEXNP0SRYS9pPoZbtqiXGA/edit

[zou3519]

If people are curious, here are the pytorch operations used by torchbench. It's roughly 30%-50% of the total pytorch API surface: https://docs.google.com/spreadsheets/d/19FAQINr_tB0PmEkJMStUMong4SXmDVW1z_sdckxl034/edit#gid=0

[ezyang]

The mode for doing this is very simple:

https://github.com/pytorch/pytorch/pull/75994/files#diff-3f788198eb89507c88f44b077881d577e0cdecc0bdd4d7cc3c94b4b4b2dba483R1498-R1528

    if TEST_WITH_COVERAGE_DB is not None:
        global COVERAGE_DB_CONN
        global COVERAGE_DB_FILE_ID
        con = sqlite3.connect(TEST_WITH_COVERAGE_DB, isolation_level=None)

        con.executescript("""
PRAGMA journal_mode = OFF;
PRAGMA synchronous = 0;
PRAGMA locking_mode = EXCLUSIVE;
PRAGMA temp_store = MEMORY;
CREATE TABLE IF NOT EXISTS files ( file_id INTEGER NOT NULL PRIMARY KEY, file TEXT NOT NULL UNIQUE );
CREATE TABLE IF NOT EXISTS tests ( test_id INTEGER NOT NULL PRIMARY KEY, test TEXT NOT NULL UNIQUE );
CREATE TABLE IF NOT EXISTS funcs ( func_id INTEGER NOT NULL PRIMARY KEY, func TEXT NOT NULL UNIQUE );
CREATE TABLE IF NOT EXISTS calls
(
    call_file INTEGER,
    call_test INTEGER,
    call_func INTEGER,
    FOREIGN KEY(call_file) REFERENCES files(file_id),
    FOREIGN KEY(call_test) REFERENCES tests(test_id),
    FOREIGN KEY(call_func) REFERENCES files(func_id),
    UNIQUE(call_file,call_test,call_func)
);
CREATE VIEW IF NOT EXISTS calls_view AS
    SELECT files.file, tests.test, funcs.func
    FROM calls
        INNER JOIN files ON files.file_id = calls.call_file
        INNER JOIN tests ON tests.test_id = calls.call_test
        INNER JOIN funcs ON funcs.func_id = calls.call_func;
CREATE TRIGGER IF NOT EXISTS calls_view_insert INSTEAD OF INSERT ON calls_view
BEGIN
  INSERT OR IGNORE INTO files (file) VALUES (NEW.file);
  INSERT OR IGNORE INTO tests (test) VALUES (NEW.test);
  INSERT OR IGNORE INTO funcs (func) VALUES (NEW.func);
  INSERT OR IGNORE INTO calls (call_file, call_test, call_func) VALUES (
    (SELECT file_id FROM files WHERE file = NEW.file),
    (SELECT test_id FROM tests WHERE test = NEW.test),
    (SELECT func_id FROM funcs WHERE func = NEW.func)
  );
END;
        """)
        con.execute("CREATE INDEX IF NOT EXISTS ix_func ON calls (call_func)")
        con.execute("CREATE INDEX IF NOT EXISTS ix_file_test ON calls (call_file, call_test)")
        COVERAGE_DB_CONN = con

...

class CoverageMode(torch.overrides.TorchFunctionMode):
    if hasattr(__main__, '__file__'):
        prefix = Path(__file__).parent.parent.parent.parent
        full_test_file = __main__.__file__
        test_file = os.path.relpath(full_test_file, prefix)
    else:
        test_file = '<stdin>'

    def __init__(self, conn, test_case):
        self.conn = conn
        self.test_case = test_case
        self.seen = set()

    def __torch_function__(self, func, types, args=(), kwargs=None):
        self.seen.add(func)
        if kwargs is None:
            kwargs = {}
        return func(*args, **kwargs)

    def commit(self):
        rows = [
            (
                self.test_file,
                self.test_case.id(),
                torch.overrides.resolve_name(func) or str(func)
            )
            for func in self.seen
        ]
        self.conn.executemany("INSERT INTO calls_view VALUES (?,?,?)", rows)
        self.conn.commit()
        self.seen.clear()

I'm slowly working on a rewrite of this but happy if someone takes it off my hands.

@suo says we can dump this info into https://rockset.com/ if you are running it in CI.

[FindHao]

Hi @zou3519 and @ezyang , could you please tell me how to use this CoverageMode in torchbench or any simple example?

[ezyang]

I never ended up cleaning this up. But the code should still work. Try just patching it in to a simple example and hack it until it works?

[FindHao]

@ezyang Thanks for your reply! I'm trying to add this feature via #1394 . I have a question about the number of total APIs. I'd like to compute the coverage rate api_used / total_apis. Do you think the size of torch.overrides.get_testing_overrides() is the correct number for total_apis?

[ezyang]

Yup! We even test that this list is complete.

[FindHao]

apis.txt
Here is an example of the API coverage. It is for model train and inference on GPU and CPU. the API coverage is 255 / 1332 = 19%. All used and missed APIs are listed in the above txt. @zou3519 , could you please tell me if the previous 30%-50% you mentioned is a guess or calculated by a script?