Torch-MLIR Reduction + Broadcast Pattern

[Someone117]

Torch-MLIR generates a pattern of reduce, element-wise operation, broadcast. In vector FHE, the rotate-and-reduce sum populates all slots in the vector with the sum, not just one element. It could be helpful to remove the broadcast operation, but linalg.reduce discards dimensions, forcing the broadcast. What would be a good way to symbolize a reduction that keeps the dimensions?

Sample from torch.layernorm:

module {
  func.func @main(%arg0: tensor<2x64x768xf32> {secret.secret}) -> tensor<2x64x768xf32> {
    %cst = arith.constant dense<0.000000e+00> : tensor<2x64xf32>
    %cst_0 = arith.constant dense<7.680000e+02> : tensor<2x64xf32>
    %reduced = linalg.reduce ins(%arg0 : tensor<2x64x768xf32>) outs(%cst : tensor<2x64xf32>) dimensions = [2] 
      (%in: f32, %init: f32) {
        %2 = arith.addf %in, %init : f32
        linalg.yield %2 : f32
      }
    %0 = arith.divf %reduced, %cst_0 : tensor<2x64xf32>
    %1 = tensor.empty() : tensor<2x64x768xf32>
    // Additional broadcast
    %broadcasted = linalg.broadcast ins(%0 : tensor<2x64xf32>) outs(%1 : tensor<2x64x768xf32>) dimensions = [2] 
    return %broadcasted : tensor<2x64x768xf32>
  }
}

Could become something like:

  func.func @main(%arg0: tensor<2x64x768xf32> {secret.secret}) -> tensor<2x64x768xf32> {
    %cst = arith.constant dense<0.000000e+00> : tensor<2x64x768xf32>
    %cst_0 = arith.constant dense<7.680000e+02> : tensor<2x64x768xf32>
    // Something similar that keeps dimensions
    %reduced = linalg.reduce ins(%arg0 : tensor<2x64x768xf32>) outs(%cst : tensor<2x64x768xf32>) dimensions = [2] 
      (%in: f32, %init: f32) {
        %2 = arith.addf %in, %init : f32
        linalg.yield %2 : f32
      }
    %0 = arith.divf %reduced, %cst_0 : tensor<2x64x768xf32>
    return %0 : tensor<2x64x768xf32>
  }
}

[asraa]

Good question! This could be written with a canonicalization pattern - you might be able to use https://mlir.llvm.org/docs/PDLL/ - we have a few files like this defined in tablegen (

heir/lib/Transforms/ApplyFolders/Patterns.td

Line 20 in c6d4eee

def AddEmptyTensor : Pat<

) that use PDLL.

It actually looks like there might be two patterns required, or maybe squash it into a single pattern if it always looks like reduce -> div -> broadcast. If it's reduce -> elementwise arith op -> broadcast, then maybe you have an operation that swaps the elementwise arith op with the broadcast, and then a pattern that converts broadcast(reduce(a)) to brodcasted constants and then the reduce.

But properly, I think that the issue is that linalg.reduce REQUIRES that the dimension be reduced. I think the only way to represent what you want (reduction on a dimension, but preserve the shape) would have to be represented with a general linalg.generic.

Maybe the right thing to do instead of high-level canonicalization is to write the kernel of the regular dimension-reducing linalg.reduce in convert-to-ciphertext-semantics that does the full slot sums - so that the materialization and layout reflects that, and then check in the linalg.broadcast lowering pattern in convert-to-ciphertext-semantics should end up as a trivial operation if the input and output layout attributes and types are the same. Does that make sense? That way, it's reflected that the broadcast is trivial since the FHE kernel naturally does the sum replication.

[j2kun]

I agree, it does seem like the broadcast pattern could be identified as a no-op later in the pipeline. In the relevant passes, however (layout-propagation, layout-optimization, and convert-to-ciphertext-semantics), we don't have support for linalg.broadcast at all.

So the problem seems to be: can broadcasts be supported so that their natural repetition semantics align with the usual slot repetition patterns in FHE, and be easily identified as no-ops (changing the layout formula but not requiring any repacking)?

If not, then it may be better to introduce a higher-level op (in tensor_ext) that describes the semantics at the right layer of abstraction.

I think the best next step here is to write down some concrete examples, derived from the IR above, of what the layout system is doing currently and what it should be doing.

[Someone117]

I have read a little bit more about CKKS rotations and worked out a few test cases. It looks like the linalg.reduce will only fill the entire ciphertext if we are reducing to a scalar.

I have started with an implementation, and I have a layout assignment that works for all linalg.broadcast ops, not just swapped ones. To check if the swap has happened in layout assignment, would it work to check:

• If the reduceOp is followed by a broadcastOp
• If the reduceOp is only used by the broadcastOp
• If the two ops have the same dimensions
• If the input to the broadcast is a scalar

Something like this would cause the broadcastOp to be a no-op:

module {
  func.func @main(%arg0: !secret.secret<tensor<8xf32>>) -> (!secret.secret<tensor<8xf32>>) {
    %cst = arith.constant dense<0.0> : tensor<f32>
    %cst_0 = arith.constant dense<2.0> : tensor<8xf32>
    %0 = tensor.empty() {secret.public} : tensor<8xf32>
    %1 = secret.generic(%arg0: !secret.secret<tensor<8xf32>>) {
    ^body(%input0: tensor<8xf32>):
      %reduced = linalg.reduce ins(%input0 : tensor<8xf32>) outs(%cst : tensor<f32>) dimensions = [0]
        (%in: f32, %init: f32) {
          %31 = arith.addf %in, %init {secret.public} : f32
          linalg.yield %31 {secret.public} : f32
        }
      %broadcasted = linalg.broadcast ins(%reduced : tensor<f32>) outs(%0 : tensor<8xf32>) dimensions = [0]
      %2 = arith.mulf %broadcasted, %cst_0 : tensor<8xf32>
      secret.yield %2: tensor<8xf32>
    } -> !secret.secret<tensor<8xf32>>
  return %1 : !secret.secret<tensor<8xf32>>
  }
}