[PATCH] D80411: [mlir][spirv] Lower allocation/deallocations of workgroup memory.

[Lei Zhang via Phabricator via llvm-commits]

antiagainst accepted this revision.
antiagainst added a comment.
This revision is now accepted and ready to land.

Cool, thanks Mahesh! Overall looks good to me.  Approval conditional on fixing the workgroup layout annotation issue.



================
Comment at: mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h:46
+  /// type. Note that it doesnt account for whether the type is legal for a
+  /// SPIR-V target (described by spirv::TargetEnvAttr). Returns null on
+  /// failure.
----------------
Returns None


================
Comment at: mlir/include/mlir/Dialect/SPIRV/SPIRVLowering.h:48
+  /// failure.
+  static Optional<int64_t> getConvertedTypeNumBytes(Type t);
+
----------------
Super nit: can we just leave out the `t`? It does not provide more information here.


================
Comment at: mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp:221
+
+    // Convert the alloc to a 1-D array of the same number of bytes. The element
+    // type might get converted to a different bitwidth type based on the target
----------------
I feel we should place this logic inside the type converter itself. This is basically getting the SPIR-V "equivalent" lowered memref type for the input memref. Letting each pattern figuring out this is not the proper separation of concerns. Fine for now given this is only used by allocation; but can we put a TODO here to remind ourselves?


================
Comment at: mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp:254
+
+    // Insert spv.global_variable for this allocation.
+    Operation *parent =
----------------
spv.globalVariable


================
Comment at: mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp:255
+    // Insert spv.global_variable for this allocation.
+    Operation *parent =
+        SymbolTable::getNearestSymbolTable(operation.getParentOp());
----------------
Could you document that this pattern requires the pass to work on module instead of function given it injects stuff to the model? 


================
Comment at: mlir/lib/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.cpp:265
+      rewriter.setInsertionPointToStart(&entryBlock);
+      auto varOps = entryBlock.getOps<spirv::GlobalVariableOp>();
+      std::string varName =
----------------
Do we need to unique the symbol manually here? I remember there is infrastructure support for it; but cannot recall off the top of my head.


================
Comment at: mlir/test/Conversion/GPUToSPIRV/load-store.mlir:61
       %13 = addi %arg4, %3 : index
+      // CHECK: %[[ZERO:.*]] = spv.constant 0 : i32
+      // CHECK: %[[OFFSET1_0:.*]] = spv.constant 0 : i32
----------------
Might just want to add `-canonicalize` to simplify here? I know it's kinda tiding two passes together but in general canonicalization is not changing too frequently and we can have more readable tests here. But up to you. :)


================
Comment at: mlir/test/Conversion/StandardToSPIRV/alloc.mlir:22
+}
+//     CHECK: spv.globalVariable @[[VAR:.+]] : !spv.ptr<!spv.struct<!spv.array<20 x f32, stride=4> [0]>, Workgroup>
+//     CHECK: func @alloc_dealloc_workgroup_mem
----------------
Ah, this should not have layout annotation at all. Variables in Workgroup should not be explicitly laid out. The layout handling is problematic in the type converter. Could you turn off layout annotation for Workgroup memrefs in TypeConverter? Just do not pass in offsets when handling Workgroup in `convertMemrefType`.


================
Comment at: mlir/test/Conversion/StandardToSPIRV/alloc.mlir:35
+
+// TODO: Uncomment this test when the extension handling correctly
+// converts an i16 type to i32 type and handles the load/stores
----------------
You mean fix? This test is not commented out at the moment.


================
Comment at: mlir/test/Conversion/StandardToSPIRV/alloc.mlir:108
+    // expected-error @+2 {{unhandled deallocation type}}
+    // expected-error @+1 {{'std.dealloc' op operand #0 must be memref of any type values, but got '!spv.ptr<!spv.struct<!spv.rtarray<f32, stride=4> [0]>, Workgroup>'}}
+    dealloc %arg0 : memref<4x?xf32, 3>
----------------
We can trim of the `but got ...` part given we are not particularly interested in verifying the detailed error here. Similarly for the following case.


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D80411/new/

https://reviews.llvm.org/D80411