Update syntax for amx.tile_muli to use two Unit attr to mark the zext case

[mehdi_amini]

Summary

This makes the annotation tied to the operand and the use of a keyword
more explicit/readable on what it means.

[aartbik]

please update integration tests too!

mlir/test/Dialect/AMX/roundtrip.mlir

33

seems to have filename at end?

[mehdi_amini]

Update integration tests

[aartbik]

mlir/test/Dialect/AMX/roundtrip.mlir

44

verify -> Verify (start of sentence)

[mehdi_amini]

fix comments in tests

mlir/include/mlir/Dialect/AMX/AMX.td

 
 def TileMulIOp : AMX_Op<"tile_muli", [NoSideEffect, AllTypesMatch<["acc", "res"]>]> {
   let summary = "tile multiplication operation (integer)";
   let description = [{
     Multiplies a "m x k" tile with a "k x n" tile and accumulates the results
     into a "m x n" destination tile. Supports all "si32 <- s/ui8 x s/ui8"
     combinations (4 bytes packed into dwords in the columns of both the
     source operand tiles; the zero or sign extension is specified with
-    the attributes). The operation is eventually lowered into one of
-    the "tdpbssd", "tdpbsud", "tdpbusd", or "tdpbuud" instructions with
-    the corresponding tile configuration.
+    the attributes and default to sign extended). The operation is eventually
+    lowered into one of the "tdpbssd", "tdpbsud", "tdpbusd", or "tdpbuud"
+    instructions with the corresponding tile configuration.
 
     Example:
 
     ```mlir
-      %0 = amx.tile_muli %a, %b, %c [true, true]
+      %0 = amx.tile_muli %a zext, %b zext, %c 
         : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
     ```
   }];
   let verifier = [{ return ::verify(*this); }];
   let arguments = (ins VectorOfRankAndType<[2], [I32, I8]>:$lhs,
                        VectorOfRankAndType<[2], [I32, I8]>:$rhs,
                        VectorOfRankAndType<[2], [I32, I8]>:$acc,
-                       BoolArrayAttr:$zext);
+                       UnitAttr:$isZextLhs,
+                       UnitAttr:$isZextRhs
+                       );
   let results = (outs VectorOfRankAndType<[2], [I32, I8]>:$res);
   let extraClassDeclaration = [{
     VectorType getLhsVectorType() {
       return lhs().getType().cast<VectorType>();
     }
     VectorType getRhsVectorType() {
       return rhs().getType().cast<VectorType>();
     }
     VectorType getVectorType() {
       return res().getType().cast<VectorType>();
     }
   }];
-  let assemblyFormat = "$lhs `,` $rhs `,` $acc $zext attr-dict `:` "
+  let assemblyFormat = "$lhs (`zext` $isZextLhs^)? `,` $rhs (`zext` $isZextRhs^)? `,` $acc attr-dict `:` "
                        "type($lhs) `,` type($rhs) `,` type($acc) ";
 }
 
 //===----------------------------------------------------------------------===//
 // AMX IntrOp definitions (LLVM compiler facing).
 //===----------------------------------------------------------------------===//
 
 //

mlir/lib/Dialect/AMX/IR/AMXDialect.cpp

