tvm icon indicating copy to clipboard operation
tvm copied to clipboard

[Unity] Check for transpose and dynamic shape in AdjustMatmulOrder

Open Lunderberg opened this issue 1 year ago • 1 comments
trafficstars

When determining whether to evaluate matrix multiplications as (A*B)*C or as A*(B*C), dynamic shapes may occur (e.g. a dynamic LoRA rank). This commit tests for these cases, and improves the arithmetic bounds used to prove which order of evaluation is preferred.

As part of the implementation, this commit also adds a utility CollectNonNegativeExpressions, exposed to the python API as relax.analysis.collect_non_negative_expresisons. This utility collects expressions within a StructInfo which must be non-negative, based on the location where they appear. For example, the size of a tensor along each dimension must be non-negative. Unlike the existing defineable_tir_vars_in_struct_info, this will include the N-2 expression in R.Tensor([N-2]).

Lunderberg avatar Feb 16 '24 17:02 Lunderberg

The unit test tests/python/relax/test_transform_adjust_matmul_order.py::TestRHSPermuteDimsWithDynamicBatch is currently failing for this PR. In order to determine whether (A*B)*C or A*(B*C) results in fewer operations, it must make use of the improved ConstIntBounds in https://github.com/apache/tvm/pull/16588. After that PR lands, TestRHSPermuteDimsWithDynamicBatch should pass.

Lunderberg avatar Feb 16 '24 17:02 Lunderberg

Rebased on top of https://github.com/apache/tvm/pull/16735, all new unit tests passing.

Lunderberg avatar Mar 25 '24 19:03 Lunderberg