Properly lower add and mul

[wonjoolee95]

Partly fixes #6589

[wonjoolee95]

I noticed that one of the SPMD unit tests that checks HLOs failed previously and seems like the HLOs are slightly changed due to this PR because it updates the lowering logic. In this simple SPMD unit test, the HLO without this change is:

ENTRY %IrToHlo.14 (p0.4: f32[1,128], p1.5: f32[1,128]) -> (f32[1,128]) {
  %p1.5 = f32[1,128]{1,0} parameter(1), sharding={replicated}
  %p0.4 = f32[1,128]{1,0} parameter(0), sharding={replicated}
  %constant.3 = f32[] constant(1)
  %broadcast.6 = f32[1,128]{1,0} broadcast(f32[] %constant.3), dimensions={}
  %multiply.7 = f32[1,128]{1,0} multiply(f32[1,128]{1,0} %p0.4, f32[1,128]{1,0} %broadcast.6)
  %add.8 = f32[1,128]{1,0} add(f32[1,128]{1,0} %p1.5, f32[1,128]{1,0} %multiply.7)
  %custom-call.9 = f32[1,128]{1,0} custom-call(f32[1,128]{1,0} %add.8), custom_call_target="Sharding", sharding={replicated}
  %constant.2 = f32[] constant(0)
  %constant.1 = f32[] constant(1)
  %multiply.10 = f32[] multiply(f32[] %constant.2, f32[] %constant.1)
  %broadcast.11 = f32[1,128]{1,0} broadcast(f32[] %multiply.10), dimensions={}
  %add.12 = f32[1,128]{1,0} add(f32[1,128]{1,0} %custom-call.9, f32[1,128]{1,0} %broadcast.11)
  ROOT %tuple.13 = (f32[1,128]{1,0}) tuple(f32[1,128]{1,0} %add.12)
}

And the updated HLO looks like:

ENTRY %IrToHlo.14 (p0.5: f32[1,128], p1.8: f32[1,128]) -> (f32[1,128]) {
  %p1.8 = f32[1,128]{1,0} parameter(1), sharding={replicated}
  %p0.5 = f32[1,128]{1,0} parameter(0), sharding={replicated}
  %constant.4 = f32[] constant(1)
  %broadcast.6 = f32[1,128]{1,0} broadcast(f32[] %constant.4), dimensions={}
  %multiply.7 = f32[1,128]{1,0} multiply(f32[1,128]{1,0} %p0.5, f32[1,128]{1,0} %broadcast.6)
  %add.9 = f32[1,128]{1,0} add(f32[1,128]{1,0} %p1.8, f32[1,128]{1,0} %multiply.7)
  %custom-call.10 = f32[1,128]{1,0} custom-call(f32[1,128]{1,0} %add.9), custom_call_target="Sharding", sharding={replicated}
  %constant.2 = f32[] constant(0)
  %constant.1 = f32[] constant(1)
  %multiply.3 = f32[] multiply(f32[] %constant.2, f32[] %constant.1)
  %broadcast.11 = f32[1,128]{1,0} broadcast(f32[] %multiply.3), dimensions={}
  %add.12 = f32[1,128]{1,0} add(f32[1,128]{1,0} %custom-call.10, f32[1,128]{1,0} %broadcast.11)
  ROOT %tuple.13 = (f32[1,128]{1,0}) tuple(f32[1,128]{1,0} %add.12)
}

The contents of the HLOs are the same, the only things that differ is the suffix numbers, ex %custom-call.9 -> %custom-call.10.

Synced offline with @yeounoh, this is mostly fine as it's just the numbering changed in the HLO.

[wonjoolee95]

Previously CI was all green, the most recent commit just rebases with master. I'll merge this now to make the rc1 branch cut.