Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[MHLO] Handle dynamic dimensions in HLO<->MHLO #21689

Merged
merged 1 commit into from
Jan 28, 2025
Merged

Conversation

copybara-service[bot]
Copy link

@copybara-service copybara-service bot commented Jan 22, 2025

[MHLO] Handle dynamic dimensions in HLO<->MHLO

  • Fix creating constant zero for ConvertOp HLO->MHLO translation
  • Fix broadcast in dim bounded lowering from MHLO->HLO
  • Don't use StableHLO verification methods on MHLO ReshapeOp with bounded dynamic outputs
$ cat /tmp/t.hlo 
HloModule main, entry_computation_layout={(pred[<=1801,1]{1,0})->pred[<=1801,1]{1,0}}

ENTRY %convert_with_predicate (Arg_0.1: pred[<=1801,1]) -> pred[<=1801,1] {
  %Arg_0.1 = pred[<=1801,1] parameter(0)
  ROOT %convert_pred = pred[<=1801,1] convert(%Arg_0.1)
}

$ xla-translate /tmp/t.hlo --hlo-text-to-mlir-hlo
func.func @main(%arg0: tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>> {
  %0 = mhlo.constant dense<false> : tensor<1801x1xi1>
  %1 = "mhlo.get_dimension_size"(%arg0) <{dimension = 0 : i64}> : (tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<i32>
  %2 = "mhlo.set_dimension_size"(%0, %1) <{dimension = 0 : i64}> : (tensor<1801x1xi1>, tensor<i32>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
  %3 = mhlo.compare  NE, %arg0, %2 : (tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>, tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
  return %3 : tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
}

Currently this fails when trying to create the mhlo.constant dense<false> that gets fed into compare since constants cannot have a bounded size.

@copybara-service copybara-service bot changed the title [MHLO] Handle dynamic dimensions when creating constants for ConvertOp translation [MHLO] Handle dynamic dimensions in HLO<->MHLO Jan 22, 2025
@copybara-service copybara-service bot force-pushed the test_718162448 branch 11 times, most recently from 1a2bf43 to d589f00 Compare January 28, 2025 19:39
- Fix creating constant zero for ConvertOp HLO->MHLO translation
- Fix broadcast in dim bounded lowering from MHLO->HLO
- Don't use StableHLO verification methods on MHLO ReshapeOp with bounded dynamic outputs

```
$ cat /tmp/t.hlo
HloModule main, entry_computation_layout={(pred[<=1801,1]{1,0})->pred[<=1801,1]{1,0}}

ENTRY %convert_with_predicate (Arg_0.1: pred[<=1801,1]) -> pred[<=1801,1] {
  %Arg_0.1 = pred[<=1801,1] parameter(0)
  ROOT %convert_pred = pred[<=1801,1] convert(%Arg_0.1)
}

$ xla-translate /tmp/t.hlo --hlo-text-to-mlir-hlo
func.func @main(%arg0: tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>> {
  %0 = mhlo.constant dense<false> : tensor<1801x1xi1>
  %1 = "mhlo.get_dimension_size"(%arg0) <{dimension = 0 : i64}> : (tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<i32>
  %2 = "mhlo.set_dimension_size"(%0, %1) <{dimension = 0 : i64}> : (tensor<1801x1xi1>, tensor<i32>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
  %3 = mhlo.compare  NE, %arg0, %2 : (tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>, tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>) -> tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
  return %3 : tensor<?x1xi1, #mhlo.type_extensions<bounds = [1801, ?]>>
}
```

Currently this fails when trying to create the `mhlo.constant dense<false>` that gets fed into compare since constants cannot have a bounded size.

PiperOrigin-RevId: 720661022
@copybara-service copybara-service bot merged commit ed20f12 into main Jan 28, 2025
1 check passed
@copybara-service copybara-service bot deleted the test_718162448 branch January 28, 2025 20:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant