Constraint on split_k on m * n

[afroz-oai]

Update constraints dict to accept callable for split_k and invoke that callable to determine split_k.

Spoke to Philippe Tillet offline, there are cases where we need to decide split_k dynamically based on the shapes of the operands.

Notably, in our case we want batch size invariance, but want to restrict the amount of memory that's otherwise allocated, i.e. [split_k, m, n] allocation for the scratch memory to do the bmk,kn->bmn multiplication, so pass that as an additional constraint max_allowable_mn which drop down to split_k = 1 if m * n > max_allowable_mn.

New contributor declaration

• I am not making a trivial change, such as fixing a typo in a comment.

• I have written a PR description following these
  rules.

• I have run pre-commit run --from-ref origin/main --to-ref HEAD.

• Select one of the following.

  • I have added tests.
    • /test for lit tests
    • /unittest for C++ tests
    • /python/test for end-to-end tests
  • [] This PR does not need a test because ADDED.

• Select one of the following.

  • I have not added any lit tests.
  • The lit tests I have added follow these best practices,
    including the "tests should be minimal" section. (Usually running Python code
    and using the instructions it generates is not minimal.)

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting

• @codex fix this CI failure
• @codex address that feedback