[MLIR][Linalg] Generate unique LibraryCallName for LinalgOps.

[kaitingwang]

Summary

When lowering LinalgToStandard for named UnaryFn/BinaryFn ops, ensure the fun name appears in the generated library name. Further, for linalg.copy to/from different address spaces, ensure the to/from address spaces are appended onto the library name for uniqueness. This fixes the lowering error with the linalg.copy testcase shown in this patch.

[ftynse]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1806

There's no guarantee this cast succeeds. You likely want a dyn_cast and propagate failure.

[kaitingwang]

Avoid performing cast here but instead, use getMemorySpaceAsInt() API.

[ftynse]

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

1806

getMemorySpaceAsInt is deprecated, and there is no guarantee it succeeds either. Furthermore, this is introducing a subtle behavior change where no suffix will be emitted for the address space 0 (this sounds desirable for the default memory space, but needs a test).

Please do a dyn_cast and just return failure() if it fails...

[kaitingwang]

Perform dyn_cast and upon failure, return failure().

[kaitingwang]

Thanks for the review! @ftynse

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

if (auto memref = t.dyn_cast<MemRefType>()) {
     ss << "view";
     for (auto size : memref.getShape())
       if (size < 0)
         ss << "sx";
       else
         ss << size << "x";
     if (failed(appendMangledType(ss, memref.getElementType())))
       return failure();
+    if (auto as = memref.getMemorySpace()) {
+      if (auto attr = as.dyn_cast<IntegerAttr>())
+        ss << "as" << attr.getInt();
+      else
+        return failure();
+    }
     return success();
   }
   if (auto vec = t.dyn_cast<VectorType>()) {
     ss << "vector";
     llvm::interleave(
         vec.getShape(), [&](int64_t i) { ss << i; }, [&]() { ss << "x"; });
     if (failed(appendMangledType(ss, vec.getElementType())))
       return failure();
     return success();
   } else if (t.isSignlessIntOrIndexOrFloat()) {
     ss << t;
     return success();
   }
   return failure();
 }
 
 std::string mlir::linalg::generateLibraryCallName(Operation *op) {
   assert(isa<LinalgOp>(op));
   std::string name(op->getName().getStringRef().str());
+  std::string fun = "";
+  for (NamedAttribute kv : op->getAttrs()) {
+    if (UnaryFnAttr ufa = kv.getValue().dyn_cast<UnaryFnAttr>()) {
+      fun = stringifyEnum(ufa.getValue()).str() + "_";
+    } else if (BinaryFnAttr bfa = kv.getValue().dyn_cast<BinaryFnAttr>()) {
+      fun = stringifyEnum(bfa.getValue()).str() + "_";
+    }
+  }
   name.reserve(128);
   std::replace(name.begin(), name.end(), '.', '_');
   llvm::raw_string_ostream ss(name);
-  ss << "_";
+  ss << "_" << fun;
   for (Type t : op->getOperandTypes()) {
     if (failed(appendMangledType(ss, t)))
       return std::string();
     ss << "_";
   }
   std::string res = ss.str();
   res.pop_back();
   return res;

mlir/test/Dialect/Linalg/library-calls.mlir

-// RUN: mlir-opt %s -convert-linalg-to-std | FileCheck %s
+// RUN: mlir-opt %s -convert-linalg-to-std -split-input-file | FileCheck %s
 
 func.func private @printMemrefF32(memref<*xf32>)
 
 // CHECK:  func private @linalg_fill_f32_viewsxsxf32(f32, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface}
 // CHECK:  func private @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32(memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) attributes {llvm.emit_c_interface}
 
 func.func @matmul(%A: memref<?x?xf32>, %B: memref<?x?xf32>) -> (memref<?x?xf32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %f0 = arith.constant 0.0 : f32
   %x = memref.dim %A, %c0 : memref<?x?xf32>
   %y = memref.dim %B, %c1 : memref<?x?xf32>
   %C = memref.alloc(%x, %y) : memref<?x?xf32>
 
   // CHECK: call @linalg_fill_f32_viewsxsxf32({{.*}}) : (f32, memref<?x?xf32, {{.*}}>)
   linalg.fill ins(%f0 : f32) outs(%C : memref<?x?xf32>)
 
   // CHECK:  call @linalg_matmul_viewsxsxf32_viewsxsxf32_viewsxsxf32({{.*}}) : (memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>, memref<?x?xf32, {{.*}}>) -> ()
   linalg.matmul ins(%A, %B: memref<?x?xf32>, memref<?x?xf32>)
                 outs(%C: memref<?x?xf32>)
   return %C : memref<?x?xf32>
 }
 
+// -----
+
+#accesses = [
+        affine_map<(d0) -> (d0)>,
+        affine_map<(d0) -> (d0)>,
+        affine_map<(d0) -> (d0)>
+    ]
+#trait = {
+    doc = "...",
+    indexing_maps = #accesses,
+    library_call = "test",
+    iterator_types = ["parallel"]
+}
+
+// CHECK: func.func private @linalg_copy_view32xf16as1_view32xf16as6(memref<32xf16, strided<[?], offset: ?>, 1>, memref<32xf16, strided<[?], offset: ?>, 6>) attributes {llvm.emit_c_interface}
+// CHECK: func.func private @linalg_copy_view32xf16as6_view32xf16as1(memref<32xf16, strided<[?], offset: ?>, 6>, memref<32xf16, strided<[?], offset: ?>, 1>) attributes {llvm.emit_c_interface}
+
+module {
+  func.func @helper(%arg7: memref<32xf16, 1>, %arg8: memref<32xf16, 1>, %arg9: memref<32xf16, 1>) {
+    %localA = memref.alloca() : memref<32xf16, 6>
+    %localB = memref.alloca() : memref<32xf16, 6>
+    %localOut = memref.alloca() : memref<32xf16, 6>
+    linalg.copy ins(%arg8 : memref<32xf16, 1>) outs(%localA : memref<32xf16, 6>)
+    linalg.copy ins(%arg9 : memref<32xf16, 1>) outs(%localB : memref<32xf16, 6>)
+
+    linalg.generic #trait
+    ins(%localA, %localB : memref<32xf16, 6>, memref<32xf16, 6>)
+    outs(%localOut : memref<32xf16, 6>) {
+      ^bb0(%0: f16, %1: f16, %2: f16) :
+        %e = arith.addf %1, %0: f16
+        linalg.yield %e : f16
+    }
+
+    linalg.copy ins(%localOut : memref<32xf16, 6>) outs(%arg7 : memref<32xf16, 1>)
+    return
+  }
+}
+
+
+// -----
+
+// CHECK: func.func private @linalg_elemwise_unary_negf_view16x8xf32_view16x8xf32(memref<16x8xf32, strided<[?, ?], offset: ?>>, memref<16x8xf32, strided<[?, ?], offset: ?>>) attributes {llvm.emit_c_interface}
+// CHECK: func.func private @linalg_elemwise_unary_negf_view16xf32_view16xf32(memref<16xf32, strided<[?], offset: ?>>, memref<16xf32, strided<[?], offset: ?>>) attributes {llvm.emit_c_interface}
+
+func.func @test_neg(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C: memref<16xf32>, %D: memref<16xf32>) {
+  linalg.elemwise_unary {fun = #linalg.unary_fn<negf>}
+                              ins(%A: memref<16x8xf32>) outs(%B: memref<16x8xf32>)
+  linalg.elemwise_unary {fun = #linalg.unary_fn<negf>}
+                              ins(%C: memref<16xf32>) outs(%D: memref<16xf32>)
+  return
+}
+
+// -----
+
+// CHECK: func.func private @linalg_elemwise_unary_exp_view16x8xf32_view16x8xf32(memref<16x8xf32, strided<[?, ?], offset: ?>>, memref<16x8xf32, strided<[?, ?], offset: ?>>) attributes {llvm.emit_c_interface}
+// CHECK: func.func private @linalg_elemwise_unary_exp_view16xf32_view16xf32(memref<16xf32, strided<[?], offset: ?>>, memref<16xf32, strided<[?], offset: ?>>) attributes {llvm.emit_c_interface}
+
+func.func @test_exp(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C: memref<16xf32>, %D: memref<16xf32>) {
+  linalg.elemwise_unary {fun = #linalg.unary_fn<exp>}
+                              ins(%A: memref<16x8xf32>) outs(%B: memref<16x8xf32>)
+  linalg.elemwise_unary {fun = #linalg.unary_fn<exp>}
+                              ins(%C: memref<16xf32>) outs(%D: memref<16xf32>)
+  return
+}
+
+// -----
+
+// CHECK: func.func private @linalg_elemwise_binary_add_view16x8xf32_view16x8xf32_view16x8xf32(memref<16x8xf32, strided<[?, ?], offset: ?>>, memref<16x8xf32, strided<[?, ?], offset: ?>>, memref<16x8xf32, strided<[?, ?], offset: ?>>) attributes {llvm.emit_c_interface}
+// CHECK: func.func private @linalg_elemwise_binary_add_view16xf32_view16xf32_view16xf32(memref<16xf32, strided<[?], offset: ?>>, memref<16xf32, strided<[?], offset: ?>>, memref<16xf32, strided<[?], offset: ?>>) attributes {llvm.emit_c_interface}
+
+func.func @test_add(%A : memref<16x8xf32>, %B: memref<16x8xf32>, %C: memref<16x8xf32>, %D: memref<16xf32>, %E: memref<16xf32>, %F: memref<16xf32>) {
+  linalg.elemwise_binary {fun = #linalg.binary_fn<add>}
+                              ins(%A, %B: memref<16x8xf32>, memref<16x8xf32>) outs(%C: memref<16x8xf32>)
+  linalg.elemwise_binary {fun = #linalg.binary_fn<add>}
+                              ins(%D, %E: memref<16xf32>, memref<16xf32>) outs(%F: memref<16xf32>)
+  return
+}