Add Max Matrix Dim checks

Author: anupamachandra

tools/clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -8028,8 +8028,6 @@ def err_hlsl_linalg_matrix_layout_is_not_transposable : Error<
   "RowMajor and ColumnMajor matrices are not transposable">;
 def err_hlsl_linalg_optimal_matrix_layout_matrix_stride_must_be_zero : Error<
   "for optimal matrix layout, matrix stride must be zero">;
-def err_hlsl_linalg_exceeds_max_matrix_dim: Error<
-  "matrix dimension for linalg operations must be less than %0">;
 def err_hlsl_linalg_matrix_dim_must_not_be_zero: Error<
   "matrix dimension must not be zero">;
 def err_hlsl_linalg_matrix_layout_invalid : Error<
@@ -8049,6 +8047,14 @@ def err_hlsl_linalg_mul_muladd_isUnsigned_for_packed_input_must_be_true : Error<
 def err_hlsl_linalg_mul_muladd_packed_input_vector_must_be_uint : Error<
   "packed input vector type must be a 32-bit unsigned int in linalg mul/muladd operations, "
   "packed formats uint8_t4_packed and sint8_t4_packed are not supported currently">;
+def err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_m: Error<
+  "matrix dimension M must be less than %0, in a linalg Mul/MulAdd operation">;
+def err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_k_unpacked_input: Error<
+  "matrix dimension K when using unpacked input vectors must be less than %0, in a "
+  "linalg Mul/MulAdd operation">;
+def err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_k_packed_input: Error<
+  "matrix dimension K when using packed input vectors must be less than %0, in a "
+  "linalg Mul/MulAdd operation">;
 
 def err_hlsl_linalg_outer_prod_acc_vector_type_mismatch : Error<
   "input vectors of outerproductaccumulate must have the same element type">;

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -11987,19 +11987,36 @@ static void CheckCommonMulandMulAddParameters(Sema &S, CallExpr *CE,
 
   // Check MatrixM and MatrixK values are less than max
   // Matrix dimension cannot exceed largest vector length in a Mul/MulAdd
-  // operation
+  // operation.
   if (MatrixMValue > DXIL::kSM69MaxVectorLength) {
     S.Diags.Report(MatrixMExpr->getExprLoc(),
-                   diag::err_hlsl_linalg_exceeds_max_matrix_dim)
+                   diag::err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_m)
         << std::to_string(DXIL::kSM69MaxVectorLength);
     return;
   }
 
-  if (MatrixKValue > DXIL::kSM69MaxVectorLength) {
-    S.Diags.Report(MatrixKExpr->getExprLoc(),
-                   diag::err_hlsl_linalg_exceeds_max_matrix_dim)
-        << std::to_string(DXIL::kSM69MaxVectorLength);
-    return;
+  // For packed input vectors 4 values are packed in a uint, so max Matrix K
+  // can be 4096
+  if (IsInputVectorPacked) {
+    const unsigned PackingFactor =
+        4; // Only supported packed formats: DATA_TYPE_(U)SINT8_T4_PACKED
+    if (MatrixKValue > DXIL::kSM69MaxVectorLength * PackingFactor) {
+      S.Diags.Report(
+          MatrixKExpr->getExprLoc(),
+          diag::
+              err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_k_packed_input)
+          << std::to_string(DXIL::kSM69MaxVectorLength * PackingFactor);
+      return;
+    }
+  } else {
+    if (MatrixKValue > DXIL::kSM69MaxVectorLength) {
+      S.Diags.Report(
+          MatrixKExpr->getExprLoc(),
+          diag::
+              err_hlsl_linalg_mul_muladd_exceeds_max_matrix_dim_k_unpacked_input)
+          << std::to_string(DXIL::kSM69MaxVectorLength);
+      return;
+    }
   }
 
   if (!IsValidVectorAndMatrixDimensions(S, CE, InputVectorSizeValue,

tools/clang/test/SemaHLSL/hlsl/linalg/builtins/mul_add_invalid.hlsl

@@ -464,7 +464,114 @@ void test_invalid_matrix_K_dimension_non_zero() {
                         matrix_stride, bias_buffer, bias_offset, bias_interpretation);
 }
 
-//Check if InputInterpretation is a constant parameter
+// Check if Matrix M dimension is less than Max
+void test_invalid_matrix_M_dimension_less_than_Max() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimK = 4;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = matrix_dimK * 4;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 4> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM_0 = 1025;
+
+  // expected-error@+3 {{matrix dimension M must be less than 1024, in a linalg Mul/MulAdd operation}}
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM_0,
+                        matrix_dimK, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+  vector<uint, 1> input_vector_1 =
+      input_vector_buffer.Load<vector<uint, 1> >(0);
+  const uint input_interpretation_1 = DataType::DATA_TYPE_UINT8_T4_PACKED;
+  const uint matrix_dimM_1 = 4097;
+
+  // expected-error@+3 {{matrix dimension M must be less than 1024, in a linalg Mul/MulAdd operation}}
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_1,
+                        is_input_unsigned, input_interpretation_1, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM_1,
+                        matrix_dimK, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+}
+
+// Check if Matrix K dimension is less than Max in unpacked input vector case
+void test_invalid_matrix_K_dimension_less_than_Max_unpacked_input_vector() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM = 4;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = 64;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 4> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimK_0 = 1025;
+
+  // expected-error@+4 {{matrix dimension K when using unpacked input vectors must be less than 1024, in a linalg Mul/MulAdd operation}}
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_0, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+  vector<uint, 4> input_vector_1 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_1 = DataType::DATA_TYPE_UINT8;
+  const uint matrix_dimK_1 = 4096;
+  // expected-error@+4 {{matrix dimension K when using unpacked input vectors must be less than 1024, in a linalg Mul/MulAdd operation}}
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_1, 
+                        is_input_unsigned, input_interpretation_1, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_1, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+}
+
+// Check if Matrix M dimension is less than Max in packed input vector case
+void test_invalid_matrix_M_dimension_less_than_Max_packed_input_vector() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM = 1024;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = 4096;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 1024> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 1024> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_UINT8_T4_PACKED;
+  const uint matrix_dimK_0 = 4097;
+
+  // expected-error@+4 {{matrix dimension K when using packed input vectors must be less than 4096, in a linalg Mul/MulAdd operation}}
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_0, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+}
+
 void test_invalid_input_interpretation_non_const() {
 
   vector<uint, 4> output_vector;

tools/clang/test/SemaHLSL/hlsl/linalg/builtins/mul_add_valid.hlsl

@@ -132,5 +132,113 @@ void test_valid_packed_input_vector_dimension() {
                         matrix_stride, bias_buffer, bias_offset, bias_interpretation);
 }
 