static LogicalResult verify(amx::TileMulFOp op) {
   Type tb = bType.getElementType();
   Type tc = cType.getElementType();
   if (!ta.isBF16() || !tb.isBF16() || !tc.isF32())
     return op.emitOpError("unsupported type combination");
   return success();
 }
 
 static LogicalResult verify(amx::TileMulIOp op) {
-  if (op.zext().size() != 2)
-    return op.emitOpError("unexpected zext length");
   VectorType aType = op.getLhsVectorType();
   VectorType bType = op.getRhsVectorType();
   VectorType cType = op.getVectorType();
   if (failed(verifyTileSize(op, aType)) || failed(verifyTileSize(op, bType)) ||
       failed(verifyTileSize(op, cType)) ||
       failed(verifyMultShape(op, aType, bType, cType, 2)))
     return failure();
   Type ta = aType.getElementType();

mlir/lib/Dialect/AMX/Transforms/LegalizeForLLVMExport.cpp

matchAndRewrite(TileMulIOp op, ArrayRef<Value> operands,
     VectorType cType = op.getVectorType();
     // Determine m x n x k tile sizes.
     std::pair<Value, Value> tsza =
         getTileSizes(rewriter, *getTypeConverter(), aType, op.getLoc());
     std::pair<Value, Value> tszb =
         getTileSizes(rewriter, *getTypeConverter(), bType, op.getLoc());
     // Replace operation with intrinsic.
     Type resType = typeConverter->convertType(cType);
-    bool zexta = op.zext()[0].cast<BoolAttr>().getValue();
-    bool zextb = op.zext()[1].cast<BoolAttr>().getValue();
+    bool zexta = op.isZextLhs();
+    bool zextb = op.isZextRhs();
     if (zexta && zextb)
       rewriter.replaceOpWithNewOp<amx::x86_amx_tdpbuud>(
           op, resType, tsza.first, tszb.second, tsza.second, adaptor.acc(),
           adaptor.lhs(), adaptor.rhs());
     else if (zexta && !zextb)
       rewriter.replaceOpWithNewOp<amx::x86_amx_tdpbusd>(
           op, resType, tsza.first, tszb.second, tsza.second, adaptor.acc(),
           adaptor.lhs(), adaptor.rhs());

mlir/test/Dialect/AMX/invalid.mlir

 
 func @multsize() {
   %0 = amx.tile_zero : vector<8x8xbf16>
   %1 = amx.tile_zero : vector<8x8xbf16>
   %2 = amx.tile_zero : vector<4x4xf32>
   // expected-error@+1 {{'amx.tile_mulf' op bad mult shape: 4 x 4 x 4}}
   %3 = amx.tile_mulf %0, %1, %2 : vector<8x8xbf16>, vector<8x8xbf16>, vector<4x4xf32>
 }
-
-// -----
-
-func @zextsize() {
-  %0 = amx.tile_zero : vector<8x8xi8>
-  %1 = amx.tile_zero : vector<8x8xi8>
-  %2 = amx.tile_zero : vector<8x8xi32>
-  // expected-error@+1 {{'amx.tile_muli' op unexpected zext length}}
-  %3 = amx.tile_muli %0, %1, %2 [true] : vector<8x8xi8>, vector<8x8xi8>, vector<8x8xi32>
-}

mlir/test/Dialect/AMX/legalize-for-llvm.mlir

 // CHECK: amx.tilestored64
 // CHECK: amx.tdpbsud
 // CHECK: amx.tilestored64
 func @muli(%arg0: memref<?x?xi8>, %arg1: memref<?x?xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_zero : vector<16x64xi8>
   %2 = amx.tile_load %arg0[%0, %0] : memref<?x?xi8> into vector<16x64xi8>
   %3 = amx.tile_load %arg1[%0, %0] : memref<?x?xi32> into vector<16x16xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, true] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   amx.tile_store %arg1[%0, %0], %4 : memref<?x?xi32>, vector<16x16xi32>
-  %5 = amx.tile_muli %1, %2, %3 [false, false] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %5 = amx.tile_muli %1, %2, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   amx.tile_store %arg1[%0, %0], %5 : memref<?x?xi32>, vector<16x16xi32>
-  %6 = amx.tile_muli %1, %2, %3 [true, false] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %6 = amx.tile_muli %1 zext, %2, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   amx.tile_store %arg1[%0, %0], %6 : memref<?x?xi32>, vector<16x16xi32>
-  %7 = amx.tile_muli %1, %2, %3 [false, true] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %7 = amx.tile_muli %1, %2 zext, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   amx.tile_store %arg1[%0, %0], %7  : memref<?x?xi32>, vector<16x16xi32>
   return
 }
 
 // CHECK-LABEL: mulf(
 // CHECK: amx.tilezero
 // CHECK: amx.tileloadd64
 // CHECK: amx.tileloadd64

mlir/test/Dialect/AMX/roundtrip.mlir

func @tmulf(%arg0: memref<?x?xbf16>, %arg1: memref<?x?xf32>) {
   amx.tile_store %arg1[%0, %0], %3 : memref<?x?xf32>, vector<16x16xf32>
   return
 }
 
 // CHECK-LABEL: tmuli
 // CHECK: %[[x:.*]] = amx.tile_load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xi8> into vector<16x64xi8>
 // CHECK: %[[y:.*]] = amx.tile_load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xi8> into vector<16x64xi8>
 // CHECK: %[[z:.*]] = amx.tile_load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xi32> into vector<16x16xi32>
-// CHECK: %[[m:.*]] = amx.tile_muli %[[x]], %[[y]], %[[z]] [true, true] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+// CHECK: %[[m:.*]] = amx.tile_muli %[[x]] zext, %[[y]] zext, %[[z]] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
 // CHECK: amx.tile_store %{{.*}}[%{{.*}}, %{{.*}}], %[[m]] : memref<?x?xi32>, vector<16x16xi32>
+// Verify the parsing/printing of the sign-extension annotation.
+// CHECK: amx.tile_muli %{{.*}}, %{{.*}} zext, %{{.*}}
+// CHECK: amx.tile_muli %{{.*}} zext, %{{.*}}, %{{.*}}
+// CHECK: amx.tile_muli %{{.*}}, %{{.*}}, %{{.*}}
 func @tmuli(%arg0: memref<?x?xi8>, %arg1: memref<?x?xi8>, %arg2: memref<?x?xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<?x?xi8> into vector<16x64xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<?x?xi8> into vector<16x64xi8>
   %3 = amx.tile_load %arg2[%0, %0] : memref<?x?xi32> into vector<16x16xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, true] : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<?x?xi32>, vector<16x16xi32>
+  // Verify the various `zext` combinations.
+  %5 = amx.tile_muli %1, %2 zext, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %6 = amx.tile_muli %1 zext, %2, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
+  %7 = amx.tile_muli %1, %2, %3 : vector<16x64xi8>, vector<16x64xi8>, vector<16x16xi32>
   return
 }

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-muli-ext.mlir

 
 func @kernel1(%arg0: memref<16x16xi8>,
               %arg1: memref<4x16xi8>,
               %arg2: memref<16x4xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<16x16xi8>  into vector<16x16xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<4x16xi8>  into vector<4x16xi8>
   %3 = amx.tile_zero : vector<16x4xi32>
-  %4 = amx.tile_muli %1, %2, %3 [false, false] : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
+  %4 = amx.tile_muli %1, %2, %3 : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<16x4xi32>, vector<16x4xi32>
   return
 }
 
 func @kernel2(%arg0: memref<16x16xi8>,
               %arg1: memref<4x16xi8>,
               %arg2: memref<16x4xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<16x16xi8>  into vector<16x16xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<4x16xi8>  into vector<4x16xi8>
   %3 = amx.tile_zero : vector<16x4xi32>
-  %4 = amx.tile_muli %1, %2, %3 [false, true] : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
+  %4 = amx.tile_muli %1, %2 zext, %3 : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<16x4xi32>, vector<16x4xi32>
   return
 }
 
 func @kernel3(%arg0: memref<16x16xi8>,
               %arg1: memref<4x16xi8>,
               %arg2: memref<16x4xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<16x16xi8>  into vector<16x16xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<4x16xi8>  into vector<4x16xi8>
   %3 = amx.tile_zero : vector<16x4xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, false] : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
