[MLIR] Add readonly,writeonly,noalias to func arg attributes

[sjw36]

Summary

Includes conversion to corresponding LLVM attributes

[ftynse]

Why can't we keep using the attributes defined in the LLVM dialect?

[sjw36]

Simplified after rebase, using translation tables.

[sjw36]

@ftynse these attributes are specific to func parameters, and not necessarily to LLVM backends (e.g. SPIRV). They are useful for dataflow analysis and would prefer not to have to include LLVMIR everywhere.

[sjw36]

[MLIR][Func] Add readonly,writeonly,noalias to func arg attributes

Includes conversion to corresponding LLVM attributes

Differential Revision: https://reviews.llvm.org/D140357

[ftynse]

There isn't sufficient context to justify the change in the commit description, and the change itself doesn't seem to do anything else than lifting function argument attributes from LLVM IR to the Func dialect level. Yes, there is SPIR-V, but the change does not include the lowering to SPIR-V or even a description of how it can be done. Blindly copying LLVM IR is acceptable in the LLVM dialect, but not at the higher-level abstractions such as Func. It is unclear to me what these attributes would even mean in the full generality of what MLIR type system allows. For example, what is a writeonly tensor? How do two SPIR-V types noalias? How do even two memrefs noalias when the default conversion to LLVM produces two pointers that precisely do alias each other?

These are useful concepts to have at this level, but they should be modeled with the sufficient generality. We quite likely want to leverage, e.g., memref's multidimensional and structured nature to capture more fine-grained information than we can about two opaque pointers with noalias. Consider posting on Discourse to discuss how to model these properly.

In the meantime, I checked the implementation, and I don't see anything that would prevent one from using llvm.noalias and the likes as an argument attribute in a func.func and regardless of the type.

[sjw36]

Understood. This is as you said just lifting these attributes, a future change would use them and I will post it as an RFC to discuss. This change can wait for that, just trying to flush some trivial parts from my backlog.

As an example to answer some of your questions, take:

func.func @conv_relu(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*xf32> {
  ...
  return %result : tensor<*xf32>
}

Since tensors are PBV, we can tag the inputs as 'readonly' and the results as 'writeonly' as we know the function cannot read and will fully write the result.

func.func @conv_relu(%arg0 : tensor<*xf32> {func.readonly, func.noalias}, %arg1 : tensor<*xf32> {func.readonly, func.noalias}) -> tensor<*xf32> {func.writeonly, func.noalias} {
  ...
  return %result : tensor<*xf32>
}

Then bufferization will take care to keep the specified attributes when moving results to out params:

func.func @conv_relu(%arg0 : memref<*xf32> {func.readonly, func.noalias}, %arg1 : memref<*xf32> {func.readonly, func.noalias}, %arg2 : memref<*xf32> {func.writeonly, func.noalias}) {
  ...
  return
}

Then one can reason about how a function will interact with the memory references that are passed to it. Further, the func can be optimized more aggressively knowing that loads/stores are independent.

mlir/include/mlir/Dialect/Func/IR/FuncOps.td

let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "func"; }
 
     //===------------------------------------------------------------------===//
     // SymbolOpInterface Methods
     //===------------------------------------------------------------------===//
 
     bool isDeclaration() { return isExternal(); }
+
+    //===------------------------------------------------------------------===//
+    // FuncOp Methods
+    //===------------------------------------------------------------------===//
+
+    /// Arg attributes for access modes
+    static StringRef getReadOnlyAttrName() { return "func.readonly"; }
+    static StringRef getWriteOnlyAttrName() { return "func.writeonly"; }
+    static StringRef getNoAliasAttrName() { return "func.noalias"; }
   }];
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

