Add CHLO decomposition pipeline and new interpreter check calls (#2439)

These are changes in preparation for refreshing all of `testdata` and
adding CHLO-based testdata files.

- Add more check ops to supported custom_calls in translate --
`@check.expect_close/almost_equal/check_eq`
- Add `--chlo-pre-serialization-pipeline` to be run on programs with
CHLO ops to decompose them along with their shape computation to
StableHLO ops.

Note: This is based on #2438
for the changes to CheckOps in that PR.

BUILD.bazel

@@ -1076,6 +1076,7 @@ cc_library(
         "@llvm-project//mlir:InferTypeOpInterface",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
         "@llvm-project//mlir:ShapeDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",

stablehlo/dialect/StablehloOps.cpp

@@ -2481,9 +2481,13 @@ LogicalResult TupleOp::inferReturnTypes(
 void CompareOp::build(OpBuilder& builder, OperationState& result, Value lhs,
                       Value rhs, ComparisonDirection comparisonDirection,
                       ComparisonType compareType) {
+  ComparisonTypeAttr comparisonTypeAttr;
+  if (compareType != ComparisonType::NOTYPE)
+    comparisonTypeAttr =
+        ComparisonTypeAttr::get(builder.getContext(), compareType);
   build(builder, result, lhs, rhs,
         ComparisonDirectionAttr::get(builder.getContext(), comparisonDirection),
-        ComparisonTypeAttr::get(builder.getContext(), compareType));
+        comparisonTypeAttr);
 }
 
 LogicalResult CompareOp::inferReturnTypeComponents(

stablehlo/testdata/sinh_float32_20_20_chlo.mlir

@@ -0,0 +1,22 @@
+// RUN: stablehlo-opt --chlo-pre-serialization-pipeline -inline %s | stablehlo-translate --interpret
+// RUN: stablehlo-opt --chlo-pre-serialization-pipeline %s | stablehlo-translate --serialize --target=current | stablehlo-translate --deserialize | stablehlo-opt > %t.0
+// RUN: stablehlo-opt --chlo-pre-serialization-pipeline %s > %t.1
+// RUN: diff %t.0 %t.1
+
+module @jit_main attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
+  func.func public @main() -> (tensor<20x20xf32> {jax.result_info = "", mhlo.layout_mode = "default"}) {
+    %0 = call @inputs() : () -> tensor<20x20xf32>
+    %1 = call @expected() : () -> tensor<20x20xf32>
+    %2 = chlo.sinh %0 : tensor<20x20xf32> -> tensor<20x20xf32>
+    stablehlo.custom_call @check.expect_almost_eq(%2, %1) {has_side_effect = true} : (tensor<20x20xf32>, tensor<20x20xf32>) -> ()
+    return %2 : tensor<20x20xf32>
+  }
+  func.func private @inputs() -> (tensor<20x20xf32> {mhlo.layout_mode = "default"}) {
+    %cst = stablehlo.constant dense<"0x8058843F8EC81BBF85E729408B75EBBE7417C23F584AA6BE030463BF300BA740131FA1C03EB377BE5F89ABC00369A7404DF865BF5142A0C0DDC3C8BF78D5663F662219C0541752C0A14F9EC047AC533E4F13293FFC9C01C0113B944029DC98BF70CA8340CE98953E07701F404BD25CBF396BF2BF15E1B9C0CC6923400ABC373EB8F9494065BF90400FDB3340ABF3B63F9D2C47408E651040C4A7203F89F74DBF84E678C020A709C08386534043E7D1BF0F7B93C0CD034E40084DC63FBE25DBC0AB4679C030D0F6BEE6C984C060445BC045A6E83E98E70BC1203918C0811F6BBF92D6B0BECE943FC022A41740C47902C09837F73E9EEC6940859E22C0C101233FC15A9BBF04F801C0A9F28ABF98560C402320823FBA33A9BFE62708C0A5EE0B40250BEEBF699559C09514283F1076A140989F9FBFAA1D4240FDFD5FC007B291BC2B3A30C02F226CBFBCD520BFD759C93F43AC04417E0ACDC0AB2287C0AAA8CD3E7736203ED6919EC09DCB2740C7B1B2C0188052C0819EE8C0F7A3B13F1D5DA940847901C06FF6013EB413C040EE3E0EC0E1BC0640F1F050C06F68B2BF5B33C6C0AEFC62C0EF2A48C026211640754B023E5DF669C03571CB3F4DA566C064E6463E971AC0BF79ED91403E4723C0924F64C0B2DC47C0BE393C3F44A108C0934CF53E70CBAC40BF14C93F4B97BEBEDA92CF3FFCE30EC0979D61BD086B533F94AC413FA8D8B6C081240B404149A6C01A6C65BE40A484C0CC4E804008FCA040F544A4BF7D05BEBF8EF94040A99FFDBF2DD445C018725A40E6AD164094428140BF263D40C44FD33EB09C1C40DFF99A3F1EDB9CBF1E3E1E400E7AD44021131DC006EF59400726F13E3A21D63ED9E143C08F8662BF1EFC5040856B9C40CDD461BF16DFB33F6C645940D8C84EBFC3C5FE3F50316C3F04813ABF0E620B41E1652EC0F27001C003C170BF71AE034021DC2FC015B662C01D8FA3BF22056E407861FABEB73188C0F51B05C0ECCDE7BF82500A3FB0ABAD3FC01F093F3F880F402ABC7EC058DEE4C0C90CEE405F36A63F50BB773FDFB50E40DF165940F9E0C13E8DC5CB3E644C4CC0E4E007BF06A9534050065EBE8008BDBF4232D8BFC6A14240B1342740C8BAA53E5D49833F8C0984BF7E8308C04D84C6BF6CE8E03FC0AA994068B78F4046EB5C40E63FFE3FE8EA3C3FBF9459C090974FBFC7041A4084C83240BD4AF43D05E236C0EBB072408B40B23F52B61DC08317AFC01876F73FAD001EC01261DDC007356040AED0614064B695401880863F9FB9943F4F8D6A3F218842409020B93FB9BDC73F65E1463F96E6F53FBBAF103FCC5D64BF343A9740046B7ABFA58D24C0C2AF5040DB9F354042AD5CBF9C4A4BC0AFDC8440EED40440A8CE343F7EC4E03F9FF8C64015B227C003B98F402539113F060725C09529A6C07E1B5540140B7640CFB88C3FE9D549BFE61429407A8812C01BE5FC3F92D40740FE12B8C02B20A23F82795C40D3393DC020871CC00ABF0940422633BF3974DCBFAE5811C0CE1E37C0DA2732BE54A2023E54CCF93E2E4B45C0BCE1893F274E49C09B0E1BC03C7FB9BF5B7F0840D2EB49BF43A0FF3ECBA79C3F313A88C05B1A1FC0C6D43A3FA90B1CC1B8E513C055418B40A20398409285924021CCF23EBB79D6BF042A153F2451C03F245099403B74FC3F85BE06C07E62F8BFA851AD40632402BFD310B6BF47C99440CD2872C063E764C05FBDB1BFB6745EBF434F6040C3D026BD62D8004152A69E3FD346E5BFB1CE18BE479A52BD69A4A43F581466404765BF3F241979C06F6DBAC08DCD143F93460DBE8A839ABFB16A4F3FBED4A73F03DE4DC0C25605C0C7143940615033C048C7A1C0079224C1CA111840BA4B07BF993489C043024F40C73279C02EEB08403F2E3BC0F1884F3F60E3423FA27DA7C0B4ED37C0824091BFF96E85C06B16013F3749E43EDC8D9CC05281E03F190386C0BFB9B9BFE859053FC0D4EC3FC9F19AC0D0BDAF3FC5CCB9C0FE434140441D85400465EFBF0DF0FEBF57842240750F4B4036EF0640F7D4743F18DA9DBFC7598B3EBB4CCFBFA01DAABF9C8880BD3CC86DC0587395C079A0CC3FAB0D514034EAF740E03CA1BD2BE402C07ABC5A4029BC90BF278A5FC0CE6A8B3F0E469340F9D355C0F54F1B40CE4A3140B8C887402B5983C0D0836CBF8D9041BFFB7921C05B83E2BF511C874018C627407FFE90C078B4813E477C05408FC88840EA9D4B401551DBBF7C4656BEBE40C840B2B59740A583CC3E"> : tensor<20x20xf32>
+    return %cst : tensor<20x20xf32>
+  }
+  func.func private @expected() -> (tensor<20x20xf32> {mhlo.layout_mode = "default"}) {
+    %cst = stablehlo.constant dense<"0x32389D3FE09325BF4D6BE240CED8F3BE9AC00A40BA3AA9BE04FB80BF3EF3B84222B399C2741F7ABEA6D4D4C2AB14BB420E1683BFE09C95C23FE912C089B4833F55A1ADC00DE054C1A6C68CC2002F553E51A3353F40436EC05B764D42D8DFBFBF1BCBF54108BC973EBEE3BF40FD3D79BFD8D44FC07D9826C33D5ACC40D4B8383E37723B41B1443842306E044120EDFB3F9563334159129740FE682B3F10EB64BF3661C3C1E59B87C09AB659415FBC1EC0A2B348C287B3474169DA0F402795EBC3DB87C4C1F23D00BF4495FDC161C775C110BDF03E8E0944C54C22ABC00DCD86BF0660B4BE4A3C1FC1CC8EA940E39B71C0C877003FB7939A4145CDC9C0543F2E3FF767C4BF1AA36FC0F1E3A7BF93938D40FDB9993FD5F7DEBF006384C08DA88C40A18448C06B626FC1566B343F35579B42DB55CCBF36B32541225384C1FEB391BC2D27FAC0268987BF54A02BBF66A513408980F944259997C3AA7108C2E43BD33E00DE203E1FEB8DC25E01DB40881705C39E3E56C1A87333C40D6AF03FA6DBC642ABBA6DC0D24F023EEB324A435BF791C0B5668140DA1251C1370DF2BFD6D774C3E5AC8AC1D33436C17587A54088A5023E4BAB9AC1884B1640E0D892C12B27483E676508C00D313F42ADEACBC013948DC1C15535C12AA64D3F376685C0A5CAFE3E765DDD429E4E1340F105C3BE94A51B40A47A93C0CCBA61BD8B486C3F60AF543FE68717C3ABE38A40779FB4C2AE5867BE756BFCC14C6EDC412A0B99423C39D5BF7FF605C01CC222414FB063C064A42FC149A37241CCFBA640EA17E341CB451941A45CD93E317CB740BEB6C33F682BC7BF614ABC401141BF431CD6B8C046B37041EE29FA3E456DDC3EB25B2AC107A280BF7B3751416DB28442242480BFB932F53F0DAB6E41240466BF18D265402C94873F08724BBFC6BF3D45CD0BF3C0B5996DC036EC8ABF6C5C7640ACB5F8C0FF138AC1C4D8D3BF56CEA441793D02BF35090DC2F4197CC0197E3EC08A24113FD014E83F9BC60F3FDFFF94402D0BD6C1C69A1FC4D5A45444FE05D93FAF1F903F0E0E9340E2896D41A88BC63EA431D13E3F6642C156590EBF712C5A41B3C45FBE14D204C0765827C03A0C2741CEF9D8408FA3A83E1F8C9B3F42BB9CBF742685C0731E10C0C5EF3340838373426B6C324203387C4116D964409B894E3FEF5F6FC1581A67BF1A16B040B336024124DFF43DF4E30AC1614CB141EEB7F13FC9B6BAC04CD7EDC2CA8F58405D93BBC06699FCC355C58441B02788419D33574299A7A03FB385B83FFA62863F0AC92641F661004044A21140887F5B3F13D355406683183FECF081BFD7A261428C2792BFD111D0C0B63D5041022B0841520A79BFB85A3FC15F2AFE41E6F97A40F837443F31BA334084CF7A4323A9DAC05F7532427023193F88A2D1C007EDB3C2222D5F4182D7BA4126D5AA3FD4675FBF887FDF40CB4F9CC08B5862401BB2834058751DC39716D13F52787A41C27319C17A3DB7C0FECF87404D2242BF1D662DC0426199C0FB680BC1480E33BE1CFD023ED6E9013F6C2C2EC14523A63F847C39C16F03B3C032CC00C09A1D8540EA845FBF8830053F8ACCC63F992E0DC2FFDFBEC00FDD4B3F486606C69EBF9FC0DE321B42FA4067427CC642420E00FC3EB2F324C0BBBF1D3FA8A50840B3D5704277896140206A81C04C325AC0A005E14214D207BFE1F9F9BF980F51427AD4AFC103E68EC120A0F0BF45877BBFE5FB844191DC26BDB56AC444CE81CA3F9D8F3AC0086019BE07B252BDEDF2D53F3F8D91418490074021FCC3C17B7929C329531D3F6AB90DBEE1DEC2BFE8DD663F1E3ADC3F5E3D47C157097DC086C60F415A4E03C178E39CC21AFA64C66AB7AA40C1AE0DBFE09011C245D64A41C34AC4C1FF0486406C9914C1A506673F0E44563F628DBBC27C310DC1FB7EB2BFA75F01C2C4A0063FE8ECEB3E244185C22B56334040BC03C2FF0D01C0D1760B3F82884640DD6E7DC25266EC3FDD2E26C3C0802341CE16004241C04AC0F52066C0C279C94090A93E41B2CF814015F38D3F9204C9BFD7138D3E7B491BC017CFE0BF379E80BDA131A4C17B7255C237D017400D71514173B690448767A1BDBC3C73C08BBE73413E9CB1BF0B6483C164AAA83F35684742BBB561C17DBDB3406B62FE40373D0B429A6AF2C14BD087BFF38A54BF832FC6C0895836C09556084248EEDA408AB139C2C418833E84A27D4075A80F424C55404173C32BC0AED757BEBF878243551065420CFFD13E"> : tensor<20x20xf32>
+    return %cst : tensor<20x20xf32>
+  }
+}