How to compute_at a stage to a not-directed related stage


#1
A = tvm.placeholder((64, 32), dtype=dtype, name="A")
C1 = tvm.compute((64, 32), lambda i, j: 7*A[i][j], name="C1")
C2 = tvm.compute((64, 64), lambda i, j: 5*A[i][j%32], name="C2")
s=create_schedule([C1.op,C2.op])

C1 and C2 are not a pair of consumer and producer, but obviousily, they have the same axis[0] derived from A.axis[0]. the fllowing IR is expected, how to generate the following IR ?

produce C2{
    for(i0,0,64){
        produce C1{
        	for(i1,0,32){
        		C1(i0,i1)=7*A(i0,i1)
	        }
	    }
        for(i2,0,32){
        	C2(i0,i2)=5*A(i0,i2)
        }
    }    
}

and we expected the interface in python to be

s[C1].compute_at(s[C2],s[C2].op.axis[0])

what key difficulties are approaching to us