Fix bugs with HostIr integration into FusionExecutorCache

[nsarka]

This PR fixes some bugs with HostIr integration into FusionExecutorCache. The list of fixes are:

1. Expressions within the SegmentedGroup were out of order. They get sorted during gpu lowering. So, I topologically sorted them after adding them to the HostIrContainer

2. CompileFusionParallel launches compileKernel with multiple threads from a pool. However, the main thread does not wait on completion before continuing. I added fix so that it will wait. (Edit: Jingyue pointed out the wait was there already. I moved the HostIr additions from previous PRs to after the wait)

3. During the preseg pass stage, Sets are added between expressions with a different sharding. These sets needed lowering into a communication. I added the missing lowering. (Edit: this is going into a separate PR)

With these changes, many (most?) of the MultiDevice tests and also DistributedTransformer tests are passing. The failing tests seem to be because of an issue with loop splitting. Here is an example:

TEST_F(MultiDeviceTest, nick_hostir_5) {
  Fusion fusion, fusion_cloned;
  FusionGuard fg(&fusion);
  TensorView* x = makeContigTensor(4);
  x->outer_split(0, 2);
  fusion.addInput(x);
  std::cout << "x transforms: " << std::endl;
  x->printTransforms();
  fg.setCurFusion(&fusion_cloned);
  IrCloner ir_cloner(&fusion_cloned);
  auto in_clone = ir_cloner.clone(fusion.inputs());
  std::cout << "cloned_x transforms: " << std::endl;
  in_clone[0]->as<TensorView>()->printTransforms();
/*
x transforms:
 logical domain : (iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4})
 contiguity: t t t t
  Outer split: iS0{i0} by factor 2 -> iS4{2}, iS5{( ceilDiv(i0, 2) )}          <----- outer split
 loop domain : (iS4{2}, iS5{( ceilDiv(i0, 2) )}, iS1{i2}, iS2{i3}, iS3{i4})
cloned_x transforms:
 logical domain : (iS0{i0}, iS1{i2}, iS2{i3}, iS3{i4})
 contiguity: t t t t                                                                                       <----- no outer split, cloner does not propagate it
 loop domain : (iS4{2}, iS5{i6}, iS1{i2}, iS2{i3}, iS3{i4})
*/
}

After cloning exprs into the hostir container, splits are not propagated.

[wujingyue]

It's great to hear many tests are passing behind the scene. Is it possible to add some tests to this PR? For example, it could be some existing tests but with host IR lowering also enabled (cf. TEST_P) or could be some new tests to test_host_ir_integration to catch potential regression?

[nsarka]

It's great to hear many tests are passing behind the scene. Is it possible to add some tests to this PR? For example, it could be some existing tests but with host IR lowering also enabled (cf. TEST_P) or could be some new tests to test_host_ir_integration to catch potential regression?

Updated with TEST_P on tests in tests/cpp/test_multidevice_lower_communication.cpp. All of these but the allgather loop split test (which I disabled) pass.

[samnordmann]

LGTM!

[nsarka]

@wujingyue I removed the host ir lowering code and replaced it with a TODO. I also updated the tests to just pass false for enable hostir lowering. The next PR i can re-enable them, and re-add the hostir lowering code.

[nsarka]

!test

[wujingyue]

LGTM with comments

[wujingyue]

Suggested change

	public testing::WithParamInterface<std::tuple<InOutMesh, bool>> {};
	public testing::WithParamInterface<std::tuple<DeviceMesh, DeviceMesh, bool>> {};

I don't see a good reason for nesting. This will probably simplify the structured bindings below.

[nsarka]

I think this is here because we want certain pairs of DeviceMesh, and not every pair that Combine(InMesh, OutMesh) would give