+// Check if Matrix M dimension is less than Max
+void test_valid_matrix_M_dimension_less_than_Max() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimK = 4;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = matrix_dimK * 4;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 4> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM_0 = 4;
+
+  // expected-no-diagnostics@+1
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM_0,
+                        matrix_dimK, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+  vector<uint, 1> input_vector_1 =
+      input_vector_buffer.Load<vector<uint, 1> >(0);
+  const uint input_interpretation_1 = DataType::DATA_TYPE_UINT8_T4_PACKED;
+  const uint matrix_dimM_1 = 4;
+
+  // expected-no-diagnostics@+1
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_1,
+                        is_input_unsigned, input_interpretation_1, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM_1,
+                        matrix_dimK, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+}
+
+// Check if Matrix K dimension is less than Max in unpacked input vector case
+void test_valid_matrix_K_dimension_less_than_Max_unpacked_input_vector() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM = 4;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = 64;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 4> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimK_0 = 4;
+
+  // expected-no-diagnostics@+1
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_0, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+  vector<uint, 4> input_vector_1 =
+      input_vector_buffer.Load<vector<uint, 4> >(0);
+  const uint input_interpretation_1 = DataType::DATA_TYPE_UINT8;
+  const uint matrix_dimK_1 = 4;
+  // expected-no-diagnostics@+1
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_1, 
+                        is_input_unsigned, input_interpretation_1, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_1, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+
+}
+
+// Check if Matrix M dimension is less than Max in packed input vector case
+void test_valid_matrix_M_dimension_less_than_Max_packed_input_vector() {
+
+  vector<uint, 4> output_vector;
+  const uint is_output_unsigned = 1;
+  const uint is_input_unsigned = 1;
+  const uint matrix_offset = 0;
+  const uint matrix_interpretation = DataType::DATA_TYPE_FLOAT32;
+  const uint matrix_dimM = 4;
+  const uint matrix_layout = MatrixLayout::MATRIX_LAYOUT_ROW_MAJOR;
+  const bool matrix_is_transposed = false;
+  const uint matrix_stride = 64;
+  const uint bias_offset = 0;
+  const uint bias_interpretation = DataType::DATA_TYPE_FLOAT32;
+
+  vector<uint, 1024> input_vector_0 =
+      input_vector_buffer.Load<vector<uint, 1024> >(0);
+  const uint input_interpretation_0 = DataType::DATA_TYPE_UINT8_T4_PACKED;
+  const uint matrix_dimK_0 = 4096;
+
+  // expected-no-diagnostics@+1
+  __builtin_MatVecMulAdd(output_vector, is_output_unsigned, input_vector_0,
+                        is_input_unsigned, input_interpretation_0, matrix_buffer,
+                        matrix_offset, matrix_interpretation, matrix_dimM,
+                        matrix_dimK_0, matrix_layout, matrix_is_transposed,
+                        matrix_stride, bias_buffer, bias_offset, bias_interpretation);
+}
+