[XLA:GPU] Fix the unpack dim calculation for I4 rewrite with non major_2_minor layouts

PiperOrigin-RevId: 715289044

Author: loislo

xla/service/gpu/BUILD

@@ -1634,7 +1634,11 @@ cc_library(
     ]) + xla_internal(["service:export_hlo"]) + if_google([
         "//xla/hlo/experimental/auto_sharding",
         "//xla/hlo/experimental/auto_sharding:auto_sharding_option",
-    ]),
+    ]) + [
+        "//xla/tsl/platform:env",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:statusor",
+    ],
 )
 
 xla_test(

xla/service/gpu/fusions/triton/BUILD

@@ -242,6 +242,9 @@ cc_library(
         "//xla/service/llvm_ir:llvm_util",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:launch_dim",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:status",
+        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",

xla/service/gpu/fusions/triton/triton_fusion_emitter_int4_device_test.cc

@@ -96,7 +96,48 @@ class PlainInt4ToPackedInt4RewritePassTest : public TritonTest {
 };
 
 TEST_F(PlainInt4ToPackedInt4RewritePassTest,
-       DotWithI4WeightsOnLhsFusedWithMultiplyByChannelScales) {
+       DotWithI4WeightsOnLhsWithNonStandardLayoutAndMultplyInEpilogue) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule hlo
+
+    fusion {
+      p_0 = s4[1,128,32]{1,2,0:E(4)} parameter(0)
+      p_0.1 = s4[1,32,128]{2,1,0:E(4)} bitcast(p_0)
+      p_0.2 = bf16[1,32,128]{2,1,0} convert(p_0.1)
+      p_0.3 = bf16[1,128,32]{1,2,0} bitcast(p_0.2)
+      p_1 = bf16[128,1,64]{2,1,0} parameter(1)
+      dot = bf16[1,32,64]{2,1,0} dot(p_0.3, p_1),
+        lhs_batch_dims={0},
+        lhs_contracting_dims={1},
+        rhs_batch_dims={1},
+        rhs_contracting_dims={0}
+      p_2 = bf16[1,1,32]{2,0,1} parameter(2)
+      p_2.1 = bf16[1,32]{1,0} bitcast(p_2)
+      p_2.2 = bf16[1,32,64]{2,1,0} broadcast(p_2.1), dimensions={0,1}
+      m = bf16[1,32,64]{2,1,0} multiply(dot, p_2.2)
+      ROOT m.1 = bf16[1,1,32,64]{3,2,1,0} bitcast(m)
+    }
+
+    ENTRY %entry_computation {
+      p_0 = s4[1,128,32]{1,2,0:E(4)} parameter(0)
+      p_1 = bf16[128,1,64]{2,1,0} parameter(1)
+      p_2 = bf16[1,1,32]{2,0,1} parameter(2)
+      ROOT gemm_fusion_dot.2 = bf16[1,1,32,64]{3,2,1,0} fusion(p_0, p_1, p_2),
+        kind=kCustom,
+        calls=fusion,
+        backend_config={
+          "fusion_backend_config":{
+            "kind":"__triton_gemm"
+          }
+        }
+    }
+  )";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(
+      kHloText, ErrorSpec{/*aabs=*/1e-5, /*arel=*/1e-5}));
+}
+
+TEST_F(PlainInt4ToPackedInt4RewritePassTest,
+       DotWithInt4WeightsOnLhsFusedWithMultiplyByChannelScales) {
   constexpr absl::string_view kHloText = R"(
     HloModule DotWithI4WeightsOnLhsFusedWithMultiplyByChannelScales
 
@@ -133,6 +174,23 @@ TEST_F(PlainInt4ToPackedInt4RewritePassTest,
       kHloText, ErrorSpec{/*aabs=*/1e-5, /*arel=*/1e-5}));
 }
 
