batchmatmul with lhs constant to fullyconnected

[glistening]

It is preferred at least for onert to use fullyconnected if lhs is const.

onert (and similarly tflite) does not allow constant for batchmatmul lhs.

See Samsung/ONE#16064.

In this case, fullyconnected provides more efficient kernels.

I tried to emit fulyconnected instead of batchmatmul.

Perhaps, op_bmm.py

def define_node(
        self,
        node: torch.fx.Node,
    ) -> circle.Operator.OperatorT:
        args = BmmArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
        input = args.input
        mat2 = args.mat2

        // TODO: check input is constant
        // Then create FULLY_CONNECTED instead of BATCH_MATMUL

        op_index = get_op_index(
            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
        )

        inputs = [input, mat2]
        outputs = [node]

        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)

However, I am not familiar with TICO yet.

Could anyone help me?

[dayo09]

@glistening Could you let me know the original code of that part? At least, with torch function you used?

[glistening]

@dayo09

I see the pattern in circle by running decode.py in #217.

@seockho-kim also found this pattern in Gemma.

[dayo09]

I see that this simple pattern reproduces the case.

torch.mm(self.weight, rhs)

I can add a pass that converts the cosnst-lhs matmul case to fc.
However, there's things to consider.

About unsupported Const-lhs matmul case

• luci-interpreter runs the case. So, it's an onert-only restriction.

About FC <-> bmm tradeoffs

        From ONE compiler's perspective:
        - BMM uses per-tensor quantization for both rhs and lhs.
        - FC uses per-channel quantization for weight and per-tensor for input.

Thus, normally FC gives better accuracy when it's quantized. However, to convert to FC, additional transpose operation needs to be added, so it will result in low latency.

If Const-rhs matmul case, transpose operation can be folded thus low latency issue doesn't arise. Currently, we do;

lhs  rhs (const)
 |     |
  \   tr
    \ |
     fc

lhs  rhs
 |     |
  \   |
    \ |
     bmm

What I suggest: use CompileConfig feature and optionally transform to fc

lhs(const)  rhs
 |            |
  \          tr
    \        |
        fc

[glistening]

@dayo09 I really appreciate your help. I've found #341 also.

Let me write two points.

First, I am working for float circle. Thus, I have no assumption or limitation caused by quantization.

Second, as you can see the graph in Samsung/ONE#16064, the batchmamul is followed by transpose. I think these two operators can be merged in 1 fullyconnected (or linear in pytorch term).

Considering two points above, I was thinking of generating 1 fullyconnected on visiting mulmat (depending on new config variable like prefer_fc_for_const_lhs_mul) rather than introducing adding pass.

But circle_serializer seems not to read config variable. Is it some design decision?

(ADD) To keep the rule of TICO, it would be good to use #341.

My approach requires:

• pass config to serializer
• generate two nodes from 1 aten node: bmm(lhs, rhs) = fullyconnected(rhs, lhs) ^T
  • aten.bmm = circle.fc + circle.transpose

[dayo09]

pass config to serializer

Do you still need it? For what purpose..?
As I told you, altering torch aten irs is much easier in passes, not serialize. Serializing maps torch irs to corresponding circle ir.

[glistening]

Do you still need it? For what purpose..?

Perhaps I will use your PR #341 as soon as it is finished.
However, as of this writing, your PR does not produce what I want yet.

BEFORE

#344

(same)

It does not solve Samsung/ONE#16064 yet.

In addition, there are other reasons to write my own patch in my way:

1. To get the needed circle in advance before you complete the PR
2. To know the AI assistant works well.

• At a first glance, it is better than I expected. It seems to work.
• After 1 hour, it is useless, it does not succeed to finish the work.
• Maybe I need to try better AI assistant.

3. To understand more about TICO.

• For example I've found TICO already generate two (or more) circle nodes. It is the opposite of my understand yesterday. batchmatmul with lhs constant to fullyconnected #339 (comment)
• In op_mm.py, 1 aten op is serialized into 2 or more circle operators.

With https://github.com/glistening/TICO/compare/prefer_fc_over_bmm_on_const_lhs?expand=1

It solves Samsung/ONE#16064.

Also, it would be better to eliminate two transposes after fullyconnected.

I will continue on onert LLM with this model, and use your PR after it is completed.

[dayo09]

@glistening

@glistening Sorry for ambiguousness in my comment. (...)

I mean, this PR doesn't support bmm to mm YET, it needs further PR. batch-1 bmm to mm is separate feature so I planned to do it in the next pr.
Copied from #341 (comment)

Could you try out #359? It's now ready :-D

[glistening]

I've confirmed 2 transpose is removed using circle2circle.

Note

You need to specify --remove_redundant_transpose .
circle2circle does not remove redundant transposes by default. 🤔

$ circle2circle --remove_redundant_transpose {in,out}.circle

 

We can get fullyconnected using either #356 or forthcoming #359.

[dayo09]

So I think this is resolved by #359
@glistening Could we close this issue?