Skip to content

Commit e96561e

Browse files
Micky774tensorflower-gardener
Micky774
authored and
tensorflower-gardener
committed
PR #10433: Fixes type annotation overload of dlpack_managed_tensor_to_buffer in python/xla_extension
Imported from GitHub PR openxla/xla#10433 Encountered bug in jax-ml/jax#20175 (see this [comment](jax-ml/jax#20175 (comment))). This adjusts the stub file to properly overload `dlpack_managed_tensor_to_buffer` so that both signatures can be checked against. Copybara import of the project: -- 75cabb5149b4a0bdd9e819fac0ea6a0ba756bff6 by Meekail Zain <zainmeekail@gmail.com>: Update Merging this change closes #10433 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#10433 from Micky774:type_update 75cabb5149b4a0bdd9e819fac0ea6a0ba756bff6 PiperOrigin-RevId: 615173747
1 parent e89c731 commit e96561e

14 files changed

Lines changed: 406 additions & 210 deletions

File tree

‎tensorflow/compiler/mlir/lite/stablehlo/tests/uniform-quantized-stablehlo-to-tfl.mlir‎

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@ func.func @dot_general_upstream_srq_per_axis_quantized_filter_with_multiple_cont
478478
// to fused tflite ops.
479479
// ============================================================================
480480

481-
// Tests that a simple per-tensor quantized `stablehlo.dot_general` is properly
481+
// Tests that a simple per-channel quantized `stablehlo.dot_general` is properly
482482
// lowered to fused `tfl.fully_connected`.
483483
// This case covers for the following quantization patterns because
484484
// activation clipping ranges take affect in scale and zp of the final
@@ -488,39 +488,39 @@ func.func @dot_general_upstream_srq_per_axis_quantized_filter_with_multiple_cont
488488
// * dot_general_with_relu6_fn
489489

490490
func.func @dot_general_srq(%arg0: tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>) -> (tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>) {
491-
%0 = stablehlo.constant() {value = dense<1> : tensor<1024x3xi8>} : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32, 2.000000e+0:0>>
492-
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>, tensor<1024x3x!quant.uniform<i8<-127:127>:f32, 2.000000e+0:0>>) -> tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>
493-
%2 = stablehlo.uniform_quantize %1 : (tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
491+
%0 = stablehlo.constant() {value = dense<1> : tensor<1024x3xi8>} : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {2.000000e+0, 2.000000e+0}>>
492+
%1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>, tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {2.000000e+0, 2.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>
493+
%2 = stablehlo.uniform_quantize %1 : (tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
494494
return %2 : tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
495495
}
496496
// CHECK-LABEL: dot_general_srq
497497
// CHECK-SAME: (%[[ARG_1:.+]]: tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
498498
// CHECK-NOT: stablehlo.dot_general
499-
// CHECK: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>, value = dense<1> : tensor<3x1024xi8>} : () -> tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>
500-
// CHECK: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x!quant.uniform<i32:f32, 2.000000e+00>>, value = dense<0> : tensor<3xi32>} : () -> tensor<3x!quant.uniform<i32:f32, 2.000000e+00>>
501-
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG_1]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+00>>, tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>, tensor<3x!quant.uniform<i32:f32, 2.000000e+00>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
499+
// CHECK: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00}>>, value = dense<1> : tensor<3x1024xi8>} : () -> tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00}>>
500+
// CHECK: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00}>>, value = dense<0> : tensor<3xi32>} : () -> tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00}>>
501+
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG_1]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+00>>, tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00}>>, tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00}>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
502502
// CHECK-NOT: tfl.batch_matmul
503503
// CHECK: return %[[FULLY_CONNECTED]]
504504

505505
// -----
506506

507-
// Tests that a fused per-tensor quantized `stablehlo.dot_general` is properly
507+
// Tests that a fused per-channel quantized `stablehlo.dot_general` is properly
508508
// lowered to fused `tfl.fully_connected`.
509509
// TODO: b/309896242 - Add more support for dynamic bias fusion cases.
510510

511511
func.func @dot_general_with_bias_same_shape_srq(%arg0: tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>) -> (tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>) {
512-
%0 = stablehlo.constant() {value = dense<1> : tensor<1024x3xi8>} : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32, 2.000000e+0:0>>
513-
%1 = stablehlo.constant() {value = dense<2> : tensor<1x3xi32>} : () -> tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>
514-
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>, tensor<1024x3x!quant.uniform<i8<-127:127>:f32, 2.000000e+0:0>>) -> tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>
515-
%3 = stablehlo.add %2, %1 : tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>
516-
%4 = stablehlo.uniform_quantize %3 : (tensor<1x3x!quant.uniform<i32:f32, 2.000000e+0:-127>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
512+
%0 = stablehlo.constant() {value = dense<1> : tensor<1024x3xi8>} : () -> tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>
513+
%1 = stablehlo.constant() {value = dense<2> : tensor<1x3xi32>} : () -> tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>
514+
%2 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0] : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+0:0>>, tensor<1024x3x!quant.uniform<i8<-127:127>:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>) -> tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>
515+
%3 = stablehlo.add %2, %1 : tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>
516+
%4 = stablehlo.uniform_quantize %3 : (tensor<1x3x!quant.uniform<i32:f32:1, {2.000000e+0, 2.000000e+0, 2.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
517517
return %4 : tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
518518
}
519519
// CHECK-LABEL: dot_general_with_bias_same_shape
520520
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+00>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
521-
// CHECK-DAG: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>, value = dense<1> : tensor<3x1024xi8>} : () -> tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>
522-
// CHECK-DAG: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x!quant.uniform<i32:f32, 2.000000e+00:-127>>, value = dense<2> : tensor<1x3xi32>} : () -> tensor<3x!quant.uniform<i32:f32, 2.000000e+00:-127>>
523-
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG_0]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+00>>, tensor<3x1024x!quant.uniform<i8<-127:127>:f32, 2.000000e+00>>, tensor<3x!quant.uniform<i32:f32, 2.000000e+00:-127>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
521+
// CHECK-DAG: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>, value = dense<1> : tensor<3x1024xi8>} : () -> tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>
522+
// CHECK-DAG: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>, value = dense<2> : tensor<1x3xi32>} : () -> tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>
523+
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG_0]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x1024x!quant.uniform<i8:f32, 1.000000e+00>>, tensor<3x1024x!quant.uniform<i8<-127:127>:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>, tensor<3x!quant.uniform<i32:f32:0, {2.000000e+00,2.000000e+00,2.000000e+00}>>) -> tensor<1x3x!quant.uniform<i8:f32, 3.000000e+00:-127>>
524524
// CHECK: return %[[FULLY_CONNECTED]]
525525

