Add a test for the Operand Upcaster HLO pass.

PiperOrigin-RevId: 738938419

Author: mrguenther

xla/hlo/transforms/tests/BUILD

@@ -13,6 +13,7 @@ lit_test_suite(
             # go/keep-sorted start
             "algebraic_simplifier.hlo",
             "cholesky_expander.hlo",
+            "operand_upcaster.hlo",
             "optimization_barrier_expander.hlo",
             "rewrite_bf16_conv_to_onednn.hlo",
             "rng_bit_generator_expander.hlo",

xla/hlo/transforms/tests/operand_upcaster.hlo

@@ -0,0 +1,79 @@
+// NOTE: Assertions have been autogenerated by hlo/tools/generate_hlo_test_checks.py
+// RUN: hlo-opt %s --split-input-file --passes=operand_upcaster | FileCheck %s
+
+// CHECK-LABEL: HloModule TestDot, entry_computation_layout={(s8[8]{0}, s16[8]{0})->s32[8]{0}}
+
+// CHECK-LABEL: ENTRY %test_dot
+// CHECK-NEXT:  %[[a:[^ ]+]] = s8[8]{0} parameter(0)
+// CHECK-NEXT:  %[[convert:[^ ]+]] = s32[8]{0} convert(%[[a]])
+// CHECK-NEXT:  %[[b:[^ ]+]] = s16[8]{0} parameter(1)
+// CHECK-NEXT:  %[[convert_1:[^ ]+]] = s32[8]{0} convert(%[[b]])
+// CHECK-NEXT:  ROOT %[[result:[^ ]+]] = s32[8]{0} dot(%[[convert]], %[[convert_1]]), lhs_contracting_dims={}, rhs_contracting_dims={}
+
+HloModule TestDot
+
+ENTRY test_dot {
+  a = s8[8] parameter(0)
+  b = s16[8] parameter(1)
+  ROOT result = s32[8] dot(a, b)
+}
+
+// -----
+
+// CHECK-LABEL: HloModule TestDotPackedNibble, entry_computation_layout={(f16[8,16]{1,0}, f16[16,8]{1,0})->f32[8,8]{1,0}}
+
+// CHECK-LABEL: ENTRY %test_dot_packed_nibble
+// CHECK-NEXT:  %[[arg_0:[^ ]+]] = f16[8,16]{1,0} parameter(0)
+// CHECK-NEXT:  %[[convert:[^ ]+]] = f32[8,16]{1,0} convert(%[[arg_0]])
+// CHECK-NEXT:  %[[arg_1:[^ ]+]] = f16[16,8]{1,0} parameter(1)
+// CHECK-NEXT:  %[[convert_1:[^ ]+]] = f32[16,8]{1,0} convert(%[[arg_1]])
+// CHECK-NEXT:  ROOT %[[dot:[^ ]+]] = f32[8,8]{1,0} dot(%[[convert]], %[[convert_1]]), lhs_contracting_dims={1}, rhs_contracting_dims={0}, operand_precision={packed_nibble,default}
+
+HloModule TestDotPackedNibble
+
+ENTRY test_dot_packed_nibble {
+  arg_0 = f16[8,16] parameter(0)
+  arg_1 = f16[16,8] parameter(1)
+  ROOT dot = f32[8,8] dot(arg_0, arg_1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, operand_precision={packed_nibble,default}
+}
+
+// -----
+
+// CHECK-LABEL: HloModule TestConvolutionPackedNibble, entry_computation_layout={(s8[3,3,7,7]{3,2,1,0}, s8[5,11,11,7]{3,2,1,0})->s32[5,11,11,7]{3,2,1,0}}
+
+// CHECK-LABEL: ENTRY %test_convolution_packed_nibble
+// CHECK-NEXT:  %[[lhs:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} parameter(1)
+// CHECK-NEXT:  %[[constant:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} broadcast(%[[constant]]), dimensions={}
+// CHECK-NEXT:  %[[shift_left:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} shift-left(%[[lhs]], %[[broadcast]])
+// CHECK-NEXT:  %[[constant_1:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast_1:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} broadcast(%[[constant_1]]), dimensions={}
+// CHECK-NEXT:  %[[shift_right_arithmetic:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} shift-right-arithmetic(%[[shift_left]], %[[broadcast_1]])
+// CHECK-NEXT:  %[[convert:[^ ]+]] = s32[5,11,11,7]{3,2,1,0} convert(%[[shift_right_arithmetic]])
+// CHECK-NEXT:  %[[rhs:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} parameter(0)
+// CHECK-NEXT:  %[[constant_3:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast_3:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} broadcast(%[[constant_3]]), dimensions={}
+// CHECK-NEXT:  %[[shift_left_1:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} shift-left(%[[rhs]], %[[broadcast_3]])
+// CHECK-NEXT:  %[[constant_4:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast_4:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} broadcast(%[[constant_4]]), dimensions={}
+// CHECK-NEXT:  %[[shift_right_arithmetic_2:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} shift-right-arithmetic(%[[shift_left_1]], %[[broadcast_4]])
+// CHECK-NEXT:  %[[convert_2:[^ ]+]] = s32[3,3,7,7]{3,2,1,0} convert(%[[shift_right_arithmetic_2]])
+// CHECK-NEXT:  %[[convolution_1:[^ ]+]] = s32[5,11,11,7]{3,2,1,0} convolution(%[[convert]], %[[convert_2]]), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f
+// CHECK-NEXT:  %[[constant_2:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast_2:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} broadcast(%[[constant_2]]), dimensions={}
+// CHECK-NEXT:  %[[shift_right_arithmetic_1:[^ ]+]] = s8[5,11,11,7]{3,2,1,0} shift-right-arithmetic(%[[lhs]], %[[broadcast_2]])
+// CHECK-NEXT:  %[[convert_1:[^ ]+]] = s32[5,11,11,7]{3,2,1,0} convert(%[[shift_right_arithmetic_1]])
+// CHECK-NEXT:  %[[constant_5:[^ ]+]] = s8[] constant(4)
+// CHECK-NEXT:  %[[broadcast_5:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} broadcast(%[[constant_5]]), dimensions={}
+// CHECK-NEXT:  %[[shift_right_arithmetic_3:[^ ]+]] = s8[3,3,7,7]{3,2,1,0} shift-right-arithmetic(%[[rhs]], %[[broadcast_5]])
+// CHECK-NEXT:  %[[convert_3:[^ ]+]] = s32[3,3,7,7]{3,2,1,0} convert(%[[shift_right_arithmetic_3]])
+// CHECK-NEXT:  %[[convolution_2:[^ ]+]] = s32[5,11,11,7]{3,2,1,0} convolution(%[[convert_1]], %[[convert_3]]), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f
+// CHECK-NEXT:  ROOT %[[add:[^ ]+]] = s32[5,11,11,7]{3,2,1,0} add(%[[convolution_1]], %[[convolution_2]])
+
+HloModule TestConvolutionPackedNibble
+
+ENTRY test_convolution_packed_nibble {
+  lhs = s8[5,11,11,7] parameter(1)
+  rhs = s8[3,3,7,7] parameter(0)
+  ROOT convolution = s32[5,11,11,7] convolution(lhs, rhs), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, operand_precision={PACKED_NIBBLE,PACKED_NIBBLE}
+}