def LLVM_Dialect : Dialect {
   let hasOperationAttrVerify = 1;
 
   let extraClassDeclaration = [{
     /// Name of the data layout attributes.
     static StringRef getDataLayoutAttrName() { return "llvm.data_layout"; }
     static StringRef getAlignAttrName() { return "llvm.align"; }
     static StringRef getNoAliasAttrName() { return "llvm.noalias"; }
     static StringRef getReadonlyAttrName() { return "llvm.readonly"; }
+    static StringRef getWriteonlyAttrName() { return "llvm.writeonly"; }
     static StringRef getNoAliasScopesAttrName() { return "noalias_scopes"; }
     static StringRef getAliasScopesAttrName() { return "alias_scopes"; }
     static StringRef getLoopAttrName() { return "llvm.loop"; }
     static StringRef getParallelAccessAttrName() { return "parallel_access"; }
     static StringRef getLoopOptionsAttrName() { return "options"; }
     static StringRef getAccessGroupsAttrName() { return "access_groups"; }
     static StringRef getStructAttrsAttrName() { return "llvm.struct_attrs"; }
     static StringRef getByValAttrName() { return "llvm.byval"; }

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

 //===- FuncToLLVM.cpp - Func to LLVM dialect conversion -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
if (ArrayAttr argAttrDicts = funcOp.getAllArgAttrs()) {
           } else if (attr.getName().getValue() ==
                      LLVM::LLVMDialect::getStructRetAttrName()) {
             convertedAttrs.push_back(rewriter.getNamedAttr(
                 LLVM::LLVMDialect::getStructRetAttrName(), convert(attr)));
           } else if (attr.getName().getValue() ==
                      LLVM::LLVMDialect::getInAllocaAttrName()) {
             convertedAttrs.push_back(rewriter.getNamedAttr(
                 LLVM::LLVMDialect::getInAllocaAttrName(), convert(attr)));
+          } else if (attr.getName().getValue() ==
+                     func::FuncOp::getReadOnlyAttrName()) {
+            assert(attr.getValue().isa<UnitAttr>());
+            convertedAttrs.push_back(rewriter.getNamedAttr(
+                LLVM::LLVMDialect::getReadonlyAttrName(), rewriter.getUnitAttr()));
+          } else if (attr.getName().getValue() ==
+                     func::FuncOp::getWriteOnlyAttrName()) {
+            assert(attr.getValue().isa<UnitAttr>());
+            convertedAttrs.push_back(rewriter.getNamedAttr(
+                LLVM::LLVMDialect::getWriteonlyAttrName(), rewriter.getUnitAttr()));
+          } else if (attr.getName().getValue() ==
+                     func::FuncOp::getNoAliasAttrName()) {
+            assert(attr.getValue().isa<UnitAttr>());
+            convertedAttrs.push_back(rewriter.getNamedAttr(
+                LLVM::LLVMDialect::getNoAliasAttrName(), rewriter.getUnitAttr()));
           } else {
             convertedAttrs.push_back(attr);
           }
         }
         auto mapping = result.getInputMapping(i);
         assert(mapping && "unexpected deletion of function argument");
         for (size_t j = 0; j < mapping->size; ++j)
           newArgAttrs[mapping->inputNo + j] =

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

 //===- ModuleTranslation.cpp - MLIR to LLVM conversion --------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
for (auto [mlirArgTy, llvmArg] :
       }
       if (auto attr = function.getArgAttrOfType<UnitAttr>(
               argIdx, LLVMDialect::getReadonlyAttrName())) {
         if (!mlirArgTy.isa<LLVM::LLVMPointerType>())
           return function.emitError(
               "llvm.readonly attribute attached to LLVM non-pointer argument");
         llvmArg.addAttr(llvm::Attribute::AttrKind::ReadOnly);
       }
+      if (auto attr = function.getArgAttrOfType<UnitAttr>(
+              argIdx, LLVMDialect::getWriteonlyAttrName())) {
+        if (!mlirArgTy.isa<LLVM::LLVMPointerType>())
+          return function.emitError(
+              "llvm.readonly attribute attached to LLVM non-pointer argument");
+        llvmArg.addAttr(llvm::Attribute::AttrKind::WriteOnly);
+      }
+
 
       if (auto attr = function.getArgAttrOfType<IntegerAttr>(
               argIdx, LLVMDialect::getAlignAttrName())) {
         // NB: Attribute already verified to be int, so check if we can indeed
         // attach the attribute to this argument, based on its type.
         if (!mlirArgTy.isa<LLVM::LLVMPointerType>())
           return function.emitError(
               "llvm.align attribute attached to LLVM non-pointer argument");

mlir/test/Conversion/FuncToLLVM/func-memref.mlir

 // RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-arith-to-llvm),convert-func-to-llvm,reconcile-unrealized-casts)" -split-input-file %s | FileCheck %s
 // RUN: mlir-opt -pass-pipeline="builtin.module(func.func(convert-arith-to-llvm),convert-func-to-llvm{use-bare-ptr-memref-call-conv=1},reconcile-unrealized-casts)" -split-input-file %s | FileCheck %s --check-prefix=BAREPTR
 
 // BAREPTR-LABEL: func @check_noalias
 // BAREPTR-SAME: %{{.*}}: !llvm.ptr<f32> {llvm.noalias}, %{{.*}}: !llvm.ptr<f32> {llvm.noalias}
-func.func @check_noalias(%static : memref<2xf32> {llvm.noalias}, %other : memref<2xf32> {llvm.noalias}) {
+func.func @check_noalias(%static : memref<2xf32> {func.noalias}, %other : memref<2xf32> {llvm.noalias}) {
+    return
+}
+
+// -----
+// BAREPTR-LABEL: func @check_readonly_writeonly
+// BAREPTR-SAME: %{{.*}}: !llvm.ptr<f32> {llvm.readonly}, %{{.*}}: !llvm.ptr<f32> {llvm.writeonly}
+func.func @check_readonly_writeonly(%static : memref<2xf32> {func.readonly}, %other : memref<2xf32> {func.writeonly}) {
     return
 }
 
 // -----
 
 // CHECK-LABEL: func @check_strided_memref_arguments(
 // CHECK-COUNT-2: !llvm.ptr<f32>
 // CHECK-COUNT-5: i64

mlir/test/Target/LLVMIR/llvmir.mlir

 
 // Function attributes: readnone
 
 // CHECK: declare void @readnone_function() #[[ATTR:[0-9]+]]
 // CHECK: attributes #[[ATTR]] = { memory(none) }
 llvm.func @readnone_function() attributes {llvm.readnone}
 
 // -----
-// CHECK: declare void @readonly_function([[PTR:.+]] readonly)
-llvm.func @readonly_function(%arg0: !llvm.ptr<f32> {llvm.readonly})
+// CHECK: declare void @readonly_writeonly_function([[PTR:.+]] readonly, [[PTR2:.+]] writeonly)
+llvm.func @readonly_writeonly_function(%arg0: !llvm.ptr<f32> {llvm.readonly}, %arg1: !llvm.ptr<f32> {llvm.writeonly})