[pull] master from tensorflow:master

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in.hlo

@@ -0,0 +1,25 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+fusion {
+  p = s32[48,16]{0,1} parameter(0)
+  ROOT bitcast = s32[16,16,3] bitcast(p)
+}
+// CHECK:     xtile.extract
+// CHECK:     tt.trans
+// CHECK:     tt.reshape
+// CHECK-NOT: tt.trans
+// CHECK:     xtile.insert
+
+ENTRY main {
+  p = s32[48,16]{0,1} parameter(0)
+  ROOT fusion = s32[16,16,3] fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+      "kind":"__triton",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["2","8","1"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in_out.hlo

@@ -0,0 +1,25 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+fusion {
+  p = s8[42,5]{0,1} parameter(0)
+  ROOT bitcast = s8[5,6,7]{1,2,0} bitcast(p)
+}
+// CHECK:     xtile.extract
+// CHECK:     tt.trans
+// CHECK:     tt.reshape
+// CHECK:     tt.trans
+// CHECK:     xtile.insert
+
+ENTRY main {
+  p = s8[42,5]{0,1} parameter(0)
+  ROOT fusion = s8[5,6,7]{1,2,0} fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+      "kind":"__triton",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["2","4","1"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_out.hlo

@@ -0,0 +1,25 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+fusion {
+  p = s8[5,42] parameter(0)
+  ROOT bitcast = s8[5,6,7]{1,2,0} bitcast(p)
+}
+// CHECK:     xtile.extract
+// CHECK-NOT: tt.trans
+// CHECK:     tt.reshape
+// CHECK:     tt.trans
+// CHECK:     xtile.insert
+
+ENTRY main {
+  p = s8[5,42] parameter(0)
+  ROOT fusion = s8[5,6,7]{1,2,0} fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+      "kind":"__triton",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["2","4","1"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo

@@ -0,0 +1,26 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+
+fusion {
+  p = f32[8,48]{1,0} parameter(0)
+  ROOT bitcast = f32[8,16,3] bitcast(p)
+}
+// CHECK:     xtile.extract
+// CHECK-NOT: tt.trans
+// CHECK:     tt.reshape
+// CHECK-NOT: tt.trans
+// CHECK:     xtile.insert
+
+ENTRY main {
+  p = f32[8,48]{1,0} parameter(0)
+  ROOT fusion = f32[8,16,3] fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+    "fusion_backend_config":{
+      "kind":"__triton",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["2","8","1"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo

@@ -0,0 +1,34 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+region {
+  param_0.1 = f32[] parameter(0)
+  param_1 = f32[] parameter(1)
+  ROOT add = f32[] maximum(param_0.1, param_1)
+}
+
+fusion {
+  param_0.2 = f32[64] parameter(0)
+  abs = f32[64] abs(param_0.2)
+  bitcast = f32[4,4,4] bitcast(abs)
+  constant = f32[] constant(0)
+  reduce = f32[4,4] reduce(bitcast, constant), dimensions={1}, to_apply=region
+  ROOT tuple = (f32[4,4], f32[64]) tuple(reduce, abs)
+}
+// CHECK-COUNT-1:  xtile.extract
+// CHECK: tt.reduce
+// CHECK-COUNT-2:  xtile.insert
+
+ENTRY main {
+  param_0.3 = f32[64] parameter(0)
+  ROOT fusion = (f32[4,4], f32[64]) fusion(param_0.3), kind=kCustom,
+    calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["1", "4"]},{"sizes":["16"]}],
+          "num_warps":"2",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo

@@ -0,0 +1,25 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+fusion {
+  p = s8[42,5]{0,1} parameter(0)
+  ROOT bitcast = s8[5,42] bitcast(p)
+}
+// CHECK:     xtile.extract
+// CHECK:     tt.trans
+// CHECK-NOT: tt.reshape
+// CHECK-NOT: tt.trans
+// CHECK:     xtile.insert
+
+ENTRY main {
+  p = s8[42,5]{0,1} parameter(0)
+  ROOT fusion = s8[5,42] fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+      "kind":"__triton",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["4","1"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_fusion.hlo

@@ -0,0 +1,30 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s
+
+fusion {
+  p0 = s32[128] parameter(0)
+  p1 = s32[128] parameter(1)
+  p2 = s32[25] parameter(2)
+  abs0 = s32[128] abs(p0)
+  neg1 = s32[128] negate(p1)
+  ROOT concatenate = s32[281] concatenate(abs0, neg1, p2),
+    dimensions={0}
+}
+// Check that we generate three branches. This is a bit of an implementation
+// detail, so it doesn't seem worth enforcing a lot here.
+// CHECK-COUNT-2: scf.if
+
+ENTRY main {
+  p0 = s32[128] parameter(0)
+  p1 = s32[128] parameter(1)
+  p2 = s32[25] parameter(2)
+  ROOT fusion = s32[281] fusion(p0, p1, p2), kind=kCustom,
+    calls=fusion, backend_config={
+    "fusion_backend_config":{
+      "kind":"__triton_nested_gemm_fusion",
+      "block_level_fusion_config":{
+        "output_tiles":[{"sizes":["32"]}],
+        "num_warps":"1",
+        "num_ctas":"1",
+        "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/bf16_small_rhs_outer.hlo

@@ -0,0 +1,23 @@
+// RUN: triton_test_correctness %s --abs_error_bound=1e-1 --rel_error_bound=1e-2
+
+fusion {
+  fdot.p0 = bf16[64,32] parameter(0)
+  fdot.p1 = bf16[32,8] parameter(1)
+  ROOT fdot.root = bf16[64,8]{1,0} dot(fdot.p0, fdot.p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+
+ENTRY main {
+  entry.p0 = bf16[64,32] parameter(0)
+  entry.p1 = bf16[32,8] parameter(1)
+  ROOT fusion = bf16[64,8] fusion(entry.p0, entry.p1),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["64","8"]}],
+          "num_warps":"4",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_concat.hlo

@@ -0,0 +1,36 @@
+// RUN: triton_test_correctness --abs_error_bound=1e-4 --rel_error_bound=1e-6 %s
+
+
+fusion {
+  p0 = f32[32,299] parameter(0)
+  p1 = f32[299,128] parameter(1)
+  p2 = f32[299,128] parameter(2)
+  p3 = f32[299,128] parameter(3)
+  p4 = f32[299,128] parameter(4)
+  lhs_cos = f32[32,299] cosine(p0)
+  abs1 = f32[299,128] abs(p1)
+  neg2 = f32[299,128] negate(p2)
+  cos4 = f32[299,128] cosine(p4)
+  concatenate = f32[299,512] concatenate(abs1, neg2, p3, cos4), dimensions={1}
+  rhs_cos = f32[299,512] cosine(concatenate)
+  ROOT dot = f32[32,512]{1,0} dot(lhs_cos, rhs_cos),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+
+ENTRY main {
+  p0 = f32[32,299] parameter(0)
+  p1 = f32[299,128] parameter(1)
+  p2 = f32[299,128] parameter(2)
+  p3 = f32[299,128] parameter(3)
+  p4 = f32[299,128] parameter(4)
+  ROOT fusion = f32[32,512] fusion(p0, p1, p2, p3, p4),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion", "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["16", "64"]}], "num_warps":"1",
+          "num_ctas":"1", "num_stages":"1"
+        }
+      }
+    }
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_major_lhs_contract.hlo

@@ -0,0 +1,27 @@
+// RUN: triton_test_correctness %s
+
+
+fusion {
+  p0 = f32[299,32] parameter(0)
+  p1 = f32[299,512] parameter(1)
+  ROOT dot = f32[32,512]{1,0} dot(p0, p1),
+    lhs_contracting_dims={0}, rhs_contracting_dims={0},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+
+ENTRY main {
+  // Take in boolean inputs for the test, in order to allow exact accumulation.
+  p0 = pred[299,32] parameter(0)
+  p1 = pred[299,512] parameter(1)
+  p0_f32 = f32[299,32] convert(p0)
+  p1_f32 = f32[299,512] convert(p1)
+  ROOT fusion = f32[32,512] fusion(p0_f32, p1_f32),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["16", "64"]}],
+          "num_warps":"1",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_minor_rhs_contract.hlo

@@ -0,0 +1,25 @@
+// RUN: triton_test_correctness %s
+
+fusion {
+  p0 = f32[32,299] parameter(0)
+  p1 = f32[512,299] parameter(1)
+  ROOT dot = f32[32,512]{1,0} dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={1},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+
+ENTRY main {
+  p0 = pred[32,299] parameter(0)
+  p1 = pred[512,299] parameter(1)
+  p0_f32 = f32[32,299] convert(p0)
+  p1_f32 = f32[512,299] convert(p1)
+  ROOT fusion = f32[32,512] fusion(p0_f32, p1_f32),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["16", "64"]}],
+          "num_warps":"1",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_unit_tile_sizes.hlo

@@ -0,0 +1,27 @@
+// RUN: triton_test_correctness %s
+
+fusion {
+  p0 = f32[2,3,32,125] parameter(0)
+  p1 = f32[2,125,3,256] parameter(1)
+  ROOT dot = f32[2,3,32,256] dot(p0, p1),
+    lhs_batch_dims={0,1}, rhs_batch_dims={0,2},
+    lhs_contracting_dims={3}, rhs_contracting_dims={1},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+
+ENTRY main {
+  // Take in boolean inputs for the test, in order to allow exact accumulation.
+  p0 = pred[2,3,32,125] parameter(0)
+  p1 = pred[2,125,3,256] parameter(1)
+  p0_f32 = f32[2,3,32,125] convert(p0)
+  p1_f32 = f32[2,125,3,256] convert(p1)
+  ROOT fusion = f32[2,3,32,256] fusion(p0_f32, p1_f32),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["1", "1", "16", "64"]}],
+          "num_warps":"1",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}

third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/masked_dot.hlo

@@ -0,0 +1,34 @@
+// RUN: hlo_to_xtileir %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
+
+
+fusion {
+  fdot.p0 = f32[32,299] parameter(0)
+  fdot.p1 = f32[299,512] parameter(1)
+  lhs.root = f32[32,299] cosine(fdot.p0)
+  frhs.root = f32[299,512] cosine(fdot.p1)
+  ROOT fdot.root = f32[32,512]{1,0} dot(lhs.root, frhs.root),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0},
+    algorithm=dot_f32_f32_f32, backend_config={sizes:[32]}
+}
+// CHECK:      %[[MASKED_OPERAND0:.*]] = scf.if
+// CHECK:        %[[SELECT0:.*]] = arith.select
+// CHECK-NEXT:   scf.yield %[[SELECT0]]
+// CHECK:      %[[MASKED_OPERAND1:.*]] = scf.if
+// CHECK:        %[[SELECT1:.*]] = arith.select
+// CHECK-NEXT:   scf.yield %[[SELECT1]]
+// CHECK:      tt.dot %[[MASKED_OPERAND0]], %[[MASKED_OPERAND1]]
+
+ENTRY main {
+  entry.p0 = f32[32,299] parameter(0)
+  entry.p1 = f32[299,512] parameter(1)
+  ROOT fusion = f32[32,512] fusion(entry.p0, entry.p1),
+    kind=kCustom, calls=fusion, backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton_nested_gemm_fusion",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["16", "64"]}],
+          "num_warps":"1",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}