+  %4 = amx.tile_muli %1 zext, %2, %3 : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<16x4xi32>, vector<16x4xi32>
   return
 }
 
 func @kernel4(%arg0: memref<16x16xi8>,
               %arg1: memref<4x16xi8>,
               %arg2: memref<16x4xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<16x16xi8>  into vector<16x16xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<4x16xi8>  into vector<4x16xi8>
   %3 = amx.tile_zero : vector<16x4xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, true] : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
+  %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<16x16xi8>, vector<4x16xi8>, vector<16x4xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<16x4xi32>, vector<16x4xi32>
   return
 }
 
 func @entry() {
   %c0 = constant 0: index
 
   // Set up memory.

mlir/test/Integration/Dialect/Vector/CPU/AMX/test-muli.mlir

 // RUN: mlir-opt %s -convert-vector-to-scf -lower-affine -convert-scf-to-std -convert-vector-to-llvm="enable-amx" -convert-std-to-llvm | \
 // RUN: mlir-translate -mlir-to-llvmir | \
 // RUN: %lli --entry-function=entry --mattr="+amx-tile,+amx-int8,+amx-bf16" --dlopen=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 // Note: To run this test, your CPU must support AMX.
 
 // Multiply into zeroed destination.
 func @kernel1(%arg0: memref<2x8xi8>,
               %arg1: memref<2x8xi8>,
         %arg2: memref<2x2xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %3 = amx.tile_zero : vector<2x2xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, true] : vector<2x8xi8>, vector<2x8xi8>, vector<2x2xi32>
+  %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<2x8xi8>, vector<2x8xi8>, vector<2x2xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<2x2xi32>, vector<2x2xi32>
   return
 }
 
 // Multiply and update into destination.
 func @kernel2(%arg0: memref<2x8xi8>,
               %arg1: memref<2x8xi8>,
         %arg2: memref<2x2xi32>) {
   %0 = constant 0 : index
   %1 = amx.tile_load %arg0[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %2 = amx.tile_load %arg1[%0, %0] : memref<2x8xi8>  into vector<2x8xi8>
   %3 = amx.tile_load %arg2[%0, %0] : memref<2x2xi32> into vector<2x2xi32>
-  %4 = amx.tile_muli %1, %2, %3 [true, true] : vector<2x8xi8>, vector<2x8xi8>, vector<2x2xi32>
+  %4 = amx.tile_muli %1 zext, %2 zext, %3 : vector<2x8xi8>, vector<2x8xi8>, vector<2x2xi32>
   amx.tile_store %arg2[%0, %0], %4 : memref<2x2xi32>, vector<2x2xi32>
   return
 }
 
 func @entry() {
   %i0 = constant 0: i32
   %c0 = constant 0: index
   %c1 = constant 1: index