Repository: secretflow/spu (language: python)

Issue:
[Compiler Optimization]: A pattern containing y/(k*sqrt(x)) is not optimized to y/k*rsqrt(x)
### 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 rewrite_div_sqrt_patterns optimization, however, I found a pattern which is not triggered as expected. Did I not use the compiler properly or was there existed some problem?

### Standalone code to reproduce the issue

```Python
# The original code containing y/(k*sqrt(x))
def func(X):
    return X / (X * jnp.sqrt(X)[:, jnp.newaxis])
config = spu.RuntimeConfig(protocol=spu_pb2.ProtocolKind.ABY3, field=spu_pb2.FieldType.FM64)
config.enable_pphlo_profile = True
sim = spsim.Simulator(3, config)
X = jnp.array([[4, 1, 2, 2], [1, 3, 9, 3], [5, 7, 5, 1]])
spu_fn = spsim.sim_jax(sim, func)
result = spu_fn(X)
print(result)
print(spu_fn.pphlo)

# The manually optimized code with the same result but better speed
def func(X):
    X = X.astype(jnp.float32)
    return X / X * jax.lax.rsqrt(X[:, jnp.newaxis])
config = spu.RuntimeConfig(protocol=spu_pb2.ProtocolKind.ABY3, field=spu_pb2.FieldType.FM64)
config.enable_pphlo_profile = True
sim = spsim.Simulator(3, config)
X = jnp.array([[4, 1, 2, 2], [1, 3, 9, 3], [5, 7, 5, 1]])
spu_fn = spsim.sim_jax(sim, func)
result = spu_fn(X)
print(result)
print(spu_fn.pphlo)
```


### Relevant log output

