diff --git a/lib/Conversion/TorchToTcp/Misc.cpp b/lib/Conversion/TorchToTcp/Misc.cpp index a58e3072..d50b9ef7 100644 --- a/lib/Conversion/TorchToTcp/Misc.cpp +++ b/lib/Conversion/TorchToTcp/Misc.cpp @@ -82,6 +82,8 @@ class ConvertAtenBroadcastLikeOps : public OpConversionPattern { SmallVector axes; SmallVector resultShape; + ArrayRef newInputShape = + input.getType().dyn_cast().getShape(); for (int64_t i = 0; i < static_cast(newDimSizes.size()); ++i) { Value newDimSize = newDimSizes[i]; @@ -111,9 +113,12 @@ class ConvertAtenBroadcastLikeOps : public OpConversionPattern { ? true : staticDimSize != inputShape[i - newLeadingDims]; + bool isInputDimBroadcastable = newInputShape[i] == 1; // Note: The order of checks in this boolean expression matters! - if (isNewDim || isDynamicDim || - (!isDimSizePreserved && doesDimSizeChange)) { + bool isOutputDimBroadcastable = + isNewDim || isDynamicDim || + (!isDimSizePreserved && doesDimSizeChange); + if (isInputDimBroadcastable && isOutputDimBroadcastable) { axes.push_back(i); newDimSize = rewriter.create( op->getLoc(), newDimSize); diff --git a/test/Conversion/TorchToTcp/misc.mlir b/test/Conversion/TorchToTcp/misc.mlir index d6c19727..b555a4dd 100644 --- a/test/Conversion/TorchToTcp/misc.mlir +++ b/test/Conversion/TorchToTcp/misc.mlir @@ -397,6 +397,27 @@ func.func @torch.aten.broadcast_to(%arg0: !torch.vtensor<[1,2,1,2],f32>) -> !tor // ----- +// CHECK-LABEL: @torch.aten.broadcast_to_with_dynamic_dim_input( +// CHECK-SAME: %[[ARG:.*]]: !torch.vtensor<[?,2736,1],f32>) -> !torch.vtensor<[?,2736,16],f32> { +// CHECK: %[[TENSOR:.*]] = torch_c.to_builtin_tensor %[[ARG]] : !torch.vtensor<[?,2736,1],f32> -> tensor +// CHECK: %[[CONSTANT:.*]] = torch.constant.int 16 +// CHECK: %[[CAST0:.*]] = torch_c.to_i64 %[[CONSTANT]] +// CHECK: %[[BROADCAST_DIM:.*]] = arith.index_cast %[[CAST0]] : i64 to index +// CHECK: %[[AFTER_BROADCAST:.*]] = tcp.broadcast %[[TENSOR]], %[[BROADCAST_DIM]] {axes = [2]} : tensor, index -> tensor +// CHECK: %[[OUT:.*]] = torch_c.from_builtin_tensor %[[AFTER_BROADCAST]] : tensor -> !torch.vtensor<[?,2736,16],f32> +// CHECK: return %[[OUT]] : !torch.vtensor<[?,2736,16],f32> +func.func @torch.aten.broadcast_to_with_dynamic_dim_input(%arg0: !torch.vtensor<[?,2736,1],f32>) -> !torch.vtensor<[?,2736,16],f32> { + %int0 = torch.constant.int 0 + %int2736 = torch.constant.int 2736 + %int16 = torch.constant.int 16 + %0 = torch.aten.size.int %arg0, %int0 : !torch.vtensor<[?,2736,1],f32>, !torch.int -> !torch.int + %1 = torch.prim.ListConstruct %0, %int2736, %int16 : (!torch.int, !torch.int, !torch.int) -> !torch.list + %2 = torch.aten.broadcast_to %arg0, %1 : !torch.vtensor<[?,2736,1],f32>, !torch.list -> !torch.vtensor<[?,2736,16],f32> + return %2 : !torch.vtensor<[?,2736,16],f32> +} + +// ----- + // CHECK-LABEL: @torch.aten.broadcast_to_dynamic_dim( // CHECK-SAME: %[[ARG0:.*]]: !torch.vtensor<[1,2],f32>, %[[ARG1:.*]]: !torch.vtensor<[?],f32>) -> !torch.vtensor<[?,2],f32> { // CHECK-DAG: %[[ARG1_T:.*]] = torch_c.to_builtin_tensor %[[ARG1]] : !torch.vtensor<[?],f32> -> tensor