526526
// -----
@@ -532,21 +532,21 @@ func.func @dot_general_with_bias_same_shape_srq(%arg0: tensor<1x1024x!quant.unif
532532
// `tfl.fully_connected` accepts a [o, i] format for rhs, which
533533
// `stablehlo.constant` op already has before the transpose.
534534

535-
// CHECK-LABEL: dot_general_upstream_srq_constant_transpose_rhs
536-
func.func @dot_general_upstream_srq_constant_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
537-
%0 = stablehlo.constant() {value = dense<1> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>
538-
%1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>
539-
%2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>
540-
%3 = stablehlo.uniform_quantize %2 : (tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
535+
// CHECK-LABEL: dot_general_srq_constant_transpose_rhs
536+
func.func @dot_general_srq_constant_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
537+
%0 = stablehlo.constant() {value = dense<1> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8:f32:0, {3.000000e+00, 3.000000e+00}>>
538+
%1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3x!quant.uniform<i8:f32:0, {3.000000e+00, 3.000000e+00}>>) -> tensor<3x2x!quant.uniform<i8:f32:1, {3.000000e+00, 3.000000e+00}>>
539+
%2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<3x2x!quant.uniform<i8:f32:1, {3.000000e+00, 3.000000e+00}>>) -> tensor<1x2x!quant.uniform<i32:f32:1, {2.000000e+00, 2.000000e+00}>>
540+
%3 = stablehlo.uniform_quantize %2 : (tensor<1x2x!quant.uniform<i32:f32:1, {2.000000e+00, 2.000000e+00}>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
541541
return %3 : tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
542542
}
543543
// CHECK-SAME: %[[ARG:.+]]: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>
544544

545545
// Checks that the `tfl.pseudo_qconst` corresponding to the `stablehlo.constant`
546546
// has the same shape.
547-
// CHECK: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {{.*}} : () -> tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>
548-
// CHECK: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {{.*}} : () -> tensor<2x!quant.uniform<i32:f32, 1.500000e+01:-23>>
549-
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>, tensor<2x!quant.uniform<i32:f32, 1.500000e+01:-23>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
547+
// CHECK: %[[QCONST_0:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<2x3x!quant.uniform<i8:f32:0, {3.000000e+00,3.000000e+00}>>, value = dense<1> : tensor<2x3xi8>} : () -> tensor<2x3x!quant.uniform<i8:f32:0, {3.000000e+00,3.000000e+00}>>
548+
// CHECK: %[[QCONST_1:.+]] = "tfl.pseudo_qconst"() {qtype = tensor<2x!quant.uniform<i32:f32:0, {1.500000e+01,1.500000e+01}>>, value = dense<0> : tensor<2xi32>} : () -> tensor<2x!quant.uniform<i32:f32:0, {1.500000e+01,1.500000e+01}>>
549+
// CHECK: %[[FULLY_CONNECTED:.+]] = "tfl.fully_connected"(%[[ARG]], %[[QCONST_0]], %[[QCONST_1]]) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<2x3x!quant.uniform<i8:f32:0, {3.000000e+00,3.000000e+00}>>, tensor<2x!quant.uniform<i32:f32:0, {1.500000e+01,1.500000e+01}>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>
550550

551551
// Also checks that the i32 -> i8 uniform quantize is absorbed into
552552
// `tfl.fully_connected`.
@@ -558,8 +558,8 @@ func.func @dot_general_upstream_srq_constant_transpose_rhs(%arg0: tensor<1x3x!qu
558558
// `stablehlo.transpose` and its operand is not a `stablehlo.constant`
559559
// (e.g. argument), the conversion to `tfl.fully_connected` doesn't happen.
560560

561-
// CHECK-LABEL: dot_general_upstream_srq_arg_transpose_rhs
562-
func.func @dot_general_upstream_srq_arg_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
561+
// CHECK-LABEL: dot_general_srq_arg_transpose_rhs
562+
func.func @dot_general_srq_arg_transpose_rhs(%arg0: tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, %arg1: tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>> {
563563
%1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<2x3x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>
564564
%2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0] : (tensor<1x3x!quant.uniform<i8:f32, 5.000000e+00:-128>>, tensor<3x2x!quant.uniform<i8:f32, 3.000000e+00:-23>>) -> tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>
565565
%3 = stablehlo.uniform_quantize %2 : (tensor<1x2x!quant.uniform<i32:f32, 2.000000e+00>>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+00:7>>

0 commit comments

Comments
 (0)