[mlir][ArmSME] Fix get_tile_id type in zero lowering

[c-rhodes]

Summary

The arm_sme.get_tile_id op returns a scalar integer but the arm_sme.zero
op lowering incorrectly uses the element type, which could be
floating-point.

[c-rhodes]

I think this op is a bit clunky and could be simplified by returning the element type of the tile rather than an int representing the bitwidth, or perhaps just the 2-D scalable vector type this is extracted from, but that's a slightly move involved refactor.

[awarzynski]

I think this op is a bit clunky and could be simplified by returning the element type of the tile rather than an int representing the bitwidth, or perhaps just the 2-D scalable vector type this is extracted from, but that's a slightly move involved refactor.

Which Op are you referring to arm_sme.zero? arm_sme.cast_tile_to_vector?

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

6

How is "8-bit" relevant here? What's meant to be 8-bit?

9

[nit] DELETME

[c-rhodes]

In D159080#4624367, @awarzynski wrote:

I think this op is a bit clunky and could be simplified by returning the element type of the tile rather than an int representing the bitwidth, or perhaps just the 2-D scalable vector type this is extracted from, but that's a slightly move involved refactor.

Which Op are you referring to arm_sme.zero? arm_sme.cast_tile_to_vector?

arm_sme.get_tile_id

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

6

How is "8-bit" relevant here? What's meant to be 8-bit?

the zero instruction mask, see https://armv8.arm.com/latest_builds/v9A/isa64/zero_za_i.xml

[c-rhodes]

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

6

How is "8-bit" relevant here? What's meant to be 8-bit?

the zero instruction mask, see https://armv8.arm.com/latest_builds/v9A/isa64/zero_za_i.xml

Apologies that's an internal link, external link: https://developer.arm.com/documentation/ddi0602/2022-06/SME-Instructions/ZERO--Zero-a-list-of-64-bit-element-ZA-tiles-

[benmxwl-arm]

My bad, LGTM!

[awarzynski]

LGTM, thanks! My comments are nits, so feel free to ignore.

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

6

Thanks, now I see what you meant!

It's a bit confusing that this is referring to "8-bit" tile masks, but there are no 8-bit tile masks here :) Perhaps:

// This test verifies the tile mask operand of the zero intrinsic zeroes
// the correct tiles. Both integer and floating-point datatypes are checked.
// Note that once lowered to ASM/machine code, these masks are narrowed to 8 bit:
//   * https://developer.arm.com/documentation/ddi0602/2022-06/SME-Instructions/ZERO--Zero-a-list-of-64-bit-element-ZA-tiles-

WDYT? Did I get it right?

[c-rhodes]

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

6