```Shell
# The profile for original code containing y/(k*sqrt(x))
[2024-02-20 09:07:41.000] [info] [api.cc:158] [Profiling] SPU execution func completed, input processing took 6.76e-07s, execution took 0.006304256s, output processing took 1.857e-06s, total time 0.006306789s.
[2024-02-20 09:07:41.000] [info] [api.cc:191] HLO profiling: total time 0.005592633
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.broadcast, executed 2 times, duration 1.7096e-05s, send bytes 0
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.convert, executed 1 times, duration 3.5814e-05s, send bytes 0
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.divide, executed 1 times, duration 0.00272042s, send bytes 24768
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.free, executed 5 times, duration 2.514e-06s, send bytes 0
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.multiply, executed 1 times, duration 0.000121707s, send bytes 1440
[2024-02-20 09:07:41.000] [info] [api.cc:194] - pphlo.sqrt, executed 1 times, duration 0.002695082s, send bytes 9072
[2024-02-20 09:07:41.000] [info] [api.cc:204] Link details: total send bytes 35280, send actions 132
.
----------------------------------------------------------------------
Ran 1 test in 0.082s

OK
[[[0.5        0.9999428  0.7070961  0.7070961 ]
  [0.4999962  0.99998856 0.7070999  0.70710754]
  [0.4999962  0.9999962  0.70710754 0.7070999 ]]

 [[0.9999962  0.57735443 0.33333206 0.5773506 ]
  [0.9999962  0.5773506  0.33333588 0.5773468 ]
  [0.9999962  0.577343   0.33333588 0.5773506 ]]

 [[0.44721603 0.37796402 0.44721985 0.9999962 ]
  [0.44721222 0.37795258 0.44721222 0.99998474]
  [0.4472084  0.37796402 0.44721222 0.99998856]]]
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<i32>>) -> tensor<3x3x4x!pphlo.secret<f32>> {
    %0 = pphlo.convert %arg0 : (tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>>
    %1 = pphlo.broadcast %0, dims = [1, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
    %2 = pphlo.sqrt %0 : tensor<3x4x!pphlo.secret<f32>>
    pphlo.free %0 : tensor<3x4x!pphlo.secret<f32>>
    %3 = pphlo.broadcast %2, dims = [0, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %2 : tensor<3x4x!pphlo.secret<f32>>
    %4 = pphlo.multiply %1, %3 : tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %3 : tensor<3x3x4x!pphlo.secret<f32>>
    %5 = pphlo.divide %1, %4 : tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %4 : tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %1 : tensor<3x3x4x!pphlo.secret<f32>>
    return %5 : tensor<3x3x4x!pphlo.secret<f32>>
  }
}

# The profile for manually optimized code with the same result but better speed
[2024-02-20 09:14:40.602] [info] [api.cc:158] [Profiling] SPU execution func completed, input processing took 8.46e-07s, execution took 0.005124268s, output processing took 1.961e-06s, total time 0.005127075s.
[2024-02-20 09:14:40.602] [info] [api.cc:191] HLO profiling: total time 0.004428919
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.broadcast, executed 2 times, duration 1.1523e-05s, send bytes 0
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.convert, executed 1 times, duration 1.8851e-05s, send bytes 0
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.divide, executed 1 times, duration 0.002416312s, send bytes 8256
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.free, executed 5 times, duration 2.126e-06s, send bytes 0
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.multiply, executed 1 times, duration 8.8427e-05s, send bytes 1440
[2024-02-20 09:14:40.602] [info] [api.cc:194] - pphlo.rsqrt, executed 1 times, duration 0.00189168s, send bytes 7152
[2024-02-20 09:14:40.602] [info] [api.cc:204] Link details: total send bytes 16848, send actions 118
.
----------------------------------------------------------------------
Ran 1 test in 0.078s

OK
[[[0.49995804 0.9999237  0.706974   0.70697784]
  [0.49996185 0.99990463 0.70697784 0.7069664 ]
  [0.49996185 0.9999237  0.706974   0.706974  ]]

 [[0.9999161  0.5772934  0.3333168  0.5772896 ]
  [0.9999199  0.57728195 0.3333168  0.57728195]
  [0.9999237  0.5772934  0.3333168  0.5772934 ]]

 [[0.44719315 0.37797165 0.44720078 0.9999237 ]
  [0.44719315 0.37796783 0.44720078 0.99991226]
  [0.44719696 0.37797165 0.44719696 0.9999237 ]]]
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<i32>>) -> tensor<3x3x4x!pphlo.secret<f32>> {
    %0 = pphlo.convert %arg0 : (tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secret<f32>>
    %1 = pphlo.divide %0, %0 : tensor<3x4x!pphlo.secret<f32>>
    %2 = pphlo.broadcast %1, dims = [1, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %1 : tensor<3x4x!pphlo.secret<f32>>
    %3 = pphlo.rsqrt %0 : tensor<3x4x!pphlo.secret<f32>>
    pphlo.free %0 : tensor<3x4x!pphlo.secret<f32>>
    %4 = pphlo.broadcast %3, dims = [0, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %3 : tensor<3x4x!pphlo.secret<f32>>
    %5 = pphlo.multiply %2, %4 : tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %4 : tensor<3x3x4x!pphlo.secret<f32>>
    pphlo.free %2 : tensor<3x3x4x!pphlo.secret<f32>>
    return %5 : tensor<3x3x4x!pphlo.secret<f32>>
  }
}
```



Discussion / hints:
Hi @zyl2016 

This is indeed a new pattern (due to jnp.newaxis inserted a broadcast). Will treat it as an enhancement.
Thank you for your prompt response! I greatly respect your willingness to acknowledge and pursue the expansion of optimization patterns!