+TEST_F(PlainInt4ToPackedInt4RewritePassTest, NonstandardLayoutInt4) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule NonstandardLayoutInt4
+
+    ENTRY main {
+      p0 = s4[64,128]{0,1} parameter(0)
+      p1 = bf16[256,64]{1,0} parameter(1)
+      ROOT %dot = bf16[128,256]{1,0} dot(s4[64,128]{0,1} p0, bf16[256,64]{1,0} p1),
+        lhs_contracting_dims={0},
+        rhs_contracting_dims={1}
+    }
+  )";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto module, GetOptimizedModule(kHloText));
+  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+}
+
 using ::testing::TestParamInfo;
 using ::testing::WithParamInterface;
 

xla/service/gpu/fusions/triton/triton_fusion_emitter_legacy_matmul.cc

@@ -61,6 +61,8 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_query.h"
 #include "xla/hlo/utils/hlo_traversal.h"
+#include "xla/layout.h"
+#include "xla/layout_util.h"
 #include "xla/literal.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 #include "xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h"
@@ -73,6 +75,7 @@ limitations under the License.
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu/matmul_indexing_utils.h"
+#include "xla/service/gpu/matmul_utils.h"
 #include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/gpu/triton_fusion_analysis.h"
 #include "xla/service/gpu/triton_tiling_propagation.h"
@@ -82,6 +85,9 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/launch_dim.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/status.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/errors.h"
@@ -1477,7 +1483,8 @@ class MatMulEmitterHelper {
             .getResult());
     if (hlo->shape().element_type() == PrimitiveType::S4 &&
         IsTritonInt4RewritesEnabled(*hlo)) {
-      tensor_ptr.getDefiningOp()->setAttr("packed_dim", GetPackedDimAttr(side));
+      tensor_ptr.getDefiningOp()->setAttr(
+          "packed_dim", GetPackedDimAttr(side, hlo->shape().layout()));
     }
     tensor_ptr = b_.create<mt::AdvanceOp>(tensor_ptr.getType(), tensor_ptr,
                                           block_offsets);
@@ -1486,16 +1493,22 @@ class MatMulEmitterHelper {
 
   // Naive implementation of the packed_dim attribute for the int4 tensors.
   // It doesn't take into account different layout schemes.
-  mlir::IntegerAttr GetPackedDimAttr(const Side& side) const {
+  mlir::IntegerAttr GetPackedDimAttr(const Side& side,
+                                     const Layout& layout) const {
     int packed_dim = 0;
+    const std::vector<int64_t> logical_to_physical =
+        LayoutUtil::MakeLogicalToPhysical(layout);
+
     if (side.scope == TritonFusionAnalysis::Scope::LHS) {
-      if (dims_.lhs_contracting_dim_idx > dims_.lhs_noncontracting_dim_idx) {
+      if (logical_to_physical[dims_.lhs_contracting_dim_idx] >
+          logical_to_physical[dims_.lhs_noncontracting_dim_idx]) {
         packed_dim = 0;
       } else {
         packed_dim = 1;
       }
     } else if (side.scope == TritonFusionAnalysis::Scope::RHS) {
-      if (dims_.rhs_contracting_dim_idx > dims_.rhs_noncontracting_dim_idx) {
+      if (logical_to_physical[dims_.rhs_contracting_dim_idx] >
+          logical_to_physical[dims_.rhs_noncontracting_dim_idx]) {
         packed_dim = 1;
       } else {
         packed_dim = 0;

xla/service/gpu/gpu_compiler.cc

@@ -263,6 +263,10 @@ limitations under the License.
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/threadpool.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
@@ -1552,7 +1556,9 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
       pipeline.AddPass<GemvRewriter>();
       pipeline.AddPass<GemmFusion>(gpu_version);
       pipeline.AddPass<GemmFusionSwapOperands>();
-      pipeline.AddPass<SimplifyInt4Dots>();
+      if (!debug_options.xla_gpu_experimental_enable_triton_i4_rewrites()) {
+        pipeline.AddPass<SimplifyInt4Dots>();
+      }
     } else if (cuda_cc != nullptr &&
                cuda_cc->major == se::CudaComputeCapability::VOLTA) {
       // Greedy pattern matching for custom kernel fusions.