I've just removed mention of 8-bit mask altogether

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

 struct ZeroOpConversion : public ConvertOpToLLVMPattern<ZeroOp> {
   using ConvertOpToLLVMPattern<ZeroOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(ZeroOp zero, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = zero.getLoc();
 
+    unsigned tileElementWidth =
+        zero.getVectorType().getElementType().getIntOrFloatBitWidth();
+
     // Get Tile ID for the `zero` intrinsic.
     auto tileId = rewriter.create<arm_sme::GetTileID>(
-        loc, zero.getVectorType().getElementType());
-
-    auto tileElementWidth = tileId.getType().getIntOrFloatBitWidth();
+        loc, rewriter.getIntegerType(tileElementWidth));
 
     // Get the base mask for tile based on the element size.
     // The base mask is just the mask to zero the first tile (of a size).
     // These masks are derived from:
     // https://developer.arm.com/documentation/ddi0602/2022-06/SME-Instructions/ZERO--Zero-a-list-of-64-bit-element-ZA-tiles-
     auto baseMaskForSize = [&] {
       switch (tileElementWidth) {
       case 8:

mlir/test/Dialect/ArmSME/tile-zero-masks.mlir

 // RUN: mlir-opt %s -convert-vector-to-llvm="enable-arm-sme" \
 // RUN:           -allocate-arm-sme-tiles -canonicalize      \
 // RUN:           -allow-unregistered-dialect                \
 // RUN: | FileCheck %s
 
+// This test verifies the tile mask operand of the zero intrinsic zeroes
+// the correct tiles. Both integer and floating-point datatypes are checked.
+
 // -----
 
 // CHECK-LABEL: zero_za_b
 func.func @zero_za_b() {
   // CHECK-DAG: %[[TILE_ID:.*]] = arith.constant 0 : i8
   // CHECK-DAG: %[[ZERO_MASK:.*]] = arith.constant 255 : i32
 
   // CHECK:      "arm_sme.intr.zero"(%[[ZERO_MASK]]) : (i32) -> ()
   // CHECK-NEXT: %[[ZERO_ZA0B:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID]] : i8 to vector<[16]x[16]xi8>
func.func @zero_za_h() {
   // CHECK-DAG: %[[ZERO_MASK_ZA0H:.*]] = arith.constant 85 : i32
   // CHECK-DAG: %[[ZERO_MASK_ZA1H:.*]] = arith.constant 170 : i32
 
   // CHECK:      "arm_sme.intr.zero"(%[[ZERO_MASK_ZA0H]]) : (i32) -> ()
   // CHECK-NEXT: %[[ZERO_ZA0H:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA0H]] : i16 to vector<[8]x[8]xi16>
   %zero_za0h = arm_sme.zero : vector<[8]x[8]xi16>
   "prevent.dce"(%zero_za0h) : (vector<[8]x[8]xi16>) -> ()
   // CHECK:      "arm_sme.intr.zero"(%[[ZERO_MASK_ZA1H]]) : (i32) -> ()
-  // CHECK-NEXT: %[[ZERO_ZA1H:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA1H]] : i16 to vector<[8]x[8]xi16>
-  %zero_za1h = arm_sme.zero : vector<[8]x[8]xi16>
-  "prevent.dce"(%zero_za1h) : (vector<[8]x[8]xi16>) -> ()
+  // CHECK-NEXT: %[[ZERO_ZA1H:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA1H]] : i16 to vector<[8]x[8]xf16>
+  %zero_za1h = arm_sme.zero : vector<[8]x[8]xf16>
+  "prevent.dce"(%zero_za1h) : (vector<[8]x[8]xf16>) -> ()
   return
 }
 
 // -----
 
 // CHECK-LABEL: zero_za_s
 func.func @zero_za_s() {
   // CHECK-DAG: %[[TILE_ID_ZA0S:.*]] = arith.constant 0 : i32
func.func @zero_za_s() {
   // CHECK-NEXT: %[[ZERO_ZA1S:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA1S]] : i32 to vector<[4]x[4]xi32>
   %zero_za1s = arm_sme.zero : vector<[4]x[4]xi32>
   "prevent.dce"(%zero_za1s) : (vector<[4]x[4]xi32>) -> ()
   // CHECK:      "arm_sme.intr.zero"(%[[ZERO_MASK_ZA2S]]) : (i32) -> ()
   // CHECK-NEXT: %[[ZERO_ZA2S:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA2S]] : i32 to vector<[4]x[4]xi32>
   %zero_za2s = arm_sme.zero : vector<[4]x[4]xi32>
   "prevent.dce"(%zero_za2s) : (vector<[4]x[4]xi32>) -> ()
   // CHECK:     "arm_sme.intr.zero"(%[[ZERO_MASK_ZA3S]]) : (i32) -> ()
-  // CHECK-NEXT: %[[ZERO_ZA3S:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA3S]] : i32 to vector<[4]x[4]xi32>
-  %zero_za3s = arm_sme.zero : vector<[4]x[4]xi32>
-  "prevent.dce"(%zero_za3s) : (vector<[4]x[4]xi32>) -> ()
+  // CHECK-NEXT: %[[ZERO_ZA3S:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA3S]] : i32 to vector<[4]x[4]xf32>
+  %zero_za3s = arm_sme.zero : vector<[4]x[4]xf32>
+  "prevent.dce"(%zero_za3s) : (vector<[4]x[4]xf32>) -> ()
   return
 }
 
 // -----
 
 // CHECK-LABEL: zero_za_d
 func.func @zero_za_d() {
   // CHECK-DAG: %[[TILE_ID_ZA0D:.*]] = arith.constant 0 : i64
func.func @zero_za_d() {
   // CHECK-NEXT: %[[ZERO_ZA5D:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA5D]] : i64 to vector<[2]x[2]xi64>
   %zero_za5d = arm_sme.zero : vector<[2]x[2]xi64>
   "prevent.dce"(%zero_za5d) : (vector<[2]x[2]xi64>) -> ()
   // CHECK:     "arm_sme.intr.zero"(%[[ZERO_MASK_ZA6D]]) : (i32) -> ()
   // CHECK-NEXT: %[[ZERO_ZA6D:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA6D]] : i64 to vector<[2]x[2]xi64>
   %zero_za6d = arm_sme.zero : vector<[2]x[2]xi64>
   "prevent.dce"(%zero_za6d) : (vector<[2]x[2]xi64>) -> ()
   // CHECK:     "arm_sme.intr.zero"(%[[ZERO_MASK_ZA7D]]) : (i32) -> ()
-  // CHECK-NEXT: %[[ZERO_ZA7D:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA7D]] : i64 to vector<[2]x[2]xi64>
-  %zero_za7d = arm_sme.zero : vector<[2]x[2]xi64>
-  "prevent.dce"(%zero_za7d) : (vector<[2]x[2]xi64>) -> ()
+  // CHECK-NEXT: %[[ZERO_ZA7D:.*]] = arm_sme.cast_tile_to_vector %[[TILE_ID_ZA7D]] : i64 to vector<[2]x[2]xf64>
+  %zero_za7d = arm_sme.zero : vector<[2]x[2]xf64>
+  "prevent.dce"(%zero_za7d) : (vector<[2]x[2]xf64>) -> ()
   return
 }