Tensor contraction 4-5x slower than `cirq.unitary` for hamiltonian simulation bloq

[anurudhp]

cirq.unitary(BloqAsCirqGate(W_e_iHt)) is 4-5x faster than W_e_iHt.tensor_contract(). Code used in:

Qualtran/qualtran/bloqs/hamiltonian_simulation/hamiltonian_simulation_by_gqsp_test.py

Lines 72 to 82 in 0d9d974

	def verify_hamiltonian_simulation_by_gqsp(
	W: QubitizationWalkOperator, H: NDArray[np.complex128], *, t: float, precision: float
	):
	N = H.shape[0]
	W_e_iHt = HamiltonianSimulationByGQSP(W, t=t, precision=precision)
	result_unitary = cirq.unitary(W_e_iHt)
	expected_top_left = scipy.linalg.expm(-1j * H * t)
	actual_top_left = result_unitary[:N, :N]
	assert_matrices_almost_equal(expected_top_left, actual_top_left, atol=1e-4)

time comparison (in seconds)

Test	cirq.unitary (s)	tensor_contract (s)	factor
passed(1e-05-2-1-1)	1.29	5.83	4.52
passed(1e-05-2-2-1)	0.98	4.21	4.29
passed(1e-05-5-1-1)	1.45	6.35	4.38
passed(1e-05-5-2-1)	1.52	8.85	5.82
passed(1e-07-2-1-1)	1.13	4.31	3.81
passed(1e-07-2-2-1)	1.19	5.57	4.68
passed(1e-07-5-1-1)	1.66	7.25	4.37
passed(1e-07-5-2-1)	1.76	10.75	6.11
passed(1e-09-2-1-1)	1.40	5.46	3.90
passed(1e-09-2-2-1)	1.45	7.17	4.94
passed(1e-09-5-1-1)	2.36	9.82	4.16
passed(1e-09-5-2-1)	2.35	12.97	5.52

Discussion: #1328 (comment)

One possible optimization is to add 0 state/effect for ancilla and then tensor contract to directly get the block encoded matrix

[mpharrigan]

I could roughly reproduce the original numbers and with some preliminary optimizations I can get

quimb thing. I included a heuristic in the tensor contraction code to try to optimize the tensor network if the contraction width is sufficiently large. This is important for some large networks to make them tractable at all, but in these test cases the "optimization" was taking several seconds but naively contracting takes like 0.08s . Removing the tensor network simplifier shaves off a couple of seconds

qualtran thing. Most of the time is spent flattening the circuit. This involves map_soqs to re-wire-up the flattened version. This involves putting a lot of soquets in dictionaries. I sped up flattening by 1) cacheing an intermediate mapping dictionary rather than reconstructing it each time 2) changing the hash function of Soquet, which was itself using a lot of time... but I did this out of order so maybe (2) isn't necessary given (1).

In the end, most of the time is still spent flattening the circuit and there's no part that's taking an obviously large amount of time, i.e. we're at the whims of Python.

These examples require flattening 9 times to get to the fully-flat form. We might be able to go faster by depth-first-searching the compute graph as we build the tensor network rather than factoring it into two steps: flatten; convert to quimb; and factoring that first step into 9 steps: flatten_once nine times. This would require more work