Fix patch:
diff --git a/libspu/compiler/passes/rewrite_div_sqrt_patterns.cc b/libspu/compiler/passes/rewrite_div_sqrt_patterns.cc
index 9d88ebe1..b8c05534 100644
--- a/libspu/compiler/passes/rewrite_div_sqrt_patterns.cc
+++ b/libspu/compiler/passes/rewrite_div_sqrt_patterns.cc
@@ -12,9 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <iostream>
-#include <limits>
-
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -70,24 +67,28 @@ struct DivRewriter : public OpRewritePattern<DivOp> {
       // Pattern 2:
       // y/(k*sqrt(x)) -> y/k*rsqrt(x)
       if (auto mulOp = denominator.getDefiningOp<MulOp>()) {
-        auto sqrtOp = mulOp.getRhs().getDefiningOp<SqrtOp>();
-        auto k = mulOp.getLhs();
-        if (sqrtOp == nullptr) {
-          sqrtOp = mulOp.getLhs().getDefiningOp<SqrtOp>();
+        Value k;
+        Operation *rsqrt =
+            rewriteSqrtIfPossible(rewriter, mulOp.getLhs().getDefiningOp());
+        if (rsqrt != nullptr) {
           k = mulOp.getRhs();
+        } else {
+          k = mulOp.getLhs();
+          rsqrt =
+              rewriteSqrtIfPossible(rewriter, mulOp.getRhs().getDefiningOp());
         }
-        if (sqrtOp) {
-          // y/k
-          auto newDiv = rewriter.create<DivOp>(
-              op.getLoc(), op->getResultTypes(), op.getLhs(), k);
-          // rsqrt(x)
-          auto newRsqrt = rewriter.create<RsqrtOp>(
-              op->getLoc(), sqrtOp->getResultTypes(), sqrtOp->getOperand(0));
-          // y/k*rsqrt(x)
-          rewriter.replaceOpWithNewOp<MulOp>(op, op.getType(), newDiv,
-                                             newRsqrt);
-          return success();
+
+        // No 1/sqrt -> rsqrt rewrite, bailout
+        if (rsqrt == nullptr) {
+          return failure();
         }
+
+        auto newDiv = rewriter.create<DivOp>(op.getLoc(), op->getResultTypes(),
+                                             op.getLhs(), k);
+        rewriter.replaceOpWithNewOp<MulOp>(op, op.getType(), newDiv,
+                                           rsqrt->getResult(0));
+
+        return success();
       }
     }
     return failure();


Test patch:
diff --git a/libspu/compiler/tests/optimizations/optimize_sqrt_to_rsqrt.mlir b/libspu/compiler/tests/optimizations/optimize_sqrt_to_rsqrt.mlir
index 3488d2e3..8ea3ec5b 100644
--- a/libspu/compiler/tests/optimizations/optimize_sqrt_to_rsqrt.mlir
+++ b/libspu/compiler/tests/optimizations/optimize_sqrt_to_rsqrt.mlir
@@ -60,3 +60,16 @@ func.func @main(%arg0: tensor<3x4x!pphlo.secret<i32>>) -> tensor<3x4x!pphlo.secr
     %7 = pphlo.divide %0, %6 : tensor<3x4x!pphlo.secret<f32>>
     return %7 : tensor<3x4x!pphlo.secret<f32>>
 }
+
+// -----
+
+func.func @main(%arg0: tensor<3x4x!pphlo.secret<f32>>, %arg1: tensor<5x6x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>> {
+    // CHECK-NOT: pphlo.sqrt
+    // CHECK: pphlo.rsqrt
+    %0 = pphlo.broadcast %arg0, dims = [1, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
+    %1 = pphlo.sqrt %arg0 : tensor<3x4x!pphlo.secret<f32>>
+    %2 = pphlo.broadcast %1, dims = [0, 2] : (tensor<3x4x!pphlo.secret<f32>>) -> tensor<3x3x4x!pphlo.secret<f32>>
+    %3 = pphlo.multiply %0, %2 : tensor<3x3x4x!pphlo.secret<f32>>
+    %4 = pphlo.divide %0, %3 : tensor<3x3x4x!pphlo.secret<f32>>
+    return %4 : tensor<3x3x4x!pphlo.secret<f32>>
+  }