[EXPERIMENTAL]: Load column major matrix with 2d block io

include/triton/Tools/Sys/GetEnv.hpp

@@ -50,6 +50,10 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_INTEL_ENABLE_INSTR_SCHED",
     "TRITON_INTEL_FAST_MATH",
     "TRITON_INTEL_REDUCE_TRANSPOSE",
+    "TRITON_INTEL_ENABLE_SIMD_REDUCE",
+    "TRITON_INTEL_ENHANCED_ACCELERATION_MATMUL",
+    "TRITON_INTEL_ENABLE_DPAS_WARP_SIZE_32",
+    "TRITONGEN_FORCE_GENISA",
     // clang-format on
 };
 

third_party/intel/include/Dialect/TritonGEN/IR/TritonGENAttrDefs.td

@@ -55,9 +55,9 @@ def TritonGEN_PrecisionTypeAttr : I32EnumAttr<"PrecisionType",
     I32EnumAttrCase<"S4",     5,  "i4">,
     I32EnumAttrCase<"S2",     6,  "i2">,
     I32EnumAttrCase<"BF8",    7,  "bf8">,
-    I32EnumAttrCase<"TF32",   8,  "tf32">,
-    I32EnumAttrCase<"BF16",   9,  "bf16">,
-    I32EnumAttrCase<"FP16",   10, "f16">
+    I32EnumAttrCase<"TF32",   10,  "tf32">,
+    I32EnumAttrCase<"BF16",   11,  "bf16">,
+    I32EnumAttrCase<"FP16",   12, "f16">
   ]> {
   let cppNamespace = "::mlir::triton::TritonGEN";
 }

third_party/intel/lib/Dialect/TritonGEN/IR/TritonGENOps.cpp

@@ -85,12 +85,15 @@ LogicalResult TritonGEN::MatrixDPASOp::verify() {
     return this->emitOpError(
         "1st operand (C) and result (D) should have the same type");
 
-  if (CTy.getNumElements() != getRc() || DTy.getNumElements() != getRc())
+  auto useGenISA = tools::getBoolEnv("TRITONGEN_FORCE_GENISA");
+
+  if (!useGenISA &&
+      (CTy.getNumElements() != getRc() || DTy.getNumElements() != getRc()))
     return this->emitOpError("the dimension for 1st operand (C) and "
                              "result (D) should match repeat count");
 
   constexpr unsigned SD = 8;
-  if (BTy.getNumElements() != SD)
+  if (!useGenISA && BTy.getNumElements() != SD)
     return this->emitOpError("the dimension for the 3rd operand (B) should "
                              "match the systolic depth of 8");
 
@@ -141,7 +144,7 @@ LogicalResult TritonGEN::MatrixDPASOp::verify() {
   case TritonGEN::PrecisionType::FP16:
   case TritonGEN::PrecisionType::U8:
   case TritonGEN::PrecisionType::S8:
-    if (ATy.getNumElements() != getRc())
+    if (!useGenISA && ATy.getNumElements() != getRc())
       return this->emitOpError("2nd operand (A) should have the same number of "
                                "elements as repeat count");
     if (!AElemTy.isInteger(16))
@@ -303,6 +306,9 @@ LogicalResult TritonGEN::Matrix2DBlockLoadOp::verify() {
   if (verify2DBlockLoadHWRestriction(*this).failed())
     return failure();
 
+  if (tools::getBoolEnv("TRITONGEN_FORCE_GENISA"))
+    return success();
+
   if (verifyMatrixInput(*this).failed())
     return failure();
 
@@ -367,6 +373,9 @@ LogicalResult TritonGEN::Matrix2DBlockStoreOp::verify() {
   if (verify2DBlockStoreHWRestriction(*this).failed())
     return failure();
 
+  if (tools::getBoolEnv("TRITONGEN_FORCE_GENISA"))
+    return success();
+
   if (verifyMatrixInput(*this).failed())
     return failure();
 

third_party/intel/lib/TritonGENToLLVM/CMakeLists.txt

@@ -1,12 +1,21 @@
+add_library(GenISAIntrinsics STATIC IMPORTED GLOBAL)
+set_target_properties(GenISAIntrinsics
+  PROPERTIES IMPORTED_LOCATION ${CMAKE_CURRENT_SOURCE_DIR}/libGenISAIntrinsics.a
+)
+
 add_triton_library(TritonGENToLLVM
   Attributes.cpp
   TritonGENToLLVMPass.cpp
+  GenIntrinsicHelper.cpp
 
   DEPENDS
   TritonGENToLLVMConversionPassIncGen
+  GenISAIntrinsics
 
   LINK_LIBS PUBLIC
+  GenISAIntrinsics
   MLIRLLVMDialect
   MLIRSPIRVDialect
   TritonIntelUtils
+  MLIRTargetLLVMIRImport
 )