[Torch] Fold aten.to.dtype on splat constants.
#4306
Conversation
|
Not sure who can review, maybe you would know @vivekkhandelwal1 @zjgarvey ? |
This commit teaches `AtenToDtypeOp::fold` to constant-fold dtype conversions when the operand is a splat `DenseElementsAttr`. Folding is done according to torch's rounding behavior, i.e. * Bool: 0 and -0.0 → false; nonzero/NaN/±Inf → true. * Float → Int: round toward zero. * Int → Float: sign-aware, rmNearestTiesToEven. * Float ↔ Float: use builtin `mlir::FloatType::getFloatSemantics()`. * Int ↔ Int: use `zextOrTrunc` / `sextOrTrunc` based on source signedness. Folding is only performed when `non_blocking == false`, `copy == false`, and `memory_format` is None.
mdazz
force-pushed
the
mdazz/add-todtype-folder
branch
from
3bf4e4b to
1d7b55b
Compare
| // int32 splat → float32 | ||
| %int_splat = torch.vtensor.literal(dense<42> : tensor<2x3xsi32>) : !torch.vtensor<[2,3],si32> | ||
| %int6 = torch.constant.int 6 // torch.float32 | ||
| // CHECK: %[[R1:.*]] = torch.vtensor.literal({{.*}} : tensor<2x3xf32>) : !torch.vtensor<[2,3],f32> |
There was a problem hiding this comment.
Can you put the actual value which I think here will be 42.0 ?
| -> !torch.vtensor<[4,4],si32> | ||
|
|
||
| // int64 splat (max int32) → int32 (trunc) | ||
| %int64_splat = torch.vtensor.literal(dense<2147483647> : tensor<10xsi64>) : !torch.vtensor<[10],si64> |
There was a problem hiding this comment.
Can this value be int32max+1 to ensure that trucation does happen in the IR being locked down?
| // float32 splat → float64 | ||
| %float32_splat = torch.vtensor.literal(dense<2.71828> : tensor<5x5xf32>) : !torch.vtensor<[5,5],f32> | ||
| %int7 = torch.constant.int 7 // torch.float64 | ||
| // CHECK: %[[R4:.*]] = torch.vtensor.literal({{.*}} : tensor<5x5xf64>) : !torch.vtensor<[5,5],f64> |
There was a problem hiding this comment.
Let's capture the actual value here too and other such places.
| // int32 splat → float32 | ||
| %int_splat = torch.vtensor.literal(dense<42> : tensor<2x3xsi32>) : !torch.vtensor<[2,3],si32> | ||
| %int6 = torch.constant.int 6 // torch.float32 | ||
| // CHECK: %[[R1:.*]] = torch.vtensor.literal({{.*}} : tensor<2x3xf32>) : !torch.vtensor<[2,3],f32> |
There was a problem hiding this comment.
Since we are not locking the values being returned from the output IR, I think we should add CHECK-NOT:torch.aten.to.dtype as well to ensure that the op is being folded.
| @@ -1762,6 +1762,78 @@ func.func @torch.aten.to.dtype$no_fold$unk_dtype(%arg0: !torch.tensor) -> !torch | |||
| return %0 : !torch.tensor | |||
| } | |||
|
|
|||
| // CHECK-LABEL: @torch.aten.to.dtype$fold_splat( | |||
There was a problem hiding this comment.
Can you add some e2e tests to ensure that torch's rounding logic is accurately captured in this implementation?
Also please fix the CI failures, we cannot merge until CI pipelines are green.
This commit teaches
AtenToDtypeOp::foldto constant-fold dtype conversions when the operand is a splatDenseElementsAttr.Folding is done according to torch's rounding behavior, i.e.
mlir::FloatType::getFloatSemantics().zextOrTrunc/sextOrTruncbased on source signedness.Folding is only performed when
non_blocking == false,copy == false, andmemory_formatis None.