Repository: secretflow/spu (language: python)

Issue:
[Compiler Optimization]:  Cannot trigger Mixed-data-type multiplication optimization
### Issue Type

Performance

### Modules Involved

SPU compiler

### Have you reproduced the bug with SPU HEAD?

Yes

### Have you searched existing issues?

Yes

### SPU Version

spu 0.7.0b0

### OS Platform and Distribution

Linux Ubuntu 22.04

### Python Version

3.10

### Compiler Version

_No response_

### Current Behavior?

Hello, sorry for disturbing you.
I am testing Mixed-data-type multiplication optimization, however, I found it not triggered as expected. Did I not use the compiler properly or was there existed some problem?

### Standalone code to reproduce the issue

```Python
def func(X, Y):
    return X * Y
config = spu.RuntimeConfig(protocol=spu_pb2.ProtocolKind.ABY3, field=spu_pb2.FieldType.FM64)
config.enable_pphlo_profile = True
config.enable_hal_profile = True
sim = spsim.Simulator(3, config)
Y = jnp.array([[4, 1, 2, 2], [1, 3, 9, 3], [5, 7, 5, 1]])
X = jnp.array([[4.0, 1.0, 2.0, 2.0], [1.0, 3.0, 9.0, 3.0], [5.0, 7.0, 5.0, 1.0]])
spu_fn = spsim.sim_jax(sim, func)
result = spu_fn(X, Y)
print(result)
print(spu_fn.pphlo)
```


### Relevant log output

```Shell
An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
[2024-02-19 08:43:08.747] [info] [api.cc:158] [Profiling] SPU execution func completed, input processing took 8.95e-07s, execution took 0.000696648s, output processing took 1.848e-06s, total time 0.000699391s.
[2024-02-19 08:43:08.747] [info] [api.cc:191] HLO profiling: total time 0.00016298700000000002
[2024-02-19 08:43:08.747] [info] [api.cc:194] - pphlo.convert, executed 1 times, duration 3.6353e-05s, send bytes 0
[2024-02-19 08:43:08.747] [info] [api.cc:194] - pphlo.free, executed 1 times, duration 5.7e-07s, send bytes 0
[2024-02-19 08:43:08.747] [info] [api.cc:194] - pphlo.multiply, executed 1 times, duration 0.000126064s, send bytes 288
[2024-02-19 08:43:08.747] [info] [api.cc:191] HAL profiling: total time 0.000148337
[2024-02-19 08:43:08.747] [info] [api.cc:194] - f_mul, executed 1 times, duration 0.00011972s, send bytes 288
[2024-02-19 08:43:08.747] [info] [api.cc:194] - int2fxp, executed 1 times, duration 2.8617e-05s, send bytes 0
[2024-02-19 08:43:08.747] [info] [api.cc:191] MPC profiling: total time 0.000114704
[2024-02-19 08:43:08.747] [info] [api.cc:194] - lshift_a, executed 1 times, duration 1.7906e-05s, send bytes 0
[2024-02-19 08:43:08.747] [info] [api.cc:194] - mul_aa, executed 1 times, duration 3.4087e-05s, send bytes 96
[2024-02-19 08:43:08.747] [info] [api.cc:194] - trunc_a, executed 1 times, duration 6.2711e-05s, send bytes 192
[2024-02-19 08:43:08.747] [info] [api.cc:204] Link details: total send bytes 288, send actions 3
.
----------------------------------------------------------------------
Ran 1 test in 0.094s

OK
[[16.  1.  4.  4.]
 [ 1.  9. 81.  9.]
 [25. 49. 25.  1.]]
module @xla_computation_func attributes {mhlo.cross_program_prefetches = [], mhlo.is_dynamic = false, mhlo.use_auto_spmd_partitioning = false} {
  func.func @main(%arg0: tensor<3x4x!pphlo.secret<f32>>, %arg1: tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>> {
    %0 = pphlo.convert %arg1 : (tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>>
    %1 = pphlo.multiply %arg0, %0 : tensor<3x4x!pphlo.secret<f32>>
    pphlo.free %0 : tensor<3x4x!pphlo.secret<f32>>
    return %1 : tensor<3x4x!pphlo.secret<f32>>
  }
}
```



Discussion / hints:
Hi @zyl2016 

Looks like a bug to me, investigating...

Fix patch:
diff --git a/libspu/dialect/pphlo/types.cc b/libspu/dialect/pphlo/types.cc
index b31e96fb..508474d9 100644
--- a/libspu/dialect/pphlo/types.cc
+++ b/libspu/dialect/pphlo/types.cc
@@ -22,17 +22,20 @@ namespace utils {
 
 // A tiny utility to handle aggregate type like tensor
 template <typename Fn, typename... Args>
-auto StripAllContainerType(const Type &t, Fn foo, Args &&...args) {
+bool StripAllContainerType(const Type &t, Fn foo, Args &&...args) {
   if (const auto &rt = t.dyn_cast<RankedTensorType>()) {
-    return foo(rt.getElementType(), std::forward<Args>(args)...);
+    return StripAllContainerType(rt.getElementType(), foo,
+                                 std::forward<Args>(args)...);
   }
 
   if (const auto &ct = t.dyn_cast<ComplexType>()) {
-    return foo(ct.getElementType(), std::forward<Args>(args)...);
+    return StripAllContainerType(ct.getElementType(), foo,
+                                 std::forward<Args>(args)...);
   }
 
   if (const auto &st = t.dyn_cast<SecretType>()) {
-    return foo(st.getBaseType(), std::forward<Args>(args)...);
+    return StripAllContainerType(st.getBaseType(), foo,
+                                 std::forward<Args>(args)...);
   }
 
   return foo(t, std::forward<Args>(args)...);


Test patch:
diff --git a/libspu/compiler/tests/optimizations/lower_mixed_type_op.mlir b/libspu/compiler/tests/optimizations/lower_mixed_type_op.mlir
index 15e13b3c..a2c5d82d 100644
--- a/libspu/compiler/tests/optimizations/lower_mixed_type_op.mlir
+++ b/libspu/compiler/tests/optimizations/lower_mixed_type_op.mlir
@@ -42,3 +42,12 @@ func.func @main(%arg0: tensor<2x2xf16>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf
     %1 = pphlo.dot %arg1, %0 : (tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<2x2xf32>
     return %1 : tensor<2x2xf32>
 }
+
+// -----
+
+func.func @main(%arg0: tensor<3x4x!pphlo.secret<f32>>, %arg1: tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>> {
+    //CHECK: pphlo.multiply %arg0, %arg1
+    %0 = pphlo.convert %arg1 : (tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>>
+    %1 = pphlo.multiply %arg0, %0 : tensor<3x4x!pphlo.secret<f32>>
+    return %1 : tensor<3x4x!pphlo.secret<f32>>
+}